Functions for reusable epidemic models

A model becomes reusable when its equations, numerical solution and outcome calculations are organised into clear functions. Then parameters can be changed without copying and editing the whole program.

Biological scenario

We want to compare three possible transmission conditions:

lower transmission: \(\beta=0.15\);
baseline transmission: \(\beta=0.30\);
higher transmission: \(\beta=0.45\).

All scenarios use the same population, initial conditions and recovery rate. Copying a complete solver three times would be repetitive and make errors more likely.

Why use functions?

Repeated program

Equations copied several times
Corrections must be repeated
Parameters may be changed inconsistently
Comparison code becomes difficult to read

Reusable functions

Each calculation is defined once
Different values are passed as arguments
Validation is applied consistently
The main experiment remains short and readable

Separate responsibilities

FUNCTION 1Equations

Calculate derivatives at one state.

FUNCTION 2Simulation

Solve the model over a time interval.

FUNCTION 3Summary

Extract peak and final outcomes.

This is called separation of responsibilities. Each function should have one clear job. A short focused function is easier to test, explain and reuse.

Anatomy of a function

def simulate_sir(beta, gamma, S0, I0, R0, end_time=160):
    """Solve one deterministic SIR scenario."""
    # calculations
    return result

defBegins the definition.

Function namesimulate_sir describes the job.

ParametersNames receiving supplied input values.

returnSends the result back to the caller.

Feature	Meaning
`"""Solve one deterministic SIR scenario."""`	A docstring explaining the function’s purpose.
`end_time=160`	A default argument. The caller may omit it or supply another value.
Indented statements	The function body. They run when the function is called.
Local variables	Names created inside the function and normally used only there.

Positional and keyword arguments

Positional call

simulate_sir(0.3, 0.1, 990, 10, 0)

Values are matched by position. The order must be remembered correctly.

Keyword call

simulate_sir(
    beta=0.3,
    gamma=0.1,
    S0=990,
    I0=10,
    R0=0
)

Values are matched by parameter name. This is longer but clearer for scientific code.

Function 1: the equations

def sir_rhs(t, y, beta, gamma):
    S, I, R = y
    N = S + I + R

    infection_flow = beta * S * I / N
    recovery_flow = gamma * I

    return [
        -infection_flow,
        infection_flow - recovery_flow,
        recovery_flow
    ]

This function knows the SIR mathematics but does not choose initial conditions, a time interval or an experiment.

Function 2: the simulation

def simulate_sir(
    beta, gamma, S0, I0, R0,
    end_time=160, points=1601
):
    # validate inputs
    # create reporting times
    # call solve_ivp
    # check success
    # return organised results

The function returns a dictionary:

return {
    "time": solution.t,
    "S": S,
    "I": I,
    "R": R,
    "beta": beta,
    "gamma": gamma
}

A dictionary stores named key–value pairs. The infectious history can then be retrieved using result["I"].

Function 3: the outcome summary

def epidemic_summary(result):
    peak_index = np.argmax(result["I"])

    return {
        "peak_I": result["I"][peak_index],
        "peak_time": result["time"][peak_index],
        "final_R": result["R"][-1]
    }

This function receives a simulation dictionary and returns a smaller dictionary containing the requested outcomes.

The reusable workflow

Define
equations and tools once

Call
with different parameters

Compare
outcomes consistently

Run three scenarios

Interactive PythonReusable SIR functions

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

def sir_rhs(t, y, beta, gamma):
    """Return SIR derivatives at one epidemic state."""
    
    S, I, R = y
    N = S + I + R
    
    infection_flow = beta * S * I / N
    recovery_flow = gamma * I
    
    return [
        -infection_flow,
        infection_flow - recovery_flow,
        recovery_flow
    ]

def simulate_sir(
    beta,
    gamma,
    S0,
    I0,
    R0,
    end_time=160,
    points=1601
):
    """Solve one deterministic SIR scenario."""
    
    values = [beta, gamma, S0, I0, R0]
    
    if min(values) < 0:
        raise ValueError("Parameters and compartments cannot be negative.")
    if S0 + I0 + R0 <= 0:
        raise ValueError("The total population must be positive.")
    if end_time <= 0:
        raise ValueError("end_time must be positive.")
    if points < 2:
        raise ValueError("points must be at least 2.")
    
    reporting_times = np.linspace(0, end_time, points)
    
    solution = solve_ivp(
        sir_rhs,
        t_span=(0, end_time),
        y0=[S0, I0, R0],
        t_eval=reporting_times,
        args=(beta, gamma),
        rtol=1e-8,
        atol=1e-10
    )
    
    if not solution.success:
        raise RuntimeError(solution.message)
    
    S, I, R = solution.y
    
    return {
        "time": solution.t,
        "S": S,
        "I": I,
        "R": R,
        "beta": beta,
        "gamma": gamma
    }

def epidemic_summary(result):
    """Extract important outcomes from one SIR simulation."""
    
    peak_index = np.argmax(result["I"])
    
    return {
        "peak_I": result["I"][peak_index],
        "peak_time": result["time"][peak_index],
        "final_R": result["R"][-1]
    }

# Scenario names and transmission parameters
scenarios = {
    "Lower transmission": 0.15,
    "Baseline": 0.30,
    "Higher transmission": 0.45
}

# Run every scenario
results = {}

for name, beta in scenarios.items():
    results[name] = simulate_sir(
        beta=beta,
        gamma=0.1,
        S0=990,
        I0=10,
        R0=0
    )

# Display the outcome table
print("Scenario                beta   Peak I   Peak day   Final R")
print("------------------------------------------------------------")

for name, result in results.items():
    summary = epidemic_summary(result)
    
    print(
        f"{name:22s}  {result['beta']:4.2f}  "
        f"{summary['peak_I']:7.2f}  "
        f"{summary['peak_time']:8.1f}  "
        f"{summary['final_R']:8.2f}"
    )

# Plot all infectious trajectories
plt.figure(figsize=(8, 4.8))

for name, result in results.items():
    plt.plot(
        result["time"],
        result["I"],
        label=name
    )

plt.xlabel("Time (days)")
plt.ylabel("Infectious population")
plt.title("Reusable SIR model: transmission scenarios")
plt.grid(alpha=0.25)
plt.legend()
plt.tight_layout()

Output

Run the code to see the result.

Understand the scenario loop

Code	Meaning
`scenarios = {...}`	Creates a dictionary connecting each scenario name to its \(\beta\) value.
`scenarios.items()`	A dictionary method returning each key and value together.
`for name, beta in ...`	Repeats the block once for each scenario and unpacks its name and value.
`results[name] = ...`	Stores each returned simulation under its scenario name.
`result["I"]`	Retrieves the infectious array from one result dictionary.

Biological interpretation

The functions ensure that only the transmission parameter changes across the three scenarios. Higher transmission generally produces faster initial growth, an earlier and larger infectious peak, and a larger final recovered population.

Because the same simulation and summary functions are reused, differences in the results arise from the supplied parameter values rather than accidental differences between copied programs.

Function-design principles

Give each function one clear responsibility.
Use descriptive names for functions, parameters and returned values.
Validate inputs near the beginning of a function.
Use docstrings to state what a function does.
Prefer keyword arguments when scientific calls contain several similar numerical values.
Return organised results instead of depending on hidden global variables.
Do not place fixed biological parameter values inside a function when they should be changeable.

Try adding a scenario

Add another dictionary entry:

"Very low transmission": 0.08

Run the same functions without changing their definitions.
Inspect whether the infectious population grows substantially.
Compare its peak and final recovered population with the other scenarios.
Explain why this demonstrates reusability.