SEI and SEIS model implementations¶

from laser.generic.SEIR import Exposed
import inspect

print(inspect.getsource(Exposed.__init__))

print(inspect.getsource(Exposed.step))

from laser.core.utils import grid
from laser.core.random import seed as set_seed
from laser.generic import Model
from laser.core import PropertySet
from laser.core.distributions import normal

initial_infectious = 10
seeds = [1, 2, 3, 5, 8, 13, 21, 34, 55, 89]
nticks = 365

scenario = grid(M=1, N=1, node_size_degs=0.1, population_fn=lambda x, y: 1_000_000)
scenario["S"] = scenario.population - initial_infectious
scenario["E"] = 0
scenario["I"] = initial_infectious

results_sei = []
for i, seed in enumerate(seeds):
    set_seed(seed)
    model = Model(scenario, params=PropertySet({"beta": 0.3, "seed": seed, "nticks": nticks}), skip_capacity=True)
    exposure_duration = normal(loc=5, scale=1.5)
    model.components = [
        Susceptible(model),
        Infectious(model),  # Infectious goes _before_ Exposed so the previous I count is propagated before the Exposed step updates it.
        ExposedSEI(model, expdurdist=exposure_duration, expdurmin=1), # ExposedSEI needs the exposure duration distribution for initializing exposed agents.
        Transmission(model, expdurdist=exposure_duration, expdurmin=1), # Transmission needs the exposure duration distribution to set the exposure timers on newly infected agents.
    ]

    model.run(f"Scenario {i+1:2}/{len(seeds)} (seed={seed:2})")
    results_sei.append((model.nodes.S, model.nodes.E, model.nodes.I))

import matplotlib.pyplot as plt

fig, ax = plt.subplots(3, 1, figsize=(10, 12), sharex=True)
for i, (S, E, I) in enumerate(results_sei):
    ax[0].plot(S.sum(axis=1), label=f"Sim {i+1}")
    ax[1].plot(E.sum(axis=1), label=f"Sim {i+1}")
    ax[2].plot(I.sum(axis=1), label=f"Sim {i+1}")
ax[0].set_ylabel("Susceptible")
ax[1].set_ylabel("Exposed")
ax[2].set_ylabel("Infectious")
ax[2].set_xlabel("Days")
ax[0].legend()
plt.suptitle("SEI Model: Single Node Suite (1,000,000 agents each)")
plt.show()

from laser.generic import SI

# Let's run 10 SI models with same R0 for comparison.
results_si = []
for i, seed in enumerate(seeds):
    set_seed(seed)
    model = Model(scenario, params=PropertySet({"beta": 0.3, "seed": seed, "nticks": nticks}), skip_capacity=True)
    exposure_duration = normal(loc=5, scale=1.5)
    model.components = [
        SI.Susceptible(model),
        SI.Infectious(model),
        SI.Transmission(model),
    ]

    model.run(f"Scenario {i+1:2}/{len(seeds)} (seed={seed:2})")
    results_si.append((model.nodes.S, model.nodes.I))

import matplotlib.pyplot as plt

fig, ax = plt.subplots(2, 1, figsize=(10, 12), sharex=True)
for i, (S, I) in enumerate(results_si):
    ax[0].plot(S.sum(axis=1), label=f"Sim {i+1}")
    ax[1].plot(I.sum(axis=1), label=f"Sim {i+1}")
ax[0].set_ylabel("Susceptible")
ax[1].set_ylabel("Infectious")
ax[1].set_xlabel("Days")
ax[0].legend()
plt.suptitle("SI Model: Single Node Suite (1,000,000 agents each)")
plt.show()

import numpy as np

import matplotlib.pyplot as plt

# Compute mean S and I for SEI model
S_sei = np.array([S for S, E, I in results_sei])  # shape: (n_sim, n_time, 1)
I_sei = np.array([I for S, E, I in results_sei])
mean_S_sei = S_sei.mean(axis=0).squeeze()
mean_I_sei = I_sei.mean(axis=0).squeeze()

# Compute mean S and I for SI model
S_si = np.array([S for S, I in results_si])
I_si = np.array([I for S, I in results_si])
mean_S_si = S_si.mean(axis=0).squeeze()
mean_I_si = I_si.mean(axis=0).squeeze()

plt.figure(figsize=(10, 6))
plt.plot(mean_S_sei, label="SEI Mean Susceptible", color="tab:blue")
plt.plot(mean_I_sei, label="SEI Mean Infectious", color="tab:orange")
plt.plot(mean_S_si, label="SI Mean Susceptible", color="tab:blue", linestyle="--")
plt.plot(mean_I_si, label="SI Mean Infectious", color="tab:orange", linestyle="--")
plt.xlabel("Days")
plt.ylabel("Number of Agents")
plt.title("Mean S and I: SEI vs SI Models")
plt.legend()
plt.show()

initial_infectious = 10
seeds = [1, 2, 3, 5, 8, 13, 21, 34, 55, 89]
nticks = 365

scenario = grid(M=1, N=1, node_size_degs=0.1, population_fn=lambda x, y: 1_000_000)
scenario["S"] = scenario.population - initial_infectious
scenario["E"] = 0
scenario["I"] = initial_infectious

results_seis = []
for i, seed in enumerate(seeds):
    set_seed(seed)
    model = Model(scenario, params=PropertySet({"beta": 0.3, "seed": seed, "nticks": nticks}), skip_capacity=True)
    exposure_duration = normal(loc=5, scale=1.5)
    infectious_duration = normal(loc=14, scale=2.0)
    model.components = [
        Susceptible(model),
        InfectiousSEIS(model, infdurdist=infectious_duration, infdurmin=1),  # Infectious goes _before_ Exposed so the previous I count is propagated before the Exposed step updates it.
        SEIR.Exposed(model, expdurdist=exposure_duration, expdurmin=1, infdurdist=infectious_duration, infdurmin=1), # ExposedSEI needs the exposure duration distribution for initializing exposed agents.
        SEIR.Transmission(model, expdurdist=exposure_duration, expdurmin=1), # Transmission needs the exposure duration distribution to set the exposure timers on newly infected agents.
    ]

    model.run(f"Scenario {i+1:2}/{len(seeds)} (seed={seed:2})")
    results_seis.append((model.nodes.S, model.nodes.E, model.nodes.I))

import matplotlib.pyplot as plt

fig, ax = plt.subplots(3, 1, figsize=(10, 12), sharex=True)
for i, (S, E, I) in enumerate(results_seis):
    ax[0].plot(S.sum(axis=1), label=f"Sim {i+1}")
    ax[1].plot(E.sum(axis=1), label=f"Sim {i+1}")
    ax[2].plot(I.sum(axis=1), label=f"Sim {i+1}")
ax[0].set_ylabel("Susceptible")
ax[1].set_ylabel("Exposed")
ax[2].set_ylabel("Infectious")
ax[2].set_xlabel("Days")
ax[0].legend()
plt.suptitle("SEIS Model: Single Node Suite (1,000,000 agents each)")
plt.show()

# # Compute mean S, E, and I for SEI model
# S_sei = np.array([S for S, E, I in results_sei])  # shape: (n_sim, n_time, 1)
E_sei = np.array([E for S, E, I in results_sei])
# I_sei = np.array([I for S, E, I in results_sei])
# mean_S_sei = S_sei.mean(axis=0).squeeze()
mean_E_sei = E_sei.mean(axis=0).squeeze()
# mean_I_sei = I_sei.mean(axis=0).squeeze()

# Compute mean S, E, and I for SEIS model
S_seis = np.array([S for S, E, I in results_seis])
E_seis = np.array([E for S, E, I in results_seis])
I_seis = np.array([I for S, E, I in results_seis])
mean_S_seis = S_seis.mean(axis=0).squeeze()
mean_E_seis = E_seis.mean(axis=0).squeeze()
mean_I_seis = I_seis.mean(axis=0).squeeze()

plt.figure(figsize=(10, 6))
plt.plot(mean_S_sei, label="SEI Mean Susceptible", color="tab:blue")
plt.plot(mean_E_sei, label="SEI Mean Exposed", color="tab:green")
plt.plot(mean_I_sei, label="SEI Mean Infectious", color="tab:orange")
plt.plot(mean_S_seis, label="SEIS Mean Susceptible", color="tab:blue", linestyle="--")
plt.plot(mean_E_seis, label="SEIS Mean Exposed", color="tab:green", linestyle="--")
plt.plot(mean_I_seis, label="SEIS Mean Infectious", color="tab:orange", linestyle="--")
plt.xlabel("Days")
plt.ylabel("Number of Agents")
plt.title("Mean S, E, and I: SEI vs SEIS Models")
plt.legend()
plt.show()