Campaign Specification Formats¶

import csv
import json
import tempfile
from pathlib import Path

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.SIR as SIR

tmpdir = Path(tempfile.mkdtemp(prefix="campaign_specs_"))
print(f"Working directory for JSON/CSV examples: {tmpdir}")

class Importation(Intervention):
    """Inject `count` infectious individuals into the targeted nodes."""

    @property
    def states(self):
        return ["S", "I", "R"]

    @property
    def properties(self):
        return []

    def execute(self, tick, who, where, params, notes):
        count  = int(params.get("count", 1))
        target = where if where is not None else list(range(len(self.model.scenario)))
        for node in target:
            self.model.states.I[tick + 1][node] += count


class Vaccination(Intervention):
    """Move a binomial-drawn fraction of susceptibles into a V compartment."""

    @property
    def states(self):
        return ["S", "V"]

    @property
    def properties(self):
        return [("new_vaccinations", self.model.params.nticks, np.int32, 0)]

    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)))
        for node in target:
            S_count = int(self.model.states.S[tick + 1][node])
            doses   = int(np.random.binomial(S_count, coverage))
            self.model.states.S[tick + 1][node]           -= doses
            self.model.states.V[tick + 1][node]           += doses
            self.model.nodes.new_vaccinations[tick][node] += doses


Campaign.register(Importation)
Campaign.register(Vaccination)
print("Registered: Importation, Vaccination")

# Disease parameters (measles-like)
R0                    = 10.0
infectious_period     = 7.0
gamma_val             = 1.0 / infectious_period
beta_val              = R0 / infectious_period
life_expectancy_years = 60.0
r_mortality_val       = 1.0 / (life_expectancy_years * 365.0)

# Network coupling
hub_coupling  = 0.05
ring_coupling = 0.01

# Simulation length: 10 years
nticks  = 10 * 365
n_nodes = 6

# Populations: hub + 5 spokes
populations = np.array([500_000] + [50_000] * (n_nodes - 1), dtype=np.int64)

# Schedule cadences shared across the three examples
import_period = 30
vacc_period   = 182
when_imports  = list(range(0, nticks, import_period))
when_vacc     = list(range(0, nticks, vacc_period))

print(f"populations    = {populations.tolist()}")
print(f"import events  = {len(when_imports)} (every {import_period} days)")
print(f"vacc events    = {len(when_vacc)} (every {vacc_period} days)")

schedule_inline = [
    {
        "who": "*",
        "what": "Importation",
        "when": when_imports,
        "where": [0],
        "parameters": {"count": 2},
        "notes": "monthly external importation into the hub",
    },
    {
        "who": ["S"],
        "what": "Vaccination",
        "when": when_vacc,
        "where": "*",
        "parameters": {"coverage": 0.5},
        "notes": "biannual 50% S\u2192V campaign",
    },
]

print(f"Inline schedule: {len(schedule_inline)} entries")
for entry in schedule_inline:
    print(f"  - {entry['what']:<14s} fires {len(entry['when']):>3d}x at {entry['where']}")

model_inline = build_and_run(populations, schedule_inline)
print(f"\nTotal doses : {model_inline.nodes.new_vaccinations.sum():>10,}")
print(f"Total cases : {model_inline.nodes.newly_infectious.sum():>10,}")

json_path = tmpdir / "scenario_a.json"
json_path.write_text(json.dumps(schedule_inline, indent=2))

# Show what the JSON file actually looks like
json_text = json_path.read_text()
preview_lines = json_text.splitlines()
if len(preview_lines) > 24:
    preview = "\n".join(preview_lines[:12] + ["        ... (truncated for display) ..."] + preview_lines[-6:])
else:
    preview = json_text
print(f"--- {json_path.name} ---")
print(preview)
print(f"--- end ({len(json_text):,} bytes, {len(preview_lines)} lines) ---")

# Load directly from the JSON path -- Campaign auto-dispatches on the suffix
model_json = build_and_run(populations, json_path)
print(f"Total doses : {model_json.nodes.new_vaccinations.sum():>10,}")
print(f"Total cases : {model_json.nodes.newly_infectious.sum():>10,}")

def _cell(value):
    """Encode a schedule field for a CSV cell.

    Lists and dicts become JSON-encoded strings; ``"*"`` stays literal; ints
    and bare strings pass through unchanged.
    """
    if isinstance(value, (list, dict)):
        return json.dumps(value)
    return value

csv_path = tmpdir / "scenario_a.csv"
with csv_path.open("w", newline="") as f:
    writer = csv.DictWriter(f, fieldnames=["who", "what", "when", "where", "parameters", "notes"])
    writer.writeheader()
    for entry in schedule_inline:
        writer.writerow({
            "who":        _cell(entry["who"]),
            "what":       entry["what"],
            "when":       _cell(entry["when"]),
            "where":      _cell(entry["where"]),
            "parameters": _cell(entry["parameters"]),
            "notes":      entry["notes"],
        })

csv_text = csv_path.read_text()
# Truncate wide cells for display readability
lines = csv_text.splitlines()
shown = []
for line in lines:
    shown.append(line if len(line) <= 140 else line[:137] + "...")
print(f"--- {csv_path.name} ({len(csv_text):,} bytes) ---")
for line in shown:
    print(line)
print("--- end ---")

model_csv = build_and_run(populations, csv_path)
print(f"Total doses : {model_csv.nodes.new_vaccinations.sum():>10,}")
print(f"Total cases : {model_csv.nodes.newly_infectious.sum():>10,}")

I_inline = model_inline.states.I.sum(axis=1)
I_json   = model_json.states.I.sum(axis=1)
I_csv    = model_csv.states.I.sum(axis=1)

assert np.array_equal(I_inline, I_json), "inline vs JSON diverge"
assert np.array_equal(I_inline, I_csv),  "inline vs CSV  diverge"

print("inline == JSON: bit-for-bit identical")
print("inline == CSV : bit-for-bit identical")

t_years = np.arange(nticks + 1) / 365.0
fig, ax = plt.subplots(figsize=(13, 4))
ax.plot(t_years, I_inline, color="firebrick",   linewidth=0.8, alpha=0.9,                  label="inline")
ax.plot(t_years, I_json,   color="steelblue",   linewidth=0.8, alpha=0.9, linestyle="--",  label="JSON")
ax.plot(t_years, I_csv,    color="forestgreen", linewidth=0.8, alpha=0.9, linestyle=":",   label="CSV")
ax.set_xlabel("Year")
ax.set_ylabel("System-wide I(t)")
ax.set_title("Scenario A under three schedule formats (curves overlap exactly)")
ax.legend()
plt.tight_layout()
plt.show()

visit_period = 60          # days between visits
cycle_period = visit_period * n_nodes   # 360 days to revisit a node

schedule_cycle = [
    {
        "who": "*",
        "what": "Importation",
        "when": when_imports,
        "where": [0],
        "parameters": {"count": 2},
        "notes": "monthly external importation into the hub",
    },
]

for node in range(n_nodes):
    node_ticks = list(range(node * visit_period, nticks, cycle_period))
    label      = "hub" if node == 0 else f"spoke {node}"
    schedule_cycle.append({
        "who":        ["S"],
        "what":       "Vaccination",
        "when":       node_ticks,
        "where":      [node],
        "parameters": {"coverage": 0.5},
        "notes":      f"cycling visit to {label}",
    })

print(f"Cycling schedule: {len(schedule_cycle)} entries")
for entry in schedule_cycle:
    visits = len(entry["when"])
    first  = entry["when"][0] if visits else "-"
    print(f"  - {entry['what']:<14s} where={str(entry['where']):<6s}  {visits:>2d} firings (first @ tick {first})")

model_cycle = build_and_run(populations, schedule_cycle)

print(f"Total doses : {model_cycle.nodes.new_vaccinations.sum():>10,}")
print(f"Total cases : {model_cycle.nodes.newly_infectious.sum():>10,}")
print("Per-node doses delivered:")
for i in range(n_nodes):
    label  = "hub" if i == 0 else f"spoke {i}"
    doses  = int(model_cycle.nodes.new_vaccinations[:, i].sum())
    visits = int((model_cycle.nodes.new_vaccinations[:, i] > 0).sum())
    print(f"  node {i} ({label:<7s})  {doses:>8,} doses across {visits} visits")

nv = model_cycle.nodes.new_vaccinations           # (nticks, n_nodes)

# Extract every (tick, node) pair where doses > 0
event_ticks, event_nodes = np.nonzero(nv)
event_doses              = nv[event_ticks, event_nodes]
event_years              = event_ticks / 365.0

# Per-node color so circles on the same row read as one rotation
node_colors = plt.cm.viridis(np.linspace(0.1, 0.9, n_nodes))
event_colors = [node_colors[n] for n in event_nodes]

# Square-root area scaling — keeps hub and spoke visits both readable.
# Matplotlib `s` is area in points^2; we want radius ∝ √doses, so area ∝ doses.
# To avoid hub circles >10× bigger than spokes (linear-area), scale on √doses instead.
max_dose = event_doses.max()
sizes    = (event_doses / max_dose) ** 0.5 * 1500.0

fig, axes = plt.subplots(2, 1, figsize=(14, 7), sharex=True,
                          gridspec_kw={"height_ratios": [1, 1.4]})

axes[0].scatter(event_years, event_nodes, s=sizes, c=event_colors,
                 edgecolors="black", linewidth=0.4, alpha=0.85)
axes[0].set_yticks(range(n_nodes))
axes[0].set_yticklabels(["hub"] + [f"spoke {i}" for i in range(1, n_nodes)])
axes[0].set_ylabel("Node")
axes[0].set_ylim(-0.6, n_nodes - 0.4)
axes[0].set_title("Cycling vaccination — each circle is one visit, area ∝ √(doses)")
axes[0].grid(True, axis="x", alpha=0.3, linestyle="--")

# Size legend: a few reference dose counts
ref_doses = [1_000, 5_000, 25_000, 100_000]
for d in ref_doses:
    if d <= max_dose * 1.2:
        s = (d / max_dose) ** 0.5 * 1500.0
        axes[0].scatter([], [], s=s, c="lightgray", edgecolors="black",
                         linewidth=0.4, label=f"{d:,} doses")
axes[0].legend(loc="upper right", scatterpoints=1, fontsize=8, framealpha=0.9,
                labelspacing=1.1, borderpad=0.8)

# Bottom: per-node infectious dynamics under cycling
t_years_full = np.arange(nticks + 1) / 365.0
for i in range(n_nodes):
    label = "hub" if i == 0 else f"spoke {i}"
    axes[1].plot(t_years_full, model_cycle.states.I[:, i], linewidth=0.7,
                  color=node_colors[i], label=label)
axes[1].set_xlabel("Year")
axes[1].set_ylabel("I(t)")
axes[1].set_title("Per-node infectious dynamics under cycling vaccination")
axes[1].legend(loc="upper right", ncol=2, fontsize=9)

plt.tight_layout()
plt.show()

print(f"Visits ranged from {event_doses.min():,} to {event_doses.max():,} doses.")

# Per-node total population — needed to convert S to a fraction
N_per_node_cycle = (model_cycle.states.S + model_cycle.states.I
                    + model_cycle.states.R + model_cycle.states.V)

fig, ax = plt.subplots(figsize=(14, 5))
ax.axhline(1.0 / R0, color="black", linestyle="--", linewidth=0.8,
           label=fr"endemic $S^*/N = 1/R_0$ = {1.0/R0:.2f}")
for i in range(n_nodes):
    label = "hub" if i == 0 else f"spoke {i}"
    ax.plot(t_years_full,
             model_cycle.states.S[:, i] / N_per_node_cycle[:, i],
             linewidth=0.8, color=node_colors[i], label=label)
ax.set_xlabel("Year")
ax.set_ylabel("Susceptible fraction S(t) / N(t)")
ax.set_title("Per-node susceptible fraction — cycling vaccination")
ax.set_ylim(bottom=0)
ax.legend(loc="upper right", ncol=2, fontsize=9)
ax.grid(True, axis="x", alpha=0.3, linestyle="--")
plt.tight_layout()
plt.show()

cycle_json_path = tmpdir / "cycling.json"
cycle_json_path.write_text(json.dumps(schedule_cycle, indent=2))

text  = cycle_json_path.read_text()
lines = text.splitlines()
if len(lines) > 30:
    preview = "\n".join(lines[:15] + [f"        ... ({len(lines) - 25} lines hidden) ..."] + lines[-10:])
else:
    preview = text
print(f"--- {cycle_json_path.name} ({len(text):,} bytes, {len(lines)} lines) ---")
print(preview)
print("--- end ---")

model_cycle_json = build_and_run(populations, cycle_json_path)

assert np.array_equal(model_cycle.states.I, model_cycle_json.states.I), "cycling: inline vs JSON diverge"
assert np.array_equal(
    model_cycle.nodes.new_vaccinations,
    model_cycle_json.nodes.new_vaccinations,
), "cycling: dose counts diverge"
print("Cycling inline == cycling JSON: bit-for-bit identical for both I(t) and doses")