How to use the wntr.sim.models.utils.Definition function in wntr

if index_over is None:
            index_over = wn.pipe_name_list + wn.valve_name_list

        for link_name in index_over:
            link = wn.get_link(link_name)
            value = 8.0 * link.minor_loss / (9.81 * math.pi**2 * link.diameter**4)
            if link_name in m.minor_loss:
                m.minor_loss[link_name].value = value
            else:
                m.minor_loss[link_name] = aml.Param(value)

            updater.add(link, 'minor_loss', minor_loss_param.update)
            updater.add(link, 'diameter', minor_loss_param.update)


class tcv_resistance_param(Definition):
    @classmethod
    def build(cls, m, wn, updater, index_over=None):
        """
        Add a tcv resistance parameter to the model

        Parameters
        ----------
        m: wntr.aml.aml.aml.Model
        wn: wntr.network.model.WaterNetworkModel
        updater: ModelUpdater
        index_over: list of str
            list of tcv names
        """
        if not hasattr(m, 'tcv_resistance'):
            m.tcv_resistance = aml.ParamDict()

import logging
from wntr.sim import aml
import wntr.network
import warnings
from wntr.utils.polynomial_interpolation import cubic_spline
from wntr.network import LinkStatus
from wntr.sim.models.utils import ModelUpdater, Definition

logger = logging.getLogger(__name__)


class mass_balance_constraint(Definition):
    @classmethod
    def build(cls, m, wn, updater, index_over=None):
        """
        Adds a mass balance to the model for the specified junctions.

        Parameters
        ----------
        m: wntr.aml.aml.aml.Model
        wn: wntr.network.model.WaterNetworkModel
        updater: ModelUpdater
        index_over: list of str
            list of junction names; default is all junctions in wn
        """
        if not hasattr(m, 'mass_balance'):
            m.mass_balance = aml.ConstraintDict()

c2 = m.pdd_poly2_coeffs_c[node_name]
                d2 = m.pdd_poly2_coeffs_d[node_name]
                con = aml.ConditionalExpression()
                con.add_condition(aml.inequality(body=h - elev - pmin, ub=0), d - d_expected*slope*(h-elev-pmin))
                con.add_condition(aml.inequality(body=h - elev - pmin - delta, ub=0), d - d_expected*(a1*(h-elev)**3 + b1*(h-elev)**2 + c1*(h-elev) + d1))
                con.add_condition(aml.inequality(body=h - elev - pnom + delta, ub=0), d - d_expected*((h-elev-pmin)/(pnom-pmin))**0.5)
                con.add_condition(aml.inequality(body=h - elev - pnom, ub=0), d - d_expected*(a2*(h-elev)**3 + b2*(h-elev)**2 + c2*(h-elev) + d2))
                con.add_final_expr(d - d_expected*(slope*(h - elev - pnom) + 1.0))
                con = aml.Constraint(con)

                m.pdd[node_name] = con

            updater.add(node, '_is_isolated', pdd_constraint.update)


class head_pump_headloss_constraint(Definition):
    @classmethod
    def build(cls, m, wn, updater, index_over=None):
        """
        Adds a headloss constraint to the model for the head curve pumps.

        Parameters
        ----------
        m: wntr.aml.aml.aml.Model
        wn: wntr.network.model.WaterNetworkModel
        updater: ModelUpdater
        index_over: list of str
            list of HeadPump names; default is all HeadPumps in wn
        """
        if not hasattr(m, 'head_pump_headloss'):
            m.head_pump_headloss = aml.ConstraintDict()

index_over = wn.pipe_name_list

        for link_name in index_over:
            link = wn.get_link(link_name)
            value = m.hw_k * link.roughness**(-1.852) * link.diameter**(-4.871) * link.length
            if link_name in m.hw_resistance:
                m.hw_resistance[link_name].value = value
            else:
                m.hw_resistance[link_name] = aml.Param(value)

            updater.add(link, 'roughness', hw_resistance_param.update)
            updater.add(link, 'diameter', hw_resistance_param.update)
            updater.add(link, 'length', hw_resistance_param.update)


class minor_loss_param(Definition):
    @classmethod
    def build(cls, m, wn, updater, index_over=None):
        """
        Add a minor loss coefficient parameter to the model

        Parameters
        ----------
        m: wntr.aml.aml.aml.Model
        wn: wntr.network.model.WaterNetworkModel
        updater: ModelUpdater
        index_over: list of str
            list of pipe names
        """
        if not hasattr(m, 'minor_loss'):
            m.minor_loss = aml.ParamDict()

m.pdd_poly2_coeffs_d[node_name].value = d2
            else:
                m.pdd_poly1_coeffs_a[node_name] = aml.Param(a1)
                m.pdd_poly1_coeffs_b[node_name] = aml.Param(b1)
                m.pdd_poly1_coeffs_c[node_name] = aml.Param(c1)
                m.pdd_poly1_coeffs_d[node_name] = aml.Param(d1)
                m.pdd_poly2_coeffs_a[node_name] = aml.Param(a2)
                m.pdd_poly2_coeffs_b[node_name] = aml.Param(b2)
                m.pdd_poly2_coeffs_c[node_name] = aml.Param(c2)
                m.pdd_poly2_coeffs_d[node_name] = aml.Param(d2)

            updater.add(node, 'minimum_pressure', pdd_poly_coeffs_param.update)
            updater.add(node, 'nominal_pressure', pdd_poly_coeffs_param.update)


class leak_poly_coeffs_param(Definition):
    @classmethod
    def build(cls, m, wn, updater, index_over=None):
        """
        Add parameters to the model for leak smoothing polynomial coefficients

        Parameters
        ----------
        m: wntr.aml.Model
        wn: wntr.network.model.WaterNetworkModel
        updater: ModelUpdater
        index_over: list of str
            list of junction names
        """
        if not hasattr(m, 'leak_poly_coeffs_a'):
            m.leak_poly_coeffs_a = aml.ParamDict()
            m.leak_poly_coeffs_b = aml.ParamDict()

m.leak_area = aml.ParamDict()

        if index_over is None:
            index_over = wn.junction_name_list + wn.tank_name_list

        for node_name in index_over:
            node = wn.get_node(node_name)
            if node_name in m.leak_area:
                m.leak_area[node_name].value = node.leak_area
            else:
                m.leak_area[node_name] = aml.Param(node.leak_area)

            updater.add(node, 'leak_area', leak_area_param.update)


class pdd_poly_coeffs_param(Definition):
    @classmethod
    def build(cls, m, wn, updater, index_over=None):
        """
        Add parameters to the model for pdd smoothing polynomial coefficients

        Parameters
        ----------
        m: wntr.aml.aml.aml.Model
        wn: wntr.network.model.WaterNetworkModel
        updater: ModelUpdater
        index_over: list of str
            list of junction names
        """
        if not hasattr(m, 'pdd_poly1_coeffs_a'):
            m.pdd_poly1_coeffs_a = aml.ParamDict()
            m.pdd_poly1_coeffs_b = aml.ParamDict()

m.leak_coeff = aml.ParamDict()

        if index_over is None:
            index_over = wn.junction_name_list + wn.tank_name_list

        for node_name in index_over:
            node = wn.get_node(node_name)
            if node_name in m.leak_coeff:
                m.leak_coeff[node_name].value = node.leak_discharge_coeff
            else:
                m.leak_coeff[node_name] = aml.Param(node.leak_discharge_coeff)

            updater.add(node, 'leak_discharge_coeff', leak_coeff_param.update)


class leak_area_param(Definition):
    @classmethod
    def build(cls, m, wn, updater, index_over=None):
        """
        Add a leak discharge coefficient parameter to the model

        Parameters
        ----------
        m: wntr.aml.Model
        wn: wntr.network.model.WaterNetworkModel
        updater: ModelUpdater
        index_over: list of str
            list of junction/tank names
        """
        if not hasattr(m, 'leak_area'):
            m.leak_area = aml.ParamDict()

m.pnom = aml.ParamDict()

        if index_over is None:
            index_over = wn.junction_name_list

        for node_name in index_over:
            node = wn.get_node(node_name)
            if node_name in m.pnom:
                m.pnom[node_name].value = node.nominal_pressure
            else:
                m.pnom[node_name] = aml.Param(node.nominal_pressure)

            updater.add(node, 'nominal_pressure', pnom_param.update)


class leak_coeff_param(Definition):
    @classmethod
    def build(cls, m, wn, updater, index_over=None):
        """
        Add a leak discharge coefficient parameter to the model

        Parameters
        ----------
        m: wntr.aml.Model
        wn: wntr.network.model.WaterNetworkModel
        updater: ModelUpdater
        index_over: list of str
            list of junction/tank names
        """
        if not hasattr(m, 'leak_coeff'):
            m.leak_coeff = aml.ParamDict()

if index_over is None:
            index_over = wn.power_pump_name_list

        for link_name in index_over:
            link = wn.get_link(link_name)
            value = link.power
            if link_name in m.pump_power:
                m.pump_power[link_name].value = value
            else:
                m.pump_power[link_name] = aml.Param(value)

            updater.add(link, 'power', pump_power_param.update)


class valve_setting_param(Definition):
    @classmethod
    def build(cls, m, wn, updater, index_over=None):
        """
        Add a valve setting parameter to the model

        Parameters
        ----------
        m: wntr.aml.aml.aml.Model
        wn: wntr.network.model.WaterNetworkModel
        updater: ModelUpdater
        index_over: list of str
            list of valve names
        """
        if not hasattr(m, 'valve_setting'):
            m.valve_setting = aml.ParamDict()

c = m.hw_c
                d = m.hw_d

                con = aml.ConditionalExpression()
                con.add_condition(aml.inequality(body=aml.abs(f), ub=m.hw_q1), -k*m.hw_m*f - aml.sign(f)*minor_k*f**m.hw_minor_exp + start_h - end_h)
                con.add_condition(aml.inequality(body=aml.abs(f), ub=m.hw_q2), -k*(a*f**3 + aml.sign(f)*b*f**2 + c*f + aml.sign(f)*d) - aml.sign(f)*minor_k*f**m.hw_minor_exp + start_h - end_h)
                con.add_final_expr(-aml.sign(f)*k*aml.abs(f)**m.hw_exp - aml.sign(f)*minor_k*f**m.hw_minor_exp + start_h - end_h)
                con = aml.Constraint(con)

            m.piecewise_hazen_williams_headloss[link_name] = con

            updater.add(link, 'status', piecewise_hazen_williams_headloss_constraint.update)
            updater.add(link, '_is_isolated', piecewise_hazen_williams_headloss_constraint.update)


class approx_hazen_williams_headloss_constraint(Definition):
    @classmethod
    def build(cls, m, wn, updater, index_over=None):
        """
        Adds a mass balance to the model for the specified junctions.

        Parameters
        ----------
        m: wntr.aml.aml.aml.Model
        wn: wntr.network.model.WaterNetworkModel
        updater: ModelUpdater
        index_over: list of str
            list of pipe names; default is all pipes in wn
        """
        if not hasattr(m, 'approx_hazen_williams_headloss'):
            m.approx_hazen_williams_headloss = aml.ConstraintDict()