Asm2d Phosphorus

def _monod(s, k):
    """Monod saturation term: s / (k + s), safe for s near zero."""
    return s / (k + s) if (k + s) > 0.0 else 0.0

def _asm1_process_rates(y):
    """Compute the 8 ASM1 biological process rates from state vector y.

    Returns array of 8 rates (rho_1 .. rho_8).
    """
    S_I, S_S, X_I, X_S, X_BH, X_BA, X_P = y[0], y[1], y[2], y[3], y[4], y[5], y[6]
    S_O, S_NO, S_NH, S_ND, X_ND, S_ALK = y[7], y[8], y[9], y[10], y[11], y[12]

    mon_ss = _monod(S_S, _K_S)
    mon_o_h = _monod(S_O, _K_OH)
    inh_o_h = _K_OH / (_K_OH + S_O) if (_K_OH + S_O) > 0.0 else 0.0
    mon_no = _monod(S_NO, _K_NO)
    mon_nh = _monod(S_NH, _K_NH)
    mon_o_a = _monod(S_O, _K_OA)

    # Process 1: aerobic growth of heterotrophs
    rho1 = _MU_H * mon_ss * mon_o_h * X_BH

    # Process 2: anoxic growth of heterotrophs
    rho2 = _MU_H * mon_ss * inh_o_h * mon_no * _ETA_G * X_BH

    # Process 3: aerobic growth of autotrophs
    rho3 = _MU_A * mon_nh * mon_o_a * X_BA

    # Process 4: decay of heterotrophs
    rho4 = _B_H * X_BH

    # Process 5: decay of autotrophs
    rho5 = _B_A * X_BA

    # Process 6: ammonification of soluble organic nitrogen
    rho6 = _K_A_AMMON * S_ND * X_BH

    # Process 7: hydrolysis of slowly biodegradable substrate
    xs_xbh_ratio = (X_S / X_BH) if X_BH > 1e-12 else 0.0
    mon_hyd = xs_xbh_ratio / (_K_X + xs_xbh_ratio) if (_K_X + xs_xbh_ratio) > 0.0 else 0.0
    hyd_switch = mon_o_h + _ETA_H * inh_o_h * mon_no
    rho7 = _K_H * mon_hyd * hyd_switch * X_BH

    # Process 8: hydrolysis of organic nitrogen
    xnd_xs_ratio = (X_ND / X_S) if X_S > 1e-12 else 0.0
    rho8 = rho7 * xnd_xs_ratio

    return np.array([rho1, rho2, rho3, rho4, rho5, rho6, rho7, rho8])

def _asm1_reaction_vector(rho):
    """Petersen matrix: convert 8 process rates to 13 state derivatives."""
    rho1, rho2, rho3, rho4, rho5, rho6, rho7, rho8 = rho

    dy = np.zeros(13)

    # dS_I/dt   = 0 (inert, only dilution)
    # dS_S/dt   = -(1/Y_H)*rho1 - (1/Y_H)*rho2 + rho7
    dy[1] = -(1.0 / _Y_H) * rho1 - (1.0 / _Y_H) * rho2 + rho7

    # dX_I/dt   = 0 (inert particulate, only dilution)
    # dX_S/dt   = (1-f_p)*rho4 + (1-f_p)*rho5 - rho7
    dy[3] = (1.0 - _F_P) * rho4 + (1.0 - _F_P) * rho5 - rho7

    # dX_BH/dt  = rho1 + rho2 - rho4
    dy[4] = rho1 + rho2 - rho4

    # dX_BA/dt  = rho3 - rho5
    dy[5] = rho3 - rho5

    # dX_P/dt   = f_p*rho4 + f_p*rho5
    dy[6] = _F_P * rho4 + _F_P * rho5

    # dS_O/dt   = -((1-Y_H)/Y_H)*rho1 - ((4.57-Y_A)/Y_A)*rho3 + KLa*(S_O_sat - S_O)
    #             (aeration handled separately in the full RHS)
    dy[7] = -((1.0 - _Y_H) / _Y_H) * rho1 - ((4.57 - _Y_A) / _Y_A) * rho3

    # dS_NO/dt  = -((1-Y_H)/(2.86*Y_H))*rho2 + (1/Y_A)*rho3
    dy[8] = -((1.0 - _Y_H) / (2.86 * _Y_H)) * rho2 + (1.0 / _Y_A) * rho3

    # dS_NH/dt  = -i_XB*rho1 - i_XB*rho2 - (i_XB + 1/Y_A)*rho3 + rho6
    dy[9] = -_I_XB * rho1 - _I_XB * rho2 - (_I_XB + 1.0 / _Y_A) * rho3 + rho6

    # dS_ND/dt  = -rho6 + rho8
    dy[10] = -rho6 + rho8

    # dX_ND/dt  = (i_XB - f_p*i_XP)*rho4 + (i_XB - f_p*i_XP)*rho5 - rho8
    dy[11] = (_I_XB - _F_P * _I_XP) * rho4 + (_I_XB - _F_P * _I_XP) * rho5 - rho8

    # dS_ALK/dt = -(i_XB/14)*rho1 + ((1-Y_H)/(14*2.86*Y_H))*rho2
    #             - (i_XB/14 + 1/(7*Y_A))*rho3 + rho6/14
    dy[12] = (
        -(_I_XB / 14.0) * rho1
        + ((1.0 - _Y_H) / (14.0 * 2.86 * _Y_H)) * rho2
        - (_I_XB / 14.0 + 1.0 / (7.0 * _Y_A)) * rho3
        + rho6 / 14.0
    )

    return dy

def _nonneg_clamp(y, dy):
    """IWA-standard non-negativity enforcement at the RHS level.

    If a state is at (or below) zero and the derivative would push it
    further negative, clamp the derivative to zero.  This prevents
    physically impossible negative concentrations without modifying the
    solver.
    """
    for i in range(len(y)):
        if y[i] <= 0.0 and dy[i] < 0.0:
            dy[i] = 0.0
    return dy

def _rhs_asm2d(t, y):
    """ASM2d CSTR: ASM1 processes + PAO bio-P processes (19-state)."""
    y_safe = np.maximum(y, 0.0)

    # -- ASM1 base processes (on the first 13 states) --
    rho = _asm1_process_rates(y_safe[:13])
    r_asm1 = _asm1_reaction_vector(rho)

    # -- PAO-specific states --
    S_O = y_safe[7]
    S_NH = y_safe[9]
    S_A = y_safe[13]
    S_PO4 = y_safe[15]
    X_PAO = y_safe[16]
    X_PHA = y_safe[17]
    X_PP = y_safe[18]

    # safe ratios to avoid division by zero
    pha_pao = (X_PHA / X_PAO) if X_PAO > 1e-12 else 0.0
    pp_pao = (X_PP / X_PAO) if X_PAO > 1e-12 else 0.0

    # Process: anaerobic PHA storage
    rho_sto = (_Q_PHA
               * _monod(S_A, _K_A_PAO)
               * (pp_pao / (_K_PP + pp_pao) if (_K_PP + pp_pao) > 0.0 else 0.0)
               * X_PAO)

    # Process: aerobic growth of PAO
    rho_pao_aer = (_MU_PAO
                   * _monod(S_O, _K_O_PAO)
                   * _monod(S_NH, _K_NH4)
                   * _monod(S_PO4, _K_P)
                   * (pha_pao / (_K_PHA + pha_pao) if (_K_PHA + pha_pao) > 0.0 else 0.0)
                   * X_PAO)

    # Process: aerobic poly-P storage
    kmax_minus_pp = max(_K_MAX - pp_pao, 0.0)
    rho_pp_aer = (_Q_PP
                  * _monod(S_O, _K_O_PAO)
                  * _monod(S_PO4, _K_PS)
                  * (pha_pao / (_K_PHA + pha_pao) if (_K_PHA + pha_pao) > 0.0 else 0.0)
                  * (kmax_minus_pp / (_K_IPP + kmax_minus_pp) if (_K_IPP + kmax_minus_pp) > 0.0 else 0.0)
                  * X_PAO)

    # Lysis processes (first-order)
    rho_lys_pao = _B_PAO * X_PAO
    rho_lys_pp = _B_PP * X_PP
    rho_lys_pha = _B_PHA * X_PHA

    # -- Build 19-state derivative --
    D_h = _D / _H_PER_D
    dy = np.zeros(19)

    # ASM1 base reactions (converted from d^-1 to h^-1)
    for i in range(13):
        dy[i] = r_asm1[i] / _H_PER_D + D_h * (_Y_IN_ASM2D[i] - y_safe[i])

    # Aeration for S_O
    dy[7] += (_KLA / _H_PER_D) * (_S_O_SAT - y_safe[7])

    # S_A (index 13): consumed by PHA storage
    dy[13] = D_h * (_Y_IN_ASM2D[13] - y_safe[13]) - rho_sto / _H_PER_D

    # S_F (index 14): dilution only (fermentation not modeled explicitly here)
    dy[14] = D_h * (_Y_IN_ASM2D[14] - y_safe[14])

    # S_PO4 (index 15): released during PHA storage, consumed during PP storage and PAO growth
    dy[15] = (D_h * (_Y_IN_ASM2D[15] - y_safe[15])
              + (_Y_PO4_PHA * rho_sto - rho_pp_aer - _Y_PP_PAO * rho_pao_aer
                 + rho_lys_pp) / _H_PER_D)

    # X_PAO (index 16): grows aerobically, decays
    dy[16] = (D_h * (_Y_IN_ASM2D[16] - y_safe[16])
              + (rho_pao_aer - rho_lys_pao) / _H_PER_D)

    # X_PHA (index 17): stored anaerobically, consumed for PAO growth, lyses
    dy[17] = (D_h * (_Y_IN_ASM2D[17] - y_safe[17])
              + (rho_sto - _Y_PHA_PAO * rho_pao_aer - rho_lys_pha) / _H_PER_D)

    # X_PP (index 18): stored aerobically, released anaerobically, lyses
    dy[18] = (D_h * (_Y_IN_ASM2D[18] - y_safe[18])
              + (rho_pp_aer - _Y_PO4_PHA * rho_sto - rho_lys_pp) / _H_PER_D)

    return _nonneg_clamp(y, dy)