Outbreak source tracing

"""Outbreak source tracing: which food explains a cluster of illness cases?

Simulates four weeks of shopper-card purchases at four grocery stores (plus a
few restaurant visits) for 1,200 card holders, and the illness onsets that
follow. Investigators add cases one by one. HeapDB ranks suspect foods per
case (top_causes) and across the cluster (link_report).

Hidden ground truth (the seeded cause):

* product:2210-spring-mix, a bagged salad, is contaminated. Each purchase
  makes the buyer ill with probability 0.15. Onset = days until eaten
  (0-1 d) + a log-normal incubation (median 1.5 d, sigma 0.5).
* Only a minority of shoppers buy it. Milk, bread and bananas are in most
  baskets, so most cases bought them too. They are innocent.
* Near-miss: product:3105-ranch-dressing is bought on 60% of salad trips. It
  is innocent but rides along with the salad.
* Noise: 2% of shoppers get sick for unrelated reasons, and 10% of trips are
  paid in cash, so they never reach the purchase log.

Learning: every case is observe()d (onset, no known cause); 8 cases get a
traceback confirmation via feedback(); shoppers whose purchases pass the
incubation window without illness are settle()d (negative evidence).

Run:
    pip install heapdb
    python 0001-outbreak-source-tracing.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 = 7
N_SHOPPERS = 1200
DAYS = 28
T0 = 1785715200.0        # 2026-08-03T00:00:00Z; the simulation runs in seconds after T0
STORES = ["store:01", "store:02", "store:03", "store:04"]

# product id -> (category, brand, probability per trip for a typical shopper)
STAPLES = {
    "product:1001-milk-1gal": ("category:dairy", "brand:hillcrest", 0.55),
    "product:1002-bread-white": ("category:bakery", "brand:oakmill", 0.45),
    "product:1003-bananas": ("category:produce", "brand:none", 0.45),
    "product:1004-eggs-12": ("category:dairy", "brand:hillcrest", 0.30),
    "product:1005-chicken-breast": ("category:meat", "brand:riverfarm", 0.25),
}
SALAD = "product:2210-spring-mix"
DRESSING = "product:3105-ranch-dressing"
SPECIAL = {
    SALAD: ("category:leafy-greens", "brand:freshvale"),
    DRESSING: ("category:condiments", "brand:pantryco"),
    "product:2211-kale-bag": ("category:leafy-greens", "brand:freshvale"),
    "product:2212-romaine-head": ("category:leafy-greens", "brand:none"),
}
OTHER = [f"product:{4000 + i}-item" for i in range(60)]
OTHER_CAT = ["category:pantry", "category:frozen", "category:snacks", "category:household"]
RESTAURANTS = [f"restaurant:r{i:02d}" for i in range(12)]

ATTACK = 0.15            # P(ill | one purchase of the contaminated salad)
SPORADIC = 0.02          # P(ill for unrelated reasons) per shopper
CASH = 0.10              # share of trips paid in cash: never logged
N_CONFIRMED = 8          # cases with a traceback / lab confirmation


def incubation(rng: random.Random) -> float:
    eat_delay = rng.uniform(0, 1) * D
    return eat_delay + math.exp(rng.gauss(math.log(1.5 * D), 0.5))


def simulate(rng: random.Random):
    """Return (logged events, onsets, true cause event key per case)."""
    events, onsets = [], {}
    for s in range(N_SHOPPERS):
        who = f"shopper-{s:04d}"
        salad_fan = rng.random() < 0.15
        home = rng.choice(STORES)
        trips = sorted(rng.uniform(0, DAYS) * D for _ in range(rng.randint(5, 12)))
        first_ill = math.inf
        for t in trips:
            t += rng.uniform(8, 20) * H - (t % D)
            store = home if rng.random() < 0.8 else rng.choice(STORES)
            basket = [p for p, (_, _, q) in STAPLES.items() if rng.random() < q]
            bought_salad = rng.random() < (0.35 if salad_fan else 0.01)
            if bought_salad:
                basket.append(SALAD)
            if rng.random() < (0.60 if bought_salad else 0.05):
                basket.append(DRESSING)
            for p in ("product:2211-kale-bag", "product:2212-romaine-head"):
                if rng.random() < 0.06:
                    basket.append(p)
            basket += rng.sample(OTHER, rng.randint(2, 5))
            cash = rng.random() < CASH
            for i, p in enumerate(basket):
                if p in STAPLES:
                    cat, brand, _ = STAPLES[p]
                elif p in SPECIAL:
                    cat, brand = SPECIAL[p]
                else:
                    n = int(p.split(":")[1].split("-")[0])
                    cat, brand = OTHER_CAT[n % 4], f"brand:b{n % 9}"
                if not cash:
                    events.append((t + i, who, "purchase", {p, cat, brand, store}))
                if p == SALAD and rng.random() < ATTACK:
                    first_ill = min(first_ill, t + incubation(rng))
        for _ in range(rng.randint(0, 3)):
            t = rng.uniform(0, DAYS) * D
            events.append((t, who, "restaurant", {rng.choice(RESTAURANTS)}))
        if rng.random() < SPORADIC:
            first_ill = min(first_ill, rng.uniform(2, DAYS) * D)
        if first_ill < DAYS * D:
            onsets[who] = first_ill
    events.sort(key=lambda e: e[0])
    return events, onsets


def make_db() -> HeapDB:
    db = HeapDB(bucket_seconds=12 * H)
    # Purchase-to-onset lag: Salmonella-like incubation, median 1.5 d. Truncated
    # at median * exp(3 sigma) (about 6.7 d): later purchases cannot explain it.
    # Each case falls ill once (repeatable=False). A large global_pseudocount lets
    # every case borrow the whole population's evidence: the case-control view.
    db.register_effect("illness", LogNormalKernel(median=1.5 * D, sigma=0.5, tail_sigmas=3.0),
                       leak=0.5, repeatable=False, cause_types={"purchase", "restaurant"},
                       global_pseudocount=1e6)
    return db


PRODUCTS = list(STAPLES) + list(SPECIAL) + OTHER


def product_ranking(db: HeapDB, min_support: float = 20):
    """Cluster-level answer: product features ranked by learned link strength."""
    return db.link_report("illness", features=PRODUCTS, min_support=min_support)


def short(feature: str) -> str:
    return feature.split(":", 1)[1]


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

    by_shopper: dict[str, list] = {}
    confirmed, n_cases, day, snapshots = [], 0, 0, []
    checkpoints = {5, 15, 30}
    for t, kind, item in stream:
        while t >= (day + 1) * D:  # daily: settle windows that have passed
            day += 1
            db.settle(T0 + day * D)
        if kind == 0:
            _, who, typ, feats = item
            ev = db.insert(who, T0 + t, typ, feats)
            by_shopper.setdefault(who, []).append(ev)
            continue
        who = item
        db.observe(who, "illness", T0 + t)
        n_cases += 1
        salad_buys = [e for e in by_shopper.get(who, [])
                      if SALAD in e.features and 0 < T0 + t - e.time < 6 * D]
        if salad_buys and len(confirmed) < N_CONFIRMED and n_cases % 3 == 0:
            db.feedback(who, "illness", T0 + t, confirmed_event_id=salad_buys[-1].id)
            confirmed.append(who)
        if n_cases in checkpoints:
            # The standing query: re-rank suspect foods now that a case was added.
            # partial=True: windows still open get negative evidence in proportion
            # to the kernel mass already elapsed; a later onset undoes it.
            db.settle(T0 + t, partial=True)
            rep = product_ranking(db, min_support=3)
            rank = [r.feature for r in rep].index(SALAD) + 1 if rep.row(SALAD) else "-"
            snapshots.append((n_cases, t, rank, [short(r.feature) for r in list(rep)[:3]]))
    db.settle(T0 + DAYS * D + 8 * D)
    rep = product_ranking(db)
    snapshots.append((n_cases, DAYS * D, [r.feature for r in rep].index(SALAD) + 1,
                      [short(r.feature) for r in list(rep)[:3]]))

    n_purch = sum(1 for e in events if e[2] == "purchase")
    buyers = {e[1] for e in events if SALAD in e[3]}
    print("SCENARIO")
    print(f"  {N_SHOPPERS} card holders, {DAYS} days, {n_purch:,} logged purchases "
          f"(cash trips missing), {len(events) - n_purch} restaurant visits")
    print(f"  {n_cases} illness cases; {len(buyers)} shoppers have a logged spring-mix purchase; "
          f"{len(confirmed)} cases confirmed by traceback")
    share = []
    for p in (SALAD, "product:1001-milk-1gal", "product:1002-bread-white", "product:1003-bananas"):
        n = sum(any(p in e.features and 0 < T0 + t - e.time < 7 * D for e in by_shopper.get(w, []))
                for w, t in onsets.items())
        share.append(f"{short(p)} {n}")
    print(f"  cases with a logged purchase in the 7 days before onset: {', '.join(share)}")

    print()
    print("STANDING QUERY: SUSPECT FOODS AS CASES ARRIVE")
    for n, t, rank, top in snapshots:
        print(f"  {n:>3} cases, day {t / D:4.1f}: spring-mix #{rank:<3} top 3: {', '.join(top)}")

    # One case, explained: a case without a lab confirmation who bought the
    # salad and the staples in the window.
    pick = None
    for who in sorted(onsets):
        if who in confirmed:
            continue
        w = [e for e in by_shopper.get(who, []) if 0 < T0 + onsets[who] - e.time < 6 * D]
        fs = [f for e in w for f in e.features]
        if fs.count(SALAD) == 1 and {DRESSING, "product:1001-milk-1gal"} <= set(fs):
            pick = who
            break
    print()
    print(f"ONE CASE, EXPLAINED ({pick}, onset day {onsets[pick] / D:.1f})")
    print(format_result(db.top_causes(pick, "illness", T0 + onsets[pick], k=3, explain=True)))

    print()
    print("ACROSS ALL CASES (product features, global links)")
    print("\n".join(str(rep).splitlines()[:10]))  # header + top 8 rows

    # Per-case accuracy: is a spring-mix purchase the top cause?
    hit = tot = 0
    for who, t in onsets.items():
        if any(SALAD in e.features and 0 < T0 + t - e.time < 6 * D for e in by_shopper.get(who, [])):
            tot += 1
            top = db.top_causes(who, "illness", T0 + t, k=1).top
            hit += top is not None and SALAD in top.event.features
    feats = [r.feature for r in rep]
    salad_rank, dress_rank = feats.index(SALAD) + 1, feats.index(DRESSING) + 1
    milk_rank = feats.index("product:1001-milk-1gal") + 1
    print()
    print("CHECK")
    print(f"  seeded cause spring-mix ranks #{salad_rank} of {len(feats)} products "
          f"(lift {rep.row(SALAD).lift:.2f}); near-miss ranch-dressing #{dress_rank} "
          f"(lift {rep.row(DRESSING).lift:.2f}); milk #{milk_rank}")
    print(f"  spring-mix purchase is the top cause for {hit}/{tot} cases with a logged "
          f"spring-mix purchase in the window")
    if salad_rank != 1 or dress_rank == 1:
        raise SystemExit(1)


if __name__ == "__main__":
    main()