Programming moment equations as an ODE system

A closed moment model becomes computable when retained moments and their derivatives are placed in a consistent state-vector function.

Choose the retained state vector

\[y(t)=\begin{pmatrix}m_1(t)\\m_2(t)\end{pmatrix}.\]

Here \(m_1=\mathbb E[I]\) and \(m_2=\mathbb E[I^2]\). The third moment is supplied by the normal closure.

Closed SIS equations

\[ m_3\approx3m_1m_2-2m_1^3, \]

\[ m_1'=(\beta-\gamma)m_1-\frac{\beta}{N}m_2, \]

\[ m_2'= \left[2(\beta-\gamma)-\frac{\beta}{N}\right]m_2 -\frac{2\beta}{N}m_3 +(\beta+\gamma)m_1. \]

Function contract

An ODE solver will call a function with current time, current moment vector and parameters. The function must return derivatives in the same order as the state vector.

def sis_moment_rhs(t, y, N, beta, gamma):
    m1, m2 = y
    ...
    return [dm1, dm2]

Interactive Python laboratory

This page evaluates and checks the derivative function at several valid moment states. Numerical integration belongs to lesson 10.

Interactive PythonProgramming the moment ODE right-hand side

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

def normal_third_moment(m1, m2):
    return 3 * m1 * m2 - 2 * m1 ** 3

def sis_moment_rhs(t, y, N, beta, gamma):
    m1, m2 = y

if N <= 0 or beta < 0 or gamma < 0:
        raise ValueError("Use biologically valid parameters.")

m3 = normal_third_moment(m1, m2)

dm1 = (
        (beta - gamma) * m1
        - beta * m2 / N
    )

dm2 = (
        (2 * (beta - gamma) - beta / N) * m2
        - 2 * beta * m3 / N
        + (beta + gamma) * m1
    )

return np.array([dm1, dm2])

N = 100
beta = 0.25
gamma = 0.15

test_states = [
    (0, 5.0, 25.0),
    (5, 8.0, 80.0),
    (10, 15.0, 300.0),
    (15, 25.0, 800.0)
]

rows = []
for time, m1, m2 in test_states:
    dm1, dm2 = sis_moment_rhs(
        time, [m1, m2], N, beta, gamma
    )
    rows.append({
        "time": time,
        "m1": m1,
        "m2": m2,
        "variance": m2 - m1 ** 2,
        "dm1/dt": dm1,
        "dm2/dt": dm2
    })

table = pd.DataFrame(rows)
print(table.round(4).to_string(index=False))

# Show how dm1 depends on both retained moments
m1_grid = np.linspace(0, 40, 81)
variance_levels = [0, 25, 100, 225]

plt.figure(figsize=(9.5, 5.1))
for variance in variance_levels:
    m2_grid = m1_grid ** 2 + variance
    dm1_values = np.array([
        sis_moment_rhs(
            0, [m1, m2], N, beta, gamma
        )[0]
        for m1, m2 in zip(m1_grid, m2_grid)
    ])
    plt.plot(m1_grid, dm1_values,
             label=f"Variance = {variance}")

plt.axhline(0, color="black", linewidth=0.8)
plt.xlabel("First moment m1")
plt.ylabel("Derivative dm1/dt")
plt.title("First-moment change depends on variance through m2")
plt.legend()
plt.grid(alpha=0.25)
plt.tight_layout()

Output

Run the code to see the result.

Programming checks

The returned derivative order matches \([m_1,m_2]\).
Initial moments from deterministic \(I_0\) are \(m_1(0)=I_0\) and \(m_2(0)=I_0^2\).
A valid moment state requires \(m_2-m_1^2\ge0\).
The closure formula is isolated in its own function so assumptions are visible.

What this lesson adds

You can now organise retained moments into a state vector, implement closure explicitly, return derivatives in solver-compatible order and test the ODE right-hand side before numerical integration.