# Eight latent states + four non-negative observables (T_*) + error_count.
SIRJPF_STATENAMES = [
"Sn", "In", "Jn", "Si", "Ii", "Ji", "F", "P",
"T_Sn", "T_In", "T_Si", "T_Ii", "error_count",
]
def sirjpf_rproc(X_, theta_, key, covars, t, dt):
"""One Euler-Maruyama step of the SIRJPF2 SDE system."""
Sn, Jn, In = X_["Sn"], X_["Jn"], X_["In"]
Si, Ji, Ii = X_["Si"], X_["Ji"], X_["Ii"]
F, P = X_["F"], X_["P"]
error_count = X_["error_count"]
sigSn, sigIn = theta_["sigSn"], theta_["sigIn"]
sigSi, sigIi = theta_["sigSi"], theta_["sigIi"]
sigJn, sigJi = theta_["sigJn"], theta_["sigJi"]
sigF, sigP = theta_["sigF"], theta_["sigP"]
theta_Sn, theta_In = theta_["theta_Sn"], theta_["theta_In"]
theta_Si, theta_Ii = theta_["theta_Si"], theta_["theta_Ii"]
theta_Jn, theta_Ji = theta_["theta_Jn"], theta_["theta_Ji"]
theta_P = theta_["theta_P"]
f_Sn, f_Si = theta_["f_Sn"], theta_["f_Si"]
rn, ri = theta_["rn"], theta_["ri"]
probn, probi = theta_["probn"], theta_["probi"]
xi = theta_["xi"]
# Fixed experimental constants
delta = 0.013 # sampling/dilution rate
mu_food = 0.37 # algal replenishment
lambda_J = 0.1 # juvenile maturation rate (both species)
# Eight independent Gaussian innovations
keys = jax.random.split(key, 8)
sqdt = jnp.sqrt(dt)
noiSn = sigSn * sqdt * jax.random.normal(keys[0])
noiIn = sigIn * sqdt * jax.random.normal(keys[1])
noiSi = sigSi * sqdt * jax.random.normal(keys[2])
noiIi = sigIi * sqdt * jax.random.normal(keys[3])
noiJn = sigJn * sqdt * jax.random.normal(keys[4])
noiJi = sigJi * sqdt * jax.random.normal(keys[5])
noiF = sigF * sqdt * jax.random.normal(keys[6])
noiP = sigP * sqdt * jax.random.normal(keys[7])
# Native species (n)
Sn_term = (lambda_J * Jn * dt
- theta_Sn * Sn * dt
- probn * f_Sn * Sn * P * dt
- delta * Sn * dt + Sn * noiSn)
Jn_term = (rn * f_Sn * F * Sn * dt
- lambda_J * Jn * dt
- theta_Jn * Jn * dt
- delta * Jn * dt + Jn * noiJn)
In_term = (probn * f_Sn * Sn * P * dt
- theta_In * In * dt
- delta * In * dt + In * noiIn)
# Invasive species (l, written `i` in code)
Si_term = (lambda_J * Ji * dt
- theta_Si * Si * dt
- probi * f_Si * Si * P * dt
- delta * Si * dt + Si * noiSi)
Ji_term = (ri * f_Si * F * Si * dt
- lambda_J * Ji * dt
- theta_Ji * Ji * dt
- delta * Ji * dt + Ji * noiJi)
Ii_term = (probi * f_Si * Si * P * dt
- theta_Ii * Ii * dt
- delta * Ii * dt + Ii * noiIi)
# Shared resources
F_term = (- f_Sn * F * (Sn + xi * In + Jn) * dt
- f_Si * F * (Si + xi * Ii + Ji) * dt
- delta * F * dt
+ mu_food * dt + F * noiF)
P_term = (30.0 * theta_In * In * dt
+ 30.0 * theta_Ii * Ii * dt
- f_Sn * (Sn + xi * In) * P * dt
- f_Si * (Si + xi * Ii) * P * dt
- theta_P * P * dt
- delta * P * dt + P * noiP)
Sn_new = Sn + Sn_term
In_new = In + In_term
Jn_new = Jn + Jn_term
Si_new = Si + Si_term
Ii_new = Ii + Ii_term
Ji_new = Ji + Ji_term
F_new = F + F_term
P_new = P + P_term
# Add the inoculum exactly once on the transition beginning at day 4.
inoculate = (t <= 4.0) & ((t + dt) > 4.0)
P_new = P_new + jnp.where(inoculate, 25.0, 0.0)
# Production reset rule with weighted error_count contributions.
def viol(x, hi):
return (x < 0.0) | (x > hi)
eps = 0.0
eps += jnp.where(viol(Sn_new, 1e5), 1.0, 0.0)
eps += jnp.where(viol(Si_new, 1e5), 1.0e6, 0.0)
eps += jnp.where(viol(F_new, 1e20), 1.0e3, 0.0)
eps += jnp.where(viol(In_new, 1e5), 1.0e-3, 0.0)
eps += jnp.where(viol(Ii_new, 1e5), 1.0e-9, 0.0)
eps += jnp.where(viol(Jn_new, 1e5), 1.0e-3, 0.0)
eps += jnp.where(viol(Ji_new, 1e5), 1.0e-9, 0.0)
eps += jnp.where(viol(P_new, 1e20) & (t > 3.9), 1.0e-6, 0.0)
Sn_new = jnp.where(viol(Sn_new, 1e5), 0.0, Sn_new)
In_new = jnp.where(viol(In_new, 1e5), 0.0, In_new)
Jn_new = jnp.where(viol(Jn_new, 1e5), 0.0, Jn_new)
Si_new = jnp.where(viol(Si_new, 1e5), 0.0, Si_new)
Ii_new = jnp.where(viol(Ii_new, 1e5), 0.0, Ii_new)
Ji_new = jnp.where(viol(Ji_new, 1e5), 0.0, Ji_new)
F_new = jnp.where(viol(F_new, 1e20), 0.0, F_new)
P_new = jnp.where(viol(P_new, 1e20) & (t > 3.9), 0.0, P_new)
return {
"Sn": Sn_new, "In": In_new, "Jn": Jn_new,
"Si": Si_new, "Ii": Ii_new, "Ji": Ji_new,
"F": F_new, "P": P_new,
"T_Sn": jnp.abs(Sn_new), "T_In": jnp.abs(In_new),
"T_Si": jnp.abs(Si_new), "T_Ii": jnp.abs(Ii_new),
"error_count": error_count + eps,
}
