Asm1 Dynamic Storm

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 _storm_influent(t):
    """Time-varying influent for storm event.

    At t=6h, S_S_in spikes to 5x baseline; X_S_in and S_NH_in increase 3x.
    Exponential decay back to baseline with tau=4h.
    """
    y_in = _Y_IN_ASM1.copy()
    if t >= 6.0:
        storm_factor_ss = 1.0 + 4.0 * np.exp(-(t - 6.0) / 4.0)
        storm_factor_3x = 1.0 + 2.0 * np.exp(-(t - 6.0) / 4.0)
        y_in[1] *= storm_factor_ss       # S_S
        y_in[3] *= storm_factor_3x       # X_S
        y_in[9] *= storm_factor_3x       # S_NH
    return y_in

def _rhs_asm1_storm(t, y):
    """ASM1 CSTR with time-varying storm influent."""
    y_safe = np.maximum(y, 0.0)

    rho = _asm1_process_rates(y_safe)
    r = _asm1_reaction_vector(rho)

    y_in = _storm_influent(t)
    D_h = _D / _H_PER_D

    dy = np.zeros(13)
    for i in range(13):
        dy[i] = r[i] / _H_PER_D + D_h * (y_in[i] - y_safe[i])

    dy[7] += (_KLA / _H_PER_D) * (_S_O_SAT - y_safe[7])

    return _nonneg_clamp(y, dy)