SIR

=SIR(s_initial, i_initial, r_initial, beta, gamma, t_start, t_end, [method])

=SIR(990, 10, 0, 0.3, 0.1, 0, 10)

=SIR(990, 10, 0, 0.3, 0.1, 0, 10, "RK23")

=SIR(990, 10, 0, 0.3, 0.1, 0, 50)

=SIR(990, 10, 0, 0.3, 0.1, 0, 10, "BDF")

from scipy.integrate import solve_ivp as scipy_solve_ivp
from typing import List, Union
 
def sir(
    s_initial: float,
    i_initial: float,
    r_initial: float,
    beta: float,
    gamma: float,
    t_start: float,
    t_end: float,
    method: str = 'RK45'
) -> Union[List[List[float]], str]:
    """
    Solves the SIR system of ordinary differential equations for infection dynamics.
 
    Args:
        s_initial: Initial number of susceptible individuals (float).
        i_initial: Initial number of infected individuals (float).
        r_initial: Initial number of recovered individuals (float).
        beta: Infection rate parameter (float).
        gamma: Recovery rate parameter (float).
        t_start: Start time of integration (float).
        t_end: End time of integration (float).
        method: Integration method to use (str): 'RK45', 'RK23', 'DOP853', 'Radau', 'BDF', 'LSODA'. Default is 'RK45'.
 
    Returns:
        2D list of solution values for S, I, and R at each time point, or an error message (str) if input is invalid.
 
    This example function is provided as-is without any representation of accuracy.
    """
    # Validate inputs
    try:
        s0 = float(s_initial)
        i0 = float(i_initial)
        r0 = float(r_initial)
        b = float(beta)
        g = float(gamma)
        t0 = float(t_start)
        tf = float(t_end)
    except Exception:
        return "Invalid input: All initial values and parameters must be numbers."
    if any(x < 0 for x in [s0, i0, r0, b, g]):
        return "Invalid input: Initial values and parameters must be non-negative."
    if tf <= t0:
        return "Invalid input: t_end must be greater than t_start."
    if method not in ['RK45', 'RK23', 'DOP853', 'Radau', 'BDF', 'LSODA']:
        return "Invalid input: method must be one of 'RK45', 'RK23', 'DOP853', 'Radau', 'BDF', 'LSODA'."
    # SIR ODE system
    def sir_ode(t, y):
        S, I, R = y
        N = S + I + R
        dSdt = -b * S * I / N
        dIdt = b * S * I / N - g * I
        dRdt = g * I
        return [dSdt, dIdt, dRdt]
    # Integrate
    try:
        sol = scipy_solve_ivp(
            sir_ode,
            [t0, tf],
            [s0, i0, r0],
            method=method
        )
    except Exception as e:
        return f"scipy.integrate.solve_ivp error: {e}"
    if not sol.success:
        return f"Integration failed: {sol.message}"
    # Return 2D list: header row, then each row is [t, S, I, R] at each time point
    header = ["t", "S", "I", "R"]
    result = [header] + [[float(sol.t[i]), float(sol.y[0][i]), float(sol.y[1][i]), float(sol.y[2][i])] for i in range(len(sol.t))]
    return result