SIR Model with Vaccination

from tys import probe, progress

def simulate(cfg: dict):
    import simpy
    env = simpy.Environment()

    population = cfg["population"]
    infected = cfg["initial_infected"]
    beta = cfg["beta"]
    gamma = cfg["gamma"]
    vacc_rate = cfg.get("vaccination_rate", 0.0)
    vaccinate = cfg.get("vaccinate", False)
    days = cfg["sim_days"]

    susceptible = population - infected
    recovered = 0.0
    vaccinated_total = 0.0

    done = env.event()

    def recorder():
        while True:
            probe("susceptible", env.now, susceptible)
            probe("infected", env.now, infected)
            probe("recovered", env.now, recovered)
            yield env.timeout(1)
    env.process(recorder())

    def dynamics():
        nonlocal susceptible, infected, recovered, vaccinated_total
        for d in range(days):
            new_infections = beta * susceptible * infected / population
            new_recoveries = gamma * infected
            if vaccinate:
                new_vaccinated = vacc_rate * susceptible
            else:
                new_vaccinated = 0.0

            susceptible = max(susceptible - new_infections - new_vaccinated, 0.0)
            infected = max(infected + new_infections - new_recoveries, 0.0)
            recovered = min(recovered + new_recoveries + new_vaccinated, population)
            vaccinated_total += new_vaccinated

            if infected < 1:
                progress(int(100 * d / days), "🎉 Epidemic nearly over")

            yield env.timeout(1)

        progress(100)
        done.succeed({
            "final_infected": infected,
            "vaccinated_total": vaccinated_total
        })

    env.process(dynamics())
    env.run(until=done)
    return done.value


def requirements():
    return {
        "builtin": ["micropip", "pyyaml"],
        "external": ["simpy==4.1.1"],
    }