Routine Immunization¶

import numpy as np
import matplotlib.pyplot as plt

from laser.core import PropertySet
from laser.core.utils import grid
from laser.generic.utils import ValuesMap
from laser.cohorts import Model, Campaign, Intervention
from laser.cohorts.vitaldynamics import NonDiseaseMortality, ConstantPopBirths
import laser.cohorts.SEIR as SEIR

# Use the built-in Vaccination intervention from laser.cohorts
try:
    from laser.cohorts.interventions.vaccination import Vaccination
    Campaign.register(Vaccination)
    print("Loaded Vaccination from laser.cohorts.interventions.vaccination")
except ImportError:
    # Fallback: minimal inline Vaccination that moves S → R
    print("Built-in Vaccination not found — using inline fallback")

    class Vaccination(Intervention):
        """Move a binomial-drawn fraction of targeted S individuals into R."""

        @property
        def states(self):
            return []

        @property
        def properties(self):
            return []

        def execute(self, tick, who, where, params, notes):
            coverage = float(params.get("coverage", 0.0))
            target   = where if where is not None else list(range(len(self.model.scenario)))
            S_idx    = self.model.states.get_state_index("S")
            R_idx    = self.model.states.get_state_index("R")
            if S_idx is None or R_idx is None:
                return
            for node in target:
                S_count  = self.model.states[tick + 1][S_idx, node]
                vaccinated = np.random.binomial(S_count, coverage)
                self.model.states[tick + 1][S_idx, node] -= vaccinated
                self.model.states[tick + 1][R_idx, node] += vaccinated

    Campaign.register(Vaccination)

# Shared parameters across all scenarios
N              = 500_000
I0             = 50
CDR            = 20                  # per 1000 per year
r_mortality    = CDR / 1000 / 365
beta           = 0.5
gamma          = 0.1                 # recovery rate (SEIR uses this as r_recovery)
sigma          = 0.2                 # progression rate E→I  (mean latent period = 5 days)
nticks         = 20 * 365
n_nodes        = 1

R0 = beta / gamma
p_critical = 1.0 - 1.0 / R0

print(f"R0              = {R0:.1f}")
print(f"Critical coverage p_c = 1 - 1/R0 = {p_critical:.2f}")

def run_seir_with_ri(coverage, N, I0, r_mortality, beta, sigma, gamma, nticks):
    """Run a single-patch SEIR model with vital dynamics and routine immunization.

    Vaccination is applied every tick to susceptibles at the given coverage.
    Uses the built-in Vaccination intervention (or the fallback S→R version).

    Args:
        coverage: Fraction of susceptibles vaccinated per tick.
        N: Initial population size.
        I0: Initial infectious count.
        r_mortality: Daily non-disease mortality rate.
        beta: Transmission rate.
        sigma: Progression rate E→I.
        gamma: Recovery rate.
        nticks: Simulation length in days.

    Returns:
        Fitted model after model.run().
    """
    scenario = grid(M=1, N=1)
    scenario["S"] = N - I0
    scenario["E"] = 0
    scenario["I"] = I0
    scenario["R"] = 0

    params = PropertySet({"nticks": nticks})
    model  = Model(scenario, params)

    # Every-tick RI schedule — "*" means all ticks, "*" means all nodes
    if coverage > 0:
        schedule = [{"who": ["S"], "what": "Vaccination", "when": "*",
                     "where": "*", "parameters": {"coverage": coverage}, "notes": "RI"}]
        campaign = Campaign(model, schedule)
    else:
        campaign = None

    betas         = ValuesMap.from_scalar(beta,  nticks, 1)
    r_recoveries  = ValuesMap.from_scalar(gamma, nticks, 1)
    r_progressions = ValuesMap.from_scalar(sigma, nticks, 1)

    components = [
        SEIR.Susceptible(model),
        SEIR.Exposed(model, r_progression=r_progressions),
        SEIR.Infectious(model, r_recovery=r_recoveries),
        SEIR.Recovered(model),
        NonDiseaseMortality(model, r_mortality=r_mortality),
        ConstantPopBirths(model),
        SEIR.Transmission(model, beta=betas),
    ]
    if campaign is not None:
        # Insert campaign before transmission so vaccinated S don't transmit
        components.insert(-1, campaign)

    model.components = components
    model.run()
    return model


# Run three scenarios
coverages = [0.0, 0.5, 0.8]
models    = {}

for cov in coverages:
    print(f"Running coverage = {cov:.1f} ...")
    models[cov] = run_seir_with_ri(
        cov, N, I0, r_mortality, beta, sigma, gamma, nticks
    )

print("All scenarios complete.")

# Plot I(t) for all three coverage scenarios
colors = {0.0: "firebrick", 0.5: "darkorange", 0.8: "steelblue"}
labels = {
    0.0: "No RI (coverage = 0)",
    0.5: "Partial RI (coverage = 0.5)",
    0.8: f"High RI (coverage = 0.8 ≈ p_c = {p_critical:.2f})",
}

t_years = np.arange(nticks + 1) / 365.0

fig, ax = plt.subplots(figsize=(12, 5))
for cov in coverages:
    I_ts = models[cov].states.I[:, 0]
    ax.plot(t_years, I_ts, color=colors[cov], linewidth=1.2, label=labels[cov])

ax.set_xlabel("Time (years)")
ax.set_ylabel("Infectious count I(t)")
ax.set_title("SEIR + vital dynamics: effect of routine immunization")
ax.legend(fontsize=9)
plt.tight_layout()
plt.show()

# Compute cumulative infections for each scenario
cumulative = {}
for cov in coverages:
    ni = models[cov].nodes.newly_infectious[:, 0]
    cumulative[cov] = int(ni.sum())
    print(f"Coverage {cov:.1f}: cumulative infections = {cumulative[cov]:>10,}")

# Bar chart
fig, ax = plt.subplots(figsize=(6, 4))
bars = ax.bar(
    [str(c) for c in coverages],
    [cumulative[c] for c in coverages],
    color=[colors[c] for c in coverages],
    edgecolor="black",
    linewidth=0.7,
)
for bar, cov in zip(bars, coverages):
    ax.text(bar.get_x() + bar.get_width() / 2,
            bar.get_height() * 1.02,
            f"{cumulative[cov]:,}",
            ha="center", va="bottom", fontsize=9)

ax.set_xlabel("RI coverage")
ax.set_ylabel("Cumulative infections (20 years)")
ax.set_title("Cumulative infections vs. routine immunization coverage")
plt.tight_layout()
plt.show()

# Reduction relative to no-RI baseline
baseline = cumulative[0.0]
for cov in coverages[1:]:
    reduction = 100.0 * (1.0 - cumulative[cov] / baseline)
    print(f"Coverage {cov:.1f}: {reduction:.1f}% reduction vs. no RI")