SIR with optimal vaccination

An SIR epidemic with the vaccination rate that minimises the burden of infections and vaccines.

• Name: sir_vaccination
• Field: biology/epidemiology
• Formalism: ordinary differential equations
• References: Lenhart and Workman 2007

Factory

@register_model(
    field='biology/epidemiology',
    tags=('optimal control', 'vaccination', 'sir', 'pontryagin', 'ode'),
    references=('Lenhart and Workman 2007',),
    defaults={'until': 30.0, 'dt': 0.2, 'solver': PontryaginSolver(objective=_cost)},
    observables=lambda system: [ObservableSpec('vaccinated', _vaccinated), ObservableSpec('cost', _cost)],
)
def sir_vaccination(
    S: float = 0.99,
    I: float = 0.01,
    *,
    beta: float = 0.5,
    gamma: float = 0.1,
    a: float = 1.0,
    b: float = 1.0,
) -> System:
    """Return an SIR epidemic with the vaccination rate that minimises the burden of infections and vaccines.

    Run with an optimal control solver, vaccination is strongest early, while many are susceptible, and the epidemic
    peaks lower than without vaccination. With no weight `a` on the infections the optimal rate is 0, the epidemic
    without control.

    Args:
        S:
            The initial susceptible fraction.

        I:
            The initial infected fraction.

        beta:
            The transmission rate.

        gamma:
            The recovery rate.

        a:
            The weight of the infected fraction in the cost.

        b:
            The weight of the squared vaccination rate in the cost.

    Returns:
        system:
            The model, with the fractions `population.S`, `population.I`, and `population.R`, the vaccination rate
            `population.u`, and the observables `vaccinated` and `cost`, the running cost.

    Examples:
        >>> from skcomplex.models.biology.epidemiology import sir_vaccination
        >>> sir_vaccination(beta=0.3).elements[0].beta
        0.3
    """
    population = VaccinatedPopulationElement(
        'population',
        S=S,
        I=I,
        R=max(1 - S - I, 0.0),
        beta=beta,
        gamma=gamma,
        a=a,
        b=b,
    )
    return System(elements=[population])

Components

class VaccinatedPopulationElement(Element):
    """Susceptible, infected, and recovered fractions of a population vaccinated at a bounded rate.

    The fractions follow `S' = -beta S I - u S`, `I' = beta S I - gamma I`, and `R' = gamma I + u S`, where the
    vaccination rate `u` is a control input between 0 and its largest value. The planner minimises the burden of the
    infections and the cost of the vaccines, `∫ (a I + b u^2 / 2) dt`, after Lenhart and Workman.

    Args:
        label:
            The label of the population.

        visual:
            How the element is drawn, or `None` for `default_visual_`.

        S:
            The initial susceptible fraction.

        I:
            The initial infected fraction.

        R:
            The initial recovered fraction.

        u:
            The vaccination rate, the control input.

        beta:
            The transmission rate.

        gamma:
            The recovery rate.

        a:
            The weight of the infected fraction in the cost.

        b:
            The weight of the squared vaccination rate in the cost.

    Examples:
        >>> from skcomplex.models.biology.epidemiology import VaccinatedPopulationElement
        >>> VaccinatedPopulationElement('population', I=0.01).I
        0.01
    """

    def __init__(
        self,
        label: str,
        visual: VisualSpec | None = None,
        *,
        S: Annotated[float, StateSpec(domain=RealSet.closed_open(0, oo))] = 0.99,
        I: Annotated[float, StateSpec(domain=RealSet.closed_open(0, oo))] = 0.01,
        R: Annotated[float, StateSpec(domain=RealSet.closed_open(0, oo))] = 0.0,
        u: Annotated[float, ParameterSpec(adaptive=True, bounds=(0.0, MAX_RATE))] = 0.0,
        beta: Annotated[float, ParameterSpec(domain=RealSet.closed_open(0, oo))] = 0.5,
        gamma: Annotated[float, ParameterSpec(domain=RealSet.closed_open(0, oo))] = 0.1,
        a: Annotated[float, ParameterSpec(domain=RealSet.closed_open(0, oo))] = 1.0,
        b: Annotated[float, ParameterSpec(domain=RealSet.open(0, oo))] = 1.0,
    ) -> None:
        super().__init__(label, visual=visual)
        self.S = S
        self.I = I
        self.R = R
        self.u = u
        self.beta = beta
        self.gamma = gamma
        self.a = a
        self.b = b

    def _get_laws(self, s: SimpleNamespace) -> list[BaseLaw]:
        """Return the flows between the fractions.

        Args:
            s:
                The namespace of the symbols the laws can use.

        Returns:
            laws:
                The laws of the component.
        """
        infection, recovery, vaccination = s.beta * s.S * s.I, s.gamma * s.I, s.u * s.S
        return [
            RateLaw(s.S, -infection - vaccination),
            RateLaw(s.I, infection - recovery),
            RateLaw(s.R, recovery + vaccination),
        ]

Equations

population.S: diff(S_population(t), t) == -beta_population*I_population(t)*S_population(t) - S_population(t)*u_population(t)
population.I: diff(I_population(t), t) == beta_population*I_population(t)*S_population(t) - gamma_population*I_population(t)
population.R: diff(R_population(t), t) == gamma_population*I_population(t) + S_population(t)*u_population(t)
population.u: diff(u_population(t), t) == 0