[QUESTION]: Is support of non-ODEs in template_from_sympy_odes is expected/planned?

[liunelson]

In the Rabiu 2024 paper (Model C of the recent hackathon), we have this system of equations:

I recall from an earlier discussion with @bgyori that a user may not need to substitute the \lambda_m, \lamda_w expressions into the ODEs. However, this is not currently supported (?) and I wonder if/when it will be.

I hope that template_model_from_sympy_odes can

• Find which equations in the list do not have a Derivative(...) on the left hand side (e.g. \lambda_m, \lambda_w, N equations)
• Try to substitute them into every ODE
• Generate the template model from this substituted ODE system
• Possibly include these non-ODE equations as observables to the model

[liunelson]

Passing the LaTeX equations extracted from the document through our equation-cleaner agent:

odes = [
  "\\frac{d S_{m}(t)}{d t} = -\\lambda_{m} * S_{m}(t)",
  "\\frac{d E_{m}(t)}{d t} = \\lambda_{m} * S_{m}(t) - \\alpha_{m} * E_{m}(t)",
  "\\frac{d I_{m}(t)}{d t} = \\alpha_{m} * E_{m}(t) - \\tau_{m} * I_{m}(t)",
  "\\frac{d R_{m}(t)}{d t} = \\tau_{m} * I_{m}(t)",
  "\\frac{d S_{w}(t)}{d t} = - \\lambda_{w} * S_{w}(t)",
  "\\frac{d E_{w}(t)}{d t} = \\lambda_{w} * S_{w}(t) - \\alpha_{w} * E_{w}(t)",
  "\\frac{d I_{w}(t)}{d t} = \\alpha_{w} * E_{w}(t) - \\tau_{w} * I_{w}(t)",
  "\\frac{d R_{w}(t)}{d t} = \\tau_{w} * I_{w}(t)",
  "\\lambda_{m} = \\frac{\\beta_{mm} * I_{m}(t)}{N_{m}} + \\frac{\\alpha_{1} * \\beta_{wm} * I_{w}(t)}{N_{m}}",
  "\\lambda_{w} = \\frac{\\alpha_{2} * \\beta_{mw} * I_{m}(t)}{N_{w}} + \\frac{\\alpha_{3} * \\beta_{ww} * I_{w}(t)}{N_{w}}",
  "N = N_{w} + N_{m}"
]

Passing only the ODEs (i.e. not the last 3 equations) through the LaTeX-SymPy agent:

import sympy
# Define time variable
t = sympy.symbols("t")
# Define time-dependent variables
S_m, E_m, I_m, R_m, S_w, E_w, I_w, R_w = sympy.symbols("S_m E_m I_m R_m S_w E_w I_w R_w", cls=sympy.Function)
# Define constant parameters
lambda_m, alpha_m, tau_m, lambda_w, alpha_w, tau_w = sympy.symbols("lambda_m alpha_m tau_m lambda_w alpha_w tau_w")
equation_output = [
    sympy.Eq(S_m(t).diff(t), -lambda_m * S_m(t)),
    sympy.Eq(E_m(t).diff(t), lambda_m * S_m(t) - alpha_m * E_m(t)),
    sympy.Eq(I_m(t).diff(t), alpha_m * E_m(t) - tau_m * I_m(t)),
    sympy.Eq(R_m(t).diff(t), tau_m * I_m(t)),
    sympy.Eq(S_w(t).diff(t), -lambda_w * S_w(t)),
    sympy.Eq(E_w(t).diff(t), lambda_w * S_w(t) - alpha_w * E_w(t)),
    sympy.Eq(I_w(t).diff(t), alpha_w * E_w(t) - tau_w * I_w(t)),
    sympy.Eq(R_w(t).diff(t), tau_w * I_w(t))
]

If I update the LaTeX-SymPy agent to handle non-ODEs, it'd probably return either one of these two options:

• Case A, the LaTeX-SymPy agent substituted the non-ODEs into the ODEs

import sympy
# Define time variable
t = sympy.symbols("t")
# Define time-dependent variables
S_m, R_m, S_w, R_w, E_m, I_m, E_w, I_w = sympy.symbols("S_m R_m S_w R_w E_m I_m E_w I_w", cls=sympy.Function)
# Define constant parameters
beta_mm, beta_wm, beta_mw, beta_ww, alpha_1, alpha_2, alpha_3, alpha_m, alpha_w, tau_m, tau_w, N_m, N_w, N = sympy.symbols("beta_mm beta_wm beta_mw beta_ww alpha_1 alpha_2 alpha_3 alpha_m alpha_w tau_m tau_w N_m N_w N")

# Define non-ODEs
lambda_m = beta_mm * I_m(t) / N_m + alpha_1 * beta_wm * I_w(t) / N_m
lambda_w = alpha_2 * beta_mw * I_m(t) / N_w + alpha_3 * beta_ww * I_w(t) / N_w
N = N_w + N_m

# Define ODEs
equation_output = [...]

• Case B, the agent is instructed to define the non-ODEs as SymPy equations

import sympy
# Define time variable
t = sympy.symbols("t")
# Define time-dependent variables
S_m, R_m, S_w, R_w, E_m, I_m, E_w, I_w = sympy.symbols("S_m R_m S_w R_w E_m I_m E_w I_w", cls=sympy.Function)
# Define constant parameters
beta_mm, beta_wm, beta_mw, beta_ww, alpha_1, alpha_2, alpha_3, alpha_m, alpha_w, tau_m, tau_w, N_m, N_w, N = sympy.symbols("beta_mm beta_wm beta_mw beta_ww alpha_1 alpha_2 alpha_3 alpha_m alpha_w tau_m tau_w N_m N_w N")

# Define all equations
equation_output = [
  ...,
  sympy.Eq(lambda_m, beta_mm * I_m(t) / N_m + alpha_1 * beta_wm * I_w(t) / N_m,
  sympy.Eq(lambda_w, alpha_2 * beta_mw * I_m(t) / N_w + alpha_3 * beta_ww * I_w(t) / N_w),
  sympy.Eq(N, N_w + N_m)
]

If we send Case A to MIRA template_model_from_sympy_odes, the model would be generated as expected, though possibly more error-prone due to more complicated LLM task.

If we send Case B, MIRA would need to do some extra work but can also define the non-ODEs into observables.

  odes = [eqn for eqn in equation_output if len(eqn.lhs.atoms(sympy.Derivative)) > 0]
  non_odes = [eqn for eqn in equation_output if len(eqn.lhs.atoms(sympy.Derivative)) = 0]
  
  # Substitutions
  odes_subs = [ode_eqn.subs([non_ode_eqn.args for non_ode_eqn in non_odes]) for ode_eqn in odes]

  # Generate the model as usual
  model = template_model_from_sympy_odes(odes_subs)
  
  # Use `non_odes` to populate observables
  model.observables = {str(eqn.lhs): Observable(name = str(eqn.lhs), expression = eqn.rhs) for eqn in non_odes}

What do you think of adding the functionality of the Case B code to MIRA?