Solving moment equations numerically

A numerical solver advances the retained moments through time; solver checks and reconstructed-moment checks are both necessary.

Initial moments

If the initial infectious count is known exactly as \(I(0)=I_0\), then

\[m_1(0)=I_0,\qquad m_2(0)=I_0^2,\qquad v(0)=0.\]

If the initial state were random, its first two moments would need to be supplied instead.

Solver inputs

Input	Purpose
`fun`	Moment derivative function.
`t_span`	Start and final time.
`y0`	Initial retained moments.
`t_eval`	Times at which results are returned.
`rtol, atol`	Relative and absolute local error controls.

Interactive Python laboratory

Interactive PythonSolving closed moment equations

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp

def moment_rhs(t, y, N, beta, gamma):
    m1, m2 = y
    m3 = 3 * m1 * m2 - 2 * m1 ** 3

dm1 = (beta - gamma) * m1 - beta * m2 / N
    dm2 = (
        (2 * (beta - gamma) - beta / N) * m2
        - 2 * beta * m3 / N
        + (beta + gamma) * m1
    )
    return [dm1, dm2]

N = 100
I0 = 5
beta = 0.25
gamma = 0.15
T = 30

output_times = np.linspace(0, T, 301)
initial_moments = [I0, I0 ** 2]

solution = solve_ivp(
    fun=moment_rhs,
    t_span=(0, T),
    y0=initial_moments,
    t_eval=output_times,
    args=(N, beta, gamma),
    method="RK45",
    rtol=1e-8,
    atol=1e-10
)

if not solution.success:
    raise RuntimeError(solution.message)

m1, m2 = solution.y
variance = m2 - m1 ** 2

if variance.min() < -1e-7:
    raise ValueError("The solution produced negative variance.")

variance = np.maximum(variance, 0)
sd = np.sqrt(variance)

selected_times = [0, 5, 10, 20, 30]
indices = [
    np.argmin(np.abs(solution.t - t))
    for t in selected_times
]
table = pd.DataFrame({
    "time": selected_times,
    "mean m1": m1[indices],
    "second moment m2": m2[indices],
    "variance": variance[indices],
    "standard deviation": sd[indices]
})

print(f"Solver success: {solution.success}")
print(f"Derivative evaluations: {solution.nfev}")
print()
print(table.round(4).to_string(index=False))

plt.figure(figsize=(9.5, 5.2))
plt.plot(solution.t, m1, color="#174c34",
         linewidth=2, label="Mean m1")
plt.fill_between(
    solution.t,
    np.maximum(m1 - sd, 0),
    m1 + sd,
    color="#b9ddc7", alpha=0.65,
    label="Mean ± one SD"
)
plt.xlabel("Time")
plt.ylabel("Infectious population")
plt.title("Numerical solution of closed SIS moment equations")
plt.legend()
plt.grid(alpha=0.25)
plt.tight_layout()

plt.figure(figsize=(9.5, 4.9))
plt.plot(solution.t, variance,
         color="#b42318", label="Variance")
plt.xlabel("Time")
plt.ylabel("Variance of I")
plt.title("Variance reconstructed from m2 − m1²")
plt.legend()
plt.grid(alpha=0.25)
plt.tight_layout()

Output

Run the code to see the result.

Why tolerate a tiny negative residual?

Variance is mathematically non-negative. A computed value such as \(-10^{-12}\) can arise from floating-point subtraction of nearly equal numbers. The code rejects materially negative variance below \(-10^{-7}\) and clips only smaller numerical residuals to zero. This is different from hiding a failed closure or solver.

Solver accuracy is not closure accuracy

Tight tolerances reduce numerical integration error in the closed ODEs. They do not make the normal closure exact. Comparison with stochastic simulations is still required.

What this lesson adds

You can now supply initial moments, solve a closed system with solve_ivp, inspect solver success, reconstruct variance and separate numerical tolerance from closure error.