diff --git a/README.md b/README.md index 4df51ab..5e452e0 100644 --- a/README.md +++ b/README.md @@ -43,7 +43,7 @@ Fusion vs. unimodal accuracy is **supporting evidence**, not the thesis. | **Cohort construction** | MIMIC-IV-Echo studies, MIMIC-IV-ECG records, MIMIC-IV demographics | Paired rows on `subject_id` within 24–48h, joined to structured LVEF, split by subject | | **Embedding extraction** | Paired cohort only (~few K–tens of K studies, not all 525K echos) | Frozen EchoJEPA-L (1024-d) and HuBERT-ECG (768-d) vectors, cached to Parquet | | **Probes** | Cached embeddings | ECG-only, echo-only (attentive), concat-MLP, and cross-attention fusion heads | -| **Deployment analyses** | Trained fused checkpoint + held-out test split | Missing-modality degradation, fairness stratification, optional calibration | +| **Deployment analyses** | Trained fused checkpoint + held-out test split | Missing-modality degradation, fairness stratification, EF≤40% calibration | Expensive forward passes run once and are cached; probe training reads only from cache and completes in minutes. Foundation-model weights are never fine-tuned. @@ -53,26 +53,73 @@ See [TECHNICAL.md](./TECHNICAL.md) for full pipeline details. ## Results -Held-out test split, n = 245. One fused cross-attention checkpoint (M09) trained on both modalities, scored under three inference-time conditions — no separate unimodal models are trained for the dropped conditions. +Held-out test split, n = 245 after dropping three non-finite HuBERT-ECG rows from the original 248-row test split. One fused cross-attention checkpoint (M09) trained on both modalities, scored under three inference-time conditions — no separate unimodal models are trained for the dropped conditions. -| Condition | Inference input | LVEF MAE | EF≤40% AUROC | -|-----------|-----------------|---------:|-------------:| -| `full` | Echo + ECG present | **10.28** | **0.766** | -| `ecg_dropped` | ECG branch masked, echo present | 15.13 | 0.750 | -| `echo_dropped` | Echo branch masked, ECG present | 18.57 | 0.693 | +| Condition | Inference input | LVEF MAE (95% CI) | EF≤40% AUROC (95% CI) | +|-----------|-----------------|-------------------:|----------------------:| +| `full` | Echo + ECG present | **10.42** (9.28–11.64) | **0.771** (0.697–0.839) | +| `ecg_dropped` | ECG branch masked, echo present | 11.22 (10.00–12.39) | 0.383 (0.279–0.493) | +| `echo_dropped` | Echo branch masked, ECG present | 20.55 (19.02–22.02) | 0.689 (0.605–0.765) | -Dropping echo costs more than dropping ECG on both metrics. AUROC holds up better than MAE under either drop, which is the graceful-degradation pattern rather than silent failure. +Dropping echo costs much more MAE than dropping ECG, while ECG-dropped AUROC falls sharply. This run is a pooled-manifest baseline: the echo branch still receives one mean-pooled vector per study, not retained clip tokens — and on pooled inputs the cross-attention block provably cannot attend across modalities, so these numbers describe a concat-equivalent model (see [TECHNICAL.md §6.4](./TECHNICAL.md#6-probes)). -**Provenance:** real run on cached EchoJEPA (`vjepa2.1-vitl-mimic-pt-100`) + HuBERT-ECG embeddings, point estimates transcribed from the E02 evaluation into `results/missing_modality.json`. Bootstrap confidence intervals need per-example predictions and are pending a canonical rerun. `results/` is gitignored, so that JSON is not in this repository — see [CONTRIBUTING.md](./CONTRIBUTING.md#reproducibility) for what reproducing these numbers takes. +**Provenance:** real run on cached EchoJEPA (`vjepa2.1-vitl-mimic-pt-100`) + HuBERT-ECG embeddings, local Mac CPU, seed 42, fusion width 256, 1,000 bootstrap resamples. The canonical fused checkpoint is the one `scripts/train_probes.py` writes to `probes/fused/cross_attn_fused.pt`, selected by validation MAE from the all-probe M10 run (`full` val MAE 10.62); manifest SHA-256 `81694c9b…`, checkpoint SHA-256 `bac18bb8…`. `results/` and `probes/` are gitignored (per-example predictions and checkpoints stay local), but sanitized aggregate copies of every result JSON plus the full checksums are committed under [docs/results/](./docs/results/) — verify a reproduction with `shasum -c docs/results/SHA256SUMS`. See [CONTRIBUTING.md](./CONTRIBUTING.md#reproducibility) for what reproducing takes. Reproduce with: ```bash -python scripts/evaluate_missing_modality.py +python scripts/evaluate_missing_modality.py \ + --manifest data/processed/echo_hubert_manifest.parquet \ + --checkpoint probes/fused/cross_attn_fused.pt \ + --embed-dim 256 \ + --echo-dim 1024 \ + --ecg-dim 768 ``` Defaults assume the standard cohort/embedding/checkpoint layout; `--help` lists the paths. `--embed-dim`, `--echo-dim`, and `--ecg-dim` must match the checkpoint architecture or the `state_dict` load fails on shape. +### In-cohort baselines + +All four probes on identical splits (seed 42, same 245-row test frame, `scripts/diagnose_baseline_gap.py`): + +| Probe | LVEF MAE (95% CI) | EF≤40% AUROC (95% CI) | +|-------|-------------------:|----------------------:| +| Cross-attention fused | **10.42** (9.28–11.64) | **0.771** (0.697–0.839) | +| Concat-MLP | 10.93 (9.77–12.22) | 0.716 (0.634–0.795) | +| Echo-only (attentive) | 11.09 (9.79–12.25) | 0.771 (0.698–0.845) | +| ECG-only (ridge) | 11.60 (10.36–12.84) | 0.671 (0.582–0.756) | + +Fusion beats ECG-only on both metrics with paired-bootstrap significance (ΔMAE −1.17, CI −1.99 to −0.33; ΔAUROC +0.100, CI +0.03 to +0.17) and is never behind either solo probe — the gap to published single-modality baselines is a cohort and label-regime property, not a fusion failure ([TECHNICAL.md §8](./TECHNICAL.md#8-metrics)). + +### Fairness + +Test-split stratification of the fused checkpoint (`scripts/evaluate_fairness.py`, same run as above; strata under n=20 are flagged in the artifact and omitted here): + +| Stratum | n | MAE `full` | MAE `echo_dropped` | AUROC `full` | +|---------|--:|-----------:|-------------------:|-------------:| +| Sex F | 144 | 10.13 | 21.19 | 0.770 | +| Sex M | 101 | 10.84 | 19.63 | 0.763 | +| Age 40–54 | 37 | 8.93 | 20.44 | 0.742 | +| Age 55–64 | 57 | 9.43 | 23.67 | 0.727 | +| Age 65–74 | 57 | 13.10 | 21.71 | 0.752 | +| Age 75–89 | 77 | 10.29 | 18.29 | 0.789 | +| Race: Black | 43 | 6.62 | 21.41 | 0.973 | +| Race: White | 175 | 11.53 | 20.13 | 0.761 | + +The sex gap is small in both conditions. The largest full-condition gaps are age 65–74 (13.10 vs 10.42 overall) and the Black–White MAE difference (6.62 vs 11.53). Note that MIMIC-IV records administrative gender and admission-reported race, both with known curation bias, and the remaining race/age strata are too small for stable estimates (flagged, not silently included, in `results/fairness/`). + +### Calibration + +EF≤40% Platt scaler fit on val predictions only, applied to test (`scripts/evaluate_calibration.py`), 10 bins: + +| Condition | ECE | +|-----------|----:| +| `full` | 0.072 | +| `echo_dropped` | 0.043 | +| `ecg_dropped` | 0.015 | + +Post-scaling ECE stays low even with a modality dropped — but read `ecg_dropped` together with its 0.383 AUROC: predictions collapse toward the training mean, so the probabilities are calibrated yet nearly uninformative. Calibration alone is not evidence the degraded model is usable. + --- ## Models & data @@ -129,8 +176,11 @@ framework is the main contribution and lives in [src/primed_ai/failure/](./src/primed_ai/failure/); the fusion pipeline described above is the substrate it is instantiated on. -Still open: PhysioNet credentialing and reserved GPU/storage, both needed for the canonical -real-data rerun. +The canonical pooled rerun is done — probe training and every deployment analysis run on a +laptop CPU in about 90 seconds total, no GPU (see [TECHNICAL.md §9](./TECHNICAL.md#9-compute-and-infrastructure)). +Still open: rebuilding the manifest with clip-level echo tokens (needs the gated HF shards), +token-level ECG re-extraction (#72), and the cohort-expansion decision (#62) — the first two +gate real cross-modal attention, the last gates statistical power. --- diff --git a/TECHNICAL.md b/TECHNICAL.md index ad3c31a..c771b82 100644 --- a/TECHNICAL.md +++ b/TECHNICAL.md @@ -51,10 +51,19 @@ Build the paired cohort **before** running any embedding extraction. Do not proc For each echo study: -1. Match on `subject_id` to the nearest ECG record within a **24–48 hour** temporal window (using MIMIC timestamps). +1. Match on `subject_id` to the nearest ECG record within a symmetric **±24 hour** window + (`build_cohort.py --window-hours 24`, the canonical build; the built manifest's + max |delta_hours| is 23.98). 2. Join to the structured **LVEF** label from MIMIC-IV-Echo. -Expected cohort size: a few thousand to tens of thousands of paired rows (exact count depends on window strictness and label availability). +Current synchronized manifest size is 1,208 paired rows. Without extracting more +EchoJEPA studies, the hard ceiling is 6,617 echo studies from subjects who also have an +ECG; larger cohorts require additional echo embedding coverage, not just a looser join. + +Window sensitivity, measured on the built manifest: tightening to ±12h keeps 767 of the +1,208 rows and ±6h keeps 515. Widening beyond ±24h cannot be measured from the manifest — +it needs a cohort-database rebuild (`scripts/run_cohort_sensitivity.py` wraps the +24h/48h/admission comparison into one command for whoever has BigQuery access). ### 3.2 Train / test split @@ -108,7 +117,10 @@ Two ECG paths exist in this repo, and they are not interchangeable: Both are 768-d after pooling, so the probes accept either. Anything downstream that says "ECG embedding" below refers to the HuBERT-ECG Parquet unless stated otherwise. -**Open decision:** pooling strategy for ECG embeddings — mean vs. attentive pooling. Resolve with a quick empirical check early in probe development. +The current HuBERT-ECG Parquet stores one pooled vector per record. Mean-vs-attentive +ECG pooling is therefore not a runnable ablation yet: attentive pooling over tiled copies +of that vector is mathematically identical to mean pooling. Token-level ECG needs a +re-extraction pass before this decision can be revisited. --- @@ -142,6 +154,20 @@ All probes are trained on cached embeddings. Recommended build order: quick conc All four predict continuous LVEF; the EF≤40% gate is derived from that output. +**What the cross-attention block can actually attend to (#72).** With the current inputs +the fusion is one-directional at best. The HuBERT-ECG side arrives as one pooled vector +tiled into identical tokens, so the `echo_to_ecg` attention output is provably independent +of the echo query — identical keys force uniform softmax weights, and identical values make +the weighted sum equal that value. That direction reduces to a fixed linear image of the +pooled ECG vector (pinned by `test_echo_to_ecg_attention_is_degenerate_on_tiled_ecg_tokens`). +`ecg_to_echo` attends genuinely only when clip-level echo tokens are retained +(`build_echo_study_embeddings --max-clips`); on the pooled manifest both directions +degenerate and the fused probe is equivalent to a concat model with per-modality linear +pre-maps. Reported pooled-manifest numbers should be read accordingly. Decision: keep the +cross-attention architecture — it fuses for real once clip-level echo tokens are in the +manifest — and treat token-level ECG (which needs a re-extraction pass from raw waveforms; +the parquet stores pooled vectors only) as the gate for restoring the second direction. + ### 6.5 Prediction targets | Target | Type | Use | @@ -174,7 +200,11 @@ For each condition, report: Plot a **degradation curve** across conditions. The key question: when echo is unavailable at inference, does the model degrade gracefully or fail silently? -**Implementation note:** at echo-dropped evaluation, zero out or mask the echo branch of the fused probe rather than retraining a separate ECG-only model. This tests the actual deployed fused model under missing input. +**Implementation note:** at dropped-modality evaluation, mask the selected fused-probe +branch at inference time (`mask_echo` / `mask_ecg`) rather than retraining a unimodal +model. This tests the deployed fused checkpoint under missing input. Learned null tokens +and branch dropout are deferred unless the missing-modality rerun shows masking is +unstable. ### 7.2 Fairness audit @@ -193,9 +223,16 @@ Report per-stratum: Flag known MIMIC gender-curation bias in documentation. This analysis is cheap and directly addresses the equity evaluation pillar. -### 7.3 Calibration (optional stretch) +### 7.3 Calibration -If time permits, assess calibration of the EF≤40% binary gate (e.g., reliability diagram, expected calibration error). A deployable risk score should produce well-calibrated probabilities, not just high AUROC. +Assess calibration of the EF≤40% binary gate (reliability diagram, expected calibration +error). A deployable risk score should produce well-calibrated probabilities, not just high +AUROC. The Platt scaler is fit on val predictions only and applied to test +(`scripts/evaluate_calibration.py`); measured numbers are in the +[README results](README.md#results). Caveat when reading dropped-condition ECE: a model +whose predictions collapse toward the training mean can look well calibrated after Platt +scaling while discriminating no better than chance — report ECE alongside AUROC, never +alone. --- @@ -203,11 +240,23 @@ If time permits, assess calibration of the EF≤40% binary gate (e.g., reliabili | Metric | Target | Baseline references | |---|---|---| -| LVEF MAE | Continuous regression | EchoJEPA: 5.97 MAE | -| EF≤40% AUROC | Binary clinical gate | ECG-FM: 0.929 AUROC | +| LVEF MAE | Continuous regression | EchoJEPA: 5.97 MAE (published, different cohort) · in-cohort echo-only probe: 11.09 | +| EF≤40% AUROC | Binary clinical gate | ECG-FM: 0.929 AUROC (published, different cohort) · in-cohort ECG-only probe: 0.671 | | Missing-modality degradation | Δ MAE / Δ AUROC across conditions | No external baseline — this is the novel result | | Fairness gap | Δ MAE / Δ AUROC across demographic strata | No external baseline | +The published solo numbers are **not** like-for-like with this cohort. On identical splits +(seed 42, 245-row test frame, `scripts/diagnose_baseline_gap.py`): fused 10.42 MAE / 0.771 +AUROC vs echo-only 11.09 / 0.771, ECG-only 11.60 / 0.671, concat 10.93 / 0.716. Paired +bootstrap deltas: fused beats ECG-only on both metrics (ΔMAE −1.17, 95% CI [−1.99, −0.33]; +ΔAUROC +0.100, CI [0.03, 0.17]) and is never behind either solo probe. The gap to the +published baselines is therefore a property of the cohort and label regime, not a fusion +failure — even the in-cohort echo-only probe (same encoder as the published 5.97) lands at +11.09. Contributing factors, measured: 821 training rows after the non-finite drop (829 in +the split); heterogeneous LVEF label sources (the `lvef_upper` fallback contributes 6 test +rows at 43.0 MAE — 18 of its 24 cohort rows carry a physiologically implausible 100.0 — +versus 5.2 MAE on `lvef_3d` rows); and the mean-pooled echo regime (§6.4). + **Pre-flight check:** confirm EF≤40% prevalence in the paired cohort is high enough for stable AUROC estimation before locking results (`scripts/check_ef40_prevalence.py`). Measured missing-modality numbers are in the [README results table](README.md#results); reproducing them is covered in [CONTRIBUTING.md](CONTRIBUTING.md#reproducibility). @@ -218,9 +267,9 @@ Measured missing-modality numbers are in the [README results table](README.md#re | Resource | Requirement | |---|---| -| GPU | H200 or equivalent; reserve early on ORCD | -| Storage | Sufficient for cached embedding tensors across full paired cohort | -| Runtime | Subsetting to paired cohort keeps extraction under ~1 hour (vs. hours for full 525K echo corpus) | +| GPU | Not needed for cached-vector probe training; needed only for future ECG token re-extraction or new EchoJEPA extraction | +| Storage | Pooled manifest 14 MB, one four-probe checkpoint set 7 MB, full results tree ~1 MB — everything fits under `data/`, `probes/`, `results/` in the repo working copy; `--max-clips` echo manifests scale linearly with retained clips | +| Runtime | All CPU, measured on an Apple M1 Pro (16 GB): four-probe M10 training 47.7s wall-clock, missing-modality eval with 1,000 bootstrap resamples 11.8s, fairness stratification 8.3s, per-condition calibration ~5s. A full train + evaluate cycle is about 90 seconds | | Reproducibility | Fixed random seeds; logged hyperparameters; versioned embedding cache | Probe training after caching is CPU/GPU-light and completes in minutes. @@ -242,12 +291,17 @@ This work differs from EchoingECG on three axes: frozen embeddings (no fine-tuni ## 11. Open technical decisions -- [ ] ECG pooling: mean vs. attentive — empirical check on validation set -- [ ] Echo↔ECG pairing window: 24h vs. 48h — sensitivity analysis or fixed choice with justification -- [ ] Multi-match resolution: when multiple ECGs fall within the window, take nearest timestamp -- [ ] EF≤40% prevalence in paired cohort — confirm before reporting AUROC -- [ ] GPU + storage budget for embedding cache on ORCD -- [ ] Missing-modality masking strategy: zero-out vs. learned null token vs. branch dropout at eval only +- [x] ECG pooling: mean — the parquet stores one pooled vector per record, and attentive + pooling over tiled copies of it is mathematically identical to mean pooling, so this + is the only runnable choice. Revisit only after a token-level ECG re-extraction (#72). +- [x] Echo↔ECG pairing window: fixed at ±24h — this is what the canonical cohort was built + with (max |delta_hours| 23.98). Within-manifest tightening loses rows fast (767 at + ±12h, 515 at ±6h); widening to 48h needs a cohort-database rebuild and is scoped + under D07 (#62) via `scripts/run_cohort_sensitivity.py`. +- [x] Multi-match resolution: when multiple ECGs fall within the window, take nearest timestamp +- [x] EF≤40% prevalence in paired cohort: train/test are reportable; validation has only 26 EF<=40 positives, so validation AUROC should be treated as unstable +- [x] GPU + storage budget for cached probe training: no GPU needed; current pooled manifest and checkpoints are laptop-scale +- [x] Missing-modality masking strategy: use inference-time branch masks; no learned null token or branch dropout for the canonical rerun --- diff --git a/configs/cohort/default.yaml b/configs/cohort/default.yaml index 5e8d262..20fb210 100644 --- a/configs/cohort/default.yaml +++ b/configs/cohort/default.yaml @@ -1,5 +1,8 @@ -pairing_window_hours: 48 +# Symmetric window in hours on EACH side of the echo (build_cohort.py --window-hours +# semantics). The canonical 1,208-row manifest was built at 24, i.e. ECG within +/-24h +# of the echo; the built manifest's max |delta_hours| is 23.98. +pairing_window_hours: 24 split_seed: 42 train_frac: 0.70 val_frac: 0.10 -test_frac: 0.20 \ No newline at end of file +test_frac: 0.20 diff --git a/configs/encoder/ecg_fm.yaml b/configs/encoder/ecg_fm.yaml index 81a7ee9..13ca2ae 100644 --- a/configs/encoder/ecg_fm.yaml +++ b/configs/encoder/ecg_fm.yaml @@ -1,6 +1,8 @@ name: ecg_fm repo_id: wanglab/ecg-fm checkpoint: mimic_iv_ecg_physionet_pretrained.pt +# Intentionally a local override: ECG-FM weights are only needed for the alternative +# extraction path, which has produced no reported number. Point at your own download. checkpoint_path: /path/to/shared/weights/mimic_iv_ecg_physionet_pretrained.pt embed_dim: 768 seq_len: 64 @@ -8,11 +10,16 @@ sample_rate: 500 signal_length: 5000 n_leads: 12 frozen: true +# ECG pooling decision (E11/#69): the HuBERT-ECG parquet stores one pooled vector per +# record, so mean pooling is the only runnable option — attentive pooling over tiled +# copies of a pooled vector is mathematically identical. Revisit only after a +# token-level ECG re-extraction (#72). pooling: mean # Set use_stub: true for offline dev/CI without fairseq_signals + checkpoint. use_stub: false # Everything above configures the ECG-FM wrapper (encoders/ecg_fm.py), an alternative # extraction path that has not produced any reported result. The probes and evaluation # scripts read the pre-extracted HuBERT-ECG embeddings below — these are the ones behind -# every reported ECG number. -hubert_ecg_parquet: /path/to/shared/ecg-hubert/mimic-iv-ecg-ve.parquet +# every reported ECG number. The repo-local canonical copy lands here via the team share; +# see docs/embeddings.md §2 and docs/echo_hubert_results.md for how to obtain it. +hubert_ecg_parquet: data/interim/hubert_ecg_embeddings.parquet diff --git a/configs/encoder/echojepa.yaml b/configs/encoder/echojepa.yaml index 8f9487b..632662a 100644 --- a/configs/encoder/echojepa.yaml +++ b/configs/encoder/echojepa.yaml @@ -1,7 +1,14 @@ name: echojepa +# Intentionally a local override: echo checkpoints live on the MIT ORCD pool under +# weights/ (docs/embeddings.md §3). Only needed to re-extract embeddings; probes read +# the pre-extracted parquet. checkpoint_path: /path/to/shared/weights/vjepa21_vitl_mimic_pt117.pt embed_dim: 1024 img_size: 256 num_frames: 16 frozen: true +# Echo pooling decision (E11/#69): study-level mean pooling is the canonical regime +# behind every reported number. Clip-level retention exists via +# build_echo_study_embeddings(..., max_clips=N); the pooled-vs-clip val comparison +# stays open until a clip-level manifest is rebuilt from the gated HF shards. pooling: mean diff --git a/configs/paths/default.yaml b/configs/paths/default.yaml index 0c013b7..702e164 100644 --- a/configs/paths/default.yaml +++ b/configs/paths/default.yaml @@ -1,5 +1,8 @@ +# Intentionally a local override: raw MIMIC access is only needed to re-extract +# embeddings or rebuild the cohort. Probe training and every evaluation script read the +# joined manifest at data/processed/echo_hubert_manifest.parquet instead. data_root: /path/to/mimic cohort_dir: cohort embedding_dir: embeddings results_dir: results -logs_dir: logs \ No newline at end of file +logs_dir: logs diff --git a/docs/results/README.md b/docs/results/README.md new file mode 100644 index 0000000..10c439b --- /dev/null +++ b/docs/results/README.md @@ -0,0 +1,18 @@ +# Sanitized result artifacts + +Aggregate-only copies of the local result JSONs backing the numbers in the top-level +README — per-example predictions and checkpoints stay local (`results/` and `probes/` +are gitignored), so this is what reviewers can diff without data access. + +Written by `scripts/export_result_bundle.py`, which strips the per-example +`predictions` / `predictions_val` blocks from the missing-modality JSON and copies the +rest verbatim. `SHA256SUMS` covers both these files and the local raw artifacts +(manifest, the four canonical checkpoints, the full missing-modality JSON); verify a +reproduction from the repo root with: + +```bash +shasum -c docs/results/SHA256SUMS +``` + +Lines for local raw artifacts fail with "No such file" until you have rebuilt them — +that is the checklist of what your reproduction still has to produce. diff --git a/docs/results/SHA256SUMS b/docs/results/SHA256SUMS new file mode 100644 index 0000000..e80a7c7 --- /dev/null +++ b/docs/results/SHA256SUMS @@ -0,0 +1,13 @@ +600e57384304eae25ba2c13b06c1fca8eacf8c24adc6e1a215d6649b70c8997e docs/results/missing_modality.metrics.json +c2d5cb552b3638e42f14ce393b103af5dc912e76873e4d22863e2c8cc7ee3f9b docs/results/fairness_metrics.json +306fbcfd7171fd3981420cea1fe4baecd49f50ed85987ba413fc4efd8856ae5f docs/results/calibration_full.json +ddbfa2eb4cf6255cc80f15fd96e73fae33c036eafdf8775d1369b0fd2ac1a14d docs/results/calibration_echo_dropped.json +20ebce7c0442c3c5f4cb57244ee1e8996c5b72741a68c2d4368de7501dd6e1e1 docs/results/calibration_ecg_dropped.json +9f78fe365b42d5cea80113d21077d63656f0d7f4ddb5c6f4423c6170c32407c3 docs/results/baseline_gap.json +54b247c296c0c08deb64d2dfd1b0e940a0b6ac7e72e408e93d3f3a9f53530560 docs/results/failure_report.json +81694c9ba20fbf4af596c942c0fa3acf76275248142204e02324681f7c7a982d data/processed/echo_hubert_manifest.parquet +d327447e3004c2b33407d64752e7ae2a6bc17815b1f77a93cae17905a1f5b57e probes/ecg/ecg_only.joblib +3b3d6cabd41dd2634dccacf83553d474685414cb193c38224613c698db645d96 probes/echo/echo_only.pt +3f33b5bf67c6f081d297add31b9486cf219b5709961204eb5b99de2e2faf00e0 probes/concat/concat_mlp.pt +bac18bb84edb93ba8ea8b60a6ad9e6b568b37919a5ac2605da1dc30242163cdb probes/fused/cross_attn_fused.pt +b1f57d9fd598d1bfe16367eab816b90e515db12cc1aafa316ef36a84fce4e4ec results/missing_modality.json diff --git a/docs/results/baseline_gap.json b/docs/results/baseline_gap.json new file mode 100644 index 0000000..cdee4ee --- /dev/null +++ b/docs/results/baseline_gap.json @@ -0,0 +1,151 @@ +{ + "task": "E10_baseline_gap_diagnosis", + "issue": 67, + "seed": 42, + "git_sha": "035a258bf2cded97722d76db09c1a6ff164e7a8a", + "manifest": "data/processed/echo_hubert_manifest.parquet", + "manifest_sha256": "81694c9ba20fbf4af596c942c0fa3acf76275248142204e02324681f7c7a982d", + "checkpoint_sha256": { + "fused": "bac18bb84edb93ba8ea8b60a6ad9e6b568b37919a5ac2605da1dc30242163cdb", + "concat": "3f33b5bf67c6f081d297add31b9486cf219b5709961204eb5b99de2e2faf00e0", + "echo_only": "3b3d6cabd41dd2634dccacf83553d474685414cb193c38224613c698db645d96", + "ecg_only": "d327447e3004c2b33407d64752e7ae2a6bc17815b1f77a93cae17905a1f5b57e" + }, + "probes_dir": "probes", + "device": "cpu", + "n_test": 245, + "n_dropped_nonfinite": 12, + "n_bootstrap": 1000, + "test_metrics": { + "fused": { + "baseline_mae": 11.701, + "mae": 10.4218, + "ef40_auroc": 0.7708, + "mae_ci95": [ + 9.2769, + 11.6448 + ], + "ef40_auroc_ci95": [ + 0.6968, + 0.8386 + ] + }, + "concat": { + "baseline_mae": 11.701, + "mae": 10.9285, + "ef40_auroc": 0.7164, + "mae_ci95": [ + 9.773, + 12.2162 + ], + "ef40_auroc_ci95": [ + 0.634, + 0.7947 + ] + }, + "echo_only": { + "baseline_mae": 11.701, + "mae": 11.0879, + "ef40_auroc": 0.7708, + "mae_ci95": [ + 9.7861, + 12.2475 + ], + "ef40_auroc_ci95": [ + 0.698, + 0.8452 + ] + }, + "ecg_only": { + "baseline_mae": 11.701, + "mae": 11.5965, + "ef40_auroc": 0.6713, + "mae_ci95": [ + 10.3604, + 12.8439 + ], + "ef40_auroc_ci95": [ + 0.5816, + 0.7556 + ] + } + }, + "paired_deltas_vs_fused": { + "fused_vs_echo_only": { + "delta_mae": -0.6661, + "delta_mae_ci95": [ + -1.5517, + 0.2733 + ], + "delta_ef40_auroc": 0.0, + "delta_ef40_auroc_ci95": [ + -0.0749, + 0.0709 + ] + }, + "fused_vs_ecg_only": { + "delta_mae": -1.1747, + "delta_mae_ci95": [ + -1.9865, + -0.3315 + ], + "delta_ef40_auroc": 0.0995, + "delta_ef40_auroc_ci95": [ + 0.0322, + 0.1686 + ] + }, + "fused_vs_concat": { + "delta_mae": -0.5067, + "delta_mae_ci95": [ + -1.0543, + 0.0807 + ], + "delta_ef40_auroc": 0.0545, + "delta_ef40_auroc_ci95": [ + 0.0032, + 0.1061 + ] + } + }, + "fused_mae_by_lvef_source": { + "biplane_lvef": { + "n": 32, + "mae": 8.2302 + }, + "lvef": { + "n": 200, + "mae": 9.9782 + }, + "lvef_3d": { + "n": 7, + "mae": 5.1668 + }, + "lvef_upper": { + "n": 6, + "mae": 43.0265 + } + }, + "fused_metrics_by_split": { + "train": { + "baseline_mae": 11.6691, + "mae": 8.719, + "ef40_auroc": 0.8865 + }, + "val": { + "baseline_mae": 12.5857, + "mae": 10.6228, + "ef40_auroc": 0.8898 + }, + "test": { + "baseline_mae": 11.701, + "mae": 10.4218, + "ef40_auroc": 0.7708 + } + }, + "published_baselines_for_context": { + "echojepa_solo_mae": 5.97, + "ecg_fm_solo_auroc": 0.929, + "note": "Different cohorts and setups; not like-for-like with this manifest." + } +} diff --git a/docs/results/calibration_ecg_dropped.json b/docs/results/calibration_ecg_dropped.json new file mode 100644 index 0000000..a0dd6d6 --- /dev/null +++ b/docs/results/calibration_ecg_dropped.json @@ -0,0 +1,91 @@ +{ + "task": "E04_ef40_calibration", + "git_sha": "035a258bf2cded97722d76db09c1a6ff164e7a8a", + "predictions_path": "results/missing_modality.json", + "prediction_source": { + "checkpoint": "probes/fused/cross_attn_fused.pt", + "checkpoint_sha256": "bac18bb84edb93ba8ea8b60a6ad9e6b568b37919a5ac2605da1dc30242163cdb", + "manifest_sha256": "81694c9ba20fbf4af596c942c0fa3acf76275248142204e02324681f7c7a982d", + "seed": 42, + "git_sha": "035a258bf2cded97722d76db09c1a6ff164e7a8a" + }, + "condition": "ecg_dropped", + "n": 245, + "n_bins": 10, + "scaler_fit_on": "val", + "ece": 0.0151, + "bins": [ + { + "bin_low": 0.0, + "bin_high": 0.1, + "count": 0, + "confidence": null, + "accuracy": null + }, + { + "bin_low": 0.1, + "bin_high": 0.2, + "count": 234, + "confidence": 0.168, + "accuracy": 0.1709 + }, + { + "bin_low": 0.2, + "bin_high": 0.3, + "count": 2, + "confidence": 0.2315, + "accuracy": 1.0 + }, + { + "bin_low": 0.3, + "bin_high": 0.4, + "count": 1, + "confidence": 0.3979, + "accuracy": 0.0 + }, + { + "bin_low": 0.4, + "bin_high": 0.5, + "count": 0, + "confidence": null, + "accuracy": null + }, + { + "bin_low": 0.5, + "bin_high": 0.6, + "count": 0, + "confidence": null, + "accuracy": null + }, + { + "bin_low": 0.6, + "bin_high": 0.7, + "count": 0, + "confidence": null, + "accuracy": null + }, + { + "bin_low": 0.7, + "bin_high": 0.8, + "count": 8, + "confidence": 0.738, + "accuracy": 0.875 + }, + { + "bin_low": 0.8, + "bin_high": 0.9, + "count": 0, + "confidence": null, + "accuracy": null + }, + { + "bin_low": 0.9, + "bin_high": 1.0, + "count": 0, + "confidence": null, + "accuracy": null + } + ], + "figure": "results/calibration/ecg_dropped/reliability.pdf", + "note": "Platt scaler fit on val predictions, applied to test." +} diff --git a/docs/results/calibration_echo_dropped.json b/docs/results/calibration_echo_dropped.json new file mode 100644 index 0000000..261dbde --- /dev/null +++ b/docs/results/calibration_echo_dropped.json @@ -0,0 +1,91 @@ +{ + "task": "E04_ef40_calibration", + "git_sha": "035a258bf2cded97722d76db09c1a6ff164e7a8a", + "predictions_path": "results/missing_modality.json", + "prediction_source": { + "checkpoint": "probes/fused/cross_attn_fused.pt", + "checkpoint_sha256": "bac18bb84edb93ba8ea8b60a6ad9e6b568b37919a5ac2605da1dc30242163cdb", + "manifest_sha256": "81694c9ba20fbf4af596c942c0fa3acf76275248142204e02324681f7c7a982d", + "seed": 42, + "git_sha": "035a258bf2cded97722d76db09c1a6ff164e7a8a" + }, + "condition": "echo_dropped", + "n": 245, + "n_bins": 10, + "scaler_fit_on": "val", + "ece": 0.0427, + "bins": [ + { + "bin_low": 0.0, + "bin_high": 0.1, + "count": 58, + "confidence": 0.0901, + "accuracy": 0.0862 + }, + { + "bin_low": 0.1, + "bin_high": 0.2, + "count": 108, + "confidence": 0.1383, + "accuracy": 0.1389 + }, + { + "bin_low": 0.2, + "bin_high": 0.3, + "count": 34, + "confidence": 0.2403, + "accuracy": 0.3529 + }, + { + "bin_low": 0.3, + "bin_high": 0.4, + "count": 15, + "confidence": 0.3418, + "accuracy": 0.4 + }, + { + "bin_low": 0.4, + "bin_high": 0.5, + "count": 12, + "confidence": 0.4528, + "accuracy": 0.3333 + }, + { + "bin_low": 0.5, + "bin_high": 0.6, + "count": 9, + "confidence": 0.5463, + "accuracy": 0.4444 + }, + { + "bin_low": 0.6, + "bin_high": 0.7, + "count": 7, + "confidence": 0.6561, + "accuracy": 0.4286 + }, + { + "bin_low": 0.7, + "bin_high": 0.8, + "count": 2, + "confidence": 0.7596, + "accuracy": 0.0 + }, + { + "bin_low": 0.8, + "bin_high": 0.9, + "count": 0, + "confidence": null, + "accuracy": null + }, + { + "bin_low": 0.9, + "bin_high": 1.0, + "count": 0, + "confidence": null, + "accuracy": null + } + ], + "figure": "results/calibration/echo_dropped/reliability.pdf", + "note": "Platt scaler fit on val predictions, applied to test." +} diff --git a/docs/results/calibration_full.json b/docs/results/calibration_full.json new file mode 100644 index 0000000..a2af434 --- /dev/null +++ b/docs/results/calibration_full.json @@ -0,0 +1,91 @@ +{ + "task": "E04_ef40_calibration", + "git_sha": "035a258bf2cded97722d76db09c1a6ff164e7a8a", + "predictions_path": "results/missing_modality.json", + "prediction_source": { + "checkpoint": "probes/fused/cross_attn_fused.pt", + "checkpoint_sha256": "bac18bb84edb93ba8ea8b60a6ad9e6b568b37919a5ac2605da1dc30242163cdb", + "manifest_sha256": "81694c9ba20fbf4af596c942c0fa3acf76275248142204e02324681f7c7a982d", + "seed": 42, + "git_sha": "035a258bf2cded97722d76db09c1a6ff164e7a8a" + }, + "condition": "full", + "n": 245, + "n_bins": 10, + "scaler_fit_on": "val", + "ece": 0.0724, + "bins": [ + { + "bin_low": 0.0, + "bin_high": 0.1, + "count": 157, + "confidence": 0.0531, + "accuracy": 0.0955 + }, + { + "bin_low": 0.1, + "bin_high": 0.2, + "count": 30, + "confidence": 0.1402, + "accuracy": 0.2667 + }, + { + "bin_low": 0.2, + "bin_high": 0.3, + "count": 20, + "confidence": 0.2543, + "accuracy": 0.35 + }, + { + "bin_low": 0.3, + "bin_high": 0.4, + "count": 6, + "confidence": 0.3643, + "accuracy": 0.1667 + }, + { + "bin_low": 0.4, + "bin_high": 0.5, + "count": 9, + "confidence": 0.4394, + "accuracy": 0.4444 + }, + { + "bin_low": 0.5, + "bin_high": 0.6, + "count": 3, + "confidence": 0.5733, + "accuracy": 0.3333 + }, + { + "bin_low": 0.6, + "bin_high": 0.7, + "count": 6, + "confidence": 0.6555, + "accuracy": 0.5 + }, + { + "bin_low": 0.7, + "bin_high": 0.8, + "count": 4, + "confidence": 0.756, + "accuracy": 0.25 + }, + { + "bin_low": 0.8, + "bin_high": 0.9, + "count": 5, + "confidence": 0.8467, + "accuracy": 0.8 + }, + { + "bin_low": 0.9, + "bin_high": 1.0, + "count": 5, + "confidence": 0.9556, + "accuracy": 1.0 + } + ], + "figure": "results/calibration/full/reliability.pdf", + "note": "Platt scaler fit on val predictions, applied to test." +} diff --git a/docs/results/failure_report.json b/docs/results/failure_report.json new file mode 100644 index 0000000..0ea4a19 --- /dev/null +++ b/docs/results/failure_report.json @@ -0,0 +1,194 @@ +{ + "task": "modality_failure_analysis", + "modalities": [ + "echo", + "ecg" + ], + "n": 245, + "threshold": 40.0, + "tolerance": 5.0, + "conditions": { + "full": { + "mae": 10.1873, + "taxonomy": { + "correct": 84, + "imprecise": 123, + "critical": 38 + } + }, + "drop_echo": { + "mae": 117.5006, + "taxonomy": { + "correct": 0, + "imprecise": 49, + "critical": 196 + } + }, + "drop_ecg": { + "mae": 16.6486, + "taxonomy": { + "correct": 28, + "imprecise": 139, + "critical": 78 + } + } + }, + "complementarity": { + "full_mae": 10.1873, + "solo_mae": { + "echo": 16.6486, + "ecg": 117.5006 + }, + "marginal_value": { + "echo": 107.3133, + "ecg": 6.4613 + }, + "fusion_gain_vs_best_solo": 6.4613, + "per_example_winners": { + "echo": 245, + "ecg": 0 + }, + "matrix": { + "modalities": [ + "echo", + "ecg" + ], + "values": [ + [ + 16.6486, + 10.1873 + ], + [ + 10.1873, + 117.5006 + ] + ] + }, + "winners_by_group": { + "sex": { + "F": { + "echo": 144, + "ecg": 0 + }, + "M": { + "echo": 101, + "ecg": 0 + } + }, + "age_band": { + "18-39": { + "echo": 12, + "ecg": 0 + }, + "40-54": { + "echo": 37, + "ecg": 0 + }, + "55-64": { + "echo": 57, + "ecg": 0 + }, + "65-74": { + "echo": 57, + "ecg": 0 + }, + "75-89": { + "echo": 77, + "ecg": 0 + }, + "90+": { + "echo": 5, + "ecg": 0 + } + }, + "race": { + "AMERICAN INDIAN/ALASKA NATIVE": { + "echo": 1, + "ecg": 0 + }, + "ASIAN": { + "echo": 1, + "ecg": 0 + }, + "ASIAN - CHINESE": { + "echo": 7, + "ecg": 0 + }, + "BLACK/AFRICAN": { + "echo": 3, + "ecg": 0 + }, + "BLACK/AFRICAN AMERICAN": { + "echo": 36, + "ecg": 0 + }, + "BLACK/CAPE VERDEAN": { + "echo": 2, + "ecg": 0 + }, + "BLACK/CARIBBEAN ISLAND": { + "echo": 2, + "ecg": 0 + }, + "HISPANIC/LATINO - DOMINICAN": { + "echo": 1, + "ecg": 0 + }, + "HISPANIC/LATINO - HONDURAN": { + "echo": 1, + "ecg": 0 + }, + "HISPANIC/LATINO - PUERTO RICAN": { + "echo": 4, + "ecg": 0 + }, + "OTHER": { + "echo": 8, + "ecg": 0 + }, + "PATIENT DECLINED TO ANSWER": { + "echo": 1, + "ecg": 0 + }, + "PORTUGUESE": { + "echo": 1, + "ecg": 0 + }, + "UNKNOWN": { + "echo": 2, + "ecg": 0 + }, + "WHITE": { + "echo": 157, + "ecg": 0 + }, + "WHITE - OTHER EUROPEAN": { + "echo": 6, + "ecg": 0 + }, + "WHITE - RUSSIAN": { + "echo": 12, + "ecg": 0 + } + } + } + }, + "dropout": { + "drop_echo": { + "induced_critical": 188, + "silent": 188, + "loud": 0, + "silent_rate": 1.0, + "mean_output_shift": 116.399, + "mae_increase": 107.3133 + }, + "drop_ecg": { + "induced_critical": 59, + "silent": 16, + "loud": 43, + "silent_rate": 0.2712, + "mean_output_shift": 14.2229, + "mae_increase": 6.4613 + } + } +} diff --git a/docs/results/fairness_metrics.json b/docs/results/fairness_metrics.json new file mode 100644 index 0000000..6d6d9bc --- /dev/null +++ b/docs/results/fairness_metrics.json @@ -0,0 +1,334 @@ +{ + "task": "E03_fairness_stratification", + "checkpoint": "probes/fused/cross_attn_fused.pt", + "checkpoint_sha256": "bac18bb84edb93ba8ea8b60a6ad9e6b568b37919a5ac2605da1dc30242163cdb", + "manifest_sha256": "81694c9ba20fbf4af596c942c0fa3acf76275248142204e02324681f7c7a982d", + "git_sha": "035a258bf2cded97722d76db09c1a6ff164e7a8a", + "n_test": 245, + "n_dropped_nonfinite_all_splits": 12, + "bias_note": "MIMIC-IV records administrative gender as 'sex' and admission-reported race; both carry known curation bias (see MIMIC-IV documentation), so strata inherit it.", + "conditions": { + "full": { + "overall": { + "n": 245, + "mae": 10.4218, + "baseline_mae": 11.701, + "ef40_auroc": 0.7708246563931695, + "small_n": false + }, + "by": { + "sex": { + "F": { + "n": 144, + "mae": 10.1294, + "baseline_mae": 10.5965, + "ef40_auroc": 0.7697466467958272, + "small_n": false + }, + "M": { + "n": 101, + "mae": 10.8387, + "baseline_mae": 13.3897, + "ef40_auroc": 0.7632632632632632, + "small_n": false + } + }, + "age_band": { + "18-39": { + "n": 12, + "mae": 8.9737, + "baseline_mae": 8.625, + "ef40_auroc": 0.7777777777777778, + "small_n": true + }, + "40-54": { + "n": 37, + "mae": 8.9349, + "baseline_mae": 8.6779, + "ef40_auroc": 0.7424242424242424, + "small_n": false + }, + "55-64": { + "n": 57, + "mae": 9.4325, + "baseline_mae": 9.7938, + "ef40_auroc": 0.726923076923077, + "small_n": false + }, + "65-74": { + "n": 57, + "mae": 13.0964, + "baseline_mae": 15.8319, + "ef40_auroc": 0.7523809523809524, + "small_n": false + }, + "75-89": { + "n": 77, + "mae": 10.2899, + "baseline_mae": 12.2098, + "ef40_auroc": 0.7891156462585034, + "small_n": false + }, + "90+": { + "n": 5, + "mae": 7.7201, + "baseline_mae": 7.84, + "ef40_auroc": 0.5, + "small_n": true + } + }, + "race": { + "ASIAN": { + "n": 8, + "mae": 10.4771, + "baseline_mae": 9.375, + "ef40_auroc": 0.0, + "small_n": true + }, + "BLACK": { + "n": 43, + "mae": 6.6185, + "baseline_mae": 8.8156, + "ef40_auroc": 0.9729729729729729, + "small_n": false + }, + "HISPANIC/LATINO": { + "n": 6, + "mae": 7.9728, + "baseline_mae": 4.6667, + "ef40_auroc": null, + "small_n": true + }, + "OTHER/UNKNOWN": { + "n": 13, + "mae": 9.2042, + "baseline_mae": 6.426, + "ef40_auroc": 0.27272727272727276, + "small_n": true + }, + "WHITE": { + "n": 175, + "mae": 11.5282, + "baseline_mae": 13.0269, + "ef40_auroc": 0.7609259259259259, + "small_n": false + } + } + } + }, + "echo_dropped": { + "overall": { + "n": 245, + "mae": 20.5507, + "baseline_mae": 11.701, + "ef40_auroc": 0.6889837567680133, + "small_n": false + }, + "by": { + "sex": { + "F": { + "n": 144, + "mae": 21.1948, + "baseline_mae": 10.5965, + "ef40_auroc": 0.6747391952309986, + "small_n": false + }, + "M": { + "n": 101, + "mae": 19.6322, + "baseline_mae": 13.3897, + "ef40_auroc": 0.6721721721721722, + "small_n": false + } + }, + "age_band": { + "18-39": { + "n": 12, + "mae": 16.3999, + "baseline_mae": 8.625, + "ef40_auroc": 0.40740740740740744, + "small_n": true + }, + "40-54": { + "n": 37, + "mae": 20.4442, + "baseline_mae": 8.6779, + "ef40_auroc": 0.6969696969696969, + "small_n": false + }, + "55-64": { + "n": 57, + "mae": 23.6714, + "baseline_mae": 9.7938, + "ef40_auroc": 0.7615384615384615, + "small_n": false + }, + "65-74": { + "n": 57, + "mae": 21.7085, + "baseline_mae": 15.8319, + "ef40_auroc": 0.626984126984127, + "small_n": false + }, + "75-89": { + "n": 77, + "mae": 18.2873, + "baseline_mae": 12.2098, + "ef40_auroc": 0.7508503401360545, + "small_n": false + }, + "90+": { + "n": 5, + "mae": 17.3801, + "baseline_mae": 7.84, + "ef40_auroc": 0.25, + "small_n": true + } + }, + "race": { + "ASIAN": { + "n": 8, + "mae": 21.7366, + "baseline_mae": 9.375, + "ef40_auroc": 0.2857142857142857, + "small_n": true + }, + "BLACK": { + "n": 43, + "mae": 21.4057, + "baseline_mae": 8.8156, + "ef40_auroc": 0.8063063063063063, + "small_n": false + }, + "HISPANIC/LATINO": { + "n": 6, + "mae": 26.8224, + "baseline_mae": 4.6667, + "ef40_auroc": null, + "small_n": true + }, + "OTHER/UNKNOWN": { + "n": 13, + "mae": 19.7079, + "baseline_mae": 6.426, + "ef40_auroc": 0.3181818181818182, + "small_n": true + }, + "WHITE": { + "n": 175, + "mae": 20.1339, + "baseline_mae": 13.0269, + "ef40_auroc": 0.6859259259259259, + "small_n": false + } + } + } + } + }, + "overall": { + "n": 245, + "mae": 10.4218, + "baseline_mae": 11.701, + "ef40_auroc": 0.7708246563931695, + "small_n": false + }, + "by": { + "sex": { + "F": { + "n": 144, + "mae": 10.1294, + "baseline_mae": 10.5965, + "ef40_auroc": 0.7697466467958272, + "small_n": false + }, + "M": { + "n": 101, + "mae": 10.8387, + "baseline_mae": 13.3897, + "ef40_auroc": 0.7632632632632632, + "small_n": false + } + }, + "age_band": { + "18-39": { + "n": 12, + "mae": 8.9737, + "baseline_mae": 8.625, + "ef40_auroc": 0.7777777777777778, + "small_n": true + }, + "40-54": { + "n": 37, + "mae": 8.9349, + "baseline_mae": 8.6779, + "ef40_auroc": 0.7424242424242424, + "small_n": false + }, + "55-64": { + "n": 57, + "mae": 9.4325, + "baseline_mae": 9.7938, + "ef40_auroc": 0.726923076923077, + "small_n": false + }, + "65-74": { + "n": 57, + "mae": 13.0964, + "baseline_mae": 15.8319, + "ef40_auroc": 0.7523809523809524, + "small_n": false + }, + "75-89": { + "n": 77, + "mae": 10.2899, + "baseline_mae": 12.2098, + "ef40_auroc": 0.7891156462585034, + "small_n": false + }, + "90+": { + "n": 5, + "mae": 7.7201, + "baseline_mae": 7.84, + "ef40_auroc": 0.5, + "small_n": true + } + }, + "race": { + "ASIAN": { + "n": 8, + "mae": 10.4771, + "baseline_mae": 9.375, + "ef40_auroc": 0.0, + "small_n": true + }, + "BLACK": { + "n": 43, + "mae": 6.6185, + "baseline_mae": 8.8156, + "ef40_auroc": 0.9729729729729729, + "small_n": false + }, + "HISPANIC/LATINO": { + "n": 6, + "mae": 7.9728, + "baseline_mae": 4.6667, + "ef40_auroc": null, + "small_n": true + }, + "OTHER/UNKNOWN": { + "n": 13, + "mae": 9.2042, + "baseline_mae": 6.426, + "ef40_auroc": 0.27272727272727276, + "small_n": true + }, + "WHITE": { + "n": 175, + "mae": 11.5282, + "baseline_mae": 13.0269, + "ef40_auroc": 0.7609259259259259, + "small_n": false + } + } + } +} diff --git a/docs/results/missing_modality.metrics.json b/docs/results/missing_modality.metrics.json new file mode 100644 index 0000000..78fbbe3 --- /dev/null +++ b/docs/results/missing_modality.metrics.json @@ -0,0 +1,96 @@ +{ + "task": "E01_missing_modality_evaluation", + "issue": 7, + "source_task": "M09_cross_attn_fused_probe", + "seed": 42, + "git_sha": "035a258bf2cded97722d76db09c1a6ff164e7a8a", + "checkpoint": "probes/fused/cross_attn_fused.pt", + "checkpoint_sha256": "bac18bb84edb93ba8ea8b60a6ad9e6b568b37919a5ac2605da1dc30242163cdb", + "manifest_sha256": "81694c9ba20fbf4af596c942c0fa3acf76275248142204e02324681f7c7a982d", + "config": { + "cohort_path": "data/processed/echo_hubert_manifest.parquet", + "echo_embedding_path": "None", + "ecg_embedding_path": "None", + "fusion_dim": 256, + "echo_dim": 1024, + "ecg_dim": 768, + "hidden": 256, + "batch_size": 64, + "n_bootstrap": 1000, + "device": "cpu" + }, + "n": { + "val": 130, + "test": 245 + }, + "test": { + "full": { + "baseline_mae": 11.701, + "mae": 10.4218, + "ef40_auroc": 0.7708 + }, + "echo_dropped": { + "baseline_mae": 11.701, + "mae": 20.5507, + "ef40_auroc": 0.689 + }, + "ecg_dropped": { + "baseline_mae": 11.701, + "mae": 11.2167, + "ef40_auroc": 0.3827 + } + }, + "bootstrap": { + "full": { + "mae_ci_low": 9.277, + "mae_ci_high": 11.6448, + "ef40_auroc_ci_low": 0.6968, + "ef40_auroc_ci_high": 0.8386 + }, + "echo_dropped": { + "mae_ci_low": 19.0182, + "mae_ci_high": 22.0169, + "ef40_auroc_ci_low": 0.605, + "ef40_auroc_ci_high": 0.7654 + }, + "ecg_dropped": { + "mae_ci_low": 9.9998, + "mae_ci_high": 12.3924, + "ef40_auroc_ci_low": 0.2793, + "ef40_auroc_ci_high": 0.4934 + } + }, + "metrics_table": [ + { + "condition": "full", + "mae": 10.4218, + "ef40_auroc": 0.7708, + "baseline_mae": 11.701, + "mae_ci_low": 9.277, + "mae_ci_high": 11.6448, + "ef40_auroc_ci_low": 0.6968, + "ef40_auroc_ci_high": 0.8386 + }, + { + "condition": "echo_dropped", + "mae": 20.5507, + "ef40_auroc": 0.689, + "baseline_mae": 11.701, + "mae_ci_low": 19.0182, + "mae_ci_high": 22.0169, + "ef40_auroc_ci_low": 0.605, + "ef40_auroc_ci_high": 0.7654 + }, + { + "condition": "ecg_dropped", + "mae": 11.2167, + "ef40_auroc": 0.3827, + "baseline_mae": 11.701, + "mae_ci_low": 9.9998, + "mae_ci_high": 12.3924, + "ef40_auroc_ci_low": 0.2793, + "ef40_auroc_ci_high": 0.4934 + } + ], + "notes": "Dropped conditions mask one branch at inference only; no unimodal retraining." +} diff --git a/scripts/diagnose_baseline_gap.py b/scripts/diagnose_baseline_gap.py new file mode 100644 index 0000000..17956da --- /dev/null +++ b/scripts/diagnose_baseline_gap.py @@ -0,0 +1,210 @@ +#!/usr/bin/env python3 +"""E10 (#67): in-cohort fused vs solo comparison on one shared test frame. + +The published baselines (EchoJEPA 5.97 MAE, ECG-FM 0.929 AUROC) come from different +cohorts, so they cannot say whether fusion itself is broken. This scores the four +trained probes on the identical 245-row test frame (non-finite rows dropped the same +way the fused eval drops them), with bootstrap CIs and paired deltas, plus an LVEF +label-source error breakdown and a fused overfit check. + +Example: + python scripts/diagnose_baseline_gap.py \ + --manifest data/processed/echo_hubert_manifest.parquet \ + --probes-dir probes --out results/baseline_gap.json +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import joblib +import numpy as np +import torch + +from primed_ai.probes import manifest as manifest_io +from primed_ai.probes.common import auroc, git_sha, regression_metrics, save_results, sha256_file +from primed_ai.probes.concat_mlp import ConcatMLPProbe +from primed_ai.probes.concat_mlp import _predict as predict_concat +from primed_ai.probes.cross_attn import ( + CrossAttnFusedProbe, + _condition_arrays, + fused_probe_loader, + prepare_fused_probe_data, +) +from primed_ai.probes.echo_only import EchoOnlyProbe +from primed_ai.probes.echo_only import _predict as predict_echo + + +def _ci95(a) -> list: + return [round(float(np.quantile(a, 0.025)), 4), round(float(np.quantile(a, 0.975)), 4)] + + +def _metrics(arrays: dict) -> dict: + m = regression_metrics(arrays["lvef"], arrays["prediction"]) + m["ef40_auroc"] = round(auroc(arrays["ef_le_40"], -np.asarray(arrays["prediction"])), 4) + return m + + +def _bootstrap(arrays: dict, *, n_bootstrap: int, seed: int) -> dict: + y = np.asarray(arrays["lvef"], dtype=np.float64) + p = np.asarray(arrays["prediction"], dtype=np.float64) + ef = np.asarray(arrays["ef_le_40"], dtype=bool) + rng = np.random.default_rng(seed) + maes, aucs = [], [] + for _ in range(n_bootstrap): + i = rng.integers(0, len(y), size=len(y)) + maes.append(np.abs(y[i] - p[i]).mean()) + a = auroc(ef[i], -p[i]) + if np.isfinite(a): + aucs.append(a) + return {"mae_ci95": _ci95(maes), "ef40_auroc_ci95": _ci95(aucs)} + + +def _paired_delta(fused: dict, other: dict, *, n_bootstrap: int, seed: int) -> dict: + """Paired bootstrap of (fused - other) on the shared rows; negative MAE delta = fusion wins.""" + y = np.asarray(fused["lvef"], dtype=np.float64) + pf = np.asarray(fused["prediction"], dtype=np.float64) + po = np.asarray(other["prediction"], dtype=np.float64) + ef = np.asarray(fused["ef_le_40"], dtype=bool) + rng = np.random.default_rng(seed) + d_mae, d_auc = [], [] + for _ in range(n_bootstrap): + i = rng.integers(0, len(y), size=len(y)) + d_mae.append(np.abs(y[i] - pf[i]).mean() - np.abs(y[i] - po[i]).mean()) + af, ao = auroc(ef[i], -pf[i]), auroc(ef[i], -po[i]) + if np.isfinite(af) and np.isfinite(ao): + d_auc.append(af - ao) + return { + "delta_mae": round(float(np.abs(y - pf).mean() - np.abs(y - po).mean()), 4), + "delta_mae_ci95": _ci95(d_mae), + "delta_ef40_auroc": round(auroc(ef, -pf) - auroc(ef, -po), 4), + "delta_ef40_auroc_ci95": _ci95(d_auc), + } + + +def main() -> None: + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--manifest", default="data/processed/echo_hubert_manifest.parquet") + ap.add_argument("--probes-dir", default="probes", help="Root holding ecg/echo/concat/fused") + ap.add_argument("--fusion-dim", type=int, default=256) + ap.add_argument("--batch-size", type=int, default=64) + ap.add_argument("--n-bootstrap", type=int, default=1000) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--out", default="results/baseline_gap.json") + args = ap.parse_args() + + device = "cpu" + probes = Path(args.probes_dir) + parts, n_dropped = prepare_fused_probe_data(args.manifest) + echo_dim, ecg_dim = manifest_io.dims(parts["test"]) + + loaders = { + s: fused_probe_loader(parts[s], embed_dim=echo_dim, batch_size=args.batch_size) + for s in ("train", "val", "test") + } + + fused = CrossAttnFusedProbe(args.fusion_dim, echo_dim=echo_dim, ecg_dim=ecg_dim).to(device) + fused.load_state_dict(torch.load(probes / "fused" / "cross_attn_fused.pt", map_location=device)) + fused.eval() + + concat = ConcatMLPProbe(echo_dim, ecg_dim).to(device) + concat.load_state_dict(torch.load(probes / "concat" / "concat_mlp.pt", map_location=device)) + concat.eval() + + echo = EchoOnlyProbe(echo_dim).to(device) + echo.load_state_dict(torch.load(probes / "echo" / "echo_only.pt", map_location=device)) + echo.eval() + + ecg_bundle = joblib.load(probes / "ecg" / "ecg_only.joblib") + X = np.vstack([np.asarray(v, dtype=np.float64) for v in parts["test"]["ecg_embedding"]]) + ecg_pred = ecg_bundle["ridge"].predict(ecg_bundle["scaler"].transform(X)) + + test = parts["test"] + arrays = { + "fused": _condition_arrays(fused, loaders["test"], device, "full"), + "concat": predict_concat(concat, loaders["test"], device), + "echo_only": predict_echo(echo, loaders["test"], device), + "ecg_only": { + "lvef": test["lvef"].to_numpy(np.float64), + "prediction": ecg_pred, + "ef_le_40": test["ef_le_40"].astype(bool).to_numpy(), + }, + } + + models = {} + for i, (name, arr) in enumerate(arrays.items()): + models[name] = { + **_metrics(arr), + **_bootstrap(arr, n_bootstrap=args.n_bootstrap, seed=args.seed + i), + } + + deltas = { + f"fused_vs_{name}": _paired_delta( + arrays["fused"], arrays[name], n_bootstrap=args.n_bootstrap, seed=args.seed + 10 + i + ) + for i, name in enumerate(("echo_only", "ecg_only", "concat")) + } + + # Label-noise check: fused error stratified by which LVEF field supplied the label + by_source = {} + if "lvef_measurement" in test.columns: + err = np.abs( + np.asarray(arrays["fused"]["lvef"]) - np.asarray(arrays["fused"]["prediction"]) + ) + for src in sorted(test["lvef_measurement"].dropna().unique()): + mask = (test["lvef_measurement"] == src).to_numpy() + by_source[str(src)] = {"n": int(mask.sum()), "mae": round(float(err[mask].mean()), 4)} + + # Overfit check: the same fused checkpoint across splits + fused_splits = { + s: _metrics(_condition_arrays(fused, loaders[s], device, "full")) + for s in ("train", "val", "test") + } + + payload = { + "task": "E10_baseline_gap_diagnosis", + "issue": 67, + "seed": args.seed, + "git_sha": git_sha(), + "manifest": str(args.manifest), + "manifest_sha256": sha256_file(args.manifest), + "checkpoint_sha256": { + "fused": sha256_file(probes / "fused" / "cross_attn_fused.pt"), + "concat": sha256_file(probes / "concat" / "concat_mlp.pt"), + "echo_only": sha256_file(probes / "echo" / "echo_only.pt"), + "ecg_only": sha256_file(probes / "ecg" / "ecg_only.joblib"), + }, + "probes_dir": str(probes), + "device": device, + "n_test": len(test), + "n_dropped_nonfinite": n_dropped, + "n_bootstrap": args.n_bootstrap, + "test_metrics": models, + "paired_deltas_vs_fused": deltas, + "fused_mae_by_lvef_source": by_source, + "fused_metrics_by_split": fused_splits, + "published_baselines_for_context": { + "echojepa_solo_mae": 5.97, + "ecg_fm_solo_auroc": 0.929, + "note": "Different cohorts and setups; not like-for-like with this manifest.", + }, + } + save_results(Path(args.out), payload) + + print(f"n_test={len(test)} (dropped {n_dropped} non-finite)") + for name, m in models.items(): + print( + f"{name:10s} MAE {m['mae']:6.2f} {m['mae_ci95']} AUROC {m['ef40_auroc']:.3f} {m['ef40_auroc_ci95']}" + ) + for name, d in deltas.items(): + print( + f"{name:20s} dMAE {d['delta_mae']:+.2f} {d['delta_mae_ci95']} dAUROC {d['delta_ef40_auroc']:+.3f} {d['delta_ef40_auroc_ci95']}" + ) + print(f"Wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/scripts/evaluate_fairness.py b/scripts/evaluate_fairness.py index f86ddab..e0029b4 100644 --- a/scripts/evaluate_fairness.py +++ b/scripts/evaluate_fairness.py @@ -8,7 +8,7 @@ --cohort cohort/paired.parquet \ --echo-embeddings embeddings/echo/pairs.parquet \ --ecg-embeddings embeddings/ecg/pairs.parquet \ - --checkpoint probes/cross_attn_fused/cross_attn_fused.pt \ + --checkpoint probes/fused/cross_attn_fused.pt \ --out results/fairness """ @@ -21,11 +21,18 @@ def main(): p = argparse.ArgumentParser(description="E03: Post-hoc fairness stratification (wrapper).") - p.add_argument("--cohort", required=True, help="Cohort path (with demographics)") - p.add_argument("--echo-embeddings", required=True, help="Echo embedding path") - p.add_argument("--ecg-embeddings", required=True, help="ECG embedding path") + p.add_argument("--cohort", help="Cohort path (with demographics)") + p.add_argument("--manifest", help="Joined manifest path; supersedes the three-table flags") + p.add_argument("--echo-embeddings", help="Echo embedding path") + p.add_argument("--ecg-embeddings", help="ECG embedding path") p.add_argument("--checkpoint", required=True, help="M09 checkpoint path") p.add_argument("--out", default="results/fairness", help="Output directory") + p.add_argument( + "--conditions", + nargs="+", + default=["full", "echo_dropped"], + help="Missing-modality conditions to stratify under", + ) # All three must match the checkpoint being loaded, or the state_dict load fails on shape. p.add_argument( "--embed-dim", type=int, default=256, help="Fusion dimension (matches the checkpoint)." @@ -36,10 +43,17 @@ def main(): p.add_argument("--device", default=None, help="Device (cuda or cpu)") args = p.parse_args() + if args.manifest: + cohort, echo, ecg = args.manifest, None, None + elif args.cohort and args.echo_embeddings and args.ecg_embeddings: + cohort, echo, ecg = args.cohort, args.echo_embeddings, args.ecg_embeddings + else: + p.error("pass --manifest, or all of --cohort/--echo-embeddings/--ecg-embeddings") + run_fairness( - args.cohort, - args.echo_embeddings, - args.ecg_embeddings, + cohort, + echo, + ecg, args.checkpoint, out_dir=args.out, embed_dim=args.embed_dim, @@ -47,6 +61,7 @@ def main(): ecg_dim=args.ecg_dim, batch_size=args.batch_size, device=args.device, + conditions=tuple(args.conditions), ) print(f"Wrote fairness outputs to: {args.out}") diff --git a/scripts/evaluate_missing_modality.py b/scripts/evaluate_missing_modality.py index 80d07df..31d6154 100644 --- a/scripts/evaluate_missing_modality.py +++ b/scripts/evaluate_missing_modality.py @@ -2,7 +2,7 @@ Example: python scripts/evaluate_missing_modality.py \ - --checkpoint probes/cross_attn_fused/cross_attn_fused.pt \ + --checkpoint probes/fused/cross_attn_fused.pt \ --cohort data/raw/cohort/paired_with_splits.parquet \ --echo-embeddings data/interim/echo_study_embeddings_vjepa2.1-vitl-mimic-pt-100.parquet \ --ecg-embeddings data/interim/hubert_ecg_embeddings.parquet \ @@ -33,7 +33,7 @@ def main() -> None: default="data/interim/echo_study_embeddings_vjepa2.1-vitl-mimic-pt-100.parquet", ) parser.add_argument("--ecg-embeddings", default="data/interim/hubert_ecg_embeddings.parquet") - parser.add_argument("--checkpoint", default="probes/cross_attn_fused/cross_attn_fused.pt") + parser.add_argument("--checkpoint", default="probes/fused/cross_attn_fused.pt") parser.add_argument("--output", "--out", dest="output", default="results/missing_modality.json") # All three must match the checkpoint being loaded, or the state_dict load fails on shape. parser.add_argument("--embed-dim", type=int, default=256, help="Shared fusion dimension.") diff --git a/scripts/export_result_bundle.py b/scripts/export_result_bundle.py new file mode 100644 index 0000000..52f8741 --- /dev/null +++ b/scripts/export_result_bundle.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python3 +"""Export a sanitized, checksummed copy of the aggregate result artifacts. + +`results/` and `probes/` are gitignored (per-example predictions and checkpoints stay +local), so reviewers had nothing committed to diff reported numbers against. This copies +the aggregate-only artifacts into docs/results/ — stripping the per-example prediction +blocks from the missing-modality JSON — and writes a SHA256SUMS covering both the bundle +and the local raw artifacts (manifest, checkpoints, full results), so an independent +reproduction can be checked file-by-file. + +Example: + python scripts/export_result_bundle.py \ + --manifest data/processed/echo_hubert_manifest.parquet \ + --results results --probes probes --out docs/results +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +from primed_ai.probes.common import sha256_file + +PER_EXAMPLE_KEYS = ("predictions", "predictions_val") + + +def main() -> None: + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--manifest", default="data/processed/echo_hubert_manifest.parquet") + ap.add_argument("--results", default="results") + ap.add_argument("--probes", default="probes") + ap.add_argument("--out", default="docs/results") + args = ap.parse_args() + + results = Path(args.results) + out = Path(args.out) + out.mkdir(parents=True, exist_ok=True) + + exported: list[Path] = [] + + def export_json(src: Path, dst_name: str, strip: tuple[str, ...] = ()) -> None: + if not src.is_file(): + print(f"skip (missing): {src}") + return + payload = json.loads(src.read_text()) + for key in strip: + payload.pop(key, None) + dst = out / dst_name + dst.write_text(json.dumps(payload, indent=2) + "\n") + exported.append(dst) + + export_json( + results / "missing_modality.json", "missing_modality.metrics.json", PER_EXAMPLE_KEYS + ) + export_json(results / "fairness" / "fairness_metrics.json", "fairness_metrics.json") + for cond in ("full", "echo_dropped", "ecg_dropped"): + export_json(results / "calibration" / cond / "calibration.json", f"calibration_{cond}.json") + export_json(results / "baseline_gap.json", "baseline_gap.json") + export_json(results / "failure" / "failure_report.json", "failure_report.json") + + # checksums: the bundle itself, plus the local raw artifacts a reproduction must match. + # Only the canonical M10 checkpoint locations — stale dirs from older layouts (e.g. + # probes/cross_attn_fused/) must not end up looking canonical in the sums. + probes = Path(args.probes) + raw = [Path(args.manifest)] + raw += [ + probes / "ecg" / "ecg_only.joblib", + probes / "echo" / "echo_only.pt", + probes / "concat" / "concat_mlp.pt", + probes / "fused" / "cross_attn_fused.pt", + ] + raw.append(results / "missing_modality.json") + lines = [] + for path in exported + [p for p in raw if p.is_file()]: + lines.append(f"{sha256_file(path)} {path.as_posix()}") + (out / "SHA256SUMS").write_text("\n".join(lines) + "\n") + print(f"exported {len(exported)} sanitized artifacts + SHA256SUMS to {out}") + + +if __name__ == "__main__": + main() diff --git a/scripts/run_cohort_sensitivity.py b/scripts/run_cohort_sensitivity.py new file mode 100644 index 0000000..2209ca4 --- /dev/null +++ b/scripts/run_cohort_sensitivity.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python3 +"""Run D07 cohort-size sensitivity funnels without downloading cohorts. + +This wraps ``build_cohort.py``'s existing SQL so the 24h / 48h / admission-mode +comparison is one command instead of several ad hoc runs. + +Example: + python scripts/run_cohort_sensitivity.py --project "$GCP_PROJECT_ID" +""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path +from typing import Mapping + +import pandas as pd +from build_cohort import build_cte_sql, get_client, run_funnel + + +def summarize_sensitivity(funnels: Mapping[str, pd.DataFrame]) -> pd.DataFrame: + """Summarize each scenario by its final funnel row.""" + rows = [] + baseline = None + for scenario, funnel in funnels.items(): + if funnel.empty: + raise ValueError(f"{scenario} funnel is empty") + first = funnel.iloc[0] + final = funnel.iloc[-1] + final_rows = int(final["n_studies"]) + if baseline is None: + baseline = final_rows + gain = final_rows - baseline + rows.append( + { + "scenario": scenario, + "final_stage": final["stage"], + "start_rows": int(first["n_studies"]), + "final_rows": final_rows, + "final_subjects": int(final["n_subjects"]), + "total_excluded": int(first["n_studies"] - final["n_studies"]), + "last_step_excluded": int(final.get("excluded_studies", 0)), + "gain_vs_first": gain, + "gain_vs_first_pct": round(gain / baseline * 100, 2) if baseline else 0.0, + } + ) + return pd.DataFrame(rows) + + +def _scenario_params(windows: list[float], include_admission: bool) -> list[dict]: + scenarios = [ + { + "name": f"window_{hours:g}h", + "pair_by": "window", + "before": hours, + "after": hours, + "require_admission": True, + } + for hours in windows + ] + if include_admission: + scenarios.append( + { + "name": "admission", + "pair_by": "admission", + "before": windows[0], + "after": windows[0], + "require_admission": True, + } + ) + return scenarios + + +def run_sensitivity( + *, + project: str | None, + windows: list[float], + output_dir: Path, + include_admission: bool = True, +) -> pd.DataFrame: + client = get_client(project) + output_dir.mkdir(parents=True, exist_ok=True) + funnels = {} + + for spec in _scenario_params(windows, include_admission): + cte = build_cte_sql( + spec["before"], + spec["after"], + 0.0, + 100.0, + spec["require_admission"], + spec["pair_by"], + ) + funnel = run_funnel(client, cte, spec["pair_by"], spec["require_admission"]) + funnels[spec["name"]] = funnel + funnel.to_csv(output_dir / f"{spec['name']}_funnel.csv", index=False) + funnel.to_json(output_dir / f"{spec['name']}_funnel.json", orient="records", indent=2) + + summary = summarize_sensitivity(funnels) + summary.to_csv(output_dir / "summary.csv", index=False) + (output_dir / "summary.json").write_text( + json.dumps(summary.to_dict(orient="records"), indent=2) + "\n" + ) + return summary + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--project", default=os.getenv("GCP_PROJECT_ID")) + parser.add_argument("--window-hours", type=float, nargs="+", default=[24.0, 48.0]) + parser.add_argument("--output-dir", type=Path, default=Path("logs/cohort_sensitivity")) + parser.add_argument("--skip-admission", action="store_true") + args = parser.parse_args() + + summary = run_sensitivity( + project=args.project, + windows=args.window_hours, + output_dir=args.output_dir, + include_admission=not args.skip_admission, + ) + print(summary.to_string(index=False)) + print(f"\nWrote sensitivity outputs to: {args.output_dir}") + + +if __name__ == "__main__": + main() diff --git a/scripts/run_failure_analysis.py b/scripts/run_failure_analysis.py index 2cd5625..84cbe71 100644 --- a/scripts/run_failure_analysis.py +++ b/scripts/run_failure_analysis.py @@ -9,6 +9,11 @@ --echo-embeddings data/interim/echo_study_embeddings_vjepa2.1-vitl-mimic-pt-100.parquet \ --ecg-embeddings data/interim/hubert_ecg_embeddings.parquet \ --out results/failure + +Joined manifest (same table the probes train from; trains on `split == "train"`, +reports on `split == "test"`): + python scripts/run_failure_analysis.py \ + --manifest data/processed/echo_hubert_manifest.parquet --out results/failure """ from __future__ import annotations @@ -65,6 +70,43 @@ def _run_demo(seed: int, out: str) -> None: _emit(report, out) +def _run_manifest(manifest_path: str, out: str) -> None: + from primed_ai.probes import manifest as manifest_io + from primed_ai.probes.common import drop_non_finite + + df = manifest_io.load(manifest_path) + if manifest_io.is_clip_level(df): + raise SystemExit( + "clip-level manifest: the ridge harness works on study vectors — " + "rebuild the manifest pooled, or mean-pool the clips first" + ) + df, n_dropped = drop_non_finite(df, ("echo_embedding", "ecg_embedding")) + if n_dropped: + print(f"dropped {n_dropped} rows with non-finite values") + embeddings = { + "echo": np.vstack(df["echo_embedding"].to_numpy()), + "ecg": np.vstack(df["ecg_embedding"].to_numpy()), + } + lvef = df["lvef"].to_numpy(float) + ef = df["ef_le_40"].to_numpy(bool) + train = (df["split"] == "train").to_numpy() + test = (df["split"] == "test").to_numpy() # strictly test; val stays out + if not train.any() or not test.any(): + raise SystemExit("manifest needs non-empty 'train' and 'test' splits") + groups = {a: df[a].to_numpy() for a in ("sex", "age_band", "race") if a in df} + predict_full = masked_ridge_predict_fn(embeddings, lvef, train) + emb_te = {m: embeddings[m][test] for m in embeddings} + g_te = {a: v[test] for a, v in groups.items()} + report = analyze_modality_failure( + emb_te, + lvef[test], + ef[test], + lambda present: predict_full(present)[test], + groups=g_te or None, + ) + _emit(report, out) + + def _run_real(cohort: str, echo: str, ecg: str, out: str) -> None: import pandas as pd @@ -101,12 +143,15 @@ def main() -> None: ) p.add_argument("--demo", action="store_true", help="run on planted synthetic data") p.add_argument("--seed", type=int, default=0) + p.add_argument("--manifest", help="joined manifest; supersedes the three-table flags") p.add_argument("--cohort") p.add_argument("--echo-embeddings") p.add_argument("--ecg-embeddings") p.add_argument("--out", default="results/failure") a = p.parse_args() - if a.demo or not (a.cohort and a.echo_embeddings and a.ecg_embeddings): + if a.manifest and not a.demo: + _run_manifest(a.manifest, a.out) + elif a.demo or not (a.cohort and a.echo_embeddings and a.ecg_embeddings): _run_demo(a.seed, a.out if a.demo else "results/failure_demo") else: _run_real(a.cohort, a.echo_embeddings, a.ecg_embeddings, a.out) diff --git a/scripts/train_probes.py b/scripts/train_probes.py index 5c1103b..2a5baed 100644 --- a/scripts/train_probes.py +++ b/scripts/train_probes.py @@ -23,6 +23,7 @@ from primed_ai.probes import manifest as manifest_io from primed_ai.probes import run_concat_mlp, run_cross_attn, run_ecg_only, run_echo_only +from primed_ai.probes.common import sha256_file from primed_ai.utils.run_manifest import save_run_metadata PROBES = ("ecg", "echo", "concat", "fused") @@ -105,12 +106,22 @@ def main() -> None: (out_root / "summary.json").write_text(json.dumps(summary, indent=2, default=str)) + checkpoint_files = { + "ecg": "ecg_only.joblib", + "echo": "echo_only.pt", + "concat": "concat_mlp.pt", + "fused": "cross_attn_fused.pt", + } echo_dim, ecg_dim = manifest_io.dims(manifest_io.load(args.manifest)) save_run_metadata( out_root, { "task": "M10_probe_training", "manifest": str(Path(args.manifest).resolve()), + "manifest_sha256": sha256_file(args.manifest), + "checkpoint_sha256": { + name: sha256_file(out_root / name / checkpoint_files[name]) for name in selected + }, "probes": list(selected), "seed": args.seed, "epochs": args.epochs, diff --git a/src/primed_ai/evaluation/calibration.py b/src/primed_ai/evaluation/calibration.py index 7638012..303cdab 100644 --- a/src/primed_ai/evaluation/calibration.py +++ b/src/primed_ai/evaluation/calibration.py @@ -15,19 +15,30 @@ import numpy as np -def platt_probabilities(scores: np.ndarray, labels: np.ndarray) -> np.ndarray: +def platt_probabilities( + scores: np.ndarray, + labels: np.ndarray, + *, + fit_scores: np.ndarray | None = None, + fit_labels: np.ndarray | None = None, +) -> np.ndarray: """Map arbitrary real-valued scores to calibrated probabilities via 1-D logistic fit. For the EF<=40% gate the natural score is the (negated) predicted LVEF; Platt scaling - learns the sign and slope. Note: an in-sample fit gives an optimistic ECE; a held-out - calibration split is preferred when data permit. + learns the sign and slope. Pass ``fit_scores``/``fit_labels`` from a held-out split to + fit there and apply to ``scores``; fitting in-sample gives an optimistic ECE. """ from sklearn.linear_model import LogisticRegression + if (fit_scores is None) != (fit_labels is None): + raise ValueError("pass fit_scores and fit_labels together, or neither") scores = np.asarray(scores, dtype=np.float64).reshape(-1, 1) - labels = np.asarray(labels).astype(int) + if fit_scores is None: + fit_scores, fit_labels = scores, labels + fit_scores = np.asarray(fit_scores, dtype=np.float64).reshape(-1, 1) + fit_labels = np.asarray(fit_labels).astype(int) lr = LogisticRegression() - lr.fit(scores, labels) + lr.fit(fit_scores, fit_labels) return lr.predict_proba(scores)[:, 1] @@ -119,21 +130,44 @@ def run_calibration( "Re-run the missing-modality evaluator with prediction dumping enabled." ) labels = np.asarray(preds["ef_le_40"]).astype(int) - # Lower predicted LVEF => higher P(EF<=40); Platt learns the mapping. - prob = platt_probabilities(-np.asarray(preds["prediction"], dtype=np.float64), labels) + # Lower predicted LVEF => higher P(EF<=40); Platt learns the mapping. Fit on the val + # split's predictions when the results JSON carries them, so test stays untouched. + val_preds = payload.get("predictions_val", {}).get(condition) + fit_kwargs = {} + if val_preds: + fit_kwargs = { + "fit_scores": -np.asarray(val_preds["prediction"], dtype=np.float64), + "fit_labels": np.asarray(val_preds["ef_le_40"]).astype(int), + } + prob = platt_probabilities( + -np.asarray(preds["prediction"], dtype=np.float64), labels, **fit_kwargs + ) out_dir = Path(out_dir) figure = plot_reliability(labels, prob, out_dir / "reliability.pdf", n_bins=n_bins) + # carry the prediction source's provenance so this artifact is traceable on its own + source = { + key: payload[key] + for key in ("checkpoint", "checkpoint_sha256", "manifest_sha256", "seed", "git_sha") + if key in payload + } result = { "task": "E04_ef40_calibration", "git_sha": git_sha(), + "predictions_path": str(predictions_path), + "prediction_source": source, "condition": condition, "n": int(len(labels)), "n_bins": n_bins, + "scaler_fit_on": "val" if val_preds else "test (in-sample)", "ece": round(expected_calibration_error(labels, prob, n_bins), 4), "bins": reliability_bins(labels, prob, n_bins), "figure": str(figure), - "note": "In-sample Platt scaling; use a held-out calibration split for a deployment estimate.", + "note": ( + "Platt scaler fit on val predictions, applied to test." + if val_preds + else "In-sample Platt scaling; use a held-out calibration split for a deployment estimate." + ), } (out_dir / "calibration.json").write_text(json.dumps(result, indent=2)) return result diff --git a/src/primed_ai/evaluation/fairness.py b/src/primed_ai/evaluation/fairness.py index 26eb282..2bc9cb7 100644 --- a/src/primed_ai/evaluation/fairness.py +++ b/src/primed_ai/evaluation/fairness.py @@ -16,14 +16,23 @@ auroc, collate_tokens, git_sha, - read_table, regression_metrics, + sha256_file, +) +from primed_ai.probes.cross_attn import ( + MISSING_MODALITY_CONDITIONS, + _condition_arrays, + fused_probe_loader, + prepare_fused_probe_data, ) -from primed_ai.probes.concat_mlp import _ensure_ecg_tokens -from primed_ai.probes.echo_only import _ensure_tokens logger = logging.getLogger(__name__) +MIMIC_BIAS_NOTE = ( + "MIMIC-IV records administrative gender as 'sex' and admission-reported race; both " + "carry known curation bias (see MIMIC-IV documentation), so strata inherit it." +) + def _safe_json(obj): # Convert NaN/inf to None recursively for JSON compatibility @@ -38,32 +47,6 @@ def _safe_json(obj): return obj -def load_and_merge( - cohort_path: str | Path, echo_path: str | Path, ecg_path: str | Path -) -> pd.DataFrame: - coh = read_table(cohort_path) - echo = read_table(echo_path) - ecg = read_table(ecg_path) - - echo_key = "echo_study_id" if "echo_study_id" in coh.columns else "subject_id" - ecg_key = "ecg_record_id" - if echo_key not in echo.columns: - logger.debug("Renaming first echo embedding column -> %s", echo_key) - echo = echo.rename(columns={echo.columns[0]: echo_key}) - if ecg_key not in ecg.columns: - logger.debug("Renaming first ecg embedding column -> %s", ecg_key) - ecg = ecg.rename(columns={ecg.columns[0]: ecg_key}) - - df = coh.merge(echo, on=echo_key, how="inner").merge(ecg, on=ecg_key, how="inner") - return df - - -def prepare_tokens(df: pd.DataFrame) -> pd.DataFrame: - df2 = _ensure_tokens(df) - df2 = _ensure_ecg_tokens(df2) - return df2 - - def load_model( checkpoint: str | Path, embed_dim: int, @@ -164,6 +147,8 @@ def _coarsen_race(value) -> str: if s not in df.columns: logger.warning("Stratum column %s not in DataFrame; filling 'unknown'", s) df[s] = "unknown" + # cohorts with missing demographics mix NaN into string columns; sorted() would choke + df[s] = df[s].where(df[s].notna(), "unknown") groups = {} for val in sorted(df[s].unique()): mask = (df[s] == val).to_numpy() @@ -247,32 +232,66 @@ def save_outputs(results: dict, out_dir: str | Path): def run_fairness( cohort_path, - echo_path, - ecg_path, - checkpoint, + echo_path=None, + ecg_path=None, + checkpoint=None, out_dir="results/fairness", embed_dim=16, echo_dim=None, ecg_dim=None, batch_size=64, device=None, + conditions=("full",), ): - df = load_and_merge(cohort_path, echo_path, ecg_path) - df = prepare_tokens(df) - test_df = df[df.get("split") == "test"].reset_index(drop=True) - if test_df.empty: - raise RuntimeError("test split is empty; ensure 'split' column contains 'test' partition") + """Stratify test predictions by demographics, per missing-modality condition. + + Omit both embedding paths to read ``cohort_path`` as a joined manifest (matching the + probes and the missing-modality eval). Either way the data goes through + ``prepare_fused_probe_data``, so non-finite rows are dropped the same way and the + stratified n lines up with the canonical rerun. + """ + if checkpoint is None: + raise ValueError("checkpoint is required") + unknown = set(conditions) - set(MISSING_MODALITY_CONDITIONS) + if unknown: + raise ValueError( + f"unknown conditions {sorted(unknown)}; pick from {MISSING_MODALITY_CONDITIONS}" + ) + + parts, n_dropped = prepare_fused_probe_data(cohort_path, echo_path, ecg_path) + test_df = parts["test"] + model, device = load_model(checkpoint, embed_dim, device, echo_dim=echo_dim, ecg_dim=ecg_dim) - y_pred, y_true, ef_flags = predict_on_df( - model, test_df, embed_dim, batch_size, device, echo_dim=echo_dim - ) + loader = fused_probe_loader(test_df, embed_dim=echo_dim or embed_dim, batch_size=batch_size) results = { "task": "E03_fairness_stratification", "checkpoint": str(checkpoint), + "checkpoint_sha256": sha256_file(checkpoint), + "manifest_sha256": sha256_file(cohort_path), "git_sha": git_sha(), "n_test": int(len(test_df)), + # aggregate count across train/val/test, not test-only — mirrors the run metadata + "n_dropped_nonfinite_all_splits": n_dropped, + "bias_note": MIMIC_BIAS_NOTE, + "conditions": {}, } - strat = compute_stratified_results(test_df, y_true, y_pred, ef_flags) - results.update(strat) - save_outputs(results, out_dir) + out = Path(out_dir) + provenance = dict(results) + del provenance["conditions"] + for condition in conditions: + arrays = _condition_arrays(model, loader, device, condition) + strat = compute_stratified_results( + test_df, + np.asarray(arrays["lvef"]), + np.asarray(arrays["prediction"]).flatten(), + np.asarray(arrays["ef_le_40"], dtype=bool), + ) + results["conditions"][condition] = strat + save_outputs({**provenance, "condition": condition, **strat}, out / condition) + # keep the full-condition metrics at top level too: aggregate.py reads overall/by there + results.update(results["conditions"].get("full", {})) + out.mkdir(parents=True, exist_ok=True) + (out / "fairness_metrics.json").write_text( + json.dumps(_safe_json(results), indent=2), encoding="utf-8" + ) return results diff --git a/src/primed_ai/evaluation/missing_modality.py b/src/primed_ai/evaluation/missing_modality.py index b1f4e9c..ac3a9dd 100644 --- a/src/primed_ai/evaluation/missing_modality.py +++ b/src/primed_ai/evaluation/missing_modality.py @@ -7,7 +7,13 @@ import numpy as np import torch -from primed_ai.probes.common import auroc, git_sha, regression_metrics, save_results +from primed_ai.probes.common import ( + auroc, + git_sha, + regression_metrics, + save_results, + sha256_file, +) from primed_ai.probes.cross_attn import ( MISSING_MODALITY_CONDITIONS, CrossAttnFusedProbe, @@ -101,6 +107,7 @@ def run( parts, n_dropped = prepare_fused_probe_data( cohort_path, echo_embedding_path, ecg_embedding_path ) + val_loader = fused_probe_loader(parts["val"], embed_dim=echo_dim, batch_size=batch_size) test_loader = fused_probe_loader(parts["test"], embed_dim=echo_dim, batch_size=batch_size) model = CrossAttnFusedProbe( @@ -112,6 +119,8 @@ def run( state = torch.load(checkpoint, map_location=device) model.load_state_dict(state) + # Val predictions ride along so E09 can fit its Platt scaler on val without rescoring. + val_predictions = predict_missing_modality(model, val_loader, device) test_predictions = predict_missing_modality(model, test_loader, device) test_metrics = { condition: _metrics_from_predictions(arrays) @@ -132,6 +141,8 @@ def run( "seed": seed, "git_sha": git_sha(), "checkpoint": str(checkpoint), + "checkpoint_sha256": sha256_file(checkpoint), + "manifest_sha256": sha256_file(cohort_path), "config": { "cohort_path": str(cohort_path), "echo_embedding_path": str(echo_embedding_path), @@ -144,13 +155,17 @@ def run( "n_bootstrap": n_bootstrap, "device": device, }, - "n": {"test": len(parts["test"])}, + "n": {"val": len(parts["val"]), "test": len(parts["test"])}, "test": test_metrics, "bootstrap": bootstrap, "predictions": { condition: _serializable_predictions(arrays) for condition, arrays in test_predictions.items() }, + "predictions_val": { + condition: _serializable_predictions(arrays) + for condition, arrays in val_predictions.items() + }, "metrics_table": [ { "condition": condition, diff --git a/src/primed_ai/probes/common.py b/src/primed_ai/probes/common.py index f1f7b6b..ee5a017 100644 --- a/src/primed_ai/probes/common.py +++ b/src/primed_ai/probes/common.py @@ -2,6 +2,7 @@ from __future__ import annotations +import hashlib import json import subprocess from pathlib import Path @@ -28,6 +29,20 @@ def git_sha() -> str: return "unknown" +def sha256_file(path: str | Path | None) -> str | None: + """Hex digest of a file, or None when it is missing — provenance must not kill a run.""" + if path is None: + return None + p = Path(path) + if not p.is_file(): + return None + h = hashlib.sha256() + with p.open("rb") as f: + for chunk in iter(lambda: f.read(1 << 20), b""): + h.update(chunk) + return h.hexdigest() + + def auroc(y_true: np.ndarray, score: np.ndarray) -> float: y_true = np.asarray(y_true, dtype=bool) if y_true.any() and (~y_true).any(): diff --git a/src/primed_ai/probes/cross_attn.py b/src/primed_ai/probes/cross_attn.py index dd545b7..35e96a6 100644 --- a/src/primed_ai/probes/cross_attn.py +++ b/src/primed_ai/probes/cross_attn.py @@ -220,7 +220,7 @@ def run( cohort_path, echo_embedding_path=None, ecg_embedding_path=None, - out_dir="probes/cross_attn_fused", + out_dir="probes/fused", *, embed_dim: int = 16, echo_dim: int | None = None, diff --git a/tests/test_build_cohort.py b/tests/test_build_cohort.py index 7337238..f6851c4 100644 --- a/tests/test_build_cohort.py +++ b/tests/test_build_cohort.py @@ -12,6 +12,7 @@ write_demographics_coverage, write_flowchart, ) +from run_cohort_sensitivity import summarize_sensitivity def test_funnel_stages_window_mode(): @@ -91,3 +92,31 @@ def test_write_flowchart(tmp_path): text = path.read_text() assert "flowchart TD" in text assert "studies = 200" in text + + +def test_summarize_sensitivity_compares_final_rows(): + funnels = { + "window_24h": pd.DataFrame( + { + "stage": ["1. Echo", "2. Paired"], + "n_studies": [1000, 200], + "n_subjects": [900, 180], + "excluded_studies": [0, 800], + } + ), + "window_48h": pd.DataFrame( + { + "stage": ["1. Echo", "2. Paired"], + "n_studies": [1000, 260], + "n_subjects": [900, 230], + "excluded_studies": [0, 740], + } + ), + } + + summary = summarize_sensitivity(funnels) + + assert summary.loc[0, "scenario"] == "window_24h" + assert summary.loc[0, "final_rows"] == 200 + assert summary.loc[1, "gain_vs_first"] == 60 + assert summary.loc[1, "gain_vs_first_pct"] == 30.0 diff --git a/tests/test_calibration.py b/tests/test_calibration.py index 49a8e23..c30d4a8 100644 --- a/tests/test_calibration.py +++ b/tests/test_calibration.py @@ -60,3 +60,44 @@ def test_run_calibration_end_to_end(tmp_path): assert res["n"] == n assert (tmp_path / "calibration" / "calibration.json").exists() assert (tmp_path / "calibration" / "reliability.pdf").exists() + + +def test_run_calibration_fits_on_val_when_present(tmp_path): + """E09 (#66): the scaler must be fit on val predictions, never on test. + + Val and test carry opposite score-label relationships, so a val-fit scaler is + confidently wrong on test (huge ECE) while an in-sample fit would look calibrated. + """ + rng = np.random.default_rng(3) + n = 100 + test_pred = rng.uniform(20, 70, size=n) + test_labels = (test_pred <= 40).astype(int).tolist() # low prediction => positive + val_pred = rng.uniform(20, 70, size=n) + val_labels = (val_pred > 40).astype(int).tolist() # inverted relationship + pj = tmp_path / "missing_modality.json" + pj.write_text( + json.dumps( + { + "checkpoint": "probes/fused/cross_attn_fused.pt", + "checkpoint_sha256": "abc123", + "seed": 42, + "predictions": { + "full": {"ef_le_40": test_labels, "prediction": test_pred.tolist()} + }, + "predictions_val": { + "full": {"ef_le_40": val_labels, "prediction": val_pred.tolist()} + }, + } + ) + ) + + res = run_calibration(pj, tmp_path / "calibration", condition="full", n_bins=10) + assert res["scaler_fit_on"] == "val" + assert res["ece"] > 0.5 + # the artifact must be traceable on its own: source checkpoint/seed ride along + assert res["predictions_path"] == str(pj) + assert res["prediction_source"] == { + "checkpoint": "probes/fused/cross_attn_fused.pt", + "checkpoint_sha256": "abc123", + "seed": 42, + } diff --git a/tests/test_fairness.py b/tests/test_fairness.py index bfb344f..ca4cad1 100644 --- a/tests/test_fairness.py +++ b/tests/test_fairness.py @@ -63,3 +63,85 @@ def test_fairness_predictions_ignore_clip_padding(): batched, _, _ = predict_on_df(model, df, EMBED, batch_size=8, device="cpu", echo_dim=ECHO) alone, _, _ = predict_on_df(model, df, EMBED, batch_size=1, device="cpu", echo_dim=ECHO) np.testing.assert_allclose(batched, alone, atol=1e-5) + + +def test_run_fairness_reads_manifest_and_stratifies_per_condition(tmp_path): + """E08 (#65): fairness must score the joined manifest under dropped conditions too. + + Synthetic fixture; only plumbing is asserted. One test row carries a non-finite ECG + vector so the stratified n must line up with the canonical rerun's drop behavior. + """ + import numpy as np + import pandas as pd + import torch + + from primed_ai.evaluation.fairness import MIMIC_BIAS_NOTE, run_fairness + from primed_ai.probes import cross_attn + + EMBED, ECHO, ECG = 16, 24, 12 + n = 30 + rng = np.random.default_rng(0) + echo = rng.standard_normal((n, ECHO)).astype("float32") + ecg = rng.standard_normal((n, ECG)).astype("float32") + lvef = np.clip(55 + 15 * echo[:, 0], 10, 80) + ecg[5, 0] = np.nan # index 5 % 3 == 2 lands in the test split below + df = pd.DataFrame( + { + "subject_id": np.arange(n), + "lvef": lvef, + "ef_le_40": lvef <= 40.0, + "split": np.array(["train", "val", "test"])[np.arange(n) % 3], + "sex": np.where(np.arange(n) % 2 == 0, "F", "M"), + "age_band": np.where(np.arange(n) < 15, "40-54", "65-74"), + "race": np.where(np.arange(n) % 3 == 0, "WHITE", "BLACK/AFRICAN AMERICAN"), + "echo_embedding": list(echo), + "ecg_embedding": list(ecg), + } + ) + path = tmp_path / "manifest.parquet" + df.to_parquet(path, index=False) + + torch.manual_seed(0) + model = cross_attn.CrossAttnFusedProbe(EMBED, echo_dim=ECHO, ecg_dim=ECG) + ckpt = tmp_path / "fused.pt" + torch.save(model.state_dict(), ckpt) + + res = run_fairness( + path, + checkpoint=ckpt, + out_dir=tmp_path / "fairness", + embed_dim=EMBED, + echo_dim=ECHO, + ecg_dim=ECG, + device="cpu", + conditions=("full", "echo_dropped"), + ) + + assert set(res["conditions"]) == {"full", "echo_dropped"} + assert res["n_dropped_nonfinite_all_splits"] == 1 + from primed_ai.probes.common import sha256_file + + assert res["checkpoint_sha256"] == sha256_file(ckpt) + assert res["manifest_sha256"] == sha256_file(path) + assert res["n_test"] == 9 # 10 test rows minus the non-finite one + assert res["bias_note"] == MIMIC_BIAS_NOTE + full = res["conditions"]["full"] + assert set(full["by"]) == {"sex", "age_band", "race"} + assert all(m["small_n"] for m in full["by"]["sex"].values()) + # aggregate.py reads overall/by at the top level of the combined payload + assert res["overall"] == full["overall"] and res["by"] == full["by"] + # masking the echo branch must actually change the predictions + assert full["overall"]["mae"] != res["conditions"]["echo_dropped"]["overall"]["mae"] + for condition in ("full", "echo_dropped"): + assert (tmp_path / "fairness" / condition / "fairness_metrics.json").exists() + assert (tmp_path / "fairness" / condition / "fairness_summary.csv").exists() + assert (tmp_path / "fairness" / "fairness_metrics.json").exists() + + +def test_run_fairness_rejects_unknown_conditions(tmp_path): + import pytest + + from primed_ai.evaluation.fairness import run_fairness + + with pytest.raises(ValueError, match="unknown conditions"): + run_fairness("whatever.parquet", checkpoint="x.pt", conditions=("upside_down",)) diff --git a/tests/test_probe_manifest.py b/tests/test_probe_manifest.py index b0e34a2..e5f5924 100644 --- a/tests/test_probe_manifest.py +++ b/tests/test_probe_manifest.py @@ -667,6 +667,14 @@ def test_missing_modality_scores_a_manifest_trained_checkpoint(tmp_path): ) assert set(res["test"]) == {"full", "echo_dropped", "ecg_dropped"} assert (tmp_path / "missing.json").exists() + # E09 fits its Platt scaler on val, so val predictions must ride along per condition + assert set(res["predictions_val"]) == {"full", "echo_dropped", "ecg_dropped"} + assert res["n"]["val"] == len(res["predictions_val"]["full"]["prediction"]) + # reproduction provenance: hashes of exactly what was scored + from primed_ai.probes.common import sha256_file + + assert res["manifest_sha256"] == sha256_file(path) + assert res["checkpoint_sha256"] == sha256_file(probe_dir / "cross_attn_fused.pt") def test_missing_modality_requires_a_checkpoint(tmp_path): @@ -708,3 +716,30 @@ def test_training_writes_a_run_manifest_recording_the_pooling_regime(tmp_path): assert meta["fusion_dim"] == 256 and meta["seed"] == 42 assert meta["manifest"] == str(clip.resolve()) assert meta["git_sha"] != "" + + +def test_echo_to_ecg_attention_is_degenerate_on_tiled_ecg_tokens(): + """Pins the #72 finding: tiling one pooled ECG vector kills echo->ECG attention. + + Synthetic fixture. Identical keys force uniform softmax weights and identical values + make the weighted sum equal that value, so ``echo_to_ecg`` output cannot depend on the + echo query. Until token-level ECG embeddings exist, the fused probe's echo->ECG + direction is a fixed linear image of the pooled ECG vector, not cross-modal attention. + """ + import torch + + from primed_ai.probes.layers import CrossAttentionFusion + + torch.manual_seed(0) + fusion = CrossAttentionFusion(ECHO_DIM).eval() + rng = np.random.default_rng(5) + ecg_tiled = torch.as_tensor( + np.tile(rng.standard_normal((2, 1, ECHO_DIM)), (1, 4, 1)), dtype=torch.float32 + ) + echo_a = torch.as_tensor(rng.standard_normal((2, 6, ECHO_DIM)), dtype=torch.float32) + echo_b = 100.0 * torch.as_tensor(rng.standard_normal((2, 6, ECHO_DIM)), dtype=torch.float32) + + with torch.no_grad(): + ctx_a, _ = fusion.echo_to_ecg(echo_a, ecg_tiled, ecg_tiled) + ctx_b, _ = fusion.echo_to_ecg(echo_b, ecg_tiled, ecg_tiled) + torch.testing.assert_close(ctx_a, ctx_b)