Hospital transmission routes

"""Hospital transmission routes: which shared unit, room, procedure or device links a cluster?

Simulates 120 days of admissions at a 7-unit hospital (one large ICU, six
wards) for about 1,500 patients. The EHR feed gives HeapDB one event per
patient-day on a unit (unit, room, pseudonymous care team) and one event per
procedure (procedure type, procedure room, and the endoscope id when one is
used). Genome sequencing places some positive cultures in one cluster; each
is a cluster_infection effect. HeapDB ranks candidate routes per patient
(top_causes) and across the cluster (link_report).

Hidden ground truth (the seeded cause):

* proc_room:ir-2, one of two interventional radiology suites, is
  contaminated. Each procedure done there infects the patient with the
  cluster strain with probability 0.20 (x1.5 for severely ill patients).
  Lag to the positive culture: log-normal, median 5 d, sigma 0.5.
* IR patients come from every unit, so the cases are spread across wards
  and a unit-by-unit review does not see one hot spot.

Distractors and an honest confounder:

* Severity. Severely ill patients start in unit:icu, get more procedures,
  and are more susceptible. They also have more cluster cases with no
  recorded route (2.5% vs 0.4%: missed exposures, community carriage). The
  ICU has the most cases, but it is not the route.
* Common procedures (CT, central lines, endoscopy with 8 scopes) and the
  IR procedure types, which are done in both suites.

Learning: every cluster case is observe()d (no labelled route); infection
prevention adjudicates 5 cases with a confirmed ir-2 exposure via
feedback(); patients whose exposure windows pass without a cluster culture
are settle()d (negative evidence).

Run:
    pip install heapdb
    python 0006-hospital-transmission-routes.py
"""

from __future__ import annotations

import math
import random

from heapdb import HeapDB, LogNormalKernel, format_result

M, H, D = 60.0, 3600.0, 86400.0
SEED = 11
T0 = 1777593600.0        # 2026-05-01T00:00:00Z; the simulation runs in seconds after T0
N_PATIENTS = 1500
DAYS = 120
WARDS = ["unit:4w", "unit:4e", "unit:5w", "unit:5e", "unit:6w", "unit:6e"]
ICU = "unit:icu"
SEEDED = "proc_room:ir-2"
IR_TYPES = ["proc:ir-drain", "proc:ir-picc", "proc:ir-biopsy"]
SCOPES = [f"scope:e{i:02d}" for i in range(1, 9)]

ATTACK = 0.20            # P(cluster infection | one procedure in the contaminated suite)
SEVERE_MULT = 1.5        # severely ill patients are more susceptible
BACKGROUND = {True: 0.025, False: 0.004}   # cluster cases with no recorded route
N_ADJUDICATED = 5        # cases infection prevention reviews and confirms


def lag(rng: random.Random) -> float:
    return math.exp(rng.gauss(math.log(5 * D), 0.5))


def simulate(rng: random.Random):
    """Return (EHR events, first cluster culture per patient, severity, home ward)."""
    events, cultures, severe_of, ward_of = [], {}, {}, {}
    for p in range(N_PATIENTS):
        who = f"pt-{p:04d}"
        severe = rng.random() < 0.25
        ward = rng.choice(WARDS)
        severe_of[who], ward_of[who] = severe, ward
        admit = rng.uniform(0, DAYS - 10) * D
        los = min(30, max(2, round(math.exp(rng.gauss(math.log(7 if severe else 5), 0.5)))))
        icu_days = rng.randint(2, 6) if severe else (rng.randint(1, 3) if rng.random() < 0.05 else 0)
        icu_days = min(icu_days, los - 1)
        first = math.inf
        room = None
        for d in range(los):
            t = admit + d * D + 7 * H
            unit = ICU if d < icu_days else ward
            if d == 0 or d == icu_days:
                n_rooms = 24 if unit == ICU else 16
                room = f"room:{unit.split(':')[1]}-{rng.randint(1, n_rooms):02d}"
            team = f"team:{unit.split(':')[1]}-{'abc'[(d // 3) % 3]}"
            events.append((t, who, "unit_day", {unit, room, team}))
            # Procedures that day (sicker patients get more of everything).
            if rng.random() < (0.15 if severe else 0.06):
                events.append((t + 3 * H, who, "procedure", {"proc:ct", f"proc_room:ct-{rng.randint(1, 2)}"}))
            if severe and d < icu_days and rng.random() < 0.25:
                events.append((t + 2 * H, who, "procedure", {"proc:central-line"}))
            if rng.random() < (0.05 if severe else 0.015):
                suite = rng.choice(["proc_room:ir-1", SEEDED])
                events.append((t + 4 * H, who, "procedure", {rng.choice(IR_TYPES), suite}))
                if suite == SEEDED and rng.random() < ATTACK * (SEVERE_MULT if severe else 1.0):
                    first = min(first, t + 4 * H + lag(rng))
            if rng.random() < 0.02:
                events.append((t + 5 * H, who, "procedure",
                               {"proc:endoscopy", "proc_room:endo-1", rng.choice(SCOPES)}))
        if rng.random() < BACKGROUND[severe]:
            first = min(first, admit + rng.uniform(1, los + 7) * D)
        if first < DAYS * D:
            cultures[who] = first
    events.sort(key=lambda e: e[0])
    return events, cultures, severe_of, ward_of


def make_db() -> HeapDB:
    db = HeapDB(bucket_seconds=D)
    # Exposure-to-culture lag for a hospital-acquired strain: days to a couple of
    # weeks. Log-normal, median 5 d, truncated at median * exp(3 sigma) (~22 d).
    # First cluster culture only (repeatable=False). A large global_pseudocount lets
    # each patient borrow the population's evidence: the cohort view of all exposures.
    db.register_effect("cluster_infection",
                       LogNormalKernel(median=5 * D, sigma=0.5, tail_sigmas=3.0),
                       leak=0.5, repeatable=False, cause_types={"unit_day", "procedure"},
                       global_pseudocount=1e6)
    return db


ROUTE_PREFIXES = ("unit:", "proc:", "proc_room:", "scope:", "team:")


def main() -> None:
    rng = random.Random(SEED)
    events, cultures, severe_of, ward_of = simulate(rng)
    db = make_db()
    stream = [(ev[0], 0, ev) for ev in events] + [(t, 1, who) for who, t in cultures.items()]
    stream.sort(key=lambda x: (x[0], x[1]))

    by_pt: dict[str, list] = {}
    adjudicated, day = [], 0
    for t, kind, item in stream:
        while t >= (day + 1) * D:  # nightly: settle exposure windows that have passed
            day += 1
            db.settle(T0 + day * D)
        if kind == 0:
            _, who, typ, feats = item
            by_pt.setdefault(who, []).append(db.insert(who, T0 + t, typ, feats))
            continue
        who = item
        db.observe(who, "cluster_infection", T0 + t)
        ir2 = [e for e in by_pt.get(who, []) if SEEDED in e.features and 0 < T0 + t - e.time < 20 * D]
        if ir2 and len(adjudicated) < N_ADJUDICATED and day > 30:  # once a review is underway
            # Infection prevention reviews the case and confirms the IR suite exposure.
            db.feedback(who, "cluster_infection", T0 + t, confirmed_event_id=ir2[-1].id)
            adjudicated.append(who)
    db.settle(T0 + (DAYS + 25) * D)

    n_proc = sum(1 for e in events if e[2] == "procedure")
    n_ir2 = sum(1 for e in events if SEEDED in e[3])
    case_wards = {ward_of[w] for w in cultures}
    print("SCENARIO")
    print(f"  {N_PATIENTS} patients, {DAYS} days, {len(events) - n_proc:,} unit-days, "
          f"{n_proc:,} procedures ({n_ir2} in ir-2)")
    print(f"  {len(cultures)} cluster cases from {len(case_wards)} home wards; "
          f"{len(adjudicated)} adjudicated by infection prevention")

    # Baseline: the unit-based review. Which unit did the most cases pass through?
    per_unit: dict[str, int] = {}
    for who in cultures:
        for u in {f for e in by_pt.get(who, []) for f in e.features if f.startswith("unit:")}:
            per_unit[u] = per_unit.get(u, 0) + 1
    top_units = sorted(per_unit.items(), key=lambda kv: (-kv[1], kv[0]))[:4]
    print()
    print("UNIT-BASED REVIEW (cases that passed through each unit)")
    print("  " + ", ".join(f"{u} {n}" for u, n in top_units))

    # One patient, explained: an unadjudicated ward patient whose window includes
    # ICU days and an ir-2 procedure.
    pick = None
    for who in sorted(cultures):
        w = [e for e in by_pt.get(who, []) if 0 < T0 + cultures[who] - e.time < 20 * D]
        fs = {f for e in w for f in e.features}
        if who not in adjudicated and {ICU, SEEDED} <= fs:
            pick = who
            break
    if pick is None:
        pick = sorted(cultures)[0]
    print()
    print(f"ONE PATIENT, EXPLAINED ({pick}, home ward {ward_of[pick]}, "
          f"culture day {cultures[pick] / D:.1f})")
    print(format_result(db.top_causes(pick, "cluster_infection", T0 + cultures[pick],
                                      k=3, explain=True)))

    routes = sorted({f for e in events for f in e[3] if f.startswith(ROUTE_PREFIXES)})
    rep = db.link_report("cluster_infection", min_support=40, features=routes)
    print()
    print("ACROSS ALL CASES (route features, global links)")
    lines = str(rep).splitlines()
    print("\n".join(lines[:10]))

    hit = tot = 0
    for who, t in cultures.items():
        if any(SEEDED in e.features and 0 < T0 + t - e.time < 20 * D for e in by_pt.get(who, [])):
            tot += 1
            top = db.top_causes(who, "cluster_infection", T0 + t, k=1).top
            hit += top is not None and SEEDED in top.event.features
    feats = [r.feature for r in rep]
    seed_rank, icu_rank = feats.index(SEEDED) + 1, feats.index(ICU) + 1
    print()
    print("CHECK")
    print(f"  seeded route ir-2 ranks #{seed_rank} of {len(feats)} route features "
          f"(lift {rep.row(SEEDED).lift:.2f})")
    print(f"  busy {ICU} ranks #{icu_rank} (lift {rep.row(ICU).lift:.2f}), "
          f"though {per_unit[ICU]} of {len(cultures)} cases passed through it")
    print(f"  ir-2 is the top cause for {hit}/{tot} cases with an ir-2 procedure in the window")
    if seed_rank != 1 or icu_rank == 1:
        raise SystemExit(1)


if __name__ == "__main__":
    main()