From b57deb4c43b993e833eb585c9313068ae0a02b1c Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Fri, 9 Oct 2026 08:02:40 -0400 Subject: [PATCH 1/4] Add --override-positions visit_end to the label-override test The concept head is supervised only at visit ends, yet the override was applied at every position with the running label (Appendix E names this as a possible reason the true value does not help). The new option applies truth/flip/random and the calibrated modes only at each visit's last position (chunk.visit_end, the flag the concept loss pools on) and leaves every other position as under none. Each mode gains a visit_end_positions block: top-1 and loss over only the visit-end positions with an observed label inside the band and a next-token target, n_positions and n_subjects, and on the left-hand mode the paired subject-clustered bootstrap differences (truth minus none, truth minus flip, flip minus none) in top-1 and loss. The default (all) is unchanged byte for byte; --hazard-heads combines with it (the landmark rows are not visit ends, documented). WP1 of the rebuttal plan gains the VM command. Co-Authored-By: Claude Fable 5.1 --- docs/ml4h2026_rebuttal_plan.md | 10 + odyssey/inference/interventions.py | 429 +++++++++++++++- .../inference/test_interventions_visit_end.py | 465 ++++++++++++++++++ 3 files changed, 886 insertions(+), 18 deletions(-) create mode 100644 tests/odyssey/inference/test_interventions_visit_end.py diff --git a/docs/ml4h2026_rebuttal_plan.md b/docs/ml4h2026_rebuttal_plan.md index 589a795f..3bb5799a 100644 --- a/docs/ml4h2026_rebuttal_plan.md +++ b/docs/ml4h2026_rebuttal_plan.md @@ -49,6 +49,16 @@ Each package says what exists, what to build, who runs it, and where the result Set `--max-shards` to the split's full shard count on each VM (37 is MIMIC's). The transformer arm keeps the lever test's TBTT view (whole history, no context truncation), so its landmark rows are the model-free set, not the packed-context set `alerts.py` scores for that backbone. - Bank: `research_journal/figure_data/{vm1,vm2}//interventions_hazard.json`. New table generator `scripts/make_lever_hazard_table.py`. +- Built 2026-10-09 (branch `rebuttal/visit-end-override`): `--override-positions visit_end` answers the reviewer's "does the true value help when the override is applied only where the concept head was supervised". The concept loss pools at each visit's last event (`chunk.visit_end`, the tokenizer's `visit_ends` flag), while the default `all` overrides every position with the running label; under `visit_end` truth/flip/random (and the calibrated modes) replace the probability only at those visit ends and every other position is the `none` model. Each mode's entry gains a `visit_end_positions` block ({n_positions, n_subjects, top1_accuracy, mean_task_loss, paired}) scored over only the visit ends that have an observed label inside the band and a next-token target (a stream's last visit has none), with `paired` on the left-hand mode (truth carries truth_minus_none and truth_minus_flip, flip carries flip_minus_none) holding 95% subject-clustered paired bootstrap intervals on the top-1 and loss differences (`--hazard-boot` resamples). The whole-stream top-1/loss are still reported. `--hazard-heads` may be combined, but the landmark rows are the first event of each 4 h bucket, not visit ends, so under `visit_end` the hazard block is mostly the un-overridden model and is not the headline. The default output is unchanged byte for byte. Command on each VM: + + ``` + python -m odyssey.inference.interventions --run-dir R --held-out-shard-dir D/held_out \ + --output-json R/interventions_band15_visitend.json --max-shards 37 --num-lanes 64 \ + --chunk-size 512 --uncertain-band 0.15 --modes none truth flip random \ + --override-positions visit_end --checkpoint checkpoint_best.pt + ``` + + Bank as `research_journal/figure_data/{vm1,vm2}//interventions_visitend.json`. Rebuttal use: if truth helps on the visit-end subset, Appendix E's explanation (the head was never trained to receive an override mid-visit) stands and the lever claim is scoped to supervised positions; if it does not, the position mismatch is ruled out as the reason. - Rebuttal use: if truth moves the 24 h death or vasopressor hazard the right way, the lever verdict changes and the paper's Q3 gets a real endpoint. If it does not, the negative result becomes like-for-like with the edit test, which is what the review asked for. Either way it answers W1. - GEMINI: same script, run by Amrit, only if time. diff --git a/odyssey/inference/interventions.py b/odyssey/inference/interventions.py index c6835303..3c3f6996 100644 --- a/odyssey/inference/interventions.py +++ b/odyssey/inference/interventions.py @@ -84,6 +84,25 @@ left-hand mode's entry. For ``flip_gated`` the heads are read from the flip-intervened features: the logit gate has no hazard analogue. Off by default, so an existing run reproduces byte for byte. + +Override positions (``--override-positions``). The concept head is +supervised only at visit ends (the concept loss pools +``chunk.visit_end``, see ``ConceptBottleneckSequenceModel.compute_streaming_loss``), +yet the default ``all`` applies the override at every position. Under +``visit_end`` the substitution (truth, flip, flip_gated, random and the +calibrated modes) is applied only at the last position of each real +visit (:func:`visit_end_positions`, the tokenizer's ``visit_ends`` flag +carried on the chunk, the same flag the loss and the alert protocol's +``visit_end`` index mode use); every other position is left exactly as +under ``none``. The zero modes edit embeddings, not probabilities, and +are not restricted. Each result then carries a ``visit_end_positions`` +block: top-1 and loss over only the visit-end positions that have an +observed label inside the band and a next-token target (the last visit +of a stream has none), with ``n_positions`` and ``n_subjects``, and on +the left-hand mode the :data:`HAZARD_PAIRS` paired subject-clustered +bootstrap differences in top-1 and loss. With ``--hazard-heads`` the +hazard block is still read at the landmark rows, which are not visit +ends, so under ``visit_end`` it mostly reflects the un-overridden model. """ import json @@ -173,6 +192,10 @@ CALIBRATED_MODES = ("truth_calibrated", "flip_calibrated") +#: Where the label override is applied: every position, or only the last +#: position of each real visit (where the concept head was supervised). +OVERRIDE_POSITIONS = ("all", "visit_end") + #: Paired hazard differences, (left, right), reported on the left mode's entry. HAZARD_PAIRS: tuple[tuple[str, str], ...] = ( ("truth", "none"), @@ -238,15 +261,24 @@ class InterventionResult: each with a subject-clustered paired bootstrap interval. ``{}`` for a mode that is never a left operand; None when hazard scoring is off.""" + visit_end_positions: dict[str, Any] | None = None + """``{n_positions, n_subjects, top1_accuracy, mean_task_loss, paired}`` + over only the visit-end positions (see :class:`VisitEndScorer`) when + ``override_positions="visit_end"``; ``paired`` holds the + :data:`HAZARD_PAIRS` differences whose left mode is this one, each a + subject-clustered paired bootstrap of the top-1 and loss difference. + None under ``override_positions="all"``.""" + def result_to_json(result: InterventionResult) -> dict[str, Any]: - """Serialise one result; the hazard keys appear only when scored. + """Serialise one result; the hazard and visit-end keys appear only when scored. - With ``--hazard-heads`` off the dict is exactly what earlier versions - wrote, so banked JSONs stay byte for byte reproducible. + With ``--hazard-heads`` off and ``--override-positions all`` the dict + is exactly what earlier versions wrote, so banked JSONs stay byte for + byte reproducible. """ out = asdict(result) - for key in ("hazard", "hazard_paired"): + for key in ("hazard", "hazard_paired", "visit_end_positions"): if out[key] is None: del out[key] return out @@ -594,6 +626,190 @@ def score_hazard_landmarks( return per_mode, paired +# --------------------------------------------------------------------------- +# Visit-end override positions +# --------------------------------------------------------------------------- + + +def visit_end_positions(chunk: StreamingChunk) -> torch.Tensor: + """``(lanes, T)`` bool: the last position of each real (hadm-bearing) visit. + + The tokenizer marks it (``PatientSequence.visit_ends``) and the + sampler carries it per chunk as ``chunk.visit_end``; it is the flag + visit-scoped concept supervision pools on and the one the alert + protocol's ``visit_end`` index mode scores, so "where the concept head + was supervised" has one definition. Events outside any visit + (``visit_ids < 0``) are never visit ends. + """ + return chunk.visit_end & (chunk.visit_ids >= 0) + + +@dataclass(frozen=True) +class VisitEndRowTable: + """Next-event outcome at every scored visit-end position of one pass.""" + + subject_ids: np.ndarray + visit_ids: np.ndarray + time_hours: np.ndarray + hits: np.ndarray + """``(n_rows,)`` float64, 1.0 where the top-1 prediction was right.""" + + losses: np.ndarray + """``(n_rows,)`` float64 next-token cross-entropy per row.""" + + @property + def n_rows(self) -> int: + """Number of scored visit-end rows.""" + return int(self.subject_ids.shape[0]) + + @property + def n_subjects(self) -> int: + """Number of distinct subjects among the rows.""" + return int(np.unique(self.subject_ids).shape[0]) + + def same_rows(self, other: "VisitEndRowTable") -> bool: + """Whether both tables hold the same (subject, visit, time) rows in order.""" + return ( + np.array_equal(self.subject_ids, other.subject_ids) + and np.array_equal(self.visit_ids, other.visit_ids) + and np.array_equal(self.time_hours, other.time_hours) + ) + + +class VisitEndScorer: + """Keep the next-event outcome at the visit-end positions of one pass. + + A row is a visit-end position (:func:`visit_end_positions`) with at + least one observed concept whose own probability lies inside the + uncertain band (every observed concept when there is no band) and a + real next-token target. That is exactly the position set the + ``visit_end`` override edits, minus the positions without a target + (the last visit of a stream, whose next token is another patient's). + The set is defined from the model's OWN probabilities, which no mode + changes, so every mode scores the same rows and they can be paired. + """ + + def __init__(self, *, uncertain_band: float | None) -> None: + """Bind the band that defines "observed and inside the band".""" + self.uncertain_band = uncertain_band + self._sids: list[np.ndarray] = [] + self._vids: list[np.ndarray] = [] + self._times: list[np.ndarray] = [] + self._hits: list[np.ndarray] = [] + self._losses: list[np.ndarray] = [] + + def eligible( + self, chunk: StreamingChunk, observed: torch.Tensor, own_probs: torch.Tensor + ) -> torch.Tensor: + """``(lanes, T)`` bool: the rows this chunk contributes.""" + entries = observed.bool() & visit_end_positions(chunk).unsqueeze(-1) + if self.uncertain_band is not None: + entries = entries & ((own_probs - 0.5).abs() < self.uncertain_band) + return entries.any(dim=-1) & chunk.real_mask + + def add_chunk( + self, chunk: StreamingChunk, logits: torch.Tensor, keep: torch.Tensor + ) -> None: + """Record top-1 hit and loss at the ``keep`` positions of this chunk.""" + if not keep.any(): + return + targets = chunk.targets[keep] + kept_logits = logits[keep] + hits = kept_logits.argmax(dim=-1) == targets + losses = F.cross_entropy( + kept_logits, targets, ignore_index=PAD_ID, reduction="none" + ) + self._sids.append(chunk.subject_ids[keep].cpu().numpy().astype(np.int64)) + self._vids.append(chunk.visit_ids[keep].cpu().numpy().astype(np.int64)) + self._times.append( + chunk.batch.aux.time_stamps[keep].cpu().numpy().astype(np.float64) + ) + self._hits.append(hits.cpu().numpy().astype(np.float64)) + self._losses.append(losses.cpu().numpy().astype(np.float64)) + + def table(self) -> VisitEndRowTable: + """Concatenate everything accumulated so far.""" + + def cat(parts: list[np.ndarray], dtype: type) -> np.ndarray: + return np.concatenate(parts) if parts else np.zeros(0, dtype) + + return VisitEndRowTable( + subject_ids=cat(self._sids, np.int64), + visit_ids=cat(self._vids, np.int64), + time_hours=cat(self._times, np.float64), + hits=cat(self._hits, np.float64), + losses=cat(self._losses, np.float64), + ) + + +def visit_end_summary(table: VisitEndRowTable) -> dict[str, Any]: + """Point top-1 and loss over the visit-end rows, with the subset size.""" + n = table.n_rows + return { + "n_positions": n, + "n_subjects": table.n_subjects, + "top1_accuracy": float(table.hits.mean()) if n else float("nan"), + "mean_task_loss": float(table.losses.mean()) if n else float("nan"), + "paired": {}, + } + + +def _mean_delta_json(d: PairedMeanDelta) -> dict[str, Any]: + return { + "point": d.point, + "ci_low": d.ci_low, + "ci_high": d.ci_high, + "separated": d.separated, + } + + +def visit_end_paired_summary( + left: VisitEndRowTable, + right: VisitEndRowTable, + *, + n_boot: int, + seed: int, +) -> dict[str, Any]: + """Paired subject-clustered bootstrap of the top-1 and loss differences.""" + if not left.same_rows(right): + raise ValueError("paired visit-end scoring needs identical rows") + subjects = left.subject_ids + top1 = paired_mean_delta(left.hits, right.hits, subjects, n_boot=n_boot, seed=seed) + loss = paired_mean_delta( + left.losses, right.losses, subjects, n_boot=n_boot, seed=seed + ) + return { + "top1_accuracy": _mean_delta_json(top1), + "mean_task_loss": _mean_delta_json(loss), + "n_positions": top1.n_rows, + "n_subjects": top1.n_subjects, + } + + +def score_visit_end_rows( + tables: dict[str, VisitEndRowTable], *, n_boot: int = 1000, seed: int = 0 +) -> dict[str, dict[str, Any]]: + """``{mode: visit_end_positions block}`` with the :data:`HAZARD_PAIRS` attached. + + Every table must hold the same rows, which the shared sampler and the + mode-independent row rule guarantee. Pairs land on the left mode's + ``paired`` entry, keyed ``_minus_``. + """ + if not tables: + return {} + first = next(iter(tables.values())) + for mode, table in tables.items(): + if not table.same_rows(first): + raise ValueError(f"mode {mode!r} scored different visit-end rows") + blocks = {mode: visit_end_summary(table) for mode, table in tables.items()} + for left, right in HAZARD_PAIRS: + if left in tables and right in tables: + blocks[left]["paired"][f"{left}_minus_{right}"] = visit_end_paired_summary( + tables[left], tables[right], n_boot=n_boot, seed=seed + ) + return blocks + + def _chunk_intervention( chunk: StreamingChunk, mode: str, @@ -606,8 +822,13 @@ def _chunk_intervention( device: str, rng: torch.Generator, uncertain_band: float | None = None, + position_mask: torch.Tensor | None = None, ) -> BottleneckIntervention | None: - """Build the per-position intervention for one chunk, or None.""" + """Build the per-position intervention for one chunk, or None. + + ``position_mask`` (``(lanes, T)`` bool), if given, restricts the + substitution to those positions on top of the observed mask. + """ if mode == "none": return None if mode == "zero_known": @@ -625,6 +846,8 @@ def _chunk_intervention( supervision=supervision, num_concepts=num_concepts, ) + if position_mask is not None: + observed = observed * position_mask.to(observed.device).unsqueeze(-1) if mode == "truth": values = labels elif mode in ("flip", "flip_gated"): @@ -661,6 +884,8 @@ def run_streaming_intervention( # noqa: PLR0912, PLR0915 -- one linear scoring calibrated_tau: float | None = None, concept_names: Sequence[str] | None = None, hazard_scorer: HazardLandmarkScorer | None = None, + override_positions: str = "all", + visit_end_scorer: VisitEndScorer | None = None, ) -> InterventionResult: """Score next-event prediction under one intervention mode. @@ -675,6 +900,13 @@ def run_streaming_intervention( # noqa: PLR0912, PLR0915 -- one linear scoring :class:`HazardLandmarkScorer`); it changes nothing in the returned next-event numbers. + ``override_positions="visit_end"`` applies the probability override + only at :func:`visit_end_positions` (see the module docstring) and + fills ``InterventionResult.visit_end_positions`` from + ``visit_end_scorer`` (one is created when not given; pass your own to + keep the row table for pairing across modes). ``"all"`` is the + unchanged default. + The identical streaming pass as :func:`~odyssey.inference.run_inference.run_streaming_inference` (same sampler, same state carrying), with the bottleneck edited per @@ -695,6 +927,14 @@ def run_streaming_intervention( # noqa: PLR0912, PLR0915 -- one linear scoring f"mode {mode!r} needs calibration_gammas (see " "odyssey.inference.concept_attribution.calibrated_gammas)" ) + if override_positions not in OVERRIDE_POSITIONS: + raise ValueError( + f"unknown override_positions {override_positions!r}; " + f"known: {OVERRIDE_POSITIONS}" + ) + at_visit_ends = override_positions == "visit_end" + if at_visit_ends and visit_end_scorer is None: + visit_end_scorer = VisitEndScorer(uncertain_band=uncertain_band) if concept_first_times is None: if mode in ("truth", "flip", "flip_gated", *CALIBRATED_MODES): logger.warning( @@ -731,6 +971,7 @@ def run_streaming_intervention( # noqa: PLR0912, PLR0915 -- one linear scoring with torch.no_grad(): for chunk in sampler: chunk = move_to_device(chunk, device) # noqa: PLW2901 + position_mask = visit_end_positions(chunk) if at_visit_ends else None intervention = ( None if mode in CALIBRATED_MODES # built below, from the model's own probs @@ -745,6 +986,7 @@ def run_streaming_intervention( # noqa: PLR0912, PLR0915 -- one linear scoring device=device, rng=rng, uncertain_band=uncertain_band, + position_mask=position_mask, ) ) if mode in CALIBRATED_MODES: @@ -766,6 +1008,10 @@ def run_streaming_intervention( # noqa: PLR0912, PLR0915 -- one linear scoring supervision=supervision, num_concepts=num_concepts, ) + if position_mask is not None: + observed = observed * position_mask.to(observed.device).unsqueeze( + -1 + ) pole = labels if mode == "truth_calibrated" else 1.0 - labels # calibration_gammas is derived from the model's own LM-head # weights, so it lives on the model's device, while `pole` @@ -812,6 +1058,22 @@ def run_streaming_intervention( # noqa: PLR0912, PLR0915 -- one linear scoring hazard_features = bottleneck_out.bottleneck if hazard_scorer is not None: hazard_scorer.add_chunk(chunk, hazard_features) + if visit_end_scorer is not None: + # The row set comes from the OBSERVED mask and the model's + # own probabilities, both the same in every mode, so the + # same positions are kept whether or not this mode edits. + _, observed_all = position_running_labels( + chunk, + concept_labels, + concept_mask, + concept_first_times, + supervision=supervision, + num_concepts=num_concepts, + ) + keep = visit_end_scorer.eligible( + chunk, observed_all, bottleneck_out.concept_probs + ) + visit_end_scorer.add_chunk(chunk, logits, keep) real = chunk.real_mask if intervention is not None and intervention.probs is not None: own = bottleneck_out.concept_probs @@ -902,6 +1164,11 @@ def run_streaming_intervention( # noqa: PLR0912, PLR0915 -- one linear scoring calibrated_tau=calibrated_tau if mode in CALIBRATED_MODES else None, n_replaced_by_concept=by_concept_n, mean_abs_displacement_by_concept=by_concept_disp, + visit_end_positions=( + visit_end_summary(visit_end_scorer.table()) + if visit_end_scorer is not None + else None + ), ) @@ -923,6 +1190,7 @@ def evaluate_interventions( hazard_boot: int = 1000, landmark_hours: float = 4.0, horizons: Sequence[float] = HORIZONS_HOURS, + override_positions: str = "all", ) -> list[InterventionResult]: """End-to-end: load a trained run, score every intervention mode. @@ -938,6 +1206,14 @@ def evaluate_interventions( landmark rows follow :mod:`odyssey.inference.alerts` exactly. The bootstrap seed is ``seed``. + ``override_positions="visit_end"`` restricts the override to visit + ends and adds the ``visit_end_positions`` block to every result, with + the :data:`HAZARD_PAIRS` paired differences (``hazard_boot`` + resamples, seed ``seed``) on the left mode's entry. The hazard block, + if also on, is still read at the landmark rows: those are the first + event of each 4 h bucket, not visit ends, so under ``visit_end`` the + hazard readout is mostly the un-overridden model's. + Data preparation matches :func:`~odyssey.inference.run_inference.evaluate_run` exactly (same normalization, binning, and label scoping from the run's own @@ -1016,8 +1292,11 @@ def evaluate_interventions( device=device, ) + _check_override_positions(override_positions, supervision) + results = [] tables: dict[str, LandmarkRiskTable] = {} + visit_end_tables: dict[str, VisitEndRowTable] = {} for mode in modes: logger.info("[interventions] scoring mode %r", mode) mode_subjects: dict[int, list[int]] | None = None @@ -1032,6 +1311,11 @@ def evaluate_interventions( landmark_hours=landmark_hours, horizons=horizons, ) + visit_end_scorer = ( + VisitEndScorer(uncertain_band=uncertain_band) + if override_positions == "visit_end" + else None + ) result = run_streaming_intervention( model, events_binned, @@ -1051,11 +1335,15 @@ def evaluate_interventions( calibrated_tau=calibrated_tau, concept_names=[c.name for c in concepts], hazard_scorer=scorer, + override_positions=override_positions, + visit_end_scorer=visit_end_scorer, ) if mode in CALIBRATED_MODES: result = replace(result, calibration_gamma=gamma_by_name) if scorer is not None: tables[mode] = scorer.table() + if visit_end_scorer is not None: + visit_end_tables[mode] = visit_end_scorer.table() results.append(result) baseline = results[0] latest = results[-1] @@ -1066,23 +1354,104 @@ def evaluate_interventions( latest.top1_accuracy - baseline.top1_accuracy, latest.mean_task_loss, ) - if hazard_targets is not None: - logger.info( - "[interventions] paired hazard bootstrap: %d resamples over %d landmark rows", - hazard_boot, - next(iter(tables.values())).n_rows if tables else 0, + return _attach_paired_blocks( + results, + hazard_tables=tables, + hazard_times=hazard_targets.times if hazard_targets is not None else None, + visit_end_tables=visit_end_tables, + n_boot=hazard_boot, + seed=seed, + ) + + +def _attach_paired_blocks( + results: list[InterventionResult], + *, + hazard_tables: dict[str, LandmarkRiskTable], + hazard_times: dict[str, EventTimes] | None, + visit_end_tables: dict[str, VisitEndRowTable], + n_boot: int, + seed: int, +) -> list[InterventionResult]: + """Attach the hazard and visit-end blocks that were scored, if any.""" + if hazard_times is not None: + results = _attach_hazard_blocks( + results, hazard_tables, hazard_times, n_boot=n_boot, seed=seed ) - per_mode, paired = score_hazard_landmarks( - tables, hazard_targets.times, n_boot=hazard_boot, seed=seed + if visit_end_tables: + results = _attach_visit_end_blocks( + results, visit_end_tables, n_boot=n_boot, seed=seed ) - results = [ - replace(r, hazard=per_mode[r.mode], hazard_paired=paired[r.mode]) - for r in results - ] - _log_hazard(results) return results +def _attach_hazard_blocks( + results: Sequence[InterventionResult], + tables: dict[str, LandmarkRiskTable], + times: dict[str, EventTimes], + *, + n_boot: int, + seed: int, +) -> list[InterventionResult]: + """Score the landmark tables, pair them, and attach the blocks per mode.""" + logger.info( + "[interventions] paired hazard bootstrap: %d resamples over %d landmark rows", + n_boot, + next(iter(tables.values())).n_rows if tables else 0, + ) + per_mode, paired = score_hazard_landmarks(tables, times, n_boot=n_boot, seed=seed) + out = [ + replace(r, hazard=per_mode[r.mode], hazard_paired=paired[r.mode]) + for r in results + ] + _log_hazard(out) + return out + + +def _check_override_positions( + override_positions: str, supervision: ConceptSupervision +) -> None: + if override_positions not in OVERRIDE_POSITIONS: + raise ValueError( + f"unknown override_positions {override_positions!r}; " + f"known: {OVERRIDE_POSITIONS}" + ) + if override_positions == "visit_end" and supervision != "visit": + logger.warning( + "[interventions] override_positions='visit_end' on a run with " + "concept_supervision=%r: the concept head was supervised at stay " + "ends, which are the last visit end of each stream only", + supervision, + ) + + +def _attach_visit_end_blocks( + results: Sequence[InterventionResult], + tables: dict[str, VisitEndRowTable], + *, + n_boot: int, + seed: int, +) -> list[InterventionResult]: + """Score the visit-end rows, pair them, and attach the block per mode.""" + logger.info( + "[interventions] paired visit-end bootstrap: %d resamples over %d rows", + n_boot, + next(iter(tables.values())).n_rows, + ) + blocks = score_visit_end_rows(tables, n_boot=n_boot, seed=seed) + out = [replace(r, visit_end_positions=blocks[r.mode]) for r in results] + for r in out: + block = r.visit_end_positions or {} + logger.info( + "[interventions] %s at visit ends: top1 %.4f, loss %.4f (n=%d)", + r.mode, + block["top1_accuracy"], + block["mean_task_loss"], + block["n_positions"], + ) + return out + + def _concept_labels( raw_events: pl.DataFrame, concepts: Sequence[Any], supervision: ConceptSupervision ) -> tuple[ConceptLabelDict, ConceptLabelDict, ConceptLabelDict]: @@ -1249,7 +1618,23 @@ def _main() -> None: "--hazard-boot", type=int, default=1000, - help="bootstrap resamples for the paired hazard differences.", + help=( + "bootstrap resamples for the paired hazard differences and for " + "the paired visit-end differences." + ), + ) + parser.add_argument( + "--override-positions", + choices=OVERRIDE_POSITIONS, + default="all", + help=( + "where truth/flip/random (and the calibrated modes) replace the " + "concept probability: 'all' positions (default, unchanged output) " + "or only each visit's last position ('visit_end'), where the " + "concept head was supervised; visit_end also adds a " + "visit_end_positions block (top-1 and loss over those positions " + "with paired subject-clustered bootstrap differences)." + ), ) parser.add_argument( "--dump-per-subject", @@ -1293,6 +1678,7 @@ def _main() -> None: calibrated_tau=args.calibrated_tau, hazard_heads=args.hazard_heads, hazard_boot=args.hazard_boot, + override_positions=args.override_positions, ) out.parent.mkdir(parents=True, exist_ok=True) out.write_text(json.dumps([result_to_json(r) for r in results], indent=2)) @@ -1321,10 +1707,13 @@ def _main() -> None: __all__ = [ "HAZARD_PAIRS", "INTERVENTION_MODES", + "OVERRIDE_POSITIONS", "HazardLandmarkScorer", "InterventionResult", "LandmarkRiskTable", "PairedMeanDelta", + "VisitEndRowTable", + "VisitEndScorer", "evaluate_interventions", "hazard_paired_summary", "hazard_summary", @@ -1333,5 +1722,9 @@ def _main() -> None: "result_to_json", "run_streaming_intervention", "score_hazard_landmarks", + "score_visit_end_rows", "subject_bootstrap_means", + "visit_end_paired_summary", + "visit_end_positions", + "visit_end_summary", ] diff --git a/tests/odyssey/inference/test_interventions_visit_end.py b/tests/odyssey/inference/test_interventions_visit_end.py new file mode 100644 index 00000000..7f518e6b --- /dev/null +++ b/tests/odyssey/inference/test_interventions_visit_end.py @@ -0,0 +1,465 @@ +"""Label override applied only at visit ends (CPU, tiny fixtures). + +The concept head is supervised at each visit's last event; the default +override edits every position. ``override_positions="visit_end"`` edits +only the supervised positions and scores them on their own. +""" + +import json +from datetime import timedelta +from pathlib import Path + +import numpy as np +import polars as pl +import pytest +import torch + +import odyssey.inference.interventions as interventions_module +from odyssey.data.streaming import PackedLaneSampler, StreamingChunk +from odyssey.inference.interventions import ( + OVERRIDE_POSITIONS, + InterventionResult, + VisitEndScorer, + _chunk_intervention, + result_to_json, + run_streaming_intervention, + score_visit_end_rows, + visit_end_paired_summary, + visit_end_positions, +) +from odyssey.models.concept_bottleneck import intervention_apply_mask +from odyssey.training.data import iter_patient_sequences +from odyssey.training.running_labels import position_running_labels +from tests.odyssey.inference.test_interventions import ( + CODES, + NUM_CONCEPTS, + T0, + _model, + _vocab, +) +from tests.odyssey.inference.test_interventions_hazard import OLD_KEYS +from tests.odyssey.inference.test_interventions_hazard import _Fixture as _HazardFixture + + +SUBJECTS = (1, 2, 3) +EVENTS_PER_VISIT = 10 + + +def _events() -> pl.DataFrame: + """Three subjects, two hourly visits each (hadm 10*sid + {1, 2}).""" + rows = [] + for sid in SUBJECTS: + for visit in (1, 2): + for i in range(EVENTS_PER_VISIT): + hour = (visit - 1) * EVENTS_PER_VISIT + i + rows.append( + ( + sid, + CODES[(sid * 3 + hour) % len(CODES)], + T0 + timedelta(hours=hour), + None, + 10 * sid + visit, + ) + ) + return pl.DataFrame( + rows, + schema={ + "subject_id": pl.Int64, + "code": pl.Utf8, + "time": pl.Datetime, + "numeric_value": pl.Float32, + "hadm_id": pl.Int64, + }, + orient="row", + ) + + +def _visit_labels() -> tuple[ + dict[tuple[int, int], torch.Tensor], + dict[tuple[int, int], torch.Tensor], + dict[tuple[int, int], torch.Tensor], +]: + """Visit-scoped labels, masks and first times; every visit observed.""" + labels: dict[tuple[int, int], torch.Tensor] = {} + masks: dict[tuple[int, int], torch.Tensor] = {} + first: dict[tuple[int, int], torch.Tensor] = {} + for sid in SUBJECTS: + for visit in (1, 2): + key = (sid, 10 * sid + visit) + labels[key] = torch.tensor([float(visit == 1), 1.0, float(sid == 2)]) + masks[key] = torch.ones(NUM_CONCEPTS) + first[key] = torch.where( + labels[key] > 0, torch.tensor(0.0), torch.tensor(float("inf")) + ) + return labels, masks, first + + +def _run( + mode: str, + *, + override_positions: str | None = None, + uncertain_band: float | None = None, + scorer: VisitEndScorer | None = None, +) -> InterventionResult: + vocab = _vocab() + labels, masks, first = _visit_labels() + kwargs = ( + {} if override_positions is None else {"override_positions": override_positions} + ) + return run_streaming_intervention( + _model(len(vocab)), + _events(), + vocab, + labels, + masks, + mode=mode, + concept_first_times=first, + supervision="visit", + num_lanes=2, + chunk_size=8, + device="cpu", + seed=0, + uncertain_band=uncertain_band, + visit_end_scorer=scorer, + **kwargs, + ) + + +def _lane_chunks() -> list[StreamingChunk]: + """Stream the whole fixture in one lane: one chunk per 20-event patient.""" + seqs = iter_patient_sequences(_events(), _vocab()) + sampler = PackedLaneSampler(seqs, num_lanes=1, chunk_size=20, reset_prob=0.0) + return list(sampler) + + +# --------------------------------------------------------------------------- +# (i) the default reproduces the old output exactly +# --------------------------------------------------------------------------- + + +def test_default_override_positions_is_all_and_keeps_the_old_schema() -> None: + for mode in ("none", "truth", "flip", "random"): + default = _run(mode) + explicit = _run(mode, override_positions="all") + assert default.visit_end_positions is None + assert json.dumps(result_to_json(default)) == json.dumps( + result_to_json(explicit) + ) + assert list(result_to_json(default)) == OLD_KEYS + # The override edits every observed position by default: 2 visits of + # 10 events per subject, all observed. + assert _run("truth").n_intervened_positions == len(SUBJECTS) * 2 * EVENTS_PER_VISIT + + +def test_none_is_the_same_pass_under_both_settings() -> None: + plain = _run("none") + restricted = _run("none", override_positions="visit_end") + assert restricted.n_predictions == plain.n_predictions + assert restricted.top1_accuracy == plain.top1_accuracy + assert restricted.mean_task_loss == plain.mean_task_loss + assert restricted.n_intervened_positions == 0 + assert restricted.visit_end_positions is not None + + +# --------------------------------------------------------------------------- +# (ii) positions that are not visit ends receive no intervention +# --------------------------------------------------------------------------- + + +def test_visit_end_positions_are_each_visits_last_event() -> None: + chunks = _lane_chunks() + # One end per (subject, visit): six in all. + assert sum(int(visit_end_positions(c).sum()) for c in chunks) == len(SUBJECTS) * 2 + for chunk in chunks: + ends = visit_end_positions(chunk)[0] + sids = chunk.subject_ids[0] + vids = chunk.visit_ids[0] + for i in torch.nonzero(ends).flatten().tolist(): + later = (vids[i + 1 :] == vids[i]) & (sids[i + 1 :] == sids[i]) + assert not later.any() + + +def test_only_visit_ends_are_intervened_and_other_logits_match_none() -> None: + labels, masks, first = _visit_labels() + model = _model(len(_vocab())).eval() + state = None + n_touched = 0 + for chunk in _lane_chunks(): + position_mask = visit_end_positions(chunk) + intervention = _chunk_intervention( + chunk, + "truth", + labels, + masks, + first, + supervision="visit", + num_concepts=NUM_CONCEPTS, + device="cpu", + rng=torch.Generator().manual_seed(0), + position_mask=position_mask, + ) + assert intervention is not None + with torch.no_grad(): + base_logits, out, _ = model( + chunk.batch, state=state, reset_mask=chunk.reset_mask + ) + edited_logits, _, state = model( + chunk.batch, + state=state, + reset_mask=chunk.reset_mask, + intervention=intervention, + ) + applied = intervention_apply_mask(intervention, out.concept_probs) + assert applied is not None + touched = applied.any(dim=-1) + # The intervention mask is exactly the visit-end mask (all observed). + assert torch.equal(touched, position_mask) + # Off the visit ends the forward is byte for byte the un-edited one. + assert torch.equal(edited_logits[~touched], base_logits[~touched]) + # On them, feeding hard 0/1 values moves the logits. + assert not torch.equal(edited_logits[touched], base_logits[touched]) + n_touched += int(touched.sum()) + assert n_touched == len(SUBJECTS) * 2 + + +def test_without_a_position_mask_the_chunk_intervention_is_unchanged() -> None: + chunk = _lane_chunks()[0] + labels, masks, first = _visit_labels() + plain = _chunk_intervention( + chunk, + "flip", + labels, + masks, + first, + supervision="visit", + num_concepts=NUM_CONCEPTS, + device="cpu", + rng=torch.Generator().manual_seed(0), + ) + assert plain is not None and plain.probs_mask is not None + real = chunk.subject_ids != -1 + assert torch.equal(plain.probs_mask.any(dim=-1), real) + + +def test_random_and_calibrated_modes_respect_the_restriction() -> None: + n_ends = len(SUBJECTS) * 2 + assert ( + _run("random", override_positions="visit_end").n_intervened_positions == n_ends + ) + assert _run("flip", override_positions="visit_end").n_intervened_positions == n_ends + vocab = _vocab() + labels, masks, first = _visit_labels() + calibrated = run_streaming_intervention( + _model(len(vocab)), + _events(), + vocab, + labels, + masks, + mode="truth_calibrated", + concept_first_times=first, + supervision="visit", + num_lanes=2, + chunk_size=8, + device="cpu", + calibration_gammas=torch.full((NUM_CONCEPTS,), 0.25), + override_positions="visit_end", + ) + assert calibrated.n_intervened_positions == n_ends + + +# --------------------------------------------------------------------------- +# (iii) the visit_end_positions block +# --------------------------------------------------------------------------- + + +def test_block_counts_visits_with_an_observed_label_and_a_target() -> None: + truth = _run("truth", override_positions="visit_end") + block = truth.visit_end_positions + assert block is not None + # Every visit end is overridden (six), but only the first visit of + # each subject has a next-token target: the second visit's last + # event is followed by another patient (or nothing). + assert truth.n_intervened_positions == len(SUBJECTS) * 2 + assert block["n_positions"] == len(SUBJECTS) + assert block["n_subjects"] == len(SUBJECTS) + assert 0.0 <= block["top1_accuracy"] <= 1.0 + assert block["mean_task_loss"] > 0.0 + assert block["paired"] == {} + assert list(result_to_json(truth)) == [*OLD_KEYS, "visit_end_positions"] + + +def test_block_rows_follow_the_band_and_match_a_single_chunk_pass() -> None: + band = 0.05 + labels, masks, first = _visit_labels() + model = _model(len(_vocab())).eval() + state = None + expected_rows = 0 + expected_overridden = 0 + for chunk in _lane_chunks(): + with torch.no_grad(): + _, out, state = model(chunk.batch, state=state, reset_mask=chunk.reset_mask) + _, observed = position_running_labels( + chunk, labels, masks, first, supervision="visit", num_concepts=NUM_CONCEPTS + ) + in_band = (out.concept_probs - 0.5).abs() < band + overridden = (observed.bool() & in_band).any(dim=-1) & visit_end_positions( + chunk + ) + expected_overridden += int(overridden.sum()) + expected_rows += int((overridden & chunk.real_mask).sum()) + # A band this narrow drops some visits, so the subset is a real subset. + assert 0 < expected_rows < len(SUBJECTS) * 2 + result = _run("truth", override_positions="visit_end", uncertain_band=band) + block = result.visit_end_positions + assert block is not None + assert block["n_positions"] == expected_rows + assert result.n_intervened_positions == expected_overridden + + +def test_none_versus_none_gives_zero_deltas() -> None: + scorer_a = VisitEndScorer(uncertain_band=None) + scorer_b = VisitEndScorer(uncertain_band=None) + _run("none", override_positions="visit_end", scorer=scorer_a) + _run("none", override_positions="visit_end", scorer=scorer_b) + left, right = scorer_a.table(), scorer_b.table() + assert left.same_rows(right) and left.n_rows == len(SUBJECTS) + paired = visit_end_paired_summary(left, right, n_boot=50, seed=0) + for metric in ("top1_accuracy", "mean_task_loss"): + assert paired[metric] == { + "point": 0.0, + "ci_low": 0.0, + "ci_high": 0.0, + "separated": False, + } + assert paired["n_positions"] == len(SUBJECTS) + assert paired["n_subjects"] == len(SUBJECTS) + + +def test_paired_blocks_land_on_the_left_mode() -> None: + tables = {} + for mode in ("none", "truth", "flip"): + scorer = VisitEndScorer(uncertain_band=None) + _run(mode, override_positions="visit_end", scorer=scorer) + tables[mode] = scorer.table() + blocks = score_visit_end_rows(tables, n_boot=50, seed=0) + assert set(blocks["truth"]["paired"]) == {"truth_minus_none", "truth_minus_flip"} + assert set(blocks["flip"]["paired"]) == {"flip_minus_none"} + assert blocks["none"]["paired"] == {} + delta = blocks["truth"]["paired"]["truth_minus_none"]["top1_accuracy"] + assert delta["point"] == pytest.approx( + blocks["truth"]["top1_accuracy"] - blocks["none"]["top1_accuracy"] + ) + assert delta["ci_low"] <= delta["point"] <= delta["ci_high"] + # The rows are mode independent, so the per-mode tables line up. + assert np.array_equal(tables["truth"].subject_ids, tables["flip"].subject_ids) + + +def test_mismatched_rows_are_refused() -> None: + scorer = VisitEndScorer(uncertain_band=None) + _run("none", override_positions="visit_end", scorer=scorer) + full = scorer.table() + half = interventions_module.VisitEndRowTable( + subject_ids=full.subject_ids[:1], + visit_ids=full.visit_ids[:1], + time_hours=full.time_hours[:1], + hits=full.hits[:1], + losses=full.losses[:1], + ) + with pytest.raises(ValueError, match="identical rows"): + visit_end_paired_summary(full, half, n_boot=5, seed=0) + with pytest.raises(ValueError, match="different visit-end rows"): + score_visit_end_rows({"none": full, "truth": half}) + + +def test_unknown_override_positions_is_rejected() -> None: + with pytest.raises(ValueError, match="unknown override_positions"): + _run("none", override_positions="landmark") + + +# --------------------------------------------------------------------------- +# (iv) CLI parsing +# --------------------------------------------------------------------------- + + +def test_cli_passes_override_positions_through( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + seen: dict[str, object] = {} + block = { + "n_positions": 3, + "n_subjects": 3, + "top1_accuracy": 0.5, + "mean_task_loss": 1.0, + "paired": {}, + } + + def fake_evaluate(*_args: object, **kwargs: object) -> list[InterventionResult]: + seen.update(kwargs) + if kwargs["override_positions"] == "visit_end": + return [InterventionResult("none", 1, 1.0, 0.5, visit_end_positions=block)] + return [InterventionResult("none", 1, 1.0, 0.5)] + + monkeypatch.setattr(interventions_module, "evaluate_interventions", fake_evaluate) + argv = [ + "prog", + "--run-dir", + "/runs/x", + "--held-out-shard-dir", + "/data/held_out", + "--output-json", + ] + default = tmp_path / "default.json" + monkeypatch.setattr("sys.argv", [*argv, str(default)]) + interventions_module._main() + assert seen["override_positions"] == "all" + assert list(json.loads(default.read_text())[0]) == OLD_KEYS + + ends = tmp_path / "ends.json" + monkeypatch.setattr( + "sys.argv", [*argv, str(ends), "--override-positions", "visit_end"] + ) + interventions_module._main() + assert seen["override_positions"] == "visit_end" + written = json.loads(ends.read_text()) + assert list(written[0]) == [*OLD_KEYS, "visit_end_positions"] + assert written[0]["visit_end_positions"] == block + + monkeypatch.setattr( + "sys.argv", [*argv, str(tmp_path / "bad.json"), "--override-positions", "x"] + ) + with pytest.raises(SystemExit): + interventions_module._main() + assert OVERRIDE_POSITIONS == ("all", "visit_end") + + +# --------------------------------------------------------------------------- +# --hazard-heads stays compatible +# --------------------------------------------------------------------------- + + +def test_hazard_heads_combine_with_the_visit_end_override() -> None: + """Both blocks are filled; the landmark rows are the usual ones.""" + fx = _HazardFixture() + hazard_only, landmark_only = fx.run("truth", hazard=True) + scorer = fx.scorer() + combined = run_streaming_intervention( + fx.model, + fx.binned, + fx.vocab, + fx.labels, + fx.mask, + mode="truth", + concept_first_times=fx.first_times, + supervision="stay", + num_lanes=2, + chunk_size=16, + device="cpu", + seed=0, + hazard_scorer=scorer, + override_positions="visit_end", + ) + assert landmark_only is not None + assert scorer.table().same_rows(landmark_only) + assert combined.visit_end_positions is not None + # Far fewer positions are overridden than under the default. + assert 0 < combined.n_intervened_positions < hazard_only.n_intervened_positions From 5bd3096116656e69d13dc947086b438f73aef611 Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Fri, 9 Oct 2026 08:05:28 -0400 Subject: [PATCH 2/4] Random-code control for the edit-attribution cohort (Table 9) A reviewer asks for a random-code control to show the concept ranking helps. Table 9's numbers came from an uncommitted scratch harness (cohort_worsen.py); only its stdout logs were banked. This adds the cohort driver to the package, `python -m odyssey.inference.concept_edit_attribution`, reproducing that log format so scripts/parse_edit_attribution_logs.py reads both, and gives it the control arm. - `--code-selection {attributed,random}` (default attributed) and `--random-seed`: the random arm draws min(top_k, pool) codes per subject uniformly from the same candidate pool the occlusion loop ranks, seeded per subject from (seed, subject_id), then builds the identical worsen edits (worsen_edits_for_codes, split out of worsen_edits_from_attribution with no behaviour change). - `--candidate-pool {mappable,all}`: the default pool is the distinct window codes that resolve to a panel signal with a worsen edit, the only codes either arm can act on; `all` ranks every window code and drops unmappable picks the way the attributed arm always did. - `--index-frac` on _index_times_by_subject for the banked headers' index_frac field. - Parser tags runs with `selection` (attributed when absent) and `random_seed`; make_edit_attribution_table gains `--with-random` (Codes column, random row under its attributed row), default output byte-identical on the banked JSON; the CI table labels random blocks. - Tests: pool filtering, seeded draw, same-count/same-subject guarantee, parser round trip, table flag. docs/ml4h2026_rebuttal_plan WP9 carries the VM command lines for both arms on MIMIC-IV and eICU. Co-Authored-By: Claude Fable 5.1 --- docs/ml4h2026_rebuttal_plan.md | 53 ++ odyssey/inference/concept_edit_attribution.py | 554 +++++++++++++++++- odyssey/inference/counterfactual.py | 20 +- scripts/make_edit_attribution_ci_table.py | 8 +- scripts/make_edit_attribution_table.py | 84 ++- scripts/parse_edit_attribution_logs.py | 25 +- .../test_concept_edit_attribution.py | 207 +++++++ .../test_make_edit_attribution_ci_table.py | 16 + .../test_make_edit_attribution_table.py | 52 ++ .../test_parse_edit_attribution_logs.py | 27 + 10 files changed, 1014 insertions(+), 32 deletions(-) diff --git a/docs/ml4h2026_rebuttal_plan.md b/docs/ml4h2026_rebuttal_plan.md index 589a795f..966d923d 100644 --- a/docs/ml4h2026_rebuttal_plan.md +++ b/docs/ml4h2026_rebuttal_plan.md @@ -134,6 +134,59 @@ Keep each block under 150 words. The response box on OpenReview is short. - Run `figures/pagecheck.py` and the awk comment check after every edit (the build has lost prose to `%` lines before). - The submitted source is `paper/ml4h/main_mixture.tex`; the old `main.tex` and its aux files were retired to `paper/ml4h/retired/` on 2026-09-25. `make_steering_table.py` and `make_specificity_table.py` now feed no table in the paper; keep them for the steering follow-up. +### WP9. Reviewer-requested controls (added 2026-10-09) + +**Random-code control for Table 9 (tab:edit-attribution).** A reviewer asks for a random-code control to show that the concept ranking helps. The arm: for each of the same 50 held-out subjects, draw the same number of codes as the attributed arm edits (top-k = 4) uniformly from the same candidate pool in the same 24-hour window, worsen them with the identical edit machinery, and report the same sign agreement. Code on branch `rebuttal/random-code-control`: `python -m odyssey.inference.concept_edit_attribution` is the cohort driver (the Table 9 numbers came from an uncommitted scratch harness, `cohort_worsen.py`, whose stdout logs are banked under `research_journal/figure_data/edit_attribution/`; this CLI reproduces its log format so `scripts/parse_edit_attribution_logs.py` reads both). New flags: `--code-selection {attributed,random}` (default `attributed`), `--random-seed` (default 0), `--candidate-pool {mappable,all}` (default `mappable`: the distinct codes in the window that resolve to a panel signal with a worsen edit, the only codes either arm can act on; the banked eICU logs show exactly four edits per subject, which is consistent with that pool), `--index-frac` (the banked qSOFA and AKI headers carry `index_frac=0.3333`; its original definition is lost with the harness, here it is the fraction of the first qualifying visit's duration; the paper caption says 24 h into a visit, so the commands below use `--index-hours 24` and the attributed arm is re-run first to confirm reproduction before the random arm is read). The JSON output keeps per-subject codes, signals and hazard deltas, which also lifts the per-subject-CI limit noted on 2026-09-09. + +Subjects: the first 50 held-out subjects in shard order whose first visit lasts at least 24 h and whose window holds at least one worsenable code; the pool does not depend on the arm, so both arms score the same subjects. The random draw is seeded per subject from `(seed, subject_id)`, so it is reproducible and independent of subject order. Note the control can only differ from the attributed arm on subjects whose pool holds more than four worsenable codes; the summary line `candidate pool larger than top_k: m/n subjects` reports how many, and the eICU and MIMIC ICU-stay windows (vitals plus labs) are where the two arms separate. + +VM recipe as in WP1 to WP4: `PY=~/odyssey/.venv/bin/python`, code from a worktree on this branch (`cd ~/odyssey && git fetch origin && git worktree add ~/odyssey_rc rebuttal/random-code-control`), `PYTHONPATH=~/odyssey_rc`. Never `uv sync` on a VM. Four held-out shards is what the alert chains score by default and gives well over 50 eligible subjects; keep `--max-shards` identical across the arms of one concept. + +VM1 (MIMIC-IV, `full_run_v10`), one line per concept and arm; run the attributed arm first: + +``` +cd ~/odyssey_rc && export PYTHONPATH=~/odyssey_rc PY=~/odyssey/.venv/bin/python +for C in sepsis3 qsofa aki_stage_3; do + for SEL in attributed random; do + $PY -m odyssey.inference.concept_edit_attribution \ + --run-dir ~/runs/full_run_v10 \ + --held-out-shard-dir ~/data/mimiciv_3.1_v1/data/held_out \ + --concept $C --max-subjects 50 --top-k 4 --lookback-hours 24 --index-hours 24 \ + --max-shards 4 --chunk-size 512 \ + --code-selection $SEL --random-seed 0 \ + --output-json ~/runs/full_run_v10/edit_attribution/cohort_${C}_${SEL}.json \ + > ~/runs/full_run_v10/edit_attribution/cohort_${C}_${SEL}.log 2>&1 + done +done +``` + +VM2 (eICU-CRD, `eicu_full_v10`; no Sepsis-3 head on that checkpoint, as the paper states): + +``` +cd ~/odyssey_rc && export PYTHONPATH=~/odyssey_rc PY=~/odyssey/.venv/bin/python +for C in qsofa aki_stage_3; do + for SEL in attributed random; do + $PY -m odyssey.inference.concept_edit_attribution \ + --run-dir ~/runs/eicu_full_v10 \ + --held-out-shard-dir ~/data/eicu_2.0_v1/data/held_out \ + --concept $C --max-subjects 50 --top-k 4 --lookback-hours 24 --index-hours 24 \ + --max-shards 4 --chunk-size 512 \ + --code-selection $SEL --random-seed 0 \ + --output-json ~/runs/eicu_full_v10/edit_attribution/cohort_${C}_${SEL}.json \ + > ~/runs/eicu_full_v10/edit_attribution/cohort_${C}_${SEL}.log 2>&1 + done +done +``` + +Optional second random seed (`--random-seed 1`, output suffix `_random_s1`) if the first draw lands near a cell boundary. Then, on the laptop, bank the logs and JSONs under `research_journal/figure_data/edit_attribution/rebuttal/` and build the two-arm table: + +``` +uv run python scripts/parse_edit_attribution_logs.py research_journal/figure_data/edit_attribution/rebuttal/cohort_*_{attributed,random}.log --out research_journal/figure_data/edit_attribution/rebuttal/cohort_results_two_arm.json +uv run python scripts/make_edit_attribution_table.py research_journal/figure_data/edit_attribution/rebuttal/cohort_results_two_arm.json --with-random --out paper/ml4h/tables/edit_attribution_random_control.tex +``` + +Without `--with-random` the table script drops the random rows and renders the published Table 9 unchanged; the CI table (`make_edit_attribution_ci_table.py`) labels a random-arm block "random codes". Read the attributed arm against the banked 87 to 97% band first; if it does not reproduce within the Wilson intervals of `tab:edit-attribution-ci`, the subject selection or index definition differs from the lost harness and the rebuttal cites the new paired pair (attributed vs random, same subjects) rather than the published row. + ## Results so far (updated 2026-09-29, 01:00 UTC) Every number below is banked under `research_journal/figure_data/` in the run directory named, and was computed on all held-out shards unless stated. Runs launched 2026-09-25 from the `rebuttal/integration` branch; VM recipe and logs are in the session memory and the chain scripts under the VM home directories. diff --git a/odyssey/inference/concept_edit_attribution.py b/odyssey/inference/concept_edit_attribution.py index 1ad4319c..209e4f7b 100644 --- a/odyssey/inference/concept_edit_attribution.py +++ b/odyssey/inference/concept_edit_attribution.py @@ -48,13 +48,27 @@ already-validated :class:`~odyssey.inference.counterfactual.ValueEdit` machinery -- the same operation the hand-specified edits use, just aimed automatically instead of by hand. + +The cohort driver at the bottom (:func:`cohort_worsen`, ``python -m +odyssey.inference.concept_edit_attribution``) runs that discovery + edit +over held-out subjects and reports sign agreement of the hazard shift. It +also carries the reviewer-requested control: ``code_selection="random"`` +draws the same number of codes per subject uniformly from the same +candidate pool the occlusion loop ranks, and pushes them through the +identical worsen-edit machinery, so the gap between the two arms is what +the concept ranking adds and nothing else. """ from __future__ import annotations +import argparse +import json import logging +import random from collections.abc import Callable, Sequence -from dataclasses import dataclass +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Literal import polars as pl @@ -65,6 +79,7 @@ HORIZONS_HOURS, ForecastReadout, ValueEdit, + _index_times_by_subject, apply_value_edits, score_record_at, ) @@ -439,17 +454,48 @@ def worsen_edits_from_attribution( resolve to the same signal (e.g. two raw SBP item codes) collapse to one edit. """ - resolver = SignalPanelResolver(source=source) chosen = [a for a in attributions[:top_k] if abs(a.delta) >= min_abs_delta] + return worsen_edits_for_codes( + [a.code for a in chosen], source=source, window_hours=window_hours + ) + + +def worsen_signal_for_code(code: str, resolver: SignalPanelResolver) -> str | None: + """Return the panel signal ``code`` resolves to, if this module can worsen it. + + ``None`` when the code is outside the panel or its signal has no entry + in :data:`WORSEN_EDIT_FOR_SIGNAL`: the one rule that decides which + codes the edit machinery can act on, shared by the attributed arm, + the random-code control and the candidate pool they both draw from. + """ + idx = resolver.resolve(code) + if idx == NO_SIGNAL: + return None + signal_name = SIGNAL_PANEL[idx][0] + return signal_name if signal_name in WORSEN_EDIT_FOR_SIGNAL else None + + +def worsen_edits_for_codes( + codes: Sequence[str], + *, + source: str = "mimic_iv", + window_hours: float = 24.0, +) -> list[ValueEdit]: + """Turn raw codes into worsen-the-signal edits, in the order given. + + The edit-construction half of :func:`worsen_edits_from_attribution`, + with no ranking of its own: a code that resolves to no worsenable + signal is dropped, and two codes on the same signal collapse to one + edit. The random-code control calls this on its uniform draw so the + two arms differ only in which codes come in. + """ + resolver = SignalPanelResolver(source=source) edits: dict[str, ValueEdit] = {} - for a in chosen: - idx = resolver.resolve(a.code) - if idx == NO_SIGNAL: - continue - signal_name = SIGNAL_PANEL[idx][0] - template = WORSEN_EDIT_FOR_SIGNAL.get(signal_name) - if template is None: + for code in codes: + signal_name = worsen_signal_for_code(code, resolver) + if signal_name is None: continue + template = WORSEN_EDIT_FOR_SIGNAL[signal_name] edits[signal_name] = ValueEdit( signal=template.signal, mode=template.mode, @@ -494,3 +540,493 @@ def score_with_worsen_edits( chunk_size=chunk_size, ) return readout, touched + + +# --------------------------------------------------------------------------- +# Cohort validation: discovered (or random) codes -> worsen edits -> hazards +# --------------------------------------------------------------------------- + +CodeSelection = Literal["attributed", "random"] +CandidatePool = Literal["mappable", "all"] + + +def mappable_candidate_codes( + raw_subject_events: pl.DataFrame, + *, + index_time: object, + lookback_hours: float, + source: str = "mimic_iv", + resolver: SignalPanelResolver | None = None, +) -> list[str]: + """Distinct codes in the window that the worsen-edit machinery can act on. + + :func:`occlude_codes`'s default pool is every distinct code with a + reading in the window; most of those (diagnoses, orders, drugs) have + no worsen edit, so the attributed arm can only ever act on the subset + that resolves to a signal in :data:`WORSEN_EDIT_FOR_SIGNAL`. This is + that subset, sorted, and it is the pool both arms of the cohort + validation draw from. + """ + resolver = resolver or SignalPanelResolver(source=source) + codes = _candidate_codes( + raw_subject_events, index_time=index_time, lookback_hours=lookback_hours + ) + return [c for c in codes if worsen_signal_for_code(c, resolver) is not None] + + +def select_random_codes( + candidate_codes: Sequence[str], + *, + n: int, + seed: int, + subject_id: int, +) -> list[str]: + """Draw ``n`` codes uniformly without replacement from ``candidate_codes``. + + Seeded per subject from ``(seed, subject_id)`` so the draw is + reproducible and independent of the order subjects are visited in. + Fewer than ``n`` candidates returns them all (shuffled), which is the + same number the attributed arm would edit from that pool. + """ + if n < 0: + raise ValueError(f"n must be non-negative, got {n}") + pool = sorted(set(candidate_codes)) + rng = random.Random(f"{seed}:{subject_id}") + return rng.sample(pool, min(n, len(pool))) + + +@dataclass +class SubjectWorsenRecord: + """One subject's discovered (or drawn) codes, edit, and hazard shift.""" + + subject_id: int + index_time: str + n_candidates: int + """Size of the candidate pool both arms draw from.""" + codes: list[str] + """Codes the arm selected (top-``k`` attributed or the random draw).""" + signals: list[str] + """Panel signals those codes resolved to, i.e. the edits applied.""" + rows_edited: int + concept_before: float + concept_after: float + delta_event_risk: dict[str, dict[str, float]] + """event -> horizon -> (worsened - factual) risk.""" + + +@dataclass +class CohortWorsenResult: + """The cohort validation's per-subject records plus how they were made.""" + + concept: str + source: str + top_k: int + lookback_hours: float + index_hours: float + index_frac: float | None + code_selection: str + random_seed: int | None + candidate_pool: str + min_abs_delta: float + subjects: list[SubjectWorsenRecord] = field(default_factory=list) + n_pool_above_top_k: int = 0 + """Subjects whose pool held more than ``top_k`` codes, i.e. the only + subjects where the random draw can differ from the attributed set.""" + + @property + def scored(self) -> list[SubjectWorsenRecord]: + """Subjects with at least one reading actually edited.""" + return [s for s in self.subjects if s.rows_edited > 0] + + def sign_agreement(self) -> dict[str, dict[str, tuple[int, int]]]: + """Return event -> horizon -> (subjects whose risk rose, subjects scored).""" + out: dict[str, dict[str, tuple[int, int]]] = {} + scored = self.scored + events = sorted({ev for s in scored for ev in s.delta_event_risk}) + for ev in events: + horizons = [ + h + for h in (f"{x:g}h" for x in HORIZONS_HOURS) + if any(h in s.delta_event_risk.get(ev, {}) for s in scored) + ] + out[ev] = {} + for h in horizons: + ds = [ + s.delta_event_risk[ev][h] + for s in scored + if h in s.delta_event_risk.get(ev, {}) + ] + out[ev][h] = (sum(1 for d in ds if d > 0), len(ds)) + return out + + +def _candidate_pool( + raw_subject_events: pl.DataFrame, + *, + index_time: object, + lookback_hours: float, + candidate_pool: CandidatePool, + source: str, + resolver: SignalPanelResolver, +) -> list[str]: + if candidate_pool == "mappable": + return mappable_candidate_codes( + raw_subject_events, + index_time=index_time, + lookback_hours=lookback_hours, + source=source, + resolver=resolver, + ) + return _candidate_codes( + raw_subject_events, index_time=index_time, lookback_hours=lookback_hours + ) + + +def _resolve_signals(codes: Sequence[str], resolver: SignalPanelResolver) -> list[str]: + seen: list[str] = [] + for code in codes: + name = worsen_signal_for_code(code, resolver) + if name is not None and name not in seen: + seen.append(name) + return seen + + +def cohort_worsen( + model: SequenceModel, + vocab: Vocabulary, + binner: QuantileBinner | None, + raw_events: pl.DataFrame, + *, + concept_name: str, + concept_names: Sequence[str], + top_k: int = 4, + lookback_hours: float = 24.0, + index_hours: float = 24.0, + index_frac: float | None = None, + max_subjects: int = 50, + code_selection: CodeSelection = "attributed", + random_seed: int = 0, + candidate_pool: CandidatePool = "mappable", + min_abs_delta: float = 0.0, + source: str = "mimic_iv", + device: str = "cpu", + chunk_size: int = 256, + log: Callable[[str], None] | None = None, +) -> CohortWorsenResult: + """Discover (or draw) codes per subject, worsen them, and re-score hazards. + + Subjects are the first ``max_subjects`` held-out subjects, in the + order the shards list them, whose first visit lasts at least + ``index_hours`` and whose candidate pool is non-empty; the pool does + not depend on the arm, so both arms score the same subjects. Per + subject, ``attributed`` ranks the pool by occlusion of the concept + head and keeps the top ``top_k`` codes; ``random`` draws + ``min(top_k, pool size)`` codes uniformly from the same pool, seeded + from ``(random_seed, subject_id)``. Both arms then build the same + worsen edits (:func:`worsen_edits_for_codes`) and read the hazard + heads before and after. ``candidate_pool="all"`` ranks every distinct + code in the window instead, the way :func:`occlude_codes` does by + default; a selected code with no worsen edit is then dropped, so a + random draw there can leave a subject with nothing to edit, which is + reported but not scored. + """ + if concept_name not in concept_names: + raise ValueError(f"{concept_name!r} not in concept_names {list(concept_names)}") + if code_selection not in ("attributed", "random"): + raise ValueError(f"unknown code_selection {code_selection!r}") + if candidate_pool not in ("mappable", "all"): + raise ValueError(f"unknown candidate_pool {candidate_pool!r}") + emit = log or (lambda _msg: None) + resolver = SignalPanelResolver(source=source) + index_times = _index_times_by_subject( + raw_events, index_hours=index_hours, index_frac=index_frac + ) + ordered = [ + int(sid) + for sid in raw_events["subject_id"].unique(maintain_order=True).to_list() + if sid in index_times + ] + result = CohortWorsenResult( + concept=concept_name, + source=source, + top_k=top_k, + lookback_hours=lookback_hours, + index_hours=index_hours, + index_frac=index_frac, + code_selection=code_selection, + random_seed=random_seed if code_selection == "random" else None, + candidate_pool=candidate_pool, + min_abs_delta=min_abs_delta, + ) + for sid in ordered: + if len(result.subjects) >= max_subjects: + break + sub = raw_events.filter(pl.col("subject_id") == sid) + index_time = index_times[sid] + pool = _candidate_pool( + sub, + index_time=index_time, + lookback_hours=lookback_hours, + candidate_pool=candidate_pool, + source=source, + resolver=resolver, + ) + if not pool: + continue + if len(pool) > top_k: + result.n_pool_above_top_k += 1 + if code_selection == "attributed": + attributions = occlusion_attribution( + model, + vocab, + binner, + sub, + index_time=index_time, + concept_name=concept_name, + concept_names=concept_names, + lookback_hours=lookback_hours, + candidate_codes=pool, + source=source, + device=device, + chunk_size=chunk_size, + ) + codes = [ + a.code for a in attributions[:top_k] if abs(a.delta) >= min_abs_delta + ] + else: + codes = select_random_codes( + pool, n=min(top_k, len(pool)), seed=random_seed, subject_id=sid + ) + edits = worsen_edits_for_codes( + codes, source=source, window_hours=lookback_hours + ) + factual = score_record_at( + model, + vocab, + binner, + sub, + index_time=index_time, + concept_names=concept_names, + source=source, + device=device, + chunk_size=chunk_size, + ) + if edits: + worsened, touched = score_with_worsen_edits( + model, + vocab, + binner, + sub, + edits, + index_time=index_time, + concept_names=concept_names, + source=source, + device=device, + chunk_size=chunk_size, + ) + else: + worsened, touched = factual, 0 + before = factual.concept_probs.get(concept_name, float("nan")) + after = worsened.concept_probs.get(concept_name, float("nan")) + record = SubjectWorsenRecord( + subject_id=sid, + index_time=str(index_time), + n_candidates=len(pool), + codes=codes, + signals=_resolve_signals(codes, resolver), + rows_edited=touched, + concept_before=before, + concept_after=after, + delta_event_risk={ + ev: {h: worsened.event_risk[ev][h] - p for h, p in hs.items()} + for ev, hs in factual.event_risk.items() + if ev in worsened.event_risk + }, + ) + result.subjects.append(record) + emit( + f"subject {sid} [{len(result.subjects)}/{max_subjects}]: " + f"{record.signals!r} concept {before:.3f}->{after:.3f}" + ) + return result + + +def format_cohort_summary(result: CohortWorsenResult) -> str: + """Render the fixed-format summary block the log parser reads. + + ``scripts/parse_edit_attribution_logs.py`` turns this into JSON. Same + layout as the banked ``cohort_*.log`` files under + ``research_journal/figure_data/edit_attribution/``, plus a + ``selection=`` tag (and ``seed=`` for the random arm) in the header + so the parser can label the arm. + """ + n = len(result.subjects) + header = ( + f"=== {result.concept} on {result.source}, n={n} subjects, top_k={result.top_k}" + ) + if result.index_frac is not None: + header += f", index_frac={result.index_frac:.4f}" + header += f", selection={result.code_selection}" + if result.code_selection == "random": + header += f", seed={result.random_seed}" + if result.candidate_pool != "mappable": + header += f", pool={result.candidate_pool}" + no_edit = n - len(result.scored) + if no_edit: + header += f", no_edit={no_edit}" + header += " ===" + saturated = sum(1 for s in result.subjects if s.concept_before >= 0.9) + freq: dict[str, int] = {} + for s in result.subjects: + for name in s.signals: + freq[name] = freq.get(name, 0) + 1 + scored = result.scored + mean_delta = ( + sum(s.concept_after - s.concept_before for s in scored) / len(scored) + if scored + else float("nan") + ) + lines = [ + header, + f"baseline concept prob >= 0.9: {saturated}/{n} subjects (saturation check)", + f"candidate pool larger than top_k: {result.n_pool_above_top_k}/{n} subjects", + f"edit signal frequency: {freq!r}", + f"mean concept-probability delta: {mean_delta:+.4f}", + "", + "Sign agreement (risk should INCREASE when discovered evidence is worsened):", + ] + for event, horizons in result.sign_agreement().items(): + for horizon, (agree, total) in horizons.items(): + pct = 100.0 * agree / total if total else 0.0 + lines.append( + f" {event:<22} {horizon:<4}: {agree:3d}/{total:3d} = {pct:5.1f}%" + ) + return "\n".join(lines) + "\n" + + +def _main() -> None: + from odyssey.data.code_normalization import maybe_normalize # noqa: PLC0415 + from odyssey.data.history_recap import maybe_history_recap # noqa: PLC0415 + from odyssey.inference.legacy_concept_pins import ( # noqa: PLC0415 + resolve_concepts_for_run, + ) + from odyssey.inference.run_inference import load_run # noqa: PLC0415 + from odyssey.training.data import load_meds_shards # noqa: PLC0415 + from odyssey.utils.device import default_device # noqa: PLC0415 + + parser = argparse.ArgumentParser( + description=( + "Automated occlusion + worsen-edit cohort validation (paper " + "tab:edit-attribution), with a random-code control arm." + ) + ) + parser.add_argument("--run-dir", required=True) + parser.add_argument("--held-out-shard-dir", required=True) + parser.add_argument( + "--concept", required=True, help="e.g. sepsis3, qsofa, aki_stage_3" + ) + parser.add_argument("--output-json", required=True) + parser.add_argument("--max-shards", type=int, default=2) + parser.add_argument("--max-subjects", type=int, default=50) + parser.add_argument("--top-k", type=int, default=4) + parser.add_argument("--lookback-hours", type=float, default=24.0) + parser.add_argument("--index-hours", type=float, default=24.0) + parser.add_argument( + "--index-frac", + type=float, + default=None, + help="index at this fraction of the first qualifying visit instead of --index-hours into it", + ) + parser.add_argument("--min-abs-delta", type=float, default=0.0) + parser.add_argument( + "--code-selection", + choices=("attributed", "random"), + default="attributed", + help="attributed: top-k by occlusion (the paper's arm); random: same count, uniform from the same pool", + ) + parser.add_argument("--random-seed", type=int, default=0) + parser.add_argument( + "--candidate-pool", + choices=("mappable", "all"), + default="mappable", + help="mappable: codes with a worsen edit (default); all: every distinct code in the window", + ) + parser.add_argument("--checkpoint", default=None) + parser.add_argument("--chunk-size", type=int, default=512) + parser.add_argument("--device", default=None) + args = parser.parse_args() + + device = args.device or default_device() + run_dir = Path(args.run_dir) + model, vocab, binner, config = load_run( + run_dir, + device=device, + checkpoint_path=run_dir / (args.checkpoint or "checkpoint_best.pt"), + ) + source = getattr(config, "source", "mimic_iv") + concept_names = [ + c.name + for c in resolve_concepts_for_run( + str(run_dir), source, getattr(config, "task_set", "v1") + ) + ] + raw = load_meds_shards(args.held_out_shard_dir, max_shards=args.max_shards) + raw = maybe_normalize( + raw, enabled=getattr(config, "normalize_medications", False), source=source + ) + raw = maybe_history_recap(raw, enabled=getattr(config, "history_recap", False)) + resolver = SignalPanelResolver(source=source) + present = { + name + for name in ( + worsen_signal_for_code(c, resolver) for c in raw["code"].unique().to_list() + ) + if name is not None + } + print( + f"worsenable signals with a reading in the loaded shards: " + f"{len(present)}/{len(WORSEN_EDIT_FOR_SIGNAL)}", + flush=True, + ) + result = cohort_worsen( + model, + vocab, + binner, + raw, + concept_name=args.concept, + concept_names=concept_names, + top_k=args.top_k, + lookback_hours=args.lookback_hours, + index_hours=args.index_hours, + index_frac=args.index_frac, + max_subjects=args.max_subjects, + code_selection=args.code_selection, + random_seed=args.random_seed, + candidate_pool=args.candidate_pool, + min_abs_delta=args.min_abs_delta, + source=source, + device=device, + chunk_size=args.chunk_size, + log=lambda msg: print(msg, flush=True), + ) + summary = format_cohort_summary(result) + print() + print(summary, end="", flush=True) + out = Path(args.output_json) + out.parent.mkdir(parents=True, exist_ok=True) + payload = asdict(result) + payload["run_dir"] = str(run_dir) + payload["sign_agreement"] = { + ev: {h: {"agree": a, "total": t} for h, (a, t) in hs.items()} + for ev, hs in result.sign_agreement().items() + } + payload["summary_text"] = summary + out.write_text(json.dumps(payload, indent=2) + "\n") + logger.info("[cohort_worsen] wrote %s", out) + + +if __name__ == "__main__": + logging.basicConfig( + level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s" + ) + _main() diff --git a/odyssey/inference/counterfactual.py b/odyssey/inference/counterfactual.py index 506c3ed2..4fdab3b1 100644 --- a/odyssey/inference/counterfactual.py +++ b/odyssey/inference/counterfactual.py @@ -365,22 +365,38 @@ def counterfactual_forecast( def _index_times_by_subject( - raw_events: pl.DataFrame, *, index_hours: float + raw_events: pl.DataFrame, + *, + index_hours: float, + index_frac: float | None = None, ) -> dict[int, object]: """Return subject -> last event time at or before ``index_hours`` into a visit. The first visit of each subject that lasts at least ``index_hours``. + With ``index_frac`` the cut sits at that fraction of the visit's + duration instead of at a fixed ``index_hours`` into it (the visit + still has to last ``index_hours`` to qualify), so a short stay is + indexed early and a long one late. """ + if index_frac is not None and not 0.0 < index_frac <= 1.0: + raise ValueError(f"index_frac must be in (0, 1], got {index_frac}") timed = raw_events.filter( pl.col("time").is_not_null() & pl.col("hadm_id").is_not_null() ) + if index_frac is None: + cut = pl.col("_start") + pl.duration(hours=index_hours) + else: + span_us = (pl.col("_end") - pl.col("_start")).dt.total_microseconds() + cut = pl.col("_start") + pl.duration( + microseconds=(span_us.cast(pl.Float64) * index_frac).cast(pl.Int64) + ) visits = ( timed.group_by("subject_id", "hadm_id") .agg(pl.col("time").min().alias("_start"), pl.col("time").max().alias("_end")) .filter((pl.col("_end") - pl.col("_start")) >= pl.duration(hours=index_hours)) .sort(["subject_id", "_start"]) .unique(subset=["subject_id"], keep="first", maintain_order=True) - .with_columns((pl.col("_start") + pl.duration(hours=index_hours)).alias("_cut")) + .with_columns(cut.alias("_cut")) ) at = ( timed.join( diff --git a/scripts/make_edit_attribution_ci_table.py b/scripts/make_edit_attribution_ci_table.py index 84569159..7e6d073a 100644 --- a/scripts/make_edit_attribution_ci_table.py +++ b/scripts/make_edit_attribution_ci_table.py @@ -86,9 +86,15 @@ def render(results: dict[str, Any]) -> str: for run in results["runs"]: concept = CONCEPT_LABELS.get(run["concept"], run["concept"]) source = SOURCE_LABELS.get(run["source"], run["source"]) + block = f"{concept}, {source}" + selection = run.get("selection") or "attributed" + if selection != "attributed": + # the random-code control arm shares concept and source with + # its attributed twin; say which block this is + block += f", {selection} codes" lines += [ "\\midrule", - f"\\multicolumn{{{1 + len(HORIZONS)}}}{{l}}{{\\emph{{{concept}, {source}}}" + f"\\multicolumn{{{1 + len(HORIZONS)}}}{{l}}{{\\emph{{{block}}}" f" ($n={run['n_subjects']}$)}} \\\\", ] for event in EVENTS: diff --git a/scripts/make_edit_attribution_table.py b/scripts/make_edit_attribution_table.py index fa549e99..f7b9ff43 100644 --- a/scripts/make_edit_attribution_table.py +++ b/scripts/make_edit_attribution_table.py @@ -9,6 +9,12 @@ ICU_ADMISSION code, at the start of the stay, so the "future" hazard this event measures is structurally non-recurring past that point -- its cell is marked, not silently averaged in as if it meant the same thing as elsewhere. + +Runs tagged ``selection: random`` (the random-code control, same subjects, +same number of codes, same edit machinery, codes drawn uniformly from the +same candidate pool) are left out by default so the published table does +not change; ``--with-random`` adds a "Codes" column and places each +random-arm row directly under its attributed row. """ from __future__ import annotations @@ -38,25 +44,72 @@ DEGENERATE = {("icu_admission", "eicu")} +SELECTION_LABELS = {"attributed": "attributed", "random": "random"} + + def _mean_pct(run: dict[str, Any], event: str) -> float | None: cells = [c["pct"] for c in run["cells"] if c["event"] == event] return sum(cells) / len(cells) if cells else None -def render(results: dict[str, Any]) -> str: - """Return the LaTeX tabular for one row per (concept, source) run.""" +def _selection(run: dict[str, Any]) -> str: + return str(run.get("selection") or "attributed") + + +def _ordered_runs( + runs: list[dict[str, Any]], *, with_random: bool +) -> list[dict[str, Any]]: + """Attributed runs in file order, each followed by its random-arm twin(s). + + Without ``with_random`` only the attributed runs are kept, so a JSON + that carries both arms renders the same table as one that carries the + attributed arm alone. + """ + attributed = [r for r in runs if _selection(r) == "attributed"] + if not with_random: + return attributed + others = [r for r in runs if _selection(r) != "attributed"] + ordered: list[dict[str, Any]] = [] + for run in attributed: + ordered.append(run) + ordered += [ + o + for o in others + if (o["concept"], o["source"]) == (run["concept"], run["source"]) + ] + matched = {id(o) for o in ordered} + ordered += [o for o in others if id(o) not in matched] + return ordered + + +def render(results: dict[str, Any], *, with_random: bool = False) -> str: + """Return the LaTeX tabular for one row per (concept, source) run. + + ``with_random`` adds a "Codes" column (attributed / random) and the + random-arm rows; the default output has neither. + """ + runs = _ordered_runs(results["runs"], with_random=with_random) + head = "Concept & Dataset & " + spec = "ll" + if with_random: + head += "Codes & " + spec += "l" lines = [ "% GENERATED by scripts/make_edit_attribution_table.py -- do not hand-edit.", "% Each event's cell averages sign agreement over the 8/24/72h horizons.", "% dag = icu_admission on eICU-CRD, structurally non-recurring there (see caption).", - "\\begin{tabular}{llr" + "r" * len(EVENTS) + "}", + ] + if with_random: + lines.append( + "% random = same subjects and code count, codes drawn uniformly from the same window pool." + ) + lines += [ + "\\begin{tabular}{" + spec + "r" + "r" * len(EVENTS) + "}", "\\toprule", - "Concept & Dataset & $n$ & " - + " & ".join(EVENT_LABELS[e] for e in EVENTS) - + " \\\\", + head + "$n$ & " + " & ".join(EVENT_LABELS[e] for e in EVENTS) + " \\\\", "\\midrule", ] - for run in results["runs"]: + for run in runs: concept = CONCEPT_LABELS.get(run["concept"], run["concept"]) source = SOURCE_LABELS.get(run["source"], run["source"]) cells = [] @@ -66,11 +119,11 @@ def render(results: dict[str, Any]) -> str: continue mean_pct = _mean_pct(run, event) cells.append(f"{mean_pct:.0f}\\%" if mean_pct is not None else "--") - lines.append( - f"{concept} & {source} & {run['n_subjects']} & " - + " & ".join(cells) - + " \\\\" - ) + row = f"{concept} & {source} & " + if with_random: + sel = _selection(run) + row += SELECTION_LABELS.get(sel, sel) + " & " + lines.append(row + f"{run['n_subjects']} & " + " & ".join(cells) + " \\\\") lines += ["\\bottomrule", "\\end{tabular}"] return "\n".join(lines) + "\n" @@ -80,10 +133,15 @@ def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("results_json", type=Path) parser.add_argument("--out", type=Path, required=True) + parser.add_argument( + "--with-random", + action="store_true", + help="add the random-code control rows and a Codes column", + ) args = parser.parse_args() results = json.loads(args.results_json.read_text()) - args.out.write_text(render(results)) + args.out.write_text(render(results, with_random=args.with_random)) print(f"wrote {args.out}") diff --git a/scripts/parse_edit_attribution_logs.py b/scripts/parse_edit_attribution_logs.py index 40fcedf5..27c9f00a 100644 --- a/scripts/parse_edit_attribution_logs.py +++ b/scripts/parse_edit_attribution_logs.py @@ -1,11 +1,15 @@ -"""Parse cohort_worsen.py's stdout logs into one structured results JSON. +"""Parse the edit-attribution cohort logs into one structured results JSON. -cohort_worsen.py (paper/ml4h/../scratch validation harness, not part of the -package) prints a fixed-format summary block per run: a header line with -the concept/source/n, an edit-signal-frequency dict, a mean concept-delta, -and one sign-agreement line per (event, horizon). This turns that text -into JSON so the paper's table/figure generators never hand-transcribe a -number from a terminal log. +The cohort validation (originally the scratch harness cohort_worsen.py, +now ``python -m odyssey.inference.concept_edit_attribution``) prints a +fixed-format summary block per run: a header line with the +concept/source/n, an edit-signal-frequency dict, a mean concept-delta, and +one sign-agreement line per (event, horizon). This turns that text into +JSON so the paper's table/figure generators never hand-transcribe a number +from a terminal log. The header may carry a ``selection=`` tag +(``attributed`` or ``random``, the reviewer-requested random-code control) +and, for the random arm, a ``seed=``; a header without the tag is the +original attributed arm. """ from __future__ import annotations @@ -26,6 +30,8 @@ _SATURATION_RE = re.compile( r"baseline concept prob >= 0\.9:\s*(?P\d+)/(?P\d+)" ) +_SELECTION_RE = re.compile(r"\bselection=(?P[a-z_]+)") +_SEED_RE = re.compile(r"\bseed=(?P\d+)") def parse_log(text: str) -> dict[str, Any]: @@ -38,6 +44,9 @@ def parse_log(text: str) -> dict[str, Any]: header = _HEADER_RE.search(text) if header is None: raise ValueError("no '=== on , n=... ===' header found") + header_line = text[header.start() : text.find("\n", header.start())] + selection = _SELECTION_RE.search(header_line) + seed = _SEED_RE.search(header_line) sat = _SATURATION_RE.search(text) cells: list[dict[str, Any]] = [] for line in text.splitlines(): @@ -57,6 +66,8 @@ def parse_log(text: str) -> dict[str, Any]: "concept": header["concept"], "source": header["source"], "n_subjects": int(header["n"]), + "selection": selection["selection"] if selection else "attributed", + "random_seed": int(seed["seed"]) if seed else None, "n_baseline_saturated": int(sat["n_sat"]) if sat else None, "n_baseline_scored": int(sat["n_total"]) if sat else None, "cells": cells, diff --git a/tests/odyssey/inference/test_concept_edit_attribution.py b/tests/odyssey/inference/test_concept_edit_attribution.py index c1efdca6..7efd4f52 100644 --- a/tests/odyssey/inference/test_concept_edit_attribution.py +++ b/tests/odyssey/inference/test_concept_edit_attribution.py @@ -553,3 +553,210 @@ def test_worsen_edit_for_signal_covers_qsofa_resp_rate() -> None: edit = WORSEN_EDIT_FOR_SIGNAL["resp_rate"] assert edit.mode == "set" assert edit.value >= 22.0 + + +# --------------------------------------------------------------------------- +# Cohort validation and the random-code control +# --------------------------------------------------------------------------- + +from odyssey.inference.concept_edit_attribution import ( # noqa: E402 + cohort_worsen, + format_cohort_summary, + mappable_candidate_codes, + select_random_codes, + worsen_edits_for_codes, +) +from odyssey.inference.counterfactual import _index_times_by_subject # noqa: E402 +from scripts.parse_edit_attribution_logs import parse_log # noqa: E402 + + +PLT = "LAB//RESULT//51265//K/uL" # platelets: panel signal with a worsen edit +GLUCOSE = "LAB//RESULT//50931//mg/dL" # panel signal WITHOUT a worsen edit +DIAG = "DIAG//ICD10//E11" # outside the panel entirely + + +def _cohort_events(n_subjects: int = 3) -> pl.DataFrame: + """``n_subjects`` records, each a 30 h visit with three worsenable codes.""" + rows: list[tuple[int, str, datetime, float | None, int]] = [] + for sid in range(1, n_subjects + 1): + hadm = 100 + sid + rows.append((sid, "HOSPITAL_ADMISSION//EMERGENCY", T0, None, hadm)) + rows.append((sid, DIAG, T0 + timedelta(hours=1), None, hadm)) + for h in range(1, 30): + rows.append((sid, SBP, T0 + timedelta(hours=h), 120.0 + sid, hadm)) + if h % 6 == 0: + rows.append((sid, CREAT, T0 + timedelta(hours=h), 1.0, hadm)) + if h % 8 == 0: + rows.append((sid, PLT, T0 + timedelta(hours=h), 200.0, hadm)) + rows.append((sid, GLUCOSE, T0 + timedelta(hours=h), 100.0, hadm)) + rows.append( + (sid, "HOSPITAL_DISCHARGE//HOME", T0 + timedelta(hours=30), None, hadm) + ) + return pl.DataFrame(rows, schema=SCHEMA, orient="row") + + +def test_mappable_candidate_codes_keeps_only_worsenable_window_codes() -> None: + events = _cohort_events(1) + pool = mappable_candidate_codes( + events, index_time=T0 + timedelta(hours=24), lookback_hours=24.0 + ) + # glucose is a panel signal but has no worsen edit; the diagnosis is + # outside the panel; both are excluded. Sorted, so the pool is stable. + assert pool == sorted([SBP, CREAT, PLT]) + + +def test_select_random_codes_draws_n_from_pool_and_is_seed_reproducible() -> None: + pool = [SBP, CREAT, PLT, "LAB//220180//mmHg"] + a = select_random_codes(pool, n=2, seed=0, subject_id=7) + b = select_random_codes(pool, n=2, seed=0, subject_id=7) + assert a == b + assert len(a) == 2 + assert set(a) <= set(pool) + assert len(set(a)) == 2 # without replacement + # a different seed or subject changes the draw somewhere in the pool + draws = { + tuple(select_random_codes(pool, n=2, seed=s, subject_id=7)) for s in range(20) + } + assert len(draws) > 1 + # fewer candidates than n: everything, same count the attributed arm gets + assert sorted(select_random_codes([SBP], n=3, seed=0, subject_id=1)) == [SBP] + with pytest.raises(ValueError, match="non-negative"): + select_random_codes(pool, n=-1, seed=0, subject_id=1) + + +def test_worsen_edits_for_codes_matches_attribution_path() -> None: + attributions = [ + CodeAttribution(code=CREAT, n_rows=4, baseline=0.5, occluded=0.1), + CodeAttribution(code=SBP, n_rows=24, baseline=0.5, occluded=0.7), + CodeAttribution(code=DIAG, n_rows=1, baseline=0.5, occluded=0.55), + CodeAttribution(code=GLUCOSE, n_rows=3, baseline=0.5, occluded=0.52), + ] + via_attribution = worsen_edits_from_attribution(attributions, top_k=4) + via_codes = worsen_edits_for_codes([a.code for a in attributions]) + assert via_codes == via_attribution + assert [e.signal for e in via_codes] == ["creatinine", "sbp_noninvasive"] + + +def test_index_frac_places_the_cut_inside_the_visit() -> None: + events = _cohort_events(1) + fixed = _index_times_by_subject(events, index_hours=24.0) + frac = _index_times_by_subject(events, index_hours=24.0, index_frac=0.5) + assert fixed[1] == T0 + timedelta(hours=24) + assert frac[1] == T0 + timedelta(hours=15) # 30 h visit, half way + with pytest.raises(ValueError, match="index_frac"): + _index_times_by_subject(events, index_hours=24.0, index_frac=0.0) + + +def _run_cohort(events: pl.DataFrame, **kwargs: object) -> tuple[list[str], object]: + vocab, model, concepts = _vocab_and_model(events) + lines: list[str] = [] + result = cohort_worsen( + model, + vocab, + None, + events, + concept_name=concepts[0], + concept_names=concepts, + top_k=2, + lookback_hours=24.0, + index_hours=24.0, + chunk_size=16, + log=lines.append, + **kwargs, # type: ignore[arg-type] + ) + return lines, result + + +def test_cohort_worsen_attributed_default_scores_every_subject() -> None: + events = _cohort_events(3) + lines, result = _run_cohort(events) + assert result.code_selection == "attributed" + assert result.random_seed is None + assert [s.subject_id for s in result.subjects] == [1, 2, 3] + assert result.n_pool_above_top_k == 3 # pool of 3 worsenable codes, top_k=2 + for s in result.subjects: + assert s.n_candidates == 3 + assert len(s.codes) == 2 + assert set(s.codes) <= {SBP, CREAT, PLT} + assert s.rows_edited > 0 + assert set(s.delta_event_risk) == {"vasopressor_start", "death"} + assert set(s.delta_event_risk["death"]) == {"8h", "24h", "72h"} + assert len(lines) == 3 + assert lines[0].startswith("subject 1 [1/50]: [") + + +def test_cohort_worsen_random_draws_same_count_from_same_pool_and_subjects() -> None: + events = _cohort_events(3) + _, attributed = _run_cohort(events) + _, random_a = _run_cohort(events, code_selection="random", random_seed=3) + _, random_b = _run_cohort(events, code_selection="random", random_seed=3) + assert random_a.code_selection == "random" + assert random_a.random_seed == 3 + # same subjects, in the same order, as the attributed arm + assert [s.subject_id for s in random_a.subjects] == [ + s.subject_id for s in attributed.subjects + ] + for att, rnd in zip(attributed.subjects, random_a.subjects): + assert rnd.n_candidates == att.n_candidates + assert len(rnd.codes) == len(att.codes) + assert set(rnd.codes) <= {SBP, CREAT, PLT} + assert rnd.rows_edited > 0 + # seed-reproducible, byte for byte + assert [s.codes for s in random_a.subjects] == [s.codes for s in random_b.subjects] + assert [s.delta_event_risk for s in random_a.subjects] == [ + s.delta_event_risk for s in random_b.subjects + ] + # and the draw actually varies with the seed somewhere in the cohort + _, random_c = _run_cohort(events, code_selection="random", random_seed=11) + assert [s.codes for s in random_a.subjects] != [s.codes for s in random_c.subjects] + + +def test_cohort_worsen_rejects_unknown_selection_and_pool() -> None: + events = _cohort_events(1) + with pytest.raises(ValueError, match="code_selection"): + _run_cohort(events, code_selection="bogus") + with pytest.raises(ValueError, match="candidate_pool"): + _run_cohort(events, candidate_pool="bogus") + + +def test_cohort_worsen_all_pool_reports_unedited_subjects() -> None: + # with every window code as a candidate, a random draw can land on + # codes with no worsen edit; those subjects are kept (same cohort) but + # not scored, and the header says how many. + events = _cohort_events(2) + _, result = _run_cohort(events, candidate_pool="all") + assert result.candidate_pool == "all" + assert all(s.n_candidates == 5 for s in result.subjects) # 3 worsenable + 2 not + text = format_cohort_summary(result) + assert "pool=all" in text.splitlines()[0] + + +def test_format_cohort_summary_round_trips_through_the_parser() -> None: + events = _cohort_events(3) + _, attributed = _run_cohort(events) + _, rnd = _run_cohort(events, code_selection="random", random_seed=5) + att_text = format_cohort_summary(attributed) + rnd_text = format_cohort_summary(rnd) + att_head = att_text.splitlines()[0] + rnd_head = rnd_text.splitlines()[0] + assert att_head.endswith("top_k=2, selection=attributed ===") + assert rnd_head.endswith("top_k=2, selection=random, seed=5 ===") + parsed_att = parse_log(att_text) + parsed_rnd = parse_log(rnd_text) + assert parsed_att["selection"] == "attributed" + assert parsed_att["random_seed"] is None + assert parsed_rnd["selection"] == "random" + assert parsed_rnd["random_seed"] == 5 + for parsed in (parsed_att, parsed_rnd): + assert parsed["n_subjects"] == 3 + assert parsed["n_baseline_scored"] == 3 + assert {(c["event"], c["horizon"]) for c in parsed["cells"]} == { + (ev, h) + for ev in ("death", "vasopressor_start") + for h in ("8h", "24h", "72h") + } + assert all(c["total"] == 3 for c in parsed["cells"]) + assert all(0 <= c["agree"] <= 3 for c in parsed["cells"]) + # the per-subject lines carry the signals, like the banked logs + assert "edit signal frequency: {" in att_text + assert "mean concept-probability delta: " in att_text diff --git a/tests/scripts/test_make_edit_attribution_ci_table.py b/tests/scripts/test_make_edit_attribution_ci_table.py index e289ebee..025e0304 100644 --- a/tests/scripts/test_make_edit_attribution_ci_table.py +++ b/tests/scripts/test_make_edit_attribution_ci_table.py @@ -98,3 +98,19 @@ def test_render_labels_each_run_with_its_n() -> None: text = render(results) assert "AKI stage 3, eICU-CRD" in text assert "$n=50$" in text + + +def test_render_labels_the_random_arm_block() -> None: + run = { + "concept": "sepsis3", + "source": "mimic_iv", + "n_subjects": 50, + "selection": "random", + "cells": [ + {"event": "death", "horizon": "8h", "agree": 25, "total": 50, "pct": 50.0} + ], + } + text = render({"runs": [run]}) + assert "Sepsis-3, MIMIC-IV, random codes" in text + run["selection"] = "attributed" + assert "random codes" not in render({"runs": [run]}) diff --git a/tests/scripts/test_make_edit_attribution_table.py b/tests/scripts/test_make_edit_attribution_table.py index 2aaa139c..a52c519d 100644 --- a/tests/scripts/test_make_edit_attribution_table.py +++ b/tests/scripts/test_make_edit_attribution_table.py @@ -109,3 +109,55 @@ def test_render_includes_n_and_labels() -> None: assert "AKI stage 3" in text assert "eICU-CRD" in text assert " 50 " in text or "& 50 &" in text + + +def _death_cells(agree: int) -> list[dict]: + return [ + { + "event": "death", + "horizon": h, + "agree": agree, + "total": 50, + "pct": 2.0 * agree, + } + for h in ("8h", "24h", "72h") + ] + + +def _two_arm_results() -> dict: + attributed = _run("sepsis3", "mimic_iv", 50, _death_cells(47)) + attributed["selection"] = "attributed" + rnd = _run("sepsis3", "mimic_iv", 50, _death_cells(26)) + rnd["selection"] = "random" + rnd["random_seed"] = 0 + # file order puts the random run first on purpose: the renderer must + # place it under its attributed twin regardless + return {"runs": [rnd, attributed]} + + +def test_render_default_drops_random_runs_and_has_no_codes_column() -> None: + text = render(_two_arm_results()) + untagged = render({"runs": [_run("sepsis3", "mimic_iv", 50, _death_cells(47))]}) + assert text == untagged + assert "Codes" not in text + assert "random" not in text + assert "94\\%" in text + assert "52\\%" not in text + + +def test_render_with_random_adds_codes_column_and_pairs_rows() -> None: + text = render(_two_arm_results(), with_random=True) + assert "Concept & Dataset & Codes & $n$" in text + assert "\\begin{tabular}{lllrrrrr}" in text + rows = [line for line in text.splitlines() if line.startswith("Sepsis-3")] + assert len(rows) == 2 + assert rows[0].startswith("Sepsis-3 & MIMIC-IV & attributed & 50 & 94\\%") + assert rows[1].startswith("Sepsis-3 & MIMIC-IV & random & 50 & 52\\%") + + +def test_render_with_random_keeps_unpaired_random_run() -> None: + rnd = _run("qsofa", "eicu", 50, _death_cells(30)) + rnd["selection"] = "random" + text = render({"runs": [rnd]}, with_random=True) + assert "qSOFA & eICU-CRD & random & 50 & 60\\%" in text + assert "qSOFA" not in render({"runs": [rnd]}) diff --git a/tests/scripts/test_parse_edit_attribution_logs.py b/tests/scripts/test_parse_edit_attribution_logs.py index b673ad4d..911f3751 100644 --- a/tests/scripts/test_parse_edit_attribution_logs.py +++ b/tests/scripts/test_parse_edit_attribution_logs.py @@ -80,3 +80,30 @@ def test_parse_log_ignores_per_subject_lines() -> None: # mistaken for sign-agreement rows. result = parse_log(SAMPLE) assert not any(c["event"].startswith("subject") for c in result["cells"]) + + +def test_parse_log_untagged_header_is_the_attributed_arm() -> None: + result = parse_log(SAMPLE) + assert result["selection"] == "attributed" + assert result["random_seed"] is None + + +def test_parse_log_reads_selection_and_seed_tags() -> None: + tagged = SAMPLE.replace( + "top_k=4, index_frac=0.3333 ===", + "top_k=4, index_frac=0.3333, selection=random, seed=7 ===", + ) + result = parse_log(tagged) + assert result["selection"] == "random" + assert result["random_seed"] == 7 + assert result["n_subjects"] == 2 + assert len(result["cells"]) == 3 + + +def test_parse_log_attributed_tag_without_seed() -> None: + tagged = SAMPLE.replace( + "top_k=4, index_frac=0.3333 ===", "top_k=4, selection=attributed ===" + ) + result = parse_log(tagged) + assert result["selection"] == "attributed" + assert result["random_seed"] is None From 76ccf6da6f0310b5a47a588484c13941a05ce426 Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Fri, 9 Oct 2026 09:39:03 -0400 Subject: [PATCH 3/4] Rebuttal plan: WP9 results (random-code control, visit-end override) Co-Authored-By: Claude Fable 5.1 --- docs/ml4h2026_rebuttal_plan.md | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/docs/ml4h2026_rebuttal_plan.md b/docs/ml4h2026_rebuttal_plan.md index 7fb3c61e..5f69ea41 100644 --- a/docs/ml4h2026_rebuttal_plan.md +++ b/docs/ml4h2026_rebuttal_plan.md @@ -254,6 +254,24 @@ Two facts found on the way. First, the banked eICU Table 13 AKI cells are under - Parameters: 32,233,101 (MIMIC-IV flagship) and 20,384,390 (eICU-CRD), the difference being the per-source next-event vocabulary head; `research_journal/figure_data/param_counts.json`. `paper/ml4h/tables/hparams.tex` generated by `scripts/make_hparams_table.py` with the GBM grid, estimator and panel sizes read from the code. - Cohort counts (`cohort_counts.json` under both flagship run directories): MIMIC-IV 291,702 / 36,463 / 36,462 subjects (train / tuning / held-out), 435,803 / 54,898 / 55,328 admissions, median stay 2.8 days, 53% female; per-subject prevalence vasopressor 5.0%, ICU admission 17.9%, AKI 18.2%, death 10.5%, Sepsis-3 12.6%, 30-day readmission 13.8%. eICU-CRD 133,084 / 16,636 / 16,635 subjects, 160,643 / 20,094 / 20,122 stays, years 2014 to 2016, median stay 5.5 days; prevalence vasopressor 27.8%, AKI 55.9%, death 8.8%, Sepsis-3 0.28%, readmission 10.5%. +### WP9, reviewer-requested controls (run 2026-10-09): done, both negative for the reviewer's hypotheses + +Both were asked for by Reviewer 3P7y after the reviews landed; both ran on the submitted flagships with the VMs restarted for one day (stopped again 13:40 UTC). + +**Random-code control for Table 9.** Logs and per-subject JSONs under `research_journal/figure_data/edit_attribution/rebuttal/{mimic,eicu}/cohort_[_k1|_k2]_{attributed,random,random_s1}.*` (top-k 4, 1 and 2; random seeds 0 and 1; same 50 subjects per concept and arm, 24-hour window, `--index-hours 24`, four held-out shards). The attributed arm at top-k 4 reproduces the published 87 to 97% band. + +- Sign agreement is near-saturated: on MIMIC-IV any worsened panel signal raises every hazard, so random codes reach 96 to 100% on AKI and death, the same as the attributed codes. At top-k 4 the MIMIC-IV pool exceeds four codes for only 6 of 50 subjects, so the two arms coincide there; the control is read at top-k 1 and 2, where MIMIC-IV separates only on the ICU-admission and vasopressor cells (attributed 86 to 94% versus random 80 to 86% at one code). +- eICU-CRD windows hold more editable codes (pool larger than k for 40 to 50 of 50 subjects): the attributed codes beat random codes by 10 to 20 points at one or two codes (AKI 90 to 98% versus 74 to 88%, vasopressor 86 to 94% versus 70 to 90%) and by 4 to 8 points at four. +- What the ranking buys is targeting and effect size: the attributed codes move the target concept's probability 1.3 to 2 times further than random codes (Sepsis-3 on MIMIC-IV +0.108 versus +0.082 at one code, AKI stage 3 on eICU-CRD +0.040 versus +0.016) and change the hazards 1.5 to 7 times more. +- Camera-ready: random-code rows go into Table 9 (`scripts/make_edit_attribution_table.py --with-random`), and the text says that sign agreement is a weak discriminator; the concept shift and the hazard magnitude are what the ranking improves. + +**Override applied only at visit ends (3P7y Q1).** `interventions_band15_visitend.json` under `vm1/full_run_v10/` and `vm2/eicu_full_v10/` (branch `rebuttal/visit-end-override`, `--override-positions visit_end`, band 0.15, modes none / truth / flip / random). The override is applied only at the last position of each visit, the position the concept loss is pooled on, and every other position is left as under `none`; the `visit_end_positions` block scores next-event top-1 and loss at the visit-end positions that have an observed concept inside the band and a next token, paired across modes with a subject-clustered bootstrap (1,000 resamples). + +- MIMIC-IV (14,848 positions, 8,764 subjects): top-1 8.03% none, 7.88% truth, 7.99% flip, 7.98% random. Truth minus none -0.15 points (-0.25 to -0.04) and +0.018 nats (+0.013 to +0.024), both separated and both worse; truth minus flip -0.11 points (-0.25 to +0.05) and -0.009 nats (-0.020 to +0.002), ties. +- eICU-CRD (3,279 positions, 2,636 subjects): top-1 27.45% none, 27.26% truth, 27.11% flip, 27.02% random. Truth minus none -0.18 points (-0.61 to +0.27), tie, with loss +0.044 nats (+0.035 to +0.053), separated and worse; truth minus flip +0.15 points (-0.46 to +0.73) and +0.016 nats (-0.003 to +0.035), ties. +- Top-1 at these positions is low on both databases because the target is the first event of the next visit. The whole-stream top-1 is unchanged to four decimals because only 15,097 (MIMIC-IV) and 18,622 (eICU-CRD) of 87M positions are edited. +- Reading: restricting the override to the supervised positions does not create a lever, so the untrained-position explanation in Appendix E does not hold and the appendix will say so, with this variant next to the all-positions result. + ### Code landed on `rebuttal/integration` Event pinning for old checkpoints (eICU checkpoints trained before PR #222 have five hazard heads; the registry now builds six because Sepsis-3 resolves; `run_pins.json` plus a legacy table), hazard-head scoring in `interventions.py` (`--hazard-heads`), `scripts/probe_channel.py` (with bank save and resume), the three GEMINI run.sh steps, `scripts/make_hparams_table.py`, `scripts/cohort_counts.py`, `scripts/panel_coverage.py`, `scripts/compare_runs_paired.py`. 1,503 tests pass. Not merged to main yet. From 890e07016d82083d5f68acef5b6cce98e920b613 Mon Sep 17 00:00:00 2001 From: Amrit Krishnan Date: Fri, 9 Oct 2026 13:59:33 -0400 Subject: [PATCH 4/4] Cover the cohort CLI and the visit-end driver path codecov/patch was at 79% against the 91% project target: the new concept_edit_attribution CLI (_main) and the run-dir driver's visit-end branch had no tests. Adds a CLI test with load_run, the concept pins and the shard loader stubbed, the cohort_worsen edge paths (unknown concept, max_subjects, empty pool, an unedited random draw under pool=all, index_frac in the header), and an evaluate_interventions test with a two-admission fixture that attaches the visit-end blocks, plus the override-positions guard and the empty-table case. Co-Authored-By: Claude Fable 5.1 --- .../test_concept_edit_attribution.py | 155 +++++++++++++++++- .../inference/test_interventions_visit_end.py | 113 ++++++++++++- 2 files changed, 266 insertions(+), 2 deletions(-) diff --git a/tests/odyssey/inference/test_concept_edit_attribution.py b/tests/odyssey/inference/test_concept_edit_attribution.py index 7efd4f52..fea59822 100644 --- a/tests/odyssey/inference/test_concept_edit_attribution.py +++ b/tests/odyssey/inference/test_concept_edit_attribution.py @@ -1,6 +1,9 @@ """Occlusion attribution: which codes a concept's probability rests on.""" +import json from datetime import datetime, timedelta +from pathlib import Path +from types import SimpleNamespace import polars as pl import pytest @@ -559,7 +562,12 @@ def test_worsen_edit_for_signal_covers_qsofa_resp_rate() -> None: # Cohort validation and the random-code control # --------------------------------------------------------------------------- +import odyssey.inference.legacy_concept_pins as pins_module # noqa: E402 +import odyssey.inference.run_inference as run_inference_module # noqa: E402 +import odyssey.training.data as data_module # noqa: E402 from odyssey.inference.concept_edit_attribution import ( # noqa: E402 + CohortWorsenResult, + _main, cohort_worsen, format_cohort_summary, mappable_candidate_codes, @@ -647,7 +655,9 @@ def test_index_frac_places_the_cut_inside_the_visit() -> None: _index_times_by_subject(events, index_hours=24.0, index_frac=0.0) -def _run_cohort(events: pl.DataFrame, **kwargs: object) -> tuple[list[str], object]: +def _run_cohort( + events: pl.DataFrame, **kwargs: object +) -> tuple[list[str], CohortWorsenResult]: vocab, model, concepts = _vocab_and_model(events) lines: list[str] = [] result = cohort_worsen( @@ -760,3 +770,146 @@ def test_format_cohort_summary_round_trips_through_the_parser() -> None: # the per-subject lines carry the signals, like the banked logs assert "edit signal frequency: {" in att_text assert "mean concept-probability delta: " in att_text + + +# --------------------------------------------------------------------------- +# cohort_worsen edge paths and the CLI entry point +# --------------------------------------------------------------------------- + + +def test_cohort_worsen_unknown_concept_raises() -> None: + events = _cohort_events(1) + vocab, model, concepts = _vocab_and_model(events) + with pytest.raises(ValueError, match="not in concept_names"): + cohort_worsen( + model, + vocab, + None, + events, + concept_name="not_a_concept", + concept_names=concepts, + chunk_size=16, + ) + + +def test_cohort_worsen_stops_at_max_subjects() -> None: + _, result = _run_cohort(_cohort_events(3), max_subjects=1) + assert [s.subject_id for s in result.subjects] == [1] + + +def test_cohort_worsen_skips_subjects_with_an_empty_pool() -> None: + # subject 3 has a 30 h visit but nothing worsenable in the window, so it + # is skipped in both arms rather than scored with nothing to edit. + events = _cohort_events(2) + extra = [ + (3, "HOSPITAL_ADMISSION//EMERGENCY", T0, None, 103), + *[(3, DIAG, T0 + timedelta(hours=h), None, 103) for h in range(1, 30)], + (3, "HOSPITAL_DISCHARGE//HOME", T0 + timedelta(hours=30), None, 103), + ] + events = pl.concat([events, pl.DataFrame(extra, schema=SCHEMA, orient="row")]) + _, attributed = _run_cohort(events) + _, rnd = _run_cohort(events, code_selection="random", random_seed=0) + assert [s.subject_id for s in attributed.subjects] == [1, 2] + assert [s.subject_id for s in rnd.subjects] == [1, 2] + + +def test_cohort_worsen_random_all_pool_can_leave_a_subject_unedited() -> None: + # seed 6 draws glucose + the diagnosis for subject 1 (neither has a + # worsen edit) and worsenable codes for subject 2; the unedited subject + # is kept, scored as a no-op, and counted in the header. + events = _cohort_events(2) + _, result = _run_cohort( + events, code_selection="random", random_seed=6, candidate_pool="all" + ) + by_sid = {s.subject_id: s for s in result.subjects} + assert set(by_sid) == {1, 2} + assert by_sid[1].rows_edited == 0 + assert by_sid[1].signals == [] + assert by_sid[1].concept_after == by_sid[1].concept_before + assert all( + v == 0.0 for hs in by_sid[1].delta_event_risk.values() for v in hs.values() + ) + assert by_sid[2].rows_edited > 0 + assert [s.subject_id for s in result.scored] == [2] + header = format_cohort_summary(result).splitlines()[0] + assert "pool=all" in header + assert header.endswith("no_edit=1 ===") + + +def test_format_cohort_summary_reports_index_frac() -> None: + _, result = _run_cohort(_cohort_events(1), index_frac=0.5) + header = format_cohort_summary(result).splitlines()[0] + assert "index_frac=0.5000" in header + assert result.index_frac == 0.5 + + +def test_cli_main_runs_the_cohort_and_writes_the_json( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + events = _cohort_events(2) + vocab, model, concepts = _vocab_and_model(events) + config = SimpleNamespace(source="mimic_iv", task_set="v1") + seen: dict[str, object] = {} + + def fake_load_run(run_dir: Path, **kwargs: object) -> tuple[object, ...]: + seen["run_dir"] = run_dir + seen["checkpoint"] = kwargs["checkpoint_path"] + return model, vocab, None, config + + monkeypatch.setattr(run_inference_module, "load_run", fake_load_run) + monkeypatch.setattr( + pins_module, + "resolve_concepts_for_run", + lambda run_dir, source, task_set: concepts_for_source(source), + ) + monkeypatch.setattr( + data_module, "load_meds_shards", lambda path, max_shards: events.clone() + ) + out = tmp_path / "nested" / "cohort_random.json" + argv = [ + "prog", + "--run-dir", + str(tmp_path / "run"), + "--held-out-shard-dir", + str(tmp_path / "held_out"), + "--concept", + concepts[0], + "--output-json", + str(out), + "--top-k", + "2", + "--max-subjects", + "2", + "--code-selection", + "random", + "--random-seed", + "3", + "--chunk-size", + "16", + "--device", + "cpu", + ] + monkeypatch.setattr("sys.argv", argv) + _main() + + assert seen["run_dir"] == tmp_path / "run" + assert seen["checkpoint"] == tmp_path / "run" / "checkpoint_best.pt" + payload = json.loads(out.read_text()) + assert payload["concept"] == concepts[0] + assert payload["code_selection"] == "random" + assert payload["random_seed"] == 3 + assert payload["top_k"] == 2 + assert payload["run_dir"] == str(tmp_path / "run") + assert [s["subject_id"] for s in payload["subjects"]] == [1, 2] + assert set(payload["sign_agreement"]) == {"vasopressor_start", "death"} + for horizons in payload["sign_agreement"].values(): + for cell in horizons.values(): + assert set(cell) == {"agree", "total"} + assert cell["total"] == 2 + parsed = parse_log(payload["summary_text"]) + assert parsed["selection"] == "random" + assert parsed["random_seed"] == 3 + printed = capsys.readouterr().out + assert "worsenable signals with a reading in the loaded shards:" in printed + assert "subject 1 [1/2]:" in printed + assert payload["summary_text"] in printed diff --git a/tests/odyssey/inference/test_interventions_visit_end.py b/tests/odyssey/inference/test_interventions_visit_end.py index 7f518e6b..e6b5847c 100644 --- a/tests/odyssey/inference/test_interventions_visit_end.py +++ b/tests/odyssey/inference/test_interventions_visit_end.py @@ -6,8 +6,9 @@ """ import json -from datetime import timedelta +from datetime import datetime, timedelta from pathlib import Path +from typing import Any import numpy as np import polars as pl @@ -16,11 +17,14 @@ import odyssey.inference.interventions as interventions_module from odyssey.data.streaming import PackedLaneSampler, StreamingChunk +from odyssey.data.value_binning import add_value_tokens +from odyssey.data.vocabulary import Vocabulary from odyssey.inference.interventions import ( OVERRIDE_POSITIONS, InterventionResult, VisitEndScorer, _chunk_intervention, + evaluate_interventions, result_to_json, run_streaming_intervention, score_visit_end_rows, @@ -30,6 +34,7 @@ from odyssey.models.concept_bottleneck import intervention_apply_mask from odyssey.training.data import iter_patient_sequences from odyssey.training.running_labels import position_running_labels +from odyssey.training.train import TrainingConfig from tests.odyssey.inference.test_interventions import ( CODES, NUM_CONCEPTS, @@ -39,6 +44,7 @@ ) from tests.odyssey.inference.test_interventions_hazard import OLD_KEYS from tests.odyssey.inference.test_interventions_hazard import _Fixture as _HazardFixture +from tests.odyssey.inference.test_interventions_hazard import _model as _hazard_model SUBJECTS = (1, 2, 3) @@ -463,3 +469,108 @@ def test_hazard_heads_combine_with_the_visit_end_override() -> None: assert combined.visit_end_positions is not None # Far fewer positions are overridden than under the default. assert 0 < combined.n_intervened_positions < hazard_only.n_intervened_positions + + +# --------------------------------------------------------------------------- +# the run-dir driver attaches the block, and the small guards +# --------------------------------------------------------------------------- + + +def test_score_visit_end_rows_of_nothing_is_empty() -> None: + assert score_visit_end_rows({}) == {} + + +def test_check_override_positions_rejects_unknown_and_warns_on_stay( + caplog: pytest.LogCaptureFixture, +) -> None: + with pytest.raises(ValueError, match="unknown override_positions"): + interventions_module._check_override_positions("landmark", "visit") + with caplog.at_level("WARNING"): + interventions_module._check_override_positions("visit_end", "visit") + assert not caplog.records + with caplog.at_level("WARNING"): + interventions_module._check_override_positions("visit_end", "stay") + assert any("concept_supervision='stay'" in r.getMessage() for r in caplog.records) + + +def _two_visit_events(n_subjects: int = 16) -> pl.DataFrame: + """Split the hazard fixture's hourly heart rates into two admissions. + + The hazard fixture has one admission per subject, so its only visit end + is the last position of each stream, which has no next-token target. + Splitting each record at hour 12 gives every subject one scorable visit + end. + """ + rows: list[tuple[int, str, datetime, float | None, int]] = [] + for sid in range(1, n_subjects + 1): + for h in range(24): + hadm = (1000 if h < 12 else 2000) + sid + hr = 130.0 if sid % 2 == 0 and h >= 12 else 80.0 + rows.append((sid, "LAB//220045//bpm", T0 + timedelta(hours=h), hr, hadm)) + return pl.DataFrame( + rows, + schema={ + "subject_id": pl.Int64, + "code": pl.Utf8, + "time": pl.Datetime("us"), + "numeric_value": pl.Float32, + "hadm_id": pl.Int64, + }, + orient="row", + ) + + +def test_evaluate_interventions_attaches_visit_end_blocks( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + raw = _two_visit_events() + vocab = Vocabulary.build(add_value_tokens(raw)["code"].to_list(), min_count=1) + model = _hazard_model(len(vocab)) + config = TrainingConfig( + train_shard_dir="/train", + tuning_shard_dir="/tuning", + output_dir="/out", + source="mimic_iv", + task_set="v1", + concept_supervision="visit", + ) + monkeypatch.setattr( + interventions_module, + "load_run", + lambda *a, **k: (model, vocab, None, config), + ) + monkeypatch.setattr( + interventions_module, "load_meds_shards", lambda *a, **k: raw.clone() + ) + held_out = tmp_path / "held_out" + held_out.mkdir() + with caplog.at_level("INFO"): + results = evaluate_interventions( + tmp_path, + held_out, + modes=["none", "truth", "flip"], + num_lanes=2, + chunk_size=16, + device="cpu", + override_positions="visit_end", + hazard_boot=20, + ) + by_mode = {r.mode: r for r in results} + assert set(by_mode) == {"none", "truth", "flip"} + assert all(r.visit_end_positions is not None for r in results) + blocks: dict[str, dict[str, Any]] = { + m: r.visit_end_positions or {} for m, r in by_mode.items() + } + n = blocks["none"]["n_positions"] + assert n > 0 + assert all(b["n_positions"] == n for b in blocks.values()) + assert set(blocks["truth"]["paired"]) == {"truth_minus_none", "truth_minus_flip"} + assert set(blocks["flip"]["paired"]) == {"flip_minus_none"} + assert blocks["none"]["paired"] == {} + assert all(r.hazard is None for r in results) + messages = [r.getMessage() for r in caplog.records] + assert any("paired visit-end bootstrap" in m for m in messages) + assert any("none at visit ends: top1" in m for m in messages) + # the JSON carries the block under the old keys + written = result_to_json(by_mode["truth"]) + assert written["visit_end_positions"] == blocks["truth"]