diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..1f64da7 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,36 @@ +name: CI + +on: + push: + branches: [main, sprints] + pull_request: + +jobs: + checks: + runs-on: ubuntu-latest + + steps: + - name: Check out repository + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Set up uv + uses: astral-sh/setup-uv@v5 + with: + enable-cache: true + + - name: Install dependencies + run: uv sync --dev --frozen + + - name: Lint + run: uv run ruff check src scripts tests + + - name: Type check + run: uv run mypy src + + - name: Test + run: uv run pytest -q diff --git a/.gitignore b/.gitignore index 2634ff6..089124c 100644 --- a/.gitignore +++ b/.gitignore @@ -195,11 +195,14 @@ cython_debug/ .pypirc # AI +.agents/ +.codex/ .cursorignore .cursorindexingignore .cursorrules .claude/ .cursor/ +.rulesync/ # Marimo marimo/_static/ @@ -208,13 +211,16 @@ __marimo__/ # Notes notes/ -docs/ +docs/internal/ # SLICES-specific /data/ outputs/ +/results/* +!/results/.gitkeep wandb/ .tensor_cache/ +logs/ # PyTorch Lightning lightning_logs/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a6c364a..c63645b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -14,7 +14,7 @@ repos: args: [--line-length=100, --target-version=py312] - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.3.4 + rev: v0.14.8 hooks: - id: ruff args: [--fix] diff --git a/EXPERIMENTS_PLAN.md b/EXPERIMENTS_PLAN.md new file mode 100644 index 0000000..2066398 --- /dev/null +++ b/EXPERIMENTS_PLAN.md @@ -0,0 +1,6 @@ +# Experiment Plan Pointer + +The active SLICES experiment plan lives at `docs/internal/EXPERIMENT_PLAN.md`. + +This file exists so tools or reviews that look for `EXPERIMENTS_PLAN.md` land on +the canonical plan rather than treating the plan as missing. diff --git a/EXPERIMENT_PLAN.md b/EXPERIMENT_PLAN.md deleted file mode 100644 index b7b4d6b..0000000 --- a/EXPERIMENT_PLAN.md +++ /dev/null @@ -1,387 +0,0 @@ -# SLICES Experiment Plan - -## Overview - -This document defines the full experiment matrix for the SLICES benchmark. The goal is to answer: - -> How do the three major SSL paradigm families — reconstruction, self-distillation, and contrastive learning — compare when applied to sparse, irregularly-sampled clinical time series under controlled conditions? - -Secondary questions addressed through ablations: - -1. Does SSL improve label efficiency over supervised baselines? -2. Do SSL representations transfer across hospital systems? -3. Do SSL paradigms differ in fairness properties? - ---- - -## 1. Experiment Matrix - -### 1.1 Independent Variables - -| Variable | Levels | Notes | -|----------|--------|-------| -| **Dataset** | MIMIC-IV, eICU, Combined | Combined = pooled pretraining corpus | -| **Paradigm** | MAE (reconstruction), JEPA (self-distillation), Contrastive, Supervised | Supervised = no pretraining baseline | -| **Downstream Task** | mortality_24h, mortality_hospital, aki_kdigo, los_remaining | 3 classification + 1 regression | -| **Seed** | 42, 123, 456 | 3 seeds for statistical significance | - -### 1.2 Controlled Variables (Held Constant) - -| Variable | Value | Rationale | -|----------|-------|-----------| -| Observation window | 48 hours | Standard in ICU benchmarks | -| Min stay | 48 hours | Ensures full observation window | -| Splits | 70/15/15 train/val/test | Patient-level, no leakage | -| Imputation | Normalize-then-zero-fill | Eliminates imputation as confound | -| Finetuning head | MLP, hidden_dims=[64] | Same head architecture across all paradigms | -| Finetuning LR | 1e-4 | Same optimization for all downstream runs | -| Finetuning epochs | 50 (patience=10) | Same budget | -| Weight decay | 0.01 | Same regularization | - -### 1.3 Paradigm-Intrinsic Differences (Documented, Not Controlled) - -These differences are inherent to each paradigm and are themselves a contribution: - -| Paradigm | Tokenization | Encoder | Rationale | -|----------|-------------|---------|-----------| -| MAE | Observation-level (1 token per observed value) | ObservationTransformer (d=128, L=4, H=8) | Avoids "mostly zeros" on sparse data | -| JEPA | Observation-level (1 token per observed value) | ObservationTransformer (d=128, L=4, H=8) | Predicts in latent space; shares encoder with MAE | -| Contrastive | Observation-level (1 token per observed value) | ObservationTransformer (d=128, L=4, H=8) | Two masked views; shares encoder with MAE | -| Supervised | Timestep-level | Transformer (d=64, L=2, H=4) | Different encoder (no paradigm constraints) | - ---- - -## 2. Training Protocol - -### 2.1 Pretraining - -| Setting | MAE | JEPA | Contrastive | -|---------|-----|------|-------------| -| Epochs | 500 | 500 | 500 | -| Batch size | 256 | 256 | 256 | -| Learning rate | 1e-3 | 1e-3 | 1e-3 | -| Scheduler | Warmup cosine (50 warmup) | Warmup cosine (50 warmup) | Warmup cosine (50 warmup) | -| Mask ratio | 0.5 (random masking) | 0.5 (random masking) | 0.5 (two random masked views) | -| Gradient clipping | 1.0 | 1.0 | 1.0 | -| Early stopping | Patience=10 on val loss | Patience=10 on val loss | Patience=10 on val loss | - -**Budget equalization**: Raw epoch counts differ because paradigms converge at different rates. For fair comparison, report **total gradient steps** and **wall-clock time** alongside epochs. If step counts diverge significantly (>2x), normalize by training the slower paradigm longer. - -### 2.2 Downstream Finetuning - -All paradigms use identical finetuning protocol: - -| Setting | Value | -|---------|-------| -| Protocol | Linear probing (freeze_encoder=true) | -| Epochs | 50 | -| Batch size | 64 | -| Learning rate | 1e-4 | -| Scheduler | Cosine decay (eta_min=1e-6) | -| Early stopping | Patience=10 on val metric (AUPRC for classification, MSE for regression) | -| Head | MLP, hidden_dims=[64], dropout=0.1, ReLU | -| Class weighting | None (report as-is; sensitivity analysis if needed) | - -**Linear probing rationale**: Isolates representation quality from finetuning optimization. Full finetuning results can be added as an appendix ablation. - -### 2.3 Supervised Baseline - -| Setting | Value | -|---------|-------| -| Epochs | 100 | -| Batch size | 64 | -| Learning rate | 1e-3 | -| Scheduler | Cosine decay | -| Early stopping | Patience=20 on val metric (AUPRC for classification, MSE for regression) | -| Encoder | Transformer (d=64, L=2, H=4), trained end-to-end | -| Head | MLP, hidden_dims=[64], dropout=0.1, ReLU | - ---- - -## 3. Evaluation Metrics - -### 3.1 Performance Metrics - -| Task | Type | Primary Metric | Secondary Metrics | -|------|------|---------------|-------------------| -| mortality_24h | Binary classification | AUPRC | AUROC, Brier score, ECE | -| mortality_hospital | Binary classification | AUPRC | AUROC, Brier score, ECE | -| aki_kdigo | Binary classification | AUPRC | AUROC | -| los_remaining | Regression | MAE | MSE, R² | - -**AUPRC as primary** for classification: ICU tasks are heavily imbalanced; AUPRC is more informative than AUROC in this regime. - -### 3.2 Fairness Metrics - -Computed on test set predictions for all classification tasks: - -| Metric | Definition | Threshold | -|--------|-----------|-----------| -| Per-group AUROC | AUROC computed per subgroup | Report gap (max - min) | -| Per-group AUPRC | AUPRC computed per subgroup | Report gap | -| Demographic parity difference | max\|P(Ŷ=1\|A=a) - P(Ŷ=1\|A=b)\| | Lower is better | -| Equalized odds difference | max(TPR gap, FPR gap) across groups | Lower is better | -| Disparate impact ratio | min_rate / max_rate | <0.8 = adverse impact | - -### 3.3 Protected Attributes - -| Attribute | Available In | Groups | Notes | -|-----------|-------------|--------|-------| -| **Sex** | MIMIC-IV, eICU | M, F | Primary; available in both datasets | -| **Age group** | MIMIC-IV, eICU | 18-44, 45-64, 65-79, 80+ | Primary; 4-bin split for granular fairness analysis | -| **Race/Ethnicity** | MIMIC-IV only | White, Black, Hispanic, Asian, Other | Secondary; MIMIC-only analysis | - -Min subgroup size: 50 patients. Groups below threshold are excluded from fairness metrics. - -### 3.4 Statistical Testing - -- Report mean ± std across 3 seeds for all metrics -- **Paired Wilcoxon signed-rank test** between paradigms (paired across seeds and tasks) -- Bonferroni correction for multiple comparisons -- Report effect sizes (Cohen's d) for key comparisons - ---- - -## 4. Main Experiments - -### 4.1 Phase 1: Pretraining (9 runs × 3 seeds = 27 runs) - -| ID | Dataset | Paradigm | Encoder | Config Override | -|----|---------|----------|---------|-----------------| -| P1 | MIMIC-IV | MAE | ObservationTransformer | ssl=mae model=observation_transformer | -| P2 | MIMIC-IV | JEPA | ObservationTransformer | ssl=jepa model=observation_transformer | -| P3 | MIMIC-IV | Contrastive | ObservationTransformer | ssl=contrastive model=observation_transformer | -| P4 | eICU | MAE | ObservationTransformer | ssl=mae model=observation_transformer | -| P5 | eICU | JEPA | ObservationTransformer | ssl=jepa model=observation_transformer | -| P6 | eICU | Contrastive | ObservationTransformer | ssl=contrastive model=observation_transformer | -| P7 | Combined | MAE | ObservationTransformer | ssl=mae model=observation_transformer | -| P8 | Combined | JEPA | ObservationTransformer | ssl=jepa model=observation_transformer | -| P9 | Combined | Contrastive | ObservationTransformer | ssl=contrastive model=observation_transformer | - -### 4.2 Phase 2: Downstream Evaluation (48 runs × 3 seeds = 144 runs) - -Each of the 9 pretrained encoders is finetuned on 4 downstream tasks (36 SSL finetuning runs). -Each of the 3 datasets gets a supervised baseline on 4 tasks (12 supervised runs). - -**Result Tables** (one per dataset): - -#### MIMIC-IV Results - -| | mortality_24h | mortality_hospital | aki_kdigo | los_remaining | -|---|---|---|---|---| -| **MAE** | AUPRC ± std | AUPRC ± std | AUPRC ± std | MAE ± std | -| **JEPA** | | | | | -| **Contrastive** | | | | | -| **Supervised** | | | | | - -Same table structure for eICU and Combined. - -### 4.3 Phase 3: Fairness Evaluation (on all 48 downstream runs) - -No additional training. Compute fairness metrics on test predictions from Phase 2. - -**Fairness Result Tables** (one per task, per dataset): - -#### Example: mortality_24h, MIMIC-IV - -| | AUROC gap (sex) | AUROC gap (age) | AUROC gap (race) | Dem. parity diff | Eq. odds diff | -|---|---|---|---|---|---| -| **MAE** | | | | | | -| **JEPA** | | | | | | -| **Contrastive** | | | | | | -| **Supervised** | | | | | | - ---- - -## 5. Ablation Experiments - -### 5.1 Label Efficiency (Few-Shot Finetuning) - -**Question**: Does SSL improve sample efficiency — how much labeled data is needed to match supervised performance? - -**Design**: Finetune with {1%, 5%, 10%, 25%, 50%, 100%} of labeled training data. - -| Scope | Tasks | Label Fractions | Runs | -|-------|-------|----------------|------| -| Full sweep | mortality_24h (anchor task) | 1%, 5%, 10%, 25%, 50%, 100% | 4 paradigms × 3 datasets × 6 fractions = 72 | -| Trend check | Remaining 3 tasks | 10%, 100% | 4 paradigms × 3 datasets × 3 tasks × 2 fractions = 72 | -| **Total** | | | **144 runs × 3 seeds = 432 runs** | - -**Expected finding**: SSL paradigms should outperform supervised at ≤10% labels. The label fraction at which supervised catches up (crossover point) is a key metric. - -**Visualization**: Learning curves (AUPRC vs label fraction) per paradigm, one plot per dataset. - -### 5.2 Cross-Dataset Transfer - -**Question**: Do SSL representations generalize across hospital systems? - -**Design**: Pretrain on source dataset, finetune on target dataset. - -| Source → Target | Paradigms | Tasks | Runs | -|-----------------|-----------|-------|------| -| MIMIC-IV → eICU | MAE, JEPA, Contrastive | All 4 | 12 | -| eICU → MIMIC-IV | MAE, JEPA, Contrastive | All 4 | 12 | -| **Total** | | | **24 runs × 3 seeds = 72 runs** | - -**Baselines for comparison** (from main experiments, no extra runs): -- In-domain pretraining (pretrain and finetune on same dataset) -- Supervised (no transfer, trained from scratch on target) - -**Expected finding**: SSL transfer should outperform supervised-from-scratch on target, especially in low-data regimes. Combined pretraining (Phase 2) should perform best. - -**Note**: Supervised has no transfer story — this asymmetry is itself a finding. - -### 5.3 Mask Ratio Sensitivity - -**Question**: How sensitive are SSL paradigms to the masking ratio? - -**Design**: Pretrain with {0.3, 0.5, 0.75} mask ratios on MIMIC-IV, finetune on mortality_24h. - -| Mask Ratio | Paradigms | Dataset | Task | Runs | -|------------|-----------|---------|------|------| -| 0.3 | MAE, JEPA, Contrastive | MIMIC-IV | mortality_24h | 3 | -| 0.5 | MAE, JEPA, Contrastive | MIMIC-IV | mortality_24h | 3 (from main experiments) | -| 0.75 | MAE, JEPA, Contrastive | MIMIC-IV | mortality_24h | 3 | -| **Total** | | | | **6 new pretraining + 6 new finetuning = 12 runs × 3 seeds = 36 runs** | - -**Note**: The 0.5 runs are reused from the main experiments (P1–P3). Only 0.3 and 0.75 require additional pretraining. - -**Visualization**: Bar chart or line plot of AUPRC vs mask ratio, one line per paradigm. - ---- - -## 6. Run Summary - -| Phase | Description | Runs | × 3 seeds | Cumulative | -|-------|-------------|------|-----------|------------| -| Phase 1 | Pretraining | 9 | 27 | 27 | -| Phase 2 | Downstream finetuning | 48 | 144 | 171 | -| Phase 3 | Fairness evaluation | 0 (eval only) | 0 | 171 | -| Ablation 5.1 | Label efficiency | 144 | 432 | 603 | -| Ablation 5.2 | Cross-dataset transfer | 24 | 72 | 675 | -| Ablation 5.3 | Mask ratio sensitivity | 12 | 36 | 711 | -| **Total** | | **237** | **711** | | - ---- - -## 7. Execution Order - -Priority ordering for incremental results and early debugging: - -### Sprint 1: Sanity Check (4 runs) -1. MIMIC-IV, all 4 paradigms (MAE, JEPA, Contrastive, Supervised), mortality_24h only, seed=42 -2. Verify: training converges, metrics are reasonable, pipeline end-to-end works -3. Check gradient step counts across paradigms — adjust epoch budgets if needed - -### Sprint 2: First Full Dataset (16 runs) -4. MIMIC-IV, all 4 paradigms × 4 tasks, seed=42 -5. Produces first complete results table for one dataset -6. Write preliminary results section - -### Sprint 3: Generalization (16 runs) -7. eICU, all 4 paradigms × 4 tasks, seed=42 -8. Cross-dataset comparison possible - -### Sprint 4: Scaling (16 runs) -9. Combined dataset, all 4 paradigms × 4 tasks, seed=42 -10. All three main result tables complete (single seed) - -### Sprint 5: Statistical Robustness (96 runs) -11. Seeds 123 and 456 for all Sprint 2–4 runs -12. Compute mean ± std, run statistical tests - -### Sprint 6: Label Efficiency Ablation (432 runs) -13. Full sweep on mortality_24h anchor -14. Trend check on remaining tasks -15. Generate learning curve plots - -### Sprint 7: Transfer Ablation (72 runs) -16. Cross-dataset transfer experiments -17. Compare against in-domain baselines - -### Sprint 8: Fairness Analysis (0 extra runs) -18. Compute fairness metrics on all Phase 2 test predictions -19. Enable `eval.fairness.enabled=true` and rerun evaluation -20. Generate fairness tables and disparity plots - ---- - -## 8. W&B Tracking - -### Project Structure - -All runs logged to W&B project `slices`. - -### Run Naming Convention - -``` -{phase}_{dataset}_{paradigm}_{task}_seed{seed} -``` - -Examples: -- `pretrain_mimic_mae_seed42` -- `finetune_eicu_jepa_mortality24h_seed123` -- `supervised_combined_aki_kdigo_seed456` -- `ablation-label_mimic_mae_mortality24h_frac10_seed42` -- `ablation-transfer_mimic2eicu_contrastive_aki_seed42` - -### Tags - -| Tag | Purpose | -|-----|---------| -| `phase:pretrain` / `phase:finetune` / `phase:supervised` | Training phase | -| `dataset:mimic` / `dataset:eicu` / `dataset:combined` | Dataset | -| `paradigm:mae` / `paradigm:jepa` / `paradigm:contrastive` / `paradigm:supervised` | SSL family | -| `task:mortality_24h` / etc. | Downstream task | -| `ablation:label-efficiency` / `ablation:transfer` | Ablation type | -| `sprint:N` | Execution sprint | -| `seed:42` / `seed:123` / `seed:456` | Random seed | - -### Key Metrics to Log - -| Phase | Metrics | -|-------|---------| -| Pretraining | train_loss, val_loss, gradient_steps, wall_clock_time | -| Finetuning | val_auroc, val_auprc, test_auroc, test_auprc, test_brier, test_ece | -| Regression | val_mae, test_mae, test_mse, test_r2 | -| Fairness | fairness/auroc_gap_sex, fairness/dem_parity_diff, fairness/eq_odds_diff | - ---- - -## 9. Expected Outputs - -### Tables for Paper - -1. **Table 1**: Main results — AUPRC (mean ± std) for each paradigm × dataset × task (3 tables or 1 large table) -2. **Table 2**: Statistical significance — p-values for pairwise paradigm comparisons -3. **Table 3**: Fairness — AUROC gap and equalized odds difference per paradigm per protected attribute -4. **Table 4**: Cross-dataset transfer — AUPRC for source→target vs in-domain vs supervised - -### Figures for Paper - -1. **Figure 1**: Architecture diagram showing paradigm-intrinsic tokenization differences -2. **Figure 2**: Learning curves — AUPRC vs label fraction per paradigm (one subplot per dataset) -3. **Figure 3**: Radar/spider chart — paradigm comparison across all tasks (one per dataset) -4. **Figure 4**: Fairness disparity heatmap — paradigm × protected attribute × task -5. **Figure 5**: Transfer gap — bar chart of (transfer AUPRC - supervised AUPRC) per paradigm - -### Appendix - -- Full finetuning results (unfreeze encoder) if time permits -- Per-seed results for transparency -- Training curves (loss vs steps) for all pretraining runs -- Hyperparameter sensitivity (if reviewer requests) - ---- - -## 10. Risk Mitigation - -| Risk | Mitigation | -|------|-----------| -| Contrastive paradigm not yet implemented | Implement before Sprint 1; if delayed, run MAE + JEPA + Supervised first | -| eICU lacks race/ethnicity data | Fairness analysis for race is MIMIC-only; document as limitation | -| Phenotyping task is MIMIC-only | Excluded from main matrix; can add as MIMIC-specific appendix | -| Unequal training budgets across paradigms | Track gradient steps; normalize if >2x difference | -| Class imbalance affects metrics | Use AUPRC as primary; report class ratios; sensitivity analysis with class_weight="balanced" | -| Compute budget exceeded | Sprint ordering ensures usable results at each checkpoint | -| Reviewer requests 5 seeds | Add 2 seeds only to contested comparisons (minor revision) | diff --git a/README.md b/README.md index 8229a0d..f8571b5 100644 --- a/README.md +++ b/README.md @@ -21,11 +21,15 @@ How do the three major SSL paradigm families — **reconstruction** (masked autoencoding), **self-distillation** (JEPA), and **contrastive learning** — compare when applied to clinical time series under controlled conditions? +The benchmark includes controlled SSL objectives, supervised training from scratch, +classical ICU baselines, and targeted extensions for temporal contrastive learning +and model-capacity studies. + ### The Comparison Triangle | Comparison | What Varies | What It Tests | |---|---|---| -| **MAE vs JEPA** | Input-space vs latent-space prediction | Same masking, same encoder input | +| **MAE vs JEPA** | Input-space vs latent-space prediction | Same encoder input, tokenization, and mask budget | | **JEPA vs Contrastive** | Local positional prediction vs global invariance | Both operate in latent space | | **MAE vs Contrastive** | Reconstruction vs discrimination | Opposite ends of the SSL spectrum | @@ -37,7 +41,10 @@ Fair comparison of SSL objectives for clinical time series is currently impossib ## SSL Paradigms -All observation-level objectives share the same `ObservationTransformerEncoder` (one token per observed measurement) and the same masking logic (`masking.py`), differing only in what they predict and how they compute loss: +The controlled SSL objectives share the same timestep-level obs-aware Transformer +encoder, obs-aware tokenization, and default mask budget. The masking strategy is +reported as part of each objective because JEPA uses block masking to avoid the +random-mask interpolation failure observed during development. | Objective | Predicts | Target | Loss | |---|---|---|---| @@ -45,15 +52,17 @@ All observation-level objectives share the same `ObservationTransformerEncoder` | **JEPA** | Latent representations at masked positions | EMA target encoder representations | MSE / Cosine | | **Contrastive** | Global embedding similarity across views | Positive pair agreement (NT-Xent) | Cross-entropy | -**SMART** (NeurIPS 2024) is also included in the codebase as a sanity check and to demonstrate the framework's extensibility — it uses its own MART encoder and element-wise masking, so it is not part of the controlled thesis experiments. +**TS2Vec** is included as a temporal-contrastive extension. **SMART** (NeurIPS +2024) remains an external reference because it swaps in its own MART encoder and +element-wise masking, so it is not part of the controlled comparison. ## Pipeline ``` RICU (R) ──→ Parquet ──→ ICUDataset ──→ SSL Pretraining ──→ Downstream Finetuning - hourly-binned dense tensors MAE / JEPA / mortality, LOS, - feature extraction + obs masks Contrastive AKI - across datasets + hourly-binned dense tensors MAE / JEPA / mortality, mortality_24h, + feature extraction + obs masks Contrastive mortality_hospital, + across datasets los_remaining, aki_kdigo ``` 1. **Extraction**: RICU (R package) harmonizes raw ICU data (MIMIC-IV, eICU) into hourly-binned parquet files @@ -90,16 +99,26 @@ uv run python -c "from slices.models.pretraining import JEPAObjective, Contrasti # R extraction (once per dataset) Rscript scripts/preprocessing/extract_with_ricu.R --dataset miiv -# Python processing +# Python processing + splits/normalization +uv run python scripts/preprocessing/build_processed_data.py --datasets miiv + +# Full benchmark rebuild from existing RICU outputs +uv run python scripts/preprocessing/build_processed_data.py --datasets miiv eicu --combined +``` + +The wrapper calls the lower-level scripts below. Use them directly when you need +to rerun only one stage: + +```bash uv run python scripts/preprocessing/extract_ricu.py dataset=miiv -# Compute splits & normalization stats uv run python scripts/preprocessing/prepare_dataset.py dataset=miiv ``` ### 2. Pretrain -Pick SSL paradigm with `ssl=`. Training budget is matched across thesis paradigms: +Pick SSL paradigm with `ssl=`. Training budget is matched across controlled +paradigms: ```bash # MAE (masked autoencoder — reconstruction baseline, default) @@ -111,7 +130,10 @@ uv run python scripts/training/pretrain.py dataset=miiv ssl=jepa # Contrastive (SimCLR-style with two masked views) uv run python scripts/training/pretrain.py dataset=miiv ssl=contrastive -# SMART (sanity check / extensibility demo — not part of thesis experiments) +# TS2Vec (temporal contrastive extension) +uv run python scripts/training/pretrain.py dataset=miiv ssl=ts2vec + +# SMART (external reference with a different encoder/tokenization contract) uv run python scripts/training/pretrain.py dataset=miiv ssl=smart model=smart ``` @@ -127,6 +149,33 @@ uv run python scripts/training/finetune.py dataset=miiv checkpoint=outputs/.../e uv run python scripts/training/supervised.py dataset=miiv ``` +### 5. Validate + +```bash +uv run pytest tests/ -q +``` + +## Benchmark Release Notes + +The public benchmark contract is the controlled comparison of MAE, JEPA, and +contrastive SSL under the shared RICU pipeline, canonical obs-aware +`TransformerEncoder`, fixed downstream tasks, and class-based experiment +metadata. TS2Vec is a temporal-contrastive extension, and SMART is an external +reference because it changes the encoder/tokenization contract. + +For thesis-scale reruns, use `scripts/internal/run_experiments.py` or +`scripts/internal/launch_thesis_tmux.sh` with an explicit `--revision`, +`--project`, and `--launch-commit`. Downstream SSL runs consume `encoder.pt`, +the last encoder from the fixed pretraining schedule, not `encoder_best_val.pt`; +downstream task evaluation still records the exact finetune checkpoint used for +test metrics and post-hoc fairness. + +Final publication export expects post-hoc fairness metrics to be written first +with `scripts/eval/evaluate_fairness.py`. The export emits per-seed rows, +aggregate mean/std/min/max/95% CI columns, derived comparison views, and +pairwise statistical tests. Do not use export or fairness escape hatches for +final benchmark tables. + ## Project Structure ``` @@ -143,15 +192,17 @@ SLICES/ │ │ └── sliding_window.py # Sliding window utilities │ ├── models/ │ │ ├── encoders/ # Backbone architectures (factory pattern) -│ │ │ ├── transformer.py # Timestep-level Transformer -│ │ │ ├── observation_transformer.py # Observation-level Transformer +│ │ │ ├── transformer.py # Canonical timestep-level obs-aware Transformer +│ │ │ ├── observation.py # Observation-level Transformer (legacy/alternate) +│ │ │ ├── gru_d.py # GRU-D baseline encoder │ │ │ └── smart.py # SMART/MART encoder │ │ ├── pretraining/ # SSL objectives (factory pattern) │ │ │ ├── masking.py # Shared observation masking │ │ │ ├── mae.py # Masked Autoencoder │ │ │ ├── jepa.py # Joint-Embedding Predictive Architecture │ │ │ ├── contrastive.py # SimCLR-style contrastive -│ │ │ └── smart.py # SMART (sanity check / extensibility demo) +│ │ │ ├── ts2vec.py # Temporal contrastive extension +│ │ │ └── smart.py # SMART (appendix-only external reference) │ │ ├── heads/ # Task heads (MLP, Linear) │ │ └── common.py # Shared utilities │ ├── training/ @@ -159,21 +210,23 @@ SLICES/ │ │ ├── finetune_module.py # FineTuneModule (Lightning) │ │ └── utils.py # Shared helpers (optimizer, scheduler, callbacks, logger, validation) │ └── eval/ -│ ├── metrics.py # AUROC, AUPRC, F1, MSE, etc. -│ ├── fairness.py # Per-group AUROC, demographic parity +│ ├── metrics.py # AUROC, AUPRC, Brier, ECE, MSE, MAE, etc. +│ ├── fairness.py # Fairness primitives (parity, odds, impact) +│ ├── fairness_evaluator.py # Patient-aware subgroup fairness reports +│ ├── statistical.py # Bootstrap, Wilcoxon, Bonferroni, Cohen's d │ └── imputation.py # SSL reconstruction quality ├── configs/ # Hydra configs -│ ├── pretrain.yaml # Unified SSL pretraining (ssl=mae/jepa/contrastive/smart) +│ ├── pretrain.yaml # Unified SSL pretraining (ssl=mae/jepa/contrastive/ts2vec/smart) │ ├── finetune.yaml │ ├── supervised.yaml │ ├── data/ # Dataset configs │ ├── model/ # Encoder configs -│ ├── ssl/ # SSL objective configs (mae, jepa, contrastive, smart) +│ ├── ssl/ # SSL objective configs (mae, jepa, contrastive, ts2vec, smart) │ └── tasks/ # Downstream task definitions ├── scripts/ # Entry point scripts │ ├── preprocessing/ │ └── training/ -└── tests/ # pytest test suite (940+ tests) +└── tests/ # pytest test suite ``` ## Configuration @@ -195,10 +248,12 @@ uv run python scripts/training/pretrain.py dataset=miiv ssl=jepa training.overfi | Group | Options | Purpose | |---|---|---| -| `ssl/` | `mae`, `jepa`, `contrastive`, `smart` | SSL objective hyperparameters | -| `model/` | `transformer`, `observation_transformer`, `smart` | Encoder architecture | +| `ssl/` | `mae`, `jepa`, `contrastive`, `ts2vec`, `smart` | SSL objective hyperparameters | +| `model/` | `transformer`, `transformer_medium`, `transformer_large`, `observation_transformer`, `smart`, `gru_d`, `linear` | Encoder architecture | | `data/` | `ricu` | Dataset and paths | -| `tasks/` | `mortality`, `los`, `aki`, ... | Downstream task definitions | +| `tasks/` | `mortality`, `mortality_24h`, `mortality_hospital`, `los_remaining`, `aki_kdigo` | Downstream task definitions | +| `protocol/` | `linear_probe`, `full_finetune` | Downstream evaluation protocol overrides | +| `eval/` | `default` | Evaluation metric defaults | ## Data Format @@ -206,9 +261,13 @@ Extracted ICU stays are stored as Parquet files: - `static.parquet` — Stay-level demographics and admission info - `timeseries.parquet` — Dense hourly-binned time-series with observation masks (T x D) -- `labels.parquet` — Task labels (mortality, LOS, AKI, etc.) +- `labels.parquet` — Configured task labels (for example `mortality_24h`, `mortality_hospital`, `aki_kdigo`, `los_remaining`, and optional `mortality`) - `metadata.yaml` — Feature names, sequence length, task definitions +For task-specific supervised evaluation, stays with missing task labels are +excluded. This matters most for `aki_kdigo`, which requires creatinine in both +the 0-24h baseline window and the 24-48h prediction window. + `ICUDataset` returns batches with: - `timeseries`: `FloatTensor (B, T, D)` — hourly-binned feature values - `mask`: `BoolTensor (B, T, D)` — True = observed, False = missing @@ -224,9 +283,9 @@ uv run pytest tests/ -v uv run pytest tests/ --cov=slices --cov-report=html --cov-report=term # Format / lint / type check -black src/ scripts/ tests/ -ruff check src/ scripts/ tests/ -mypy src/ +uv run black src/ scripts/ tests/ +uv run ruff check src/ scripts/ tests/ +uv run mypy src/ ``` ## Key Design Decisions @@ -234,15 +293,20 @@ mypy src/ - **RICU-based extraction**: Data harmonization across datasets handled by RICU (R). Python reads the output. - **Normalize-then-zero-fill**: Single imputation strategy (z-normalize, fill missing with 0). Eliminates imputation as a confound. - **Observation masks**: Missingness tracked separately; SSL objectives use this for masking. -- **Shared masking**: MAE, JEPA, and Contrastive share identical masking code (`masking.py`) for fair comparison. +- **Controlled masking budget**: MAE, JEPA, and Contrastive share the same default mask ratio through the common `masking.py` infrastructure; objective-specific masking strategies are reported explicitly. +- **Reported HP sensitivity**: Contrastive robustness rows that disable complementary masks are labeled as view/mask sensitivity, not pure mask-ratio sensitivity. - **Patient-level splits**: No data leakage between train/val/test. - **Config-driven ablations**: Change one YAML default to switch paradigm, encoder, or task. ## Extending SLICES - +The framework uses factory patterns throughout, making it straightforward to add new components. The concrete extension points live in: -The framework uses factory patterns throughout, making it straightforward to add new components. See `EXTENDING_SLICES.md` (coming soon) for details on adding new downstream tasks, SSL objectives, encoder architectures, and datasets. +- `src/slices/models/encoders/factory.py` and `configs/model/` +- `src/slices/models/pretraining/factory.py` and `configs/ssl/` +- `src/slices/models/heads/factory.py` +- `src/slices/data/tasks/` for package label semantics used during extraction +- `configs/tasks/` for Hydra task configs used during training/evaluation ## References diff --git a/configs/data/ricu.yaml b/configs/data/ricu.yaml index e4f3690..2f83cae 100644 --- a/configs/data/ricu.yaml +++ b/configs/data/ricu.yaml @@ -38,6 +38,10 @@ feature_set: core # Feature categories (not used by RICU extractor, kept for config compatibility) categories: null +# Benchmark observation window / cohort criterion +seq_length_hours: 24 +min_stay_hours: 24 + # Tasks to extract labels for tasks: - mortality_24h diff --git a/configs/debug.yaml b/configs/debug.yaml index ac1c379..996dd27 100644 --- a/configs/debug.yaml +++ b/configs/debug.yaml @@ -5,7 +5,7 @@ # uv run python scripts/debug/export_snapshots.py # # # Override data paths: -# uv run python scripts/debug/export_snapshots.py data.processed_dir=data/processed/mimic-iv +# uv run python scripts/debug/export_snapshots.py data.processed_dir=data/processed/miiv # # # Override RICU data paths: # uv run python scripts/debug/export_snapshots.py \ @@ -21,6 +21,8 @@ defaults: - data: ricu - _self_ +dataset: miiv + # Output directory for debug artifacts (defaults to processed_dir/debug_snapshots) debug_output_dir: null @@ -41,7 +43,7 @@ max_samples: 5000 # Max samples for embedding extraction # Output options plots: false # Generate visualization plots -max_hours: 48 # Max hours for timeseries export +max_hours: 24 # Max hours for timeseries export flatten_timeseries: true # Convert nested arrays to long format # General diff --git a/configs/eval/default.yaml b/configs/eval/default.yaml index 293f114..0d0ce60 100644 --- a/configs/eval/default.yaml +++ b/configs/eval/default.yaml @@ -8,7 +8,7 @@ # brier_score, ece (expected calibration error) # multiclass: auroc, auprc, accuracy, f1, precision, recall # multilabel: auroc, auprc, accuracy, f1 -# regression: (not yet implemented) +# regression: mse, mae, r2 # # Clinical metrics explanation: # - auroc: Area under ROC curve (discrimination) @@ -21,13 +21,19 @@ # Metric configuration metrics: # Metrics to compute (null = use defaults for task type) - # Defaults: binary -> [auroc, auprc, brier_score], multiclass -> [auroc, accuracy] + # Defaults: + # binary -> [auroc, auprc, accuracy, f1, precision, recall, specificity, + # brier_score, ece] + # multiclass -> [auroc, accuracy] + # regression -> [mse, mae, r2] names: null # Decision threshold for binary classification threshold: 0.5 -# Fairness analysis configuration (not yet implemented) +# Fairness analysis configuration +# Enable in-training fairness eval via eval.fairness.enabled=true, or use +# scripts/eval/evaluate_fairness.py for post-run batch evaluation. # fairness: # enabled: false # protected_attributes: diff --git a/configs/evaluate_imputation.yaml b/configs/evaluate_imputation.yaml index 7658c87..9c19b5e 100644 --- a/configs/evaluate_imputation.yaml +++ b/configs/evaluate_imputation.yaml @@ -5,22 +5,24 @@ # Example usage: # uv run python scripts/eval/evaluate_imputation.py \ # checkpoint=outputs/encoder.pt \ -# data.processed_dir=data/processed/mimic-iv +# data.processed_dir=data/processed/miiv defaults: - data: ricu - _self_ +dataset: miiv + # Checkpoint path (one of these must be provided) checkpoint: null # encoder.pt file -pretrain_checkpoint: null # full pretrain .ckpt (includes MAE decoder) +pretrain_checkpoint: null # full pretrain .ckpt; evaluation still trains a probe decoder # Masking configuration masking: strategies: [random, feature_block, temporal_block] mask_ratio: 0.15 -# Reconstruction head (for non-MAE models only) +# Reconstruction head probe (trained for all checkpoint types) reconstruction_head: max_epochs: 10 lr: 1.0e-3 diff --git a/configs/finetune.yaml b/configs/finetune.yaml index b317ba2..2588f34 100644 --- a/configs/finetune.yaml +++ b/configs/finetune.yaml @@ -1,14 +1,15 @@ # Downstream Task Finetuning Configuration # # This config orchestrates finetuning a pretrained encoder on downstream tasks. -# By default, the encoder is frozen and only the task head is trained (linear probing). +# By default, the encoder is frozen and only the task head is trained. Use +# protocol=linear_probe or protocol=full_finetune for benchmark-safe manual launches. # # Example usage: -# # Linear probing (frozen encoder, default) -# uv run python scripts/training/finetune.py dataset=miiv checkpoint=outputs/encoder.pt +# # Linear probe (frozen encoder, linear head, 50 epochs, lr=1e-4) +# uv run python scripts/training/finetune.py dataset=miiv checkpoint=outputs/encoder.pt protocol=linear_probe # -# # Full finetuning (train encoder + head) -# uv run python scripts/training/finetune.py dataset=miiv checkpoint=outputs/encoder.pt training.freeze_encoder=false +# # Full finetune (unfrozen encoder, MLP head, 100 epochs, lr=3e-4) +# uv run python scripts/training/finetune.py dataset=miiv checkpoint=outputs/encoder.pt protocol=full_finetune # # # Gradual unfreezing (freeze for N epochs, then unfreeze) # uv run python scripts/training/finetune.py dataset=miiv checkpoint=outputs/encoder.pt training.unfreeze_epoch=5 @@ -21,9 +22,10 @@ defaults: - data: ricu - - model: observation_transformer + - model: transformer - tasks: mortality_24h - _self_ + - optional protocol: null # Project info project_name: slices @@ -35,28 +37,33 @@ dataset: miiv # miiv | eicu | combined # Paradigm — auto-detected from checkpoint in finetune.py; override with paradigm=jepa etc. paradigm: mae -# Experiment sprint (for W&B filtering, e.g. sprint=1) -sprint: null +# Final rerun corpus metadata +experiment_class: null +experiment_subtype: null +revision: null +rerun_reason: null # Random seed for reproducibility seed: 42 -# Checkpoint path (encoder weights from pretraining) -# Can be either: +# Checkpoint path (encoder weights from pretraining). Provide exactly one: # - encoder.pt file: checkpoint=outputs/encoder.pt # - full .ckpt file: pretrain_checkpoint=outputs/ssl-epoch=099.ckpt checkpoint: null pretrain_checkpoint: null +# Resume this downstream trainer from an interrupted run. +ckpt_path: null + # Output directory (Hydra will create timestamped subdirs) output_dir: ${hydra:runtime.output_dir} checkpoint_dir: ${output_dir}/checkpoints -# Encoder configuration inherited from configs/model/observation_transformer.yaml +# Encoder configuration inherited from configs/model/transformer.yaml # For SSL finetuning, the v3 checkpoint overrides this with the saved encoder config. # Task configuration inherited from configs/tasks/*.yaml -# Override with: task=mortality_hospital, task=los_remaining, etc. +# Override with: tasks=mortality_hospital, tasks=los_remaining, etc. # Label fraction for label-efficiency ablations # 1.0 = use all training data (default) @@ -80,11 +87,15 @@ training: # Use learned missing token from pretraining use_missing_token: true + # If best-checkpoint reload fails before test-time evaluation, optionally + # fall back to the final in-memory model instead of failing closed. + allow_best_ckpt_fallback: false + # Class weighting for imbalanced tasks # null: no weighting (default) # "balanced": auto-compute inverse frequency weights from training labels # [w0, w1, ...]: explicit per-class weights - class_weight: null + class_weight: balanced # Compute settings accelerator: auto # auto | cpu | gpu | mps @@ -96,11 +107,15 @@ training: accumulate_grad_batches: 1 # Early stopping - # Monitor/mode are set automatically by task type in setup_callbacks(): + # Monitor/mode default to task.primary_metric when provided. + # Otherwise setup_callbacks() falls back to: # classification -> val/auprc (max), regression -> val/mse (min) # Override here only if you need a non-default metric. early_stopping_patience: 10 + # Label smoothing for classification tasks + label_smoothing: 0.1 + # Debug: overfit on N batches (0 = disabled) overfit_batches: 0 @@ -109,7 +124,7 @@ training: optimizer: name: adamw lr: 1.0e-4 # Lower LR for finetuning - weight_decay: 0.01 + weight_decay: 0.05 # Optional learning rate scheduler scheduler: @@ -129,7 +144,7 @@ logging: use_wandb: true wandb_project: ${project_name} wandb_entity: null # Set to your W&B username/team - run_name: finetune_${dataset}_${paradigm}_${task.task_name}_seed${seed} + run_name: ${experiment_class}_finetune_${dataset}_${paradigm}_${task.task_name}_seed${seed} wandb_group: finetune_${dataset}_${paradigm}_${task.task_name} wandb_tags: - "phase:finetune" @@ -137,6 +152,7 @@ logging: - "paradigm:${paradigm}" - "task:${task.task_name}" - "seed:${seed}" + - "lr:${optimizer.lr}" log_every_n_steps: 10 # Hydra configuration diff --git a/configs/gru_d.yaml b/configs/gru_d.yaml new file mode 100644 index 0000000..384d265 --- /dev/null +++ b/configs/gru_d.yaml @@ -0,0 +1,107 @@ +# GRU-D Baseline Training Configuration +# +# Trains GRU-D from scratch as a non-SSL baseline. GRU-D handles missing data +# natively via learnable decay, so no missing token is needed. +# +# Uses the same FineTuneModule infrastructure as supervised training. +# +# Example usage: +# uv run python scripts/training/supervised.py --config-name gru_d dataset=miiv +# uv run python scripts/training/supervised.py --config-name gru_d dataset=eicu tasks=mortality_hospital + +defaults: + - data: ricu + - model: gru_d + - tasks: mortality_24h + - _self_ + +# Project info +project_name: slices +experiment_name: gru_d + +# Dataset identifier +dataset: miiv + +# Paradigm +paradigm: gru_d +protocol: full_finetune + +# Final rerun corpus metadata +experiment_class: null +experiment_subtype: null +revision: null +rerun_reason: null + +# Random seed +seed: 42 + +# Resume from checkpoint +ckpt_path: null + +# Output directory +output_dir: ${hydra:runtime.output_dir} +checkpoint_dir: ${output_dir}/checkpoints + +# Label fraction for label-efficiency ablations +label_fraction: 1.0 + +# Training configuration +training: + max_epochs: 100 + batch_size: 64 + freeze_encoder: false + use_missing_token: false + accelerator: auto + devices: auto + precision: 32 + class_weight: balanced + gradient_clip_val: 1.0 + accumulate_grad_batches: 1 + early_stopping_patience: 10 + label_smoothing: 0.1 + overfit_batches: 0 + +# Optimizer +optimizer: + name: adamw + lr: 3.0e-4 + weight_decay: 0.05 + +# Scheduler +scheduler: + name: cosine + T_max: ${training.max_epochs} + eta_min: 1.0e-6 + +# Evaluation +eval: + fairness: + enabled: false + protected_attributes: [gender, age_group, race] + min_subgroup_size: 50 + +# Logging +logging: + use_wandb: true + wandb_project: ${project_name} + wandb_entity: null + run_name: ${experiment_class}_gru_d_${dataset}_${task.task_name}_seed${seed} + wandb_group: gru_d_${dataset}_${task.task_name} + wandb_tags: + - "phase:baseline" + - "dataset:${dataset}" + - "paradigm:gru_d" + - "task:${task.task_name}" + - "seed:${seed}" + - "lr:${optimizer.lr}" + log_every_n_steps: 10 + +# Hydra +hydra: + run: + dir: outputs/gru_d/${now:%Y-%m-%d}/${now:%H-%M-%S} + sweep: + dir: multirun/gru_d/${now:%Y-%m-%d}/${now:%H-%M-%S} + subdir: ${hydra.job.num} + +deterministic: false diff --git a/configs/model/gru_d.yaml b/configs/model/gru_d.yaml new file mode 100644 index 0000000..41f3e4c --- /dev/null +++ b/configs/model/gru_d.yaml @@ -0,0 +1,13 @@ +# @package encoder +# GRU-D Encoder Configuration +# +# GRU-D (Che et al. 2018) is a standard RNN baseline for clinical time series +# with missing data. Uses learnable decay to handle missingness natively. +# No EncoderWithMissingToken wrapper needed. + +name: gru_d +d_model: 64 +n_layers: 1 +n_gru_layers: 1 +dropout: 0.1 +max_seq_length: 168 diff --git a/configs/model/observation_transformer.yaml b/configs/model/observation_transformer.yaml index b120b09..67649f8 100644 --- a/configs/model/observation_transformer.yaml +++ b/configs/model/observation_transformer.yaml @@ -14,7 +14,7 @@ name: observation_transformer # Input/Output dimensions d_model: 128 -max_seq_length: 168 # Default; overridden at runtime by datamodule.get_seq_length() (48h) +max_seq_length: 168 # Default; overridden at runtime by datamodule.get_seq_length() (24h benchmark) # Transformer architecture n_layers: 4 diff --git a/configs/model/transformer.yaml b/configs/model/transformer.yaml index 1d4d000..2854cb1 100644 --- a/configs/model/transformer.yaml +++ b/configs/model/transformer.yaml @@ -5,15 +5,8 @@ # Training configs (supervised.yaml, pretrain.yaml, finetune.yaml) inherit # from this via Hydra defaults and override only what differs. # -# Common configurations: -# -# SMART-like (default, ~100K parameters): -# encoder.d_model=32 encoder.n_layers=2 encoder.n_heads=4 encoder.d_ff=128 -# Matches SMART (NeurIPS 2024) architecture for missing-aware EHR modeling. -# Prevents overfitting on small positive classes typical in ICU prediction. -# -# Medium (~500K parameters): -# encoder.d_model=128 encoder.n_layers=4 encoder.n_heads=8 encoder.d_ff=512 +# All paradigms (SSL and supervised) share the same architecture (~174K params +# with obs_aware=True) to eliminate model capacity as a confounding variable. name: transformer @@ -21,7 +14,7 @@ name: transformer # Input/Output dimensions # d_input is typically set automatically from data config d_model: 64 # Model dimension (embedding size) -max_seq_length: 168 # Default; overridden at runtime by datamodule.get_seq_length() (48h) +max_seq_length: 168 # Default; overridden at runtime by datamodule.get_seq_length() (24h benchmark) # Transformer architecture n_layers: 2 # Number of transformer encoder layers @@ -45,3 +38,7 @@ use_positional_encoding: true # Add sinusoidal positional encoding # - last: Use last valid timestep # - none: Return per-timestep embeddings (required for MAE pretraining) pooling: mean + +# Observation-aware mode: concat (values, obs_mask) -> MLP -> d_model +# Enabled for all paradigms (SSL and supervised) for fair comparison +obs_aware: true diff --git a/configs/model/transformer_large.yaml b/configs/model/transformer_large.yaml new file mode 100644 index 0000000..b1fa545 --- /dev/null +++ b/configs/model/transformer_large.yaml @@ -0,0 +1,28 @@ +# @package encoder +# Transformer Encoder — Large (256-dim, 4-layer) +# +# Scaled-up variant for the capacity ablation pilot. +# ~3.2M params vs ~112K for the base transformer. +# All other settings identical to transformer.yaml. + +name: transformer + +# Input/Output dimensions +d_model: 256 +max_seq_length: 168 # Overridden at runtime by datamodule + +# Transformer architecture +n_layers: 4 +n_heads: 8 # 256 / 8 = 32 per head +d_ff: 1024 # 4 * d_model + +# Regularization +dropout: 0.1 + +# Architecture choices (same as base) +activation: gelu +layer_norm_eps: 1.0e-5 +prenorm: true +use_positional_encoding: true +pooling: mean +obs_aware: true diff --git a/configs/model/transformer_medium.yaml b/configs/model/transformer_medium.yaml new file mode 100644 index 0000000..d4aa114 --- /dev/null +++ b/configs/model/transformer_medium.yaml @@ -0,0 +1,28 @@ +# @package encoder +# Transformer Encoder — Medium (128-dim, 4-layer) +# +# Scaled-up variant for the capacity ablation pilot. +# ~816K params vs ~112K for the base transformer. +# All other settings identical to transformer.yaml. + +name: transformer + +# Input/Output dimensions +d_model: 128 +max_seq_length: 168 # Overridden at runtime by datamodule + +# Transformer architecture +n_layers: 4 +n_heads: 8 # 128 / 8 = 16 per head +d_ff: 512 # 4 * d_model + +# Regularization +dropout: 0.1 + +# Architecture choices (same as base) +activation: gelu +layer_norm_eps: 1.0e-5 +prenorm: true +use_positional_encoding: true +pooling: mean +obs_aware: true diff --git a/configs/pretrain.yaml b/configs/pretrain.yaml index 5e1d0c7..2b4dd06 100644 --- a/configs/pretrain.yaml +++ b/configs/pretrain.yaml @@ -12,7 +12,7 @@ defaults: - data: ricu - - model: observation_transformer + - model: transformer - ssl: mae - _self_ @@ -22,8 +22,11 @@ project_name: slices # Dataset identifier — drives data paths via ricu.yaml interpolation dataset: miiv # miiv | eicu | combined -# Experiment sprint (for W&B filtering, e.g. sprint=1) -sprint: null +# Final rerun corpus metadata +experiment_class: null +experiment_subtype: null +revision: null +rerun_reason: null # Random seed for reproducibility seed: 42 @@ -32,6 +35,9 @@ seed: 42 output_dir: ${hydra:runtime.output_dir} checkpoint_dir: ${output_dir}/checkpoints +# Resume from an interrupted trainer checkpoint. +ckpt_path: null + # Data configuration overrides # Base settings inherited from configs/data/ricu.yaml data: @@ -41,29 +47,32 @@ data: # Sliding window configuration for longer sequences # Enable after re-extracting with longer seq_length_hours (e.g., 168h) enable_sliding_windows: false # Set true after re-extraction - window_size: 48 # Window size in hours (48h = typical ICU prediction horizon) - window_stride: 24 # Stride for overlapping windows (50% overlap) + window_size: 24 # Match the benchmark observation window when enabled + window_stride: 12 # 50% overlap for 24h windows -# Encoder overrides (base from configs/model/observation_transformer.yaml: d=128, L=4, H=8) +# Encoder overrides for SSL pretraining +# Architecture (d_model, n_layers, etc.) inherited from configs/model/transformer.yaml encoder: pooling: none # Required for SSL - need per-token outputs (vs 'mean' for tasks) -# Training configuration (same budget across thesis paradigms for fair comparison) +# Training configuration (same budget across paradigms for fair comparison) training: - max_epochs: 500 + max_epochs: 100 batch_size: 256 # Compute settings accelerator: auto # auto | cpu | gpu | mps devices: auto # auto | 1 | [0,1] | etc. - precision: 32 # 32 | 16 | bf16 + precision: 32 # 32 | 16-mixed | bf16-mixed # Regularization gradient_clip_val: 1.0 - accumulate_grad_batches: 1 + accumulate_grad_batches: 1 # single batch, no accumulation - # Early stopping - early_stopping_patience: 10 + # No early stopping for pretraining: self-distillation methods (JEPA) have + # rising val loss as representations get richer, making val loss unreliable + # as a stopping criterion. Fixed-schedule training follows I-JEPA, DINO, BYOL. + early_stopping_patience: null # Debug: overfit on N batches (0 = disabled) overfit_batches: 0 @@ -77,7 +86,7 @@ optimizer: # Optional learning rate scheduler scheduler: name: warmup_cosine # cosine | step | plateau | warmup_cosine - warmup_epochs: 50 + warmup_epochs: 10 max_epochs: ${training.max_epochs} eta_min: 1.0e-6 @@ -87,13 +96,14 @@ logging: use_wandb: true wandb_project: ${project_name} wandb_entity: null # Set to your W&B username/team - run_name: pretrain_${dataset}_${ssl.name}_seed${seed} + run_name: ${experiment_class}_pretrain_${dataset}_${ssl.name}_seed${seed} wandb_group: pretrain_${dataset}_${ssl.name} wandb_tags: - "phase:pretrain" - "dataset:${dataset}" - "paradigm:${ssl.name}" - "seed:${seed}" + - "lr:${optimizer.lr}" log_every_n_steps: 50 # Hydra configuration diff --git a/configs/protocol/full_finetune.yaml b/configs/protocol/full_finetune.yaml new file mode 100644 index 0000000..f68be34 --- /dev/null +++ b/configs/protocol/full_finetune.yaml @@ -0,0 +1,21 @@ +# @package _global_ +# Full finetune: train encoder and downstream head for practical utility. + +protocol: full_finetune + +training: + freeze_encoder: false + max_epochs: 100 + early_stopping_patience: 10 + +optimizer: + lr: 3.0e-4 + +scheduler: + T_max: ${training.max_epochs} + +task: + head_type: mlp + hidden_dims: [64] + dropout: 0.3 + activation: relu diff --git a/configs/protocol/linear_probe.yaml b/configs/protocol/linear_probe.yaml new file mode 100644 index 0000000..fd7150c --- /dev/null +++ b/configs/protocol/linear_probe.yaml @@ -0,0 +1,20 @@ +# @package _global_ +# Linear probe: frozen encoder, linear head, representation-quality evaluation. + +protocol: linear_probe + +training: + freeze_encoder: true + max_epochs: 50 + early_stopping_patience: 10 + +optimizer: + lr: 1.0e-4 + +scheduler: + T_max: ${training.max_epochs} + +task: + head_type: linear + hidden_dims: [] + dropout: 0.0 diff --git a/configs/ssl/contrastive.yaml b/configs/ssl/contrastive.yaml index 06714ec..e8ef1ff 100644 --- a/configs/ssl/contrastive.yaml +++ b/configs/ssl/contrastive.yaml @@ -1,21 +1,29 @@ # @package ssl # Contrastive (SimCLR-style) SSL Configuration # -# Observation-level tokenization: two different random masks create two "views" -# of the same sample. NT-Xent loss pulls views of the same sample together -# and pushes different samples apart in a learned projection space. +# Timestep-level tokenization: two timestep masks create two "views" of the +# same sample. In the default instance mode, NT-Xent pulls mean-pooled sequence +# embeddings together and pushes different samples apart in projection space. # # Key difference from MAE: discriminative (not reconstructive) # Key difference from JEPA: global invariance (not local positional prediction) name: contrastive -# Masking parameters (same as MAE/JEPA for fair comparison) +# Mode: "temporal" (per-timestep overlap pairs) or "instance" (mean-pool, SimCLR-style) +mode: instance + +# Masking budget matches the core MAE/JEPA default; the two-view strategy is +# objective-specific. mask_ratio: 0.5 # Experiment plan default (ablate over 0.3/0.5/0.75) # Projection head parameters -proj_hidden_dim: 512 # Hidden dimension of projection MLP -proj_output_dim: 128 # Output dimension (contrastive embedding space) +proj_hidden_dim: 256 # Hidden dimension of projection MLP (4 * encoder d_model) +proj_output_dim: 64 # Output dimension (contrastive embedding space, = encoder d_model) + +# NT-Xent temperature (0.07 suits instance-level with ~511 negatives at B=256) +temperature: 0.07 -# NT-Xent temperature -temperature: 0.1 +# Complementary masks: view 2 = ~view 1, ensuring zero overlap between views. +# Forces encoder to learn abstract semantics from non-overlapping windows. +complementary_masks: true diff --git a/configs/ssl/jepa.yaml b/configs/ssl/jepa.yaml index 464a883..3c0031a 100644 --- a/configs/ssl/jepa.yaml +++ b/configs/ssl/jepa.yaml @@ -1,29 +1,37 @@ # @package ssl # JEPA (Joint-Embedding Predictive Architecture) SSL Configuration # -# Observation-level tokenization: same masking as MAE, but predicts -# latent representations (d_model vectors) instead of raw input values. -# Uses EMA target encoder for stable training targets. +# Timestep-level tokenization: same encoder/tokenization interface and mask +# budget as MAE, but predicts latent timestep representations (d_model vectors) +# instead of raw features. +# Uses an EMA target encoder to produce stable teacher targets. # # Key difference from MAE: latent-space prediction (not input-space) # Key difference from Contrastive: local positional prediction (not global invariance) name: jepa -# Masking parameters (same as MAE for fair comparison) +# Masking parameters mask_ratio: 0.5 # Experiment plan default (ablate over 0.3/0.5/0.75) +# Block masking is intentional for JEPA: random timestep masks let the predictor +# interpolate from adjacent visible timesteps and previously caused collapse. +mask_strategy: block # "block" for contiguous segments, "random" for per-timestep +mask_n_blocks: 3 # Number of contiguous blocks (only for block strategy) -# Predictor parameters (mirrors MAE decoder for fairness) -predictor_d_model: 128 +# Predictor parameters — bottlenecked relative to encoder (50% d_model). +# Unlike MAE decoder (which reconstructs raw features, d_model→D), JEPA predictor +# predicts in latent space (d_model→d_model). The bottleneck forces information +# routing through the online encoder's context representations. +predictor_d_model: 32 # 50% of encoder d_model=64 predictor_n_layers: 2 -predictor_n_heads: 4 -predictor_d_ff: 512 +predictor_n_heads: 4 # 32/4 = 8 per head +predictor_d_ff: 128 # 4 * predictor_d_model predictor_dropout: 0.1 # Momentum encoder parameters # Target encoder updated via EMA: target = m * target + (1-m) * online # Momentum increases linearly from base to final during training -momentum_base: 0.996 +momentum_base: 0.999 momentum_final: 1.0 # Loss type: mse or cosine diff --git a/configs/ssl/mae.yaml b/configs/ssl/mae.yaml index a30afe9..c6fb5d0 100644 --- a/configs/ssl/mae.yaml +++ b/configs/ssl/mae.yaml @@ -1,19 +1,19 @@ # @package ssl # MAE (Masked Autoencoder) SSL Configuration # -# Observation-level tokenization: each token = one observed (timestep, feature, value). -# Encoder processes only visible tokens (25%). Decoder reconstructs masked tokens. -# No MISSING_TOKEN needed -- missingness is handled intrinsically by only -# tokenizing observed values. +# Timestep-level tokenization: each token = one hourly timestep containing all +# features plus the observation mask via the encoder's obs-aware projection. +# The encoder processes only visible timesteps; the decoder reconstructs masked +# timestep features from the remaining context. name: mae # Masking parameters mask_ratio: 0.5 # Experiment plan default (ablate over 0.3/0.5/0.75) -# Decoder parameters (lighter than encoder) -decoder_d_model: 128 # Decoder hidden dimension +# Decoder parameters (matches encoder d_model) +decoder_d_model: 64 # Decoder hidden dimension (= encoder d_model) decoder_n_layers: 2 # Number of decoder transformer layers decoder_n_heads: 4 # Number of decoder attention heads -decoder_d_ff: 512 # Decoder feedforward dimension +decoder_d_ff: 256 # Decoder feedforward dimension (4 * decoder_d_model) decoder_dropout: 0.1 diff --git a/configs/ssl/ts2vec.yaml b/configs/ssl/ts2vec.yaml new file mode 100644 index 0000000..7ac06a3 --- /dev/null +++ b/configs/ssl/ts2vec.yaml @@ -0,0 +1,35 @@ +# @package ssl +# TS2Vec-style Temporal Contrastive SSL Configuration +# +# Addresses the augmentation limitation of the standard contrastive objective +# by combining timestamp masking with input-level Gaussian noise and +# hierarchical temporal contrastive loss. +# +# Key differences from standard contrastive: +# - Input-level noise augmentation (not just masking) +# - Independent masks (not complementary) +# - Per-timestep temporal contrastive loss (not instance-level) +# - Hierarchical max-pool over multiple temporal scales +# +# Reference: Yue et al., "TS2Vec", AAAI 2022. + +name: ts2vec + +# Masking (same default for comparability with other objectives) +mask_ratio: 0.5 + +# Input-level augmentation +noise_scale: 0.01 # Gaussian noise std applied to observed values (0 = masking only) +crop_ratio: 1.0 # Sub-sequence crop ratio (1.0 = no crop, 0.5 = half-length) + +# Projection head parameters (match standard contrastive for fair comparison) +proj_hidden_dim: 256 # Hidden dimension of per-timestep projection MLP +proj_output_dim: 64 # Output dimension (contrastive embedding space) + +# Temperature for temporal NT-Xent (lower than instance-level because the +# per-timestep similarity matrix is larger and sparser) +temperature: 0.05 + +# Hierarchical temporal pooling: number of max-pool scales [1, 2, 4, ...] +# 4 scales at T=48 gives pool sizes [1, 2, 4, 8] +n_hierarchical_scales: 4 diff --git a/configs/supervised.yaml b/configs/supervised.yaml index aa59d20..0c8ae9a 100644 --- a/configs/supervised.yaml +++ b/configs/supervised.yaml @@ -32,9 +32,13 @@ dataset: miiv # miiv | eicu | combined # Paradigm — always "supervised" for this config paradigm: supervised +protocol: full_finetune -# Experiment sprint (for W&B filtering, e.g. sprint=1) -sprint: null +# Final rerun corpus metadata +experiment_class: null +experiment_subtype: null +revision: null +rerun_reason: null # Random seed for reproducibility seed: 42 @@ -47,9 +51,10 @@ output_dir: ${hydra:runtime.output_dir} checkpoint_dir: ${output_dir}/checkpoints # Encoder configuration inherited from configs/model/transformer.yaml +# Same architecture as SSL paradigms (d=64, L=2, H=4, obs_aware=True) for fair comparison # Task configuration inherited from configs/tasks/*.yaml -# Override with: task=mortality_hospital, task=los_remaining, etc. +# Override with: tasks=mortality_hospital, tasks=los_remaining, etc. # Label fraction for label-efficiency ablations # 1.0 = use all training data (default) @@ -67,6 +72,10 @@ training: # and a random one would confound SSL vs supervised comparison) use_missing_token: false + # If best-checkpoint reload fails before test-time evaluation, optionally + # fall back to the final in-memory model instead of failing closed. + allow_best_ckpt_fallback: false + # Compute settings accelerator: auto # auto | cpu | gpu | mps devices: auto @@ -76,17 +85,21 @@ training: # null: no weighting (default) # "balanced": auto-compute inverse frequency weights from training labels # [w0, w1, ...]: explicit per-class weights - class_weight: null + class_weight: balanced # Regularization gradient_clip_val: 1.0 accumulate_grad_batches: 1 # Early stopping - # Monitor/mode are set automatically by task type in setup_callbacks(): + # Monitor/mode default to task.primary_metric when provided. + # Otherwise setup_callbacks() falls back to: # classification -> val/auprc (max), regression -> val/mse (min) # Override here only if you need a non-default metric. - early_stopping_patience: 20 # More patience for training from scratch + early_stopping_patience: 10 + + # Label smoothing for classification tasks + label_smoothing: 0.1 # Debug: overfit on N batches (0 = disabled) overfit_batches: 0 @@ -95,8 +108,8 @@ training: # Higher learning rate than finetuning since we're training from scratch optimizer: name: adamw - lr: 1.0e-3 # Higher LR for training from scratch - weight_decay: 0.01 + lr: 3.0e-4 + weight_decay: 0.05 # Learning rate scheduler scheduler: @@ -116,7 +129,7 @@ logging: use_wandb: true wandb_project: ${project_name} wandb_entity: null - run_name: supervised_${dataset}_${task.task_name}_seed${seed} + run_name: ${experiment_class}_supervised_${dataset}_${task.task_name}_seed${seed} wandb_group: supervised_${dataset}_${task.task_name} wandb_tags: - "phase:supervised" @@ -124,6 +137,7 @@ logging: - "paradigm:supervised" - "task:${task.task_name}" - "seed:${seed}" + - "lr:${optimizer.lr}" log_every_n_steps: 10 # Hydra configuration diff --git a/configs/tasks/aki_kdigo.yaml b/configs/tasks/aki_kdigo.yaml index c2b5e11..e953bbc 100644 --- a/configs/tasks/aki_kdigo.yaml +++ b/configs/tasks/aki_kdigo.yaml @@ -1,13 +1,31 @@ # @package task # AKI prediction using KDIGO criteria (binary classification) -# Predict if patient will develop acute kidney injury +# Predict if patient will develop acute kidney injury after the 24h observation window task_name: aki_kdigo task_type: binary +observation_window_hours: 24 +prediction_window_hours: 24 +gap_hours: 0 +label_sources: + - stays + - timeseries +label_params: + creatinine_col: crea + baseline_window_hours: 24 + absolute_rise_threshold: 0.3 + relative_rise_threshold: 1.5 + relative_window_hours: 168 +quality_checks: + # AKI is evaluated on the creatinine-observed filtered cohort. eICU has + # materially higher missingness than MIMIC, so warn only above the current + # expected cross-dataset envelope. + max_missing_percentage: 35.0 +primary_metric: auprc n_classes: null head_type: mlp hidden_dims: [64] -dropout: 0.1 +dropout: 0.3 activation: relu use_layer_norm: false projection_dim: null diff --git a/configs/tasks/los_remaining.yaml b/configs/tasks/los_remaining.yaml index a7168a0..464b817 100644 --- a/configs/tasks/los_remaining.yaml +++ b/configs/tasks/los_remaining.yaml @@ -1,13 +1,20 @@ # @package task # Remaining length of stay prediction (regression) -# Predict remaining days in ICU after the observation window +# Predict remaining days in ICU after the 24h observation window task_name: los_remaining task_type: regression +observation_window_hours: 24 +prediction_window_hours: null +gap_hours: 0 +label_sources: + - stays +label_params: {} +primary_metric: mae n_classes: null head_type: mlp hidden_dims: [64] -dropout: 0.1 +dropout: 0.3 activation: relu use_layer_norm: false projection_dim: null diff --git a/configs/tasks/mortality.yaml b/configs/tasks/mortality.yaml new file mode 100644 index 0000000..44a23d9 --- /dev/null +++ b/configs/tasks/mortality.yaml @@ -0,0 +1,21 @@ +# @package task +# ICU mortality prediction (binary classification) +# Predict if patient will die during the remainder of the ICU stay after the 24h observation window + +task_name: mortality +task_type: binary +observation_window_hours: 24 +prediction_window_hours: -1 +gap_hours: 0 +label_sources: + - stays + - mortality_info +label_params: {} +primary_metric: auprc +n_classes: null +head_type: mlp +hidden_dims: [64] +dropout: 0.3 +activation: relu +use_layer_norm: false +projection_dim: null diff --git a/configs/tasks/mortality_24h.yaml b/configs/tasks/mortality_24h.yaml index cd21f21..37d64c8 100644 --- a/configs/tasks/mortality_24h.yaml +++ b/configs/tasks/mortality_24h.yaml @@ -1,13 +1,24 @@ # @package task # 24-hour mortality prediction (binary classification) # Predict if patient will die within 24h after the observation window +# +# Timeline: |---- observation (24h) ----|---- prediction (24h) ----| +# intime obs_end pred_end task_name: mortality_24h task_type: binary +observation_window_hours: 24 +prediction_window_hours: 24 +gap_hours: 0 +label_sources: + - stays + - mortality_info +label_params: {} +primary_metric: auprc n_classes: null head_type: mlp hidden_dims: [64] -dropout: 0.1 +dropout: 0.3 activation: relu use_layer_norm: false projection_dim: null diff --git a/configs/tasks/mortality_hospital.yaml b/configs/tasks/mortality_hospital.yaml index 37caa8c..0f3847e 100644 --- a/configs/tasks/mortality_hospital.yaml +++ b/configs/tasks/mortality_hospital.yaml @@ -1,13 +1,21 @@ # @package task # Hospital mortality prediction (binary classification) -# Predict if patient will die before hospital discharge +# Predict if patient will die before hospital discharge after the 24h observation window task_name: mortality_hospital task_type: binary +observation_window_hours: 24 +prediction_window_hours: null +gap_hours: 0 +label_sources: + - stays + - mortality_info +label_params: {} +primary_metric: auprc n_classes: null head_type: mlp hidden_dims: [64] -dropout: 0.1 +dropout: 0.3 activation: relu use_layer_norm: false projection_dim: null diff --git a/configs/xgboost.yaml b/configs/xgboost.yaml new file mode 100644 index 0000000..a0a6009 --- /dev/null +++ b/configs/xgboost.yaml @@ -0,0 +1,86 @@ +# XGBoost Baseline Configuration +# +# Non-neural baseline using hand-crafted tabular features from ICU time series. +# Reuses ICUDataModule for data/splits and logs to W&B. +# +# Example usage: +# uv run python scripts/training/xgboost_baseline.py dataset=miiv +# uv run python scripts/training/xgboost_baseline.py dataset=eicu tasks=mortality_hospital + +defaults: + - data: ricu + - tasks: mortality_24h + - _self_ + +# Project info +project_name: slices +experiment_name: xgboost + +# Dataset identifier +dataset: miiv + +# Paradigm +paradigm: xgboost +protocol: full_finetune + +# Final rerun corpus metadata +experiment_class: null +experiment_subtype: null +revision: null +rerun_reason: null + +# Random seed +seed: 42 + +# Output directory +output_dir: ${hydra:runtime.output_dir} + +# Label fraction for label-efficiency ablations +label_fraction: 1.0 + +# XGBoost hyperparameters +xgboost: + n_estimators: 1000 + max_depth: 6 + learning_rate: 0.1 + subsample: 0.8 + colsample_bytree: 0.8 + min_child_weight: 1 + early_stopping_rounds: 20 + # "balanced" resolves to XGBoost's native binary ratio: n_negative / n_positive. + # This is intentionally stronger than neural sqrt-balanced class weights. + scale_pos_weight: balanced + n_jobs: 4 # cap per-process CPU use when many baselines run in parallel + +# Data loading +training: + batch_size: 256 # for feature extraction only + +# Evaluation +eval: + fairness: + enabled: false + protected_attributes: [gender, age_group, race] + min_subgroup_size: 50 + +# Logging +logging: + use_wandb: true + wandb_project: ${project_name} + wandb_entity: null + run_name: ${experiment_class}_xgboost_${dataset}_${task.task_name}_seed${seed} + wandb_group: xgboost_${dataset}_${task.task_name} + wandb_tags: + - "phase:baseline" + - "dataset:${dataset}" + - "paradigm:xgboost" + - "task:${task.task_name}" + - "seed:${seed}" + +# Hydra +hydra: + run: + dir: outputs/xgboost/${now:%Y-%m-%d}/${now:%H-%M-%S} + sweep: + dir: multirun/xgboost/${now:%Y-%m-%d}/${now:%H-%M-%S} + subdir: ${hydra.job.num} diff --git a/docs/EVAL_FAIRNESS.md b/docs/EVAL_FAIRNESS.md new file mode 100644 index 0000000..b61ecd4 --- /dev/null +++ b/docs/EVAL_FAIRNESS.md @@ -0,0 +1,219 @@ +# Fairness Evaluation + +This document describes the current fairness pipeline in SLICES. + +It focuses on implementation status and data semantics rather than freezing a +specific numeric table. The canonical outputs are the `fairness/*` keys written +to W&B summaries and the parquet exports produced by `scripts/export_results.py`. + +## Current Scope + +- Standalone fairness evaluation runs via `scripts/eval/evaluate_fairness.py`. +- The standalone script is revision-scoped. Pass `--revision ` so reruns + do not get mixed together. +- Default benchmark fairness corpus: + - Experiment classes `core_ssl_benchmark`, `label_efficiency`, + `cross_dataset_transfer`, `hp_robustness`, `capacity_study`, + `classical_baselines`, `ts2vec_extension`, and + `smart_external_reference` + - Phases `finetune`, `supervised`, and `baseline` +- The script re-runs inference from the checkpoint provenance recorded by the + original training run, then writes fairness metrics back to that same W&B run. +- Existing fairness summaries are skipped only when they contain current + metadata for the evaluated artifact path/source, checkpoint provenance, + script version, protected attributes, and minimum subgroup size. Otherwise + they are recomputed. + +## Protected Attributes + +| Attribute | Availability | Current behavior | +|---|---|---| +| `gender` | All datasets | Encoded from the static table; unknown values are excluded | +| `age_group` | All datasets | Age is binned into `18-44`, `45-64`, `65-79`, `80+` | +| `race` | MIMIC rows only | Raw race strings are canonicalized into `White`, `Black`, `Hispanic`, `Asian`, `Other` | + +Important details: + +- eICU race is excluded from the fairness sweep. +- For the combined dataset, race fairness is computed only on rows whose + `source_dataset` indicates MIMIC/MIIV provenance. +- Missing demographic values are mapped to `unknown` and excluded from the + fairness computations. +- Minimum subgroup size is enforced on unique patients, not stays. + +## Pipeline + +### 1. Metric primitives + +`src/slices/eval/fairness.py` provides the pure numerical metrics: + +- `demographic_parity_difference` +- `equalized_odds_difference` +- `disparate_impact_ratio` + +These functions operate only on predictions, labels, and group IDs. + +### 2. Clinical-aware evaluator + +`src/slices/eval/fairness_evaluator.py` joins predictions with demographics from +`static.parquet`, canonicalizes groups, filters undersized subgroups, and +computes per-attribute reports. + +For binary tasks it logs: + +- `per_group_auroc` +- `per_group_auprc` +- `worst_group_auroc` +- `worst_group_auprc` +- `auroc_gap` +- `auprc_gap` +- `demographic_parity_diff` +- `equalized_odds_diff` +- `disparate_impact_ratio` +- `n_valid_groups` +- `n_metric_valid_groups` +- `group_sizes` +- `group_sample_sizes` + +`n_valid_groups` counts size-valid subgroups. `n_metric_valid_groups` counts +the subset with both outcome classes, which is required for comparable +AUROC/AUPRC. Worst-group AUROC/AUPRC and their gaps are written as `NaN` unless +at least two size-valid groups have both classes. + +For regression tasks it logs: + +- `per_group_mse` +- `per_group_mae` +- `per_group_r2` +- `worst_group_mse` +- `worst_group_mae` +- `mse_gap` +- `mae_gap` +- `group_sizes` +- `group_sample_sizes` + +### 3. Post-run batch evaluation + +`scripts/eval/evaluate_fairness.py`: + +1. Fetches finished W&B runs matching the requested class/paradigm/dataset, + phase, and revision filters. +2. Requires a `revision` tag so different reruns are not mixed. +3. Resolves checkpoint provenance from the run summary instead of heuristically + guessing from disk. +4. Re-runs test inference. +5. Flattens the nested report into W&B-safe `fairness/*` keys. +6. Writes `_fairness_*` metadata that identifies the artifact and evaluation + settings used to produce the summary. + +### 4. Canonical result export + +`scripts/export_results.py` is the canonical reporting path. It exports: + +- `results/per_seed_results.parquet` +- `results/aggregated_results.parquet` +- `results/statistical_tests.parquet` + +Those exports include fairness columns when the runs already have `fairness/*` +summary metrics. Publication validation requires dataset/task-appropriate +fairness summaries for downstream rows: gender and age for all datasets, and +race for MIMIC/combined rows. eICU race is not required. + +Publication validation also requires fresh `_fairness_*` metadata. Missing +metadata, old script/schema versions, changed protected attributes, changed +minimum subgroup size, artifact-source mismatches, and artifact-path mismatches +are treated as validation failures unless `--allow-incomplete` is passed for an +exploratory export. + +## Design Decisions + +1. Threshold-dependent classification fairness metrics use a fixed threshold of + `0.5`. +2. Race fairness follows the benchmark's canonical five-bin schema rather than + raw source strings. +3. Subgroup inclusion is patient-based, with the default threshold set to `50` + patients. +4. The standalone default targets the full downstream benchmark corpus, + including the canonical classical baselines. + +## How To Refresh Fairness Results + +1. Run the fairness sweep for a specific revision. For final publication, force + recomputation so stale summaries are replaced even if old `fairness/*` keys + exist: + +```bash +uv run python scripts/eval/evaluate_fairness.py \ + --project slices-thesis \ + --revision thesis-v1 \ + --entity \ + --force +``` + +2. Re-export the result tables: + +```bash +uv run python scripts/export_results.py \ + --project slices-thesis \ + --revision thesis-v1 \ + --entity +``` + +Do not pass `--allow-incomplete`, `--allow-extraction-failures`, +`--allow-duplicate-fingerprints`, or `--allow-multiple-revisions` for final +publication exports. Those flags are for exploratory cleanup only. + +## Final Publication Sequence + +After Bucket 1 and Bucket 2 are merged into the branch used for final reruns: + +1. Set the final revision name used by the launcher and export commands: + +```bash +export REVISION=thesis-v1 +``` + +2. Run or resume the final training corpus with the project launcher for that + revision. Confirm W&B contains the expected finished runs for the class-based + corpus. + +3. Recompute fairness summaries with metadata: + +```bash +uv run python scripts/eval/evaluate_fairness.py \ + --project slices-thesis \ + --revision "$REVISION" \ + --entity \ + --force +``` + +4. Export publication tables: + +```bash +uv run python scripts/export_results.py \ + --project slices-thesis \ + --revision "$REVISION" \ + --entity \ + --output-dir results +``` + +5. Run the publication reporting checks: + +```bash +uv run pytest \ + tests/test_metrics.py \ + tests/test_export_results.py \ + tests/test_fairness_evaluator.py \ + tests/test_evaluate_fairness.py +``` + +## Source Of Truth + +Use these in order: + +1. W&B run summaries with `fairness/*` +2. `results/*.parquet` generated by `scripts/export_results.py` +3. This document for implementation semantics only + +This file documents implementation semantics; use exported result tables for +numeric reporting. diff --git a/docs/internal/EXPERIMENT_PLAN.md b/docs/internal/EXPERIMENT_PLAN.md new file mode 100644 index 0000000..655ad09 --- /dev/null +++ b/docs/internal/EXPERIMENT_PLAN.md @@ -0,0 +1,279 @@ +# SLICES Final Experiment Plan + +This plan defines the final thesis rerun corpus for SLICES. The corpus is +organized by scientific experiment class, not by historical launch order. + +Final release-ready W&B runs live in the final thesis project, default +`slices-thesis`, and use a mandatory revision tag such as `thesis-v1`. +Exploratory runs remain in older projects/outputs and are not mixed into final +export or fairness jobs. + +## Benchmark Contract + +SLICES compares SSL paradigms for sparse, irregular ICU time series under +controlled conditions. + +Controlled SSL objectives: + +- `mae` +- `jepa` +- `contrastive` + +Controlled SSL objectives share: + +- same RICU-derived data pipeline +- same canonical `TransformerEncoder` +- same obs-aware timestep tokenization +- same masking budget where configs intend parity +- same training budget unless the experiment class is an explicit ablation + +Masking strategy is not forced to be identical across objectives. MAE uses +random timestep masking, JEPA uses block timestep masking to avoid the +random-mask interpolation failure observed during development, and Contrastive +uses two masked views. The controlled invariant is the mask budget/interface, not +the exact mask geometry. + +`ts2vec_extension` is a temporal-contrastive extension. It provides context for +the contrastive result but is not one of the three core thesis vertices. + +`smart_external_reference` is an appendix external reference because it uses +`SMARTEncoder` / MART and element-wise masking, not the shared controlled +encoder contract. + +Supervised Transformer, GRU-D, XGBoost, and linear views are baselines or +contextual comparisons, not SSL paradigms. + +## Fixed Benchmark Choices + +Fixed preprocessing/evaluation choices are benchmark invariants and should live +in `src/slices/constants.py` when code needs them: + +- 24 hour observation window and minimum stay +- patient-level splits, 70/15/15 +- normalization stats from train split only +- normalize then zero-fill +- preserve observation mask separately from values +- block model-input leakage from `los_hosp`, `los_icu`, and `dur_var` +- default labels: `mortality_24h`, `mortality_hospital`, `aki_kdigo`, + `los_remaining` + +The optional ICU-mortality task `mortality` is not part of the primary final +matrix. + +## Metadata Contract + +Every final run must carry class-based metadata: + +| Field | Meaning | +|---|---| +| `experiment_class` | Scientific experiment family | +| `experiment_subtype` | Optional finer-grained family, e.g. `lr_sensitivity` | +| `revision` | Rerun corpus identifier, e.g. `thesis-v1` | +| `phase` | `pretrain`, `finetune`, `supervised`, or `baseline` | +| `protocol` | `linear_probe` or `full_finetune` where applicable | +| `dataset` | `miiv`, `eicu`, or `combined` | +| `paradigm` | `mae`, `jepa`, `contrastive`, `supervised`, `gru_d`, `xgboost`, `ts2vec`, `smart` | +| `task` | Downstream task where applicable | +| `seed` | Random seed | + +Final runs must not require historical launch-group metadata. W&B tags and +exports use `experiment_class:*` as the canonical family key. + +Example run names: + +```text +core_ssl_benchmark_pretrain_miiv_mae_seed42 +core_ssl_benchmark_probe_miiv_mae_mortality_24h_seed42 +label_efficiency_finetune_eicu_jepa_mortality_24h_seed123_frac010 +hp_robustness_finetune_miiv_mae_mortality_24h_seed456_lr00005 +``` + +Exact suffixes may differ, but the experiment class must be visible without +project-history knowledge. + +## Experiment Classes + +The final launched run budget keeps the controlled benchmark structure and uses +task-specific low-label fractions where fixed random splits are scientifically +defensible. + +| Experiment class | Scope | Launched runs | +|---|---:|---:| +| `core_ssl_benchmark` | MAE, JEPA, contrastive, supervised Transformer; 3 datasets; 4 tasks; 5 seeds; Protocol A/B for SSL | 465 | +| `label_efficiency` | SSL Protocol A/B plus supervised Transformer at low-label fractions; 5 seeds | 1155 | +| `cross_dataset_transfer` | MIMIC-IV to eICU and eICU to MIMIC-IV; SSL Protocol B; 5 seeds | 120 | +| `hp_robustness` | LR robustness plus MAE/JEPA mask-ratio and contrastive view/mask sensitivity on MIMIC-IV `mortality_24h`; 5 seeds | 150 | +| `capacity_study` | Larger MAE and supervised Transformer encoders on MIIV `mortality_24h`; 5 seeds | 100 | +| `classical_baselines` | XGBoost and GRU-D full-label plus label-efficiency context; 5 seeds | 330 | +| `ts2vec_extension` | TS2Vec temporal contrastive extension; 3 datasets; 4 tasks; 5 seeds; Protocol A/B | 135 | +| `smart_external_reference` | SMART external SSL reference; 3 datasets; 4 tasks; 5 seeds; Protocol A/B | 135 | +| **Total** | Includes appendix SMART reference | **2590** | + +The formal thesis corpus excluding `smart_external_reference` has 2455 launched +runs. + +Low-label policy: + +- `mortality_24h` starts at 5% labels. The 1% fixed random subsets can + undersample positives under the final seeds. +- `mortality_hospital` gets the full 1%, 5%, 10%, 25%, 50% low-label curve. +- `aki_kdigo` and `los_remaining` remain contextual 10% low-label points. +- Capacity and classical-baseline `mortality_24h` slices also start at 5%. + +## Derived Comparison Views + +Some analysis plots need comparison rows that are not relaunched in the target +class. These are explicit export-time views, not launch-time tags: + +- label-efficiency curves include 100 percent endpoints from + `core_ssl_benchmark` +- capacity plots include default-size baselines from `core_ssl_benchmark` and + `label_efficiency` +- classical baseline context compares `classical_baselines` against neural + full-label and low-label rows +- TS2Vec comparison pairs `ts2vec_extension` against core `contrastive` + +The canonical per-seed export keeps each run in its true `experiment_class`. + +## Launching + +Main final rerun: + +```bash +uv run python scripts/internal/run_experiments.py run \ + --experiment-class \ + core_ssl_benchmark \ + label_efficiency \ + cross_dataset_transfer \ + hp_robustness \ + capacity_study \ + classical_baselines \ + ts2vec_extension \ + smart_external_reference \ + --project slices-thesis \ + --revision thesis-v1 \ + --launch-commit \ + --entity +``` + +Use the reviewed git commit hash for `--launch-commit`. The tmux launcher fills +this from `HEAD` and refuses tracked dirty launches by default; direct launcher +invocations should be equally explicit so retries cannot mix same-revision +artifacts produced by different code. + +Dry-run count check: + +```bash +uv run python scripts/internal/run_experiments.py run \ + --experiment-class \ + core_ssl_benchmark \ + label_efficiency \ + cross_dataset_transfer \ + hp_robustness \ + capacity_study \ + classical_baselines \ + ts2vec_extension \ + smart_external_reference \ + --project slices-thesis \ + --revision thesis-v1 \ + --launch-commit \ + --entity dummy \ + --dry-run +``` + +Expected dry-run count: 2590 launched runs. + +## Fairness + +Fairness is evaluated only on downstream-producing runs. Default classes: + +```text +core_ssl_benchmark +label_efficiency +cross_dataset_transfer +hp_robustness +capacity_study +classical_baselines +ts2vec_extension +smart_external_reference +``` + +Run fairness after training: + +```bash +uv run python scripts/eval/evaluate_fairness.py \ + --project slices-thesis \ + --revision thesis-v1 \ + --entity \ + --force +``` + +Revision filtering is mandatory to avoid mixing rerun corpora. Use `--force` +for final publication refreshes so old fairness summaries are replaced with +current artifact/source and evaluation-setting metadata. + +## Export + +Publication export: + +```bash +uv run python scripts/export_results.py \ + --project slices-thesis \ + --revision thesis-v1 \ + --entity \ + --output-dir results/slices-thesis_thesis-v1 +``` + +The export writes canonical per-run/aggregated tables plus derived comparison +views for label efficiency, capacity, classical context, and TS2Vec comparison. +Aggregated metric columns include seed mean, standard deviation, min/max, and +95% confidence intervals for finite seed values. + +Do not use publication escape hatches such as `--allow-incomplete`, +`--allow-extraction-failures`, `--allow-duplicate-fingerprints`, or +`--allow-multiple-revisions` for final tables. + +## Checkpoint Policy + +SSL downstream runs use `encoder.pt`, the last encoder from the fixed pretraining +schedule. They intentionally do not use `encoder_best_val.pt`, because SSL +validation loss is not a reliable early-stopping criterion for every paradigm. +Downstream finetune and supervised runs still evaluate their best downstream +checkpoint when reporting test metrics, and record checkpoint provenance for +post-hoc fairness evaluation. + +## Validation Checklist + +Before launching final runs: + +- generated matrix has exactly 2590 launched runs including SMART +- generated matrix has exactly 2455 launched runs excluding SMART +- every generated run has `experiment_class` +- no generated command contains historical launch-group overrides +- every generated command for a final rerun contains `revision=` +- every downstream SSL run has exactly one valid pretrain dependency +- supervised, GRU-D, and XGBoost runs have no pretrain dependency +- transfer runs depend on source-dataset core pretrains +- TS2Vec and SMART downstream runs depend on their own class pretrains +- no duplicate scientific fingerprints exist +- export groups by `experiment_class` +- fairness defaults are class-based and do not fetch pretraining runs +- final launch/export/fairness commands target the final W&B project +- final launch/retry commands include `--launch-commit` + +Focused regression suite: + +```bash +uv run pytest \ + tests/test_config_schemas.py \ + tests/test_export_results.py \ + tests/test_evaluate_fairness.py \ + tests/test_fixes.py +``` + +## Archival Note + +Earlier internal planning used execution batches as a convenience while the +matrix was still being explored. Those batch labels are retired for the final +rerun corpus. The final system exposes experiment classes as the public API and +derives dependencies from the generated run graph. diff --git a/examples/eval_metrics_example.py b/examples/eval_metrics_example.py deleted file mode 100644 index 3acda0f..0000000 --- a/examples/eval_metrics_example.py +++ /dev/null @@ -1,226 +0,0 @@ -"""Example: Using the eval module for configurable metrics. - -This example demonstrates: -1. Creating metric configurations -2. Building metrics for different task types -3. Using metrics in a training-like loop -4. Computing and displaying results - -Run: - uv run python examples/eval_metrics_example.py -""" - -import torch -from slices.eval import MetricConfig, build_metrics, get_default_metrics - - -def example_binary_classification(): - """Example: Binary classification with custom metrics.""" - print("=" * 60) - print("Example 1: Binary Classification (Mortality Prediction)") - print("=" * 60) - - # Create metric configuration - config = MetricConfig( - task_type="binary", - metrics=["auroc", "auprc", "accuracy", "f1"], - ) - - # Build metrics for validation - val_metrics = build_metrics(config, prefix="val") - print(f"\nConfigured metrics: {config.metrics}") - print(f"MetricCollection: {val_metrics}\n") - - # Simulate predictions and labels - torch.manual_seed(42) - _n_samples = 100 # noqa: F841 - for documentation - - # Generate realistic mortality prediction data - # Positive class (mortality) has higher predicted probabilities - probs_positive = torch.clamp(torch.randn(20) * 0.2 + 0.7, 0, 1) # ~20% mortality - probs_negative = torch.clamp(torch.randn(80) * 0.2 + 0.3, 0, 1) # ~80% survival - - probs = torch.cat([probs_positive, probs_negative]) - labels = torch.cat([torch.ones(20), torch.zeros(80)]).long() - - # Update metrics - val_metrics.update(probs, labels) - - # Compute results - results = val_metrics.compute() - - print("Results on validation set:") - for metric_name, value in results.items(): - print(f" {metric_name}: {value:.4f}") - - print() - - -def example_multiclass_classification(): - """Example: Multiclass classification with defaults.""" - print("=" * 60) - print("Example 2: Multiclass Classification (Diagnosis Prediction)") - print("=" * 60) - - # Use default metrics for multiclass - test_metrics = get_default_metrics("multiclass", n_classes=5, prefix="test") - print(f"\nDefault metrics: {list(test_metrics.keys())}\n") - - # Simulate predictions (batch_size=50, n_classes=5) - torch.manual_seed(42) - logits = torch.randn(50, 5) - probs = torch.softmax(logits, dim=1) - labels = torch.randint(0, 5, (50,)) - - # Update and compute - test_metrics.update(probs, labels) - results = test_metrics.compute() - - print("Results on test set:") - for metric_name, value in results.items(): - print(f" {metric_name}: {value:.4f}") - - print() - - -def example_batch_accumulation(): - """Example: Accumulating metrics across batches.""" - print("=" * 60) - print("Example 3: Batch Accumulation (Training Loop)") - print("=" * 60) - - # Create metrics - config = MetricConfig(task_type="binary", metrics=["auroc", "auprc"]) - train_metrics = build_metrics(config, prefix="train") - - # Simulate training batches - torch.manual_seed(42) - n_batches = 10 - batch_size = 32 - - print(f"\nProcessing {n_batches} batches of size {batch_size}...\n") - - for batch_idx in range(n_batches): - # Simulate batch predictions - probs = torch.rand(batch_size) - labels = torch.randint(0, 2, (batch_size,)) - - # Update metrics (accumulates across batches) - train_metrics.update(probs, labels) - - if batch_idx % 3 == 0: - print(f" Processed batch {batch_idx + 1}/{n_batches}") - - # Compute final metrics - results = train_metrics.compute() - - print(f"\nFinal metrics (accumulated over {n_batches} batches):") - for metric_name, value in results.items(): - print(f" {metric_name}: {value:.4f}") - - # Reset for next epoch - train_metrics.reset() - print("\n✓ Metrics reset for next epoch") - - print() - - -def example_configuration_validation(): - """Example: Configuration validation and error handling.""" - print("=" * 60) - print("Example 4: Configuration Validation") - print("=" * 60) - - # Valid configuration - try: - config = MetricConfig( - task_type="binary", - metrics=["auroc", "auprc"], - ) - print("\n✓ Valid config created:", config) - except ValueError as e: - print(f"\n✗ Error: {e}") - - # Invalid task type - try: - config = MetricConfig(task_type="invalid_type") - print("\n✓ Config created") - except ValueError as e: - print(f"\n✗ Invalid task type error (expected): {e}") - - # Invalid metric for task - try: - config = MetricConfig( - task_type="binary", - metrics=["nonexistent_metric"], - ) - print("\n✓ Config created") - except ValueError as e: - print(f"\n✗ Invalid metric error (expected): {e}") - - print() - - -def example_minimal_vs_comprehensive(): - """Example: Minimal vs comprehensive metrics.""" - print("=" * 60) - print("Example 5: Minimal vs Comprehensive Metrics") - print("=" * 60) - - # Minimal (defaults) - faster computation - minimal_config = MetricConfig(task_type="binary") - print(f"\nMinimal (default): {minimal_config.metrics}") - - # Comprehensive - more detailed evaluation - comprehensive_config = MetricConfig( - task_type="binary", - metrics=["auroc", "auprc", "accuracy", "f1"], - ) - print(f"Comprehensive: {comprehensive_config.metrics}") - - # Build and compare - torch.manual_seed(42) - probs = torch.rand(100) - labels = torch.randint(0, 2, (100,)) - - # Minimal metrics - minimal_metrics = build_metrics(minimal_config, prefix="val") - minimal_metrics.update(probs, labels) - minimal_results = minimal_metrics.compute() - - print("\nMinimal results:") - for name, value in minimal_results.items(): - print(f" {name}: {value:.4f}") - - # Comprehensive metrics - comprehensive_metrics = build_metrics(comprehensive_config, prefix="val") - comprehensive_metrics.update(probs, labels) - comprehensive_results = comprehensive_metrics.compute() - - print("\nComprehensive results:") - for name, value in comprehensive_results.items(): - print(f" {name}: {value:.4f}") - - print("\n💡 Use minimal for fast iteration, comprehensive for final evaluation") - print() - - -def main(): - """Run all examples.""" - print("\n" + "=" * 60) - print("SLICES Evaluation Module Examples") - print("=" * 60 + "\n") - - example_binary_classification() - example_multiclass_classification() - example_batch_accumulation() - example_configuration_validation() - example_minimal_vs_comprehensive() - - print("=" * 60) - print("All examples completed!") - print("=" * 60) - - -if __name__ == "__main__": - main() diff --git a/examples/mae_pretraining_example.py b/examples/mae_pretraining_example.py deleted file mode 100644 index 2acd205..0000000 --- a/examples/mae_pretraining_example.py +++ /dev/null @@ -1,406 +0,0 @@ -"""Example: MAE pretraining with TransformerEncoder and ICUDataset. - -This example demonstrates how to: -1. Load ICU data with the dataset -2. Configure a transformer encoder -3. Set up MAE (Masked Autoencoder) SSL objective -4. Run a simple training loop -5. Switch between different masking strategies - -The MAE objective learns representations by: -- Masking portions of the input time-series -- Encoding the masked input -- Reconstructing the original values -- Minimizing reconstruction error on masked positions -""" - -from pathlib import Path - -import torch -import torch.nn as nn -from slices.data.dataset import ICUDataset -from slices.models.encoders import TransformerConfig, TransformerEncoder -from slices.models.pretraining import MAEConfig, MAEObjective, build_ssl_objective - - -def create_mae_model( - d_input: int, - d_model: int = 128, - n_layers: int = 4, - n_heads: int = 8, - mask_ratio: float = 0.15, - mask_strategy: str = "random", -) -> nn.Module: - """Create MAE model with transformer encoder. - - Args: - d_input: Input feature dimension. - d_model: Transformer model dimension. - n_layers: Number of transformer layers. - n_heads: Number of attention heads. - mask_ratio: Fraction of input to mask. - mask_strategy: Masking strategy (random, block, timestep, feature). - - Returns: - MAE objective module. - """ - # Configure encoder (must use pooling='none' for MAE) - encoder_config = TransformerConfig( - d_input=d_input, - d_model=d_model, - n_layers=n_layers, - n_heads=n_heads, - d_ff=4 * d_model, - dropout=0.1, - pooling="none", # Required for MAE - need per-timestep outputs - use_positional_encoding=True, - prenorm=True, - ) - encoder = TransformerEncoder(encoder_config) - - # Configure MAE objective - # Note: The MAE now uses a two-token system: - # - MISSING_TOKEN: replaces genuinely missing positions (obs_mask=False) - # - MASK_TOKEN: replaces SSL-masked positions (obs_mask=True AND selected for masking) - # Loss is only computed on MASK_TOKEN positions. - mae_config = MAEConfig( - name="mae", - mask_ratio=mask_ratio, - mask_strategy=mask_strategy, - min_block_size=3, - max_block_size=10, - decoder_d_model=d_model // 2, # Lighter decoder - decoder_n_layers=2, - decoder_n_heads=4, - decoder_d_ff=2 * d_model, - decoder_dropout=0.1, - loss_on_observed_only=True, # Only penalize on originally observed values - norm_target=False, # Whether to normalize reconstruction targets - ) - - # Build MAE objective - mae = MAEObjective(encoder, mae_config) - - return mae - - -def train_mae_epoch( - model: nn.Module, - dataloader: torch.utils.data.DataLoader, - optimizer: torch.optim.Optimizer, - device: torch.device, -) -> dict: - """Run one epoch of MAE training. - - Args: - model: MAE objective module. - dataloader: Training dataloader. - optimizer: Optimizer. - device: Device to train on. - - Returns: - Dict of average metrics for the epoch. - """ - model.train() - - total_loss = 0.0 - total_metrics = {} - n_batches = 0 - - for batch in dataloader: - # Move batch to device - timeseries = batch["timeseries"].to(device) # (B, T, D) - mask = batch["mask"].to(device) # (B, T, D) - - # Forward pass - loss, metrics = model(timeseries, mask) - - # Backward pass - optimizer.zero_grad() - loss.backward() - optimizer.step() - - # Accumulate metrics - total_loss += loss.item() - for key, value in metrics.items(): - if key not in total_metrics: - total_metrics[key] = 0.0 - total_metrics[key] += value.item() if torch.is_tensor(value) else value - - n_batches += 1 - - # Average metrics - avg_metrics = { - "loss": total_loss / n_batches, - **{k: v / n_batches for k, v in total_metrics.items()}, - } - - return avg_metrics - - -def main(): - """Main example demonstrating MAE pretraining.""" - - # Configuration - data_dir = Path("data/processed/mimic-iv-demo") - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - print("=" * 80) - print("MAE Pretraining Example") - print("=" * 80) - - # Check if data exists - if not data_dir.exists(): - print(f"\n❌ Data directory not found: {data_dir}") - print("Please run setup_mimic_iv.py first to extract data.") - return - - # Load dataset - print(f"\n1. Loading dataset from {data_dir}...") - dataset = ICUDataset( - data_dir=data_dir, - task_name=None, # No task labels needed for pretraining - normalize=True, - ) - - print(f" ✓ Loaded {len(dataset)} ICU stays") - print(f" ✓ Feature dimension: {dataset.n_features}") - print(f" ✓ Sequence length: {dataset.seq_length}") - - # Create dataloader - batch_size = 32 - dataloader = torch.utils.data.DataLoader( - dataset, - batch_size=batch_size, - shuffle=True, - num_workers=0, - ) - print(f" ✓ Created dataloader with batch_size={batch_size}") - - # ========================================================================= - # Example 1: Random masking (BERT-style) - # ========================================================================= - print("\n2. Creating MAE model with RANDOM masking...") - mae_random = create_mae_model( - d_input=dataset.n_features, - d_model=128, - n_layers=4, - n_heads=8, - mask_ratio=0.15, # Mask 15% of positions - mask_strategy="random", - ).to(device) - - n_params = sum(p.numel() for p in mae_random.parameters()) - print(f" ✓ Model created with {n_params:,} parameters") - print(" ✓ Masking: 15% random positions") - - # Create optimizer - optimizer = torch.optim.AdamW(mae_random.parameters(), lr=1e-3) - - # Train for a few batches - print("\n3. Training with random masking...") - mae_random.train() - - for i, batch in enumerate(dataloader): - if i >= 3: # Just 3 batches for demo - break - - timeseries = batch["timeseries"].to(device) - mask = batch["mask"].to(device) - - loss, metrics = mae_random(timeseries, mask) - - optimizer.zero_grad() - loss.backward() - optimizer.step() - - print( - f" Batch {i+1}: loss={loss.item():.4f}, " - f"masked_loss={metrics['mae_recon_loss_masked'].item():.4f}, " - f"visible_loss={metrics['mae_recon_loss_visible'].item():.4f}" - ) - - # ========================================================================= - # Example 2: Block masking (for temporal structure learning) - # ========================================================================= - print("\n4. Creating MAE model with BLOCK masking...") - mae_block = create_mae_model( - d_input=dataset.n_features, - d_model=128, - n_layers=4, - n_heads=8, - mask_ratio=0.15, - mask_strategy="block", # Mask contiguous time blocks - ).to(device) - - print(" ✓ Model created with block masking") - print(" ✓ Masking: 15% in contiguous blocks (3-10 timesteps)") - - optimizer = torch.optim.AdamW(mae_block.parameters(), lr=1e-3) - - print("\n5. Training with block masking...") - mae_block.train() - - for i, batch in enumerate(dataloader): - if i >= 3: - break - - timeseries = batch["timeseries"].to(device) - mask = batch["mask"].to(device) - - loss, metrics = mae_block(timeseries, mask) - - optimizer.zero_grad() - loss.backward() - optimizer.step() - - print( - f" Batch {i+1}: loss={loss.item():.4f}, " - f"mask_ratio={metrics['mae_mask_ratio_actual']:.3f}" - ) - - # ========================================================================= - # Example 3: Timestep masking (mask entire timesteps) - # ========================================================================= - print("\n6. Creating MAE model with TIMESTEP masking...") - mae_timestep = create_mae_model( - d_input=dataset.n_features, - d_model=128, - n_layers=4, - n_heads=8, - mask_ratio=0.15, - mask_strategy="timestep", # Mask entire timesteps - ).to(device) - - print(" ✓ Model created with timestep masking") - print(" ✓ Masking: 15% of timesteps (all features)") - - optimizer = torch.optim.AdamW(mae_timestep.parameters(), lr=1e-3) - - print("\n7. Training with timestep masking...") - mae_timestep.train() - - for i, batch in enumerate(dataloader): - if i >= 3: - break - - timeseries = batch["timeseries"].to(device) - mask = batch["mask"].to(device) - - loss, metrics = mae_timestep(timeseries, mask) - - optimizer.zero_grad() - loss.backward() - optimizer.step() - - print( - f" Batch {i+1}: loss={loss.item():.4f}, " f"obs_ratio={metrics['mae_obs_ratio']:.3f}" - ) - - # ========================================================================= - # Example 4: Feature masking (mask entire features) - # ========================================================================= - print("\n8. Creating MAE model with FEATURE masking...") - mae_feature = create_mae_model( - d_input=dataset.n_features, - d_model=128, - n_layers=4, - n_heads=8, - mask_ratio=0.15, - mask_strategy="feature", # Mask entire features - ).to(device) - - print(" ✓ Model created with feature masking") - print(" ✓ Masking: 15% of features (all timesteps)") - - optimizer = torch.optim.AdamW(mae_feature.parameters(), lr=1e-3) - - print("\n9. Training with feature masking...") - mae_feature.train() - - for i, batch in enumerate(dataloader): - if i >= 3: - break - - timeseries = batch["timeseries"].to(device) - mask = batch["mask"].to(device) - - loss, metrics = mae_feature(timeseries, mask) - - optimizer.zero_grad() - loss.backward() - optimizer.step() - - print(f" Batch {i+1}: loss={loss.item():.4f}") - - # ========================================================================= - # Example 5: Using factory function for easy switching - # ========================================================================= - print("\n10. Using factory function for easy SSL objective switching...") - - # Create encoder - encoder_config = TransformerConfig( - d_input=dataset.n_features, - d_model=128, - n_layers=4, - n_heads=8, - pooling="none", - ) - encoder = TransformerEncoder(encoder_config).to(device) - - # Configure MAE - mae_config = MAEConfig( - name="mae", - mask_ratio=0.20, - mask_strategy="block", - ) - - # Build using factory - ssl_objective = build_ssl_objective(encoder, mae_config).to(device) - - print(f" ✓ Created SSL objective: {mae_config.name}") - print(f" ✓ Strategy: {mae_config.mask_strategy}") - print(f" ✓ Mask ratio: {mae_config.mask_ratio}") - - # Train for one batch - optimizer = torch.optim.AdamW(ssl_objective.parameters(), lr=1e-3) - - batch = next(iter(dataloader)) - timeseries = batch["timeseries"].to(device) - mask = batch["mask"].to(device) - - loss, metrics = ssl_objective(timeseries, mask) - optimizer.zero_grad() - loss.backward() - optimizer.step() - - print(f" ✓ Trained one batch: loss={loss.item():.4f}") - - # ========================================================================= - # Summary - # ========================================================================= - print("\n" + "=" * 80) - print("Summary") - print("=" * 80) - print("\n✓ MAE objective implemented with 4 masking strategies:") - print(" 1. Random: Uniform random masking (BERT-style)") - print(" 2. Block: Contiguous time blocks (temporal structure)") - print(" 3. Timestep: Entire timesteps (all features)") - print(" 4. Feature: Entire features (all timesteps)") - print("\n✓ Key features:") - print(" - Two-token system: MISSING_TOKEN and MASK_TOKEN") - print(" - MISSING_TOKEN replaces genuinely missing positions") - print(" - MASK_TOKEN replaces SSL-masked positions (for reconstruction)") - print(" - Configurable mask ratio and strategy") - print(" - Lightweight decoder for reconstruction") - print(" - Loss computed only on MASK_TOKEN positions") - print(" - Easy to switch objectives via factory pattern") - print("\n✓ Next steps:") - print(" - Implement other SSL objectives (contrastive, JEPA, etc.)") - print(" - Add to scripts/pretrain.py for full training") - print(" - Evaluate learned representations on downstream tasks") - print("=" * 80) - - -if __name__ == "__main__": - main() diff --git a/examples/transformer_encoder_example.py b/examples/transformer_encoder_example.py deleted file mode 100644 index 0af8d0c..0000000 --- a/examples/transformer_encoder_example.py +++ /dev/null @@ -1,273 +0,0 @@ -"""Example: Using the Transformer Encoder for ICU Time-Series - -This script demonstrates how to: -1. Create a transformer encoder from config -2. Process ICU time-series data with observation and padding masks -3. Use different pooling strategies -4. Integrate with PyTorch training loops -""" - -import torch -import torch.nn as nn -from slices.models.encoders import TransformerConfig, TransformerEncoder - - -def basic_usage(): - """Basic usage: Create encoder and process time-series.""" - print("=== Basic Usage ===") - - # Create config - config = TransformerConfig( - d_input=35, # Number of input features - d_model=128, # Model dimension - n_layers=4, # Number of transformer layers - n_heads=8, # Number of attention heads - d_ff=512, # Feedforward dimension - max_seq_length=168, # 7 days in hours - pooling="mean", # Mean pooling for sequence-level representation - ) - - # Create encoder - encoder = TransformerEncoder(config) - print(f"Created transformer with {sum(p.numel() for p in encoder.parameters()):,} parameters") - - # Simulate ICU data: batch_size=16, seq_length=48 hours, features=35 - batch_size = 16 - seq_length = 48 - x = torch.randn(batch_size, seq_length, config.d_input) - - # Forward pass - embeddings = encoder(x) - print(f"Input shape: {x.shape}") - print(f"Output shape (mean pooling): {embeddings.shape}") - print() - - -def with_observation_mask(): - """Using observation mask to handle missing values.""" - print("=== With Observation Mask (Missing Values) ===") - - config = TransformerConfig( - d_input=35, - d_model=128, - n_layers=2, - n_heads=8, - pooling="mean", - ) - encoder = TransformerEncoder(config) - - batch_size = 8 - seq_length = 48 - x = torch.randn(batch_size, seq_length, config.d_input) - - # Create observation mask (True = observed, False = missing) - # Simulate 30% missing values - obs_mask = torch.rand(batch_size, seq_length, config.d_input) > 0.3 - - print(f"Data shape: {x.shape}") - print( - f"Missing values: {(~obs_mask).sum().item()} / {obs_mask.numel()} " - f"({100 * (~obs_mask).sum().item() / obs_mask.numel():.1f}%)" - ) - - # Note: Missing values should be imputed before passing to encoder - # The mask is currently used for logging/analysis only - embeddings = encoder(x, mask=obs_mask) - print(f"Output embeddings: {embeddings.shape}") - print() - - -def with_padding_mask(): - """Using padding mask for variable-length sequences.""" - print("=== With Padding Mask (Variable-Length Sequences) ===") - - config = TransformerConfig( - d_input=35, - d_model=128, - n_layers=2, - n_heads=8, - pooling="mean", - ) - encoder = TransformerEncoder(config) - - batch_size = 8 - max_length = 72 - x = torch.randn(batch_size, max_length, config.d_input) - - # Create padding mask (True = valid, False = padding) - # Simulate variable-length sequences - padding_mask = torch.ones(batch_size, max_length, dtype=torch.bool) - sequence_lengths = torch.randint(24, max_length, (batch_size,)) - - for i, length in enumerate(sequence_lengths): - padding_mask[i, length:] = False - - print(f"Data shape: {x.shape}") - print(f"Sequence lengths: {sequence_lengths.tolist()}") - - # Forward pass with padding mask - embeddings = encoder(x, padding_mask=padding_mask) - print(f"Output embeddings: {embeddings.shape}") - print("Note: Mean pooling correctly averages only over valid timesteps") - print() - - -def different_pooling_strategies(): - """Comparing different pooling strategies.""" - print("=== Different Pooling Strategies ===") - - batch_size = 4 - seq_length = 48 - x = torch.randn(batch_size, seq_length, 35) - - pooling_strategies = ["mean", "max", "cls", "last", "none"] - - for pooling in pooling_strategies: - config = TransformerConfig( - d_input=35, - d_model=128, - n_layers=2, - n_heads=8, - pooling=pooling, - ) - encoder = TransformerEncoder(config) - - embeddings = encoder(x) - print(f"Pooling: {pooling:6s} -> Output shape: {embeddings.shape}") - - print() - - -def prenorm_vs_postnorm(): - """Comparing Pre-LN and Post-LN architectures.""" - print("=== Pre-LN vs Post-LN Transformer ===") - - x = torch.randn(8, 48, 35) - - # Pre-LN (modern, more stable) - config_prenorm = TransformerConfig( - d_input=35, - d_model=128, - n_layers=4, - n_heads=8, - prenorm=True, - ) - encoder_prenorm = TransformerEncoder(config_prenorm) - - # Post-LN (original transformer) - config_postnorm = TransformerConfig( - d_input=35, - d_model=128, - n_layers=4, - n_heads=8, - prenorm=False, - ) - encoder_postnorm = TransformerEncoder(config_postnorm) - - embeddings_prenorm = encoder_prenorm(x) - embeddings_postnorm = encoder_postnorm(x) - - print(f"Pre-LN output shape: {embeddings_prenorm.shape}") - print(f"Post-LN output shape: {embeddings_postnorm.shape}") - print("Pre-LN is recommended for deeper models (more stable training)") - print() - - -def training_loop_integration(): - """Integration with PyTorch training loop.""" - print("=== Training Loop Integration ===") - - # Setup - config = TransformerConfig( - d_input=35, - d_model=128, - n_layers=4, - n_heads=8, - pooling="mean", - ) - encoder = TransformerEncoder(config) - - # Simple downstream task: binary classification from embeddings - classifier = nn.Sequential( - nn.Linear(128, 64), - nn.ReLU(), - nn.Dropout(0.1), - nn.Linear(64, 1), - ) - - # Optimizer - model = nn.Sequential(encoder, classifier) - optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) - criterion = nn.BCEWithLogitsLoss() - - # Simulate training batch - batch_size = 16 - x = torch.randn(batch_size, 48, 35) - labels = torch.randint(0, 2, (batch_size, 1)).float() - - # Training step - model.train() - optimizer.zero_grad() - - logits = model(x) - loss = criterion(logits, labels) - - loss.backward() - optimizer.step() - - print(f"Batch size: {batch_size}") - print(f"Loss: {loss.item():.4f}") - print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}") - print() - - -def memory_efficient_inference(): - """Memory-efficient inference with torch.no_grad().""" - print("=== Memory-Efficient Inference ===") - - config = TransformerConfig( - d_input=35, - d_model=128, - n_layers=4, - n_heads=8, - pooling="mean", - ) - encoder = TransformerEncoder(config) - encoder.eval() # Set to evaluation mode (disables dropout) - - # Large batch for inference - batch_size = 128 - x = torch.randn(batch_size, 48, 35) - - # No gradient computation for inference - with torch.no_grad(): - embeddings = encoder(x) - - print(f"Batch size: {batch_size}") - print(f"Output shape: {embeddings.shape}") - print("Using torch.no_grad() saves memory by not storing gradients") - print() - - -def main(): - """Run all examples.""" - print("=" * 70) - print("Transformer Encoder Examples for ICU Time-Series") - print("=" * 70) - print() - - basic_usage() - with_observation_mask() - with_padding_mask() - different_pooling_strategies() - prenorm_vs_postnorm() - training_loop_integration() - memory_efficient_inference() - - print("=" * 70) - print("Examples completed!") - print("=" * 70) - - -if __name__ == "__main__": - main() diff --git a/examples/transformer_with_dataset_example.py b/examples/transformer_with_dataset_example.py deleted file mode 100644 index 2dfd77c..0000000 --- a/examples/transformer_with_dataset_example.py +++ /dev/null @@ -1,282 +0,0 @@ -"""Example: Integrating Transformer Encoder with ICUDataset - -This script demonstrates how to use the transformer encoder with real ICU data -loaded from the ICUDataset. It shows the complete pipeline from data loading -to embedding extraction. -""" - -from pathlib import Path - -import torch -from slices.models.encoders import TransformerConfig, TransformerEncoder - - -def check_data_availability() -> bool: - """Check if extracted MIMIC-IV demo data is available.""" - data_dir = Path("data/processed/mimic-iv-demo") - return data_dir.exists() and (data_dir / "timeseries.parquet").exists() - - -def basic_integration(): - """Basic integration: Load dataset and process with transformer.""" - print("=== Basic Integration with ICUDataset ===") - - if not check_data_availability(): - print("⚠️ MIMIC-IV demo data not found.") - print(" Run: uv run python scripts/setup_mimic_iv.py") - print(" This example requires extracted ICU data.\n") - return False - - from slices.data import ICUDataset - - # Load dataset - dataset = ICUDataset( - "data/processed/mimic-iv-demo", - task_name="mortality_24h", - normalize=True, - ) - - print(f"Dataset loaded: {len(dataset)} samples") - print(f"Features: {dataset.n_features}") - print(f"Sequence length: {dataset.seq_length}") - print(f"Available tasks: {dataset.task_names}") - - # Get a sample - sample = dataset[0] - print("\nSample structure:") - print(f" timeseries: {sample['timeseries'].shape}") - print(f" mask: {sample['mask'].shape}") - print(f" label: {sample['label'].item()}") - print(f" static: {list(sample['static'].keys())}") - - # Create transformer - config = TransformerConfig( - d_input=dataset.n_features, - d_model=128, - n_layers=4, - n_heads=8, - d_ff=512, - max_seq_length=dataset.seq_length, - pooling="mean", - ) - encoder = TransformerEncoder(config) - print(f"\nTransformer created: {sum(p.numel() for p in encoder.parameters()):,} parameters") - - # Process single sample - x = sample["timeseries"].unsqueeze(0) # Add batch dimension - mask = sample["mask"].unsqueeze(0) - - encoder.eval() - with torch.no_grad(): - embedding = encoder(x, mask=mask) - - print(f"\nSingle sample embedding: {embedding.shape}") - print(f"Embedding norm: {embedding.norm().item():.4f}") - - return True - - -def batch_processing(): - """Process a batch of samples.""" - print("=== Batch Processing ===") - - if not check_data_availability(): - print("⚠️ MIMIC-IV demo data not found. Skipping.\n") - return False - - from slices.data import ICUDataset - from torch.utils.data import DataLoader - - # Load dataset - dataset = ICUDataset( - "data/processed/mimic-iv-demo", - task_name="mortality_24h", - normalize=True, - ) - - # Create dataloader - dataloader = DataLoader( - dataset, - batch_size=16, - shuffle=True, - num_workers=0, # Use 0 for debugging, >0 for production - ) - - # Create transformer - config = TransformerConfig( - d_input=dataset.n_features, - d_model=128, - n_layers=4, - n_heads=8, - pooling="mean", - ) - encoder = TransformerEncoder(config) - encoder.eval() - - # Process first batch - batch = next(iter(dataloader)) - x = batch["timeseries"] # (B, T, D) - mask = batch["mask"] # (B, T, D) - labels = batch["label"] # (B,) - - print(f"Batch shape: {x.shape}") - print(f"Labels: {labels.tolist()}") - - with torch.no_grad(): - embeddings = encoder(x, mask=mask) - - print(f"Embeddings shape: {embeddings.shape}") - print(f"Mean embedding norm: {embeddings.norm(dim=1).mean().item():.4f}") - - return True - - -def with_datamodule(): - """Use with ICUDataModule (Lightning).""" - print("=== Integration with ICUDataModule ===") - - if not check_data_availability(): - print("⚠️ MIMIC-IV demo data not found. Skipping.\n") - return False - - from slices.data import ICUDataModule - - # Create datamodule - datamodule = ICUDataModule( - processed_dir="data/processed/mimic-iv-demo", - task_name="mortality_24h", - batch_size=32, - num_workers=0, - train_ratio=0.7, - val_ratio=0.15, - test_ratio=0.15, - seed=42, - ) - - # Setup splits - datamodule.setup() - - # Get split information - split_info = datamodule.get_split_info() - print(f"Train stays: {split_info['train_stays']} ({split_info['train_patients']} patients)") - print(f"Val stays: {split_info['val_stays']} ({split_info['val_patients']} patients)") - print(f"Test stays: {split_info['test_stays']} ({split_info['test_patients']} patients)") - - # Get train dataloader - train_loader = datamodule.train_dataloader() - - # Create transformer - config = TransformerConfig( - d_input=datamodule.get_feature_dim(), - d_model=128, - n_layers=4, - n_heads=8, - pooling="mean", - ) - encoder = TransformerEncoder(config) - - # Simulate training loop - encoder.train() - optimizer = torch.optim.AdamW(encoder.parameters(), lr=1e-4) - - batch = next(iter(train_loader)) - x = batch["timeseries"] - mask = batch["mask"] - _labels = batch["label"] # Available for task heads - - # Forward pass - embeddings = encoder(x, mask=mask) - - # Dummy loss (normally you'd have a task head) - loss = embeddings.pow(2).mean() # Just for demonstration - - # Backward pass - optimizer.zero_grad() - loss.backward() - optimizer.step() - - print("\nTraining step completed:") - print(f" Batch size: {x.size(0)}") - print(f" Loss: {loss.item():.4f}") - print(f" Embeddings mean: {embeddings.mean().item():.4f}") - - return True - - -def different_sequence_lengths(): - """Handle sequences with different lengths (padding).""" - print("=== Variable-Length Sequences ===") - - if not check_data_availability(): - print("⚠️ MIMIC-IV demo data not found. Skipping.\n") - return False - - from slices.data import ICUDataset - - # Load dataset with shorter sequences - dataset = ICUDataset( - "data/processed/mimic-iv-demo", - seq_length=24, # Use only first 24 hours - normalize=True, - ) - - print(f"Dataset with seq_length={dataset.seq_length}") - - # Create transformer (with longer max_seq_length) - config = TransformerConfig( - d_input=dataset.n_features, - d_model=128, - n_layers=2, - n_heads=8, - max_seq_length=72, # Can handle up to 72 hours - pooling="mean", - ) - encoder = TransformerEncoder(config) - - # Get sample - sample = dataset[0] - x = sample["timeseries"].unsqueeze(0) - mask = sample["mask"].unsqueeze(0) - - print(f"Input shape: {x.shape}") - - encoder.eval() - with torch.no_grad(): - embedding = encoder(x, mask=mask) - - print(f"Output shape: {embedding.shape}") - print("✓ Transformer handles variable sequence lengths correctly\n") - - return True - - -def main(): - """Run all examples.""" - print("=" * 70) - print("Transformer + ICUDataset Integration Examples") - print("=" * 70) - print() - - # Track which examples ran successfully - success = [] - - success.append(basic_integration()) - success.append(batch_processing()) - success.append(with_datamodule()) - success.append(different_sequence_lengths()) - - print("=" * 70) - - if any(success): - print("✓ Examples completed successfully!") - else: - print("⚠️ No examples ran (MIMIC-IV demo data not found)") - print("\nTo run these examples:") - print("1. Download MIMIC-IV demo dataset") - print("2. Run: uv run python scripts/setup_mimic_iv.py") - - print("=" * 70) - - -if __name__ == "__main__": - main() diff --git a/pyproject.toml b/pyproject.toml index 1f23e88..15101ae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,19 +15,25 @@ dependencies = [ "omegaconf>=2.3", "polars>=0.19", "pyarrow>=14.0", - "duckdb>=0.9", "wandb>=0.15", "numpy>=1.24", + "pandas>=2.0", + "pydantic>=2.0", + "PyYAML>=6.0", + "scipy>=1.10", "scikit-learn>=1.3", "rich>=13.0", - "matplotlib>=3.10.7", - "seaborn>=0.13.2", - "pre-commit>=4.5.0", "portalocker>=2.8.0", - "m4-infra>=0.0.0.dev0", + "torchmetrics>=1.0", + "tqdm>=4.0", + "xgboost>=2.0", ] [project.optional-dependencies] +viz = [ + "matplotlib>=3.10.7", + "umap-learn>=0.5", +] dev = [ "pytest>=7.0", "pytest-cov>=4.0", @@ -37,8 +43,8 @@ dev = [ "pre-commit>=3.0", ] -[tool.setuptools.packages.find] -where = ["src"] +[tool.hatch.build.targets.wheel] +packages = ["src/slices"] [tool.black] line-length = 100 @@ -77,7 +83,11 @@ disable_error_code = [ [dependency-groups] dev = [ + "black>=23.0", + "mypy>=1.0", "pytest>=9.0.1", + "pytest-cov>=4.0", "python-pptx>=1.0.2", "ruff>=0.14.8", + "types-PyYAML>=6.0", ] diff --git a/results/.gitkeep b/results/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/scripts/analyze_labels.py b/scripts/analyze_labels.py index 588d5c6..133a463 100644 --- a/scripts/analyze_labels.py +++ b/scripts/analyze_labels.py @@ -4,22 +4,23 @@ This script provides comprehensive statistics about task labels including: - Available tasks and their types - Class distributions (counts and percentages) +- Continuous-label summaries for regression tasks - Missing label counts - Per-split statistics (train/val/test) - Class imbalance metrics Example usage: # Analyze labels in processed directory - uv run python scripts/analyze_labels.py data/processed/mimic-iv-demo + uv run python scripts/analyze_labels.py data/processed/miiv # Analyze specific task - uv run python scripts/analyze_labels.py data/processed/mimic-iv-demo --task mortality_24h + uv run python scripts/analyze_labels.py data/processed/miiv --task mortality_24h # Show per-split statistics - uv run python scripts/analyze_labels.py data/processed/mimic-iv-demo --splits + uv run python scripts/analyze_labels.py data/processed/miiv --splits # Export to JSON - uv run python scripts/analyze_labels.py data/processed/mimic-iv-demo --output stats.json + uv run python scripts/analyze_labels.py data/processed/miiv --output stats.json """ import argparse @@ -31,6 +32,8 @@ import polars as pl import yaml +from slices.data.utils import get_package_data_dir + def load_labels(processed_dir: Path) -> pl.DataFrame: """Load labels from parquet file. @@ -68,6 +71,21 @@ def load_metadata(processed_dir: Path) -> Dict[str, Any]: return yaml.safe_load(f) or {} +def load_task_config(task_name: str) -> Dict[str, Any]: + """Load task metadata from the package YAML.""" + config_path = get_package_data_dir() / "tasks" / f"{task_name}.yaml" + if not config_path.exists(): + return {} + + with open(config_path) as f: + return yaml.safe_load(f) or {} + + +def load_task_quality_checks(task_name: str) -> Dict[str, Any]: + """Load optional warning thresholds for a task from the package YAML.""" + return load_task_config(task_name).get("quality_checks", {}) or {} + + def get_task_columns(labels_df: pl.DataFrame) -> List[str]: """Get task column names (excluding stay_id). @@ -179,9 +197,45 @@ def compute_auroc_ci_width( return lower, upper +def infer_task_type(valid_labels: pl.Series, configured_task_type: Optional[str] = None) -> str: + """Infer task type when the task YAML does not provide it.""" + if configured_task_type: + return configured_task_type + + values = valid_labels.to_numpy() + if len(values) == 0: + return "unknown" + + finite = values[np.isfinite(values.astype(float))] + if len(finite) == 0: + return "unknown" + + unique = np.unique(finite) + integer_like = np.allclose(unique, np.round(unique)) + if integer_like and len(unique) <= 20: + return "binary" if set(unique.astype(int)).issubset({0, 1}) else "multiclass" + return "regression" + + +def summarize_numeric_labels(labels: np.ndarray) -> Dict[str, float]: + """Return descriptive statistics for continuous labels.""" + labels = labels.astype(float) + return { + "mean": round(float(np.mean(labels)), 4), + "std": round(float(np.std(labels, ddof=1)), 4) if len(labels) > 1 else 0.0, + "min": round(float(np.min(labels)), 4), + "q25": round(float(np.percentile(labels, 25)), 4), + "median": round(float(np.median(labels)), 4), + "q75": round(float(np.percentile(labels, 75)), 4), + "max": round(float(np.max(labels)), 4), + } + + def analyze_task( labels_df: pl.DataFrame, task_name: str, + task_type: Optional[str] = None, + quality_checks: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """Analyze a single task's label distribution. @@ -211,13 +265,36 @@ def analyze_task( "valid_labels": valid_count, "missing_labels": missing_count, "missing_percentage": (missing_count / total_stays * 100) if total_stays > 0 else 0, + "quality_checks": quality_checks or {}, + "quality_warnings": [], } if valid_count == 0: + stats["task_type"] = task_type or "unknown" stats["class_distribution"] = {} stats["n_classes"] = 0 return stats + inferred_task_type = infer_task_type(valid_labels, task_type) + stats["task_type"] = inferred_task_type + + if inferred_task_type == "regression": + values = valid_labels.to_numpy().astype(float) + stats["label_summary"] = summarize_numeric_labels(values) + stats["n_classes"] = None + stats["class_distribution"] = {} + if quality_checks: + max_missing_percentage = quality_checks.get("max_missing_percentage") + if max_missing_percentage is not None and stats["missing_percentage"] > float( + max_missing_percentage + ): + stats["quality_warnings"].append( + "Missing label rate " + f"{stats['missing_percentage']:.1f}% exceeds the configured " + f"maximum of {float(max_missing_percentage):.1f}%." + ) + return stats + # Get unique classes and their counts class_counts = valid_labels.value_counts().sort("count", descending=True) @@ -254,12 +331,24 @@ def analyze_task( 1: round(valid_count / (2 * positive_count), 4) if positive_count > 0 else 1.0, } + if quality_checks: + max_missing_percentage = quality_checks.get("max_missing_percentage") + if max_missing_percentage is not None and stats["missing_percentage"] > float( + max_missing_percentage + ): + stats["quality_warnings"].append( + "Missing label rate " + f"{stats['missing_percentage']:.1f}% exceeds the configured " + f"maximum of {float(max_missing_percentage):.1f}%." + ) + return stats def analyze_splits( processed_dir: Path, task_name: str, + task_type: Optional[str] = None, train_ratio: float = 0.7, val_ratio: float = 0.15, test_ratio: float = 0.15, @@ -310,7 +399,10 @@ def analyze_splits( for idx in indices: sample = datamodule.dataset[idx] if "label" in sample and sample["label"] is not None: - labels.append(int(sample["label"])) + value = sample["label"] + if hasattr(value, "item"): + value = value.item() + labels.append(float(value)) if not labels: split_stats[split_name] = { @@ -320,8 +412,21 @@ def analyze_splits( } continue - labels_array = np.array(labels) - labels_series = pl.Series(labels) + labels_array = np.array(labels, dtype=float) + split_task_type = infer_task_type(pl.Series(labels_array), task_type) + + if split_task_type == "regression": + split_stats[split_name] = { + "n_samples": len(indices), + "n_valid_labels": len(labels), + "task_type": split_task_type, + "label_summary": summarize_numeric_labels(labels_array), + "class_distribution": {}, + } + continue + + class_labels = labels_array.astype(int) + labels_series = pl.Series(class_labels) class_counts = labels_series.value_counts().sort("count", descending=True) class_distribution = {} @@ -336,6 +441,7 @@ def analyze_splits( split_stats[split_name] = { "n_samples": len(indices), "n_valid_labels": len(labels), + "task_type": split_task_type, "class_distribution": class_distribution, } @@ -385,7 +491,23 @@ def print_task_stats(stats: Dict[str, Any]) -> None: print(f" Total stays: {stats['total_stays']:,}") print(f" Valid labels: {stats['valid_labels']:,}") print(f" Missing labels: {stats['missing_labels']:,} ({stats['missing_percentage']:.1f}%)") - print(f" Number of classes: {stats['n_classes']}") + print(f" Task type: {stats.get('task_type', 'unknown')}") + if stats.get("n_classes") is not None: + print(f" Number of classes: {stats['n_classes']}") + + if stats.get("quality_warnings"): + print("\n Quality Warnings:") + for warning in stats["quality_warnings"]: + print(f" - {warning}") + + if stats.get("label_summary"): + summary = stats["label_summary"] + print("\n Label Summary:") + print(f" Mean: {summary['mean']:.4f}") + print(f" Std: {summary['std']:.4f}") + print(f" Median: {summary['median']:.4f}") + print(f" IQR: {summary['q25']:.4f} - {summary['q75']:.4f}") + print(f" Range: {summary['min']:.4f} - {summary['max']:.4f}") if stats["class_distribution"]: print("\n Class Distribution:") @@ -414,6 +536,28 @@ def print_split_stats(split_stats: Dict[str, Dict[str, Any]], task_name: str) -> print(f"\n Per-Split Statistics for '{task_name}':") print(f" {'─' * 60}") + if any("label_summary" in split_stats.get(s, {}) for s in ["train", "val", "test"]): + print(f" {'Split':<10} {'Samples':>10} {'Mean':>12} {'Median':>12} {'IQR':>21}") + print(f" {'-' * 10} {'-' * 10} {'-' * 12} {'-' * 12} {'-' * 21}") + + for split_name in ["train", "val", "test"]: + if split_name not in split_stats: + continue + stats = split_stats[split_name] + summary = stats.get("label_summary") + if not summary: + print( + f" {split_name:<10} {stats['n_samples']:>10,} " + f"{'N/A':>12} {'N/A':>12} {'N/A':>21}" + ) + continue + iqr = f"{summary['q25']:.2f}-{summary['q75']:.2f}" + print( + f" {split_name:<10} {stats['n_samples']:>10,} " + f"{summary['mean']:>12.2f} {summary['median']:>12.2f} {iqr:>21}" + ) + return + # Basic counts header print(f" {'Split':<10} {'Samples':>10} {'Positive':>12} {'Negative':>12} {'Pos Rate':>12}") print(f" {'-' * 10} {'-' * 10} {'-' * 12} {'-' * 12} {'-' * 12}") @@ -598,7 +742,13 @@ def main(): print("=" * 60) for task_name in task_columns: - stats = analyze_task(labels_df, task_name) + task_config = load_task_config(task_name) + stats = analyze_task( + labels_df, + task_name, + task_type=task_config.get("task_type"), + quality_checks=task_config.get("quality_checks", {}) or {}, + ) all_stats[task_name] = stats print_task_stats(stats) @@ -612,10 +762,12 @@ def main(): test_ratio = val_ratio for task_name in task_columns: + task_config = load_task_config(task_name) try: split_stats = analyze_splits( args.processed_dir, task_name, + task_type=task_config.get("task_type"), train_ratio=args.train_ratio, val_ratio=val_ratio, test_ratio=test_ratio, diff --git a/scripts/debug/explore_missingness.py b/scripts/debug/explore_missingness.py index d869be6..55f652d 100644 --- a/scripts/debug/explore_missingness.py +++ b/scripts/debug/explore_missingness.py @@ -10,7 +10,7 @@ Example usage: uv run python scripts/debug/explore_missingness.py \\ - --processed-dir data/processed/mimic-iv-demo \\ + --processed-dir data/processed/miiv \\ --top-n 10 """ @@ -21,6 +21,8 @@ import polars as pl import yaml +from slices.constants import SEQ_LENGTH_HOURS + def load_data(data_dir: Path) -> tuple[pl.DataFrame, pl.DataFrame, dict]: """Load timeseries, static data, and metadata.""" @@ -333,7 +335,7 @@ def main(): parser.add_argument( "--processed-dir", type=Path, - default=Path("data/processed/mimic-iv-demo"), + default=Path("data/processed/miiv"), help="Path to processed data directory (with static.parquet, " "timeseries.parquet, metadata.yaml)", ) @@ -349,7 +351,7 @@ def main(): timeseries_df, static_df, metadata = load_data(args.processed_dir) feature_names = metadata["feature_names"] - seq_length = metadata.get("seq_length_hours", 48) + seq_length = metadata.get("seq_length_hours", SEQ_LENGTH_HOURS) n_stays = len(timeseries_df) n_features = len(feature_names) diff --git a/scripts/debug/export_snapshots.py b/scripts/debug/export_snapshots.py new file mode 100644 index 0000000..ec398d7 --- /dev/null +++ b/scripts/debug/export_snapshots.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python +"""Export CSV snapshots from a processed SLICES dataset. + +Example usage: + uv run python scripts/debug/export_snapshots.py + uv run python scripts/debug/export_snapshots.py data.processed_dir=data/processed/eicu + uv run python scripts/debug/export_snapshots.py stay_ids='[123,456]' +""" + +from pathlib import Path + +import hydra +import polars as pl +from omegaconf import DictConfig + +from slices.debug.snapshots import create_snapshots_from_processed + + +def resolve_stay_ids(processed_dir: Path, cfg: DictConfig) -> list[int]: + """Return explicit stay IDs or a small deterministic sample from static.parquet.""" + if cfg.get("stay_ids"): + return [int(stay_id) for stay_id in cfg.stay_ids] + + static_path = processed_dir / "static.parquet" + if not static_path.exists(): + raise FileNotFoundError( + f"Cannot infer stay_ids because {static_path} does not exist. " + "Pass stay_ids='[id1,id2]' or generate the processed dataset first." + ) + + n_stays = int(cfg.get("sentinel", {}).get("n_per_stratum", 3)) + static_df = pl.read_parquet(static_path, columns=["stay_id"]).head(max(n_stays, 1)) + return [int(stay_id) for stay_id in static_df["stay_id"].to_list()] + + +@hydra.main(version_base=None, config_path="../../configs", config_name="debug") +def main(cfg: DictConfig) -> None: + processed_dir = Path(cfg.data.processed_dir) + output_dir = ( + Path(cfg.debug_output_dir) + if cfg.get("debug_output_dir") + else processed_dir / "debug_snapshots" + ) + + stay_ids = resolve_stay_ids(processed_dir, cfg) + exported = create_snapshots_from_processed( + processed_dir=processed_dir, + stay_ids=stay_ids, + output_dir=output_dir, + include_labels=True, + flatten_timeseries=bool(cfg.get("flatten_timeseries", True)), + ) + + print(f"Exported {len(exported)} snapshot artifacts to {output_dir}") + for name, path in exported.items(): + print(f" {name}: {path}") + + +if __name__ == "__main__": + main() diff --git a/scripts/debug/inspect_embeddings.py b/scripts/debug/inspect_embeddings.py index ffac1aa..ac455ba 100644 --- a/scripts/debug/inspect_embeddings.py +++ b/scripts/debug/inspect_embeddings.py @@ -7,7 +7,7 @@ # Analyze embeddings from a checkpoint uv run python scripts/debug/inspect_embeddings.py \ checkpoint=outputs/encoder.pt \ - processed_dir=data/processed/mimic-iv-demo + processed_dir=data/processed/miiv # Load pre-computed embeddings uv run python scripts/debug/inspect_embeddings.py \ @@ -16,7 +16,7 @@ # Generate visualizations uv run python scripts/debug/inspect_embeddings.py \ checkpoint=outputs/encoder.pt \ - processed_dir=data/processed/mimic-iv-demo \ + processed_dir=data/processed/miiv \ plots=true """ @@ -26,6 +26,7 @@ import numpy as np import polars as pl from omegaconf import DictConfig + from slices.debug import ( EmbeddingQualityReport, analyze_embeddings, @@ -89,6 +90,7 @@ def extract_embeddings_from_checkpoint( Tuple of (embeddings, stay_ids, labels). """ import torch + from slices.data.datamodule import ICUDataModule from slices.models.encoders import build_encoder @@ -111,7 +113,7 @@ def extract_embeddings_from_checkpoint( # Lightning checkpoint format from slices.training import SSLPretrainModule - model = SSLPretrainModule.load_from_checkpoint(checkpoint_path) + model = SSLPretrainModule.load_from_checkpoint(checkpoint_path, weights_only=False) encoder = model.encoder else: raise ValueError(f"Unknown checkpoint format. Keys: {checkpoint.keys()}") diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py new file mode 100644 index 0000000..7a57de3 --- /dev/null +++ b/scripts/eval/evaluate_fairness.py @@ -0,0 +1,1327 @@ +#!/usr/bin/env python3 +"""Post-run fairness evaluation for SLICES experiment runs. + +Queries W&B for finished downstream runs, reconstructs the exact evaluation +artifact used for each run, runs inference on the test set, computes fairness +metrics via FairnessEvaluator, and writes results back to the same W&B run's +summary. + +Designed for batch evaluation of the benchmark fairness corpus across finetune, +supervised, and classical baseline runs. Supports resumability via +--skip-existing (default) and scoping via --experiment-class/--paradigm/--dataset filters. + +Usage: + # Evaluate the benchmark fairness corpus for one explicit revision + uv run python scripts/eval/evaluate_fairness.py \ + --project slices-thesis --revision thesis-v1 --entity + + # Scope to specific experiment class/dataset + uv run python scripts/eval/evaluate_fairness.py \ + --project slices-thesis --revision thesis-v1 --entity \ + --experiment-class core_ssl_benchmark --dataset miiv + + # Preview which runs would be evaluated + uv run python scripts/eval/evaluate_fairness.py \ + --project slices-thesis --revision thesis-v1 --entity --dry-run + + # Override paths (e.g., different machine than training) + uv run python scripts/eval/evaluate_fairness.py \ + --project slices-thesis --revision thesis-v1 --entity \ + --outputs-root /mnt/data/outputs --data-root /mnt/data + + # Recompute fairness for runs that already have metrics + uv run python scripts/eval/evaluate_fairness.py \ + --project slices-thesis --revision thesis-v1 --entity --force + + # Debug with a single run + uv run python scripts/eval/evaluate_fairness.py \ + --project slices-thesis --revision thesis-v1 --entity --max-runs 1 +""" +from __future__ import annotations + +import argparse +import json +import logging +import os +import re +import sys +import time +from pathlib import Path +from typing import Any, Optional + +import torch + +from slices.constants import ( + FULL_FINETUNE_PROTOCOL, + canonical_downstream_protocol, + downstream_protocol_from_freeze, +) +from slices.constants import ( + THESIS_TASKS as BENCHMARK_THESIS_TASKS, +) +from slices.eval.fairness_evaluator import flatten_fairness_report +from slices.eval.fairness_metadata import ( + EVAL_ARTIFACT_SHA256_KEY, + FAIRNESS_ARTIFACT_PATH_KEY, + FAIRNESS_ARTIFACT_SHA256_KEY, + FAIRNESS_ARTIFACT_SOURCE_KEY, + FAIRNESS_CHECKPOINT_SOURCE_KEY, + FAIRNESS_CLEAR_PREFIXES, + FAIRNESS_DEFAULT_MIN_SUBGROUP_SIZE, + FAIRNESS_DEFAULT_PROTECTED_ATTRIBUTES, + FAIRNESS_MIN_SUBGROUP_SIZE_KEY, + FAIRNESS_PROTECTED_ATTRIBUTES_KEY, + FAIRNESS_SCHEMA_VERSION_KEY, + FAIRNESS_SCRIPT_VERSION, + FAIRNESS_SCRIPT_VERSION_KEY, + FAIRNESS_SUMMARY_SCHEMA_VERSION, + canonical_artifact_id, + decode_protected_attributes, + encode_protected_attributes, + file_sha256, + normalize_protected_attributes, +) + +log = logging.getLogger("evaluate_fairness") + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +DEFAULT_EXPERIMENT_CLASSES = [ + "core_ssl_benchmark", + "label_efficiency", + "cross_dataset_transfer", + "hp_robustness", + "capacity_study", + "classical_baselines", + "ts2vec_extension", + "smart_external_reference", +] +DEFAULT_PHASES = ["finetune", "supervised", "baseline"] +DEFAULT_PROTECTED_ATTRIBUTES = FAIRNESS_DEFAULT_PROTECTED_ATTRIBUTES +THESIS_TASKS = set(BENCHMARK_THESIS_TASKS) +BINARY_FAIRNESS_REQUIRED_METRICS = [ + "n_valid_groups", + "n_metric_valid_groups", + "worst_group_auroc", + "worst_group_auprc", + "auroc_gap", + "auprc_gap", + "demographic_parity_diff", + "equalized_odds_diff", + "disparate_impact_ratio", +] +REGRESSION_FAIRNESS_REQUIRED_METRICS = [ + "n_valid_groups", + "worst_group_mse", + "worst_group_mae", + "mse_gap", + "mae_gap", +] +FAIRNESS_MATRIX_KEY_COLUMNS = [ + "experiment_class", + "experiment_type", + "experiment_subtype", + "paradigm", + "dataset", + "task", + "protocol", + "label_fraction", + "model_size", + "source_dataset", + "phase", + "upstream_pretrain_lr", + "upstream_pretrain_mask_ratio", + "seed", +] + + +# --------------------------------------------------------------------------- +# W&B helpers +# --------------------------------------------------------------------------- + + +def _retry(fn, max_retries=3, base_delay=5): + """Retry with exponential backoff on transient errors.""" + for attempt in range(max_retries + 1): + try: + return fn() + except Exception as e: + err_str = str(e).lower() + is_retryable = any( + kw in err_str + for kw in ["timeout", "timed out", "connection", "429", "500", "502", "503"] + ) + if not is_retryable or attempt == max_retries: + raise + delay = base_delay * (2**attempt) + log.warning("Retry %d/%d after %ds: %r", attempt + 1, max_retries, delay, e) + time.sleep(delay) + + +def _run_config_filter_value(run, key: str): + config = dict(getattr(run, "config", {}) or {}) + if key == "paradigm": + ssl_cfg = config.get("ssl") or {} + return config.get("paradigm") or ssl_cfg.get("name") + return config.get(key) + + +def _run_matches_filter(run, key: str, tag_prefix: str, values: Optional[list[str]]) -> bool: + """Match a W&B run by config first, falling back to tags.""" + if not values: + return True + + allowed = {str(value) for value in values} + config_value = _run_config_filter_value(run, key) + if config_value is not None: + return str(config_value) in allowed + + tags = set(getattr(run, "tags", []) or []) + return any(f"{tag_prefix}:{value}" in tags for value in allowed) + + +def fetch_eval_runs( + project: str, + entity: Optional[str], + experiment_classes: Optional[list[str]], + paradigms: Optional[list[str]], + datasets: Optional[list[str]], + phases: list[str], + revisions: Optional[list[str]], +) -> list: + """Fetch finished evaluation runs from W&B matching filters.""" + import wandb + + api = wandb.Api(timeout=300) + path = f"{entity}/{project}" if entity else project + + filters: dict = {"state": "finished"} + + log.info("Fetching runs from %s with filters: %s", path, json.dumps(filters, default=str)) + runs_iter = api.runs(path, filters=filters or {}, order="-created_at") + + runs = [] + for run in runs_iter: + if not _run_matches_filter(run, "experiment_class", "experiment_class", experiment_classes): + continue + if not _run_matches_filter(run, "paradigm", "paradigm", paradigms): + continue + if not _run_matches_filter(run, "dataset", "dataset", datasets): + continue + if not _run_matches_filter(run, "phase", "phase", phases): + continue + if not _run_matches_filter(run, "revision", "revision", revisions): + continue + + runs.append(run) + + log.info("Fetched %d runs.", len(runs)) + return runs + + +def _tag_value(tags: list[str], prefix: str) -> str | None: + full_prefix = f"{prefix}:" + for tag in tags: + if tag.startswith(full_prefix): + return tag.split(":", 1)[1] + return None + + +def _matrix_key_value(column: str, value): + if value is None or value == "": + return "__none__" + try: + if value != value: + return "__none__" + except Exception: + pass + if column == "seed": + return int(value) + if column in {"label_fraction", "upstream_pretrain_lr", "upstream_pretrain_mask_ratio"}: + return round(float(value), 12) + return str(value) + + +def _matrix_key(row: dict) -> tuple: + return tuple( + _matrix_key_value(column, row.get(column)) for column in FAIRNESS_MATRIX_KEY_COLUMNS + ) + + +def _format_matrix_key(key: tuple) -> str: + return ", ".join( + f"{column}={value}" for column, value in zip(FAIRNESS_MATRIX_KEY_COLUMNS, key, strict=True) + ) + + +def _run_matrix_row(run) -> dict[str, Any]: + """Extract the expected-matrix identity columns from a W&B run.""" + config = dict(getattr(run, "config", {}) or {}) + tags = list(getattr(run, "tags", []) or []) + phase = config.get("phase") or _tag_value(tags, "phase") + paradigm = config.get("paradigm") or _get_nested(config, "ssl.name") + if phase is None: + if paradigm in {"gru_d", "xgboost"}: + phase = "baseline" + elif paradigm == "supervised": + phase = "supervised" + else: + phase = "finetune" + + freeze = _get_nested(config, "training.freeze_encoder") + protocol = canonical_downstream_protocol(config.get("protocol") or _tag_value(tags, "protocol")) + if protocol is None: + protocol = downstream_protocol_from_freeze(freeze) + if protocol is None and phase == "baseline": + protocol = FULL_FINETUNE_PROTOCOL + + experiment_class = config.get("experiment_class") or _tag_value(tags, "experiment_class") + return { + "experiment_class": experiment_class, + "experiment_type": experiment_class, + "experiment_subtype": config.get("experiment_subtype"), + "paradigm": paradigm, + "dataset": config.get("dataset"), + "task": _get_nested(config, "task.task_name"), + "protocol": protocol, + "label_fraction": config.get("label_fraction", 1.0), + "model_size": config.get("model_size") or "default", + "source_dataset": config.get("source_dataset"), + "phase": phase, + "upstream_pretrain_lr": config.get("upstream_pretrain_lr"), + "upstream_pretrain_mask_ratio": config.get("upstream_pretrain_mask_ratio"), + "seed": config.get("seed"), + } + + +def fairness_matrix_coverage_issues( + runs: list, + experiment_classes: list[str] | None, + paradigms: list[str] | None, + datasets: list[str] | None, + phases: list[str] | None, + include_extension_tasks: bool = False, +) -> list[str]: + """Return missing/unexpected matrix identities in the fairness input corpus.""" + if include_extension_tasks: + return [] + + from scripts.export_results import build_expected_matrix_df + + expected_df = build_expected_matrix_df( + experiment_class=experiment_classes, + paradigm=paradigms, + dataset=datasets, + phase=phases, + ) + if not expected_df.empty and "task" in expected_df.columns: + expected_df = expected_df[expected_df["task"].isin(THESIS_TASKS)] + if expected_df.empty: + return [] + + expected_keys = {_matrix_key(row.to_dict()) for _, row in expected_df.iterrows()} + observed_rows = [_run_matrix_row(run) for run in runs] + observed_rows = [row for row in observed_rows if row.get("task") in THESIS_TASKS] + observed_keys = {_matrix_key(row) for row in observed_rows} + + missing = sorted(expected_keys - observed_keys) + unexpected = sorted(observed_keys - expected_keys) + issues: list[str] = [] + if missing: + issues.append( + f"fairness input is missing {len(missing)}/{len(expected_keys)} expected " + "downstream matrix rows" + ) + for key in missing[:20]: + issues.append(f" missing {_format_matrix_key(key)}") + if len(missing) > 20: + issues.append(f" ... {len(missing) - 20} more missing rows") + + if unexpected: + issues.append( + f"fairness input has {len(unexpected)} downstream rows outside the expected matrix" + ) + for key in unexpected[:20]: + issues.append(f" unexpected {_format_matrix_key(key)}") + if len(unexpected) > 20: + issues.append(f" ... {len(unexpected) - 20} more unexpected rows") + + return issues + + +def _expected_fairness_attributes(run, protected_attributes: list[str]) -> list[str]: + """Return requested fairness attributes that are meaningful for the run dataset.""" + dataset = str(run.config.get("dataset", "")).lower() + attrs = list(dict.fromkeys(protected_attributes)) + if dataset == "eicu": + attrs = [attr for attr in attrs if attr != "race"] + return attrs + + +def _has_summary_value(summary: dict[str, Any], key: str) -> bool: + """Return true when a fairness summary key was written, even if undefined.""" + return key in summary and summary.get(key) is not None + + +def _task_type_for_run(run) -> str: + """Infer fairness metric family from W&B config.""" + config = run.config or {} + task_type = _get_nested(config, "task.task_type") + if task_type: + return str(task_type).lower() + + task_name = _get_nested(config, "task.task_name", "") + if str(task_name).startswith("los"): + return "regression" + return "binary" + + +def _required_fairness_metrics_for_run(run) -> list[str]: + """Return required aggregate fairness metrics for this run's task type.""" + if _task_type_for_run(run) == "regression": + return REGRESSION_FAIRNESS_REQUIRED_METRICS + return BINARY_FAIRNESS_REQUIRED_METRICS + + +def _expected_fairness_artifact_source(run) -> str | None: + """Return the artifact-source metadata expected for a run.""" + paradigm = str((run.config or {}).get("paradigm", "")).lower() + if paradigm == "xgboost": + return "xgboost_model" + + eval_source = (run.summary_metrics or {}).get("_eval_checkpoint_source") + if eval_source == "best": + return "recorded_best" + if eval_source == "final": + return "recorded_final" + return None + + +def _expected_fairness_checkpoint_source(run) -> str | None: + """Return checkpoint-provenance metadata expected for a run.""" + paradigm = str((run.config or {}).get("paradigm", "")).lower() + if paradigm == "xgboost": + return "xgboost_model" + source = (run.summary_metrics or {}).get("_eval_checkpoint_source") + return str(source) if source in {"best", "final"} else None + + +def _expected_fairness_artifact_id(run) -> str | None: + """Return expected evaluated-artifact identity when it can be inferred.""" + config = run.config or {} + summary = run.summary_metrics or {} + output_dir = config.get("output_dir", "") + paradigm = str(config.get("paradigm", "")).lower() + + if paradigm == "xgboost": + if not output_dir: + return None + return canonical_artifact_id(Path(output_dir) / "xgboost_model.json") + + eval_source = summary.get("_eval_checkpoint_source") + if eval_source == "best": + best_path = summary.get("_best_ckpt_path") + return canonical_artifact_id(best_path) if best_path else None + if eval_source == "final" and output_dir: + return canonical_artifact_id(Path(output_dir) / "checkpoints" / "last.ckpt") + return None + + +def build_fairness_summary_metadata( + run, + artifact_path: Path, + artifact_source: str, + protected_attributes: list[str], + min_subgroup_size: int, +) -> dict[str, Any]: + """Build versioned metadata that makes fairness summaries auditable.""" + checkpoint_source = _expected_fairness_checkpoint_source(run) or artifact_source + return { + FAIRNESS_SCHEMA_VERSION_KEY: FAIRNESS_SUMMARY_SCHEMA_VERSION, + FAIRNESS_SCRIPT_VERSION_KEY: FAIRNESS_SCRIPT_VERSION, + FAIRNESS_ARTIFACT_PATH_KEY: str(artifact_path), + FAIRNESS_ARTIFACT_SHA256_KEY: file_sha256(artifact_path), + FAIRNESS_ARTIFACT_SOURCE_KEY: artifact_source, + FAIRNESS_CHECKPOINT_SOURCE_KEY: checkpoint_source, + FAIRNESS_PROTECTED_ATTRIBUTES_KEY: encode_protected_attributes(protected_attributes), + FAIRNESS_MIN_SUBGROUP_SIZE_KEY: int(min_subgroup_size), + } + + +def fairness_summary_metadata_issues( + run, + protected_attributes: list[str], + min_subgroup_size: int = FAIRNESS_DEFAULT_MIN_SUBGROUP_SIZE, +) -> list[str]: + """Return freshness problems for existing fairness summary metadata.""" + summary = run.summary_metrics or {} + issues: list[str] = [] + + if summary.get(FAIRNESS_SCHEMA_VERSION_KEY) != FAIRNESS_SUMMARY_SCHEMA_VERSION: + issues.append("missing or stale fairness schema version") + if summary.get(FAIRNESS_SCRIPT_VERSION_KEY) != FAIRNESS_SCRIPT_VERSION: + issues.append("missing or stale fairness script version") + + actual_attrs = decode_protected_attributes(summary.get(FAIRNESS_PROTECTED_ATTRIBUTES_KEY)) + expected_attrs = normalize_protected_attributes(protected_attributes) + if actual_attrs != expected_attrs: + issues.append( + f"protected attributes mismatch: expected={expected_attrs}, actual={actual_attrs}" + ) + + try: + actual_min_subgroup_size = int(float(summary.get(FAIRNESS_MIN_SUBGROUP_SIZE_KEY))) + except (TypeError, ValueError): + actual_min_subgroup_size = None + if actual_min_subgroup_size != int(min_subgroup_size): + issues.append( + "min subgroup size mismatch: " + f"expected={int(min_subgroup_size)}, actual={actual_min_subgroup_size}" + ) + + expected_artifact_source = _expected_fairness_artifact_source(run) + actual_artifact_source = summary.get(FAIRNESS_ARTIFACT_SOURCE_KEY) + if expected_artifact_source is not None and actual_artifact_source != expected_artifact_source: + issues.append( + "artifact source mismatch: " + f"expected={expected_artifact_source}, actual={actual_artifact_source}" + ) + + expected_checkpoint_source = _expected_fairness_checkpoint_source(run) + actual_checkpoint_source = summary.get(FAIRNESS_CHECKPOINT_SOURCE_KEY) + if ( + expected_checkpoint_source is not None + and actual_checkpoint_source != expected_checkpoint_source + ): + issues.append( + "checkpoint source mismatch: " + f"expected={expected_checkpoint_source}, actual={actual_checkpoint_source}" + ) + + actual_artifact_id = canonical_artifact_id(summary.get(FAIRNESS_ARTIFACT_PATH_KEY)) + expected_artifact_id = _expected_fairness_artifact_id(run) + if not actual_artifact_id: + issues.append("missing fairness artifact path") + elif expected_artifact_id is not None and actual_artifact_id != expected_artifact_id: + issues.append( + "artifact path mismatch: " + f"expected={expected_artifact_id}, actual={actual_artifact_id}" + ) + + actual_artifact_sha256 = summary.get(FAIRNESS_ARTIFACT_SHA256_KEY) + expected_artifact_sha256 = summary.get(EVAL_ARTIFACT_SHA256_KEY) + if not actual_artifact_sha256: + issues.append("missing fairness artifact sha256") + elif expected_artifact_sha256 and actual_artifact_sha256 != expected_artifact_sha256: + issues.append( + "artifact sha256 mismatch: " + f"expected={expected_artifact_sha256}, actual={actual_artifact_sha256}" + ) + elif expected_artifact_id is not None and not expected_artifact_sha256: + issues.append("missing evaluation artifact sha256") + + return issues + + +def has_fairness_metrics( + run, + protected_attributes: list[str], + min_subgroup_size: int = FAIRNESS_DEFAULT_MIN_SUBGROUP_SIZE, +) -> bool: + """Check whether all requested dataset-appropriate fairness keys exist. + + Uses ``summary_metrics`` (populated from the batch query) instead of + ``summary._json_dict`` which triggers a per-run GraphQL reload. Existing + summaries must also carry current metadata so stale fairness outputs are + recomputed instead of skipped. + """ + try: + sm = run.summary_metrics or {} + expected_attrs = _expected_fairness_attributes(run, protected_attributes) + if not expected_attrs: + return False + required_metrics = _required_fairness_metrics_for_run(run) + for attr in expected_attrs: + prefix = f"fairness/{attr}/" + for metric_name in required_metrics: + if not _has_summary_value(sm, f"{prefix}{metric_name}"): + return False + return not fairness_summary_metadata_issues(run, protected_attributes, min_subgroup_size) + except Exception: + return False + + +def missing_fairness_report_requirements( + run, + report: dict[str, Any], + protected_attributes: list[str], +) -> list[str]: + """Return missing requested fairness attributes or required aggregate metrics.""" + expected_attrs = _expected_fairness_attributes(run, protected_attributes) + required_metrics = _required_fairness_metrics_for_run(run) + missing: list[str] = [] + + for attr in expected_attrs: + if attr not in report: + missing.append(f"{attr}: no valid fairness groups") + continue + + flat = flatten_fairness_report({attr: report[attr]}) + prefix = f"fairness/{attr}/" + for metric_name in required_metrics: + key = f"{prefix}{metric_name}" + if not _has_summary_value(flat, key): + missing.append(f"{attr}: missing {metric_name}") + + return missing + + +def write_fairness_to_wandb( + run_path: str, + fairness_flat: dict[str, Any], + dry_run: bool = False, + clear_existing: bool = True, +) -> None: + """Write fairness metrics to W&B run summary. + + Existing fairness summaries are replaced by default so stale subgroup keys + from older script versions or settings cannot survive a recomputation and + leak into publication exports. + """ + import wandb + + if dry_run: + mode = "replace" if clear_existing else "write" + log.info( + " [DRY RUN] Would %s %d fairness keys on %s", + mode, + len(fairness_flat), + run_path, + ) + return + + def _do_update(): + api = wandb.Api(timeout=120) + run = api.run(run_path) + if clear_existing: + try: + existing_summary = dict(run.summary) + except (TypeError, ValueError): + existing_summary = dict(getattr(run.summary, "_json_dict", {}) or {}) + for key in existing_summary: + if key.startswith(FAIRNESS_CLEAR_PREFIXES) and key not in fairness_flat: + run.summary[key] = None + run.summary.update(fairness_flat) + run.summary.save() + + _retry(_do_update) + log.info(" Wrote %d fairness keys to W&B run %s", len(fairness_flat), run_path) + + +# --------------------------------------------------------------------------- +# Checkpoint discovery +# --------------------------------------------------------------------------- + + +def _resolve_ckpt_dir(output_dir: str, outputs_root: Optional[str] = None) -> Path: + """Resolve the checkpoint directory, applying outputs_root rebase if needed.""" + return _resolve_output_dir(output_dir, outputs_root) / "checkpoints" + + +def _resolve_output_dir(output_dir: str, outputs_root: Optional[str] = None) -> Path: + """Resolve a run output directory, applying outputs_root rebase if needed.""" + if outputs_root: + rel = output_dir + if rel.startswith("outputs/"): + rel = rel[len("outputs/") :] + elif rel.startswith("/") and "/outputs/" in rel: + rel = rel.split("/outputs/", 1)[1] + return Path(outputs_root) / rel + return Path(output_dir) + + +def _resolve_logged_checkpoint_path( + checkpoint_path: str, + outputs_root: Optional[str] = None, +) -> Path: + """Resolve a checkpoint path recorded in W&B summary metadata.""" + path = Path(checkpoint_path) + if outputs_root is None: + return path + + checkpoint_str = str(checkpoint_path) + if checkpoint_str.startswith("outputs/"): + return Path(outputs_root) / checkpoint_str[len("outputs/") :] + if checkpoint_str.startswith("/") and "/outputs/" in checkpoint_str: + return Path(outputs_root) / checkpoint_str.split("/outputs/", 1)[1] + if path.is_absolute(): + return path + return Path(outputs_root) / checkpoint_str + + +def find_best_checkpoint( + output_dir: str, + outputs_root: Optional[str] = None, + task_type: str = "binary", +) -> Optional[Path]: + """Find the best .ckpt file in the run's checkpoint directory. + + Search strategy: + 1. Look for metric-named .ckpt files directly in checkpoints/ (standard case) + 2. Look for .ckpt files in subdirectories (Lightning creates subdirs when the + monitor metric contains '/' — e.g. val/auprc becomes a directory separator) + 3. Fall back to last.ckpt + + Args: + output_dir: Run output directory (from W&B config). + outputs_root: Optional rebase root for checkpoint paths. + task_type: Task type — determines whether higher or lower metric + values are better. Classification monitors val/auprc (higher=better), + regression monitors val/mse (lower=better). + """ + ckpt_dir = _resolve_ckpt_dir(output_dir, outputs_root) + + if not ckpt_dir.exists(): + log.warning(" Checkpoint dir not found: %s", ckpt_dir) + return None + + # Strategy 1: Direct .ckpt files (excluding last*.ckpt variants) + ckpts = [p for p in ckpt_dir.glob("*.ckpt") if not p.name.startswith("last")] + + # Strategy 2: .ckpt files in subdirectories (Lightning '/' in metric name issue) + if not ckpts: + ckpts = list(ckpt_dir.glob("*/*.ckpt")) + + # Parse metric values from filenames if we found any candidates + if ckpts: + metric_pattern = re.compile(r"[-=](\d+\.\d+)\.ckpt$") + scored = [] + for p in ckpts: + match = metric_pattern.search(p.name) + if match: + scored.append((float(match.group(1)), p)) + + if scored: + # Regression: lower is better (val/mse); classification: higher (val/auprc) + pick_lowest = task_type == "regression" + scored.sort(key=lambda x: x[0], reverse=not pick_lowest) + best = scored[0][1] + log.info(" Best checkpoint: %s (metric=%.4f)", best.name, scored[0][0]) + return best + + # Non-last checkpoints exist but no parseable metric — pick most recent + ckpts.sort(key=lambda p: p.stat().st_mtime, reverse=True) + log.info(" Using most recent checkpoint: %s", ckpts[0].name) + return ckpts[0] + + # Strategy 3: Fall back to last.ckpt + last = ckpt_dir / "last.ckpt" + if last.exists(): + log.info(" Using last.ckpt (no best checkpoint found)") + return last + + log.warning(" No checkpoints found in %s", ckpt_dir) + return None + + +def resolve_evaluation_checkpoint( + run, + outputs_root: Optional[str] = None, + task_type: str = "binary", +) -> tuple[Optional[Path], str]: + """Resolve the checkpoint that was actually used for logged test metrics. + + Uses the run's recorded evaluation provenance when present: + - `_eval_checkpoint_source=best` -> recorded `_best_ckpt_path` + - `_eval_checkpoint_source=final` -> `last.ckpt` + + Fails closed for runs that lack recorded provenance entirely. + """ + summary = dict(run.summary_metrics or {}) + output_dir = run.config.get("output_dir", "") + eval_source = summary.get("_eval_checkpoint_source") + + if eval_source == "best": + best_path = summary.get("_best_ckpt_path", "") + if not best_path: + raise FileNotFoundError( + "Run recorded _eval_checkpoint_source=best but did not persist _best_ckpt_path." + ) + ckpt_path = _resolve_logged_checkpoint_path(best_path, outputs_root) + if not ckpt_path.exists(): + raise FileNotFoundError( + f"Recorded best checkpoint not found: {ckpt_path} " + f"(from summary path {best_path!r})" + ) + return ckpt_path, "recorded_best" + + if eval_source == "final": + ckpt_path = _resolve_ckpt_dir(output_dir, outputs_root) / "last.ckpt" + if not ckpt_path.exists(): + raise FileNotFoundError( + f"Run recorded final-model evaluation but last.ckpt was not found at {ckpt_path}" + ) + return ckpt_path, "recorded_final" + + if eval_source == "failed": + best_path = summary.get("_best_ckpt_path", "") + error = summary.get("_best_ckpt_error", "unknown error") + raise RuntimeError( + "Training recorded a checkpoint-selection failure " + f"(best_ckpt={best_path!r}, error={error!r})." + ) + + raise RuntimeError( + "Run " + f"{run.id} lacks recorded checkpoint provenance (_eval_checkpoint_source). " + "Fairness evaluation now requires explicit provenance so it cannot silently " + "re-evaluate a different checkpoint than the logged test metrics." + ) + + +def resolve_evaluation_artifact( + run, + outputs_root: Optional[str] = None, + task_type: str = "binary", +) -> tuple[Path, str]: + """Resolve the saved artifact needed to reproduce a run's test-time predictions.""" + paradigm = str(run.config.get("paradigm", "")).lower() + if paradigm == "xgboost": + model_path = _resolve_output_dir(run.config.get("output_dir", ""), outputs_root) + model_path = model_path / "xgboost_model.json" + if not model_path.exists(): + raise FileNotFoundError(f"Saved XGBoost model not found: {model_path}") + return model_path, "xgboost_model" + + ckpt_path, ckpt_source = resolve_evaluation_checkpoint(run, outputs_root, task_type) + if ckpt_path is None: + raise FileNotFoundError(f"Could not resolve evaluation checkpoint for run {run.id}") + return ckpt_path, ckpt_source + + +def validate_evaluation_artifact_digest(run, artifact_path: Path) -> None: + """Fail before inference when the local artifact differs from logged test metadata.""" + summary = dict(run.summary_metrics or {}) + expected_id = _expected_fairness_artifact_id(run) + expected_sha256 = summary.get(EVAL_ARTIFACT_SHA256_KEY) + if expected_id is not None and not expected_sha256: + raise RuntimeError( + f"Run {run.id} is missing {EVAL_ARTIFACT_SHA256_KEY}; cannot verify " + "that fairness inference uses the artifact evaluated for test metrics." + ) + + if not expected_sha256: + return + + actual_sha256 = file_sha256(artifact_path) + if actual_sha256 != expected_sha256: + raise RuntimeError( + "Evaluation artifact sha256 mismatch before fairness inference: " + f"expected={expected_sha256}, actual={actual_sha256}, path={artifact_path}" + ) + + +# --------------------------------------------------------------------------- +# Model + data reconstruction +# --------------------------------------------------------------------------- + + +def _get_nested(config: dict, dotted_key: str, default=None): + """Get a value from a nested dict using a dotted key path.""" + parts = dotted_key.split(".") + val = config + for p in parts: + if isinstance(val, dict) and p in val: + val = val[p] + else: + return default + return val + + +def _run_task_name(run) -> str | None: + return _get_nested(run.config or {}, "task.task_name") + + +def filter_thesis_task_runs(runs: list) -> list: + """Keep only fixed thesis tasks for publication fairness evaluation.""" + return [run for run in runs if _run_task_name(run) in THESIS_TASKS] + + +def build_datamodule( + wandb_config: dict, + batch_size: int = 64, + data_root: Optional[str] = None, +): + """Build ICUDataModule from W&B run config.""" + from slices.data.datamodule import ICUDataModule + + processed_dir = _get_nested(wandb_config, "data.processed_dir", "") + if data_root: + parts = Path(processed_dir).parts + if len(parts) >= 2 and parts[-2] == "processed": + processed_dir = str(Path(data_root) / parts[-2] / parts[-1]) + else: + processed_dir = str(Path(data_root) / parts[-1]) + + task_name = _get_nested(wandb_config, "task.task_name", "mortality_24h") + seed = wandb_config.get("seed", 42) + label_fraction = wandb_config.get("label_fraction", 1.0) + + datamodule = ICUDataModule( + processed_dir=processed_dir, + task_name=task_name, + batch_size=batch_size, + num_workers=4, + seed=seed, + label_fraction=label_fraction, + ) + datamodule.setup() + return datamodule + + +def build_model(wandb_config: dict, checkpoint_path: Path, datamodule): + """Build FineTuneModule from W&B config and load checkpoint weights.""" + from omegaconf import OmegaConf + + from slices.training import FineTuneModule + + cfg = OmegaConf.create(wandb_config) + OmegaConf.set_struct(cfg, False) + cfg.encoder.d_input = datamodule.get_feature_dim() + cfg.encoder.max_seq_length = datamodule.get_seq_length() + + class_weight = _get_nested(wandb_config, "training.class_weight") + if class_weight == "balanced" or class_weight is None: + cfg.training.class_weight = None + OmegaConf.set_struct(cfg, True) + + model = FineTuneModule( + config=cfg, + checkpoint_path=None, + pretrain_checkpoint_path=None, + ) + + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + model.load_state_dict(checkpoint["state_dict"]) + return model + + +# --------------------------------------------------------------------------- +# Fairness evaluation +# --------------------------------------------------------------------------- + + +def evaluate_run_fairness( + model: torch.nn.Module, + datamodule, + protected_attributes: list[str], + min_subgroup_size: int, + device: str, +) -> dict[str, Any]: + """Run fairness evaluation on a single run.""" + from slices.eval.inference import run_inference + + model = model.to(device) + predictions, labels, stay_ids = run_inference( + model, + datamodule.test_dataloader(), + device=device, + ) + + task_type = getattr(model, "task_type", "binary") + report = evaluate_predictions_fairness( + predictions, + labels, + stay_ids, + datamodule, + protected_attributes, + min_subgroup_size, + task_type, + ) + return report + + +def evaluate_predictions_fairness( + predictions: torch.Tensor, + labels: torch.Tensor, + stay_ids: list[int], + datamodule, + protected_attributes: list[str], + min_subgroup_size: int, + task_type: str, +) -> dict[str, Any]: + """Run fairness evaluation from materialized predictions and labels.""" + from slices.eval.fairness_evaluator import FairnessEvaluator + + evaluator = FairnessEvaluator( + static_df=datamodule.dataset.static_df, + protected_attributes=protected_attributes, + min_subgroup_size=min_subgroup_size, + task_type=task_type, + dataset_name=getattr(getattr(datamodule, "processed_dir", None), "name", None), + ) + return evaluator.evaluate(predictions, labels, stay_ids) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def _resolve_device(device_str: str) -> str: + if device_str == "auto": + if torch.cuda.is_available(): + return "cuda" + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return "mps" + return "cpu" + return device_str + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Post-run fairness evaluation for SLICES experiment runs.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, + ) + parser.add_argument( + "--project", + default=os.environ.get("WANDB_PROJECT", "slices-thesis"), + help="W&B project name (default: $WANDB_PROJECT or 'slices-thesis')", + ) + parser.add_argument( + "--entity", + default=os.environ.get("WANDB_ENTITY"), + help="W&B entity (default: $WANDB_ENTITY)", + ) + parser.add_argument( + "--experiment-class", + nargs="+", + help="Filter to experiment class(es). Default: downstream-producing benchmark classes", + ) + parser.add_argument("--paradigm", nargs="+", help="Filter to paradigm(s)") + parser.add_argument("--dataset", nargs="+", help="Filter to dataset(s)") + parser.add_argument( + "--revision", + nargs="+", + help=( + "Filter to revision tag(s). Required unless REVISION or " + "WANDB_REVISION is set in the environment." + ), + ) + parser.add_argument( + "--phase", + nargs="+", + default=DEFAULT_PHASES, + help=f"Filter to phase(s) (default: {DEFAULT_PHASES})", + ) + parser.add_argument( + "--outputs-root", + default=None, + help="Override outputs directory root (rebase checkpoint paths)", + ) + parser.add_argument( + "--data-root", + default=None, + help="Override data directory root (rebase data.processed_dir)", + ) + parser.add_argument("--batch-size", type=int, default=64, help="Inference batch size") + parser.add_argument("--device", default="auto", help="Device for inference (auto/cpu/cuda/mps)") + parser.add_argument( + "--protected-attributes", + nargs="+", + default=DEFAULT_PROTECTED_ATTRIBUTES, + help=f"Attributes to evaluate (default: {DEFAULT_PROTECTED_ATTRIBUTES})", + ) + parser.add_argument( + "--min-subgroup-size", + type=int, + default=FAIRNESS_DEFAULT_MIN_SUBGROUP_SIZE, + help="Min patients per subgroup", + ) + parser.add_argument("--dry-run", action="store_true", help="List runs without processing") + parser.add_argument( + "--include-extension-tasks", + action="store_true", + help="Include task runs outside the fixed thesis task set.", + ) + parser.add_argument( + "--allow-incomplete", + action="store_true", + help="Exit successfully even if no runs match, fairness groups are skipped, or runs fail.", + ) + parser.add_argument( + "--skip-existing", + action="store_true", + default=True, + help="Skip runs that already have fairness metrics (default: True)", + ) + parser.add_argument( + "--force", + action="store_true", + help="Recompute even if fairness metrics already exist", + ) + parser.add_argument("--max-runs", type=int, default=None, help="Limit runs (for debugging)") + parser.add_argument("-v", "--verbose", action="store_true", help="Verbose logging") + args = parser.parse_args() + + if not args.revision: + env_revision = os.environ.get("REVISION") or os.environ.get("WANDB_REVISION") + if env_revision: + args.revision = [env_revision] + else: + parser.error( + "--revision is required to avoid mixing reruns. " + "Pass --revision or set REVISION/WANDB_REVISION." + ) + + return args + + +def main() -> None: + args = parse_args() + + logging.basicConfig( + level=logging.DEBUG if args.verbose else logging.INFO, + format="%(asctime)s %(name)s %(levelname)s: %(message)s", + datefmt="%H:%M:%S", + ) + + device = _resolve_device(args.device) + log.info("Device: %s", device) + + # Default to the benchmark fairness classes if not specified. + experiment_classes = args.experiment_class or DEFAULT_EXPERIMENT_CLASSES + log.info("Experiment classes: %s", experiment_classes) + + # Fetch runs from W&B + runs = _retry( + lambda: fetch_eval_runs( + project=args.project, + entity=args.entity, + experiment_classes=experiment_classes, + paradigms=args.paradigm, + datasets=args.dataset, + phases=args.phase, + revisions=args.revision, + ) + ) + + if not args.include_extension_tasks: + before = len(runs) + runs = filter_thesis_task_runs(runs) + dropped = before - len(runs) + if dropped: + log.info("Dropped %d runs outside thesis tasks: %s", dropped, sorted(THESIS_TASKS)) + + if not runs: + print("No runs found matching filters.", file=sys.stderr) + sys.exit(0 if args.allow_incomplete else 1) + + coverage_issues = fairness_matrix_coverage_issues( + runs, + experiment_classes=experiment_classes, + paradigms=args.paradigm, + datasets=args.dataset, + phases=args.phase, + include_extension_tasks=args.include_extension_tasks, + ) + if coverage_issues: + print("Fairness matrix coverage validation failed:", file=sys.stderr) + for issue in coverage_issues: + print(issue, file=sys.stderr) + if not args.allow_incomplete: + sys.exit(1) + + # Filter out runs that already have fairness metrics (unless --force) + if args.skip_existing and not args.force: + before = len(runs) + runs = [ + r + for r in runs + if not has_fairness_metrics( + r, + args.protected_attributes, + args.min_subgroup_size, + ) + ] + skipped = before - len(runs) + if skipped: + log.info("Skipped %d runs with fresh existing fairness metrics.", skipped) + if not runs: + log.info("No pending runs after --skip-existing filtering.") + return + + if args.max_runs: + runs = runs[: args.max_runs] + + # Sort by (dataset, task, seed) to maximize datamodule reuse + def _sort_key(r): + cfg = r.config + return ( + cfg.get("dataset", ""), + _get_nested(cfg, "task.task_name", ""), + cfg.get("seed", 0), + ) + + runs.sort(key=_sort_key) + + print(f"\nRuns to evaluate: {len(runs)}") + if args.dry_run: + print("\n[DRY RUN] Listing runs:\n") + for i, r in enumerate(runs): + cfg = r.config + task_type = _get_nested(cfg, "task.task_type", "binary") + try: + artifact, artifact_source = resolve_evaluation_artifact( + r, args.outputs_root, task_type + ) + ckpt_str = f"{artifact} [{artifact_source}]" + except Exception as e: + ckpt_str = f"ERROR: {e}" + print( + f" {i + 1:3d}. {r.name or r.id} " + f"[{cfg.get('dataset', '?')}/{_get_nested(cfg, 'task.task_name', '?')}/" + f"seed{cfg.get('seed', '?')}] " + f"ckpt={ckpt_str}" + ) + print(f"\nTotal: {len(runs)} runs") + return + + # Process runs + results = {"processed": 0, "skipped": 0, "failed": 0, "errors": []} + prev_dm_key: Optional[tuple] = None + datamodule = None + + for i, run in enumerate(runs): + cfg = run.config + run_desc = run.name or run.id + ds = cfg.get("dataset", "?") + task = _get_nested(cfg, "task.task_name", "?") + seed = cfg.get("seed", "?") + log.info("[%d/%d] %s (%s/%s/seed%s)", i + 1, len(runs), run_desc, ds, task, seed) + + try: + # 1. Resolve the saved evaluation artifact + task_type = _get_nested(cfg, "task.task_type", "binary") + artifact_path, artifact_source = resolve_evaluation_artifact( + run, args.outputs_root, task_type + ) + validate_evaluation_artifact_digest(run, artifact_path) + log.info(" Evaluation artifact: %s (%s)", artifact_path, artifact_source) + + # 2. Reconstruct model + data (reuse datamodule if same dataset/task/seed) + dm_key = (ds, task, seed, cfg.get("label_fraction", 1.0)) + if dm_key != prev_dm_key: + if datamodule is not None: + del datamodule + datamodule = build_datamodule(cfg, args.batch_size, args.data_root) + prev_dm_key = dm_key + + paradigm = str(cfg.get("paradigm", "")).lower() + model = None + if paradigm == "xgboost": + from slices.eval.inference import run_xgboost_inference + + predictions, labels, stay_ids = run_xgboost_inference( + artifact_path, + task_type, + datamodule.dataset, + datamodule.test_indices, + ) + else: + from slices.eval.inference import run_inference + + model = build_model(cfg, artifact_path, datamodule) + model = model.to(device) + predictions, labels, stay_ids = run_inference( + model, + datamodule.test_dataloader(), + device=device, + ) + + # 3. Evaluate fairness + report = evaluate_predictions_fairness( + predictions, + labels, + stay_ids, + datamodule, + args.protected_attributes, + args.min_subgroup_size, + task_type, + ) + + if not report: + log.warning(" No fairness results (no valid attribute groups)") + results["skipped"] += 1 + continue + + missing_requirements = missing_fairness_report_requirements( + run, + report, + args.protected_attributes, + ) + if missing_requirements: + message = "Incomplete fairness results: " + "; ".join(missing_requirements) + log.warning(" %s", message) + results["skipped"] += 1 + results["errors"].append((run.id, run_desc, message)) + continue + + # 4. Flatten and write back + fairness_flat = flatten_fairness_report(report) + fairness_flat.update( + build_fairness_summary_metadata( + run, + artifact_path, + artifact_source, + args.protected_attributes, + args.min_subgroup_size, + ) + ) + run_path = ( + f"{args.entity}/{args.project}/{run.id}" + if args.entity + else f"{args.project}/{run.id}" + ) + write_fairness_to_wandb( + run_path, + fairness_flat, + args.dry_run, + ) + results["processed"] += 1 + + # Free model memory + if model is not None: + del model + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + except Exception as e: + results["failed"] += 1 + results["errors"].append((run.id, run_desc, str(e))) + log.error(" FAILED: %s", e, exc_info=args.verbose) + continue + + # Summary + print("\n" + "=" * 60) + print("Fairness Evaluation Summary") + print("=" * 60) + print(f" Processed: {results['processed']}") + print(f" Skipped: {results['skipped']}") + print(f" Failed: {results['failed']}") + if results["errors"]: + print("\n Errors:") + for entry in results["errors"]: + run_id, run_name, err = entry + print(f" {run_name} ({run_id}): {err}") + print("=" * 60) + + if not args.allow_incomplete and (results["failed"] > 0 or results["skipped"] > 0): + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/scripts/eval/evaluate_imputation.py b/scripts/eval/evaluate_imputation.py index 3ee71a5..799a09e 100644 --- a/scripts/eval/evaluate_imputation.py +++ b/scripts/eval/evaluate_imputation.py @@ -7,12 +7,12 @@ # Evaluate MAE checkpoint with all strategies uv run python scripts/eval/evaluate_imputation.py \ checkpoint=outputs/pretrain/encoder.pt \ - data.processed_dir=data/processed/mimic-iv + data.processed_dir=data/processed/miiv - # Evaluate full pretrain checkpoint (MAE with decoder) + # Evaluate a full pretrain checkpoint (reuses the encoder, trains a probe decoder) uv run python scripts/eval/evaluate_imputation.py \ pretrain_checkpoint=outputs/pretrain/ssl-last.ckpt \ - data.processed_dir=data/processed/mimic-iv + data.processed_dir=data/processed/miiv # Single strategy uv run python scripts/eval/evaluate_imputation.py \ @@ -20,12 +20,18 @@ masking.strategies=[random] """ +import json +from pathlib import Path + import hydra import lightning.pytorch as pl import torch from omegaconf import DictConfig, OmegaConf -from slices.data.datamodule import ICUDataModule +from torch.utils.data import DataLoader, Subset + +from slices.data.datamodule import ICUDataModule, icu_collate_fn from slices.eval.imputation import ImputationEvaluator +from slices.training.utils import validate_data_prerequisites @hydra.main( @@ -46,6 +52,8 @@ def main(cfg: DictConfig) -> None: print("\nConfiguration:") print(OmegaConf.to_yaml(cfg)) + validate_data_prerequisites(cfg.data.processed_dir, cfg.dataset) + # Set random seed pl.seed_everything(cfg.seed, workers=True) @@ -100,19 +108,20 @@ def main(cfg: DictConfig) -> None: ) # ========================================================================= - # 3. Train decoder (if not MAE) + # 3. Train reconstruction decoder probe # ========================================================================= - if not pretrain_ckpt: - print("\n" + "=" * 80) - print("3. Training Reconstruction Decoder") - print("=" * 80) - - recon_cfg = cfg.get("reconstruction_head", {}) - evaluator.train_decoder( - dataloader=datamodule.train_dataloader(), - max_epochs=recon_cfg.get("max_epochs", 10), - lr=recon_cfg.get("lr", 1e-3), - ) + print("\n" + "=" * 80) + print("3. Training Reconstruction Decoder") + print("=" * 80) + if pretrain_ckpt: + print(" Using the MAE checkpoint to initialize the encoder, then fitting a probe decoder.") + + recon_cfg = cfg.get("reconstruction_head", {}) + evaluator.train_decoder( + dataloader=datamodule.train_dataloader(), + max_epochs=recon_cfg.get("max_epochs", 10), + lr=recon_cfg.get("lr", 1e-3), + ) # ========================================================================= # 4. Evaluate @@ -124,6 +133,16 @@ def main(cfg: DictConfig) -> None: strategies = cfg.masking.strategies mask_ratio = cfg.masking.mask_ratio test_loader = datamodule.test_dataloader() + train_stats_loader = DataLoader( + Subset(datamodule.dataset, datamodule.train_indices), + batch_size=cfg.get("batch_size", 64), + shuffle=False, + num_workers=cfg.data.get("num_workers", 4), + pin_memory=datamodule.pin_memory, + collate_fn=icu_collate_fn, + drop_last=False, + ) + train_feature_stds = evaluator.compute_feature_stds(train_stats_loader) # Setup W&B logger if configured logger = None @@ -146,7 +165,12 @@ def main(cfg: DictConfig) -> None: print(f"\n Strategy: {strategy} (mask_ratio={mask_ratio})") print(" " + "-" * 40) - results = evaluator.evaluate(test_loader, mask_strategy=strategy, mask_ratio=mask_ratio) + results = evaluator.evaluate( + test_loader, + mask_strategy=strategy, + mask_ratio=mask_ratio, + feature_stds=train_feature_stds, + ) all_results[strategy] = results print(f" NRMSE overall: {results['nrmse_overall']:.4f}") @@ -186,10 +210,22 @@ def main(cfg: DictConfig) -> None: f"{results['mae_overall']:>10.4f}" ) + output_dir = Path(cfg.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + results_path = output_dir / "imputation_results.json" + payload = { + "checkpoint": encoder_ckpt, + "pretrain_checkpoint": pretrain_ckpt, + "config": OmegaConf.to_container(cfg, resolve=True), + "results": all_results, + } + results_path.write_text(json.dumps(payload, indent=2, sort_keys=True, default=str)) + if logger: wandb.finish() print(f"\n Output directory: {cfg.output_dir}") + print(f" Results: {results_path}") print("\n" + "=" * 80) diff --git a/scripts/export_results.py b/scripts/export_results.py new file mode 100644 index 0000000..a3452a4 --- /dev/null +++ b/scripts/export_results.py @@ -0,0 +1,2019 @@ +#!/usr/bin/env python3 +"""Canonical class-based results export pipeline for SLICES. + +Pulls final rerun-corpus W&B runs, extracts config plus test metrics, and writes +publication-oriented parquet files: + + - results/per_seed_results.parquet + - results/aggregated_results.parquet + - results/statistical_tests.parquet + - results/label_efficiency_curves.parquet + - results/capacity_study_comparison.parquet + - results/classical_context.parquet + - results/ts2vec_vs_core_contrastive.parquet + +Aggregated metric columns include mean, standard deviation, min/max, and 95% +confidence intervals across finite seed values. + +Usage: + uv run python scripts/export_results.py \ + --project slices-thesis --revision thesis-v1 --entity + uv run python scripts/export_results.py \ + --project slices-thesis --revision thesis-v1 --entity \ + --experiment-class core_ssl_benchmark label_efficiency +""" +from __future__ import annotations + +import argparse +import json +import math +import os +import re +import sys +import time +from itertools import combinations +from pathlib import Path + +import pandas as pd +import wandb +from scipy import stats as scipy_stats + +from slices.constants import ( + FULL_FINETUNE_PROTOCOL, + canonical_downstream_protocol, + downstream_protocol_from_freeze, +) +from slices.constants import ( + THESIS_TASKS as BENCHMARK_THESIS_TASKS, +) +from slices.eval.fairness_metadata import ( + EVAL_ARTIFACT_PATH_KEY, + EVAL_ARTIFACT_SHA256_KEY, + FAIRNESS_ARTIFACT_PATH_KEY, + FAIRNESS_ARTIFACT_SHA256_KEY, + FAIRNESS_ARTIFACT_SOURCE_KEY, + FAIRNESS_CHECKPOINT_SOURCE_KEY, + FAIRNESS_DEFAULT_MIN_SUBGROUP_SIZE, + FAIRNESS_DEFAULT_PROTECTED_ATTRIBUTES, + FAIRNESS_METADATA_COLUMNS, + FAIRNESS_MIN_SUBGROUP_SIZE_KEY, + FAIRNESS_PROTECTED_ATTRIBUTES_KEY, + FAIRNESS_SCHEMA_VERSION_KEY, + FAIRNESS_SCRIPT_VERSION, + FAIRNESS_SCRIPT_VERSION_KEY, + FAIRNESS_SUMMARY_SCHEMA_VERSION, + canonical_artifact_id, + decode_protected_attributes, + normalize_protected_attributes, +) +from slices.eval.statistical import ( + bonferroni_correction, + cohens_d, + paired_wilcoxon_signed_rank, +) + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +EXPERIMENT_CLASSES = [ + "core_ssl_benchmark", + "label_efficiency", + "cross_dataset_transfer", + "hp_robustness", + "capacity_study", + "classical_baselines", + "ts2vec_extension", + "smart_external_reference", +] + +CONTEXTUAL_STAT_CLASSES = { + "classical_context_full", + "classical_context_label_efficiency", +} + +FIXED_SEED_EXPERIMENT_CLASSES = set(EXPERIMENT_CLASSES) +EXPECTED_FIXED_SEEDS = {42, 123, 456, 789, 1011} +CAPACITY_LABEL_FRACTIONS = {0.05, 0.1, 0.5} +CAPACITY_PARADIGMS = {"mae", "supervised"} +THESIS_TASKS = set(BENCHMARK_THESIS_TASKS) + +FINGERPRINT = [ + "experiment_class", + "experiment_type", + "experiment_subtype", + "paradigm", + "dataset", + "task", + "protocol", + "label_fraction", + "model_size", + "source_dataset", + "phase", + "upstream_pretrain_lr", + "upstream_pretrain_mask_ratio", +] + +FINGERPRINT_AUDIT_COLUMNS = [ + *FINGERPRINT, + "lr", + "mask_ratio", +] + +TEST_METRICS = [ + "test/auroc", + "test/auprc", + "test/accuracy", + "test/f1", + "test/precision", + "test/recall", + "test/specificity", + "test/brier_score", + "test/ece", + "test/mse", + "test/mae", + "test/r2", + "test/loss", + "fairness/gender/worst_group_auroc", + "fairness/gender/worst_group_auprc", + "fairness/gender/auroc_gap", + "fairness/gender/auprc_gap", + "fairness/gender/demographic_parity_diff", + "fairness/gender/equalized_odds_diff", + "fairness/gender/disparate_impact_ratio", + "fairness/gender/worst_group_mse", + "fairness/gender/worst_group_mae", + "fairness/gender/mse_gap", + "fairness/gender/mae_gap", + "fairness/gender/n_valid_groups", + "fairness/gender/n_metric_valid_groups", + "fairness/age_group/worst_group_auroc", + "fairness/age_group/worst_group_auprc", + "fairness/age_group/auroc_gap", + "fairness/age_group/auprc_gap", + "fairness/age_group/demographic_parity_diff", + "fairness/age_group/equalized_odds_diff", + "fairness/age_group/disparate_impact_ratio", + "fairness/age_group/worst_group_mse", + "fairness/age_group/worst_group_mae", + "fairness/age_group/mse_gap", + "fairness/age_group/mae_gap", + "fairness/age_group/n_valid_groups", + "fairness/age_group/n_metric_valid_groups", + "fairness/race/worst_group_auroc", + "fairness/race/worst_group_auprc", + "fairness/race/auroc_gap", + "fairness/race/auprc_gap", + "fairness/race/demographic_parity_diff", + "fairness/race/equalized_odds_diff", + "fairness/race/disparate_impact_ratio", + "fairness/race/worst_group_mse", + "fairness/race/worst_group_mae", + "fairness/race/mse_gap", + "fairness/race/mae_gap", + "fairness/race/n_valid_groups", + "fairness/race/n_metric_valid_groups", +] + +VAL_METRICS = [ + "val/auroc", + "val/auprc", + "val/loss", +] + +ALL_METRICS = TEST_METRICS + VAL_METRICS +PERFORMANCE_TEST_METRICS = [metric for metric in TEST_METRICS if metric.startswith("test/")] + +MODEL_VARIANTS = { + (64, 2): "default", + (128, 4): "medium", + (256, 4): "large", +} + +PRIMARY_TEST_METRIC_BY_TASK = { + "mortality_24h": "test/auprc", + "mortality_hospital": "test/auprc", + "aki_kdigo": "test/auprc", + "los_remaining": "test/mae", +} + +LOWER_IS_BETTER_METRICS = { + "test/loss", + "test/mae", + "test/mse", +} + +EVAL_PHASES = ["finetune", "supervised", "baseline"] +FAIRNESS_ATTRIBUTES = FAIRNESS_DEFAULT_PROTECTED_ATTRIBUTES +FAIRNESS_SUMMARY_KEY_COLUMN = "_fairness_summary_keys" +BINARY_FAIRNESS_REQUIRED_METRICS = [ + "n_valid_groups", + "n_metric_valid_groups", + "worst_group_auroc", + "worst_group_auprc", + "auroc_gap", + "auprc_gap", + "demographic_parity_diff", + "equalized_odds_diff", + "disparate_impact_ratio", +] +REGRESSION_FAIRNESS_REQUIRED_METRICS = [ + "n_valid_groups", + "worst_group_mse", + "worst_group_mae", + "mse_gap", + "mae_gap", +] + + +# --------------------------------------------------------------------------- +# W&B Fetching +# --------------------------------------------------------------------------- + + +def _run_config_filter_value(run, key: str): + config = dict(getattr(run, "config", {}) or {}) + if key == "paradigm": + ssl_cfg = config.get("ssl") or {} + return config.get("paradigm") or ssl_cfg.get("name") + return config.get(key) + + +def _run_matches_filter(run, key: str, tag_prefix: str, values: list[str] | None) -> bool: + """Match a W&B run by config first, falling back to tags. + + Some historical runs have complete config metadata but incomplete tags. Fetch + broadly and apply this client-side predicate so publication exports do not + silently drop runs before config fallback can see them. + """ + if not values: + return True + + allowed = {str(value) for value in values} + config_value = _run_config_filter_value(run, key) + if config_value is not None: + return str(config_value) in allowed + + tags = set(getattr(run, "tags", []) or []) + return any(f"{tag_prefix}:{value}" in tags for value in allowed) + + +def fetch_all_runs( + project: str, + entity: str | None = None, + state: str = "finished", + experiment_class: list[str] | None = None, + paradigm: list[str] | None = None, + dataset: list[str] | None = None, + phase: list[str] | None = None, + revision: list[str] | None = None, +) -> list: + """Fetch runs from W&B with server-side filters where W&B supports them.""" + api = wandb.Api(timeout=300) + path = f"{entity}/{project}" if entity else project + + filters: dict = {} + if state: + filters["state"] = state + + print(f"Fetching runs from {path}...", file=sys.stderr) + print(f" Server-side filters: {json.dumps(filters, default=str)}", file=sys.stderr) + runs_iter = api.runs(path, filters=filters or {}, order="-created_at") + + runs = [] + for run in runs_iter: + if not _run_matches_filter(run, "experiment_class", "experiment_class", experiment_class): + continue + if not _run_matches_filter(run, "paradigm", "paradigm", paradigm): + continue + if not _run_matches_filter(run, "dataset", "dataset", dataset): + continue + if not _run_matches_filter(run, "phase", "phase", phase): + continue + if not _run_matches_filter(run, "revision", "revision", revision): + continue + runs.append(run) + + print(f" Fetched {len(runs)} runs.", file=sys.stderr) + return runs + + +# --------------------------------------------------------------------------- +# Config & Metric Extraction +# --------------------------------------------------------------------------- + + +def _get_nested(config: dict, dotted_key: str, default=None): + parts = dotted_key.split(".") + val = config + for part in parts: + if isinstance(val, dict) and part in val: + val = val[part] + else: + return default + return val + + +def _retry(fn, max_retries=3, base_delay=5): + for attempt in range(max_retries + 1): + try: + return fn() + except Exception as exc: + err_str = str(exc).lower() + is_retryable = any( + word in err_str + for word in ["timeout", "timed out", "connection", "429", "500", "502", "503"] + ) + if not is_retryable or attempt == max_retries: + raise + delay = base_delay * (2**attempt) + print( + f" Retry {attempt + 1}/{max_retries} after {delay}s: {exc!r}", + file=sys.stderr, + ) + time.sleep(delay) + + +def _load_run_data(run) -> tuple[dict, dict, str, str, str, list[str], str, str]: + config = dict(run.config) + summary = dict(run.summary_metrics or {}) + return ( + config, + summary, + run.id, + run.url, + run.name or "", + list(run.tags), + run.group, + run.created_at or "", + ) + + +_LR_DECODE = { + "00002": 2e-4, + "00005": 5e-4, + "0002": 2e-3, +} +_MR_DECODE = {"03": 0.3, "075": 0.75} + + +def _tag_value(tags: list[str], prefix: str) -> str | None: + full_prefix = f"{prefix}:" + for tag in tags: + if tag.startswith(full_prefix): + return tag.split(":", 1)[1] + return None + + +def _recover_source_dataset(run_name: str, config: dict) -> str | None: + source_dataset = config.get("source_dataset") + if source_dataset in {"miiv", "eicu", "combined"}: + return source_dataset + for candidate in [config.get("output_dir", ""), run_name]: + match = re.search(r"_from_(miiv|eicu|combined)(?:_|$)", candidate or "") + if match: + return match.group(1) + return None + + +def _recover_pretrain_metadata( + run_name: str, + config: dict, +) -> tuple[float | None, float | None, str | None]: + """Recover upstream pretrain metadata from explicit config or encoded names.""" + up_lr = config.get("upstream_pretrain_lr") + up_mr = config.get("upstream_pretrain_mask_ratio") + subtype = config.get("experiment_subtype") + + if up_lr is None and subtype in {"lr_sensitivity", "lr_ablation"}: + up_lr = _get_nested(config, "optimizer.lr") + if up_mr is None and subtype in { + "mask_ratio_sensitivity", + "view_mask_sensitivity", + "mask_ablation", + }: + up_mr = _get_nested(config, "ssl.mask_ratio") + if up_lr is not None or up_mr is not None: + return up_lr, up_mr, subtype + + output_dir = config.get("output_dir", run_name) + lr_match = re.search(r"_lr([^_]+)", output_dir) + if lr_match: + up_lr = _LR_DECODE.get(lr_match.group(1)) + if up_lr is not None: + subtype = "lr_sensitivity" + + mr_match = re.search(r"_mask_?ratio([^_]+)", output_dir) + if mr_match: + up_mr = _MR_DECODE.get(mr_match.group(1)) + if up_mr is not None: + subtype = subtype or "mask_ratio_sensitivity" + + return up_lr, up_mr, subtype + + +def _infer_model_size(config: dict) -> str: + explicit = config.get("model_size") + if explicit: + return explicit + model_markers = { + str(value).lower() + for value in [ + config.get("paradigm"), + _get_nested(config, "ssl.name"), + _get_nested(config, "encoder.name"), + ] + if value is not None + } + if "smart" in model_markers: + return "default" + d_model = _get_nested(config, "encoder.d_model") + n_layers = _get_nested(config, "encoder.n_layers") + if d_model is None: + return "default" + variant = MODEL_VARIANTS.get((d_model, n_layers)) + if variant is not None: + return variant + if d_model == 64: + return "default" + if n_layers is None: + return f"d{d_model}" + return f"d{d_model}_L{n_layers}" + + +def _metric_value_or_nan(value) -> float: + """Coerce numeric W&B summary values to export floats.""" + if isinstance(value, bool) or not isinstance(value, (int, float)): + return float("nan") + if math.isnan(value) or math.isinf(value): + return float("nan") + return float(value) + + +def extract_run(run, metric_keys: list[str]) -> dict: + """Extract one W&B run to a class-based result row.""" + config, summary, run_id, run_url, run_name, tags, group, created_at = _retry( + lambda: _load_run_data(run) + ) + + experiment_class = config.get("experiment_class") or _tag_value(tags, "experiment_class") + if not experiment_class: + raise RuntimeError( + f"Run {run_id} has no experiment_class in config or tags; final exports fail closed." + ) + + phase = config.get("phase") or _tag_value(tags, "phase") + if phase is None: + paradigm = config.get("paradigm", "") + if paradigm in {"supervised", "gru_d", "xgboost"}: + phase = "baseline" if paradigm in {"gru_d", "xgboost"} else "supervised" + else: + phase = "finetune" + + freeze = _get_nested(config, "training.freeze_encoder") + protocol = canonical_downstream_protocol(config.get("protocol") or _tag_value(tags, "protocol")) + if protocol is None: + protocol = downstream_protocol_from_freeze(freeze) + if protocol is None and (config.get("paradigm") == "xgboost" or phase == "baseline"): + protocol = FULL_FINETUNE_PROTOCOL + + source_dataset = _recover_source_dataset(run_name, config) + up_lr, up_mr, experiment_subtype = _recover_pretrain_metadata(run_name, config) + + metrics = {} + for key in metric_keys: + metrics[key] = _metric_value_or_nan(summary.get(key, None)) + for key, value in summary.items(): + if key.startswith("fairness/") and isinstance(value, (int, float)): + metrics[key] = _metric_value_or_nan(value) + fairness_summary_keys = sorted( + key for key, value in summary.items() if key.startswith("fairness/") and value is not None + ) + + paradigm = config.get("paradigm") or _get_nested(config, "ssl.name") + lr = _get_nested(config, "optimizer.lr", default=None) + if paradigm == "xgboost": + lr = _get_nested(config, "xgboost.learning_rate", default=lr) + row = { + "wandb_run_id": run_id, + "wandb_run_url": run_url, + "wandb_run_name": run_name, + "wandb_group": group, + "created_at": created_at, + "experiment_class": experiment_class, + "experiment_type": experiment_class, + "experiment_subtype": experiment_subtype, + "paradigm": paradigm, + "dataset": config.get("dataset", None), + "task": _get_nested(config, "task.task_name", default=None), + "seed": config.get("seed", None), + "protocol": protocol, + "label_fraction": config.get("label_fraction", 1.0), + "lr": lr, + "mask_ratio": _get_nested(config, "ssl.mask_ratio", default=None), + "model_size": _infer_model_size(config), + "source_dataset": source_dataset, + "revision": config.get("revision", None) or _tag_value(tags, "revision"), + "launch_commit": config.get("launch_commit", None) or _tag_value(tags, "commit"), + "phase": phase, + "upstream_pretrain_lr": up_lr, + "upstream_pretrain_mask_ratio": up_mr, + "_output_dir": config.get("output_dir", None), + "_eval_checkpoint_source": summary.get("_eval_checkpoint_source", None), + EVAL_ARTIFACT_PATH_KEY: summary.get(EVAL_ARTIFACT_PATH_KEY, None), + EVAL_ARTIFACT_SHA256_KEY: summary.get(EVAL_ARTIFACT_SHA256_KEY, None), + "_best_ckpt_path": summary.get("_best_ckpt_path", None), + "_best_ckpt_load_ok": summary.get("_best_ckpt_load_ok", None), + "_best_ckpt_error": summary.get("_best_ckpt_error", None), + FAIRNESS_SUMMARY_KEY_COLUMN: json.dumps(fairness_summary_keys), + } + for key in FAIRNESS_METADATA_COLUMNS: + row[key] = summary.get(key, None) + row.update(metrics) + return row + + +# --------------------------------------------------------------------------- +# DataFrame Construction +# --------------------------------------------------------------------------- + + +def _fillna_for_grouping(df: pd.DataFrame, columns: list[str]) -> pd.DataFrame: + work = df.copy() + for col in columns: + if col in work.columns: + work[col] = work[col].fillna("__none__") + return work + + +def _assert_no_ambiguous_fingerprint_collisions(df: pd.DataFrame) -> None: + collisions = [] + fingerprint = [col for col in [*FINGERPRINT, "seed"] if col in df.columns] + identity_cols = [col for col in FINGERPRINT_AUDIT_COLUMNS if col in df.columns] + if not fingerprint: + return + + work = _fillna_for_grouping(df, list(dict.fromkeys(fingerprint + identity_cols))) + for key, group in work.groupby(fingerprint, dropna=False): + if len(group) <= 1: + continue + if len(group[identity_cols].drop_duplicates()) > 1: + if not isinstance(key, tuple): + key = (key,) + collisions.append((dict(zip(fingerprint, key)), len(group))) + + if collisions: + preview = [] + for fingerprint_values, count in collisions[:5]: + desc = ", ".join(f"{key}={value}" for key, value in fingerprint_values.items()) + preview.append(f"{desc} ({count} rows)") + raise RuntimeError( + "Ambiguous export fingerprint would collapse distinct runs:\n " + "\n ".join(preview) + ) + + +def build_per_seed_df( + runs: list, + allow_extraction_failures: bool = False, + allow_duplicate_fingerprints: bool = False, +) -> pd.DataFrame: + rows = [] + failed = [] + for index, run in enumerate(runs): + if (index + 1) % 100 == 0: + print(f" Processing run {index + 1}/{len(runs)}...", file=sys.stderr) + try: + rows.append(extract_run(run, ALL_METRICS)) + except Exception as exc: + run_id = getattr(run, "id", "unknown") + failed.append(run_id) + print(f" FAILED to extract run {run_id}: {exc!r}", file=sys.stderr) + + if failed: + print(f" WARNING: {len(failed)} runs failed extraction and were skipped.", file=sys.stderr) + if not allow_extraction_failures: + preview = ", ".join(failed[:10]) + raise RuntimeError( + f"{len(failed)} W&B runs failed extraction and were skipped: {preview}. " + "Pass --allow-extraction-failures only for exploratory exports." + ) + + df = pd.DataFrame(rows) + if df.empty: + return df + + if "seed" in df.columns: + df["seed"] = pd.to_numeric(df["seed"], errors="coerce").astype("Int64") + if "label_fraction" in df.columns: + df["label_fraction"] = pd.to_numeric(df["label_fraction"], errors="coerce").fillna(1.0) + if "lr" in df.columns: + df["lr"] = pd.to_numeric(df["lr"], errors="coerce") + if "mask_ratio" in df.columns: + df["mask_ratio"] = pd.to_numeric(df["mask_ratio"], errors="coerce") + + df = df.sort_values("created_at", ascending=False) + before = len(df) + _assert_no_ambiguous_fingerprint_collisions(df) + dedup_cols = [col for col in [*FINGERPRINT, "seed"] if col in df.columns] + duplicate_rows = df[df.duplicated(subset=dedup_cols, keep=False)] + if not duplicate_rows.empty and not allow_duplicate_fingerprints: + preview = [] + for _, row in duplicate_rows.head(10).iterrows(): + desc = ", ".join(f"{col}={row.get(col)}" for col in dedup_cols) + run_id = row.get("wandb_run_id", "unknown") + preview.append(f"{desc}, wandb_run_id={run_id}") + raise RuntimeError( + "Duplicate exact export fingerprints would be silently collapsed:\n " + + "\n ".join(preview) + + "\nPass --allow-duplicate-fingerprints only for exploratory cleanup exports." + ) + df = df.drop_duplicates(subset=dedup_cols, keep="first") + after = len(df) + if before != after: + print( + f" Deduplicated: {before} -> {after} rows " + f"({before - after} older duplicate runs removed).", + file=sys.stderr, + ) + + sort_cols = [ + col + for col in ["experiment_class", "paradigm", "dataset", "task", "seed"] + if col in df.columns + ] + if sort_cols: + df = df.sort_values(sort_cols, na_position="last").reset_index(drop=True) + + counts = df["experiment_class"].value_counts().sort_index() + print("\n Row counts by experiment_class after dedup:", file=sys.stderr) + for experiment_class, count in counts.items(): + print(f" {experiment_class}: {count}", file=sys.stderr) + print(f" TOTAL: {len(df)}", file=sys.stderr) + return df + + +def filter_thesis_tasks( + per_seed_df: pd.DataFrame, + include_extension_tasks: bool = False, +) -> pd.DataFrame: + """Restrict publication exports to the fixed thesis task set by default.""" + if include_extension_tasks or per_seed_df.empty or "task" not in per_seed_df.columns: + return per_seed_df + + task_mask = per_seed_df["task"].isna() | per_seed_df["task"].isin(THESIS_TASKS) + dropped = per_seed_df.loc[~task_mask] + if not dropped.empty: + counts = dropped["task"].value_counts(dropna=False).to_dict() + print( + " Dropped extension-task rows outside the thesis task set: " + f"{len(dropped)} ({counts})", + file=sys.stderr, + ) + return per_seed_df.loc[task_mask].reset_index(drop=True) + + +def _mean_ci95(values: pd.Series) -> tuple[float, float]: + """Return a two-sided 95% CI for a seed mean using Student's t interval.""" + n = len(values) + if n < 2: + return float("nan"), float("nan") + std = values.std(ddof=1) + if pd.isna(std): + return float("nan"), float("nan") + if std == 0: + mean = float(values.mean()) + return mean, mean + + mean = float(values.mean()) + half_width = float(scipy_stats.t.ppf(0.975, df=n - 1) * std / math.sqrt(n)) + return mean - half_width, mean + half_width + + +def _aggregate_group(df: pd.DataFrame, fingerprint_cols: list[str]) -> pd.DataFrame: + metric_cols = [ + col + for col in df.columns + if col in ALL_METRICS or (isinstance(col, str) and col.startswith("fairness/")) + ] + work = _fillna_for_grouping(df, fingerprint_cols) + grouped = work.groupby(fingerprint_cols, dropna=False) + rows = [] + + for fingerprint, group in grouped: + if isinstance(fingerprint, str): + fingerprint = (fingerprint,) + row = dict(zip(fingerprint_cols, fingerprint)) + for col in fingerprint_cols: + if row.get(col) == "__none__": + row[col] = None + + row["n_seeds"] = len(group) + row["wandb_run_ids"] = json.dumps(group["wandb_run_id"].tolist()) + row["seed_list"] = json.dumps(sorted(group["seed"].dropna().astype(int).unique().tolist())) + row["experiment_class_list"] = json.dumps( + sorted(group["experiment_class"].dropna().unique().tolist()) + ) + + for metric in metric_cols: + values = group[metric].dropna() + row[f"{metric}/n"] = int(len(values)) + if len(values) > 0: + row[f"{metric}/mean"] = values.mean() + row[f"{metric}/std"] = values.std(ddof=1) if len(values) > 1 else 0.0 + row[f"{metric}/min"] = values.min() + row[f"{metric}/max"] = values.max() + ci_lower, ci_upper = _mean_ci95(values) + row[f"{metric}/ci95_lower"] = ci_lower + row[f"{metric}/ci95_upper"] = ci_upper + else: + row[f"{metric}/mean"] = float("nan") + row[f"{metric}/std"] = float("nan") + row[f"{metric}/min"] = float("nan") + row[f"{metric}/max"] = float("nan") + row[f"{metric}/ci95_lower"] = float("nan") + row[f"{metric}/ci95_upper"] = float("nan") + rows.append(row) + + return pd.DataFrame(rows) + + +def build_aggregated_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: + if per_seed_df.empty: + return pd.DataFrame() + fingerprint_cols = [col for col in FINGERPRINT if col in per_seed_df.columns] + agg_df = _aggregate_group(per_seed_df, fingerprint_cols) + sort_cols = [ + col for col in ["experiment_class", "paradigm", "dataset", "task"] if col in agg_df.columns + ] + if sort_cols: + agg_df = agg_df.sort_values(sort_cols, na_position="last").reset_index(drop=True) + return agg_df + + +# --------------------------------------------------------------------------- +# Derived Comparison Views +# --------------------------------------------------------------------------- + + +def _copy_for_view(df: pd.DataFrame, view_name: str, target_class: str) -> pd.DataFrame: + view = df.copy() + view["source_experiment_class"] = view["experiment_class"] + view["experiment_class"] = target_class + view["experiment_type"] = target_class + view["comparison_view"] = view_name + return view + + +def build_label_efficiency_curve_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: + if per_seed_df.empty: + return pd.DataFrame() + label_fraction = pd.to_numeric(per_seed_df["label_fraction"], errors="coerce").fillna(1.0) + own = per_seed_df[per_seed_df["experiment_class"] == "label_efficiency"].copy() + endpoint = per_seed_df[ + (per_seed_df["experiment_class"] == "core_ssl_benchmark") + & (label_fraction == 1.0) + & (per_seed_df["phase"].isin(["finetune", "supervised"])) + & (per_seed_df["source_dataset"].isna()) + & (per_seed_df["model_size"] == "default") + ].copy() + parts = [] + if not own.empty: + parts.append(_copy_for_view(own, "label_efficiency_with_core_endpoint", "label_efficiency")) + if not endpoint.empty: + parts.append( + _copy_for_view(endpoint, "label_efficiency_with_core_endpoint", "label_efficiency") + ) + return pd.concat(parts, ignore_index=True) if parts else pd.DataFrame() + + +def build_capacity_study_comparison_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: + if per_seed_df.empty: + return pd.DataFrame() + label_fraction = pd.to_numeric(per_seed_df["label_fraction"], errors="coerce").fillna(1.0) + own = per_seed_df[per_seed_df["experiment_class"] == "capacity_study"].copy() + default_baseline = per_seed_df[ + (per_seed_df["experiment_class"].isin(["core_ssl_benchmark", "label_efficiency"])) + & (per_seed_df["dataset"] == "miiv") + & (per_seed_df["task"] == "mortality_24h") + & (per_seed_df["paradigm"].isin(CAPACITY_PARADIGMS)) + & (per_seed_df["model_size"] == "default") + & (label_fraction.isin(CAPACITY_LABEL_FRACTIONS | {1.0})) + ].copy() + parts = [] + if not own.empty: + parts.append(_copy_for_view(own, "capacity_with_default_baseline", "capacity_study")) + if not default_baseline.empty: + parts.append( + _copy_for_view(default_baseline, "capacity_with_default_baseline", "capacity_study") + ) + return pd.concat(parts, ignore_index=True) if parts else pd.DataFrame() + + +def build_classical_context_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: + if per_seed_df.empty: + return pd.DataFrame() + label_fraction = pd.to_numeric(per_seed_df["label_fraction"], errors="coerce").fillna(1.0) + full_finetune = ( + per_seed_df["protocol"].map(canonical_downstream_protocol) == FULL_FINETUNE_PROTOCOL + ) + full = per_seed_df[ + full_finetune + & (label_fraction == 1.0) + & (per_seed_df["experiment_class"].isin(["core_ssl_benchmark", "classical_baselines"])) + ].copy() + low_label = per_seed_df[ + full_finetune + & (label_fraction < 1.0) + & (per_seed_df["experiment_class"].isin(["label_efficiency", "classical_baselines"])) + ].copy() + + parts = [] + if not full.empty: + parts.append(_copy_for_view(full, "classical_context_full", "classical_context_full")) + if not low_label.empty: + parts.append( + _copy_for_view( + low_label, + "classical_context_label_efficiency", + "classical_context_label_efficiency", + ) + ) + return pd.concat(parts, ignore_index=True) if parts else pd.DataFrame() + + +# --------------------------------------------------------------------------- +# Statistical Tables +# --------------------------------------------------------------------------- + + +def _primary_metric_for_task(task_name: str | None) -> str | None: + if task_name is None: + return None + if task_name in PRIMARY_TEST_METRIC_BY_TASK: + return PRIMARY_TEST_METRIC_BY_TASK[task_name] + if task_name.startswith("los"): + return "test/mae" + return "test/auprc" + + +def _format_task_seed_pairs(keys: set[tuple[str, int]]) -> str: + return json.dumps([{"task": task, "seed": int(seed)} for task, seed in sorted(keys)]) + + +def _metric_pairs_for_paradigm( + scope_group: pd.DataFrame, + paradigm: str, + value_col: str = "primary_metric_value", +) -> pd.DataFrame: + pairs = scope_group[scope_group["paradigm"] == paradigm][["task", "seed", value_col]].rename( + columns={value_col: "value"} + ) + pairs = pairs.dropna(subset=["task", "seed", "value"]).copy() + if not pairs.empty: + pairs["seed"] = pairs["seed"].astype(int) + pairs["value"] = pairs["value"].astype(float) + return pairs + + +def _paired_metric_frame( + scope_group: pd.DataFrame, + paradigm_a: str, + paradigm_b: str, + value_col: str = "primary_metric_value", +) -> tuple[pd.DataFrame, dict[str, int | str]]: + pairs_a = _metric_pairs_for_paradigm(scope_group, paradigm_a, value_col=value_col) + pairs_b = _metric_pairs_for_paradigm(scope_group, paradigm_b, value_col=value_col) + keys_a = set(zip(pairs_a["task"], pairs_a["seed"], strict=True)) + keys_b = set(zip(pairs_b["task"], pairs_b["seed"], strict=True)) + shared_keys = keys_a & keys_b + union_keys = keys_a | keys_b + paired = pairs_a.rename(columns={"value": "value_a"}).merge( + pairs_b.rename(columns={"value": "value_b"}), + on=["task", "seed"], + how="inner", + ) + coverage = { + "n_task_seed_pairs_a": len(keys_a), + "n_task_seed_pairs_b": len(keys_b), + "n_shared_task_seed_pairs": len(shared_keys), + "n_union_task_seed_pairs": len(union_keys), + "missing_task_seed_pairs_a": _format_task_seed_pairs(keys_b - keys_a), + "missing_task_seed_pairs_b": _format_task_seed_pairs(keys_a - keys_b), + } + return paired, coverage + + +def _paired_stat_row( + scope: dict, + paired: pd.DataFrame, + coverage: dict[str, int | str], + paradigm_a: str, + paradigm_b: str, + higher_is_better: bool, +) -> dict: + values_a = paired["value_a"].astype(float).tolist() + values_b = paired["value_b"].astype(float).tolist() + if higher_is_better: + improvement = [a - b for a, b in zip(values_a, values_b, strict=True)] + else: + improvement = [b - a for a, b in zip(values_a, values_b, strict=True)] + wilcoxon = paired_wilcoxon_signed_rank(improvement, [0.0] * len(improvement)) + effect_size = cohens_d(improvement, [0.0] * len(improvement), paired=True) + mean_improvement = sum(improvement) / len(improvement) + return { + **scope, + "paradigm_a": paradigm_a, + "paradigm_b": paradigm_b, + "n_pairs": int(len(paired)), + "n_tasks": int(paired["task"].nunique()), + "task_list": json.dumps(sorted(paired["task"].unique().tolist())), + "seed_list": json.dumps(sorted(paired["seed"].dropna().astype(int).unique().tolist())), + **coverage, + "score_a_mean": float(sum(values_a) / len(values_a)), + "score_b_mean": float(sum(values_b) / len(values_b)), + "mean_improvement": float(mean_improvement), + "median_improvement": float(pd.Series(improvement).median()), + "better_paradigm": ( + paradigm_a if mean_improvement > 0 else paradigm_b if mean_improvement < 0 else None + ), + "wilcoxon_statistic": wilcoxon["statistic"], + "wilcoxon_z": wilcoxon["z_score"], + "p_value": wilcoxon["p_value"], + "n_nonzero_pairs": int(wilcoxon["n_nonzero_pairs"]), + "cohens_d": float(effect_size), + "significant_at_005": bool( + not math.isnan(wilcoxon["p_value"]) and wilcoxon["p_value"] < 0.05 + ), + } + + +def _add_contextual_classical_stat_rows(per_seed_df: pd.DataFrame) -> pd.DataFrame: + required = {"experiment_class", "protocol", "label_fraction", "paradigm"} + if per_seed_df.empty or not required.issubset(per_seed_df.columns): + return per_seed_df + + parts = [per_seed_df] + context = build_classical_context_df(per_seed_df) + if not context.empty: + parts.append(context) + return pd.concat(parts, ignore_index=True) if len(parts) > 1 else per_seed_df + + +def build_statistical_tests_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: + if per_seed_df.empty or "experiment_class" not in per_seed_df.columns: + return pd.DataFrame() + + per_seed_df = _add_contextual_classical_stat_rows(per_seed_df) + rows = [] + + for experiment_class in sorted(per_seed_df["experiment_class"].dropna().unique()): + subset = per_seed_df[per_seed_df["experiment_class"] == experiment_class].copy() + if subset.empty or "paradigm" not in subset.columns or "task" not in subset.columns: + continue + subset["primary_metric_name"] = subset["task"].map(_primary_metric_for_task) + subset["primary_metric_value"] = [ + row.get(metric_name, float("nan")) if metric_name is not None else float("nan") + for row, metric_name in zip( + subset.to_dict("records"), + subset["primary_metric_name"], + strict=True, + ) + ] + + base_scope_cols = [ + col + for col in FINGERPRINT + if col in subset.columns and col not in {"paradigm", "task", "phase"} + ] + if "primary_metric_name" not in base_scope_cols: + base_scope_cols.append("primary_metric_name") + + scope_defs = [ + ("omnibus_primary_metric", base_scope_cols, 2), + ("per_task", [*base_scope_cols, "task"], 1), + ] + for comparison_scope, raw_scope_cols, min_tasks in scope_defs: + scope_cols = list(dict.fromkeys(col for col in raw_scope_cols if col in subset.columns)) + work = _fillna_for_grouping(subset, scope_cols) + for scope_key, scope_group in work.groupby(scope_cols, dropna=False): + if scope_group["task"].nunique() < min_tasks: + continue + if not isinstance(scope_key, tuple): + scope_key = (scope_key,) + scope = dict(zip(scope_cols, scope_key)) + for col, value in list(scope.items()): + if value == "__none__": + scope[col] = None + scope["comparison_scope"] = comparison_scope + + paradigms = sorted( + paradigm + for paradigm in scope_group["paradigm"].dropna().unique().tolist() + if paradigm != "__none__" + ) + if len(paradigms) < 2: + continue + + higher_is_better = scope.get("primary_metric_name") not in LOWER_IS_BETTER_METRICS + family_rows = [] + for paradigm_a, paradigm_b in combinations(paradigms, 2): + if experiment_class in CONTEXTUAL_STAT_CLASSES: + a_classical = paradigm_a in {"gru_d", "xgboost"} + b_classical = paradigm_b in {"gru_d", "xgboost"} + if a_classical == b_classical: + continue + paired, coverage = _paired_metric_frame(scope_group, paradigm_a, paradigm_b) + if paired.empty: + continue + family_rows.append( + _paired_stat_row( + scope=scope, + paired=paired, + coverage=coverage, + paradigm_a=paradigm_a, + paradigm_b=paradigm_b, + higher_is_better=higher_is_better, + ) + ) + + corrected = bonferroni_correction([row["p_value"] for row in family_rows]) + for row, corrected_p in zip(family_rows, corrected, strict=True): + row["p_value_bonferroni"] = corrected_p + row["significant_bonferroni_005"] = bool( + not math.isnan(corrected_p) and corrected_p < 0.05 + ) + rows.append(row) + + if not rows: + return pd.DataFrame() + stats_df = pd.DataFrame(rows) + sort_cols = [ + col + for col in [ + "experiment_class", + "comparison_scope", + "dataset", + "task", + "protocol", + "label_fraction", + "primary_metric_name", + "paradigm_a", + "paradigm_b", + ] + if col in stats_df.columns + ] + if sort_cols: + stats_df = stats_df.sort_values(sort_cols, na_position="last").reset_index(drop=True) + return stats_df + + +def build_ts2vec_vs_core_contrastive_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: + required = {"experiment_class", "paradigm", "task", "seed"} + if per_seed_df.empty or not required.issubset(per_seed_df.columns): + return pd.DataFrame() + + subset = per_seed_df[ + ( + (per_seed_df["experiment_class"] == "ts2vec_extension") + & (per_seed_df["paradigm"] == "ts2vec") + ) + | ( + (per_seed_df["experiment_class"] == "core_ssl_benchmark") + & (per_seed_df["paradigm"] == "contrastive") + ) + ].copy() + if subset.empty: + return pd.DataFrame() + + subset["primary_metric_name"] = subset["task"].map(_primary_metric_for_task) + subset["primary_metric_value"] = [ + row.get(metric_name, float("nan")) if metric_name is not None else float("nan") + for row, metric_name in zip( + subset.to_dict("records"), + subset["primary_metric_name"], + strict=True, + ) + ] + + scope_cols = [ + col + for col in [ + "dataset", + "protocol", + "label_fraction", + "model_size", + "source_dataset", + "primary_metric_name", + ] + if col in subset.columns + ] + rows = [] + work = _fillna_for_grouping(subset, scope_cols) + for scope_key, scope_group in work.groupby(scope_cols, dropna=False): + if not isinstance(scope_key, tuple): + scope_key = (scope_key,) + scope = dict(zip(scope_cols, scope_key)) + for col, value in list(scope.items()): + if value == "__none__": + scope[col] = None + if not {"contrastive", "ts2vec"}.issubset(set(scope_group["paradigm"].dropna())): + continue + paired, coverage = _paired_metric_frame(scope_group, "ts2vec", "contrastive") + if paired.empty: + continue + higher_is_better = scope.get("primary_metric_name") not in LOWER_IS_BETTER_METRICS + rows.append( + _paired_stat_row( + scope={ + **scope, + "comparison_type": "ts2vec_vs_core_contrastive", + "experiment_class_a": "ts2vec_extension", + "experiment_class_b": "core_ssl_benchmark", + }, + paired=paired, + coverage=coverage, + paradigm_a="ts2vec", + paradigm_b="contrastive", + higher_is_better=higher_is_better, + ) + ) + + if not rows: + return pd.DataFrame() + corrected = bonferroni_correction([row["p_value"] for row in rows]) + for row, corrected_p in zip(rows, corrected, strict=True): + row["p_value_bonferroni"] = corrected_p + row["significant_bonferroni_005"] = bool(not math.isnan(corrected_p) and corrected_p < 0.05) + return pd.DataFrame(rows) + + +# --------------------------------------------------------------------------- +# Validation +# --------------------------------------------------------------------------- + + +def _is_missing_export_value(value) -> bool: + """Return whether a scalar export value should be treated as missing.""" + if value is None: + return True + if isinstance(value, str): + return value.strip() == "" + try: + missing = pd.isna(value) + except (TypeError, ValueError): + return False + return bool(missing) if isinstance(missing, bool) else False + + +def _is_true_export_value(value) -> bool: + """Return whether a W&B summary value represents boolean true.""" + if value is True: + return True + if isinstance(value, str): + return value.strip().lower() == "true" + if isinstance(value, (int, float)) and not isinstance(value, bool): + return value == 1 + return False + + +def _row_requires_checkpoint_provenance(row: pd.Series) -> bool: + """Return whether an exported evaluation row should carry checkpoint provenance.""" + if row.get("phase") not in EVAL_PHASES: + return False + if str(row.get("paradigm", "")).lower() == "xgboost": + return False + return True + + +def _checkpoint_provenance_issues(per_seed_df: pd.DataFrame) -> list[tuple[pd.Series, str]]: + """Find evaluation rows whose logged test metrics lack reproducible checkpoint metadata.""" + if per_seed_df.empty: + return [] + + issues = [] + for _, row in per_seed_df.iterrows(): + if not _row_requires_checkpoint_provenance(row): + continue + + source = row.get("_eval_checkpoint_source") + if _is_missing_export_value(source): + issues.append((row, "missing _eval_checkpoint_source")) + continue + + if source == "failed": + error = row.get("_best_ckpt_error") + reason = "recorded checkpoint-selection failure" + if not _is_missing_export_value(error): + reason += f": {error}" + issues.append((row, reason)) + continue + + if source not in {"best", "final"}: + issues.append((row, f"unrecognized _eval_checkpoint_source={source!r}")) + continue + + if _is_missing_export_value(row.get(EVAL_ARTIFACT_SHA256_KEY)): + issues.append((row, f"missing {EVAL_ARTIFACT_SHA256_KEY}")) + continue + + if source == "best": + if _is_missing_export_value(row.get("_best_ckpt_path")): + issues.append((row, "best-checkpoint evaluation missing _best_ckpt_path")) + continue + if not _is_true_export_value(row.get("_best_ckpt_load_ok")): + issues.append((row, "best-checkpoint evaluation did not record load success")) + + return issues + + +def _fairness_attributes_for_row(row: pd.Series) -> list[str]: + """Return protected attributes expected in publication exports for a row.""" + dataset = str(row.get("dataset", "")).lower() + attrs = list(FAIRNESS_ATTRIBUTES) + if dataset == "eicu": + attrs = [attr for attr in attrs if attr != "race"] + return attrs + + +def _fairness_required_metrics_for_row(row: pd.Series) -> list[str]: + """Return required aggregate fairness metrics for the row's task family.""" + task = str(row.get("task", "") or "").lower() + if task.startswith("los"): + return REGRESSION_FAIRNESS_REQUIRED_METRICS + return BINARY_FAIRNESS_REQUIRED_METRICS + + +def _fairness_present_key_set(row: pd.Series) -> set[str] | None: + """Return W&B fairness summary keys known to have existed, if available.""" + value = row.get(FAIRNESS_SUMMARY_KEY_COLUMN) + if _is_missing_export_value(value): + return None + if isinstance(value, str): + try: + parsed = json.loads(value) + except json.JSONDecodeError: + return None + elif isinstance(value, (list, tuple, set)): + parsed = value + else: + return None + return {str(key) for key in parsed} + + +def _row_has_fairness_metric(row: pd.Series, key: str) -> bool: + """Check whether a required fairness metric was written for this row.""" + present_keys = _fairness_present_key_set(row) + if present_keys is not None: + return key in present_keys + if key not in row.index: + return False + return not _is_missing_export_value(row.get(key)) + + +def _fairness_completeness_issues(per_seed_df: pd.DataFrame) -> list[tuple[pd.Series, list[str]]]: + """Find evaluation rows missing dataset/task-appropriate fairness summaries.""" + if per_seed_df.empty: + return [] + + issues = [] + for _, row in per_seed_df.iterrows(): + if row.get("phase") not in EVAL_PHASES: + continue + if _is_missing_export_value(row.get("task")): + continue + + missing_keys = [] + for attr in _fairness_attributes_for_row(row): + for metric_name in _fairness_required_metrics_for_row(row): + key = f"fairness/{attr}/{metric_name}" + if not _row_has_fairness_metric(row, key): + missing_keys.append(key) + + if missing_keys: + issues.append((row, missing_keys)) + + return issues + + +def _expected_fairness_checkpoint_source_for_row(row: pd.Series) -> str | None: + """Return checkpoint provenance expected by fairness metadata for a row.""" + if str(row.get("paradigm", "")).lower() == "xgboost": + return "xgboost_model" + + source = row.get("_eval_checkpoint_source") + return str(source) if source in {"best", "final"} else None + + +def _expected_fairness_artifact_source_for_row(row: pd.Series) -> str | None: + """Return evaluated artifact source expected by fairness metadata for a row.""" + if str(row.get("paradigm", "")).lower() == "xgboost": + return "xgboost_model" + + source = row.get("_eval_checkpoint_source") + if source == "best": + return "recorded_best" + if source == "final": + return "recorded_final" + return None + + +def _expected_fairness_artifact_id_for_row(row: pd.Series) -> str | None: + """Return expected artifact identity when export metadata makes it knowable.""" + output_dir = row.get("_output_dir") + paradigm = str(row.get("paradigm", "")).lower() + + if paradigm == "xgboost": + if _is_missing_export_value(output_dir): + return None + return canonical_artifact_id(Path(str(output_dir)) / "xgboost_model.json") + + source = row.get("_eval_checkpoint_source") + if source == "best": + best_path = row.get("_best_ckpt_path") + if _is_missing_export_value(best_path): + return None + return canonical_artifact_id(best_path) + if source == "final" and not _is_missing_export_value(output_dir): + return canonical_artifact_id(Path(str(output_dir)) / "checkpoints" / "last.ckpt") + return None + + +def _coerce_int_metadata(value) -> int | None: + if _is_missing_export_value(value): + return None + try: + return int(float(value)) + except (TypeError, ValueError): + return None + + +def _fairness_metadata_staleness_issues( + per_seed_df: pd.DataFrame, +) -> list[tuple[pd.Series, list[str]]]: + """Find evaluation rows whose fairness summaries lack current provenance metadata.""" + if per_seed_df.empty: + return [] + + issues = [] + for _, row in per_seed_df.iterrows(): + if row.get("phase") not in EVAL_PHASES: + continue + if _is_missing_export_value(row.get("task")): + continue + + row_issues: list[str] = [] + if row.get(FAIRNESS_SCHEMA_VERSION_KEY) != FAIRNESS_SUMMARY_SCHEMA_VERSION: + row_issues.append("missing or stale fairness schema version") + if row.get(FAIRNESS_SCRIPT_VERSION_KEY) != FAIRNESS_SCRIPT_VERSION: + row_issues.append("missing or stale fairness script version") + + actual_attrs = decode_protected_attributes(row.get(FAIRNESS_PROTECTED_ATTRIBUTES_KEY)) + expected_attrs = normalize_protected_attributes(FAIRNESS_ATTRIBUTES) + if actual_attrs != expected_attrs: + row_issues.append( + f"protected attributes mismatch: expected={expected_attrs}, actual={actual_attrs}" + ) + + actual_min_subgroup_size = _coerce_int_metadata(row.get(FAIRNESS_MIN_SUBGROUP_SIZE_KEY)) + if actual_min_subgroup_size != FAIRNESS_DEFAULT_MIN_SUBGROUP_SIZE: + row_issues.append( + "min subgroup size mismatch: " + f"expected={FAIRNESS_DEFAULT_MIN_SUBGROUP_SIZE}, " + f"actual={actual_min_subgroup_size}" + ) + + expected_artifact_source = _expected_fairness_artifact_source_for_row(row) + actual_artifact_source = row.get(FAIRNESS_ARTIFACT_SOURCE_KEY) + if ( + expected_artifact_source is not None + and actual_artifact_source != expected_artifact_source + ): + row_issues.append( + "artifact source mismatch: " + f"expected={expected_artifact_source}, actual={actual_artifact_source}" + ) + + expected_checkpoint_source = _expected_fairness_checkpoint_source_for_row(row) + actual_checkpoint_source = row.get(FAIRNESS_CHECKPOINT_SOURCE_KEY) + if ( + expected_checkpoint_source is not None + and actual_checkpoint_source != expected_checkpoint_source + ): + row_issues.append( + "checkpoint source mismatch: " + f"expected={expected_checkpoint_source}, actual={actual_checkpoint_source}" + ) + + actual_artifact_id = canonical_artifact_id(row.get(FAIRNESS_ARTIFACT_PATH_KEY)) + expected_artifact_id = _expected_fairness_artifact_id_for_row(row) + if not actual_artifact_id: + row_issues.append("missing fairness artifact path") + elif expected_artifact_id is not None and actual_artifact_id != expected_artifact_id: + row_issues.append( + "artifact path mismatch: " + f"expected={expected_artifact_id}, actual={actual_artifact_id}" + ) + + actual_artifact_sha256 = row.get(FAIRNESS_ARTIFACT_SHA256_KEY) + expected_artifact_sha256 = row.get(EVAL_ARTIFACT_SHA256_KEY) + if _is_missing_export_value(actual_artifact_sha256): + row_issues.append("missing fairness artifact sha256") + elif ( + not _is_missing_export_value(expected_artifact_sha256) + and actual_artifact_sha256 != expected_artifact_sha256 + ): + row_issues.append( + "artifact sha256 mismatch: " + f"expected={expected_artifact_sha256}, actual={actual_artifact_sha256}" + ) + elif expected_artifact_id is not None and _is_missing_export_value( + expected_artifact_sha256 + ): + row_issues.append("missing evaluation artifact sha256") + + if row_issues: + issues.append((row, row_issues)) + + return issues + + +def _mixed_revision_group_warnings(per_seed_df: pd.DataFrame) -> list[str]: + """Warn when one aggregated scientific config draws seeds from multiple revisions.""" + if per_seed_df.empty or "revision" not in per_seed_df.columns: + return [] + + fingerprint_cols = [col for col in FINGERPRINT if col in per_seed_df.columns] + if not fingerprint_cols: + return [] + + warnings = [] + work = _fillna_for_grouping(per_seed_df, fingerprint_cols + ["revision"]) + for key, group in work.groupby(fingerprint_cols, dropna=False): + revisions = sorted( + revision + for revision in group["revision"].dropna().astype(str).unique().tolist() + if revision != "__none__" + ) + if len(revisions) <= 1: + continue + if not isinstance(key, tuple): + key = (key,) + desc = _format_matrix_key(key, fingerprint_cols) + warnings.append( + "WARNING: export group mixes multiple revisions; publication exports should " + f"use one revision. {desc}; revisions={revisions}" + ) + + return warnings + + +def _export_text_or_none(value) -> str | None: + if _is_missing_export_value(value): + return None + text = str(value).strip() + return text or None + + +def _launch_commit_homogeneity_warnings(per_seed_df: pd.DataFrame) -> list[str]: + """Warn when one revision is exported from missing or mixed launch commits.""" + if per_seed_df.empty or "revision" not in per_seed_df.columns: + return [] + + if "launch_commit" not in per_seed_df.columns: + return [ + "WARNING: export rows lack launch_commit metadata; final revision provenance " + "cannot be audited." + ] + + warnings = [] + for revision_value, group in per_seed_df.groupby("revision", dropna=False): + revision = _export_text_or_none(revision_value) + if revision is None: + warnings.append( + f"WARNING: {len(group)} exported rows are missing revision metadata; " + "launch-commit homogeneity cannot be scoped." + ) + continue + + commits = sorted( + { + commit + for commit in (_export_text_or_none(value) for value in group["launch_commit"]) + if commit is not None + } + ) + missing = int(group["launch_commit"].map(_export_text_or_none).isna().sum()) + if missing: + warnings.append( + f"WARNING: revision={revision} has {missing} runs missing launch_commit; " + "final provenance cannot be audited." + ) + if len(commits) > 1: + warnings.append( + f"WARNING: revision={revision} mixes multiple launch commits; " f"commits={commits}" + ) + + return warnings + + +def _expected_row_from_run(run) -> dict: + return { + "experiment_class": run.experiment_class, + "experiment_type": run.experiment_class, + "experiment_subtype": run.experiment_subtype, + "paradigm": run.paradigm, + "dataset": run.dataset, + "task": run.task, + "seed": run.seed, + "protocol": run.protocol, + "label_fraction": run.label_fraction, + "model_size": run.model_size or "default", + "source_dataset": run.source_dataset, + "phase": run.phase, + "upstream_pretrain_lr": run.upstream_pretrain_lr, + "upstream_pretrain_mask_ratio": run.upstream_pretrain_mask_ratio, + } + + +def build_expected_matrix_df( + experiment_class: list[str] | None = None, + paradigm: list[str] | None = None, + dataset: list[str] | None = None, + phase: list[str] | None = None, +) -> pd.DataFrame: + """Build the expected per-seed evaluation matrix for the selected export scope.""" + from scripts.internal.run_experiments import generate_all_runs + + experiment_class_set = set(experiment_class) if experiment_class else None + paradigm_set = set(paradigm) if paradigm else None + dataset_set = set(dataset) if dataset else None + phase_set = set(phase) if phase else set(EVAL_PHASES) + + rows = [] + for run in generate_all_runs(): + if run.phase not in EVAL_PHASES: + continue + row = _expected_row_from_run(run) + if experiment_class_set and row["experiment_class"] not in experiment_class_set: + continue + if paradigm_set and row["paradigm"] not in paradigm_set: + continue + if dataset_set and row["dataset"] not in dataset_set: + continue + if phase_set and row["phase"] not in phase_set: + continue + rows.append(row) + + return pd.DataFrame(rows) + + +def _matrix_key_value(column: str, value): + if _is_missing_export_value(value): + return "__none__" + if column == "seed": + return int(value) + if column in {"label_fraction", "upstream_pretrain_lr", "upstream_pretrain_mask_ratio"}: + return round(float(value), 12) + return str(value) + + +def _matrix_keys(df: pd.DataFrame, columns: list[str]) -> set[tuple]: + if df.empty: + return set() + keys = set() + for _, row in df.iterrows(): + keys.add(tuple(_matrix_key_value(col, row.get(col)) for col in columns)) + return keys + + +def _format_matrix_key(key: tuple, columns: list[str]) -> str: + return ", ".join(f"{col}={value}" for col, value in zip(columns, key, strict=True)) + + +def _matrix_coverage_warnings( + per_seed_df: pd.DataFrame, + expected_matrix_df: pd.DataFrame | None, +) -> list[str]: + if expected_matrix_df is None or expected_matrix_df.empty: + return [] + + key_columns = [*FINGERPRINT, "seed"] + missing_columns = [col for col in key_columns if col not in per_seed_df.columns] + if missing_columns: + return [ + "WARNING: export is missing matrix key columns needed for coverage validation: " + + ", ".join(missing_columns) + ] + + expected_keys = _matrix_keys(expected_matrix_df, key_columns) + observed_keys = _matrix_keys(per_seed_df, key_columns) + missing = sorted(expected_keys - observed_keys) + unexpected = sorted(observed_keys - expected_keys) + + warnings = [] + if missing: + warnings.append( + f"WARNING: {len(missing)}/{len(expected_keys)} expected matrix evaluation rows " + "are absent from export:" + ) + for key in missing[:20]: + warnings.append(f" missing {_format_matrix_key(key, key_columns)}") + if len(missing) > 20: + warnings.append(f" ... {len(missing) - 20} more missing rows") + + if unexpected: + warnings.append( + f"WARNING: {len(unexpected)} exported evaluation rows are outside the expected matrix:" + ) + for key in unexpected[:20]: + warnings.append(f" unexpected {_format_matrix_key(key, key_columns)}") + if len(unexpected) > 20: + warnings.append(f" ... {len(unexpected) - 20} more unexpected rows") + + return warnings + + +def validate( + per_seed_df: pd.DataFrame, + aggregated_df: pd.DataFrame, + statistical_df: pd.DataFrame | None = None, + expected_seeds: set[int] | None = None, + expected_matrix_df: pd.DataFrame | None = None, +) -> list[str]: + warnings = [] + expected_seeds = set(EXPECTED_FIXED_SEEDS if expected_seeds is None else expected_seeds) + + warnings.extend(_matrix_coverage_warnings(per_seed_df, expected_matrix_df)) + + if "experiment_class" in aggregated_df.columns: + fixed_seed = aggregated_df[ + aggregated_df["experiment_class"].isin(FIXED_SEED_EXPERIMENT_CLASSES) + ] + bad_seed_rows = [] + for _, row in fixed_seed.iterrows(): + seeds = set(json.loads(row["seed_list"])) if pd.notna(row.get("seed_list")) else set() + if seeds != expected_seeds: + bad_seed_rows.append((row, seeds)) + if bad_seed_rows: + warnings.append( + f"WARNING: {len(bad_seed_rows)}/{len(fixed_seed)} fixed-seed configs do not " + f"have the expected seed set {sorted(expected_seeds)}:" + ) + for row, seeds in bad_seed_rows: + desc = ", ".join( + f"{col}={row[col]}" + for col in ["paradigm", "dataset", "task", "protocol"] + if col in row and row[col] is not None + ) + warnings.append( + f" experiment_class={row['experiment_class']}, {desc} - " + f"seeds={sorted(seeds)}, missing={sorted(expected_seeds - seeds)}, " + f"unexpected={sorted(seeds - expected_seeds)}" + ) + + if statistical_df is not None and not statistical_df.empty: + coverage_cols = {"n_shared_task_seed_pairs", "n_union_task_seed_pairs"} + if coverage_cols.issubset(statistical_df.columns): + incomplete = statistical_df[ + statistical_df["n_shared_task_seed_pairs"] + < statistical_df["n_union_task_seed_pairs"] + ] + if len(incomplete) > 0: + warnings.append( + f"WARNING: {len(incomplete)}/{len(statistical_df)} statistical comparisons " + "have incomplete paired task/seed coverage:" + ) + + test_cols = [col for col in PERFORMANCE_TEST_METRICS if col in per_seed_df.columns] + if test_cols: + all_nan = per_seed_df[test_cols].isna().all(axis=1) + if all_nan.any(): + warnings.append( + f"WARNING: {all_nan.sum()} runs have no test metrics at all. " + "These may be pretraining runs or crashed evaluations." + ) + + if not per_seed_df.empty and "task" in per_seed_df.columns: + missing_primary = [] + for _, row in per_seed_df.iterrows(): + if row.get("phase") not in EVAL_PHASES: + continue + primary_metric = _primary_metric_for_task(row.get("task")) + if primary_metric is None or primary_metric not in per_seed_df.columns: + continue + if pd.isna(row.get(primary_metric)): + missing_primary.append((row, primary_metric)) + if missing_primary: + warnings.append( + f"WARNING: {len(missing_primary)} evaluation runs are missing their " + "primary test metric." + ) + + checkpoint_issues = _checkpoint_provenance_issues(per_seed_df) + if checkpoint_issues: + warnings.append( + f"WARNING: {len(checkpoint_issues)} evaluation runs have missing or failed " + "checkpoint provenance. Publication exports fail closed by default; use " + "--allow-incomplete only for exploratory exports." + ) + for row, reason in checkpoint_issues[:10]: + warnings.append( + " " + + ", ".join( + f"{col}={row.get(col)}" + for col in [ + "wandb_run_id", + "paradigm", + "dataset", + "task", + "seed", + "phase", + ] + if col in row + ) + + f" - {reason}" + ) + + fairness_issues = _fairness_completeness_issues(per_seed_df) + if fairness_issues: + warnings.append( + f"WARNING: {len(fairness_issues)} evaluation runs are missing required " + "fairness summary metrics. Run scripts/eval/evaluate_fairness.py before " + "publication export; use --allow-incomplete only for exploratory exports." + ) + for row, missing_keys in fairness_issues[:10]: + preview = ", ".join(missing_keys[:8]) + if len(missing_keys) > 8: + preview += f", ... {len(missing_keys) - 8} more" + warnings.append( + " " + + ", ".join( + f"{col}={row.get(col)}" + for col in [ + "wandb_run_id", + "paradigm", + "dataset", + "task", + "seed", + "phase", + ] + if col in row + ) + + f" - missing={preview}" + ) + + fairness_metadata_issues = _fairness_metadata_staleness_issues(per_seed_df) + if fairness_metadata_issues: + warnings.append( + f"WARNING: {len(fairness_metadata_issues)} evaluation runs have missing or " + "stale fairness summary metadata. Re-run scripts/eval/evaluate_fairness.py " + "with the publication defaults before export; use --allow-incomplete only " + "for exploratory exports." + ) + for row, row_issues in fairness_metadata_issues[:10]: + preview = "; ".join(row_issues[:6]) + if len(row_issues) > 6: + preview += f"; ... {len(row_issues) - 6} more" + warnings.append( + " " + + ", ".join( + f"{col}={row.get(col)}" + for col in [ + "wandb_run_id", + "paradigm", + "dataset", + "task", + "seed", + "phase", + ] + if col in row + ) + + f" - {preview}" + ) + + warnings.extend(_mixed_revision_group_warnings(per_seed_df)) + warnings.extend(_launch_commit_homogeneity_warnings(per_seed_df)) + + print("\nValidation summary:", file=sys.stderr) + print(f" Total runs: {len(per_seed_df)}", file=sys.stderr) + print(f" Unique configs: {len(aggregated_df)}", file=sys.stderr) + if "experiment_class" in aggregated_df.columns: + for experiment_class, group in aggregated_df.groupby("experiment_class"): + seed_dist = dict(group["n_seeds"].value_counts().sort_index()) + print( + f" {experiment_class}: {len(group)} configs, seeds: {seed_dist}", file=sys.stderr + ) + print(f" Warnings: {len(warnings)}", file=sys.stderr) + return warnings + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Export SLICES experiment results from W&B to parquet files.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, + ) + parser.add_argument( + "--project", + default=os.environ.get("WANDB_PROJECT", "slices-thesis"), + help="W&B project name (default: WANDB_PROJECT env var or 'slices-thesis')", + ) + parser.add_argument( + "--entity", + default=os.environ.get("WANDB_ENTITY"), + help="W&B entity name (default: WANDB_ENTITY env var)", + ) + parser.add_argument( + "--experiment-class", + nargs="+", + choices=EXPERIMENT_CLASSES, + help="Filter to specific experiment class(es)", + ) + parser.add_argument("--paradigm", nargs="+", help="Filter to specific paradigm(s)") + parser.add_argument("--dataset", nargs="+", help="Filter to specific dataset(s)") + parser.add_argument( + "--phase", + nargs="+", + default=EVAL_PHASES, + help=f"Filter to specific phase(s) (default: {EVAL_PHASES})", + ) + parser.add_argument("--revision", nargs="+", help="Filter to specific revision tag(s)") + parser.add_argument( + "--allow-multiple-revisions", + action="store_true", + help=( + "Permit exporting more than one revision in a single exploratory run. " + "Publication exports should use exactly one revision." + ), + ) + parser.add_argument("--state", default="finished", help="Run state filter (default: finished)") + parser.add_argument("--output-dir", default="results", help="Output directory") + parser.add_argument( + "--expected-seeds", + nargs="+", + type=int, + default=sorted(EXPECTED_FIXED_SEEDS), + help=( + "Expected exact seed set for fixed-seed configs " + f"(default: {sorted(EXPECTED_FIXED_SEEDS)})" + ), + ) + parser.add_argument( + "--allow-incomplete", + action="store_true", + help="Exit successfully even if no runs match or validation warnings are emitted.", + ) + parser.add_argument( + "--allow-extraction-failures", + action="store_true", + help=( + "Skip W&B runs that fail row extraction. Use only for exploratory exports; " + "publication exports fail closed by default." + ), + ) + parser.add_argument( + "--allow-duplicate-fingerprints", + action="store_true", + help=( + "Keep the newest duplicate exact scientific fingerprint. Use only for " + "exploratory cleanup exports; publication exports fail closed by default." + ), + ) + parser.add_argument( + "--include-extension-tasks", + action="store_true", + help=( + "Include task rows outside the fixed thesis task set. By default, " + f"primary thesis exports are restricted to {sorted(THESIS_TASKS)}." + ), + ) + args = parser.parse_args() + + if not args.revision: + env_revision = os.environ.get("REVISION") or os.environ.get("WANDB_REVISION") + if env_revision: + args.revision = [env_revision] + else: + parser.error( + "--revision is required to avoid mixing reruns. " + "Pass --revision or set REVISION/WANDB_REVISION." + ) + if len(args.revision) > 1 and not args.allow_multiple_revisions: + parser.error( + "Multiple --revision values are disabled for publication exports. " + "Run one revision at a time, or pass --allow-multiple-revisions for exploratory export." + ) + + return args + + +def main() -> None: + args = parse_args() + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + runs = fetch_all_runs( + project=args.project, + entity=args.entity, + state=args.state, + experiment_class=args.experiment_class, + paradigm=args.paradigm, + dataset=args.dataset, + phase=args.phase, + revision=args.revision, + ) + if not runs: + print("No runs found matching filters. Exiting.", file=sys.stderr) + sys.exit(0 if args.allow_incomplete else 1) + + print(f"\nBuilding per-seed DataFrame from {len(runs)} runs...", file=sys.stderr) + per_seed_df = build_per_seed_df( + runs, + allow_extraction_failures=args.allow_extraction_failures, + allow_duplicate_fingerprints=args.allow_duplicate_fingerprints, + ) + per_seed_df = filter_thesis_tasks( + per_seed_df, + include_extension_tasks=args.include_extension_tasks, + ) + print(f" Shape: {per_seed_df.shape}", file=sys.stderr) + + if per_seed_df.empty: + print("No rows remain after task filtering. Exiting.", file=sys.stderr) + sys.exit(0 if args.allow_incomplete else 1) + + print("\nBuilding aggregated DataFrame...", file=sys.stderr) + aggregated_df = build_aggregated_df(per_seed_df) + print(f" Shape: {aggregated_df.shape}", file=sys.stderr) + + print("\nBuilding derived comparison views...", file=sys.stderr) + label_efficiency_df = build_label_efficiency_curve_df(per_seed_df) + capacity_df = build_capacity_study_comparison_df(per_seed_df) + classical_context_df = build_classical_context_df(per_seed_df) + + print("\nBuilding statistical significance table...", file=sys.stderr) + statistical_df = build_statistical_tests_df(per_seed_df) + print(f" Shape: {statistical_df.shape}", file=sys.stderr) + + print("\nBuilding TS2Vec vs core contrastive table...", file=sys.stderr) + ts2vec_contrastive_df = build_ts2vec_vs_core_contrastive_df(per_seed_df) + print(f" Shape: {ts2vec_contrastive_df.shape}", file=sys.stderr) + + expected_matrix_df = build_expected_matrix_df( + experiment_class=args.experiment_class, + paradigm=args.paradigm, + dataset=args.dataset, + phase=args.phase, + ) + + coverage_parts = [df for df in [statistical_df, ts2vec_contrastive_df] if not df.empty] + statistical_coverage_df = ( + pd.concat(coverage_parts, ignore_index=True) if coverage_parts else pd.DataFrame() + ) + warnings = validate( + per_seed_df, + aggregated_df, + statistical_df=statistical_coverage_df, + expected_seeds=set(args.expected_seeds), + expected_matrix_df=expected_matrix_df, + ) + for warning in warnings: + print(warning, file=sys.stderr) + + if warnings and not args.allow_incomplete: + print("Validation failed; no parquet outputs were written.", file=sys.stderr) + sys.exit(1) + + outputs = { + "per_seed_results.parquet": per_seed_df, + "aggregated_results.parquet": aggregated_df, + "statistical_tests.parquet": statistical_df, + "label_efficiency_curves.parquet": label_efficiency_df, + "capacity_study_comparison.parquet": capacity_df, + "classical_context.parquet": classical_context_df, + "ts2vec_vs_core_contrastive.parquet": ts2vec_contrastive_df, + } + for name, df in outputs.items(): + path = output_dir / name + df.to_parquet(path, index=False) + print(f" {path} ({len(df)} rows)", file=sys.stderr) + + print("\n--- Quick Summary ---") + print( + f"Runs: {len(per_seed_df)}, Configs: {len(aggregated_df)}, " + f"StatTests: {len(statistical_df)}, " + f"TS2VecContrastiveTests: {len(ts2vec_contrastive_df)}" + ) + if "experiment_class" in aggregated_df.columns: + for experiment_class, group in aggregated_df.groupby("experiment_class"): + seed_dist = dict(group["n_seeds"].value_counts().sort_index()) + print(f" {experiment_class}: {len(group)} configs, seeds: {seed_dist}") + + +if __name__ == "__main__": + main() diff --git a/scripts/internal/auto_shutdown.sh b/scripts/internal/auto_shutdown.sh new file mode 100755 index 0000000..3e86d11 --- /dev/null +++ b/scripts/internal/auto_shutdown.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +# ============================================================================= +# Auto-shutdown: powers off the VM after the configured idle threshold. +# Install as a cron job on the GCP VM (see bottom of script). +# ============================================================================= +set -euo pipefail + +IDLE_THRESHOLD_MIN=59 +STAMP_FILE="/tmp/slices_last_training_activity" + +# Check if any final-run process or coding agent session is running +if pgrep -f "scripts/(training/(pretrain|finetune|supervised|xgboost_baseline)|internal/run_experiments|preprocessing/prepare_dataset|eval/evaluate_fairness|export_results)\.py" > /dev/null 2>&1 \ + || pgrep -f "claude" > /dev/null 2>&1 \ + || pgrep -f "codex" > /dev/null 2>&1; then + # Activity detected; update timestamp and exit + date +%s > "$STAMP_FILE" + exit 0 +fi + +# No training running. Check when the last one was seen. +if [ ! -f "$STAMP_FILE" ]; then + # No stamp file = script just installed, start the clock now + date +%s > "$STAMP_FILE" + exit 0 +fi + +LAST_ACTIVE=$(cat "$STAMP_FILE") +NOW=$(date +%s) +IDLE_SEC=$(( NOW - LAST_ACTIVE )) +IDLE_MIN=$(( IDLE_SEC / 60 )) + +if [ "$IDLE_MIN" -ge "$IDLE_THRESHOLD_MIN" ]; then + echo "$(date): No training activity for ${IDLE_MIN}m. Shutting down." + logger "slices-auto-shutdown: idle ${IDLE_MIN}m, shutting down" + sudo shutdown -h now +else + echo "$(date): Idle ${IDLE_MIN}m / ${IDLE_THRESHOLD_MIN}m threshold. Staying up." +fi diff --git a/scripts/internal/launch_thesis_tmux.sh b/scripts/internal/launch_thesis_tmux.sh new file mode 100755 index 0000000..72a88b2 --- /dev/null +++ b/scripts/internal/launch_thesis_tmux.sh @@ -0,0 +1,339 @@ +#!/usr/bin/env bash +set -euo pipefail + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +SESSION_NAME="${SESSION_NAME:-slices-thesis}" +WANDB_PROJECT="${WANDB_PROJECT:-slices-thesis}" +WANDB_ENTITY="${WANDB_ENTITY:-}" +REVISION="${REVISION:-thesis-v1}" +REASON="${REASON:-thesis-final class-based rerun corpus}" +LAUNCH_COMMIT="${LAUNCH_COMMIT:-}" +SKIP_LAUNCH_GIT_CHECK="${SKIP_LAUNCH_GIT_CHECK:-0}" +ALLOW_DIRTY="${ALLOW_DIRTY:-0}" +EXPECTED_FEATURES="${EXPECTED_FEATURES:-84}" +VALIDATE_PROCESSED_ARTIFACTS="${VALIDATE_PROCESSED_ARTIFACTS:-1}" +PURGE_RUNTIME_CACHES="${PURGE_RUNTIME_CACHES:-1}" +PARALLEL_MAIN="${PARALLEL_MAIN:-4}" +PARALLEL_APPENDIX="${PARALLEL_APPENDIX:-4}" +BATCH_SIZE_FAIRNESS="${BATCH_SIZE_FAIRNESS:-64}" +DEVICE_FAIRNESS="${DEVICE_FAIRNESS:-auto}" +INCLUDE_SMART_REFERENCE="${INCLUDE_SMART_REFERENCE:-1}" +INCLUDE_TS2VEC_EXTENSION="${INCLUDE_TS2VEC_EXTENSION:-1}" +INCLUDE_CAPACITY_STUDY="${INCLUDE_CAPACITY_STUDY:-1}" +VALIDATE_FINAL_CORPUS="${VALIDATE_FINAL_CORPUS:-1}" +EXPECTED_TOTAL_RUNS="${EXPECTED_TOTAL_RUNS:-2590}" +RUN_EXPORT="${RUN_EXPORT:-1}" +STATUS_INTERVAL="${STATUS_INTERVAL:-60}" +RESULTS_DIR="${RESULTS_DIR:-results/${WANDB_PROJECT}_${REVISION}}" +LOG_DIR="${LOG_DIR:-logs/runner}" + +if ! command -v tmux >/dev/null 2>&1; then + echo "tmux is required but not installed." >&2 + exit 1 +fi + +if tmux has-session -t "$SESSION_NAME" 2>/dev/null; then + echo "tmux session '$SESSION_NAME' already exists." >&2 + exit 1 +fi + +if [[ -z "$WANDB_ENTITY" ]]; then + echo "WANDB_ENTITY must be set for thesis runs." >&2 + exit 1 +fi + +WANDB_MODE_VALUE="${WANDB_MODE:-online}" +WANDB_MODE_LC="$(printf '%s' "$WANDB_MODE_VALUE" | tr '[:upper:]' '[:lower:]')" +if [[ "$WANDB_MODE_LC" != "online" ]]; then + echo "WANDB_MODE must be unset or 'online' for thesis runs; got '$WANDB_MODE_VALUE'." >&2 + exit 1 +fi +export WANDB_MODE=online + +if [[ "$SKIP_LAUNCH_GIT_CHECK" == "1" ]]; then + LAUNCH_COMMIT="${LAUNCH_COMMIT:-unchecked}" +else + if [[ -z "$LAUNCH_COMMIT" ]]; then + LAUNCH_COMMIT="$(git -C "$REPO_ROOT" rev-parse --verify HEAD)" + fi + + if ! RESOLVED_LAUNCH_COMMIT="$(git -C "$REPO_ROOT" rev-parse --verify "$LAUNCH_COMMIT^{commit}" 2>/dev/null)"; then + echo "LAUNCH_COMMIT is not available in this checkout: $LAUNCH_COMMIT" >&2 + exit 1 + fi + LAUNCH_COMMIT="$RESOLVED_LAUNCH_COMMIT" + + CURRENT_COMMIT="$(git -C "$REPO_ROOT" rev-parse --verify HEAD)" + if [[ "$CURRENT_COMMIT" != "$LAUNCH_COMMIT" ]]; then + echo "Refusing to launch from a different commit." >&2 + echo " current: $CURRENT_COMMIT" >&2 + echo " expected: $LAUNCH_COMMIT" >&2 + exit 1 + fi + + if [[ "$ALLOW_DIRTY" != "1" ]]; then + if ! git -C "$REPO_ROOT" diff --quiet || ! git -C "$REPO_ROOT" diff --cached --quiet; then + echo "Refusing to launch with tracked uncommitted changes." >&2 + echo "Commit the reviewed state, or set ALLOW_DIRTY=1 for an explicit local dry run." >&2 + exit 1 + fi + fi +fi + +validate_processed_artifacts() { + uv run python - "$EXPECTED_FEATURES" <<'PY' +from pathlib import Path +import sys + +import polars as pl +import yaml + +expected_features = int(sys.argv[1]) +base = Path("data/processed") +datasets = ("miiv", "eicu", "combined") +required_files = ( + "metadata.yaml", + "static.parquet", + "timeseries.parquet", + "labels.parquet", + "splits.yaml", + "normalization_stats.yaml", +) +errors = [] + +for dataset in datasets: + path = base / dataset + missing = [name for name in required_files if not (path / name).exists()] + if missing: + errors.append(f"{dataset}: missing {', '.join(missing)}") + continue + + with open(path / "metadata.yaml") as f: + metadata = yaml.safe_load(f) or {} + + feature_names = metadata.get("feature_names") or [] + n_features = metadata.get("n_features", len(feature_names)) + if n_features != expected_features: + errors.append(f"{dataset}: expected {expected_features} features, found {n_features}") + + n_stays = metadata.get("n_stays") + for filename in ("static.parquet", "timeseries.parquet", "labels.parquet"): + height = pl.scan_parquet(path / filename).select(pl.len()).collect().item() + if n_stays is not None and height != n_stays: + errors.append(f"{dataset}: {filename} has {height} rows, metadata has {n_stays}") + + if dataset in {"miiv", "eicu"} and not metadata.get("zero_observation_stays_excluded"): + errors.append(f"{dataset}: metadata does not confirm zero-observation exclusion") + + print(f"{dataset}: {n_stays:,} stays, {n_features} features") + +if errors: + print("\nProcessed artifact readiness check failed:", file=sys.stderr) + for error in errors: + print(f" - {error}", file=sys.stderr) + raise SystemExit(1) +PY +} + +purge_runtime_caches() { + if [[ "$PURGE_RUNTIME_CACHES" != "1" ]]; then + return 0 + fi + if [[ ! -d "$REPO_ROOT/data/processed" ]]; then + return 0 + fi + + find "$REPO_ROOT/data/processed" -maxdepth 2 -type d -name ".tensor_cache" -prune -exec rm -rf {} + + find "$REPO_ROOT/data/processed" -maxdepth 2 -type f -name "normalization_stats_*.yaml" -delete +} + +echo "Launch commit: $LAUNCH_COMMIT" +if [[ "$VALIDATE_PROCESSED_ARTIFACTS" == "1" ]]; then + echo "Validating processed artifacts..." + ( + cd "$REPO_ROOT" + validate_processed_artifacts + ) + echo "Processed artifacts are ready." +fi +purge_runtime_caches + +mkdir -p "$LOG_DIR" + +main_classes=( + core_ssl_benchmark + label_efficiency + cross_dataset_transfer + hp_robustness + classical_baselines +) +fairness_classes=("${main_classes[@]}") +appendix_classes=() + +if [[ "$INCLUDE_TS2VEC_EXTENSION" == "1" ]]; then + main_classes+=(ts2vec_extension) + fairness_classes+=(ts2vec_extension) +fi + +if [[ "$INCLUDE_CAPACITY_STUDY" == "1" ]]; then + main_classes+=(capacity_study) + fairness_classes+=(capacity_study) +fi + +if [[ "$INCLUDE_SMART_REFERENCE" == "1" ]]; then + appendix_classes+=(smart_external_reference) + fairness_classes+=(smart_external_reference) +fi + +warmup_classes=("${main_classes[@]}") +if [[ "$INCLUDE_SMART_REFERENCE" == "1" ]]; then + warmup_classes+=(smart_external_reference) +fi + +class_expected_count() { + case "$1" in + core_ssl_benchmark) echo 465 ;; + label_efficiency) echo 1155 ;; + cross_dataset_transfer) echo 120 ;; + hp_robustness) echo 150 ;; + capacity_study) echo 100 ;; + classical_baselines) echo 330 ;; + ts2vec_extension) echo 135 ;; + smart_external_reference) echo 135 ;; + *) return 1 ;; + esac +} + +validate_final_corpus() { + if [[ "$VALIDATE_FINAL_CORPUS" != "1" ]]; then + echo "WARNING: final corpus validation disabled by VALIDATE_FINAL_CORPUS=$VALIDATE_FINAL_CORPUS." >&2 + return 0 + fi + + local selected=("${main_classes[@]}" "${appendix_classes[@]}") + local seen=" " + local total=0 + local class count + + for class in "${selected[@]}"; do + case "$seen" in + *" $class "*) + echo "Duplicate experiment class in final corpus: $class" >&2 + exit 1 + ;; + esac + seen="${seen}${class} " + if ! count="$(class_expected_count "$class")"; then + echo "Unknown experiment class in final corpus: $class" >&2 + exit 1 + fi + total=$((total + count)) + done + + if (( total != EXPECTED_TOTAL_RUNS )); then + echo "Final corpus preflight failed: selected $total runs, expected $EXPECTED_TOTAL_RUNS." >&2 + echo "Selected classes:${seen}" >&2 + echo "Unset INCLUDE_* overrides for the thesis launch, or set VALIDATE_FINAL_CORPUS=0 only for exploratory runs." >&2 + exit 1 + fi + + echo "Final corpus preflight: $total planned runs across ${#selected[@]} classes." +} + +validate_final_corpus + +quote_cmd() { + printf "%q " "$@" +} + +run_args=(--project "$WANDB_PROJECT" --entity "$WANDB_ENTITY") +export_args=(--project "$WANDB_PROJECT" --entity "$WANDB_ENTITY") +fairness_args=(--project "$WANDB_PROJECT" --entity "$WANDB_ENTITY") + +run_args+=(--revision "$REVISION" --reason "$REASON" --launch-commit "$LAUNCH_COMMIT") +export_args+=(--revision "$REVISION" --output-dir "$RESULTS_DIR") +if [[ -n "$REVISION" ]]; then + fairness_args+=(--revision "$REVISION") +fi +export_args+=(--experiment-class "${fairness_classes[@]}") +fairness_args+=(--experiment-class "${fairness_classes[@]}" --batch-size "$BATCH_SIZE_FAIRNESS" --device "$DEVICE_FAIRNESS" --force) + +warmup_cmd=(uv run python scripts/internal/run_experiments.py warmup --experiment-class "${warmup_classes[@]}") +main_cmd=(uv run python scripts/internal/run_experiments.py run --experiment-class "${main_classes[@]}" --parallel "$PARALLEL_MAIN" "${run_args[@]}") +fairness_cmd=(uv run python scripts/eval/evaluate_fairness.py "${fairness_args[@]}") +export_cmd=(uv run python scripts/export_results.py "${export_args[@]}") + +if ((${#appendix_classes[@]} > 0)); then + appendix_cmd=(uv run python scripts/internal/run_experiments.py run --experiment-class "${appendix_classes[@]}" --parallel "$PARALLEL_APPENDIX" "${run_args[@]}") +else + appendix_cmd=() +fi + +timestamp="$(date +%Y%m%d-%H%M%S)" +runner_script="$LOG_DIR/thesis-run-${timestamp}.sh" +runner_log="$LOG_DIR/thesis-run-${timestamp}.log" +status_log="$LOG_DIR/thesis-status-${timestamp}.log" + +printf -v warmup_line "%q " "${warmup_cmd[@]}" +printf -v main_line "%q " "${main_cmd[@]}" +printf -v fairness_line "%q " "${fairness_cmd[@]}" +printf -v export_line "%q " "${export_cmd[@]}" + +appendix_block="" +if ((${#appendix_cmd[@]} > 0)); then + printf -v appendix_line "%q " "${appendix_cmd[@]}" + appendix_block="${appendix_line}"$'\n' +fi + +export_block="" +if [[ "$RUN_EXPORT" == "1" ]]; then + export_block="${export_line}"$'\n' +fi + +cat > "$runner_script" <&2 + exit 1 + fi +fi + +uv sync --dev --locked +${warmup_line} +${main_line} +${appendix_block}${fairness_line} +${export_block}echo "Finished: \$(date -Iseconds)" +EOF + +chmod +x "$runner_script" + +status_loop="cd $(printf "%q" "$REPO_ROOT"); while true; do clear; date; echo; " +status_loop+="uv run python scripts/internal/run_experiments.py status " +status_loop+="--revision $(printf "%q" "$REVISION"); " +status_loop+="sleep $(printf "%q" "$STATUS_INTERVAL"); done" +status_cmd=(bash -lc "$status_loop") +printf -v status_line "%q " "${status_cmd[@]}" + +tmux new-session -d -s "$SESSION_NAME" -n run "bash $(printf "%q" "$runner_script") 2>&1 | tee $(printf "%q" "$runner_log")" +tmux new-window -t "$SESSION_NAME" -n status "$status_line" +tmux pipe-pane -o -t "$SESSION_NAME:status" "cat >> $(printf "%q" "$status_log")" + +echo "Created tmux session: $SESSION_NAME" +echo "Run log: $runner_log" +echo "Status log: $status_log" +echo "Attach with: tmux attach -t $SESSION_NAME" diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py new file mode 100644 index 0000000..0988c91 --- /dev/null +++ b/scripts/internal/run_experiments.py @@ -0,0 +1,2202 @@ +#!/usr/bin/env python3 +"""Class-based experiment runner for the final SLICES thesis rerun corpus. + +Generates the thesis experiment matrix by scientific experiment class, resolves +pretrain dependencies, and executes runs with crash recovery and resumable state. + +Usage: + uv run python scripts/internal/run_experiments.py warmup \ + --experiment-class core_ssl_benchmark label_efficiency + uv run python scripts/internal/run_experiments.py run \ + --experiment-class core_ssl_benchmark label_efficiency \ + --project slices-thesis --revision thesis-v1 --entity \ + --launch-commit + uv run python scripts/internal/run_experiments.py run \ + --experiment-class core_ssl_benchmark --dry-run \ + --project slices-thesis --revision thesis-v1 --entity \ + --launch-commit + uv run python scripts/internal/run_experiments.py status \ + --experiment-class core_ssl_benchmark --revision thesis-v1 + uv run python scripts/internal/run_experiments.py retry --failed --skipped \ + --experiment-class core_ssl_benchmark --revision thesis-v1 \ + --project slices-thesis --entity --launch-commit +""" +from __future__ import annotations + +import argparse +import json +import os +import re +import shlex +import signal +import subprocess +import sys +import tempfile +import time +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path + +from slices.constants import ( + FULL_FINETUNE_PROTOCOL, + THESIS_TASKS, + downstream_protocol_from_freeze, +) + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +EXPERIMENT_CLASSES = [ + "core_ssl_benchmark", + "label_efficiency", + "cross_dataset_transfer", + "hp_robustness", + "capacity_study", + "classical_baselines", + "ts2vec_extension", + "smart_external_reference", +] + +DEFAULT_EXPERIMENT_CLASSES = list(EXPERIMENT_CLASSES) + +DOWNSTREAM_EXPERIMENT_CLASSES = [ + "core_ssl_benchmark", + "label_efficiency", + "cross_dataset_transfer", + "hp_robustness", + "capacity_study", + "classical_baselines", + "ts2vec_extension", + "smart_external_reference", +] + +EXPECTED_CLASS_COUNTS = { + "core_ssl_benchmark": 465, + "label_efficiency": 1155, + "cross_dataset_transfer": 120, + "hp_robustness": 150, + "capacity_study": 100, + "classical_baselines": 330, + "ts2vec_extension": 135, + "smart_external_reference": 135, +} + +SSL_PARADIGMS = ["mae", "jepa", "contrastive"] +DATASETS = ["miiv", "eicu", "combined"] +TASKS = list(THESIS_TASKS) +SEEDS_EXTENDED = [42, 123, 456, 789, 1011] +LABEL_FRACTIONS_FULL = [0.01, 0.05, 0.1, 0.25, 0.5] +LABEL_FRACTIONS_MORTALITY_24H = [0.05, 0.1, 0.25, 0.5] +LABEL_FRACTIONS_TREND = [0.1] +LABEL_FRACTIONS_CAPACITY = [0.05, 0.1, 0.5] + +MODEL_SIZES = { + "medium": { + "model": "transformer_medium", + "ssl_scale": { + "mae": { + "ssl.decoder_d_model": 128, + "ssl.decoder_n_layers": 2, + "ssl.decoder_n_heads": 8, + "ssl.decoder_d_ff": 512, + }, + }, + }, + "large": { + "model": "transformer_large", + "ssl_scale": { + "mae": { + "ssl.decoder_d_model": 256, + "ssl.decoder_n_layers": 2, + "ssl.decoder_n_heads": 8, + "ssl.decoder_d_ff": 1024, + }, + }, + }, +} + +LR_ROBUSTNESS = [2e-4, 5e-4, 2e-3] +MASK_RATIO_ROBUSTNESS = [0.3, 0.75] +VIEW_MASK_EXPERIMENT_SUBTYPE = "view_mask_sensitivity" +TRANSFER_PAIRS = [("miiv", "eicu"), ("eicu", "miiv")] + +STATE_FILE = Path("outputs/experiment_state.json") +LOG_DIR = Path("logs/runner") +LAUNCH_IDENTITY_FILE = ".runner_launch_identity.json" +LAUNCH_IDENTITY_KEYS = ("revision", "wandb_project", "wandb_entity", "launch_commit") + +# --------------------------------------------------------------------------- +# Data Model +# --------------------------------------------------------------------------- + + +@dataclass +class Run: + id: str + experiment_class: str + run_type: str # "pretrain" | "finetune" | "supervised" | "gru_d" | "xgboost" + paradigm: str + dataset: str + seed: int + output_dir: str + run_key: str | None = None + depends_on: list[str] = field(default_factory=list) + task: str | None = None + label_fraction: float = 1.0 + freeze_encoder: bool | None = None + extra_overrides: dict = field(default_factory=dict) + source_dataset: str | None = None + upstream_pretrain_lr: float | None = None + upstream_pretrain_mask_ratio: float | None = None + experiment_subtype: str | None = None + model_size: str | None = None + + def __post_init__(self) -> None: + if self.run_key is None: + self.run_key = self.id + + @property + def phase(self) -> str: + if self.run_type == "pretrain": + return "pretrain" + if self.run_type == "supervised": + return "supervised" + if self.run_type in {"gru_d", "xgboost"}: + return "baseline" + return "finetune" + + @property + def protocol(self) -> str | None: + protocol = downstream_protocol_from_freeze(self.freeze_encoder) + if protocol is not None: + return protocol + if self.run_type in {"supervised", "gru_d", "xgboost"}: + return FULL_FINETUNE_PROTOCOL + return None + + @property + def protocol_selector(self) -> str | None: + return self.protocol + + def build_command(self, runs_by_id: dict[str, "Run"]) -> list[str]: + """Build the subprocess command for this run.""" + if self.run_type == "pretrain": + return self._pretrain_cmd() + if self.run_type == "finetune": + return self._finetune_cmd(runs_by_id) + if self.run_type == "supervised": + return self._supervised_cmd() + if self.run_type == "gru_d": + return self._gru_d_cmd() + if self.run_type == "xgboost": + return self._xgboost_cmd() + raise ValueError(f"Unknown run_type: {self.run_type}") + + def _metadata_overrides(self) -> list[str]: + overrides = [ + f"experiment_class={self.experiment_class}", + f"+phase={self.phase}", + ] + if self.experiment_subtype is not None: + overrides.append(f"experiment_subtype={self.experiment_subtype}") + if self.protocol_selector is not None and self.run_type == "finetune": + overrides.append(f"protocol={self.protocol_selector}") + if self.model_size is not None: + overrides.append(f"+model_size={self.model_size}") + for k, v in self.extra_overrides.items(): + overrides.append(f"{k}={v}") + return overrides + + def _append_resume(self, cmd: list[str]) -> None: + last_ckpt = Path(self.output_dir) / "checkpoints" / "last.ckpt" + if last_ckpt.exists() and _can_resume_from_last_checkpoint(self): + cmd.append(f"ckpt_path={last_ckpt}") + + def _pretrain_cmd(self) -> list[str]: + cmd = [ + "uv", + "run", + "python", + "scripts/training/pretrain.py", + f"dataset={self.dataset}", + f"ssl={self.paradigm}", + f"seed={self.seed}", + f"hydra.run.dir={self.output_dir}", + ] + self._append_resume(cmd) + cmd.extend(self._metadata_overrides()) + return cmd + + def _finetune_cmd(self, runs_by_id: dict[str, "Run"]) -> list[str]: + pretrain_ids = [dep_id for dep_id in self.depends_on if dep_id in runs_by_id] + if len(pretrain_ids) != 1: + raise ValueError(f"Finetune run {self.id} must have exactly one pretrain dependency") + pretrain_dir = runs_by_id[pretrain_ids[0]].output_dir + + cmd = [ + "uv", + "run", + "python", + "scripts/training/finetune.py", + f"dataset={self.dataset}", + f"checkpoint={pretrain_dir}/encoder.pt", + f"tasks={self.task}", + f"seed={self.seed}", + f"hydra.run.dir={self.output_dir}", + ] + self._append_resume(cmd) + if self.label_fraction < 1.0: + cmd.append(f"label_fraction={self.label_fraction}") + if self.source_dataset is not None: + cmd.append(f"+source_dataset={self.source_dataset}") + if self.upstream_pretrain_lr is not None: + cmd.append(f"+upstream_pretrain_lr={self.upstream_pretrain_lr}") + if self.upstream_pretrain_mask_ratio is not None: + cmd.append(f"+upstream_pretrain_mask_ratio={self.upstream_pretrain_mask_ratio}") + cmd.extend(self._metadata_overrides()) + return cmd + + def _supervised_cmd(self) -> list[str]: + cmd = [ + "uv", + "run", + "python", + "scripts/training/supervised.py", + f"dataset={self.dataset}", + f"tasks={self.task}", + f"seed={self.seed}", + f"hydra.run.dir={self.output_dir}", + ] + self._append_resume(cmd) + if self.label_fraction < 1.0: + cmd.append(f"label_fraction={self.label_fraction}") + cmd.extend(self._metadata_overrides()) + return cmd + + def _gru_d_cmd(self) -> list[str]: + cmd = [ + "uv", + "run", + "python", + "scripts/training/supervised.py", + "--config-name", + "gru_d", + f"dataset={self.dataset}", + f"tasks={self.task}", + f"seed={self.seed}", + f"hydra.run.dir={self.output_dir}", + ] + self._append_resume(cmd) + if self.label_fraction < 1.0: + cmd.append(f"label_fraction={self.label_fraction}") + cmd.extend(self._metadata_overrides()) + return cmd + + def _xgboost_cmd(self) -> list[str]: + cmd = [ + "uv", + "run", + "python", + "scripts/training/xgboost_baseline.py", + f"dataset={self.dataset}", + f"tasks={self.task}", + f"seed={self.seed}", + f"hydra.run.dir={self.output_dir}", + ] + if self.label_fraction < 1.0: + cmd.append(f"label_fraction={self.label_fraction}") + cmd.extend(self._metadata_overrides()) + return cmd + + +# --------------------------------------------------------------------------- +# Matrix Generation +# --------------------------------------------------------------------------- + + +def _short_value(value) -> str: + return str(value).replace(".", "") + + +def _name_suffix(extra: dict | None) -> str: + if not extra: + return "" + parts = [] + for key, value in sorted(extra.items()): + short_key = key.split(".")[-1].replace("_", "") + parts.append(f"{short_key}{_short_value(value)}") + return "_" + "_".join(parts) + + +def _pretrain_key(paradigm: str, dataset: str, seed: int, extra: dict | None = None) -> str: + return f"pretrain_{paradigm}_{dataset}_seed{seed}{_name_suffix(extra)}" + + +def _output_dir(experiment_class: str, run_key: str) -> str: + return f"outputs/{experiment_class}/{run_key}" + + +class MatrixBuilder: + """Generate the final thesis matrix by experiment class.""" + + def __init__(self) -> None: + self.runs: list[Run] = [] + self.pretrain_index: dict[tuple[str, str, int, tuple[tuple[str, str], ...]], Run] = {} + + def _pretrain_index_key( + self, + paradigm: str, + dataset: str, + seed: int, + extra: dict | None = None, + ) -> tuple[str, str, int, tuple[tuple[str, str], ...]]: + extra_items = tuple(sorted((str(k), str(v)) for k, v in (extra or {}).items())) + return paradigm, dataset, seed, extra_items + + def _add_run(self, run: Run) -> Run: + self.runs.append(run) + return run + + def _add_pretrain( + self, + experiment_class: str, + paradigm: str, + dataset: str, + seed: int, + extra: dict | None = None, + experiment_subtype: str | None = None, + model_size: str | None = None, + ) -> Run: + extra = extra or {} + index_key = self._pretrain_index_key(paradigm, dataset, seed, extra) + if index_key in self.pretrain_index: + return self.pretrain_index[index_key] + + run_key = _pretrain_key(paradigm, dataset, seed, extra) + run = Run( + id=f"{experiment_class}_{run_key}", + run_key=run_key, + experiment_class=experiment_class, + run_type="pretrain", + paradigm=paradigm, + dataset=dataset, + seed=seed, + output_dir=_output_dir(experiment_class, run_key), + extra_overrides=extra, + experiment_subtype=experiment_subtype, + model_size=model_size, + ) + self._add_run(run) + self.pretrain_index[index_key] = run + return run + + def _get_pretrain( + self, + paradigm: str, + dataset: str, + seed: int, + extra: dict | None = None, + ) -> Run: + index_key = self._pretrain_index_key(paradigm, dataset, seed, extra) + if index_key not in self.pretrain_index: + raise KeyError(f"Pretrain not found: {_pretrain_key(paradigm, dataset, seed, extra)}") + return self.pretrain_index[index_key] + + def _add_finetune( + self, + experiment_class: str, + paradigm: str, + dataset: str, + seed: int, + task: str, + freeze: bool, + pretrain_run: Run, + label_fraction: float = 1.0, + extra: dict | None = None, + source_dataset: str | None = None, + name_extra: dict | None = None, + upstream_pretrain_lr: float | None = None, + upstream_pretrain_mask_ratio: float | None = None, + experiment_subtype: str | None = None, + model_size: str | None = None, + ) -> Run: + prefix = "probe" if freeze else "finetune" + run_key = f"{prefix}_{paradigm}_{task}_{dataset}_seed{seed}" + if source_dataset: + run_key += f"_from_{source_dataset}" + if label_fraction < 1.0: + run_key += f"_frac{_short_value(label_fraction)}" + run_key += _name_suffix(name_extra) + run_key += _name_suffix(extra) + + return self._add_run( + Run( + id=f"{experiment_class}_{run_key}", + run_key=run_key, + experiment_class=experiment_class, + run_type="finetune", + paradigm=paradigm, + dataset=dataset, + seed=seed, + output_dir=_output_dir(experiment_class, run_key), + depends_on=[pretrain_run.id], + task=task, + label_fraction=label_fraction, + freeze_encoder=freeze, + extra_overrides=extra or {}, + source_dataset=source_dataset, + upstream_pretrain_lr=upstream_pretrain_lr, + upstream_pretrain_mask_ratio=upstream_pretrain_mask_ratio, + experiment_subtype=experiment_subtype, + model_size=model_size, + ) + ) + + def _add_supervised( + self, + experiment_class: str, + dataset: str, + seed: int, + task: str, + label_fraction: float = 1.0, + extra: dict | None = None, + model_size: str | None = None, + ) -> Run: + run_key = f"supervised_{task}_{dataset}_seed{seed}" + if model_size: + run_key += f"_{model_size}" + if label_fraction < 1.0: + run_key += f"_frac{_short_value(label_fraction)}" + return self._add_run( + Run( + id=f"{experiment_class}_{run_key}", + run_key=run_key, + experiment_class=experiment_class, + run_type="supervised", + paradigm="supervised", + dataset=dataset, + seed=seed, + output_dir=_output_dir(experiment_class, run_key), + task=task, + label_fraction=label_fraction, + extra_overrides=extra or {}, + model_size=model_size, + ) + ) + + def _add_gru_d( + self, + experiment_class: str, + dataset: str, + seed: int, + task: str, + label_fraction: float = 1.0, + ) -> Run: + run_key = f"gru_d_{task}_{dataset}_seed{seed}" + if label_fraction < 1.0: + run_key += f"_frac{_short_value(label_fraction)}" + return self._add_run( + Run( + id=f"{experiment_class}_{run_key}", + run_key=run_key, + experiment_class=experiment_class, + run_type="gru_d", + paradigm="gru_d", + dataset=dataset, + seed=seed, + output_dir=_output_dir(experiment_class, run_key), + task=task, + label_fraction=label_fraction, + ) + ) + + def _add_xgboost( + self, + experiment_class: str, + dataset: str, + seed: int, + task: str, + label_fraction: float = 1.0, + ) -> Run: + run_key = f"xgboost_{task}_{dataset}_seed{seed}" + if label_fraction < 1.0: + run_key += f"_frac{_short_value(label_fraction)}" + return self._add_run( + Run( + id=f"{experiment_class}_{run_key}", + run_key=run_key, + experiment_class=experiment_class, + run_type="xgboost", + paradigm="xgboost", + dataset=dataset, + seed=seed, + output_dir=_output_dir(experiment_class, run_key), + task=task, + label_fraction=label_fraction, + ) + ) + + def build_core_ssl_benchmark(self) -> None: + experiment_class = "core_ssl_benchmark" + for seed in SEEDS_EXTENDED: + for dataset in DATASETS: + for paradigm in SSL_PARADIGMS: + pretrain = self._add_pretrain(experiment_class, paradigm, dataset, seed) + for task in TASKS: + self._add_finetune( + experiment_class, paradigm, dataset, seed, task, False, pretrain + ) + self._add_finetune( + experiment_class, paradigm, dataset, seed, task, True, pretrain + ) + for task in TASKS: + self._add_supervised(experiment_class, dataset, seed, task) + + def build_label_efficiency(self) -> None: + experiment_class = "label_efficiency" + for seed in SEEDS_EXTENDED: + for dataset in DATASETS: + for paradigm in SSL_PARADIGMS: + pretrain = self._get_pretrain(paradigm, dataset, seed) + for frac in LABEL_FRACTIONS_MORTALITY_24H: + self._add_finetune( + experiment_class, + paradigm, + dataset, + seed, + "mortality_24h", + False, + pretrain, + frac, + ) + self._add_finetune( + experiment_class, + paradigm, + dataset, + seed, + "mortality_24h", + True, + pretrain, + frac, + ) + for task in TASKS[1:]: + fractions = ( + LABEL_FRACTIONS_FULL + if task == "mortality_hospital" + else LABEL_FRACTIONS_TREND + ) + for frac in fractions: + self._add_finetune( + experiment_class, + paradigm, + dataset, + seed, + task, + False, + pretrain, + frac, + ) + self._add_finetune( + experiment_class, + paradigm, + dataset, + seed, + task, + True, + pretrain, + frac, + ) + for frac in LABEL_FRACTIONS_MORTALITY_24H: + self._add_supervised(experiment_class, dataset, seed, "mortality_24h", frac) + for task in TASKS[1:]: + fractions = ( + LABEL_FRACTIONS_FULL + if task == "mortality_hospital" + else LABEL_FRACTIONS_TREND + ) + for frac in fractions: + self._add_supervised(experiment_class, dataset, seed, task, frac) + + def build_cross_dataset_transfer(self) -> None: + experiment_class = "cross_dataset_transfer" + for seed in SEEDS_EXTENDED: + for source_dataset, target_dataset in TRANSFER_PAIRS: + for paradigm in SSL_PARADIGMS: + pretrain = self._get_pretrain(paradigm, source_dataset, seed) + for task in TASKS: + self._add_finetune( + experiment_class, + paradigm, + target_dataset, + seed, + task, + False, + pretrain, + source_dataset=source_dataset, + ) + + def build_hp_robustness(self) -> None: + experiment_class = "hp_robustness" + task = "mortality_24h" + dataset = "miiv" + for seed in SEEDS_EXTENDED: + for paradigm in SSL_PARADIGMS: + for lr in LR_ROBUSTNESS: + extra = {"optimizer.lr": lr} + subtype = "lr_sensitivity" + pretrain = self._add_pretrain( + experiment_class, + paradigm, + dataset, + seed, + extra, + experiment_subtype=subtype, + ) + self._add_finetune( + experiment_class, + paradigm, + dataset, + seed, + task, + False, + pretrain, + name_extra=extra, + upstream_pretrain_lr=lr, + experiment_subtype=subtype, + ) + for mask_ratio in MASK_RATIO_ROBUSTNESS: + extra = {"ssl.mask_ratio": mask_ratio} + subtype = "mask_ratio_sensitivity" + if paradigm == "contrastive": + # Contrastive complementary views make the per-view mask + # budgets asymmetric away from 0.5. The robustness slice + # therefore uses independent views and is labeled as a + # view/mask sensitivity test rather than a pure + # mask-ratio sweep. + extra["ssl.complementary_masks"] = False + subtype = VIEW_MASK_EXPERIMENT_SUBTYPE + pretrain = self._add_pretrain( + experiment_class, + paradigm, + dataset, + seed, + extra, + experiment_subtype=subtype, + ) + self._add_finetune( + experiment_class, + paradigm, + dataset, + seed, + task, + False, + pretrain, + name_extra=extra, + upstream_pretrain_mask_ratio=mask_ratio, + experiment_subtype=subtype, + ) + + def build_capacity_study(self) -> None: + experiment_class = "capacity_study" + dataset = "miiv" + task = "mortality_24h" + for seed in SEEDS_EXTENDED: + for model_size, size_cfg in MODEL_SIZES.items(): + model_extra = {"model": size_cfg["model"]} + pretrain_extra = {**model_extra, **size_cfg["ssl_scale"]["mae"]} + pretrain = self._add_pretrain( + experiment_class, + "mae", + dataset, + seed, + pretrain_extra, + model_size=model_size, + ) + for frac in LABEL_FRACTIONS_CAPACITY: + self._add_finetune( + experiment_class, + "mae", + dataset, + seed, + task, + False, + pretrain, + frac, + extra=model_extra, + name_extra={"size": model_size}, + model_size=model_size, + ) + self._add_finetune( + experiment_class, + "mae", + dataset, + seed, + task, + True, + pretrain, + frac, + extra=model_extra, + name_extra={"size": model_size}, + model_size=model_size, + ) + self._add_supervised( + experiment_class, + dataset, + seed, + task, + frac, + extra=model_extra, + model_size=model_size, + ) + + def build_classical_baselines(self) -> None: + experiment_class = "classical_baselines" + for seed in SEEDS_EXTENDED: + for dataset in DATASETS: + for task in TASKS: + self._add_xgboost(experiment_class, dataset, seed, task) + self._add_gru_d(experiment_class, dataset, seed, task) + for frac in LABEL_FRACTIONS_MORTALITY_24H: + self._add_xgboost(experiment_class, dataset, seed, "mortality_24h", frac) + self._add_gru_d(experiment_class, dataset, seed, "mortality_24h", frac) + for task in TASKS[1:]: + for frac in LABEL_FRACTIONS_TREND: + self._add_xgboost(experiment_class, dataset, seed, task, frac) + self._add_gru_d(experiment_class, dataset, seed, task, frac) + + def build_ts2vec_extension(self) -> None: + experiment_class = "ts2vec_extension" + for seed in SEEDS_EXTENDED: + for dataset in DATASETS: + pretrain = self._add_pretrain(experiment_class, "ts2vec", dataset, seed) + for task in TASKS: + self._add_finetune( + experiment_class, "ts2vec", dataset, seed, task, False, pretrain + ) + self._add_finetune( + experiment_class, "ts2vec", dataset, seed, task, True, pretrain + ) + + def build_smart_external_reference(self) -> None: + experiment_class = "smart_external_reference" + pretrain_extra = {"model": "smart"} + finetune_extra = {"model": "smart"} + model_size = "default" + for seed in SEEDS_EXTENDED: + for dataset in DATASETS: + pretrain = self._add_pretrain( + experiment_class, + "smart", + dataset, + seed, + pretrain_extra, + model_size=model_size, + ) + for task in TASKS: + self._add_finetune( + experiment_class, + "smart", + dataset, + seed, + task, + False, + pretrain, + extra=finetune_extra, + model_size=model_size, + ) + self._add_finetune( + experiment_class, + "smart", + dataset, + seed, + task, + True, + pretrain, + extra=finetune_extra, + model_size=model_size, + ) + + def build_all(self) -> list[Run]: + self.build_core_ssl_benchmark() + self.build_label_efficiency() + self.build_cross_dataset_transfer() + self.build_hp_robustness() + self.build_capacity_study() + self.build_classical_baselines() + self.build_ts2vec_extension() + self.build_smart_external_reference() + return self.runs + + +def generate_all_runs() -> list[Run]: + return MatrixBuilder().build_all() + + +def scientific_fingerprint(run: Run) -> tuple: + """Scientific identity used by matrix snapshot tests.""" + pretrain_overrides = tuple(sorted((str(k), str(v)) for k, v in run.extra_overrides.items())) + return ( + run.run_type, + run.experiment_class, + run.experiment_subtype, + run.paradigm, + run.dataset, + run.task, + run.seed, + run.protocol, + run.label_fraction, + run.source_dataset, + run.model_size, + run.upstream_pretrain_lr, + run.upstream_pretrain_mask_ratio, + pretrain_overrides if run.run_type == "pretrain" else (), + ) + + +def apply_revision(runs: list[Run], revision: str, reason: str | None = None) -> list[Run]: + """Inject revision metadata into IDs, output dirs, dependencies, and overrides.""" + old_to_new: dict[str, str] = {} + + for run in runs: + old_id = run.id + run.id = f"{run.experiment_class}_rev-{revision}_{run.run_key}" + old_to_new[old_id] = run.id + run.output_dir = _output_dir( + f"{run.experiment_class}_rev-{revision}", run.run_key or old_id + ) + run.extra_overrides["revision"] = revision + if reason: + run.extra_overrides["rerun_reason"] = reason + + for run in runs: + run.depends_on = [old_to_new.get(dep_id, dep_id) for dep_id in run.depends_on] + + return runs + + +def apply_wandb_target( + runs: list[Run], + project: str | None = None, + entity: str | None = None, +) -> list[Run]: + """Inject W&B project/entity overrides into generated runs.""" + if project is None and entity is None: + return runs + for run in runs: + if project is not None: + run.extra_overrides["project_name"] = project + run.extra_overrides["logging.wandb_project"] = project + if entity is not None: + run.extra_overrides["logging.wandb_entity"] = entity + return runs + + +def apply_launch_commit(runs: list[Run], launch_commit: str | None = None) -> list[Run]: + """Inject the exact git commit used for launch into each Hydra config.""" + if not launch_commit: + return runs + for run in runs: + run.extra_overrides["+launch_commit"] = launch_commit + return runs + + +def _git_stdout(args: list[str]) -> str: + result = subprocess.run( + ["git", *args], + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0: + message = result.stderr.strip() or result.stdout.strip() or shlex.join(["git", *args]) + raise RuntimeError(message) + return result.stdout.strip() + + +def _git_quiet(args: list[str]) -> bool: + result = subprocess.run( + ["git", *args], + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + text=True, + check=False, + ) + if result.returncode in {0, 1}: + return result.returncode == 0 + message = result.stderr.strip() or shlex.join(["git", *args]) + raise RuntimeError(message) + + +def _resolve_launch_commit(launch_commit: str) -> str: + """Resolve a user-provided commit-ish to an immutable git commit hash.""" + return _git_stdout(["rev-parse", "--verify", f"{launch_commit}^{{commit}}"]) + + +def _validated_final_launch_commit(launch_commit: str) -> tuple[str | None, str | None]: + """Return (resolved_commit, error) for final launch provenance checks.""" + try: + resolved_launch_commit = _resolve_launch_commit(launch_commit) + current_commit = _git_stdout(["rev-parse", "--verify", "HEAD"]) + if current_commit != resolved_launch_commit: + return ( + None, + "final launch commit does not match the current checkout " + f"(current={current_commit}, launch_commit={resolved_launch_commit})", + ) + if not _git_quiet(["diff", "--quiet"]) or not _git_quiet(["diff", "--cached", "--quiet"]): + return ( + None, + "final run/retry requires a clean tracked worktree. " + "Commit or stash tracked changes first; dry runs may be used for local audits.", + ) + except RuntimeError as exc: + return None, f"could not validate final launch git provenance: {exc}" + return resolved_launch_commit, None + + +def _validate_clean_final_launch_state(launch_commit: str) -> str | None: + """Return an error string if a final launch would have ambiguous provenance.""" + _, error = _validated_final_launch_commit(launch_commit) + return error + + +def _validate_wandb_online_mode() -> str | None: + mode = os.environ.get("WANDB_MODE") + if mode and mode.lower() != "online": + return f"final run/retry requires WANDB_MODE to be unset or 'online', got {mode!r}" + return None + + +def _skip_launch_git_check_enabled() -> bool: + """Return whether the explicit launcher escape hatch should bypass git checks.""" + return os.environ.get("SLICES_SKIP_LAUNCH_GIT_CHECK", "").lower() in { + "1", + "true", + "yes", + } or os.environ.get("SKIP_LAUNCH_GIT_CHECK", "").lower() in {"1", "true", "yes"} + + +def validate_direct_final_launch_policy(args, parser: argparse.ArgumentParser) -> None: + """Require auditable provenance for direct final run/retry invocations.""" + if args.command not in {"run", "retry"}: + return + + launch_commit = getattr(args, "launch_commit", None) + is_run_dry_run = args.command == "run" and getattr(args, "dry_run", False) + if is_run_dry_run: + if not launch_commit: + print( + "WARNING: dry run has no launch provenance. Final run/retry requires " + "--launch-commit or SLICES_LAUNCH_COMMIT.", + file=sys.stderr, + ) + else: + try: + args.launch_commit = _resolve_launch_commit(str(launch_commit)) + except RuntimeError as exc: + parser.error(f"could not validate dry-run launch git provenance: {exc}") + return + + if not launch_commit: + parser.error( + "final run/retry requires --launch-commit or SLICES_LAUNCH_COMMIT " "for W&B provenance" + ) + + wandb_mode_error = _validate_wandb_online_mode() + if wandb_mode_error: + parser.error(wandb_mode_error) + + if _skip_launch_git_check_enabled(): + print( + "WARNING: skipping final launch git provenance validation because " + "SLICES_SKIP_LAUNCH_GIT_CHECK/SKIP_LAUNCH_GIT_CHECK is enabled.", + file=sys.stderr, + ) + return + + resolved_launch_commit, error = _validated_final_launch_commit(str(launch_commit)) + if error: + parser.error(error) + args.launch_commit = resolved_launch_commit + + +# --------------------------------------------------------------------------- +# State Management +# --------------------------------------------------------------------------- + + +def load_state() -> dict: + if STATE_FILE.exists(): + return json.loads(STATE_FILE.read_text()) + return {"version": 1, "runs": {}} + + +def save_state(state: dict) -> None: + STATE_FILE.parent.mkdir(parents=True, exist_ok=True) + fd, tmp = tempfile.mkstemp(dir=STATE_FILE.parent, suffix=".tmp") + try: + with os.fdopen(fd, "w") as f: + json.dump(state, f, indent=2) + os.replace(tmp, STATE_FILE) + except BaseException: + try: + os.unlink(tmp) + except OSError: + pass + raise + + +def _normalize_identity_value(value) -> str | None: + if value is None: + return None + text = str(value) + return text if text else None + + +def _normalize_launch_identity(identity: dict | None) -> dict[str, str | None]: + if not isinstance(identity, dict): + identity = {} + return {key: _normalize_identity_value(identity.get(key)) for key in LAUNCH_IDENTITY_KEYS} + + +def _override_value(overrides: dict, *keys: str) -> str | None: + for key in keys: + value = overrides.get(key) + if value is not None and str(value): + return str(value) + return None + + +def _run_launch_identity(run: Run) -> dict[str, str | None]: + return _normalize_launch_identity( + { + "revision": _override_value(run.extra_overrides, "revision"), + "wandb_project": _override_value( + run.extra_overrides, + "logging.wandb_project", + "project_name", + ), + "wandb_entity": _override_value(run.extra_overrides, "logging.wandb_entity"), + "launch_commit": _run_launch_commit(run), + } + ) + + +def _command_override(command: str | None, *keys: str) -> str | None: + if not command: + return None + try: + tokens = shlex.split(command) + except ValueError: + tokens = str(command).split() + + key_set = set(keys) + for token in tokens: + if "=" not in token: + continue + key, value = token.split("=", 1) + key = key.lstrip("+") + if key in key_set and value: + return value + return None + + +def _state_launch_identity(info: dict | None) -> dict[str, str | None]: + if not isinstance(info, dict): + return _normalize_launch_identity(None) + + identity = {} + if isinstance(info.get("launch_identity"), dict): + identity.update(info["launch_identity"]) + + for key in LAUNCH_IDENTITY_KEYS: + if not identity.get(key) and info.get(key): + identity[key] = info[key] + + command = str(info.get("command") or "") + if not identity.get("revision"): + identity["revision"] = _command_override(command, "revision") + if not identity.get("wandb_project"): + identity["wandb_project"] = _command_override( + command, + "logging.wandb_project", + "project_name", + ) + if not identity.get("wandb_entity"): + identity["wandb_entity"] = _command_override(command, "logging.wandb_entity") + if not identity.get("launch_commit"): + identity["launch_commit"] = _state_launch_commit(info) + + return _normalize_launch_identity(identity) + + +def _has_scoped_launch_identity(identity: dict[str, str | None]) -> bool: + return any(identity.get(key) is not None for key in LAUNCH_IDENTITY_KEYS) + + +def _launch_identity_matches( + expected: dict[str, str | None], + actual: dict[str, str | None], +) -> bool: + expected = _normalize_launch_identity(expected) + actual = _normalize_launch_identity(actual) + return all(expected[key] == actual[key] for key in LAUNCH_IDENTITY_KEYS) + + +def _launch_identity_changed_keys( + expected: dict[str, str | None], + actual: dict[str, str | None], +) -> list[str]: + expected = _normalize_launch_identity(expected) + actual = _normalize_launch_identity(actual) + return [key for key in LAUNCH_IDENTITY_KEYS if expected[key] != actual[key]] + + +def _format_launch_identity(identity: dict[str, str | None]) -> str: + identity = _normalize_launch_identity(identity) + return ", ".join( + f"{key}={identity[key] if identity[key] is not None else 'none'}" + for key in LAUNCH_IDENTITY_KEYS + ) + + +def _read_output_launch_identity(output_dir: Path) -> dict[str, str | None] | None: + marker = output_dir / LAUNCH_IDENTITY_FILE + if not marker.exists(): + return None + try: + payload = json.loads(marker.read_text()) + except (OSError, json.JSONDecodeError): + return _normalize_launch_identity(None) + if isinstance(payload, dict) and isinstance(payload.get("launch_identity"), dict): + return _normalize_launch_identity(payload["launch_identity"]) + if isinstance(payload, dict): + return _normalize_launch_identity(payload) + return _normalize_launch_identity(None) + + +def _write_output_launch_identity(run: Run) -> None: + output_dir = Path(run.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + payload = { + "run_id": run.id, + "launch_identity": _run_launch_identity(run), + "written_at": datetime.now(timezone.utc).isoformat(), + } + fd, tmp = tempfile.mkstemp(dir=output_dir, prefix=f"{LAUNCH_IDENTITY_FILE}.", suffix=".tmp") + try: + with os.fdopen(fd, "w") as f: + json.dump(payload, f, indent=2) + os.replace(tmp, output_dir / LAUNCH_IDENTITY_FILE) + except BaseException: + try: + os.unlink(tmp) + except OSError: + pass + raise + + +def _can_resume_from_last_checkpoint(run: Run) -> bool: + expected = _run_launch_identity(run) + actual = _read_output_launch_identity(Path(run.output_dir)) + if actual is None: + return not _has_scoped_launch_identity(expected) + return _launch_identity_matches(expected, actual) + + +def get_run_status(state: dict, run_id: str) -> str: + return state["runs"].get(run_id, {}).get("status", "pending") + + +def set_run_status(state: dict, run_id: str, status: str, **kwargs) -> None: + state["runs"].setdefault(run_id, {}) + state["runs"][run_id]["status"] = status + state["runs"][run_id].update(kwargs) + + +def _completed_pretrain_missing_encoder(run: Run) -> bool: + """Return whether a supposedly completed pretrain lacks its encoder artifact.""" + return run.run_type == "pretrain" and not (Path(run.output_dir) / "encoder.pt").exists() + + +def _mark_completed_pretrain_missing_encoder( + run: Run, + runs: list[Run], + state: dict, + required_by: str | None = None, +) -> None: + encoder_path = Path(run.output_dir) / "encoder.pt" + reason = f"completed pretrain missing encoder.pt at {encoder_path}" + now = datetime.now(timezone.utc).isoformat() + suffix = f"; required by {required_by}" if required_by else "" + print(f" FAILED {run.id}: {reason}{suffix}") + set_run_status(state, run.id, "failed", finished_at=now, exit_code=None, reason=reason) + _propagate_failure(run.id, runs, state) + + +def _revalidate_completed_pretrain_artifacts(runs: list[Run], state: dict) -> None: + """Fail stale completed pretrains whose encoder artifact is no longer present.""" + for run in runs: + if get_run_status(state, run.id) == "completed" and _completed_pretrain_missing_encoder( + run + ): + _mark_completed_pretrain_missing_encoder(run, runs, state) + + +def _dependencies_ready( + run: Run, + runs: list[Run], + runs_by_id: dict[str, Run], + state: dict, +) -> bool: + """Check dependency state and revalidate completed pretrain artifacts.""" + for dep_id in run.depends_on: + if get_run_status(state, dep_id) != "completed": + return False + dep = runs_by_id.get(dep_id) + if dep is not None and _completed_pretrain_missing_encoder(dep): + _mark_completed_pretrain_missing_encoder(dep, runs, state, required_by=run.id) + return False + return True + + +def _run_launch_commit(run: Run) -> str | None: + commit = run.extra_overrides.get("+launch_commit") + return str(commit) if commit else None + + +def _state_launch_commit(info: dict) -> str | None: + commit = info.get("launch_commit") + if commit: + return str(commit) + command = str(info.get("command") or "") + match = re.search(r"(?:^|\s)\+?launch_commit=([^\s]+)", command) + return match.group(1) if match else None + + +def _safe_path_component(value: str | None) -> str: + text = str(value or "unknown") + return re.sub(r"[^A-Za-z0-9._-]+", "-", text)[:32] or "unknown" + + +def _quarantine_stale_output_dir( + run: Run, + actual_identity: str | None, + expected_identity: str, +) -> str | None: + output_dir = Path(run.output_dir) + if not output_dir.exists(): + return None + + parent = output_dir.parent + timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + actual = _safe_path_component(actual_identity) + expected = _safe_path_component(expected_identity) + stem = f"{output_dir.name}.stale-{actual}-to-{expected}-{timestamp}" + candidate = parent / stem + suffix = 1 + while candidate.exists(): + suffix += 1 + candidate = parent / f"{stem}-{suffix}" + + output_dir.rename(candidate) + return str(candidate) + + +def reset_state_for_launch_identity_mismatch(runs: list[Run], state: dict) -> None: + """Do not trust state or artifacts from a different launch target.""" + reset = 0 + quarantined = 0 + for run in runs: + expected_identity = _run_launch_identity(run) + info = state["runs"].get(run.id) + status = info.get("status", "pending") if info else "pending" + actual_identity = _state_launch_identity(info) + reasons = [] + + if info and not _launch_identity_matches(expected_identity, actual_identity): + changed_keys = _launch_identity_changed_keys(expected_identity, actual_identity) + if changed_keys == ["launch_commit"]: + reasons.append( + "launch_commit changed from " + f"{actual_identity['launch_commit'] or 'unknown'} " + f"to {expected_identity['launch_commit'] or 'none'}" + ) + else: + reasons.append( + "launch identity changed from " + f"{_format_launch_identity(actual_identity)} to " + f"{_format_launch_identity(expected_identity)}" + ) + + output_dir = Path(run.output_dir) + output_identity = _read_output_launch_identity(output_dir) if output_dir.exists() else None + if output_dir.exists(): + if output_identity is None: + if _has_scoped_launch_identity(expected_identity): + reasons.append("output directory has no launch identity marker") + elif not _launch_identity_matches(expected_identity, output_identity): + reasons.append( + "output directory launch identity is " + f"{_format_launch_identity(output_identity)}, expected " + f"{_format_launch_identity(expected_identity)}" + ) + + if not reasons: + continue + + if status == "running": + print( + f" WARNING {run.id}: state/output identity mismatch while run is marked " + "running; leaving it untouched." + ) + continue + + quarantined_output_dir = _quarantine_stale_output_dir( + run, + _format_launch_identity(output_identity or actual_identity), + _format_launch_identity(expected_identity), + ) + if quarantined_output_dir: + quarantined += 1 + + state["runs"][run.id] = { + "status": "pending", + "launch_identity": expected_identity, + "reset_reason": "; ".join(dict.fromkeys(reasons)), + } + if quarantined_output_dir: + state["runs"][run.id]["quarantined_output_dir"] = quarantined_output_dir + reset += 1 + + if reset: + print(f"Reset {reset} stale run state entries for launch identity mismatch.") + if quarantined: + print(f"Quarantined {quarantined} stale output directories before relaunch.") + + +def reset_state_for_launch_commit_mismatch(runs: list[Run], state: dict) -> None: + """Backward-compatible wrapper for launch identity reconciliation.""" + reset_state_for_launch_identity_mismatch(runs, state) + + +def is_pid_alive(pid: int) -> bool: + try: + os.kill(pid, 0) + return True + except (OSError, ProcessLookupError): + return False + + +def _pid_matches_command(pid: int, expected_command: str | None) -> bool: + if not expected_command: + return True + try: + result = subprocess.run( + ["ps", "-p", str(pid), "-o", "command="], + capture_output=True, + text=True, + check=False, + ) + except OSError: + return True + if result.returncode != 0: + return False + actual_command = result.stdout.strip() + return bool(actual_command) and ( + expected_command == actual_command or expected_command in actual_command + ) + + +def recover_stale_running(state: dict) -> int: + """Reset running entries whose PIDs are dead or no longer match the run.""" + recovered = 0 + for run_id, info in state["runs"].items(): + if info.get("status") == "running": + pid = info.get("pid") + if ( + pid is None + or not is_pid_alive(int(pid)) + or not _pid_matches_command(int(pid), info.get("command")) + ): + info["status"] = "pending" + info.pop("pid", None) + info.pop("command", None) + print(f" Recovered stale run: {run_id}") + recovered += 1 + return recovered + + +RUN_SLOT_COSTS = { + "pretrain": 4, + "finetune": 1, + "supervised": 1, + "gru_d": 1, + "xgboost": 1, +} + + +def _slot_cost(run: Run, slot_budget: int) -> int: + return min(RUN_SLOT_COSTS.get(run.run_type, 1), slot_budget) + + +def _select_ready_runs( + ready: list[Run], + active_run_ids: set[str], + runs_by_id: dict[str, Run], + slot_budget: int, +) -> list[Run]: + active_slots = sum(_slot_cost(runs_by_id[run_id], slot_budget) for run_id in active_run_ids) + available_slots = slot_budget - active_slots + selected: list[Run] = [] + + indexed_ready = list(enumerate(ready)) + indexed_ready.sort(key=lambda item: (-_slot_cost(item[1], slot_budget), item[0])) + for _, run in indexed_ready: + cost = _slot_cost(run, slot_budget) + if cost <= available_slots: + selected.append(run) + available_slots -= cost + return selected + + +# --------------------------------------------------------------------------- +# Scheduler +# --------------------------------------------------------------------------- + + +def run_scheduler(runs: list[Run], state: dict, parallel: int, dry_run: bool) -> int: + if parallel < 1: + raise ValueError(f"--parallel must be >= 1, got {parallel}") + + runs_by_id = {run.id: run for run in runs} + active: dict[str, subprocess.Popen] = {} + shutting_down = False + + def handle_signal(signum, frame): + nonlocal shutting_down + if shutting_down: + return + shutting_down = True + print(f"\nReceived signal {signum}, shutting down gracefully...") + for proc in active.values(): + try: + proc.terminate() + except OSError: + pass + deadline = time.time() + 30 + for proc in list(active.values()): + try: + proc.wait(timeout=max(0, deadline - time.time())) + except subprocess.TimeoutExpired: + proc.kill() + for run_id in active: + set_run_status(state, run_id, "pending") + save_state(state) + print("State saved. Interrupted runs reset to pending.") + sys.exit(1) + + signal.signal(signal.SIGINT, handle_signal) + signal.signal(signal.SIGTERM, handle_signal) + + LOG_DIR.mkdir(parents=True, exist_ok=True) + + if dry_run: + _print_dry_run(runs, runs_by_id) + return 0 + + reset_state_for_launch_identity_mismatch(runs, state) + _revalidate_completed_pretrain_artifacts(runs, state) + + print(f"Scheduler started (slot_budget={parallel})") + print(f"State file: {STATE_FILE}\n") + + while not shutting_down: + finished = [] + for run_id, proc in active.items(): + exit_code = proc.poll() + if exit_code is not None: + finished.append((run_id, exit_code)) + + for run_id, exit_code in finished: + proc = active.pop(run_id) + log_fh = getattr(proc, "_log_fh", None) + if log_fh: + log_fh.close() + run = runs_by_id[run_id] + now = datetime.now(timezone.utc).isoformat() + elapsed = _format_elapsed(state, run_id, now) + + if exit_code == 0: + if ( + run.run_type == "pretrain" + and not (Path(run.output_dir) / "encoder.pt").exists() + ): + print(f" FAILED {run_id}: exit 0 but encoder.pt missing {elapsed}") + set_run_status(state, run_id, "failed", finished_at=now, exit_code=exit_code) + _propagate_failure(run_id, runs, state) + continue + print(f" DONE {run_id} {elapsed}") + set_run_status(state, run_id, "completed", finished_at=now, exit_code=0) + else: + print(f" FAILED {run_id} (exit {exit_code}) {elapsed}") + set_run_status(state, run_id, "failed", finished_at=now, exit_code=exit_code) + _propagate_failure(run_id, runs, state) + + ready = [] + for run in runs: + if get_run_status(state, run.id) != "pending": + continue + if run.id in active: + continue + if _dependencies_ready(run, runs, runs_by_id, state): + ready.append(run) + + launch_batch = _select_ready_runs(ready, set(active), runs_by_id, parallel) + for run in launch_batch: + cmd = run.build_command(runs_by_id) + log_file = LOG_DIR / f"{run.id}.log" + now = datetime.now(timezone.utc).isoformat() + print(f" START {run.id}") + _write_output_launch_identity(run) + log_fh = open(log_file, "w") + proc = subprocess.Popen( + cmd, + stdout=log_fh, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + proc._log_fh = log_fh # type: ignore[attr-defined] + active[run.id] = proc + status_updates = { + "started_at": now, + "pid": proc.pid, + "log_file": str(log_file), + "command": shlex.join(cmd), + "launch_identity": _run_launch_identity(run), + } + for key, value in status_updates["launch_identity"].items(): + if value is not None: + status_updates[key] = value + set_run_status(state, run.id, "running", **status_updates) + + save_state(state) + + if not active and not ready: + break + time.sleep(2) + + for proc in active.values(): + log_fh = getattr(proc, "_log_fh", None) + if log_fh: + log_fh.close() + + print() + _print_summary(runs, state) + return _scheduler_exit_code(runs, state) + + +def _format_elapsed(state: dict, run_id: str, now_iso: str) -> str: + started = state["runs"].get(run_id, {}).get("started_at") + if not started: + return "" + try: + t0 = datetime.fromisoformat(started) + t1 = datetime.fromisoformat(now_iso) + secs = int((t1 - t0).total_seconds()) + if secs < 60: + return f"({secs}s)" + if secs < 3600: + return f"({secs // 60}m{secs % 60:02d}s)" + hours, rem = divmod(secs, 3600) + minutes, seconds = divmod(rem, 60) + return f"({hours}h{minutes:02d}m{seconds:02d}s)" + except (TypeError, ValueError): + return "" + + +def _propagate_failure(failed_id: str, runs: list[Run], state: dict) -> None: + queue = [failed_id] + while queue: + dep_id = queue.pop(0) + for run in runs: + if dep_id in run.depends_on and get_run_status(state, run.id) == "pending": + set_run_status(state, run.id, "skipped", reason=f"dependency {dep_id} failed") + print(f" SKIP {run.id} (dep {dep_id} failed)") + queue.append(run.id) + + +def _print_dry_run(runs: list[Run], runs_by_id: dict[str, Run]) -> None: + print(f"DRY RUN: {len(runs)} runs\n") + for run in runs: + deps = ", ".join(run.depends_on) if run.depends_on else "(none)" + cmd = shlex.join(run.build_command(runs_by_id)) + print(f"[{run.id}]") + print(f" experiment_class={run.experiment_class} type={run.run_type} deps={deps}") + print(f" dir={run.output_dir}") + print(f" cmd: {cmd}\n") + + +# --------------------------------------------------------------------------- +# Status / Retry +# --------------------------------------------------------------------------- + + +def _extract_revision_from_id(run_id: str) -> str | None: + match = re.search(r"_rev-([^_]+)_", run_id) + return match.group(1) if match else None + + +def _extract_experiment_class_from_id(run_id: str) -> str | None: + for experiment_class in sorted(EXPERIMENT_CLASSES, key=len, reverse=True): + if run_id.startswith(f"{experiment_class}_"): + return experiment_class + return None + + +def print_status( + experiment_class_filter: list[str] | None = None, + revision_filter: str | None = None, +) -> None: + all_runs = generate_all_runs() + if revision_filter: + all_runs = apply_revision(all_runs, revision_filter) + state = load_state() + if recover_stale_running(state): + save_state(state) + generated_ids = {run.id for run in all_runs} + + groups: dict[tuple[str, str | None], list[str]] = {} + for run in all_runs: + groups.setdefault((run.experiment_class, revision_filter), []).append(run.id) + + for run_id in state.get("runs", {}): + if run_id not in generated_ids and "_rev-" in run_id: + experiment_class = _extract_experiment_class_from_id(run_id) + revision = _extract_revision_from_id(run_id) + if revision_filter and revision != revision_filter: + continue + if experiment_class and revision: + groups.setdefault((experiment_class, revision), []).append(run_id) + + group_keys = sorted(groups, key=lambda key: (key[0], key[1] is not None, key[1] or "")) + if experiment_class_filter: + allowed = set(experiment_class_filter) + group_keys = [key for key in group_keys if key[0] in allowed] + + print( + f"{'Experiment Class':>28} | {'Total':>5} | {'Done':>4} | {'Run':>3} | " + f"{'Fail':>4} | {'Pend':>4} | {'Skip':>4}" + ) + print("-" * 76) + + totals = {"total": 0, "completed": 0, "running": 0, "failed": 0, "pending": 0, "skipped": 0} + for experiment_class, revision in group_keys: + run_ids = groups[(experiment_class, revision)] + label = f"{experiment_class}/{revision}" if revision else experiment_class + counts = { + "total": len(run_ids), + "completed": 0, + "running": 0, + "failed": 0, + "pending": 0, + "skipped": 0, + } + for run_id in run_ids: + status = get_run_status(state, run_id) + counts[status if status in counts else "pending"] += 1 + print( + f"{label:>28} | {counts['total']:>5} | {counts['completed']:>4} | " + f"{counts['running']:>3} | {counts['failed']:>4} | " + f"{counts['pending']:>4} | {counts['skipped']:>4}" + ) + for key in totals: + totals[key] += counts[key] + + print("-" * 76) + print( + f"{'TOTAL':>28} | {totals['total']:>5} | {totals['completed']:>4} | " + f"{totals['running']:>3} | {totals['failed']:>4} | " + f"{totals['pending']:>4} | {totals['skipped']:>4}" + ) + + +def _print_summary(runs: list[Run], state: dict) -> None: + counts = {"completed": 0, "failed": 0, "skipped": 0, "pending": 0} + for run in runs: + status = get_run_status(state, run.id) + counts[status] = counts.get(status, 0) + 1 + total = len(runs) + print( + f"Summary: {counts['completed']}/{total} completed, " + f"{counts['failed']} failed, {counts['skipped']} skipped, {counts['pending']} pending" + ) + + +def _scheduler_exit_code(runs: list[Run], state: dict) -> int: + failed = [run.id for run in runs if get_run_status(state, run.id) == "failed"] + skipped = [run.id for run in runs if get_run_status(state, run.id) == "skipped"] + running = [run.id for run in runs if get_run_status(state, run.id) == "running"] + pending = [run.id for run in runs if get_run_status(state, run.id) == "pending"] + if failed or skipped or running or pending: + print( + "Scheduler incomplete: " + f"{len(failed)} failed, {len(skipped)} skipped due to dependency failure, " + f"{len(running)} running, {len(pending)} pending" + ) + return 1 + return 0 + + +def _collect_dependency_closure( + runs: list[Run], + all_by_id: dict[str, Run], +) -> tuple[list[Run], dict[str, set[str]], set[str]]: + deps: list[Run] = [] + required_by: dict[str, set[str]] = {} + missing: set[str] = set() + seen: set[str] = set() + queue = [(dep_id, run.id) for run in runs for dep_id in run.depends_on] + while queue: + dep_id, parent_id = queue.pop(0) + required_by.setdefault(dep_id, set()).add(parent_id) + if dep_id in seen: + continue + seen.add(dep_id) + dep = all_by_id.get(dep_id) + if dep is None: + missing.add(dep_id) + continue + deps.append(dep) + queue.extend((child_dep_id, dep.id) for child_dep_id in dep.depends_on) + return deps, required_by, missing + + +def _collect_skipped_dependent_closure( + runs: list[Run], + all_runs: list[Run], + state: dict, +) -> list[Run]: + """Find skipped downstream runs that will remain skipped after retry --failed.""" + downstream: dict[str, list[Run]] = {} + for run in all_runs: + for dep_id in run.depends_on: + downstream.setdefault(dep_id, []).append(run) + + start_ids = {run.id for run in runs} + queue = list(start_ids) + seen: set[str] = set() + skipped_dependents: list[Run] = [] + while queue: + dep_id = queue.pop(0) + for run in downstream.get(dep_id, []): + if run.id in seen or run.id in start_ids: + continue + seen.add(run.id) + if get_run_status(state, run.id) == "skipped": + skipped_dependents.append(run) + queue.append(run.id) + return skipped_dependents + + +def _retry_command_suggestion(args, *, include_failed: bool, include_skipped: bool) -> str: + cmd = ["uv", "run", "python", "scripts/internal/run_experiments.py", "retry"] + if include_failed: + cmd.append("--failed") + if include_skipped: + cmd.append("--skipped") + if args.experiment_class: + cmd.append("--experiment-class") + cmd.extend(args.experiment_class) + if args.revision: + cmd.extend(["--revision", str(args.revision)]) + if args.reason: + cmd.extend(["--reason", str(args.reason)]) + if args.project: + cmd.extend(["--project", str(args.project)]) + if args.entity: + cmd.extend(["--entity", str(args.entity)]) + launch_commit = getattr(args, "launch_commit", None) + if launch_commit: + cmd.extend(["--launch-commit", str(launch_commit)]) + cmd.extend(["--parallel", str(args.parallel)]) + return shlex.join(cmd) + + +def _fail_for_blocked_retry_dependencies( + blocked: list[Run], + missing: set[str], + state: dict, + required_by: dict[str, set[str]], + args, +) -> None: + print("Cannot retry the selected runs because required dependencies are blocked.") + if blocked: + print("\nBlocked dependencies:") + for dep in sorted(blocked, key=lambda run: (get_run_status(state, run.id), run.id)): + dependents = ", ".join(sorted(required_by.get(dep.id, set()))) + print(f" {dep.id} [{get_run_status(state, dep.id)}] required by: {dependents}") + if missing: + print("\nMissing generated dependencies:") + for dep_id in sorted(missing): + dependents = ", ".join(sorted(required_by.get(dep_id, set()))) + print(f" {dep_id} required by: {dependents}") + suggested = _retry_command_suggestion( + args, + include_failed=args.failed + or any(get_run_status(state, dep.id) == "failed" for dep in blocked), + include_skipped=args.skipped + or any(get_run_status(state, dep.id) == "skipped" for dep in blocked), + ) + print("\nSuggested command:") + print(f" {suggested}") + raise SystemExit(1) + + +def _warn_for_skipped_dependents_after_failed_retry( + runs_to_retry: list[Run], + all_runs: list[Run], + state: dict, + args, +) -> list[Run]: + if not args.failed or args.skipped: + return [] + + skipped_dependents = _collect_skipped_dependent_closure(runs_to_retry, all_runs, state) + if not skipped_dependents: + return [] + + print( + "WARNING: retry --failed will leave skipped downstream runs selected for " + "accounting only. Prefer retry --failed --skipped for final recovery." + ) + print("\nSkipped dependents:") + for run in skipped_dependents[:20]: + print(f" {run.id} [skipped]") + if len(skipped_dependents) > 20: + print(f" ... {len(skipped_dependents) - 20} more") + suggested = _retry_command_suggestion(args, include_failed=True, include_skipped=True) + print("\nSuggested command:") + print(f" {suggested}") + return skipped_dependents + + +def _expand_retry_with_skipped_dependents( + runs_to_retry: list[Run], + all_runs: list[Run], + state: dict, +) -> list[Run]: + """Include downstream skipped runs that were blocked by the selected retry roots.""" + selected_ids = {run.id for run in runs_to_retry} + return [ + run + for run in _collect_skipped_dependent_closure(runs_to_retry, all_runs, state) + if run.id not in selected_ids + ] + + +def _filter_requested_runs( + all_runs: list[Run], + experiment_classes: list[str], +) -> list[Run]: + requested = list(dict.fromkeys(experiment_classes)) + runs = [ + run + for experiment_class in requested + for run in all_runs + if run.experiment_class == experiment_class + ] + run_ids = {run.id for run in runs} + all_by_id = {run.id: run for run in all_runs} + deps_needed = { + dep_id + for run in runs + for dep_id in run.depends_on + if dep_id not in run_ids and dep_id in all_by_id + } + deps = [run for run in all_runs if run.id in deps_needed] + return deps + runs + + +def cmd_run(args): + all_runs = generate_all_runs() + state = load_state() + recover_stale_running(state) + + runs = _filter_requested_runs(all_runs, args.experiment_class) + if not runs: + print(f"No runs found for experiment class(es): {', '.join(args.experiment_class)}") + return 0 + + if args.revision: + runs = apply_revision(runs, args.revision, args.reason) + runs = apply_wandb_target(runs, args.project, args.entity) + launch_commit = getattr(args, "launch_commit", None) + runs = apply_launch_commit(runs, launch_commit) + + print(f"Experiment class(es) {', '.join(args.experiment_class)}: {len(runs)} runs") + if args.revision: + print(f"Revision: {args.revision}" + (f" ({args.reason})" if args.reason else "")) + if launch_commit: + print(f"Launch commit: {launch_commit}") + if args.project or args.entity: + target = f"{args.entity}/{args.project}" if args.entity else args.project + print(f"W&B target: {target}") + return run_scheduler(runs, state, args.parallel, args.dry_run) + + +def cmd_status(args): + print_status(args.experiment_class, args.revision) + + +def cmd_retry(args): + all_runs = generate_all_runs() + class_filter = set(args.experiment_class) if args.experiment_class else None + + if args.revision: + all_runs = apply_revision(all_runs, args.revision, args.reason) + + all_runs = apply_wandb_target(all_runs, args.project, args.entity) + launch_commit = getattr(args, "launch_commit", None) + all_runs = apply_launch_commit(all_runs, launch_commit) + state = load_state() + recover_stale_running(state) + + candidate_runs = [ + run for run in all_runs if class_filter is None or run.experiment_class in class_filter + ] + runs_to_retry = [] + for run in candidate_runs: + status = get_run_status(state, run.id) + if args.failed and status == "failed": + runs_to_retry.append(run) + elif args.skipped and status == "skipped": + runs_to_retry.append(run) + + if not runs_to_retry: + print("No runs to retry.") + return 0 + + if args.skipped: + skipped_dependents = _expand_retry_with_skipped_dependents( + runs_to_retry, + all_runs, + state, + ) + if skipped_dependents: + print("Including skipped downstream dependents required by the selected retry roots:") + for run in skipped_dependents[:20]: + print(f" {run.id} [skipped]") + if len(skipped_dependents) > 20: + print(f" ... {len(skipped_dependents) - 20} more") + runs_to_retry.extend(skipped_dependents) + + all_by_id = {run.id: run for run in all_runs} + dependencies, required_by, missing = _collect_dependency_closure(runs_to_retry, all_by_id) + blocked = [ + dep + for dep in dependencies + if (get_run_status(state, dep.id) == "failed" and not args.failed) + or (get_run_status(state, dep.id) == "skipped" and not args.skipped) + ] + unresolved_missing = { + dep_id for dep_id in missing if get_run_status(state, dep_id) != "completed" + } + if blocked or unresolved_missing: + _fail_for_blocked_retry_dependencies(blocked, unresolved_missing, state, required_by, args) + + accounting_only = _warn_for_skipped_dependents_after_failed_retry( + runs_to_retry, + all_runs, + state, + args, + ) + + runs_by_id: dict[str, Run] = {} + for run in dependencies + runs_to_retry + accounting_only: + runs_by_id.setdefault(run.id, run) + for run in dependencies + runs_to_retry: + if get_run_status(state, run.id) in {"failed", "skipped"}: + set_run_status(state, run.id, "pending") + + save_state(state) + print(f"Retrying {len(runs_by_id)} runs") + if args.project or args.entity: + target = f"{args.entity}/{args.project}" if args.entity else args.project + print(f"W&B target: {target}") + if launch_commit: + print(f"Launch commit: {launch_commit}") + return run_scheduler(list(runs_by_id.values()), state, args.parallel, dry_run=False) + + +def cmd_warmup(args): + if args.revision: + print(f"Note: --revision={args.revision} ignored (warmup is revision-independent)") + + runs = _filter_requested_runs(generate_all_runs(), args.experiment_class) + datasets = dict.fromkeys(run.dataset for run in runs) + + print(f"Warming up tensor caches for experiment class(es) {', '.join(args.experiment_class)}") + print(f" {len(datasets)} unique datasets to cache\n") + + failed = [] + for index, dataset in enumerate(datasets, 1): + processed_dir = f"data/processed/{dataset}" + print(f"[{index}/{len(datasets)}] {dataset}") + result = subprocess.run( + [ + "uv", + "run", + "python", + "-c", + "from slices.data.dataset import ICUDataset; " + f"ds = ICUDataset(data_dir={processed_dir!r}, task_name=None, normalize=False); " + "print(len(ds))", + ], + capture_output=True, + text=True, + ) + if result.returncode == 0: + n_samples = result.stdout.strip().split("\n")[-1] + print(f" -> Cached ({int(n_samples):,} samples)") + else: + stderr_tail = result.stderr.strip().split("\n")[-3:] + print(f" -> ERROR: {' '.join(stderr_tail)}") + failed.append(dataset) + + if failed: + print(f"\nWarmup failed for dataset(s): {', '.join(failed)}") + return 1 + print("\nWarmup complete. Raw tensor caches saved to data/processed//.tensor_cache/") + return 0 + + +def main() -> None: + parser = argparse.ArgumentParser(description="SLICES class-based experiment runner") + sub = parser.add_subparsers(dest="command", required=True) + + p_run = sub.add_parser("run", help="Run experiments for given experiment classes") + p_run.add_argument("--experiment-class", nargs="+", required=True, choices=EXPERIMENT_CLASSES) + p_run.add_argument( + "--parallel", + type=int, + default=4, + help="Max scheduler slots (default: 4). Pretrains consume 4 slots, other runs 1.", + ) + p_run.add_argument("--dry-run", action="store_true", help="Print runs without executing") + p_run.add_argument("--revision", type=str, default=None, help="Revision name") + p_run.add_argument("--reason", type=str, default=None, help="Reason for rerun") + p_run.add_argument( + "--launch-commit", + type=str, + default=os.environ.get("SLICES_LAUNCH_COMMIT"), + help="Exact git commit hash for launch provenance (default: SLICES_LAUNCH_COMMIT)", + ) + p_run.add_argument( + "--project", + type=str, + default=os.environ.get("WANDB_PROJECT"), + help="W&B project override for launched runs (default: WANDB_PROJECT env var)", + ) + p_run.add_argument( + "--entity", + type=str, + default=os.environ.get("WANDB_ENTITY"), + help="W&B entity override for launched runs (default: WANDB_ENTITY env var)", + ) + + p_status = sub.add_parser("status", help="Show experiment status") + p_status.add_argument("--experiment-class", nargs="*", choices=EXPERIMENT_CLASSES, default=None) + p_status.add_argument("--revision", type=str, default=None, help="Show status for one revision") + + p_retry = sub.add_parser("retry", help="Retry failed/skipped runs") + p_retry.add_argument("--failed", action="store_true", help="Retry failed runs") + p_retry.add_argument( + "--skipped", + action="store_true", + help="Retry dependency-skipped runs; prefer with --failed for final recovery", + ) + p_retry.add_argument( + "--parallel", + type=int, + default=4, + help="Max scheduler slots (default: 4). Pretrains consume 4 slots, other runs 1.", + ) + p_retry.add_argument("--experiment-class", nargs="+", choices=EXPERIMENT_CLASSES, default=None) + p_retry.add_argument("--revision", type=str, default=None, help="Revision name to retry") + p_retry.add_argument("--reason", type=str, default=None, help="Reason for rerun") + p_retry.add_argument( + "--launch-commit", + type=str, + default=os.environ.get("SLICES_LAUNCH_COMMIT"), + help="Exact git commit hash for launch provenance (default: SLICES_LAUNCH_COMMIT)", + ) + p_retry.add_argument( + "--project", + type=str, + default=os.environ.get("WANDB_PROJECT"), + help="W&B project override for relaunched runs (default: WANDB_PROJECT env var)", + ) + p_retry.add_argument( + "--entity", + type=str, + default=os.environ.get("WANDB_ENTITY"), + help="W&B entity override for relaunched runs (default: WANDB_ENTITY env var)", + ) + + p_warmup = sub.add_parser( + "warmup", help="Pre-build tensor caches to avoid OOM during parallel runs" + ) + p_warmup.add_argument( + "--experiment-class", nargs="+", required=True, choices=EXPERIMENT_CLASSES + ) + p_warmup.add_argument( + "--revision", type=str, default=None, help="Ignored (warmup is revision-independent)" + ) + p_warmup.add_argument("--reason", type=str, default=None, help="Ignored for warmup") + + args = parser.parse_args() + + if getattr(args, "reason", None) and not getattr(args, "revision", None): + parser.error("--reason requires --revision") + if args.command == "run": + if not args.revision: + parser.error("run requires --revision to tag run IDs, output dirs, and W&B metadata") + if not args.project: + parser.error( + "run requires --project or WANDB_PROJECT to avoid logging to config defaults" + ) + if not args.entity: + parser.error("run requires --entity or WANDB_ENTITY to make W&B ownership explicit") + if args.command == "retry": + if not args.revision: + parser.error("retry requires --revision to select the revisioned state namespace") + if not args.experiment_class: + parser.error("retry requires --experiment-class to scope which classes to revise") + if not args.project: + parser.error( + "retry requires --project or WANDB_PROJECT to avoid logging to config defaults" + ) + if not args.entity: + parser.error("retry requires --entity or WANDB_ENTITY to make W&B ownership explicit") + if not args.failed and not args.skipped: + parser.error("retry requires --failed and/or --skipped") + + validate_direct_final_launch_policy(args, parser) + + if args.command == "run": + exit_code = cmd_run(args) + elif args.command == "status": + exit_code = cmd_status(args) + elif args.command == "retry": + exit_code = cmd_retry(args) + elif args.command == "warmup": + exit_code = cmd_warmup(args) + else: + exit_code = 0 + sys.exit(exit_code or 0) + + +if __name__ == "__main__": + main() diff --git a/scripts/measure_flops.py b/scripts/measure_flops.py new file mode 100644 index 0000000..b0c77cf --- /dev/null +++ b/scripts/measure_flops.py @@ -0,0 +1,335 @@ +"""Offline FLOPs measurement for each SSL paradigm. + +Instantiates each SSL objective with the same configs used in training runs, +profiles a single forward pass using PyTorch's FlopCounterMode, and estimates +total training FLOPs by multiplying per-step FLOPs by gradient steps from W&B. + +Usage: + uv run python scripts/measure_flops.py + uv run python scripts/measure_flops.py --no-wandb # skip W&B query + uv run python scripts/measure_flops.py --n-features 84 # override feature dim + uv run python scripts/measure_flops.py --sparsity 0.7 # fraction missing +""" + +import argparse +import os +from pathlib import Path +from typing import Dict, Tuple + +import torch +import torch.nn as nn +from omegaconf import OmegaConf +from torch.utils.flop_counter import FlopCounterMode + +from slices.constants import SEQ_LENGTH_HOURS +from slices.models.encoders.factory import build_encoder +from slices.models.pretraining.factory import build_ssl_objective, get_ssl_config_class + +# --------------------------------------------------------------------------- +# Paradigm configs — loaded from the YAML defaults used in actual runs +# --------------------------------------------------------------------------- + +DEFAULT_N_FEATURES = 84 +CONFIG_DIR = Path(__file__).resolve().parents[1] / "configs" +PARADIGMS = ["mae", "jepa", "contrastive", "ts2vec"] + + +def load_yaml_config(path: Path) -> dict: + """Load a Hydra YAML file as a plain dict without the config-group name.""" + cfg = OmegaConf.to_container(OmegaConf.load(path), resolve=True) + if not isinstance(cfg, dict): + raise TypeError(f"Expected mapping in {path}") + cfg.pop("name", None) + return cfg + + +ENCODER_CONFIG = { + **load_yaml_config(CONFIG_DIR / "model" / "transformer.yaml"), + "d_input": DEFAULT_N_FEATURES, + "max_seq_length": SEQ_LENGTH_HOURS, + "pooling": "none", # Required for SSL pretraining + "obs_aware": True, # All controlled SSL paradigms use obs-aware tokenization +} + +SSL_CONFIGS: Dict[str, dict] = { + name: load_yaml_config(CONFIG_DIR / "ssl" / f"{name}.yaml") for name in PARADIGMS +} + + +def count_parameters(model: nn.Module) -> Tuple[int, int]: + """Return (trainable_params, total_params).""" + trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) + total = sum(p.numel() for p in model.parameters()) + return trainable, total + + +def measure_forward_flops( + model: nn.Module, + x: torch.Tensor, + obs_mask: torch.Tensor, +) -> int: + """Measure forward-pass FLOPs for an SSL objective.""" + model.eval() + with torch.no_grad(): + flop_counter = FlopCounterMode(display=False) + with flop_counter: + model(x, obs_mask) + return flop_counter.get_total_flops() + + +def build_objective(name: str, encoder_config: dict) -> nn.Module: + """Build an SSL objective with a fresh encoder.""" + encoder = build_encoder("transformer", encoder_config) + ssl_config_cls = get_ssl_config_class(name) + ssl_config = ssl_config_cls(**SSL_CONFIGS[name]) + return build_ssl_objective(encoder, ssl_config) + + +def format_flops(flops: int) -> str: + """Format FLOPs as human-readable string.""" + if flops >= 1e12: + return f"{flops / 1e12:.2f} TFLOPs" + if flops >= 1e9: + return f"{flops / 1e9:.2f} GFLOPs" + if flops >= 1e6: + return f"{flops / 1e6:.2f} MFLOPs" + return f"{flops:,} FLOPs" + + +def format_params(n: int) -> str: + """Format parameter count as human-readable string.""" + if n >= 1e6: + return f"{n / 1e6:.2f}M" + if n >= 1e3: + return f"{n / 1e3:.1f}K" + return str(n) + + +def query_wandb_runs(entity: str, project: str, revision: str | None = None) -> Dict[str, dict]: + """Query W&B for pretrain run gradient steps and wall-clock time. + + Returns dict keyed by SSL name -> {grad_steps, wall_clock_s, n_runs}. + """ + import wandb + + api = wandb.Api(timeout=300) + path = f"{entity}/{project}" if entity else project + + # Server-side filter: only finished pretrain runs from the selected corpus. + required_tags = ["phase:pretrain"] + if revision: + required_tags.append(f"revision:{revision}") + filters = {"tags": {"$all": required_tags}, "state": "finished"} + runs = api.runs(path, filters=filters, order="-created_at") + + results: Dict[str, list] = {} + for run in runs: + config = dict(run.config) + summary = dict(run.summary_metrics or {}) + + ssl_name = config.get("ssl", {}).get("name") or config.get("ssl_name") + if not ssl_name: + continue + + grad_steps = summary.get("train/gradient_steps") or summary.get("trainer/global_step") + wall_clock = summary.get("train/wall_clock_seconds") + + if grad_steps is None: + continue + + if ssl_name not in results: + results[ssl_name] = [] + results[ssl_name].append( + { + "grad_steps": int(grad_steps), + "wall_clock_s": float(wall_clock) if wall_clock is not None else None, + "run_name": run.name, + } + ) + + # Aggregate per paradigm + aggregated = {} + for name, run_list in results.items(): + steps = [r["grad_steps"] for r in run_list] + clocks = [r["wall_clock_s"] for r in run_list if r["wall_clock_s"] is not None] + aggregated[name] = { + "avg_grad_steps": sum(steps) / len(steps), + "avg_wall_clock_s": sum(clocks) / len(clocks) if clocks else None, + "n_runs": len(run_list), + } + + return aggregated + + +def main(): + parser = argparse.ArgumentParser(description="Measure FLOPs per SSL paradigm") + parser.add_argument( + "--n-features", + type=int, + default=DEFAULT_N_FEATURES, + help=f"Number of input features (d_input; default: {DEFAULT_N_FEATURES})", + ) + parser.add_argument( + "--seq-length", type=int, default=SEQ_LENGTH_HOURS, help="Sequence length (T)" + ) + parser.add_argument("--batch-size", type=int, default=1, help="Batch size for profiling") + parser.add_argument( + "--sparsity", type=float, default=0.7, help="Fraction of missing values (0-1)" + ) + parser.add_argument("--no-wandb", action="store_true", help="Skip W&B query") + parser.add_argument( + "--wandb-entity", type=str, default=os.environ.get("WANDB_ENTITY", ""), help="W&B entity" + ) + parser.add_argument( + "--wandb-project", + type=str, + default=os.environ.get("WANDB_PROJECT", "slices-thesis"), + help="W&B project", + ) + parser.add_argument( + "--revision", + type=str, + default=os.environ.get("SLICES_REVISION", "thesis-v1"), + help="Revision tag to query in W&B (use empty string to disable)", + ) + args = parser.parse_args() + + # Build encoder config with CLI overrides + encoder_config = { + **ENCODER_CONFIG, + "d_input": args.n_features, + "max_seq_length": args.seq_length, + } + + # Create synthetic input with realistic sparsity + B, T, D = args.batch_size, args.seq_length, args.n_features + torch.manual_seed(42) + x = torch.randn(B, T, D) + obs_mask = (torch.rand(B, T, D) > args.sparsity).bool() + # Ensure at least one observation per timestep to avoid degenerate cases + for b in range(B): + for t in range(T): + if not obs_mask[b, t].any(): + obs_mask[b, t, 0] = True + + print(f"Input shape: ({B}, {T}, {D}), sparsity: {args.sparsity:.0%}") + print(f"Avg observations per sample: {obs_mask.float().sum() / B:.0f} / {T * D}") + print() + + # Query W&B if requested + wandb_data = {} + if not args.no_wandb: + try: + revision = args.revision or None + target = ( + f"{args.wandb_entity}/{args.wandb_project}" + if args.wandb_entity + else args.wandb_project + ) + print(f"Querying W&B target: {target}, revision: {revision or 'unfiltered'}") + wandb_data = query_wandb_runs(args.wandb_entity, args.wandb_project, revision) + print(f"Found W&B data for: {list(wandb_data.keys())}") + print() + except Exception as e: + print(f"W&B query failed ({e}), continuing without it.\n") + + # Measure each paradigm + rows = [] + for name in PARADIGMS: + objective = build_objective(name, encoder_config) + trainable, total = count_parameters(objective) + + fwd_flops = measure_forward_flops(objective, x.clone(), obs_mask.clone()) + + # Backward ≈ 2× forward; total step = fwd + bwd = 3× fwd + # Exception: JEPA target encoder is forward-only (no gradient) + if name == "jepa": + # The FlopCounterMode already captures both online + target encoder forward. + # But backward only applies to online encoder + predictor (not target). + # Approximate: target encoder ≈ online encoder forward FLOPs. + # So: total_step ≈ fwd_flops (online+target) + 2 × (fwd_flops - target_fwd) + # Since target ≈ online encoder, target_fwd ≈ fwd_flops * (encoder_share). + # Simpler: fwd already includes target. bwd ≈ 2 × fwd_without_target. + # We estimate target encoder as ~40% of total forward (encoder is the bulk). + # Conservative: just use 3× as upper bound like others. + step_flops = fwd_flops * 3 + else: + step_flops = fwd_flops * 3 + + wb = wandb_data.get(name, {}) + avg_steps = wb.get("avg_grad_steps") + avg_clock = wb.get("avg_wall_clock_s") + n_runs = wb.get("n_runs", 0) + + total_flops = step_flops * avg_steps if avg_steps else None + + rows.append( + { + "name": name, + "trainable": trainable, + "total_params": total, + "fwd_flops": fwd_flops, + "step_flops": step_flops, + "avg_steps": avg_steps, + "total_flops": total_flops, + "avg_clock": avg_clock, + "n_runs": n_runs, + } + ) + + # Print results table + cols = [ + f"{'Paradigm':<14}", + f"{'Params':>10}", + f"{'FLOPs/step (fwd)':>20}", + f"{'FLOPs/step (fwd+bwd)':>22}", + f"{'Grad steps':>12}", + f"{'Total FLOPs':>14}", + f"{'Wall-clock':>12}", + f"{'Runs':>5}", + ] + header = " ".join(cols) + print(header) + print("-" * len(header)) + + for r in rows: + steps_str = f"{r['avg_steps']:.0f}" if r["avg_steps"] else "N/A" + total_str = format_flops(r["total_flops"]) if r["total_flops"] else "N/A" + clock_str = f"{r['avg_clock']:.0f}s" if r["avg_clock"] else "N/A" + runs_str = str(r["n_runs"]) if r["n_runs"] else "-" + + print( + f"{r['name']:<14} " + f"{format_params(r['trainable']):>10} " + f"{format_flops(r['fwd_flops']):>20} " + f"{format_flops(r['step_flops']):>22} " + f"{steps_str:>12} " + f"{total_str:>14} " + f"{clock_str:>12} " + f"{runs_str:>5}" + ) + + # Also print raw numbers for programmatic use + print("\n--- Raw values (for tables/plots) ---") + raw_cols = [ + f"{'Paradigm':<14}", + f"{'Trainable':>12}", + f"{'FWD FLOPs':>14}", + f"{'Step FLOPs':>14}", + f"{'Avg Steps':>12}", + f"{'Total FLOPs':>18}", + ] + print(" ".join(raw_cols)) + for r in rows: + print( + f"{r['name']:<14} " + f"{r['trainable']:>12,} " + f"{r['fwd_flops']:>14,} " + f"{r['step_flops']:>14,} " + f"{r['avg_steps'] or 0:>12.0f} " + f"{r['total_flops'] or 0:>18,.0f}" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/preprocessing/build_processed_data.py b/scripts/preprocessing/build_processed_data.py new file mode 100644 index 0000000..fd8676e --- /dev/null +++ b/scripts/preprocessing/build_processed_data.py @@ -0,0 +1,137 @@ +"""Build processed SLICES datasets from existing RICU parquet exports. + +This is a convenience orchestrator for the common local rebuild path: + + uv run python scripts/preprocessing/build_processed_data.py --datasets miiv eicu --combined + +It assumes ``data/ricu_output/{dataset}`` already exists. Use +``scripts/setup_and_extract.sh`` when you also need dependency installation or +the upstream RICU R export step. +""" + +from __future__ import annotations + +import argparse +import shlex +import subprocess +import sys +from pathlib import Path +from typing import Sequence + +SUPPORTED_DATASETS = ("miiv", "eicu") + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + """Parse CLI arguments.""" + parser = argparse.ArgumentParser( + description="Run Python extraction, preparation, and optional combined-dataset build.", + ) + parser.add_argument( + "--datasets", + nargs="+", + choices=SUPPORTED_DATASETS, + default=list(SUPPORTED_DATASETS), + help="Base datasets to process. Defaults to miiv eicu.", + ) + parser.add_argument( + "--combined", + action="store_true", + help="Also build data/processed/combined from miiv and eicu.", + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Seed used for split generation. Defaults to 42.", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print commands without executing them.", + ) + return parser.parse_args(argv) + + +def _append_unique(values: list[str], value: str) -> None: + if value not in values: + values.append(value) + + +def build_commands( + datasets: Sequence[str], + *, + build_combined: bool, + seed: int, + python_executable: str = sys.executable, +) -> list[list[str]]: + """Build the subprocess command list for the requested processing plan.""" + ordered_datasets: list[str] = [] + for dataset in datasets: + _append_unique(ordered_datasets, dataset) + + if build_combined: + for dataset in SUPPORTED_DATASETS: + _append_unique(ordered_datasets, dataset) + + commands: list[list[str]] = [] + for dataset in ordered_datasets: + commands.append( + [ + python_executable, + "scripts/preprocessing/extract_ricu.py", + f"dataset={dataset}", + ] + ) + commands.append( + [ + python_executable, + "scripts/preprocessing/prepare_dataset.py", + f"dataset={dataset}", + f"seed={seed}", + ] + ) + + if build_combined: + commands.append( + [ + python_executable, + "scripts/preprocessing/create_combined_dataset.py", + "--source", + "data/processed/miiv", + "data/processed/eicu", + "--names", + "miiv", + "eicu", + "--output", + "data/processed/combined", + "--seed", + str(seed), + ] + ) + + return commands + + +def run_commands(commands: Sequence[Sequence[str]], *, cwd: Path, dry_run: bool) -> None: + """Print and optionally execute commands.""" + for command in commands: + print(f"$ {shlex.join(command)}", flush=True) + if not dry_run: + subprocess.run(command, cwd=cwd, check=True) + + +def main(argv: Sequence[str] | None = None) -> int: + """CLI entry point.""" + args = parse_args(argv) + project_root = Path(__file__).resolve().parents[2] + commands = build_commands( + args.datasets, + build_combined=args.combined, + seed=args.seed, + ) + run_commands(commands, cwd=project_root, dry_run=args.dry_run) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/preprocessing/create_combined_dataset.py b/scripts/preprocessing/create_combined_dataset.py index 4eb49c9..db910a4 100644 --- a/scripts/preprocessing/create_combined_dataset.py +++ b/scripts/preprocessing/create_combined_dataset.py @@ -2,7 +2,8 @@ Merges MIMIC-IV and eICU (or any two datasets) into a single processed directory for combined pretraining. Handles stay_id collision by adding -a dataset-specific offset to ensure globally unique IDs. +a dataset-specific offset to ensure globally unique IDs, then prepares +the merged dataset for training by default. Usage: uv run python scripts/preprocessing/create_combined_dataset.py \ @@ -23,8 +24,11 @@ import polars as pl import yaml -# Offset applied to stay_id and patient_id for the second dataset -# to avoid collisions. Large enough to exceed any real ID range. +from slices.constants import MIN_STAY_HOURS, SEQ_LENGTH_HOURS, THESIS_TASKS +from slices.data.preparation import prepare_processed_dataset + +# Offset applied to stay_id for the second dataset when needed to avoid +# collisions. Large enough to exceed any real ID range. DATASET_ID_OFFSET = 100_000_000 @@ -50,19 +54,31 @@ def load_dataset(processed_dir: Path) -> dict: def offset_ids(df: pl.DataFrame, offset: int) -> pl.DataFrame: - """Add offset to stay_id and patient_id columns.""" + """Add an offset to the stay_id column when collision resolution is needed.""" exprs = [] if "stay_id" in df.columns: exprs.append((pl.col("stay_id") + offset).alias("stay_id")) - if "patient_id" in df.columns: - exprs.append((pl.col("patient_id") + offset).alias("patient_id")) if exprs: return df.with_columns(exprs) return df +def namespace_patient_ids(df: pl.DataFrame, dataset_name: str) -> pl.DataFrame: + """Namespace patient_id by source dataset to preserve combined split integrity.""" + if "patient_id" not in df.columns: + return df + + patient_expr = ( + pl.when(pl.col("patient_id").is_null()) + .then(None) + .otherwise(pl.format("{}:{}", pl.lit(dataset_name), pl.col("patient_id").cast(pl.Utf8))) + .alias("patient_id") + ) + return df.with_columns(patient_expr) + + def validate_no_id_collision(static_a: pl.DataFrame, static_b: pl.DataFrame) -> None: - """Verify no stay_id overlap between two datasets.""" + """Verify no stay_id or patient_id overlap between two datasets.""" ids_a = set(static_a["stay_id"].to_list()) ids_b = set(static_b["stay_id"].to_list()) overlap = ids_a & ids_b @@ -73,15 +89,34 @@ def validate_no_id_collision(static_a: pl.DataFrame, static_b: pl.DataFrame) -> "applying the dataset offset." ) + if "patient_id" in static_a.columns and "patient_id" in static_b.columns: + patients_a = {str(value) for value in static_a["patient_id"].drop_nulls().to_list()} + patients_b = {str(value) for value in static_b["patient_id"].drop_nulls().to_list()} + patient_overlap = patients_a & patients_b + if patient_overlap: + raise ValueError( + f"patient_id collision detected after namespacing: {len(patient_overlap)} " + f"overlapping IDs (e.g., {list(patient_overlap)[:5]})." + ) + def validate_feature_compatibility(meta_a: dict, meta_b: dict) -> None: - """Verify both datasets have the same feature set.""" - features_a = set(meta_a.get("feature_names", [])) - features_b = set(meta_b.get("feature_names", [])) + """Verify both datasets share the same preprocessing contract.""" + features_a = list(meta_a.get("feature_names", [])) + features_b = list(meta_b.get("feature_names", [])) if features_a != features_b: - only_a = features_a - features_b - only_b = features_b - features_a + set_a = set(features_a) + set_b = set(features_b) + if set_a == set_b: + raise ValueError( + "Feature order mismatch between datasets. " + "The same features appear in a different order, which would silently " + "corrupt combined tensors." + ) + + only_a = set_a - set_b + only_b = set_b - set_a raise ValueError( f"Feature mismatch between datasets.\n" f" Only in dataset A: {only_a}\n" @@ -89,6 +124,89 @@ def validate_feature_compatibility(meta_a: dict, meta_b: dict) -> None: "Both datasets must have the same features for combined training." ) + invariant_fields = ( + "feature_set", + "seq_length_hours", + "min_stay_hours", + "label_horizon_hours", + ) + mismatches = [] + for field in invariant_fields: + value_a = meta_a.get(field) + value_b = meta_b.get(field) + if value_a != value_b: + mismatches.append((field, value_a, value_b)) + + if mismatches: + details = "\n".join( + f" {field}: dataset A={value_a}, dataset B={value_b}" + for field, value_a, value_b in mismatches + ) + raise ValueError( + "Dataset preprocessing invariants do not match.\n" + f"{details}\n" + "Re-extract both datasets with the same preprocessing settings before combining." + ) + + +def validate_frame_schema(name: str, df_a: pl.DataFrame, df_b: pl.DataFrame) -> None: + """Verify matching column order and dtypes for a dataframe pair.""" + if df_a.columns != df_b.columns: + raise ValueError( + f"{name} schema mismatch between datasets.\n" + f" Dataset A columns: {df_a.columns}\n" + f" Dataset B columns: {df_b.columns}\n" + "Column order must match exactly before concatenation." + ) + + dtype_mismatches = [ + (col, df_a[col].dtype, df_b[col].dtype) + for col in df_a.columns + if df_a[col].dtype != df_b[col].dtype + ] + if dtype_mismatches: + details = "\n".join( + f" {col}: dataset A={dtype_a}, dataset B={dtype_b}" + for col, dtype_a, dtype_b in dtype_mismatches + ) + raise ValueError( + f"{name} dtype mismatch between datasets.\n{details}\n" + "Re-extract both datasets with matching schemas before combining." + ) + + +def validate_required_task_coverage( + meta_a: dict, + meta_b: dict, + labels_a: pl.DataFrame, + labels_b: pl.DataFrame, + names: tuple[str, str], + required_tasks: tuple[str, ...] = THESIS_TASKS, +) -> None: + """Fail closed if either source dataset lacks a thesis task or manifest entry.""" + required = set(required_tasks) + task_sets = [set(meta_a.get("task_names", [])), set(meta_b.get("task_names", []))] + manifests = [meta_a.get("label_manifest", {}) or {}, meta_b.get("label_manifest", {}) or {}] + label_columns = [set(labels_a.columns), set(labels_b.columns)] + + messages = [] + for name, tasks, manifest, columns in zip(names, task_sets, manifests, label_columns): + missing_tasks = sorted(required - tasks) + missing_manifest = sorted(required - set(manifest)) + missing_columns = sorted(required - columns) + if missing_tasks: + messages.append(f"{name} metadata task_names missing {missing_tasks}") + if missing_manifest: + messages.append(f"{name} label_manifest missing {missing_manifest}") + if missing_columns: + messages.append(f"{name} labels.parquet missing columns {missing_columns}") + + if messages: + raise ValueError( + "Combined dataset creation requires both source datasets to contain all " + "thesis tasks before merge:\n " + "\n ".join(messages) + ) + def merge_labels(labels_a: pl.DataFrame, labels_b: pl.DataFrame) -> pl.DataFrame: """Merge label DataFrames, handling column mismatches.""" @@ -132,6 +250,17 @@ def main(): default=None, help="Dataset names (for metadata). Defaults to directory names.", ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Split seed to use when preparing the combined dataset (default: 42).", + ) + parser.add_argument( + "--skip-prepare", + action="store_true", + help="Only write merged parquet/metadata files and skip splits/normalization prep.", + ) args = parser.parse_args() source_a = Path(args.source[0]) @@ -161,9 +290,17 @@ def main(): data_b = load_dataset(source_b) print(f" {len(data_b['static'])} stays, {len(data_b['timeseries'])} timeseries rows") + # Namespace patient_id for both datasets up front so patient-level split + # generation cannot merge unrelated cross-dataset patients with the same + # raw identifier. + data_a["static"] = namespace_patient_ids(data_a["static"], names[0]) + data_b["static"] = namespace_patient_ids(data_b["static"], names[1]) + print(" Patient IDs namespaced by source dataset.") + # Validate feature compatibility print("\nValidating feature compatibility...") validate_feature_compatibility(data_a["metadata"], data_b["metadata"]) + validate_frame_schema("timeseries", data_a["timeseries"], data_b["timeseries"]) print(" Features match.") # Check for natural ID collisions @@ -189,6 +326,15 @@ def main(): data_a["static"] = data_a["static"].with_columns(pl.lit(names[0]).alias("source_dataset")) data_b["static"] = data_b["static"].with_columns(pl.lit(names[1]).alias("source_dataset")) + # Harmonize column dtypes across datasets (e.g. patient_id: Int32 vs String) + for col in set(data_a["static"].columns) & set(data_b["static"].columns): + dt_a = data_a["static"][col].dtype + dt_b = data_b["static"][col].dtype + if dt_a != dt_b: + print(f" Harmonizing column '{col}': {dt_a} vs {dt_b} → String") + data_a["static"] = data_a["static"].with_columns(pl.col(col).cast(pl.Utf8)) + data_b["static"] = data_b["static"].with_columns(pl.col(col).cast(pl.Utf8)) + # Merge print("\nMerging...") static_combined = pl.concat([data_a["static"], data_b["static"]]) @@ -207,6 +353,75 @@ def main(): tasks_a = set(meta_a.get("task_names", [])) tasks_b = set(meta_b.get("task_names", [])) common_tasks = sorted(tasks_a & tasks_b) + validate_required_task_coverage( + meta_a, + meta_b, + data_a["labels"], + data_b["labels"], + names, + ) + + # Merge label manifests from source datasets for freshness checking. + # Both source datasets must have manifests, and for each common task + # the builder_version and config_hash must agree — otherwise the + # combined labels would mix stale and fresh data. + manifest_a = meta_a.get("label_manifest", {}) + manifest_b = meta_b.get("label_manifest", {}) + + if not manifest_a: + print( + f"\n ERROR: {names[0]} has no label_manifest in metadata.yaml.\n" + f" Re-run extraction for {names[0]} before combining.", + file=sys.stderr, + ) + sys.exit(1) + if not manifest_b: + print( + f"\n ERROR: {names[1]} has no label_manifest in metadata.yaml.\n" + f" Re-run extraction for {names[1]} before combining.", + file=sys.stderr, + ) + sys.exit(1) + + combined_manifest = {} + for task in common_tasks: + entry_a = manifest_a.get(task) + entry_b = manifest_b.get(task) + + if entry_a is None and entry_b is None: + continue + if entry_a is None or entry_b is None: + present = names[0] if entry_a else names[1] + missing = names[1] if entry_a else names[0] + print( + f"\n ERROR: Task '{task}' has a manifest entry in {present} " + f"but not in {missing}.\n" + f" Re-run extraction for {missing} with task '{task}'.", + file=sys.stderr, + ) + sys.exit(1) + + if entry_a.get("builder_version") != entry_b.get("builder_version"): + print( + f"\n ERROR: builder_version mismatch for task '{task}':\n" + f" {names[0]}: {entry_a.get('builder_version')}\n" + f" {names[1]}: {entry_b.get('builder_version')}\n" + f" Re-run extraction for both datasets.", + file=sys.stderr, + ) + sys.exit(1) + + if entry_a.get("config_hash") != entry_b.get("config_hash"): + print( + f"\n ERROR: config_hash mismatch for task '{task}':\n" + f" {names[0]}: {entry_a.get('config_hash')}\n" + f" {names[1]}: {entry_b.get('config_hash')}\n" + f" Both datasets must be extracted with the same task config.", + file=sys.stderr, + ) + sys.exit(1) + + combined_manifest[task] = entry_a combined_metadata = { "dataset": "combined", @@ -215,15 +430,38 @@ def main(): "feature_set": meta_a.get("feature_set", "core"), "feature_names": meta_a.get("feature_names", []), "n_features": meta_a.get("n_features", 0), - "seq_length_hours": meta_a.get("seq_length_hours", 48), - "min_stay_hours": meta_a.get("min_stay_hours", 48), + "seq_length_hours": meta_a.get("seq_length_hours", SEQ_LENGTH_HOURS), + "input_seq_length_hours": meta_a.get( + "input_seq_length_hours", + meta_a.get("seq_length_hours", SEQ_LENGTH_HOURS), + ), + "label_horizon_hours": meta_a.get("label_horizon_hours"), + "raw_export_horizon_hours_by_dataset": { + names[0]: meta_a.get("raw_export_horizon_hours"), + names[1]: meta_b.get("raw_export_horizon_hours"), + }, + "required_raw_export_horizon_hours_by_dataset": { + names[0]: meta_a.get("required_raw_export_horizon_hours"), + names[1]: meta_b.get("required_raw_export_horizon_hours"), + }, + "label_quality_stats_by_dataset": { + names[0]: meta_a.get("label_quality_stats", {}), + names[1]: meta_b.get("label_quality_stats", {}), + }, + "min_stay_hours": meta_a.get("min_stay_hours", MIN_STAY_HOURS), "task_names": common_tasks, + "label_manifest": combined_manifest, "n_stays": len(static_combined), "stays_per_dataset": { names[0]: len(data_a["static"]), names[1]: len(data_b["static"]), }, "id_offset_applied": DATASET_ID_OFFSET if natural_overlap else 0, + "patient_id_namespaced": True, + "patient_id_namespace_by_dataset": { + names[0]: names[0], + names[1]: names[1], + }, } # Save @@ -243,6 +481,16 @@ def main(): yaml.dump(combined_metadata, f, default_flow_style=False) print(" metadata.yaml") + if args.skip_prepare: + print("\nSkipping combined dataset preparation (--skip-prepare).") + else: + print("\nPreparing combined dataset artifacts...") + prepare_processed_dataset( + processed_dir=output_dir, + seed=args.seed, + dataset_name="combined", + ) + print(f"\nCombined dataset created at {output_dir}") print(f" Total stays: {len(static_combined):,}") print(f" Common tasks: {common_tasks}") diff --git a/scripts/preprocessing/extract_ricu.py b/scripts/preprocessing/extract_ricu.py index 4567354..c3fd8fe 100644 --- a/scripts/preprocessing/extract_ricu.py +++ b/scripts/preprocessing/extract_ricu.py @@ -2,7 +2,8 @@ Two-step pipeline: # Step 1: R extraction (once per dataset) - Rscript scripts/preprocessing/extract_with_ricu.R --dataset miiv + Rscript scripts/preprocessing/extract_with_ricu.R --dataset miiv \ + --raw_export_horizon_hours 48 # Step 2: Python processing -> produces final SLICES format uv run python scripts/preprocessing/extract_ricu.py dataset=miiv @@ -20,6 +21,7 @@ import hydra from omegaconf import DictConfig + from slices.constants import MIN_STAY_HOURS, SEQ_LENGTH_HOURS from slices.data.extractors.base import ExtractorConfig from slices.data.extractors.ricu import RicuExtractor @@ -45,12 +47,12 @@ def main(cfg: DictConfig) -> None: print("\nRun the R extraction first:") print( " Rscript scripts/preprocessing/extract_with_ricu.R " - f"--dataset miiv --output_dir {ricu_dir}" + f"--dataset miiv --output_dir {ricu_dir} --raw_export_horizon_hours 48" ) sys.exit(1) # Build config kwargs, only overriding tasks if explicitly specified - # so that ExtractorConfig defaults (mortality_24h, mortality_hospital) are used + # so that ExtractorConfig defaults cover the fixed thesis task surface. config_kwargs: dict = { "parquet_root": str(cfg.data.ricu_output_dir), "output_dir": str(cfg.data.processed_dir), diff --git a/scripts/preprocessing/extract_with_ricu.R b/scripts/preprocessing/extract_with_ricu.R index 0e5432b..360cad5 100644 --- a/scripts/preprocessing/extract_with_ricu.R +++ b/scripts/preprocessing/extract_with_ricu.R @@ -5,7 +5,8 @@ #' RICU handles concept lookup, unit harmonization, hourly binning, and gap filling. #' #' Output files: -#' ricu_timeseries.parquet - stay_id, hour, {concept}, {concept}_mask, ... +#' ricu_timeseries.parquet/ - directory of parquet part files (chunked by stay) +#' each part: stay_id, hour, {concept}, {concept}_mask, ... #' ricu_stays.parquet - stay_id, patient_id, intime, outtime, los_days, ... #' ricu_mortality.parquet - stay_id, date_of_death, hospital_expire_flag, ... #' ricu_diagnoses.parquet - stay_id, icd_code, icd_version @@ -15,7 +16,7 @@ #' Rscript scripts/preprocessing/extract_with_ricu.R \ #' --dataset miiv \ #' --output_dir data/ricu_output/miiv \ -#' --seq_length_hours 48 +#' --raw_export_horizon_hours 48 # Auto-install missing packages required_packages <- c("ricu", "arrow", "yaml", "data.table", "optparse", "units") @@ -26,6 +27,7 @@ if (length(missing) > 0) { install.packages(missing, repos = "https://cloud.r-project.org") } + suppressPackageStartupMessages({ library(ricu) library(arrow) @@ -43,8 +45,10 @@ option_list <- list( help = "RICU source name (miiv, eicu, hirid, aumc, mimic, sic)"), make_option("--output_dir", type = "character", default = NULL, help = "Output directory for parquet files"), - make_option("--seq_length_hours", type = "integer", default = 48L, - help = "Max hours per stay [default: %default]"), + make_option("--raw_export_horizon_hours", type = "integer", default = NULL, + help = "Max hours of raw timeseries to export per stay [default: 72]"), + make_option("--seq_length_hours", type = "integer", default = NULL, + help = "DEPRECATED alias for --raw_export_horizon_hours"), make_option("--raw_data_dir", type = "character", default = NULL, help = "Path to raw CSV files for ricu import (auto-detected if not set)") ) @@ -55,7 +59,29 @@ VALID_DATASETS <- c("miiv", "eicu", "hirid", "aumc", "mimic", "sic", "mimic_demo", "eicu_demo") # Batch size for concept loading — controls peak memory during extraction. -CONCEPT_BATCH_SIZE <- 8L +CONCEPT_BATCH_SIZE <- 4L + +# Tables actually needed for concept extraction and admin/mortality/diagnosis +# data. Derived from ricu:::tbl_cfg — tables referenced by time-series, +# admin, mortality, and diagnosis concepts. Everything else (emar_detail, +# pharmacy, poe, prescriptions, …) is unused and can take 20+ min to import +# for no benefit. +ESSENTIAL_TABLES <- list( + miiv = c("chartevents", "labevents", "inputevents", "outputevents", + "procedureevents", "ingredientevents", "datetimeevents", + "icustays", "patients", "admissions", "transfers", + "d_labitems", "diagnoses_icd", "d_icd_diagnoses"), + mimic = c("chartevents", "labevents", "inputevents", "outputevents", + "procedureevents", "ingredientevents", "datetimeevents", + "icustays", "patients", "admissions", "transfers", + "d_labitems", "diagnoses_icd", "d_icd_diagnoses"), + # eicu requires all tables: load_concepts() internally calls + # load_dictionary() which checks full source availability. + # Omitted from this list so all tables are imported (see fallback below). + hirid = c("observations", "pharma", "general"), + aumc = c("numericitems", "drugitems", "procedureorderitems", + "admissions", "listitems") +) # --------------------------------------------------------------------------- # Helpers @@ -122,7 +148,8 @@ discover_ts_concepts <- function(dataset, dict) { # 2. Extract time-series # --------------------------------------------------------------------------- -extract_timeseries <- function(dataset, concepts, seq_length_hours, dict) { +extract_timeseries <- function(dataset, concepts, seq_length_hours, dict, + output_path) { message("[2/6] Extracting time-series data in batches...") # Pre-compute stay windows (shared across all batches) @@ -137,7 +164,70 @@ extract_timeseries <- function(dataset, concepts, seq_length_hours, dict) { ceiling(seq_len(n_concepts) / CONCEPT_BATCH_SIZE) ) n_batches <- length(batch_indices) - batch_results <- vector("list", n_batches) + + # Helper: load a set of concepts into a data.table with stay_id, hour, and + # value + mask columns. Returns NULL on failure. + load_batch <- function(cnames, dataset, dict, wins, seq_length_hours) { + tryCatch({ + ts_data <- load_concepts( + cnames, + src = dataset, + concepts = dict, + interval = hours(1L), + merge_data = TRUE, + verbose = FALSE + ) + ts_dense <- fill_gaps(ts_data, limits = wins) + idx_name <- index_var(ts_dense) + ts_capped <- ts_dense[get(idx_name) < hours(seq_length_hours)] + batch_concept_cols <- data_vars(ts_capped) + for (col in batch_concept_cols) { + mask_name <- paste0(col, "_mask") + set(ts_capped, j = mask_name, value = !is.na(ts_capped[[col]])) + } + dt <- as.data.table(ts_capped) + id_name <- id_var(ts_capped) + dt[, hour := as.integer(as.numeric(get(idx_name), units = "hours"))] + setnames(dt, id_name, "stay_id") + dt[, (idx_name) := NULL] + rm(ts_data, ts_dense, ts_capped) + gc() + dt + }, error = function(e) { + warning(sprintf(" Failed to load concepts [%s]: %s", + paste(cnames, collapse = ", "), conditionMessage(e))) + NULL + }) + } + + # --- Phase 0: Discover stay IDs and assign to chunks ------------------- + # We need stay IDs before processing batches so we can pre-split each + # batch by stay chunk on first read (avoiding re-reading batch files). + message(" Discovering stay IDs...") + all_stay_ids <- sort(unique(as.data.table(wins)[[id_var(wins)]])) + STAY_CHUNK_SIZE <- 10000L + # Map each stay_id to its chunk index for fast lookup + chunk_assignment <- ceiling(seq_along(all_stay_ids) / STAY_CHUNK_SIZE) + names(chunk_assignment) <- as.character(all_stay_ids) + n_chunks <- max(chunk_assignment) + n_stays <- length(all_stay_ids) + message(sprintf(" %d stays → %d chunks of ≤%d stays.", + n_stays, n_chunks, STAY_CHUNK_SIZE)) + + # --- Phase 1: Extract concept batches and pre-split by stay chunk ------ + # Each concept batch is loaded once, split by stay chunk, and written to + # tmp//batch_.parquet. This way each batch's data + # is read from RICU exactly once (not re-read per chunk in Phase 2). + tmp_dir <- file.path(tempdir(), "ricu_merge") + unlink(tmp_dir, recursive = TRUE) + for (ci in seq_len(n_chunks)) { + dir.create(file.path(tmp_dir, sprintf("chunk_%04d", ci)), + recursive = TRUE, showWarnings = FALSE) + } + + skipped_concepts <- character(0) + all_concept_cols <- character(0) + batches_written <- 0L for (b in seq_len(n_batches)) { batch_concepts <- concepts[batch_indices[[b]]] @@ -145,69 +235,120 @@ extract_timeseries <- function(dataset, concepts, seq_length_hours, dict) { b, n_batches, length(batch_concepts), paste(batch_concepts, collapse = ", "))) - # Load this batch — RICU handles unit conversion, hourly binning, median agg - ts_data <- load_concepts( - batch_concepts, - src = dataset, - concepts = dict, - interval = hours(1L), - merge_data = TRUE, - verbose = FALSE - ) - - # Fill gaps to create a dense hourly grid (capped at seq_length_hours) - ts_dense <- fill_gaps(ts_data, limits = wins) - - # Cap at seq_length_hours (belt-and-suspenders) - idx_name <- index_var(ts_dense) - ts_capped <- ts_dense[get(idx_name) < hours(seq_length_hours)] + dt <- load_batch(batch_concepts, dataset, dict, wins, seq_length_hours) + + if (is.null(dt) && length(batch_concepts) > 1) { + # Retry concepts individually to isolate the broken one(s) + message(" Batch failed — retrying concepts individually...") + individual_dts <- list() + for (cname in batch_concepts) { + cdt <- load_batch(cname, dataset, dict, wins, seq_length_hours) + if (is.null(cdt)) { + message(sprintf(" Skipping concept '%s' (unsupported for this dataset).", cname)) + skipped_concepts <- c(skipped_concepts, cname) + } else { + individual_dts[[length(individual_dts) + 1]] <- cdt + } + rm(cdt); gc() + } + if (length(individual_dts) > 0) { + dt <- individual_dts[[1]] + if (length(individual_dts) > 1) { + for (j in 2:length(individual_dts)) { + dt <- merge(dt, individual_dts[[j]], + by = c("stay_id", "hour"), all = TRUE) + } + } + rm(individual_dts); gc() + } else { + dt <- NULL + } + } else if (is.null(dt)) { + message(sprintf(" Skipping concept '%s' (unsupported for this dataset).", + batch_concepts)) + skipped_concepts <- c(skipped_concepts, batch_concepts) + } - # Derive boolean observation masks for this batch's concepts - batch_concept_cols <- data_vars(ts_capped) - for (col in batch_concept_cols) { - mask_name <- paste0(col, "_mask") - set(ts_capped, j = mask_name, value = !is.na(ts_capped[[col]])) + if (!is.null(dt)) { + # Track concept columns from this batch + batch_cols <- names(dt) + batch_mask <- batch_cols[grepl("_mask$", batch_cols)] + all_concept_cols <- c(all_concept_cols, sub("_mask$", "", batch_mask)) + + # Split by stay chunk and write to per-chunk directories + dt[, .chunk_idx := chunk_assignment[as.character(stay_id)]] + for (ci in unique(dt$.chunk_idx)) { + chunk_dt <- dt[.chunk_idx == ci] + chunk_dt[, .chunk_idx := NULL] + chunk_path <- file.path(tmp_dir, sprintf("chunk_%04d", ci), + sprintf("batch_%03d.parquet", b)) + write_parquet(chunk_dt, chunk_path) + rm(chunk_dt) + } + batches_written <- batches_written + 1L } + rm(dt); gc() + } - # Convert to data.table for efficient merging - dt <- as.data.table(ts_capped) + rm(wins); gc() - # Integer hour column from the difftime index - id_name <- id_var(ts_capped) - dt[, hour := as.integer(as.numeric(get(idx_name), units = "hours"))] - setnames(dt, id_name, "stay_id") - dt[, (idx_name) := NULL] + if (length(skipped_concepts) > 0) { + message(sprintf(" Skipped %d concepts: %s", + length(skipped_concepts), + paste(skipped_concepts, collapse = ", "))) + } - batch_results[[b]] <- dt + if (batches_written == 0) { + stop("All concept batches failed. Cannot continue.") + } - # Free intermediate objects - rm(ts_data, ts_dense, ts_capped, dt) + # --- Phase 2: Merge pre-split batches per chunk ------------------------- + # For each stay chunk, read its small batch files, merge wide, and write + # one output parquet part. Peak memory ≈ one chunk's worth of data in + # narrow form + one chunk wide (~1-2 GB for 10k stays). + message(sprintf(" Merging %d chunks...", n_chunks)) + dir.create(output_path, showWarnings = FALSE, recursive = TRUE) + total_rows <- 0L + + for (ci in seq_len(n_chunks)) { + chunk_dir <- file.path(tmp_dir, sprintf("chunk_%04d", ci)) + batch_files <- list.files(chunk_dir, pattern = "\\.parquet$", + full.names = TRUE) + if (length(batch_files) == 0) next + + chunk_merged <- NULL + for (bf in batch_files) { + batch_dt <- as.data.table(read_parquet(bf)) + setkeyv(batch_dt, c("stay_id", "hour")) + + if (is.null(chunk_merged)) { + chunk_merged <- batch_dt + } else { + chunk_merged <- merge(chunk_merged, batch_dt, + by = c("stay_id", "hour"), all = TRUE) + } + rm(batch_dt) + } gc() - } - # Merge all batches by (stay_id, hour) - message(" Merging batches...") - merged <- batch_results[[1]] - if (n_batches > 1) { - for (b in 2:n_batches) { - merged <- merge(merged, batch_results[[b]], - by = c("stay_id", "hour"), all = TRUE) - batch_results[[b]] <- NULL - gc() + total_rows <- total_rows + nrow(chunk_merged) + + out_path <- file.path(output_path, sprintf("part_%04d.parquet", ci)) + write_parquet(chunk_merged, out_path) + rm(chunk_merged); gc() + + if (ci %% 5 == 0 || ci == n_chunks) { + message(sprintf(" Chunk %d/%d done.", ci, n_chunks)) } } - rm(batch_results, wins) - gc() - - # Collect concept column names (value columns paired with _mask columns) - all_cols <- names(merged) - mask_cols <- all_cols[grepl("_mask$", all_cols)] - concept_cols <- sub("_mask$", "", mask_cols) message(sprintf(" Time-series: %d rows, %d concepts, %d stays.", - nrow(merged), length(concept_cols), - length(unique(merged$stay_id)))) - list(df = merged, concept_cols = concept_cols) + total_rows, length(all_concept_cols), n_stays)) + + # Clean up temp files + unlink(tmp_dir, recursive = TRUE) + + list(concept_cols = all_concept_cols, n_stays = n_stays) } # --------------------------------------------------------------------------- @@ -476,9 +617,10 @@ extract_mortality_data <- function(dataset, dict) { extract_mortality_generic(dataset, dict) ) - # Ensure expected columns - expected <- c("stay_id", "date_of_death", "hospital_expire_flag", - "dischtime", "discharge_location") + # Ensure expected columns (new precision-aware schema + legacy) + expected <- c("stay_id", "date_of_death", "death_time", "death_date", + "death_time_precision", "death_source", + "hospital_expire_flag", "dischtime", "discharge_location") for (col in expected) { if (!col %in% names(result)) result[[col]] <- NA } @@ -497,6 +639,7 @@ extract_mortality_miiv <- function(dataset, dict) { df <- data.frame(stay_id = icu$stay_id, stringsAsFactors = FALSE) + # Get date-only dod from patients table if (!is.null(pat)) { pat_sub <- pat[, intersect(c("subject_id", "dod"), names(pat)), drop = FALSE] @@ -507,10 +650,11 @@ extract_mortality_miiv <- function(dataset, dict) { names(df)[names(df) == "dod"] <- "date_of_death" } + # Get deathtime + other columns from admissions table if (!is.null(adm)) { adm_cols <- intersect( c("hadm_id", "hospital_expire_flag", "dischtime", - "discharge_location"), + "discharge_location", "deathtime"), names(adm) ) if (length(adm_cols) > 1) { @@ -523,6 +667,38 @@ extract_mortality_miiv <- function(dataset, dict) { } } + # Build precision-aware death schema: + # - death_time: exact timestamp from admissions.deathtime (preferred) + # - death_date: date-only from patients.dod + # - death_time_precision: "timestamp" / "date" / "unknown" + # - death_source: which column the death info comes from + has_deathtime <- "deathtime" %in% names(df) & !is.na(df$deathtime) + has_dod <- "date_of_death" %in% names(df) & !is.na(df$date_of_death) + + df$death_time <- as.POSIXct(rep(NA, nrow(df)), tz = "UTC") + df$death_date <- as.Date(rep(NA, nrow(df))) + df$death_time_precision <- NA_character_ + df$death_source <- NA_character_ + + # Prefer admissions.deathtime (exact timestamp) + if ("deathtime" %in% names(df)) { + idx_dt <- !is.na(df$deathtime) + df$death_time[idx_dt] <- df$deathtime[idx_dt] + df$death_time_precision[idx_dt] <- "timestamp" + df$death_source[idx_dt] <- "admissions.deathtime" + } + + # Fall back to patients.dod (date-only) for rows without deathtime + if ("date_of_death" %in% names(df)) { + idx_dod_only <- is.na(df$death_time) & !is.na(df$date_of_death) + df$death_date[idx_dod_only] <- as.Date(df$date_of_death[idx_dod_only]) + df$death_time_precision[idx_dod_only] <- "date" + df$death_source[idx_dod_only] <- "patients.dod" + } + + # Clean up intermediate column + df$deathtime <- NULL + df } @@ -550,7 +726,7 @@ extract_mortality_mimic <- function(dataset, dict) { if (!is.null(adm)) { adm_cols <- intersect( c("hadm_id", "hospital_expire_flag", "dischtime", - "discharge_location"), + "discharge_location", "deathtime"), names(adm) ) if (length(adm_cols) > 1) { @@ -564,6 +740,28 @@ extract_mortality_mimic <- function(dataset, dict) { } } + # Build precision-aware death schema (same logic as miiv) + df$death_time <- as.POSIXct(rep(NA, nrow(df)), tz = "UTC") + df$death_date <- as.Date(rep(NA, nrow(df))) + df$death_time_precision <- NA_character_ + df$death_source <- NA_character_ + + if ("deathtime" %in% names(df)) { + idx_dt <- !is.na(df$deathtime) + df$death_time[idx_dt] <- df$deathtime[idx_dt] + df$death_time_precision[idx_dt] <- "timestamp" + df$death_source[idx_dt] <- "admissions.deathtime" + } + + if ("date_of_death" %in% names(df)) { + idx_dod_only <- is.na(df$death_time) & !is.na(df$date_of_death) + df$death_date[idx_dod_only] <- as.Date(df$date_of_death[idx_dod_only]) + df$death_time_precision[idx_dod_only] <- "date" + df$death_source[idx_dod_only] <- "patients.dod" + } + + df$deathtime <- NULL + df } @@ -577,23 +775,125 @@ extract_mortality_eicu <- function(dataset, dict) { names(pat)[1] } - # eICU uses status strings and offsets rather than timestamps - hosp_flag <- if ("hospitaldischargestatus" %in% names(pat)) { - as.integer(tolower(pat$hospitaldischargestatus) == "expired") - } else { - NA_integer_ + # eICU uses status/location strings and offsets rather than timestamps. + # Preserve the hospital mortality outcome as tri-state: + # 1 = death evidence from discharge status/location + # 0 = explicit alive hospital discharge status + # NA = unknown outcome + # This avoids converting missing/unknown statuses into false survivors, and + # it recovers rows where hospitaldischargestatus is missing but discharge + # location says Death. + normalize_text <- function(x) { + out <- tolower(trimws(as.character(x))) + out[out %in% c("", "na", "nan", "null", "unknown")] <- NA_character_ + out + } + get_text_col <- function(col) { + if (col %in% names(pat)) { + as.character(pat[[col]]) + } else { + rep(NA_character_, nrow(pat)) + } } + get_numeric_col <- function(col) { + if (col %in% names(pat)) { + suppressWarnings(as.numeric(pat[[col]])) + } else { + rep(NA_real_, nrow(pat)) + } + } + + hospital_status <- get_text_col("hospitaldischargestatus") + unit_status <- get_text_col("unitdischargestatus") + hospital_location <- get_text_col("hospitaldischargelocation") + unit_location <- get_text_col("unitdischargelocation") + + hospital_status_norm <- normalize_text(hospital_status) + unit_status_norm <- normalize_text(unit_status) + hospital_location_norm <- normalize_text(hospital_location) + unit_location_norm <- normalize_text(unit_location) + + death_values <- c("expired", "death", "died", "deceased", "dead") + alive_values <- c("alive") + + hospital_status_death <- hospital_status_norm %in% death_values + unit_status_death <- unit_status_norm %in% death_values + hospital_location_death <- hospital_location_norm %in% death_values + unit_location_death <- unit_location_norm %in% death_values + + death_evidence <- hospital_status_death | unit_status_death | + hospital_location_death | unit_location_death + explicit_alive <- hospital_status_norm %in% alive_values + + hosp_flag <- rep(NA_integer_, nrow(pat)) + hosp_flag[explicit_alive] <- 0L + hosp_flag[death_evidence] <- 1L + + discharge_location <- hospital_location + use_unit_location <- is.na(discharge_location) | trimws(discharge_location) == "" + discharge_location[use_unit_location] <- unit_location[use_unit_location] + death_location <- ifelse(hospital_location_death, hospital_location, + ifelse(unit_location_death, unit_location, NA_character_)) + use_death_location <- death_evidence & !is.na(death_location) + discharge_location[use_death_location] <- death_location[use_death_location] + + # Derive date_of_death/death_time from discharge offsets for death-evidence + # rows. Hospital death evidence prefers hospitaldischargeoffset; unit death + # evidence can use unitdischargeoffset if hospital offset is unavailable. + # We use the same synthetic epoch as the stays extraction so that windowed + # mortality tasks (e.g., mortality_24h) can compute whether death falls + # within the prediction window. + epoch <- as.POSIXct("2000-01-01 00:00:00", tz = "UTC") + hospital_offset <- get_numeric_col("hospitaldischargeoffset") + unit_offset <- get_numeric_col("unitdischargeoffset") + + death_offset <- rep(NA_real_, nrow(pat)) + hospital_death_evidence <- hospital_status_death | hospital_location_death + unit_death_evidence <- unit_status_death | unit_location_death + death_offset[hospital_death_evidence] <- hospital_offset[hospital_death_evidence] + needs_unit_offset <- is.na(death_offset) & unit_death_evidence + death_offset[needs_unit_offset] <- unit_offset[needs_unit_offset] + needs_hospital_fallback <- is.na(death_offset) & death_evidence + death_offset[needs_hospital_fallback] <- hospital_offset[needs_hospital_fallback] + needs_unit_fallback <- is.na(death_offset) & death_evidence + death_offset[needs_unit_fallback] <- unit_offset[needs_unit_fallback] + + dod <- as.POSIXct(rep(NA, nrow(pat)), tz = "UTC") + has_death_offset <- death_evidence & !is.na(death_offset) + dod[has_death_offset] <- epoch + death_offset[has_death_offset] * 60 + + # Also derive dischtime for all rows from hospital offset, falling back to + # unit offset when hospital discharge timing is unavailable. + discharge_offset <- hospital_offset + missing_discharge_offset <- is.na(discharge_offset) + discharge_offset[missing_discharge_offset] <- unit_offset[missing_discharge_offset] + dischtime <- as.POSIXct(rep(NA, nrow(pat)), tz = "UTC") + has_discharge_offset <- !is.na(discharge_offset) + dischtime[has_discharge_offset] <- epoch + discharge_offset[has_discharge_offset] * 60 + + death_source <- rep(NA_character_, nrow(pat)) + death_source[hospital_status_death] <- "patient.hospitaldischargestatus" + death_source[is.na(death_source) & hospital_location_death] <- + "patient.hospitaldischargelocation" + death_source[is.na(death_source) & unit_status_death] <- "patient.unitdischargestatus" + death_source[is.na(death_source) & unit_location_death] <- "patient.unitdischargelocation" df <- data.frame( stay_id = pat[[id_col]], - date_of_death = NA, + date_of_death = dod, hospital_expire_flag = hosp_flag, - dischtime = NA, - discharge_location = if ("unitdischargelocation" %in% names(pat)) { - pat$unitdischargelocation - } else { - NA_character_ - }, + dischtime = dischtime, + discharge_location = discharge_location, + # Precision-aware schema: eICU timestamps are derived from offsets + # (minute-level precision) when an offset is available. + death_time = dod, + death_date = as.Date(rep(NA, nrow(pat))), + death_time_precision = ifelse( + has_death_offset, + "timestamp", + ifelse(death_evidence, "unknown", NA_character_) + ), + death_source = death_source, stringsAsFactors = FALSE ) df @@ -605,12 +905,17 @@ extract_mortality_generic <- function(dataset, dict) { if (!"death" %in% avail) { message(" Warning: 'death' concept not available; returning empty mortality.") wins <- stay_windows(dataset, interval = hours(1L)) + n <- nrow(as.data.frame(wins)) return(data.frame( stay_id = as.data.frame(wins)[[id_var(wins)]], - date_of_death = NA, + date_of_death = as.POSIXct(rep(NA, n), tz = "UTC"), hospital_expire_flag = NA_integer_, - dischtime = NA, + dischtime = as.POSIXct(rep(NA, n), tz = "UTC"), discharge_location = NA_character_, + death_time = as.POSIXct(rep(NA, n), tz = "UTC"), + death_date = as.Date(rep(NA, n)), + death_time_precision = NA_character_, + death_source = NA_character_, stringsAsFactors = FALSE )) } @@ -618,13 +923,18 @@ extract_mortality_generic <- function(dataset, dict) { death_raw <- load_concepts("death", src = dataset, concepts = dict) death_df <- as.data.frame(death_raw) id_name <- id_var(death_raw) + n <- nrow(death_df) df <- data.frame( stay_id = death_df[[id_name]], - date_of_death = NA, + date_of_death = as.POSIXct(rep(NA, n), tz = "UTC"), hospital_expire_flag = as.integer(death_df$death), - dischtime = NA, + dischtime = as.POSIXct(rep(NA, n), tz = "UTC"), discharge_location = NA_character_, + death_time = as.POSIXct(rep(NA, n), tz = "UTC"), + death_date = as.Date(rep(NA, n)), + death_time_precision = NA_character_, + death_source = NA_character_, stringsAsFactors = FALSE ) df @@ -759,15 +1069,16 @@ extract_diagnoses_eicu <- function(dataset) { # --------------------------------------------------------------------------- write_metadata <- function(output_dir, dataset, concept_cols, - seq_length_hours, n_stays) { + raw_export_horizon_hours, n_stays) { message("[6/6] Writing metadata...") metadata <- list( - dataset = dataset, - feature_names = as.list(concept_cols), - n_features = length(concept_cols), - seq_length_hours = seq_length_hours, - n_stays = n_stays, - ricu_version = as.character(packageVersion("ricu")) + dataset = dataset, + feature_names = as.list(concept_cols), + n_features = length(concept_cols), + seq_length_hours = raw_export_horizon_hours, + raw_export_horizon_hours = raw_export_horizon_hours, + n_stays = n_stays, + ricu_version = as.character(packageVersion("ricu")) ) write_yaml(metadata, file.path(output_dir, "ricu_metadata.yaml")) } @@ -777,14 +1088,38 @@ write_metadata <- function(output_dir, dataset, concept_cols, # --------------------------------------------------------------------------- main <- function(opts) { - dataset <- opts$dataset - output_dir <- opts$output_dir - seq_length_hours <- opts$seq_length_hours + dataset <- opts$dataset + output_dir <- opts$output_dir + + raw_export_horizon_hours <- opts$raw_export_horizon_hours + deprecated_seq_length_hours <- opts$seq_length_hours + if (is.null(raw_export_horizon_hours)) { + if (!is.null(deprecated_seq_length_hours)) { + raw_export_horizon_hours <- deprecated_seq_length_hours + warning( + "--seq_length_hours is deprecated; use --raw_export_horizon_hours instead.", + call. = FALSE + ) + } else { + raw_export_horizon_hours <- 72L + } + } else if (!is.null(deprecated_seq_length_hours) && + raw_export_horizon_hours != deprecated_seq_length_hours) { + stop( + paste0( + "Conflicting values provided for --raw_export_horizon_hours and ", + "--seq_length_hours. Use only --raw_export_horizon_hours." + ) + ) + } # Validate arguments if (is.null(dataset) || is.null(output_dir)) { stop("Both --dataset and --output_dir are required.") } + if (raw_export_horizon_hours <= 0) { + stop("--raw_export_horizon_hours must be a positive integer.") + } if (!dataset %in% VALID_DATASETS) { stop(sprintf("Invalid dataset '%s'. Valid: %s", dataset, paste(VALID_DATASETS, collapse = ", "))) @@ -857,12 +1192,21 @@ main <- function(opts) { ) ) - # ricu requires ALL tables to be present before it considers a source - # "available" (e.g. load_dictionary(src=...) will fail otherwise). - # Import every table, but one at a time to control peak memory. - message(sprintf(" Importing %d tables one at a time...", length(all_tables))) - for (tbl_name in all_tables) { - message(sprintf(" Importing: %s", tbl_name)) + # Only import tables actually needed for concept extraction and admin + # data. We bypass RICU's availability gate with parse_ricu_dictionary() + # later, so we don't need ALL tables — just the ones concepts reference. + fam <- dataset_family(dataset) + if (fam %in% names(ESSENTIAL_TABLES)) { + tables_to_import <- intersect(all_tables, ESSENTIAL_TABLES[[fam]]) + } else { + tables_to_import <- all_tables + } + n_import <- length(tables_to_import) + message(sprintf(" Importing %d/%d essential tables...", + n_import, length(all_tables))) + for (i in seq_along(tables_to_import)) { + tbl_name <- tables_to_import[i] + message(sprintf(" [%d/%d] Importing: %s", i, n_import, tbl_name)) tryCatch({ import_src(dataset, tables = tbl_name) }, error = function(e) { @@ -887,8 +1231,8 @@ main <- function(opts) { } dir.create(output_dir, recursive = TRUE, showWarnings = FALSE) - message(sprintf("=== RICU extraction: dataset=%s, output=%s, hours=%d ===", - dataset, output_dir, seq_length_hours)) + message(sprintf("=== RICU extraction: dataset=%s, output=%s, raw_hours=%d ===", + dataset, output_dir, raw_export_horizon_hours)) # Parse the concept dictionary once, bypassing RICU's availability gate # which can fail even after successful import_src(). @@ -900,13 +1244,12 @@ main <- function(opts) { stop("No time-series concepts found. Check ricu setup.") } - # 2. Extract timeseries - ts_result <- extract_timeseries(dataset, concepts, seq_length_hours, dict) + # 2. Extract timeseries (writes parquet directly to avoid full-table merge) + ts_path <- file.path(output_dir, "ricu_timeseries.parquet") + ts_result <- extract_timeseries(dataset, concepts, raw_export_horizon_hours, dict, + output_path = ts_path) concept_cols <- ts_result$concept_cols - n_stays <- length(unique(ts_result$df$stay_id)) - # arrow::write_parquet writes data.table directly — no as.data.frame() needed - write_parquet(ts_result$df, file.path(output_dir, "ricu_timeseries.parquet")) - rm(ts_result); gc() + n_stays <- ts_result$n_stays # 3. Extract stays stays <- extract_stays_data(dataset, dict) @@ -925,7 +1268,7 @@ main <- function(opts) { # 6. Write metadata write_metadata(output_dir, dataset, concept_cols, - seq_length_hours, n_stays) + raw_export_horizon_hours, n_stays) message("=== Extraction complete. ===") message(sprintf(" Output: %s", output_dir)) diff --git a/scripts/preprocessing/prepare_dataset.py b/scripts/preprocessing/prepare_dataset.py index b3ff47c..ab61276 100644 --- a/scripts/preprocessing/prepare_dataset.py +++ b/scripts/preprocessing/prepare_dataset.py @@ -1,290 +1,22 @@ -"""Prepare dataset splits and normalization statistics. +"""Prepare dataset splits and normalization statistics.""" -This script computes patient-level train/val/test splits and normalization -statistics on the training set only. Run this ONCE after extraction, before -training. - -The output files (splits.yaml, normalization_stats.yaml) are required for -reproducible training runs and prevent data leakage from val/test sets. - -Usage: - uv run python scripts/preprocessing/prepare_dataset.py dataset=miiv - uv run python scripts/preprocessing/prepare_dataset.py dataset=eicu -""" - -import os from pathlib import Path -from tempfile import NamedTemporaryFile import hydra -import numpy as np -import polars as pl -import yaml from omegaconf import DictConfig -from slices.constants import TEST_RATIO, TRAIN_RATIO, VAL_RATIO - - -def _atomic_yaml_write(path: Path, data: dict) -> None: - """Write YAML atomically using temp file + rename.""" - with NamedTemporaryFile(dir=path.parent, suffix=".yaml", mode="w", delete=False) as tmp: - tmp_path = tmp.name - yaml.dump(data, tmp, default_flow_style=False) - os.replace(tmp_path, path) - - -def compute_patient_splits( - static_df: pl.DataFrame, - timeseries_df: pl.DataFrame, - train_ratio: float, - val_ratio: float, - test_ratio: float, - seed: int, -) -> dict: - """Compute patient-level splits. - - All stays from the same patient go to the same split. - - Returns: - Dictionary with split information including indices and patient lists. - """ - # Validate ratios - total = train_ratio + val_ratio + test_ratio - if not np.isclose(total, 1.0): - raise ValueError(f"Split ratios must sum to 1.0, got {total}") - - # Get stay_ids in timeseries order (this is the canonical ordering) - stay_ids = timeseries_df["stay_id"].to_list() - - # Get stay_id -> patient_id mapping - stay_to_patient = dict(zip(static_df["stay_id"].to_list(), static_df["patient_id"].to_list())) - - # Get unique patients (sorted for deterministic ordering across Python runs) - unique_patients = sorted(set(stay_to_patient.values())) - n_patients = len(unique_patients) - - print(f" Total patients: {n_patients:,}") - print(f" Total stays: {len(stay_ids):,}") - - # Warn if patient_id == stay_id for all stays (e.g. HiRID, SICdb) - if n_patients == len(stay_ids) and n_patients > 0: - patient_set = set(stay_to_patient.values()) - stay_set = set(stay_ids) - if patient_set == stay_set: - print( - " WARNING: patient_id == stay_id for all stays. This dataset likely " - "lacks true patient-level identifiers (e.g. HiRID, SICdb). " - "Patient-level split cannot prevent leakage from repeat ICU admissions." - ) - - # Shuffle patients deterministically - rng = np.random.RandomState(seed) - patient_indices = np.arange(n_patients) - rng.shuffle(patient_indices) - shuffled_patients = [unique_patients[i] for i in patient_indices] - - # Split patients - n_train = int(n_patients * train_ratio) - n_val = int(n_patients * val_ratio) - - train_patients = set(shuffled_patients[:n_train]) - val_patients = set(shuffled_patients[n_train : n_train + n_val]) - test_patients = set(shuffled_patients[n_train + n_val :]) - - # Verify no overlap - assert train_patients.isdisjoint(val_patients), "Train/val overlap!" - assert train_patients.isdisjoint(test_patients), "Train/test overlap!" - assert val_patients.isdisjoint(test_patients), "Val/test overlap!" - - # Map to stay indices - train_indices = [] - val_indices = [] - test_indices = [] - - for idx, stay_id in enumerate(stay_ids): - patient_id = stay_to_patient[stay_id] - if patient_id in train_patients: - train_indices.append(idx) - elif patient_id in val_patients: - val_indices.append(idx) - else: - test_indices.append(idx) - - return { - "seed": seed, - "train_ratio": train_ratio, - "val_ratio": val_ratio, - "test_ratio": test_ratio, - "train_patients": sorted(train_patients), - "val_patients": sorted(val_patients), - "test_patients": sorted(test_patients), - "train_indices": train_indices, - "val_indices": val_indices, - "test_indices": test_indices, - "train_stays": len(train_indices), - "val_stays": len(val_indices), - "test_stays": len(test_indices), - } - - -def compute_normalization_stats( - timeseries_df: pl.DataFrame, - train_indices: list, - feature_names: list, - seq_length: int, -) -> dict: - """Compute mean and std for each feature on training set only. - Uses vectorized numpy operations for efficiency. - """ - n_features = len(feature_names) - n_train = len(train_indices) - - # Extract training data into numpy arrays for vectorized computation - raw_timeseries = timeseries_df["timeseries"].to_list() - raw_masks = timeseries_df["mask"].to_list() - - # Accumulators - sums = np.zeros(n_features) - sq_sums = np.zeros(n_features) - counts = np.zeros(n_features) - - for progress, idx in enumerate(train_indices): - if (progress + 1) % 5000 == 0: - print(f" Processing {progress + 1:,}/{n_train:,} training samples...") - - # Convert to numpy arrays for vectorized ops - ts_arr = np.array(raw_timeseries[idx][:seq_length], dtype=np.float64) - mask_arr = np.array(raw_masks[idx][:seq_length], dtype=bool) - - # Valid = mask is True AND value is finite - valid = mask_arr & np.isfinite(ts_arr) - - # Zero out invalid entries so they don't affect sums - ts_valid = np.where(valid, ts_arr, 0.0) - - # Accumulate per-feature sums along time axis - counts += valid.sum(axis=0) - sums += ts_valid.sum(axis=0) - sq_sums += (ts_valid**2).sum(axis=0) - - # Compute mean and std - means = np.zeros(n_features) - stds = np.ones(n_features) - - safe_counts = np.maximum(counts, 1) - means = sums / safe_counts - variance = (sq_sums - counts * means**2) / np.maximum(safe_counts - 1, 1) - stds = np.sqrt(np.maximum(variance, 0.0)) - stds = np.where(stds > 1e-6, stds, 1.0) - # Zero out means for features with no observations - means = np.where(counts > 0, means, 0.0) - - return { - "feature_means": means.tolist(), - "feature_stds": stds.tolist(), - "feature_names": feature_names, - "train_indices": train_indices, - "observation_counts": counts.tolist(), - } +from slices.data.preparation import prepare_processed_dataset @hydra.main(version_base=None, config_path="../../configs", config_name="config") def main(cfg: DictConfig) -> None: - """Prepare dataset splits and normalization statistics. - - Usage: - uv run python scripts/preprocessing/prepare_dataset.py dataset=miiv - uv run python scripts/preprocessing/prepare_dataset.py dataset=eicu - """ - print("=" * 70) - print(f"Dataset Preparation — {cfg.dataset}") - print("=" * 70) - - processed_dir = Path(cfg.data.processed_dir) - if not processed_dir.exists(): - raise FileNotFoundError( - f"Processed directory not found: {processed_dir}\n" - f"Run first: uv run python scripts/preprocessing/extract_ricu.py dataset={cfg.dataset}" - ) - - # Load metadata - print("\n1. Loading metadata...") - with open(processed_dir / "metadata.yaml") as f: - metadata = yaml.safe_load(f) - - feature_names = metadata["feature_names"] - seq_length = metadata["seq_length_hours"] - print(f" Features: {len(feature_names)}") - print(f" Sequence length: {seq_length} hours") - - # Load data - print("\n2. Loading parquet files...") - static_df = pl.read_parquet(processed_dir / "static.parquet") - timeseries_df = pl.read_parquet(processed_dir / "timeseries.parquet") - print(f" Loaded {len(timeseries_df):,} stays") - - # Compute splits (ratios are benchmark invariants from constants.py) - print("\n3. Computing patient-level splits...") - splits = compute_patient_splits( - static_df=static_df, - timeseries_df=timeseries_df, - train_ratio=TRAIN_RATIO, - val_ratio=VAL_RATIO, - test_ratio=TEST_RATIO, + """Prepare dataset splits and normalization statistics.""" + prepare_processed_dataset( + processed_dir=Path(cfg.data.processed_dir), seed=cfg.seed, + dataset_name=cfg.dataset, ) - print(f" Train: {splits['train_stays']:,} stays ({len(splits['train_patients']):,} patients)") - print(f" Val: {splits['val_stays']:,} stays ({len(splits['val_patients']):,} patients)") - print(f" Test: {splits['test_stays']:,} stays ({len(splits['test_patients']):,} patients)") - - # Save splits FIRST - splits_path = processed_dir / "splits.yaml" - # Save full patient lists to enable cache validation in datamodule - # (patient lists are required to verify split consistency) - splits_to_save = { - "seed": splits["seed"], - "train_ratio": splits["train_ratio"], - "val_ratio": splits["val_ratio"], - "test_ratio": splits["test_ratio"], - "train_indices": splits["train_indices"], - "val_indices": splits["val_indices"], - "test_indices": splits["test_indices"], - "train_stays": splits["train_stays"], - "val_stays": splits["val_stays"], - "test_stays": splits["test_stays"], - # Full patient lists for datamodule cache validation - "train_patients": splits["train_patients"], - "val_patients": splits["val_patients"], - "test_patients": splits["test_patients"], - } - _atomic_yaml_write(splits_path, splits_to_save) - print(f"\n Saved: {splits_path}") - - # Compute normalization stats on train set only - print("\n4. Computing normalization statistics (train set only)...") - stats = compute_normalization_stats( - timeseries_df=timeseries_df, - train_indices=splits["train_indices"], - feature_names=feature_names, - seq_length=seq_length, - ) - - # Save normalization stats - stats_path = processed_dir / "normalization_stats.yaml" - _atomic_yaml_write(stats_path, stats) - print(f" Saved: {stats_path}") - - # Summary - print("\n" + "=" * 70) - print("Preparation Complete") - print("=" * 70) - print(f"\nOutput directory: {processed_dir}") - print(" - splits.yaml") - print(" - normalization_stats.yaml") - print("\nYou can now run training:") - print(f" uv run python scripts/training/pretrain.py dataset={cfg.dataset} ssl=mae") - if __name__ == "__main__": main() diff --git a/scripts/sanity_checks/sc_mae_pretraining.py b/scripts/sanity_checks/sc_mae_pretraining.py index 33ee8db..21bdbd4 100644 --- a/scripts/sanity_checks/sc_mae_pretraining.py +++ b/scripts/sanity_checks/sc_mae_pretraining.py @@ -15,13 +15,13 @@ Usage: # Default: random masks (generalization test) - uv run python scripts/sanity_checks/sc_mae_pretraining.py data/processed/mimic-iv + uv run python scripts/sanity_checks/sc_mae_pretraining.py data/processed/miiv # Fixed mask (true overfitting test) - uv run python scripts/sanity_checks/sc_mae_pretraining.py data/processed/mimic-iv --fixed-mask + uv run python scripts/sanity_checks/sc_mae_pretraining.py data/processed/miiv --fixed-mask # Test encoder-decoder without SSL masking - uv run python scripts/sanity_checks/sc_mae_pretraining.py data/processed/mimic-iv --autoencoder + uv run python scripts/sanity_checks/sc_mae_pretraining.py data/processed/miiv --autoencoder """ import sys @@ -35,11 +35,12 @@ import torch import torch.nn as nn import yaml +from torch.utils.data import DataLoader, TensorDataset + from slices.data.transforms import create_ssl_mask from slices.models.encoders import build_encoder from slices.models.pretraining import build_ssl_objective from slices.models.pretraining.mae import MAEConfig -from torch.utils.data import DataLoader, TensorDataset SEED = 42 @@ -62,7 +63,7 @@ def load_subset( ) -> Tuple[torch.Tensor, torch.Tensor, int, int]: """Load a small subset using prepared normalization stats. - Applies forward-fill imputation to match the actual pipeline behavior. + Applies normalize-then-zero-fill preprocessing to match the actual pipeline. Args: processed_dir: Path to processed data directory. @@ -115,16 +116,8 @@ def load_subset( ts_tensor[t, f] = float("nan") mask_tensor[t, f] = bool(mask_data[t][f]) - # Forward-fill imputation (matches actual pipeline) - for f in range(n_features): - last_valid = means[f].item() - for t in range(seq_length): - if not torch.isnan(ts_tensor[t, f]): - last_valid = ts_tensor[t, f].item() - else: - ts_tensor[t, f] = last_valid - ts_tensor = (ts_tensor - means) / stds + ts_tensor = torch.nan_to_num(ts_tensor, nan=0.0) timeseries_list.append(ts_tensor) mask_list.append(mask_tensor) @@ -487,13 +480,13 @@ def forward(self, x, obs_mask): epilog=""" Examples: # Test generalization learning (random masks) - uv run python scripts/sanity_checks/sc_mae_pretraining.py data/processed/mimic-iv + uv run python scripts/sanity_checks/sc_mae_pretraining.py data/processed/miiv # Test architecture can overfit (fixed mask) - uv run python scripts/sanity_checks/sc_mae_pretraining.py data/processed/mimic-iv --fixed-mask + uv run python scripts/sanity_checks/sc_mae_pretraining.py data/processed/miiv --fixed-mask # Test encoder-decoder without SSL - uv run python scripts/sanity_checks/sc_mae_pretraining.py data/processed/mimic-iv --autoencoder + uv run python scripts/sanity_checks/sc_mae_pretraining.py data/processed/miiv --autoencoder """, ) parser.add_argument("processed_dir", help="Path to processed data directory") diff --git a/scripts/sanity_checks/sc_supervised_learning.py b/scripts/sanity_checks/sc_supervised_learning.py index 3d80c9a..2f121a9 100644 --- a/scripts/sanity_checks/sc_supervised_learning.py +++ b/scripts/sanity_checks/sc_supervised_learning.py @@ -7,7 +7,7 @@ Run prepare_dataset.py first to generate splits.yaml and normalization_stats.yaml Usage: - uv run python scripts/sanity_checks/sc_supervised_learning.py data/processed/mimic-iv-demo + uv run python scripts/sanity_checks/sc_supervised_learning.py data/processed/miiv """ import sys @@ -19,9 +19,10 @@ import torch import torch.nn as nn import yaml +from torch.utils.data import DataLoader, TensorDataset + from slices.models.encoders import build_encoder from slices.models.heads import TaskHeadConfig, build_task_head -from torch.utils.data import DataLoader, TensorDataset SEED = 42 @@ -193,8 +194,14 @@ def forward(self, x, mask): if __name__ == "__main__": if len(sys.argv) < 2: - print("Usage: uv run python scripts/sanity_check.py ") - print("Example: uv run python scripts/sanity_check.py data/processed/mimic-iv-demo") + print( + "Usage: uv run python scripts/sanity_checks/sc_supervised_learning.py " + "" + ) + print( + "Example: uv run python scripts/sanity_checks/sc_supervised_learning.py " + "data/processed/miiv" + ) sys.exit(1) success = run_sanity_check(sys.argv[1]) diff --git a/scripts/setup_and_extract.sh b/scripts/setup_and_extract.sh index 715b1e7..852d52f 100755 --- a/scripts/setup_and_extract.sh +++ b/scripts/setup_and_extract.sh @@ -5,10 +5,12 @@ # for one or more ICU datasets. Targets Debian/Ubuntu (e.g., GCP VMs). # # Usage: -# ./scripts/setup_and_extract.sh # All datasets (miiv + eicu) +# ./scripts/setup_and_extract.sh # All datasets (miiv + eicu + combined) # ./scripts/setup_and_extract.sh miiv # Single dataset # ./scripts/setup_and_extract.sh miiv eicu # Multiple datasets +# ./scripts/setup_and_extract.sh combined # Build combined + its source datasets # ./scripts/setup_and_extract.sh --skip-deps miiv # Skip dependency installation +# ./scripts/setup_and_extract.sh --force-processed miiv # Regenerate processed artifacts set -euo pipefail @@ -31,21 +33,51 @@ error() { echo -e "${RED}[ERROR]${NC} $1"; } # Parse arguments # --------------------------------------------------------------------------- SKIP_DEPS=false -DATASETS=() +FORCE_PROCESSED=false +EXPECTED_FEATURES="${EXPECTED_FEATURES:-84}" +REQUESTED_DATASETS=() +BASE_DATASETS=() + +append_base_dataset() { + local value="$1" + if [ ${#BASE_DATASETS[@]} -gt 0 ]; then + for existing in "${BASE_DATASETS[@]}"; do + if [ "$existing" = "$value" ]; then + return 0 + fi + done + fi + BASE_DATASETS+=("$value") +} for arg in "$@"; do case "$arg" in --skip-deps) SKIP_DEPS=true ;; - miiv|eicu) DATASETS+=("$arg") ;; + --force-processed) FORCE_PROCESSED=true ;; + miiv|eicu|combined) REQUESTED_DATASETS+=("$arg") ;; *) error "Unknown argument: $arg" - echo "Usage: $0 [--skip-deps] [miiv] [eicu]" + echo "Usage: $0 [--skip-deps] [--force-processed] [miiv] [eicu] [combined]" exit 1 ;; esac done # Default to all datasets if none specified -if [ ${#DATASETS[@]} -eq 0 ]; then - DATASETS=(miiv eicu) +if [ ${#REQUESTED_DATASETS[@]} -eq 0 ]; then + REQUESTED_DATASETS=(miiv eicu combined) +fi + +BUILD_COMBINED=false + +for ds in "${REQUESTED_DATASETS[@]}"; do + case "$ds" in + miiv|eicu) append_base_dataset "$ds" ;; + combined) BUILD_COMBINED=true ;; + esac +done + +if [ "$BUILD_COMBINED" = true ]; then + append_base_dataset "miiv" + append_base_dataset "eicu" fi # Dataset name -> raw data directory mapping (matches extract_with_ricu.R DEFAULT_RAW_PATHS) @@ -57,6 +89,71 @@ raw_dir_for() { esac } +processed_artifacts_exist() { + local processed_dir="$1" + [ -f "$processed_dir/timeseries.parquet" ] && \ + [ -f "$processed_dir/static.parquet" ] && \ + [ -f "$processed_dir/labels.parquet" ] && \ + [ -f "$processed_dir/metadata.yaml" ] +} + +processed_data_ready() { + local processed_dir="$1" + local dataset="$2" + + uv run python - "$processed_dir" "$dataset" "$EXPECTED_FEATURES" <<'PY' +from pathlib import Path +import sys + +import polars as pl +import yaml + +processed_dir = Path(sys.argv[1]) +dataset = sys.argv[2] +expected_features = int(sys.argv[3]) + +with open(processed_dir / "metadata.yaml") as f: + metadata = yaml.safe_load(f) or {} + +feature_names = metadata.get("feature_names") or [] +n_features = metadata.get("n_features", len(feature_names)) +if n_features != expected_features: + raise SystemExit(f"{dataset}: expected {expected_features} features, found {n_features}") + +n_stays = metadata.get("n_stays") +for filename in ("static.parquet", "timeseries.parquet", "labels.parquet"): + height = pl.scan_parquet(processed_dir / filename).select(pl.len()).collect().item() + if n_stays is not None and height != n_stays: + raise SystemExit(f"{dataset}: {filename} has {height} rows, metadata has {n_stays}") + +if dataset != "combined" and not metadata.get("zero_observation_stays_excluded"): + raise SystemExit(f"{dataset}: metadata does not confirm zero-observation exclusion") + +print(f"{dataset}: {n_stays:,} stays, {n_features} features") +PY +} + +reset_processed_artifacts() { + local processed_dir="$1" + + rm -f "$processed_dir/timeseries.parquet" + rm -f "$processed_dir/static.parquet" + rm -f "$processed_dir/labels.parquet" + rm -f "$processed_dir/metadata.yaml" + rm -f "$processed_dir/splits.yaml" + rm -f "$processed_dir/normalization_stats.yaml" + rm -f "$processed_dir"/normalization_stats_*.yaml + rm -f "$processed_dir/dataset_metadata.yaml" + rm -rf "$processed_dir/.tensor_cache" +} + +clear_runtime_caches() { + local processed_dir="$1" + + rm -f "$processed_dir"/normalization_stats_*.yaml + rm -rf "$processed_dir/.tensor_cache" +} + # --------------------------------------------------------------------------- # Ensure we're in the project root # --------------------------------------------------------------------------- @@ -70,7 +167,7 @@ info "Project root: $PROJECT_ROOT" # --------------------------------------------------------------------------- section "Validating raw data" -for ds in "${DATASETS[@]}"; do +for ds in "${BASE_DATASETS[@]}"; do raw_dir="$(raw_dir_for "$ds")" if [ ! -d "$raw_dir" ]; then error "Raw data directory not found: $raw_dir" @@ -128,7 +225,7 @@ if [ "$SKIP_DEPS" = false ]; then # Python dependencies section "Installing Python dependencies" - uv sync --dev + uv sync --dev --frozen else info "Skipping dependency installation (--skip-deps)" fi @@ -136,7 +233,7 @@ fi # --------------------------------------------------------------------------- # Step 2: Run extraction pipeline for each dataset # --------------------------------------------------------------------------- -for ds in "${DATASETS[@]}"; do +for ds in "${BASE_DATASETS[@]}"; do section "Processing dataset: $ds" ricu_output="data/ricu_output/$ds" @@ -154,11 +251,20 @@ for ds in "${DATASETS[@]}"; do fi # --- Python extraction --- - if [ -f "$processed_dir/timeseries.parquet" ] && \ - [ -f "$processed_dir/static.parquet" ] && \ - [ -f "$processed_dir/labels.parquet" ] && \ - [ -f "$processed_dir/metadata.yaml" ]; then - info "Processed data already exists: $processed_dir (skipping Python extraction)" + if processed_artifacts_exist "$processed_dir"; then + if processed_data_ready "$processed_dir" "$ds"; then + info "Processed data already exists and is ${EXPECTED_FEATURES}-feature ready: $processed_dir" + elif [ "$FORCE_PROCESSED" = true ]; then + warn "Processed data is stale; regenerating: $processed_dir" + reset_processed_artifacts "$processed_dir" + info "Running Python extraction for $ds..." + uv run python scripts/preprocessing/extract_ricu.py dataset="$ds" + info "Python extraction complete: $processed_dir" + else + error "Processed data exists but is not launch-ready: $processed_dir" + echo " Re-run with --force-processed to regenerate from RICU output." + exit 1 + fi else info "Running Python extraction for $ds..." uv run python scripts/preprocessing/extract_ricu.py dataset="$ds" @@ -166,24 +272,61 @@ for ds in "${DATASETS[@]}"; do fi # --- Dataset preparation (splits + normalization) --- - if [ -f "$processed_dir/splits.yaml" ] && \ + info "Running dataset preparation for $ds (refreshing split/normalization provenance)..." + uv run python scripts/preprocessing/prepare_dataset.py dataset="$ds" + clear_runtime_caches "$processed_dir" + info "Dataset preparation complete" + + info "Dataset $ds ready!" +done + +if [ "$BUILD_COMBINED" = true ]; then + section "Processing dataset: combined" + + processed_dir="data/processed/combined" + + if processed_artifacts_exist "$processed_dir" && \ + [ -f "$processed_dir/splits.yaml" ] && \ [ -f "$processed_dir/normalization_stats.yaml" ]; then - info "Splits & normalization already exist: $processed_dir (skipping preparation)" + if processed_data_ready "$processed_dir" "combined"; then + info "Combined dataset already exists and is ${EXPECTED_FEATURES}-feature ready: $processed_dir" + info "Refreshing combined split/normalization provenance..." + uv run python scripts/preprocessing/prepare_dataset.py dataset="combined" + clear_runtime_caches "$processed_dir" + elif [ "$FORCE_PROCESSED" = true ]; then + warn "Combined data is stale; regenerating: $processed_dir" + reset_processed_artifacts "$processed_dir" + info "Creating and preparing combined dataset..." + uv run python scripts/preprocessing/create_combined_dataset.py \ + --source data/processed/miiv data/processed/eicu \ + --names miiv eicu \ + --output "$processed_dir" + clear_runtime_caches "$processed_dir" + info "Combined dataset creation complete: $processed_dir" + else + error "Combined data exists but is not launch-ready: $processed_dir" + echo " Re-run with --force-processed to regenerate from source processed datasets." + exit 1 + fi else - info "Running dataset preparation for $ds..." - uv run python scripts/preprocessing/prepare_dataset.py dataset="$ds" - info "Dataset preparation complete" + info "Creating and preparing combined dataset..." + uv run python scripts/preprocessing/create_combined_dataset.py \ + --source data/processed/miiv data/processed/eicu \ + --names miiv eicu \ + --output "$processed_dir" + clear_runtime_caches "$processed_dir" + info "Combined dataset creation complete: $processed_dir" fi - info "Dataset $ds ready!" -done + info "Dataset combined ready!" +fi # --------------------------------------------------------------------------- # Summary # --------------------------------------------------------------------------- section "Pipeline Complete" -for ds in "${DATASETS[@]}"; do +for ds in "${REQUESTED_DATASETS[@]}"; do processed_dir="data/processed/$ds" echo -e " ${GREEN}$ds${NC}: $processed_dir/" for f in timeseries.parquet static.parquet labels.parquet metadata.yaml splits.yaml normalization_stats.yaml; do diff --git a/scripts/training/finetune.py b/scripts/training/finetune.py index 17742b4..e2d191e 100644 --- a/scripts/training/finetune.py +++ b/scripts/training/finetune.py @@ -28,20 +28,88 @@ training.unfreeze_epoch=5 """ +from pathlib import Path + import hydra import lightning.pytorch as pl import torch from omegaconf import DictConfig, OmegaConf + from slices.data.datamodule import ICUDataModule +from slices.eval.fairness_metadata import ( + EVAL_ARTIFACT_PATH_KEY, + EVAL_ARTIFACT_SHA256_KEY, + file_sha256, +) from slices.training import FineTuneModule from slices.training.utils import ( + WandbEntityNotFoundError, + report_and_validate_train_label_support, + resolve_balanced_class_weights, run_fairness_evaluation, setup_finetune_callbacks, setup_wandb_logger, + train_label_support_summary, validate_data_prerequisites, ) +def _evaluated_artifact_metadata( + cfg: DictConfig, + eval_checkpoint_source: str, + best_ckpt: str | None, +) -> dict[str, str]: + """Return path and digest for the artifact used to produce test metrics.""" + if eval_checkpoint_source == "best" and best_ckpt: + artifact_path = Path(best_ckpt) + elif eval_checkpoint_source == "final": + artifact_path = Path(cfg.get("checkpoint_dir", "checkpoints")) / "last.ckpt" + else: + return {EVAL_ARTIFACT_PATH_KEY: "", EVAL_ARTIFACT_SHA256_KEY: ""} + + artifact_sha256 = file_sha256(artifact_path) if artifact_path.exists() else "" + return { + EVAL_ARTIFACT_PATH_KEY: str(artifact_path), + EVAL_ARTIFACT_SHA256_KEY: artifact_sha256, + } + + +def _detect_paradigm_from_checkpoint(path: str, *, full_checkpoint: bool) -> str | None: + """Infer the SSL paradigm recorded in an encoder or Lightning checkpoint.""" + checkpoint = torch.load( + path, + map_location="cpu", + weights_only=not full_checkpoint, + ) + + if not isinstance(checkpoint, dict): + return None + + if "ssl_name" in checkpoint: + return checkpoint["ssl_name"] + + if not full_checkpoint: + return None + + hyper_parameters = checkpoint.get("hyper_parameters") or {} + config = hyper_parameters.get("config") or {} + if isinstance(config, DictConfig): + config = OmegaConf.to_container(config, resolve=True) + + if not isinstance(config, dict): + return None + + ssl_config = config.get("ssl") or {} + if isinstance(ssl_config, DictConfig): + ssl_config = OmegaConf.to_container(ssl_config, resolve=True) + + if isinstance(ssl_config, dict) and ssl_config.get("name") is not None: + return str(ssl_config["name"]) + + paradigm = config.get("paradigm") + return str(paradigm) if paradigm is not None else None + + @hydra.main(version_base=None, config_path="../../configs", config_name="finetune") def main(cfg: DictConfig) -> None: """Run downstream task finetuning.""" @@ -52,27 +120,45 @@ def main(cfg: DictConfig) -> None: print("\nConfiguration:") print(OmegaConf.to_yaml(cfg)) - # Validate checkpoint + # Validate checkpoint source. + if cfg.checkpoint is not None and cfg.pretrain_checkpoint is not None: + raise ValueError( + "Provide exactly one checkpoint source: use 'checkpoint' for encoder.pt " + "or 'pretrain_checkpoint' for a full Lightning pretrain .ckpt, not both." + ) if cfg.checkpoint is None and cfg.pretrain_checkpoint is None: raise ValueError( "Must provide either 'checkpoint' (encoder.pt) or " "'pretrain_checkpoint' (full .ckpt file)" ) - # Auto-detect paradigm from encoder checkpoint metadata + # Auto-detect paradigm from checkpoint metadata + detected_paradigm = None if cfg.checkpoint is not None: - ckpt = torch.load(cfg.checkpoint, map_location="cpu", weights_only=True) - if isinstance(ckpt, dict) and "ssl_name" in ckpt: - detected = ckpt["ssl_name"] - if cfg.paradigm != detected: - print(f"\n Auto-detected paradigm from checkpoint: {detected}") - OmegaConf.set_struct(cfg, False) - cfg.paradigm = detected - OmegaConf.set_struct(cfg, True) - del ckpt + detected_paradigm = _detect_paradigm_from_checkpoint( + cfg.checkpoint, + full_checkpoint=False, + ) + elif cfg.pretrain_checkpoint is not None: + detected_paradigm = _detect_paradigm_from_checkpoint( + cfg.pretrain_checkpoint, + full_checkpoint=True, + ) + + if detected_paradigm and cfg.paradigm != detected_paradigm: + print(f"\n Auto-detected paradigm from checkpoint: {detected_paradigm}") + OmegaConf.set_struct(cfg, False) + cfg.paradigm = detected_paradigm + OmegaConf.set_struct(cfg, True) # Validate data prerequisites - validate_data_prerequisites(cfg.data.processed_dir, cfg.dataset) + task_name = cfg.task.get("task_name", "mortality_24h") + validate_data_prerequisites( + cfg.data.processed_dir, + cfg.dataset, + task_names=[task_name], + task_configs=[cfg.task], + ) # Set random seed for reproducibility pl.seed_everything(cfg.seed, workers=True) @@ -84,8 +170,6 @@ def main(cfg: DictConfig) -> None: print("1. Setting up DataModule") print("=" * 80) - task_name = cfg.task.get("task_name", "mortality_24h") - datamodule = ICUDataModule( processed_dir=cfg.data.processed_dir, task_name=task_name, @@ -114,18 +198,36 @@ def main(cfg: DictConfig) -> None: ) label_stats = datamodule.get_label_statistics() + train_label_stats = datamodule.get_train_label_statistics() + task_type = cfg.task.get("task_type", "binary") if task_name in label_stats: stats = label_stats[task_name] print(f"\n Label distribution for '{task_name}':") print(f" - Total samples: {stats['total']}") - print( - f" - Positive: {stats.get('positive', 'N/A')} " - f"({stats.get('prevalence', 0)*100:.1f}%)" - ) - print( - f" - Negative: {stats.get('negative', 'N/A')} " - f"({(1 - stats.get('prevalence', 0))*100:.1f}%)" - ) + if stats.get("task_type") == "regression": + print(f" - Mean: {stats.get('mean', 0.0):.4f}") + print(f" - Std: {stats.get('std', 0.0):.4f}") + print(f" - Min: {stats.get('min', 0.0):.4f}") + print(f" - Max: {stats.get('max', 0.0):.4f}") + else: + print( + f" - Positive: {stats.get('positive', 'N/A')} " + f"({stats.get('prevalence', 0)*100:.1f}%)" + ) + print( + f" - Negative: {stats.get('negative', 'N/A')} " + f"({(1 - stats.get('prevalence', 0))*100:.1f}%)" + ) + + train_support_stats = report_and_validate_train_label_support( + datamodule=datamodule, + task_name=task_name, + task_type=task_type, + dataset=cfg.dataset, + seed=cfg.seed, + label_fraction=cfg.get("label_fraction", 1.0), + min_train_positives=cfg.get("min_train_positives", 3), + ) # ========================================================================= # 2. Create Finetune Module @@ -140,15 +242,24 @@ def main(cfg: DictConfig) -> None: # Resolve "balanced" class weights from label distribution if cfg.training.get("class_weight") == "balanced": - if task_name in label_stats: - stats = label_stats[task_name] - n_pos = stats.get("positive", 1) - n_neg = stats.get("negative", 1) - n_total = n_pos + n_neg - cfg.training.class_weight = [n_total / (2 * n_neg), n_total / (2 * n_pos)] - print(f"\n Balanced class weights: {cfg.training.class_weight}") + resolved_class_weight = resolve_balanced_class_weights( + task_name=task_name, + task_type=task_type, + train_label_stats=train_label_stats, + ) + if resolved_class_weight is not None: + cfg.training.class_weight = resolved_class_weight + n_total = int(train_label_stats[task_name]["total"]) + print(f"\n Training-split labels used for class weighting: {n_total}") + print(f"\n sqrt(balanced) class weights: {cfg.training.class_weight}") else: - print(f"\n Warning: No label stats for '{task_name}', skipping class weighting") + if task_type == "regression": + print(f"\n class_weight='balanced' ignored for regression task '{task_name}'") + else: + print( + f"\n Warning: No train-split label stats for '{task_name}', " + "skipping class weighting" + ) cfg.training.class_weight = None OmegaConf.set_struct(cfg, True) @@ -180,7 +291,13 @@ def main(cfg: DictConfig) -> None: print("=" * 80) callbacks = setup_finetune_callbacks(cfg, checkpoint_prefix="finetune") - logger = setup_wandb_logger(cfg) + try: + logger = setup_wandb_logger(cfg) + except WandbEntityNotFoundError as exc: + print(f"\nError: {exc}") + raise SystemExit(1) from exc + if logger: + logger.experiment.summary.update(train_label_support_summary(train_support_stats)) trainer = pl.Trainer( max_epochs=cfg.training.max_epochs, @@ -214,7 +331,10 @@ def main(cfg: DictConfig) -> None: print("4. Starting Training") print("=" * 80) - trainer.fit(model, datamodule=datamodule) + ckpt_path = cfg.get("ckpt_path", None) + if ckpt_path: + print(f" Resuming trainer state from: {ckpt_path}") + trainer.fit(model, datamodule=datamodule, ckpt_path=ckpt_path) # ========================================================================= # 5. Test @@ -224,6 +344,9 @@ def main(cfg: DictConfig) -> None: print("=" * 80) best_ckpt = callbacks[0].best_model_path if hasattr(callbacks[0], "best_model_path") else None + eval_checkpoint_source = "none" + best_ckpt_load_ok = False + best_ckpt_error = None if best_ckpt: print(f"\n Best checkpoint: {best_ckpt}") @@ -231,11 +354,37 @@ def main(cfg: DictConfig) -> None: checkpoint = torch.load(best_ckpt, map_location="cpu", weights_only=False) model.load_state_dict(checkpoint["state_dict"]) print(" - Loaded best checkpoint weights") + eval_checkpoint_source = "best" + best_ckpt_load_ok = True except Exception as e: - print(f" - Warning: Could not load checkpoint ({e}), using final model") + best_ckpt_error = str(e) + allow_fallback = cfg.training.get("allow_best_ckpt_fallback", False) + if allow_fallback: + print( + f" - WARNING: Could not load best checkpoint ({e})," + " falling back to final model" + ) + eval_checkpoint_source = "final" + else: + # Log failure to W&B before crashing + if logger: + logger.experiment.summary.update( + { + "_eval_checkpoint_source": "failed", + "_best_ckpt_path": best_ckpt, + "_best_ckpt_load_ok": False, + "_best_ckpt_error": best_ckpt_error, + } + ) + logger.experiment.finish() + raise RuntimeError( + f"Best checkpoint load failed: {e}. " + f"Set training.allow_best_ckpt_fallback=true to use final model instead." + ) from e test_results = trainer.test(model, datamodule=datamodule) else: print("\n No best checkpoint found, testing with final model") + eval_checkpoint_source = "final" test_results = trainer.test(model, datamodule=datamodule) # ========================================================================= @@ -261,6 +410,15 @@ def main(cfg: DictConfig) -> None: if logger: logger.experiment.summary.update(test_results[0]) + logger.experiment.summary.update( + { + "_eval_checkpoint_source": eval_checkpoint_source, + **_evaluated_artifact_metadata(cfg, eval_checkpoint_source, best_ckpt), + "_best_ckpt_path": best_ckpt or "", + "_best_ckpt_load_ok": best_ckpt_load_ok, + "_best_ckpt_error": best_ckpt_error or "", + } + ) print(f"\n Output directory: {cfg.output_dir}") print(f" - Checkpoints: {cfg.get('checkpoint_dir', 'checkpoints')}") diff --git a/scripts/training/pretrain.py b/scripts/training/pretrain.py index 03c91ca..075b946 100644 --- a/scripts/training/pretrain.py +++ b/scripts/training/pretrain.py @@ -17,11 +17,14 @@ import hydra import lightning.pytorch as pl +import torch from omegaconf import DictConfig, OmegaConf + from slices.data.config_schemas import DataConfig from slices.data.datamodule import ICUDataModule from slices.training import SSLPretrainModule from slices.training.utils import ( + WandbEntityNotFoundError, setup_pretrain_callbacks, setup_wandb_logger, validate_data_prerequisites, @@ -29,9 +32,10 @@ # SSL -> compatible model encoder mappings _SSL_MODEL_COMPAT = { - "mae": {"observation_transformer"}, - "jepa": {"observation_transformer"}, - "contrastive": {"observation_transformer"}, + "mae": {"transformer"}, + "jepa": {"transformer"}, + "contrastive": {"transformer"}, + "ts2vec": {"transformer"}, "smart": {"smart"}, } @@ -149,7 +153,11 @@ def main(cfg: DictConfig) -> None: print("=" * 80) callbacks = setup_pretrain_callbacks(cfg) - logger = setup_wandb_logger(cfg) + try: + logger = setup_wandb_logger(cfg) + except WandbEntityNotFoundError as exc: + print(f"\nError: {exc}") + raise SystemExit(1) from exc trainer = pl.Trainer( max_epochs=cfg.training.max_epochs, @@ -184,28 +192,39 @@ def main(cfg: DictConfig) -> None: print("4. Starting Training") print("=" * 80) - trainer.fit(model, datamodule=datamodule) + ckpt_path = cfg.get("ckpt_path", None) + if ckpt_path: + print(f" Resuming trainer state from: {ckpt_path}") + trainer.fit(model, datamodule=datamodule, ckpt_path=ckpt_path) # ========================================================================= - # 5. Save Encoder (from best checkpoint) + # 5. Save Encoder # ========================================================================= print("\n" + "=" * 80) print("5. Saving Encoder") print("=" * 80) + output_dir = Path(cfg.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + # Primary: save encoder from last epoch (model state after trainer.fit) + # Fixed-schedule training follows I-JEPA/DINO/BYOL convention. + encoder_path = output_dir / "encoder.pt" + model.save_encoder(str(encoder_path)) + print(f"\n Encoder (last epoch) saved to: {encoder_path}") + print(" - Use this for downstream fine-tuning") + + # Secondary: save encoder from best-val-loss checkpoint (robustness check) checkpoint_callback = callbacks[0] best_ckpt = checkpoint_callback.best_model_path if best_ckpt: - print(f"\n Loading best checkpoint: {best_ckpt}") + print(f"\n Best val/loss checkpoint: {best_ckpt}") print(f" Best val/loss: {checkpoint_callback.best_model_score:.4f}") - model = SSLPretrainModule.load_from_checkpoint(best_ckpt) - - encoder_path = Path(cfg.output_dir) / "encoder.pt" - encoder_path.parent.mkdir(parents=True, exist_ok=True) - model.save_encoder(str(encoder_path)) - - print(f"\n Encoder saved to: {encoder_path}") - print(" - Use this for downstream fine-tuning") + best_state = torch.load(best_ckpt, map_location="cpu", weights_only=False) + model.load_state_dict(best_state["state_dict"]) + encoder_best_path = output_dir / "encoder_best_val.pt" + model.save_encoder(str(encoder_best_path)) + print(f" Encoder (best val) saved to: {encoder_best_path}") # ========================================================================= # Summary @@ -214,13 +233,11 @@ def main(cfg: DictConfig) -> None: print("Training Complete!") print("=" * 80) - if best_ckpt: - print(f"\n Best checkpoint: {best_ckpt}") - print(f" - Best val/loss: {checkpoint_callback.best_model_score:.4f}") - print(f"\n Output directory: {cfg.output_dir}") print(f" - Checkpoints: {cfg.get('checkpoint_dir', 'checkpoints')}") - print(f" - Encoder weights: {encoder_path}") + print(f" - Encoder (last epoch): {encoder_path}") + if best_ckpt: + print(f" - Encoder (best val): {output_dir / 'encoder_best_val.pt'}") if logger: print( diff --git a/scripts/training/supervised.py b/scripts/training/supervised.py index 061d337..a8cd6ce 100644 --- a/scripts/training/supervised.py +++ b/scripts/training/supervised.py @@ -30,60 +30,110 @@ import lightning.pytorch as pl import torch from omegaconf import DictConfig, OmegaConf + from slices.data.datamodule import ICUDataModule +from slices.eval.fairness_metadata import ( + EVAL_ARTIFACT_PATH_KEY, + EVAL_ARTIFACT_SHA256_KEY, + file_sha256, +) from slices.models.encoders import EncoderWithMissingToken from slices.training import FineTuneModule from slices.training.utils import ( + WandbEntityNotFoundError, + report_and_validate_train_label_support, + resolve_balanced_class_weights, run_fairness_evaluation, save_encoder_checkpoint, setup_finetune_callbacks, setup_wandb_logger, + train_label_support_summary, validate_data_prerequisites, ) +def _evaluated_artifact_metadata( + cfg: DictConfig, + eval_checkpoint_source: str, + best_ckpt: str | None, +) -> dict[str, str]: + """Return path and digest for the artifact used to produce test metrics.""" + if eval_checkpoint_source == "best" and best_ckpt: + artifact_path = Path(best_ckpt) + elif eval_checkpoint_source == "final": + artifact_path = Path(cfg.get("checkpoint_dir", "checkpoints")) / "last.ckpt" + else: + return {EVAL_ARTIFACT_PATH_KEY: "", EVAL_ARTIFACT_SHA256_KEY: ""} + + artifact_sha256 = file_sha256(artifact_path) if artifact_path.exists() else "" + return { + EVAL_ARTIFACT_PATH_KEY: str(artifact_path), + EVAL_ARTIFACT_SHA256_KEY: artifact_sha256, + } + + +def _collect_dataset_labels(dataset) -> torch.Tensor: + """Collect labels from a dataset into a flat tensor.""" + labels = [] + for i in range(len(dataset)): + sample = dataset[i] + if "label" not in sample: + continue + labels.append(torch.as_tensor(sample["label"], dtype=torch.float32).reshape(-1)) + + if not labels: + return torch.empty(0, dtype=torch.float32) + + return torch.cat(labels, dim=0) + + def compute_baseline_metrics(datamodule, task_name: str, task_type: str = "binary") -> dict: """Compute baseline metrics for comparison. + Baselines are fit on the train split and evaluated on the test split. + For classification tasks, computes AUROC for random predictions and - majority-class predictions. For regression tasks, computes mean/median - predictor baselines. + majority-class predictions. For regression tasks, computes train-fit + mean/median predictor baselines on the test labels. """ + del task_name baselines = {} + train_dataset = datamodule.train_dataloader().dataset test_dataset = datamodule.test_dataloader().dataset - labels = [] - for i in range(len(test_dataset)): - sample = test_dataset[i] - if "label" in sample: - labels.append(sample["label"]) - if not labels: + train_labels = _collect_dataset_labels(train_dataset) + test_labels = _collect_dataset_labels(test_dataset) + + if len(train_labels) == 0 or len(test_labels) == 0: return baselines - labels_tensor = torch.tensor(labels) - n_samples = len(labels_tensor) + n_samples = len(test_labels) baselines["test/n_samples"] = n_samples if task_type == "regression": - mean_label = labels_tensor.mean().item() - median_label = labels_tensor.median().item() - baselines["test/label_mean"] = mean_label - baselines["test/label_std"] = labels_tensor.std().item() - baselines["baseline/mean_predictor_mse"] = (labels_tensor - mean_label).pow(2).mean().item() + mean_label = train_labels.mean().item() + median_label = train_labels.median().item() + baselines["baseline/train_label_mean"] = mean_label + baselines["baseline/train_label_std"] = train_labels.std(unbiased=False).item() + baselines["baseline/train_label_median"] = median_label + baselines["baseline/mean_predictor_mse"] = (test_labels - mean_label).pow(2).mean().item() baselines["baseline/median_predictor_mae"] = ( - (labels_tensor - median_label).abs().mean().item() + (test_labels - median_label).abs().mean().item() ) else: - n_positive = labels_tensor.sum().item() + n_positive = test_labels.sum().item() n_negative = n_samples - n_positive positive_ratio = n_positive / n_samples + train_positive_ratio = train_labels.mean().item() + majority_class = 1.0 if train_positive_ratio >= 0.5 else 0.0 baselines["test/n_positive"] = n_positive baselines["test/n_negative"] = n_negative baselines["test/positive_ratio"] = positive_ratio + baselines["baseline/train_positive_ratio"] = train_positive_ratio baselines["baseline/random_auroc"] = 0.5 - majority_class_accuracy = max(positive_ratio, 1 - positive_ratio) + majority_class_accuracy = (test_labels == majority_class).float().mean().item() baselines["baseline/majority_accuracy"] = majority_class_accuracy baselines["baseline/trivial_auroc"] = 0.5 @@ -138,7 +188,13 @@ def main(cfg: DictConfig) -> None: OmegaConf.set_struct(cfg, True) # Validate data prerequisites - validate_data_prerequisites(cfg.data.processed_dir, cfg.dataset) + task_name_for_validation = cfg.task.get("task_name", "mortality_24h") + validate_data_prerequisites( + cfg.data.processed_dir, + cfg.dataset, + task_names=[task_name_for_validation], + task_configs=[cfg.task], + ) # Set random seed for reproducibility pl.seed_everything(cfg.seed, workers=True) @@ -180,18 +236,36 @@ def main(cfg: DictConfig) -> None: ) label_stats = datamodule.get_label_statistics() + train_label_stats = datamodule.get_train_label_statistics() + task_type = cfg.task.get("task_type", "binary") if task_name in label_stats: stats = label_stats[task_name] print(f"\n Label distribution for '{task_name}':") print(f" - Total samples: {stats['total']}") - print( - f" - Positive: {stats.get('positive', 'N/A')} " - f"({stats.get('prevalence', 0)*100:.1f}%)" - ) - print( - f" - Negative: {stats.get('negative', 'N/A')} " - f"({(1 - stats.get('prevalence', 0))*100:.1f}%)" - ) + if stats.get("task_type") == "regression": + print(f" - Mean: {stats.get('mean', 0.0):.4f}") + print(f" - Std: {stats.get('std', 0.0):.4f}") + print(f" - Min: {stats.get('min', 0.0):.4f}") + print(f" - Max: {stats.get('max', 0.0):.4f}") + else: + print( + f" - Positive: {stats.get('positive', 'N/A')} " + f"({stats.get('prevalence', 0)*100:.1f}%)" + ) + print( + f" - Negative: {stats.get('negative', 'N/A')} " + f"({(1 - stats.get('prevalence', 0))*100:.1f}%)" + ) + + train_support_stats = report_and_validate_train_label_support( + datamodule=datamodule, + task_name=task_name, + task_type=task_type, + dataset=cfg.dataset, + seed=cfg.seed, + label_fraction=cfg.get("label_fraction", 1.0), + min_train_positives=cfg.get("min_train_positives", 3), + ) # ========================================================================= # 2. Create Model (Training from Scratch) @@ -206,15 +280,24 @@ def main(cfg: DictConfig) -> None: # Resolve "balanced" class weights from label distribution if cfg.training.get("class_weight") == "balanced": - if task_name in label_stats: - stats = label_stats[task_name] - n_pos = stats.get("positive", 1) - n_neg = stats.get("negative", 1) - n_total = n_pos + n_neg - cfg.training.class_weight = [n_total / (2 * n_neg), n_total / (2 * n_pos)] - print(f"\n Balanced class weights: {cfg.training.class_weight}") + resolved_class_weight = resolve_balanced_class_weights( + task_name=task_name, + task_type=task_type, + train_label_stats=train_label_stats, + ) + if resolved_class_weight is not None: + cfg.training.class_weight = resolved_class_weight + n_total = int(train_label_stats[task_name]["total"]) + print(f"\n Training-split labels used for class weighting: {n_total}") + print(f"\n sqrt(balanced) class weights: {cfg.training.class_weight}") else: - print(f"\n Warning: No label stats for '{task_name}', skipping class weighting") + if task_type == "regression": + print(f"\n class_weight='balanced' ignored for regression task '{task_name}'") + else: + print( + f"\n Warning: No train-split label stats for '{task_name}', " + "skipping class weighting" + ) cfg.training.class_weight = None OmegaConf.set_struct(cfg, True) @@ -244,7 +327,13 @@ def main(cfg: DictConfig) -> None: print("=" * 80) callbacks = setup_finetune_callbacks(cfg, checkpoint_prefix="supervised") - logger = setup_wandb_logger(cfg) + try: + logger = setup_wandb_logger(cfg) + except WandbEntityNotFoundError as exc: + print(f"\nError: {exc}") + raise SystemExit(1) from exc + if logger: + logger.experiment.summary.update(train_label_support_summary(train_support_stats)) trainer = pl.Trainer( max_epochs=cfg.training.max_epochs, @@ -293,6 +382,9 @@ def main(cfg: DictConfig) -> None: print("=" * 80) best_ckpt = callbacks[0].best_model_path if hasattr(callbacks[0], "best_model_path") else None + eval_checkpoint_source = "none" + best_ckpt_load_ok = False + best_ckpt_error = None if best_ckpt: print(f"\n Loading best checkpoint: {best_ckpt}") @@ -300,8 +392,34 @@ def main(cfg: DictConfig) -> None: checkpoint = torch.load(best_ckpt, map_location="cpu", weights_only=False) model.load_state_dict(checkpoint["state_dict"]) print(" - Loaded best checkpoint weights") + eval_checkpoint_source = "best" + best_ckpt_load_ok = True except Exception as e: - print(f" - Warning: Could not load checkpoint ({e}), using final model") + best_ckpt_error = str(e) + allow_fallback = cfg.training.get("allow_best_ckpt_fallback", False) + if allow_fallback: + print( + f" - WARNING: Could not load best checkpoint ({e})," + " falling back to final model" + ) + eval_checkpoint_source = "final" + else: + if logger: + logger.experiment.summary.update( + { + "_eval_checkpoint_source": "failed", + "_best_ckpt_path": best_ckpt, + "_best_ckpt_load_ok": False, + "_best_ckpt_error": best_ckpt_error, + } + ) + logger.experiment.finish() + raise RuntimeError( + f"Best checkpoint load failed: {e}. " + f"Set training.allow_best_ckpt_fallback=true to use final model instead." + ) from e + else: + eval_checkpoint_source = "final" encoder_path = save_encoder_weights(model, cfg, cfg.output_dir) print(f"\n Encoder saved to: {encoder_path}") @@ -386,10 +504,19 @@ def main(cfg: DictConfig) -> None: else: print(" -> WARNING: Model performs at or below random baseline!") - # Log test results and baseline metrics to W&B summary + # Log test results, baseline metrics, and checkpoint provenance to W&B summary if logger: if test_results: logger.experiment.summary.update(test_results[0]) + logger.experiment.summary.update( + { + "_eval_checkpoint_source": eval_checkpoint_source, + **_evaluated_artifact_metadata(cfg, eval_checkpoint_source, best_ckpt), + "_best_ckpt_path": best_ckpt or "", + "_best_ckpt_load_ok": best_ckpt_load_ok, + "_best_ckpt_error": best_ckpt_error or "", + } + ) if baseline_metrics: logger.experiment.summary.update(baseline_metrics) diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py new file mode 100644 index 0000000..7b6c1e8 --- /dev/null +++ b/scripts/training/xgboost_baseline.py @@ -0,0 +1,423 @@ +"""XGBoost baseline for ICU prediction tasks. + +Non-neural baseline that extracts hand-crafted tabular features from ICU time +series and trains XGBoost. Reuses ICUDataModule for data loading and splits, +and logs metrics to W&B in the same format as neural baselines. + +Example usage: + uv run python scripts/training/xgboost_baseline.py dataset=miiv + uv run python scripts/training/xgboost_baseline.py dataset=eicu tasks=mortality_hospital + uv run python scripts/training/xgboost_baseline.py dataset=miiv logging.use_wandb=false +""" + +from pathlib import Path + +import hydra +import numpy as np +import torch +from omegaconf import DictConfig, OmegaConf +from sklearn.metrics import ( + accuracy_score, + average_precision_score, + brier_score_loss, + confusion_matrix, + f1_score, + mean_absolute_error, + mean_squared_error, + precision_score, + r2_score, + recall_score, + roc_auc_score, +) +from xgboost import XGBClassifier, XGBRegressor + +from slices.constants import canonical_downstream_protocol +from slices.data.datamodule import ICUDataModule +from slices.eval.fairness_metadata import ( + EVAL_ARTIFACT_PATH_KEY, + EVAL_ARTIFACT_SHA256_KEY, + file_sha256, +) +from slices.eval.inference import extract_tabular_features + + +def _xgboost_eval_metric(task_type: str) -> str: + """Return the eval metric aligned with the benchmark primary metric.""" + return "mae" if task_type == "regression" else "aucpr" + + +def _add_wandb_tag(tags: list[str], tag: str | None) -> None: + """Append a W&B tag once, ignoring unset values.""" + if not tag: + return + if tag not in tags: + tags.append(tag) + + +def _resolve_scale_pos_weight(scale_pos_weight, y_train) -> float | None: + """Resolve XGBoost's native positive-class weight. + + ``balanced`` intentionally follows XGBoost's standard binary weighting: + ``n_negative / n_positive``. This differs from the neural sqrt-balanced + class weights, whose effective positive/negative ratio is milder. + """ + if scale_pos_weight is None: + return None + if scale_pos_weight == "balanced": + y_train = np.asarray(y_train) + n_pos = float(y_train.sum()) + n_neg = float(len(y_train) - n_pos) + return n_neg / max(n_pos, 1.0) + return float(scale_pos_weight) + + +def _binary_ece(y_true, y_pred_proba, n_bins: int = 15) -> float: + """Compute L1 expected calibration error with uniform probability bins.""" + y_true = np.asarray(y_true, dtype=float) + y_pred_proba = np.asarray(y_pred_proba, dtype=float) + if y_true.size == 0: + return float("nan") + + bin_edges = np.linspace(0.0, 1.0, n_bins + 1) + bin_ids = np.digitize(y_pred_proba, bin_edges[1:-1], right=False) + + ece = 0.0 + for bin_idx in range(n_bins): + in_bin = bin_ids == bin_idx + if not np.any(in_bin): + continue + bin_weight = in_bin.mean() + confidence = y_pred_proba[in_bin].mean() + accuracy = y_true[in_bin].mean() + ece += bin_weight * abs(confidence - accuracy) + return float(ece) + + +def _binary_classification_metrics( + y_true, + y_pred_proba, + threshold: float = 0.5, +) -> dict[str, float]: + """Compute binary publication metrics for XGBoost test predictions.""" + y_pred_class = (np.asarray(y_pred_proba) >= threshold).astype(int) + tn, fp, _, _ = confusion_matrix(y_true, y_pred_class, labels=[0, 1]).ravel() + specificity = tn / (tn + fp) if (tn + fp) > 0 else 0.0 + + return { + "test/auroc": roc_auc_score(y_true, y_pred_proba), + "test/auprc": average_precision_score(y_true, y_pred_proba), + "test/accuracy": accuracy_score(y_true, y_pred_class), + "test/f1": f1_score(y_true, y_pred_class, zero_division=0), + "test/precision": precision_score(y_true, y_pred_class, zero_division=0), + "test/recall": recall_score(y_true, y_pred_class, zero_division=0), + "test/specificity": specificity, + "test/brier_score": brier_score_loss(y_true, y_pred_proba), + "test/ece": _binary_ece(y_true, y_pred_proba), + } + + +def _build_wandb_tags(cfg: DictConfig) -> list[str] | None: + """Build W&B tags in parity with neural training runs.""" + tags = list(cfg.logging.get("wandb_tags", [])) + if cfg.get("experiment_class") is not None: + _add_wandb_tag(tags, f"experiment_class:{cfg.experiment_class}") + if cfg.get("experiment_subtype") is not None: + _add_wandb_tag(tags, f"experiment_subtype:{cfg.experiment_subtype}") + if cfg.get("revision") is not None: + _add_wandb_tag(tags, f"revision:{cfg.revision}") + if cfg.get("rerun_reason") is not None: + tag = f"rerun-reason:{cfg.rerun_reason}" + if len(tag) > 64: + tag = tag[:61] + "..." + _add_wandb_tag(tags, tag) + if cfg.get("launch_commit") is not None: + _add_wandb_tag(tags, f"commit:{str(cfg.launch_commit)[:12]}") + if cfg.get("phase") is not None: + _add_wandb_tag(tags, f"phase:{cfg.phase}") + if cfg.get("dataset") is not None: + _add_wandb_tag(tags, f"dataset:{cfg.dataset}") + if cfg.get("paradigm") is not None: + _add_wandb_tag(tags, f"paradigm:{cfg.paradigm}") + task_name = cfg.get("task", {}).get("task_name") + if task_name is not None: + _add_wandb_tag(tags, f"task:{task_name}") + if cfg.get("seed") is not None: + _add_wandb_tag(tags, f"seed:{cfg.seed}") + if cfg.get("protocol") is not None: + _add_wandb_tag(tags, f"protocol:{canonical_downstream_protocol(cfg.protocol)}") + if cfg.get("label_fraction", 1.0) < 1.0: + _add_wandb_tag(tags, f"label_fraction:{cfg.label_fraction}") + _add_wandb_tag(tags, "ablation:label-efficiency") + return tags or None + + +def _wandb_run_name_and_group(cfg: DictConfig, task_name: str) -> tuple[str, str]: + run_name = cfg.logging.get("run_name", f"xgboost_{cfg.dataset}_{task_name}") + group_name = cfg.logging.get("wandb_group", f"xgboost_{cfg.dataset}_{task_name}") + if cfg.get("label_fraction", 1.0) < 1.0: + frac_str = str(cfg.label_fraction).replace(".", "") + run_name += f"_frac{frac_str}" + group_name += f"_frac{frac_str}" + return run_name, group_name + + +def _init_wandb_run(cfg: DictConfig, task_name: str): + """Initialize W&B before expensive baseline setup so early failures are visible.""" + if not cfg.logging.get("use_wandb", False): + return None, None + + import wandb + + run_name, group_name = _wandb_run_name_and_group(cfg, task_name) + wandb_entity = cfg.logging.get("wandb_entity", None) + try: + run = wandb.init( + project=cfg.logging.wandb_project, + entity=wandb_entity, + name=run_name, + group=group_name, + tags=_build_wandb_tags(cfg), + config=OmegaConf.to_container(cfg, resolve=True), + ) + except wandb.errors.CommError as exc: + message = str(exc) + if wandb_entity and "entity" in message and "not found" in message: + from slices.training.utils import WandbEntityNotFoundError + + raise WandbEntityNotFoundError( + f"W&B entity '{wandb_entity}' could not be found. " + "Remove logging.wandb_entity to use the signed-in account, " + "or provide a valid W&B username/team." + ) from None + raise + return wandb, run + + +@hydra.main(version_base=None, config_path="../../configs", config_name="xgboost") +def main(cfg: DictConfig) -> None: + """Train XGBoost baseline.""" + print("=" * 80) + print("XGBoost Baseline") + print("=" * 80) + + print("\nConfiguration:") + print(OmegaConf.to_yaml(cfg)) + + # ========================================================================= + # 1. Setup DataModule + # ========================================================================= + print("\n" + "=" * 80) + print("1. Setting up DataModule") + print("=" * 80) + + from slices.training.utils import ( + WandbEntityNotFoundError, + report_and_validate_train_label_support, + train_label_support_summary, + validate_data_prerequisites, + ) + + task_name = cfg.task.get("task_name", "mortality_24h") + task_type = cfg.task.get("task_type", "binary") + try: + wandb_module, wandb_run = _init_wandb_run(cfg, task_name) + except WandbEntityNotFoundError as exc: + print(f"\nError: {exc}") + raise SystemExit(1) from exc + + # Validate data prerequisites including label freshness + validate_data_prerequisites( + cfg.data.processed_dir, + cfg.dataset, + task_names=[task_name], + task_configs=[cfg.task], + ) + + datamodule = ICUDataModule( + processed_dir=cfg.data.processed_dir, + task_name=task_name, + batch_size=cfg.training.batch_size, + num_workers=0, # feature extraction is single-threaded + seed=cfg.seed, + label_fraction=cfg.get("label_fraction", 1.0), + ) + datamodule.setup() + + print(f"\n Task: {task_name} ({task_type})") + print(f" Train: {len(datamodule.train_indices)} stays") + print(f" Val: {len(datamodule.val_indices)} stays") + print(f" Test: {len(datamodule.test_indices)} stays") + + train_support_stats = report_and_validate_train_label_support( + datamodule=datamodule, + task_name=task_name, + task_type=task_type, + dataset=cfg.dataset, + seed=cfg.seed, + label_fraction=cfg.get("label_fraction", 1.0), + min_train_positives=cfg.get("min_train_positives", 3), + ) + + # ========================================================================= + # 2. Extract Tabular Features + # ========================================================================= + print("\n" + "=" * 80) + print("2. Extracting tabular features") + print("=" * 80) + + X_train, y_train = extract_tabular_features(datamodule.dataset, datamodule.train_indices) + X_val, y_val = extract_tabular_features(datamodule.dataset, datamodule.val_indices) + X_test, y_test = extract_tabular_features(datamodule.dataset, datamodule.test_indices) + + print(f" Train: {X_train.shape}") + print(f" Val: {X_val.shape}") + print(f" Test: {X_test.shape}") + + # ========================================================================= + # 3. Train XGBoost + # ========================================================================= + print("\n" + "=" * 80) + print("3. Training XGBoost") + print("=" * 80) + + xgb_cfg = cfg.xgboost + early_stopping_rounds = xgb_cfg.get("early_stopping_rounds", 20) + n_jobs = xgb_cfg.get("n_jobs", 4) + if n_jobs is None: + n_jobs = 4 + common_params = { + "n_estimators": xgb_cfg.n_estimators, + "max_depth": xgb_cfg.max_depth, + "learning_rate": xgb_cfg.learning_rate, + "subsample": xgb_cfg.subsample, + "colsample_bytree": xgb_cfg.colsample_bytree, + "min_child_weight": xgb_cfg.min_child_weight, + "early_stopping_rounds": early_stopping_rounds, + "random_state": cfg.seed, + "n_jobs": n_jobs, + } + + resolved_scale_pos_weight = None + if task_type == "regression": + model = XGBRegressor( + **common_params, + eval_metric=_xgboost_eval_metric(task_type), + ) + else: + requested_scale_pos_weight = xgb_cfg.get("scale_pos_weight", None) + resolved_scale_pos_weight = _resolve_scale_pos_weight(requested_scale_pos_weight, y_train) + if resolved_scale_pos_weight is not None: + was_struct = OmegaConf.is_struct(cfg) + OmegaConf.set_struct(cfg, False) + cfg.xgboost.scale_pos_weight = resolved_scale_pos_weight + OmegaConf.set_struct(cfg, was_struct) + if requested_scale_pos_weight == "balanced": + print( + " scale_pos_weight (native balanced n_neg/n_pos): " + f"{resolved_scale_pos_weight:.6g}" + ) + else: + print(f" scale_pos_weight: {resolved_scale_pos_weight:.6g}") + model = XGBClassifier( + **common_params, + scale_pos_weight=resolved_scale_pos_weight, + eval_metric=_xgboost_eval_metric(task_type), + ) + + model.fit( + X_train, + y_train, + eval_set=[(X_val, y_val)], + verbose=50, + ) + + # ========================================================================= + # 4. Evaluate + # ========================================================================= + print("\n" + "=" * 80) + print("4. Evaluating on test set") + print("=" * 80) + + metrics = {} + + if task_type == "regression": + y_pred = model.predict(X_test) + metrics["test/mse"] = mean_squared_error(y_test, y_pred) + metrics["test/mae"] = mean_absolute_error(y_test, y_pred) + metrics["test/r2"] = r2_score(y_test, y_pred) + else: + y_pred_proba = model.predict_proba(X_test)[:, 1] + metrics.update(_binary_classification_metrics(y_test, y_pred_proba)) + if resolved_scale_pos_weight is not None: + metrics["xgboost/scale_pos_weight"] = resolved_scale_pos_weight + + print("\n Test Results:") + for key, value in metrics.items(): + print(f" {key}: {value:.4f}") + + # ========================================================================= + # 5. Optional Fairness Evaluation + # ========================================================================= + fairness_cfg = cfg.get("eval", {}).get("fairness", {}) + fairness_report = None + if fairness_cfg.get("enabled", False): + print("\n" + "=" * 80) + print("5. Fairness Evaluation") + print("=" * 80) + + from slices.eval.fairness_evaluator import FairnessEvaluator, flatten_fairness_report + + # Get stay_ids for test set + test_stay_ids = [datamodule.dataset.stay_ids[i] for i in datamodule.test_indices] + + evaluator = FairnessEvaluator( + static_df=datamodule.dataset.static_df, + protected_attributes=list( + fairness_cfg.get("protected_attributes", ["gender", "age_group"]) + ), + min_subgroup_size=fairness_cfg.get("min_subgroup_size", 50), + task_type=task_type, + dataset_name=getattr(getattr(datamodule, "processed_dir", None), "name", None), + ) + + if task_type == "regression": + predictions_tensor = torch.tensor(y_pred, dtype=torch.float32) + else: + predictions_tensor = torch.tensor(y_pred_proba, dtype=torch.float32) + labels_tensor = torch.tensor(y_test, dtype=torch.float32) + + fairness_report = evaluator.evaluate(predictions_tensor, labels_tensor, test_stay_ids) + evaluator.print_report(fairness_report) + + # ========================================================================= + # 6. Save Model + # ========================================================================= + output_dir = Path(cfg.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + model_path = output_dir / "xgboost_model.json" + model.save_model(str(model_path)) + print(f"\n Model saved to: {model_path}") + + # ========================================================================= + # 7. Log to W&B + # ========================================================================= + if wandb_run is not None: + wandb_run.summary.update(train_label_support_summary(train_support_stats)) + wandb_run.summary.update(metrics) + wandb_run.summary.update( + { + EVAL_ARTIFACT_PATH_KEY: str(model_path), + EVAL_ARTIFACT_SHA256_KEY: file_sha256(model_path), + } + ) + if fairness_report: + wandb_run.summary.update(flatten_fairness_report(fairness_report)) + wandb_module.finish() + + print("\n" + "=" * 80) + print("XGBoost Baseline Complete!") + print("=" * 80) + + +if __name__ == "__main__": + main() diff --git a/src/slices/constants.py b/src/slices/constants.py index 20e26b0..33d4b55 100644 --- a/src/slices/constants.py +++ b/src/slices/constants.py @@ -11,8 +11,61 @@ # ============================================================================= # Observation Window # ============================================================================= -SEQ_LENGTH_HOURS: int = 48 -MIN_STAY_HOURS: int = 48 +SEQ_LENGTH_HOURS: int = 24 # Model input window +MIN_STAY_HOURS: int = 24 +# Benchmark label horizon: 24h observation + up to 24h forward prediction +# (e.g. mortality_24h / AKI from hours 24-48). +LABEL_HORIZON_HOURS: int = 48 + +# Thesis task surface. These are fixed benchmark labels, not sweep knobs. +THESIS_TASKS: tuple[str, ...] = ( + "mortality_24h", + "mortality_hospital", + "aki_kdigo", + "los_remaining", +) + +# Public downstream evaluation regimes. Legacy "A"/"B" values are normalized +# at export/evaluation boundaries so historical W&B runs remain comparable. +LINEAR_PROBE_PROTOCOL: str = "linear_probe" +FULL_FINETUNE_PROTOCOL: str = "full_finetune" +DOWNSTREAM_PROTOCOLS: tuple[str, str] = ( + LINEAR_PROBE_PROTOCOL, + FULL_FINETUNE_PROTOCOL, +) + + +def canonical_downstream_protocol(value: object) -> str | None: + """Return the public downstream protocol name for current or legacy values.""" + if value is None: + return None + text = str(value).strip() + if not text: + return None + key = text.lower().replace("-", "_") + if key in {"a", "protocol_a", "probe", "linear_probe"}: + return LINEAR_PROBE_PROTOCOL + if key in { + "b", + "protocol_b", + "finetune", + "fine_tune", + "full_finetune", + "full_fine_tune", + "fully_finetuned", + }: + return FULL_FINETUNE_PROTOCOL + return key + + +def downstream_protocol_from_freeze(freeze_encoder: object) -> str | None: + """Infer the downstream protocol from the encoder freeze setting.""" + if freeze_encoder is True: + return LINEAR_PROBE_PROTOCOL + if freeze_encoder is False: + return FULL_FINETUNE_PROTOCOL + return None + # ============================================================================= # Extraction @@ -35,7 +88,65 @@ # ============================================================================= # Feature Blocklist # ============================================================================= -# RICU concepts that leak downstream task labels and must be excluded from -# model input. los_hosp / los_icu are updated hourly by RICU and directly -# reveal the length-of-stay answer. -FEATURE_BLOCKLIST: frozenset = frozenset({"los_hosp", "los_icu"}) +# RICU concepts that must be excluded from model input because they either +# directly leak labels, encode future treatment duration, or are high-level +# summaries derived from lower-level features already present in the benchmark. +DIRECT_LEAKAGE_FEATURES: frozenset[str] = frozenset( + { + # Updated hourly by RICU and directly reveal length-of-stay targets. + "los_hosp", + "los_icu", + # Internal ricu helper variable, not a clinical concept. + "dur_var", + } +) + +FUTURE_DERIVED_FEATURES: frozenset[str] = frozenset( + { + # Medication duration concepts use stop/end times and can encode + # post-observation-window treatment duration in rows timestamped <24h. + "dobu_dur", + "dopa_dur", + "epi_dur", + "norepi_dur", + # Derived from vasopressor rate + duration concepts. + "dobu60", + "dopa60", + "epi60", + "norepi60", + "vaso_ind", + # Derived sepsis/cardio-SOFA features downstream of the duration concepts + # or suspected-infection windows; keep raw ingredients instead. + "sofa_cardio", + "sofa", + "susp_inf", + "sep3", + } +) + +DERIVED_SUMMARY_FEATURES: frozenset[str] = frozenset( + { + # Ratios or dose-equivalent summaries; raw components/rates are retained. + "pafi", + "safi", + "norepi_equiv", + # Windowed or composite clinical indicators/scores; raw components are + # retained where available. + "vent_ind", + "gcs", + "urine24", + "sofa_resp", + "sofa_coag", + "sofa_liver", + "sofa_cns", + "sofa_renal", + "qsofa", + "sirs", + "news", + "mews", + } +) + +FEATURE_BLOCKLIST: frozenset[str] = ( + DIRECT_LEAKAGE_FEATURES | FUTURE_DERIVED_FEATURES | DERIVED_SUMMARY_FEATURES +) diff --git a/src/slices/data/config_schemas.py b/src/slices/data/config_schemas.py index fe6ffdf..91d1c1a 100644 --- a/src/slices/data/config_schemas.py +++ b/src/slices/data/config_schemas.py @@ -21,6 +21,7 @@ PIN_MEMORY, SEQ_LENGTH_HOURS, TEST_RATIO, + THESIS_TASKS, TRAIN_RATIO, VAL_RATIO, ) @@ -180,6 +181,7 @@ class DataConfig(BaseModel): # Data Paths # ========================================================================== csv_root: Optional[str] = None # Raw CSV path (only for convert_csv_to_parquet) + ricu_output_dir: Optional[str] = None # RICU parquet output path parquet_root: Optional[str] = None # Parquet files (input for extraction) processed_dir: str # Extracted features (output of extraction, input for training) @@ -191,7 +193,7 @@ class DataConfig(BaseModel): feature_set: Literal["core", "extended"] = "core" # Feature set to extract categories: Optional[List[str]] = None # Feature categories (null = all) extraction_batch_size: int = EXTRACTION_BATCH_SIZE - tasks: List[str] = Field(default_factory=lambda: ["mortality_24h", "mortality_hospital"]) + tasks: List[str] = Field(default_factory=lambda: list(THESIS_TASKS)) # Config directory paths (auto-detected if null) tasks_dir: Optional[str] = None @@ -203,14 +205,19 @@ class DataConfig(BaseModel): pin_memory: bool = PIN_MEMORY # Patient-level split ratios - train_ratio: float = TRAIN_RATIO - val_ratio: float = VAL_RATIO - test_ratio: float = TEST_RATIO + train_ratio: float = Field(default=TRAIN_RATIO, ge=0.0, le=1.0) + val_ratio: float = Field(default=VAL_RATIO, ge=0.0, le=1.0) + test_ratio: float = Field(default=TEST_RATIO, ge=0.0, le=1.0) # Preprocessing applied during training normalize: bool = NORMALIZE - model_config = {"extra": "allow"} # Allow additional fields from Hydra + # Optional pretraining windowing controls + enable_sliding_windows: bool = False + window_size: Optional[int] = None + window_stride: Optional[int] = None + + model_config = {"extra": "forbid"} @model_validator(mode="after") def validate_split_ratios(self) -> "DataConfig": @@ -235,3 +242,11 @@ def validate_positive_int(cls, v: int) -> int: if v <= 0: raise ValueError("Value must be positive") return v + + @field_validator("window_size", "window_stride") + @classmethod + def validate_optional_positive_int(cls, v: Optional[int]) -> Optional[int]: + """Validate optional window dimensions when configured.""" + if v is not None and v <= 0: + raise ValueError("Value must be positive") + return v diff --git a/src/slices/data/datamodule.py b/src/slices/data/datamodule.py index d976b24..6504a2a 100644 --- a/src/slices/data/datamodule.py +++ b/src/slices/data/datamodule.py @@ -7,13 +7,11 @@ import logging from pathlib import Path -from typing import Any, Dict, List, Optional, Set, Tuple, Union +from typing import Any, Dict, List, Optional, Union import lightning.pytorch as L import numpy as np -import polars as pl import torch -import yaml from torch.utils.data import DataLoader, Subset from slices.constants import ( @@ -25,13 +23,15 @@ ) from slices.data.dataset import ICUDataset from slices.data.sliding_window import SlidingWindowDataset +from slices.data.splits import ( + compute_patient_level_splits, + save_global_split_info, + subsample_train_indices, +) # Module-level logger logger = logging.getLogger(__name__) -# Constants -TQDM_MIN_ITEMS = 1000 # Only show progress bar for collections larger than this - def icu_collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]: """Collate function for batching ICU samples. @@ -83,7 +83,7 @@ class ICUDataModule(L.LightningDataModule): Example: >>> dm = ICUDataModule( - ... processed_dir="data/processed/mimic-iv-demo", + ... processed_dir="data/processed/miiv", ... task_name="mortality_24h", ... batch_size=32, ... ) @@ -163,389 +163,28 @@ def __init__( self.window_stride = window_stride # Validate ratios + ratios = { + "train_ratio": train_ratio, + "val_ratio": val_ratio, + "test_ratio": test_ratio, + } + for name, ratio in ratios.items(): + if ratio < 0.0 or ratio > 1.0: + raise ValueError(f"{name} must be in [0, 1], got {ratio}") + total_ratio = train_ratio + val_ratio + test_ratio if not np.isclose(total_ratio, 1.0): raise ValueError(f"Split ratios must sum to 1.0, got {total_ratio}") # Will be set in setup() self.dataset: Optional[ICUDataset] = None + self.total_raw_stays: int = 0 + self.normalization_train_indices: List[int] = [] + self.full_train_indices: List[int] = [] self.train_indices: List[int] = [] self.val_indices: List[int] = [] self.test_indices: List[int] = [] - def _load_cached_splits( - self, static_df: pl.DataFrame, stay_ids: List[int] - ) -> Optional[Tuple[List[int], List[int], List[int]]]: - """Load cached splits from splits.yaml if valid. - - Validates that cached splits match current parameters (seed, ratios) and - that the patient lists are consistent with current data. - - Args: - static_df: Static dataframe with stay_id -> patient_id mapping. - stay_ids: List of stay_ids in order from timeseries parquet. - - Returns: - Tuple of (train_indices, val_indices, test_indices) if cache is valid, - None otherwise. - """ - split_path = self.processed_dir / "splits.yaml" - if not split_path.exists(): - return None - - try: - with open(split_path) as f: - cached = yaml.safe_load(f) - - # Validate parameters match - if ( - cached.get("seed") != self.seed - or not np.isclose(cached.get("train_ratio", 0), self.train_ratio) - or not np.isclose(cached.get("val_ratio", 0), self.val_ratio) - or not np.isclose(cached.get("test_ratio", 0), self.test_ratio) - ): - logger.debug("Cached splits have different parameters, recomputing") - return None - - # Get patient lists from cache - train_patients = set(cached.get("train_patients", [])) - val_patients = set(cached.get("val_patients", [])) - test_patients = set(cached.get("test_patients", [])) - - if not train_patients or not val_patients or not test_patients: - logger.debug("Cached splits missing patient lists, recomputing") - return None - - # Validate patient sets are disjoint - if not ( - train_patients.isdisjoint(val_patients) - and train_patients.isdisjoint(test_patients) - and val_patients.isdisjoint(test_patients) - ): - logger.debug("Cached splits have overlapping patients, recomputing") - return None - - # Get stay_id -> patient_id mapping - stay_to_patient = dict( - zip(static_df["stay_id"].to_list(), static_df["patient_id"].to_list()) - ) - - # Validate all patients in data are accounted for - current_patients = set(stay_to_patient.values()) - cached_patients = train_patients | val_patients | test_patients - - if current_patients != cached_patients: - logger.debug( - f"Cached splits have different patients " - f"(cached: {len(cached_patients)}, current: {len(current_patients)}), " - f"recomputing" - ) - return None - - # Reconstruct indices from patient lists - train_indices = [] - val_indices = [] - test_indices = [] - - for idx, stay_id in enumerate(stay_ids): - patient_id = stay_to_patient.get(stay_id) - if patient_id in train_patients: - train_indices.append(idx) - elif patient_id in val_patients: - val_indices.append(idx) - elif patient_id in test_patients: - test_indices.append(idx) - else: - logger.debug(f"Stay {stay_id} has unknown patient {patient_id}, recomputing") - return None - - logger.info( - f"Loaded cached splits: " - f"{len(train_indices)} train, {len(val_indices)} val, {len(test_indices)} test" - ) - return train_indices, val_indices, test_indices - - except Exception as e: - logger.debug(f"Failed to load cached splits: {e}, recomputing") - return None - - def _filter_stays_with_missing_labels( - self, stay_ids: List[int], labels_df: pl.DataFrame - ) -> Tuple[List[int], Set[int]]: - """Filter out stays with missing labels for the configured task. - - This MUST be called BEFORE computing splits to ensure indices remain consistent - between DataModule and Dataset. - - Args: - stay_ids: List of all stay_ids from timeseries. - labels_df: Labels dataframe with task columns. - - Returns: - Tuple of (filtered_stay_ids, excluded_stay_ids_set). - """ - if self.task_name is None: - return stay_ids, set() - - logger.debug(f"Checking labels for task '{self.task_name}'") - - # Detect multi-label tasks: columns prefixed with "{task_name}_" - multilabel_cols = [c for c in labels_df.columns if c.startswith(f"{self.task_name}_")] - is_multilabel = len(multilabel_cols) > 0 and self.task_name not in labels_df.columns - - # Get stays with valid labels for this task - if not is_multilabel and self.task_name not in labels_df.columns: - raise ValueError( - f"Task '{self.task_name}' not found in labels. " f"Available: {labels_df.columns}" - ) - - # Find stays with non-null labels - if is_multilabel: - valid_labels_df = labels_df.filter( - pl.all_horizontal([pl.col(c).is_not_null() for c in multilabel_cols]) - ) - else: - valid_labels_df = labels_df.filter(pl.col(self.task_name).is_not_null()) - valid_stay_ids = set(valid_labels_df["stay_id"].to_list()) - - # Filter stay_ids maintaining order - filtered_stay_ids = [sid for sid in stay_ids if sid in valid_stay_ids] - excluded_stay_ids = set(stay_ids) - valid_stay_ids - - if excluded_stay_ids: - pct_excluded = len(excluded_stay_ids) / len(stay_ids) * 100 - logger.warning( - f"Excluding {len(excluded_stay_ids):,} stays ({pct_excluded:.1f}%) " - f"with missing '{self.task_name}' labels. Remaining: {len(filtered_stay_ids):,}" - ) - - return filtered_stay_ids, excluded_stay_ids - - def _get_patient_level_splits(self) -> Tuple[List[int], List[int], List[int]]: - """Create patient-level train/val/test splits from parquet files. - - First checks for cached splits in splits.yaml. If found and parameters match, - loads splits from cache. Otherwise, computes splits from scratch. - - Loads static and timeseries data directly without requiring the full dataset - to be initialized. This is called BEFORE dataset creation to ensure - normalization statistics use only training data (prevents data leakage). - - IMPORTANT: Filters out stays with missing labels BEFORE computing splits - to ensure indices remain consistent with the Dataset after it loads. - - Uses deterministic shuffling of patient_id to assign patients to splits. - All stays from a patient go to the same split. - - Returns: - Tuple of (train_indices, val_indices, test_indices). - """ - logger.info("Loading data for split computation...") - - # Load only needed columns for efficiency - static_path = self.processed_dir / "static.parquet" - logger.debug(f"Loading static data from {static_path.name}") - self._static_df = pl.read_parquet(static_path, columns=["stay_id", "patient_id"]) - - # Load only stay_id column from timeseries (much faster than full load) - timeseries_path = self.processed_dir / "timeseries.parquet" - logger.debug(f"Loading stay_ids from {timeseries_path.name}") - timeseries_df = pl.read_parquet(timeseries_path, columns=["stay_id"]) - all_stay_ids = timeseries_df["stay_id"].to_list() - - # Load labels to filter out stays with missing labels - labels_path = self.processed_dir / "labels.parquet" - logger.debug(f"Loading labels from {labels_path.name}") - self._labels_df = pl.read_parquet(labels_path) - - # CRITICAL: Filter stays with missing labels BEFORE computing splits - # This ensures indices computed here match the filtered Dataset - stay_ids, self._excluded_stay_ids = self._filter_stays_with_missing_labels( - all_stay_ids, self._labels_df - ) - - # Store for later use - self._all_stay_ids = all_stay_ids - self._filtered_stay_ids = stay_ids - - # Try to load cached splits first - logger.debug("Checking for cached splits") - cached_splits = self._load_cached_splits(self._static_df, stay_ids) - if cached_splits is not None: - return cached_splits - - logger.info("Computing patient-level splits...") - - # Get stay_id -> patient_id mapping (only for filtered stays) - logger.debug("Building stay-to-patient mapping") - stay_to_patient = dict( - zip(self._static_df["stay_id"].to_list(), self._static_df["patient_id"].to_list()) - ) - - # Get unique patients (sorted for deterministic ordering across Python runs) - filtered_patient_ids = {stay_to_patient[sid] for sid in stay_ids if sid in stay_to_patient} - unique_patients = sorted(filtered_patient_ids) - n_patients = len(unique_patients) - logger.debug(f"Found {n_patients:,} unique patients") - - # Warn if patient_id == stay_id for all stays (e.g. HiRID, SICdb) - # This means the dataset lacks true patient-level IDs, so multiple ICU - # stays from the same real patient may leak across splits. - if n_patients == len(stay_ids) and n_patients > 0: - filtered_stay_set = set(stay_ids) - if filtered_patient_ids == filtered_stay_set: - logger.warning( - "patient_id == stay_id for all stays. This dataset likely lacks " - "true patient-level identifiers (e.g. HiRID, SICdb). " - "Patient-level split cannot prevent leakage from repeat ICU admissions." - ) - - # Shuffle patients deterministically using seed - logger.debug(f"Shuffling patients (seed={self.seed})") - rng = np.random.RandomState(self.seed) - patient_indices = np.arange(n_patients) - rng.shuffle(patient_indices) - shuffled_patients = [unique_patients[i] for i in patient_indices] - - # Split patients - n_train = int(n_patients * self.train_ratio) - n_val = int(n_patients * self.val_ratio) - - train_patients = set(shuffled_patients[:n_train]) - val_patients = set(shuffled_patients[n_train : n_train + n_val]) - test_patients = set(shuffled_patients[n_train + n_val :]) - logger.debug( - f"Split: {len(train_patients):,} train, " - f"{len(val_patients):,} val, {len(test_patients):,} test patients" - ) - - # Verify no patient overlap between splits (data leakage check) - assert train_patients.isdisjoint( - val_patients - ), "Patient leakage detected: train/val splits have overlapping patients" - assert train_patients.isdisjoint( - test_patients - ), "Patient leakage detected: train/test splits have overlapping patients" - assert val_patients.isdisjoint( - test_patients - ), "Patient leakage detected: val/test splits have overlapping patients" - - # Verify all patients are accounted for in exactly one split - all_patients_in_splits = train_patients | val_patients | test_patients - all_unique_patients = set(unique_patients) - missing_patients = all_unique_patients - all_patients_in_splits - extra_patients = all_patients_in_splits - all_unique_patients - - assert not missing_patients, ( - f"Patient split validation failed: {len(missing_patients)} patients " - f"not assigned to any split. First 5: {list(missing_patients)[:5]}" - ) - assert not extra_patients, ( - f"Patient split validation failed: {len(extra_patients)} patients " - f"in splits but not in data. First 5: {list(extra_patients)[:5]}" - ) - - # Map back to stay indices - logger.debug("Mapping patients to stay indices") - train_indices = [] - val_indices = [] - test_indices = [] - - for idx, stay_id in enumerate(stay_ids): - patient_id = stay_to_patient.get(stay_id) - if patient_id is None: - raise ValueError( - f"Stay {stay_id} has no patient_id mapping. " - "Patient IDs are required for patient-level splits." - ) - if patient_id in train_patients: - train_indices.append(idx) - elif patient_id in val_patients: - val_indices.append(idx) - elif patient_id in test_patients: - test_indices.append(idx) - else: - raise ValueError( - f"Stay {stay_id} has patient_id {patient_id} which is not " - "in any split. This indicates a bug in split computation." - ) - - # Final validation: all stays should be accounted for - total_assigned = len(train_indices) + len(val_indices) + len(test_indices) - assert total_assigned == len(stay_ids), ( - f"Split validation failed: {total_assigned} stays assigned but " - f"{len(stay_ids)} total stays in dataset" - ) - - logger.info( - f"Splits computed: {len(train_indices):,} train, " - f"{len(val_indices):,} val, {len(test_indices):,} test stays" - ) - - return train_indices, val_indices, test_indices - - def _subsample_train_indices(self, train_indices: List[int]) -> List[int]: - """Subsample training indices for label-efficiency ablations. - - Uses a separate RNG seeded with self.seed to ensure deterministic - subsampling. Always selects at least 1 sample. - - Args: - train_indices: Full list of training indices. - - Returns: - Subsampled list of training indices. - """ - n_full = len(train_indices) - n_subsample = max(1, int(n_full * self.label_fraction)) - - rng = np.random.RandomState(self.seed) - subsample_idx = rng.choice(n_full, size=n_subsample, replace=False) - subsample_idx.sort() # Maintain original order - subsampled = [train_indices[i] for i in subsample_idx] - - logger.info( - f"Label fraction={self.label_fraction}: using {n_subsample:,}/{n_full:,} " - f"training samples ({self.label_fraction * 100:.1f}%)" - ) - - return subsampled - - def _save_split_info(self) -> None: - """Save split information to file for reproducibility.""" - if self.dataset is None: - return - - static_df = self.dataset.static_df - stay_to_patient = dict( - zip(static_df["stay_id"].to_list(), static_df["patient_id"].to_list()) - ) - - train_patients = sorted( - {stay_to_patient[self.dataset.stay_ids[i]] for i in self.train_indices} - ) - val_patients = sorted({stay_to_patient[self.dataset.stay_ids[i]] for i in self.val_indices}) - test_patients = sorted( - {stay_to_patient[self.dataset.stay_ids[i]] for i in self.test_indices} - ) - - split_info = { - "seed": self.seed, - "train_ratio": self.train_ratio, - "val_ratio": self.val_ratio, - "test_ratio": self.test_ratio, - "train_patients": train_patients, - "val_patients": val_patients, - "test_patients": test_patients, - "train_stays": len(self.train_indices), - "val_stays": len(self.val_indices), - "test_stays": len(self.test_indices), - } - - split_path = self.processed_dir / "splits.yaml" - with open(split_path, "w") as f: - yaml.dump(split_info, f, default_flow_style=False) - def setup(self, stage: Optional[str] = None) -> None: """Set up datasets for train/val/test. @@ -565,23 +204,50 @@ def setup(self, stage: Optional[str] = None) -> None: # CRITICAL: Get splits FIRST before creating dataset # This also filters stays with missing labels to ensure index consistency logger.debug("[Step 1/3] Computing patient-level splits") - self.train_indices, self.val_indices, self.test_indices = self._get_patient_level_splits() + ( + self.train_indices, + self.val_indices, + self.test_indices, + self._static_df, + self._labels_df, + self._all_stay_ids, + self._filtered_stay_ids, + self._excluded_stay_ids, + self._normalization_train_indices, + ) = compute_patient_level_splits( + self.processed_dir, + self.task_name, + self.seed, + self.train_ratio, + self.val_ratio, + self.test_ratio, + ) + + # Save full train indices for normalization BEFORE subsampling. + # Label statistics and optimization subsets are task-filtered, but + # normalization must use the full unfiltered train split so AKI + # finetuning/supervised runs share the SSL pretraining normalizer. + self.full_train_indices = list(self.train_indices) + self.normalization_train_indices = list(self._normalization_train_indices) + self.total_raw_stays = len(self._all_stay_ids) # Subsample training indices for label-efficiency ablations if self.label_fraction < 1.0: - self.train_indices = self._subsample_train_indices(self.train_indices) + self.train_indices = subsample_train_indices( + self.train_indices, self.label_fraction, self.seed + ) - # Create dataset with training indices for normalization - # IMPORTANT: Use handle_missing_labels='raise' because we already filtered - # missing labels in _get_patient_level_splits. If any are still missing, - # it indicates a bug. + # Create dataset with task-filtered training indices for provenance and + # full-cohort training indices for normalization. + # self.train_indices (possibly subsampled) is used by train_dataloader. logger.debug("[Step 2/3] Creating ICUDataset") self.dataset = ICUDataset( data_dir=self.processed_dir, task_name=self.task_name, seq_length=self.seq_length, normalize=self.normalize, - train_indices=self.train_indices, + train_indices=self.full_train_indices, + normalization_train_indices=self.normalization_train_indices, # Use 'raise' since we pre-filtered - any missing labels now is a bug handle_missing_labels="raise" if self.task_name else "filter", # Pass excluded stays so Dataset can validate consistency @@ -599,14 +265,37 @@ def setup(self, stage: Optional[str] = None) -> None: "This indicates an index consistency bug." ) - # Save split information for reproducibility - logger.debug("[Step 3/3] Saving split information") - self._save_split_info() + # Save split information only when the canonical prep artifact is absent. + # Prepared datasets already own a stable splits.yaml, and downstream runs + # must not rewrite it after task filtering or label-efficiency subsampling. + logger.debug("[Step 3/3] Preserving canonical split information") + splits_path = self.processed_dir / "splits.yaml" + if not splits_path.exists(): + save_global_split_info( + processed_dir=self.processed_dir, + static_df=self._static_df, + stay_ids=self._all_stay_ids, + seed=self.seed, + train_ratio=self.train_ratio, + val_ratio=self.val_ratio, + test_ratio=self.test_ratio, + dataset=self.dataset if self.label_fraction < 1.0 else None, + train_subset_indices=self.train_indices if self.label_fraction < 1.0 else None, + label_fraction=self.label_fraction, + ) + else: + logger.debug("Canonical splits.yaml already exists; leaving it unchanged") + + # Free temporary data used only during setup — Dataset holds its own copies + del self._static_df, self._labels_df + del self._all_stay_ids, self._filtered_stay_ids, self._excluded_stay_ids + del self._normalization_train_indices logger.info( f"DataModule setup complete: " f"Train={len(self.train_indices):,}, Val={len(self.val_indices):,}, " - f"Test={len(self.test_indices):,} stays" + f"Test={len(self.test_indices):,}, " + f"NormalizerTrain={len(self.normalization_train_indices):,} stays" ) def train_dataloader(self) -> DataLoader: @@ -738,6 +427,8 @@ def get_split_info(self) -> Dict[str, Any]: return { # Stay counts "train_stays": len(self.train_indices), + "full_train_stays": len(self.full_train_indices), + "normalization_train_stays": len(self.normalization_train_indices), "val_stays": len(self.val_indices), "test_stays": len(self.test_indices), "total_stays": total_stays, @@ -748,6 +439,14 @@ def get_split_info(self) -> Dict[str, Any]: "total_patients": total_patients, # Actual ratios (for verification) "actual_train_ratio": len(self.train_indices) / total_stays if total_stays > 0 else 0, + "actual_full_train_ratio": ( + len(self.full_train_indices) / total_stays if total_stays > 0 else 0 + ), + "actual_normalization_train_ratio": ( + len(self.normalization_train_indices) / self.total_raw_stays + if self.total_raw_stays > 0 + else 0 + ), "actual_val_ratio": len(self.val_indices) / total_stays if total_stays > 0 else 0, "actual_test_ratio": len(self.test_indices) / total_stays if total_stays > 0 else 0, } @@ -757,3 +456,18 @@ def get_label_statistics(self) -> Dict[str, Dict[str, Any]]: if self.dataset is None: raise RuntimeError("Call setup() before get_label_statistics()") return self.dataset.get_label_statistics() + + def get_train_label_statistics(self, use_full_train: bool = False) -> Dict[str, Dict[str, Any]]: + """Return label statistics for the train split. + + Args: + use_full_train: When True, compute statistics on the full patient-level + train split before any label-efficiency subsampling. When False, + compute statistics on the optimization subset actually used for + training in this run. + """ + if self.dataset is None: + raise RuntimeError("Call setup() before get_train_label_statistics()") + + train_indices = self.full_train_indices if use_full_train else self.train_indices + return self.dataset.get_label_statistics(indices=train_indices) diff --git a/src/slices/data/dataset.py b/src/slices/data/dataset.py index 12f7899..8d246e3 100644 --- a/src/slices/data/dataset.py +++ b/src/slices/data/dataset.py @@ -4,24 +4,36 @@ returns (timeseries, mask, labels, static_features) tuples for training. """ -import hashlib import logging -import warnings +from collections import Counter from pathlib import Path from typing import Any, Dict, List, Optional, Set, Tuple, Union -import numpy as np import polars as pl import torch import yaml from torch.utils.data import Dataset from tqdm import tqdm +from slices.data.tensor_cache import ( + load_cached_tensors, + load_normalization_stats, + save_cached_tensors, + save_dataset_metadata, + save_normalization_stats, +) +from slices.data.tensor_preprocessing import ( + apply_normalization_and_imputation, + compute_normalization_stats, + compute_single_sample_stages, + convert_raw_to_tensors, + extract_tensors_from_dataframe, +) + # Module-level logger logger = logging.getLogger(__name__) # Constants -MIN_STD_THRESHOLD = 1e-6 # Minimum standard deviation to avoid division by zero LARGE_DATASET_WARNING_THRESHOLD = 100_000 # Warn for datasets larger than this TQDM_MIN_ITEMS = 1000 # Only show progress bar for collections larger than this @@ -86,7 +98,7 @@ class ICUDataset(Dataset): metadata.yaml # Feature names, task names, etc. Example: - >>> dataset = ICUDataset("data/processed/mimic-iv-demo") + >>> dataset = ICUDataset("data/processed/miiv") >>> sample = dataset[0] >>> sample["timeseries"].shape # (seq_length, n_features) torch.Size([48, 9]) @@ -101,6 +113,7 @@ def __init__( seq_length: Optional[int] = None, # TODO: Add support for different sequence lengths normalize: bool = True, # TODO: Add support for different normalization strategies train_indices: Optional[List[int]] = None, + normalization_train_indices: Optional[List[int]] = None, handle_missing_labels: str = "filter", _excluded_stay_ids: Optional[Set[int]] = None, ) -> None: @@ -111,10 +124,17 @@ def __init__( task_name: Name of the task for label extraction (e.g., 'mortality_24h'). If None, no labels are returned. seq_length: Override sequence length (uses metadata default if None). - normalize: Whether to normalize features (z-score per feature). + normalize: Whether to z-score features before imputation. When False, + data remains in original units and missing values are + imputed from the available feature means. train_indices: Optional list of indices for training set. If provided, - normalization statistics are computed only on these samples. - This prevents data leakage from val/test sets. + records the task-filtered training split associated with this + dataset. When normalization_train_indices is not provided, these + indices are also used for normalization for backwards compatibility. + normalization_train_indices: Optional list of full raw-cohort training indices + used for normalization statistics. DataModule passes this + separately for task-filtered downstream runs so AKI uses the same + dataset/seed normalizer as SSL pretraining. handle_missing_labels: How to handle stays with missing labels when task_name is specified. Options: - 'filter': Remove samples with missing labels (default) @@ -127,6 +147,11 @@ def __init__( self.task_name = task_name self.normalize = normalize self.train_indices = train_indices + self.normalization_train_indices = ( + normalization_train_indices + if normalization_train_indices is not None + else train_indices + ) self.handle_missing_labels = handle_missing_labels self._excluded_stay_ids = _excluded_stay_ids or set() @@ -156,6 +181,9 @@ def __init__( self.n_features: int = len(self.feature_names) self.seq_length: int = seq_length or self.metadata["seq_length_hours"] self.task_names: List[str] = self.metadata.get("task_names", []) + self.task_types: Dict[str, str] = self._resolve_task_types() + self._timeseries_tensor: torch.Tensor + self._mask_tensor: torch.Tensor # Validate task_name if provided if task_name is not None and task_name not in self.task_names: @@ -171,169 +199,198 @@ def __init__( self._load_data() # Save dataset metadata (including removed samples) for reproducibility - self._save_dataset_metadata() + save_dataset_metadata( + self.data_dir, self.task_name, self.handle_missing_labels, self.removed_samples + ) + + @staticmethod + def _task_config_dirs() -> List[Path]: + """Return candidate task-config directories for task-type resolution.""" + repo_tasks = Path(__file__).resolve().parents[3] / "configs" / "tasks" + package_tasks = Path(__file__).resolve().parent / "tasks" + return [repo_tasks, package_tasks] + + def _resolve_task_type_from_config(self, task_name: str) -> Optional[str]: + """Resolve task type from the checked-in task configuration.""" + for task_dir in self._task_config_dirs(): + config_path = task_dir / f"{task_name}.yaml" + if not config_path.exists(): + continue + + with open(config_path) as f: + config = yaml.safe_load(f) or {} + + task_type = config.get("task_type") + if isinstance(task_type, str): + return task_type + + return None + + def _resolve_task_types(self) -> Dict[str, str]: + """Resolve task types from metadata with a config-file fallback.""" + label_manifest = self.metadata.get("label_manifest") or {} + task_types: Dict[str, str] = {} + + for task_name in self.task_names: + task_type = None + + manifest_entry = label_manifest.get(task_name) + if isinstance(manifest_entry, dict): + manifest_task_type = manifest_entry.get("task_type") + if isinstance(manifest_task_type, str): + task_type = manifest_task_type + + if task_type is None: + task_type = self._resolve_task_type_from_config(task_name) + + if task_type is None: + logger.warning( + "Could not resolve task_type for '%s'; defaulting statistics to binary.", + task_name, + ) + task_type = "binary" + + task_types[task_name] = task_type + + return task_types def _load_data(self) -> None: - """Load data from Parquet files into memory.""" + """Load data from Parquet files into memory. + + Uses a two-phase approach for efficiency: + 1. Load raw tensors from cache (or extract from parquet on cache miss) + 2. Normalize using full-cohort train indices, then filter excluded stays + + The raw tensor cache is shared across all runs for the same dataset, + regardless of task, seed, or label_fraction. This reduces cache size + from ~140GB (one per run config) to ~0.8GB (one per dataset). + """ logger.info("Loading data from Parquet files...") - # Load timeseries (dense format with nested lists) timeseries_path = self.data_dir / "timeseries.parquet" - logger.debug(f"Loading timeseries from {timeseries_path.name}") - self.timeseries_df = pl.read_parquet(timeseries_path) - - # Warn for large datasets that may cause memory issues - n_stays = len(self.timeseries_df) - if n_stays > LARGE_DATASET_WARNING_THRESHOLD: - logger.warning( - f"Large dataset detected ({n_stays:,} stays). " - "Loading entire dataset into memory." - ) - # Load static features + # Load static features and labels (always needed) static_path = self.data_dir / "static.parquet" logger.debug(f"Loading static features from {static_path.name}") self.static_df = pl.read_parquet(static_path) - # Load labels labels_path = self.data_dir / "labels.parquet" logger.debug(f"Loading labels from {labels_path.name}") self.labels_df = pl.read_parquet(labels_path) - # Pre-filter excluded stays if provided (from DataModule) - if self._excluded_stay_ids: - logger.debug(f"Pre-filtering {len(self._excluded_stay_ids):,} excluded stays") - self.timeseries_df = self.timeseries_df.filter( - ~pl.col("stay_id").is_in(list(self._excluded_stay_ids)) - ) - self.static_df = self.static_df.filter( - ~pl.col("stay_id").is_in(list(self._excluded_stay_ids)) + # Phase 1: Get raw tensors (from cache or parquet) + cached_tensors = load_cached_tensors(self.data_dir, self.seq_length, self.n_features) + if cached_tensors is not None: + timeseries_tensor = cached_tensors["timeseries_tensor"] + masks_tensor = cached_tensors["mask_tensor"] + # Load stay_ids from parquet (lightweight, just one column) + all_stay_ids = pl.read_parquet(timeseries_path, columns=["stay_id"])[ + "stay_id" + ].to_list() + logger.info("Using cached raw tensors") + else: + # Extract from parquet and save raw cache + logger.debug(f"Loading timeseries from {timeseries_path.name}") + timeseries_df = pl.read_parquet(timeseries_path) + + n_stays = len(timeseries_df) + if n_stays > LARGE_DATASET_WARNING_THRESHOLD: + logger.warning( + f"Large dataset detected ({n_stays:,} stays). " + "Loading entire dataset into memory." + ) + + all_stay_ids = timeseries_df["stay_id"].to_list() + + timeseries_tensor, masks_tensor = extract_tensors_from_dataframe( + timeseries_df, self.seq_length, self.n_features ) - self.labels_df = self.labels_df.filter( - ~pl.col("stay_id").is_in(list(self._excluded_stay_ids)) + del timeseries_df + + # Save raw tensors to cache (shared across all runs for this dataset) + save_cached_tensors( + self.data_dir, + timeseries_tensor, + masks_tensor, + self.seq_length, + self.n_features, ) - logger.debug(f"Filtered down to {len(self.timeseries_df):,} stays") - # Create stay_id -> index mapping + # Phase 2: Normalize before any task-label filtering. This keeps + # downstream AKI scaling aligned with SSL pretraining because both use + # the same full-cohort train indices for a dataset/seed. + cached_stats = load_normalization_stats( + self.data_dir, + self.normalization_train_indices, + self.normalize, + ) + self._precompute_tensors( + precomputed_tensors=(timeseries_tensor, masks_tensor), + train_indices=self.normalization_train_indices, + cached_stats=cached_stats, + ) + + # Phase 3: Apply stay exclusion filtering (task-specific, e.g. missing labels) + self.stay_ids = all_stay_ids + if self._excluded_stay_ids: + logger.debug(f"Filtering {len(self._excluded_stay_ids):,} excluded stays") + keep_indices = [ + i for i, sid in enumerate(all_stay_ids) if sid not in self._excluded_stay_ids + ] + idx_tensor = torch.tensor(keep_indices, dtype=torch.long) + self._timeseries_tensor = self._timeseries_tensor[idx_tensor] + self._mask_tensor = self._mask_tensor[idx_tensor] + self.stay_ids = [all_stay_ids[i] for i in keep_indices] + + # Filter static and labels DataFrames to match + excluded_list = list(self._excluded_stay_ids) + self.static_df = self.static_df.filter(~pl.col("stay_id").is_in(excluded_list)) + self.labels_df = self.labels_df.filter(~pl.col("stay_id").is_in(excluded_list)) + logger.debug(f"Filtered down to {len(self.stay_ids):,} stays") + + # Build stay_id -> index mapping logger.debug("Building stay_id index mapping") - self.stay_ids = self.timeseries_df["stay_id"].to_list() self.stay_id_to_idx = {sid: idx for idx, sid in enumerate(self.stay_ids)} logger.info(f"Loaded {len(self.stay_ids):,} stays") - # Try to load cached preprocessed tensors first (big speedup on subsequent runs) - cached_tensors = self._load_cached_tensors(self.train_indices) - if cached_tensors is not None: - # Use cached tensors (stacked format) - self._timeseries_tensor = cached_tensors["timeseries_tensor"] - self._mask_tensor = cached_tensors["mask_tensor"] - self.feature_means = cached_tensors["feature_means"] - self.feature_stds = cached_tensors["feature_stds"] - logger.info("Using cached preprocessed tensors") - else: - # Pre-extract raw arrays - raw_timeseries = self.timeseries_df["timeseries"].to_list() - raw_masks = self.timeseries_df["mask"].to_list() - - # Try to load existing normalization stats (for reproducibility) - cached_stats = self._load_normalization_stats(self.train_indices) - - # Pre-compute tensors for all samples (much faster __getitem__) - self._precompute_tensors(raw_timeseries, raw_masks, self.train_indices, cached_stats) - - # Save preprocessed tensors to cache for next run - self._save_cached_tensors(self.train_indices) - # Pre-compute labels and static features self._precompute_labels_and_static() def _precompute_tensors( self, - raw_timeseries: List[List[List[float]]], - raw_masks: List[List[List[bool]]], + raw_timeseries: Optional[List[List[List[float]]]] = None, + raw_masks: Optional[List[List[List[bool]]]] = None, train_indices: Optional[List[int]] = None, cached_stats: Optional[Dict[str, Any]] = None, + precomputed_tensors: Optional[tuple] = None, ) -> None: """Pre-compute all tensors at initialization for fast __getitem__. This converts raw nested lists to tensors and applies imputation/normalization once, rather than on every access. Uses vectorized operations for speed. - IMPORTANT: If train_indices is provided, normalization statistics are computed - ONLY on training samples to prevent data leakage from validation/test sets. - - Normalization statistics are cached to metadata for reproducibility across - dataset reloads. - Args: raw_timeseries: List of timeseries arrays (n_samples x seq_len x n_features). + Not needed if precomputed_tensors is provided. raw_masks: List of mask arrays (n_samples x seq_len x n_features). - train_indices: Optional list of indices for training set. If provided, - normalization stats computed only on these samples. - cached_stats: Optional cached normalization statistics. If provided, uses - these instead of recomputing (for reproducibility). + Not needed if precomputed_tensors is provided. + train_indices: Optional list of indices for training set. + cached_stats: Optional cached normalization statistics. + precomputed_tensors: Optional tuple of (timeseries_tensor, masks_tensor) + already extracted. Skips the list-to-tensor conversion step. """ - n_samples = len(raw_timeseries) - logger.info(f"Preprocessing {n_samples:,} samples...") - - # ===================================================================== - # Step 1: Vectorized conversion from nested lists to tensors - # ===================================================================== - logger.debug("[1/3] Converting to tensors...") - - # Convert all samples to numpy arrays in batches for memory efficiency - batch_size = 10000 - all_timeseries = [] - all_masks = [] - - n_batches = (n_samples + batch_size - 1) // batch_size - batch_iter = range(0, n_samples, batch_size) - if n_batches >= 10: # Only show progress for many batches - batch_iter = tqdm(batch_iter, desc=" Converting batches", unit="batch") - - for batch_start in batch_iter: - batch_end = min(batch_start + batch_size, n_samples) - batch_ts = [] - batch_mask = [] - - for i in range(batch_start, batch_end): - ts_data = raw_timeseries[i] - mask_data = raw_masks[i] - actual_len = len(ts_data) - - # Pad or truncate to seq_length - if actual_len >= self.seq_length: - # Truncate - ts_arr = np.array(ts_data[: self.seq_length], dtype=np.float32) - mask_arr = np.array(mask_data[: self.seq_length], dtype=bool) - else: - # Pad with NaN / False - ts_arr = np.full((self.seq_length, self.n_features), np.nan, dtype=np.float32) - mask_arr = np.zeros((self.seq_length, self.n_features), dtype=bool) - ts_arr[:actual_len] = np.array(ts_data, dtype=np.float32) - mask_arr[:actual_len] = np.array(mask_data, dtype=bool) - - batch_ts.append(ts_arr) - batch_mask.append(mask_arr) - - all_timeseries.extend(batch_ts) - all_masks.extend(batch_mask) - - # Stack into single arrays - timeseries_np = np.stack(all_timeseries) # (n_samples, seq_len, n_features) - masks_np = np.stack(all_masks) # (n_samples, seq_len, n_features) - - # Convert to tensors - timeseries_tensor = torch.from_numpy(timeseries_np) # (n_samples, seq_len, n_features) - masks_tensor = torch.from_numpy(masks_np) # (n_samples, seq_len, n_features) - - # Free numpy arrays - del timeseries_np, masks_np, all_timeseries, all_masks + if precomputed_tensors is not None: + timeseries_tensor, masks_tensor = precomputed_tensors + else: + timeseries_tensor, masks_tensor = convert_raw_to_tensors( + raw_timeseries, raw_masks, self.seq_length, self.n_features + ) + del raw_timeseries, raw_masks - # ===================================================================== - # Step 2: Compute normalization statistics (vectorized) - # ===================================================================== - logger.debug("[2/3] Computing normalization statistics...") + n_samples = timeseries_tensor.shape[0] + logger.info(f"Preprocessing {n_samples:,} samples...") + # Step 2: Compute preprocessing statistics used for normalization/imputation self.feature_means = torch.zeros(self.n_features) self.feature_stds = torch.ones(self.n_features) @@ -343,324 +400,51 @@ def _precompute_tensors( self.feature_stds = torch.tensor(cached_stats["feature_stds"], dtype=torch.float32) logger.debug("Using cached normalization statistics") elif train_indices is not None: - # Compute feature means (always needed for imputation) and stds (for normalization) - # CRITICAL: Use only training samples to prevent data leakage - train_ts = timeseries_tensor[train_indices] # (n_train, seq_len, n_features) - train_masks = masks_tensor[train_indices] - - # Vectorized computation of mean and std per feature - # Reshape to (n_samples * seq_len, n_features) for easier computation - flat_ts = train_ts.reshape(-1, self.n_features) # (n_train * seq_len, n_features) - flat_masks = train_masks.reshape(-1, self.n_features) - - # Create combined mask: observed AND not NaN - valid_mask = flat_masks & ~torch.isnan(flat_ts) - - # Vectorized mean computation using masked tensor operations - # Replace invalid values with 0 for sum, count valid entries - masked_ts = torch.where(valid_mask, flat_ts, torch.zeros_like(flat_ts)) - valid_counts = valid_mask.sum(dim=0).float() # (n_features,) - feature_sums = masked_ts.sum(dim=0) # (n_features,) - - # Compute means (avoid div by zero) - self.feature_means = torch.where( - valid_counts > 0, - feature_sums / valid_counts, - torch.zeros(self.n_features), + self.feature_means, self.feature_stds = compute_normalization_stats( + timeseries_tensor, masks_tensor, train_indices, self.n_features, self.normalize ) - - if self.normalize: - # Vectorized std computation - # Compute squared deviations from mean - deviations = torch.where( - valid_mask, - (flat_ts - self.feature_means.unsqueeze(0)) ** 2, - torch.zeros_like(flat_ts), - ) - variance = torch.where( - valid_counts > 1, - deviations.sum(dim=0) / (valid_counts - 1), # Bessel's correction - torch.ones(self.n_features), - ) - self.feature_stds = torch.sqrt(variance) - - # Clamp minimum std to avoid division by zero - self.feature_stds = torch.clamp(self.feature_stds, min=MIN_STD_THRESHOLD) - # Save computed stats for reproducibility - self._save_normalization_stats(train_indices) - logger.debug(f"Computed normalization stats for {self.n_features} features") + save_normalization_stats( + self.data_dir, + self.feature_means, + self.feature_stds, + self.feature_names, + train_indices, + self.normalize, + ) elif self.normalize: raise ValueError( "train_indices must be provided when normalize=True to prevent data leakage. " "Pass train_indices from your data splits, or set normalize=False." ) - - # ===================================================================== - # Step 3: Normalize then impute missing values (vectorized) - # ===================================================================== - logger.debug("[3/3] Applying normalization and imputation...") - - if self.normalize: - # Normalize first (NaN positions stay NaN through arithmetic) - timeseries_tensor = (timeseries_tensor - self.feature_means) / self.feature_stds - # After z-score normalization, 0 = feature mean in original space. - # This is the most neutral default: "no information, assume population average." - timeseries_tensor = torch.nan_to_num(timeseries_tensor, nan=0.0) else: - # Without normalization, impute with feature means (not 0). - # Zero-filling in original space creates physiologically impossible values - # (e.g., 0 heart rate, 0 blood pressure). Feature means are a better default. - for f in range(self.n_features): - nan_mask = torch.isnan(timeseries_tensor[:, :, f]) - timeseries_tensor[:, :, f][nan_mask] = self.feature_means[f] + all_indices = list(range(n_samples)) + self.feature_means, self.feature_stds = compute_normalization_stats( + timeseries_tensor, + masks_tensor, + all_indices, + self.n_features, + normalize=False, + ) + logger.debug( + "Computed feature means on the loaded dataset for normalize=False imputation" + ) + + # Step 3: Normalize then impute missing values + timeseries_tensor = apply_normalization_and_imputation( + timeseries_tensor, + self.feature_means, + self.feature_stds, + self.normalize, + self.n_features, + ) # Keep tensors stacked for memory efficiency and faster access - # __getitem__ will index directly into these tensors self._timeseries_tensor = timeseries_tensor # (n_samples, seq_len, n_features) self._mask_tensor = masks_tensor # (n_samples, seq_len, n_features) logger.info("Preprocessing complete") - def _load_normalization_stats( - self, current_train_indices: Optional[List[int]] - ) -> Optional[Dict[str, Any]]: - """Load cached normalization statistics from file if they exist and match current split. - - Validates that cached statistics were computed on the same training set to prevent - data leakage from validation/test sets. - - Args: - current_train_indices: List of training indices for current split, - or None for unsupervised. - - Returns: - Dictionary with 'feature_means' and 'feature_stds' if file exists and split matches, - None otherwise. - """ - if not self.normalize: - return None - - stats_path = self.data_dir / "normalization_stats.yaml" - if not stats_path.exists(): - return None - - try: - with open(stats_path) as f: - stats = yaml.safe_load(f) - - # Validate that cached stats were computed on the same training split - cached_train_indices = stats.get("train_indices") - current_train_set = set(current_train_indices) if current_train_indices else None - cached_train_set = set(cached_train_indices) if cached_train_indices else None - - if current_train_set != cached_train_set: - warnings.warn( - f"Cached normalization stats were computed on a different training split. " - f"Cached: {len(cached_train_set) if cached_train_set else 0} samples, " - f"Current: {len(current_train_set) if current_train_set else 0} samples. " - "Recomputing statistics to prevent data leakage.", - UserWarning, - ) - return None - - return stats - except Exception as e: - warnings.warn( - f"Failed to load normalization stats from {stats_path}: {e}. " - "Will recompute statistics.", - UserWarning, - ) - return None - - def _save_normalization_stats(self, train_indices: Optional[List[int]]) -> None: - """Save computed normalization statistics to file for reproducibility. - - Args: - train_indices: Optional list of training indices used to compute stats. - """ - stats_path = self.data_dir / "normalization_stats.yaml" - - stats = { - "feature_means": self.feature_means.tolist(), - "feature_stds": self.feature_stds.tolist(), - "feature_names": self.feature_names, - "train_indices": train_indices, - "normalize": self.normalize, - } - - try: - with open(stats_path, "w") as f: - yaml.dump(stats, f, default_flow_style=False) - except Exception as e: - warnings.warn( - f"Failed to save normalization stats to {stats_path}: {e}", - UserWarning, - ) - - def _get_tensor_cache_key(self, train_indices: Optional[List[int]]) -> str: - """Generate a hash key for tensor caching based on preprocessing parameters. - - The cache key includes: - - normalize flag - - seq_length - - hash of train_indices (for normalization stats consistency) - - hash of _excluded_stay_ids (for filtering consistency) - - Args: - train_indices: Optional list of training indices. - - Returns: - Hash string to use as cache identifier. - """ - # Create a deterministic string representation of parameters - params = { - "normalize": self.normalize, - "seq_length": self.seq_length, - "n_features": self.n_features, - } - - # Hash train_indices if provided (too large to store directly) - if train_indices is not None: - indices_str = ",".join(map(str, sorted(train_indices))) - indices_hash = hashlib.md5(indices_str.encode()).hexdigest()[:8] - params["train_indices_hash"] = indices_hash - else: - params["train_indices_hash"] = "none" - - # Hash _excluded_stay_ids if provided (ensures cache invalidation when filtering changes) - if self._excluded_stay_ids: - excluded_str = ",".join(map(str, sorted(self._excluded_stay_ids))) - excluded_hash = hashlib.md5(excluded_str.encode()).hexdigest()[:8] - params["excluded_stays_hash"] = excluded_hash - else: - params["excluded_stays_hash"] = "none" - - # Create overall hash - params_str = str(sorted(params.items())) - return hashlib.md5(params_str.encode()).hexdigest()[:12] - - def _get_tensor_cache_path(self, train_indices: Optional[List[int]]) -> Path: - """Get the path to the tensor cache file. - - Args: - train_indices: Optional list of training indices. - - Returns: - Path to the tensor cache file. - """ - cache_key = self._get_tensor_cache_key(train_indices) - cache_dir = self.data_dir / ".tensor_cache" - return cache_dir / f"tensors_{cache_key}.pt" - - def _load_cached_tensors(self, train_indices: Optional[List[int]]) -> Optional[Dict[str, Any]]: - """Load cached preprocessed tensors if they exist and are valid. - - Args: - train_indices: Optional list of training indices. - - Returns: - Dictionary with cached tensors and metadata if valid, None otherwise. - """ - cache_path = self._get_tensor_cache_path(train_indices) - if not cache_path.exists(): - return None - - try: - logger.debug(f"Loading cached tensors from {cache_path.name}") - cached = torch.load(cache_path, weights_only=True) - - # Validate cache metadata - if cached.get("n_features") != self.n_features: - logger.debug("Cached tensors have different feature count, recomputing") - return None - if cached.get("seq_length") != self.seq_length: - logger.debug("Cached tensors have different sequence length, recomputing") - return None - if cached.get("normalize") != self.normalize: - logger.debug("Cached tensors have different normalize flag, recomputing") - return None - - # Validate tensor shapes (support both old list and new stacked format) - timeseries = cached.get("timeseries_tensor") or cached.get("timeseries_tensors") - masks = cached.get("mask_tensor") or cached.get("mask_tensors") - if timeseries is None or masks is None: - logger.debug("Cached tensors missing data, recomputing") - return None - - # Convert old list format to stacked if needed - if isinstance(timeseries, list): - logger.debug("Converting cached tensors from list to stacked format") - timeseries = torch.stack(timeseries) - masks = torch.stack(masks) - cached["timeseries_tensor"] = timeseries - cached["mask_tensor"] = masks - - n_samples = timeseries.shape[0] if hasattr(timeseries, "shape") else len(timeseries) - logger.info(f"Loaded {n_samples:,} cached samples") - return cached - - except Exception as e: - logger.debug(f"Failed to load cached tensors: {e}, recomputing") - return None - - def _save_cached_tensors(self, train_indices: Optional[List[int]]) -> None: - """Save preprocessed tensors to cache file. - - Args: - train_indices: Optional list of training indices. - """ - cache_path = self._get_tensor_cache_path(train_indices) - cache_dir = cache_path.parent - cache_dir.mkdir(exist_ok=True) - - cache_data = { - "timeseries_tensor": self._timeseries_tensor, - "mask_tensor": self._mask_tensor, - "feature_means": self.feature_means, - "feature_stds": self.feature_stds, - "n_features": self.n_features, - "seq_length": self.seq_length, - "normalize": self.normalize, - } - - try: - logger.debug(f"Saving tensors to cache: {cache_path.name}") - torch.save(cache_data, cache_path) - except Exception as e: - logger.warning(f"Failed to save tensor cache to {cache_path}: {e}") - - def _save_dataset_metadata(self) -> None: - """Save dataset metadata including removed samples for reproducibility. - - Creates or updates dataset_metadata.yaml with information about: - - Number of samples used (original vs. final after filtering) - - Removed samples and reasons for removal - - Task name and label handling strategy - """ - if not self.removed_samples: - return # Nothing to save if no samples were removed - - metadata_path = self.data_dir / "dataset_metadata.yaml" - - metadata = { - "task_name": self.task_name, - "handle_missing_labels": self.handle_missing_labels, - "removed_samples_count": len(self.removed_samples), - "removed_samples": [ - {"stay_id": stay_id, "reason": reason} for stay_id, reason in self.removed_samples - ], - } - - try: - with open(metadata_path, "w") as f: - yaml.dump(metadata, f, default_flow_style=False) - except Exception as e: - warnings.warn( - f"Failed to save dataset metadata to {metadata_path}: {e}", - UserWarning, - ) - def _precompute_labels_and_static(self) -> None: """Pre-compute labels and static features for fast __getitem__ access. @@ -768,8 +552,10 @@ def _precompute_labels_and_static(self) -> None: self._labels_tensor = None indices_to_keep = list(range(len(self.stay_ids))) - # Pre-compute static features using vectorized Polars operations - logger.debug("Extracting static features") + # Pre-compute safe static features for batch payloads. Future-derived + # fields such as los_days stay available on static_df for cohort/fairness + # analysis but are intentionally not exposed to model-facing batches. + logger.debug("Extracting safe static features") self._static_data = [] for stay_id in self.stay_ids: static_row = static_by_stay.get(stay_id, {}) @@ -777,7 +563,6 @@ def _precompute_labels_and_static(self) -> None: { "age": static_row.get("age"), "gender": static_row.get("gender"), - "los_days": static_row.get("los_days"), } ) logger.info(f"Pre-computed {len(self._static_data):,} static feature records") @@ -831,64 +616,124 @@ def get_task_names(self) -> List[str]: """Return list of available task names.""" return self.task_names - def get_label_statistics(self) -> Dict[str, Dict[str, Any]]: - """Compute label statistics for each task. + def get_label_statistics( + self, indices: Optional[List[int]] = None + ) -> Dict[str, Dict[str, Any]]: + """Compute label statistics for each task or subset. + + Returns task-type-aware statistics: + - binary: {total, positive, negative, prevalence} + - regression: {total, mean, std, min, max} + - multiclass: {total, n_classes, class_counts} + - multilabel: per-subtask prevalence plus aggregate mean prevalence - For single-label tasks, returns {total, positive, negative, prevalence}. - For multi-label tasks (detected by prefixed columns like {task_name}_{subtask}), - returns per-subtask prevalence and an aggregate mean prevalence. + Args: + indices: Optional dataset indices to restrict the computation to. + When None, computes statistics over the full dataset. Returns: - Dict mapping task_name -> {count, positive, negative, prevalence, ...} + Dict mapping task_name -> task-specific summary statistics. """ + labels_df = self.labels_df + if indices is not None: + subset_stay_ids = [self.stay_ids[i] for i in indices] + labels_df = labels_df.filter(pl.col("stay_id").is_in(subset_stay_ids)) + stats: Dict[str, Dict[str, Any]] = {} for task_name in self.task_names: - if task_name in self.labels_df.columns: - labels = self.labels_df[task_name].drop_nulls() - positive = (labels == 1).sum() - total = len(labels) + task_type = self.task_types.get(task_name, "binary") + multilabel_cols = [c for c in labels_df.columns if c.startswith(f"{task_name}_")] + + if task_type == "multilabel" and multilabel_cols: + subtask_stats = {} + prevalences = [] + for col in multilabel_cols: + col_labels = labels_df[col].drop_nulls() + pos = (col_labels == 1).sum() + tot = len(col_labels) + prev = pos / tot if tot > 0 else 0.0 + subtask_stats[col] = { + "task_type": "binary", + "total": tot, + "positive": pos, + "negative": tot - pos, + "prevalence": prev, + } + prevalences.append(prev) stats[task_name] = { - "total": total, - "positive": positive, - "negative": total - positive, - "prevalence": positive / total if total > 0 else 0.0, + "task_type": "multilabel", + "total": len(labels_df), + "n_labels": len(multilabel_cols), + "mean_prevalence": ( + sum(prevalences) / len(prevalences) if prevalences else 0.0 + ), + "subtasks": subtask_stats, } - else: - # Check for multi-label columns (e.g., {task_name}_{subtask}, ...) - multilabel_cols = [ - c for c in self.labels_df.columns if c.startswith(f"{task_name}_") - ] - if multilabel_cols: - subtask_stats = {} - prevalences = [] - for col in multilabel_cols: - col_labels = self.labels_df[col].drop_nulls() - pos = (col_labels == 1).sum() - tot = len(col_labels) - prev = pos / tot if tot > 0 else 0.0 - subtask_stats[col] = { - "total": tot, - "positive": pos, - "negative": tot - pos, - "prevalence": prev, - } - prevalences.append(prev) + elif task_name in labels_df.columns: + labels = labels_df[task_name].drop_nulls() + total = len(labels) + + if task_type == "regression": + labels_float = labels.cast(pl.Float64) + std = labels_float.std() stats[task_name] = { - "total": len(self.labels_df), - "n_labels": len(multilabel_cols), - "mean_prevalence": ( - sum(prevalences) / len(prevalences) if prevalences else 0.0 - ), - "subtasks": subtask_stats, + "task_type": "regression", + "total": total, + "mean": float(labels_float.mean()) if total > 0 else 0.0, + "std": float(std) if std is not None else 0.0, + "min": float(labels_float.min()) if total > 0 else 0.0, + "max": float(labels_float.max()) if total > 0 else 0.0, } + elif task_type == "multiclass": + class_counts = { + str(label): count + for label, count in sorted(Counter(labels.to_list()).items()) + } + stats[task_name] = { + "task_type": "multiclass", + "total": total, + "n_classes": len(class_counts), + "class_counts": class_counts, + } + else: + positive = (labels == 1).sum() + stats[task_name] = { + "task_type": "binary", + "total": total, + "positive": positive, + "negative": total - positive, + "prevalence": positive / total if total > 0 else 0.0, + } + elif multilabel_cols: + subtask_stats = {} + prevalences = [] + for col in multilabel_cols: + col_labels = labels_df[col].drop_nulls() + pos = (col_labels == 1).sum() + tot = len(col_labels) + prev = pos / tot if tot > 0 else 0.0 + subtask_stats[col] = { + "task_type": "binary", + "total": tot, + "positive": pos, + "negative": tot - pos, + "prevalence": prev, + } + prevalences.append(prev) + stats[task_name] = { + "task_type": "multilabel", + "total": len(labels_df), + "n_labels": len(multilabel_cols), + "mean_prevalence": sum(prevalences) / len(prevalences) if prevalences else 0.0, + "subtasks": subtask_stats, + } return stats def get_preprocessing_stages(self, idx: int) -> Dict[str, Dict[str, torch.Tensor]]: """Get intermediate preprocessing stages for a single sample. - This method re-runs the preprocessing pipeline for a specific sample, - capturing the tensor at each stage. Useful for debugging to see exactly - what transformations are applied to the data. + Reloads raw data from Parquet on demand (the in-memory DataFrame + is freed after tensor loading to save ~3-5 GB RAM). Debug only. The stages are: - grid: Raw 2D tensor (seq_length, n_features) with NaN for missing @@ -903,7 +748,7 @@ def get_preprocessing_stages(self, idx: int) -> Dict[str, Dict[str, torch.Tensor - 'mask': The observation mask (same for all stages) Example: - >>> dataset = ICUDataset("data/processed/mimic-iv-demo") + >>> dataset = ICUDataset("data/processed/miiv") >>> stages = dataset.get_preprocessing_stages(0) >>> stages['grid']['timeseries'].shape torch.Size([48, 9]) @@ -914,46 +759,23 @@ def get_preprocessing_stages(self, idx: int) -> Dict[str, Dict[str, torch.Tensor """ stay_id = self.stay_ids[idx] - # Get raw nested list data from original DataFrame - row = self.timeseries_df.filter(pl.col("stay_id") == stay_id).row(0, named=True) + # Reload from parquet (timeseries_df freed after init) + timeseries_df = pl.read_parquet(self.data_dir / "timeseries.parquet") + if self._excluded_stay_ids: + timeseries_df = timeseries_df.filter( + ~pl.col("stay_id").is_in(list(self._excluded_stay_ids)) + ) + + row = timeseries_df.filter(pl.col("stay_id") == stay_id).row(0, named=True) raw_ts = row["timeseries"] raw_mask = row["mask"] - # ===================================================================== - # Stage 1: GRID - Convert to tensor with padding/truncation - # ===================================================================== - actual_len = len(raw_ts) - if actual_len >= self.seq_length: - grid_ts = np.array(raw_ts[: self.seq_length], dtype=np.float32) - mask_arr = np.array(raw_mask[: self.seq_length], dtype=bool) - else: - grid_ts = np.full((self.seq_length, self.n_features), np.nan, dtype=np.float32) - mask_arr = np.zeros((self.seq_length, self.n_features), dtype=bool) - grid_ts[:actual_len] = np.array(raw_ts, dtype=np.float32) - mask_arr[:actual_len] = np.array(raw_mask, dtype=bool) - - grid_tensor = torch.from_numpy(grid_ts) - mask_tensor = torch.from_numpy(mask_arr) - - # ===================================================================== - # Stage 2: NORMALIZED - z-score normalize then impute - # ===================================================================== - if self.normalize: - normalized_tensor = (grid_tensor - self.feature_means) / self.feature_stds - normalized_tensor = torch.nan_to_num(normalized_tensor, nan=0.0) - else: - normalized_tensor = grid_tensor.clone() - for f in range(self.n_features): - nan_mask = torch.isnan(normalized_tensor[:, f]) - normalized_tensor[:, f][nan_mask] = self.feature_means[f] - - return { - "grid": { - "timeseries": grid_tensor, - "mask": mask_tensor, - }, - "normalized": { - "timeseries": normalized_tensor, - "mask": mask_tensor, - }, - } + return compute_single_sample_stages( + raw_ts, + raw_mask, + self.seq_length, + self.n_features, + self.feature_means, + self.feature_stds, + self.normalize, + ) diff --git a/src/slices/data/extractors/base.py b/src/slices/data/extractors/base.py index ca9601f..02781bb 100644 --- a/src/slices/data/extractors/base.py +++ b/src/slices/data/extractors/base.py @@ -11,7 +11,7 @@ from contextlib import contextmanager from pathlib import Path from tempfile import NamedTemporaryFile -from typing import Callable, Dict, Generator, List, Optional +from typing import Any, Callable, Dict, Generator, List, Optional import polars as pl import portalocker @@ -19,7 +19,7 @@ from pydantic import BaseModel, Field, field_validator from rich.console import Console -from slices.constants import EXTRACTION_BATCH_SIZE, MIN_STAY_HOURS, SEQ_LENGTH_HOURS +from slices.constants import EXTRACTION_BATCH_SIZE, MIN_STAY_HOURS, SEQ_LENGTH_HOURS, THESIS_TASKS from slices.data.labels import LabelBuilder, LabelBuilderFactory, LabelConfig from slices.data.utils import get_package_data_dir @@ -47,7 +47,7 @@ class ExtractorConfig(BaseModel): seq_length_hours: int = Field(default=SEQ_LENGTH_HOURS, gt=0) feature_set: str = "core" # core | extended tasks_dir: Optional[str] = None - tasks: List[str] = Field(default_factory=lambda: ["mortality_24h", "mortality_hospital"]) + tasks: List[str] = Field(default_factory=lambda: list(THESIS_TASKS)) min_stay_hours: int = Field(default=MIN_STAY_HOURS, ge=0) batch_size: int = Field(default=EXTRACTION_BATCH_SIZE, gt=0) categories: Optional[List[str]] = None # e.g., ["vitals_dense"] for subset extraction @@ -92,6 +92,7 @@ def __init__(self, config: ExtractorConfig) -> None: self.config = config self.parquet_root = Path(config.parquet_root) self.output_dir = Path(config.output_dir) + self._label_quality_stats: Dict[str, Dict[str, Any]] = {} # Validate parquet directory exists if not self.parquet_root.exists(): @@ -132,8 +133,7 @@ def _load_task_configs(self, task_names: List[str]) -> List[LabelConfig]: for task_name in task_names: config_file = tasks_path / f"{task_name}.yaml" if not config_file.exists(): - console.print(f"[yellow]Warning: Task config not found: {config_file}[/yellow]") - continue + raise FileNotFoundError(f"Task config not found: {config_file}") with open(config_file) as f: config_dict = yaml.safe_load(f) @@ -147,6 +147,25 @@ def _load_task_configs(self, task_names: List[str]) -> List[LabelConfig]: return task_configs + def _build_label_manifest( + self, task_configs: Optional[List[LabelConfig]] = None + ) -> Dict[str, Dict]: + """Build the semantic label manifest for the current extraction config.""" + manifest: Dict[str, Dict] = {} + for task_config in task_configs or self._load_task_configs(self.config.tasks or []): + builder = LabelBuilderFactory.create(task_config) + manifest[task_config.task_name] = { + "builder_version": builder.SEMANTIC_VERSION, + "config_hash": LabelBuilder.config_hash(task_config), + "label_sources": sorted(task_config.label_sources), + "supported_datasets": ( + sorted(task_config.supported_datasets) + if task_config.supported_datasets is not None + else None + ), + } + return manifest + def _validate_observation_window(self, task_config: LabelConfig) -> None: """Validate that task's observation_window_hours matches extraction seq_length_hours. @@ -186,6 +205,14 @@ def _get_dataset_name(self) -> str: """ pass + def _get_upstream_source_signature(self) -> Optional[dict]: + """Return a fingerprint of upstream inputs used for extraction. + + Subclasses can override this to make resume safety depend on upstream + source identity as well as Python-side extraction config. + """ + return None + @abstractmethod def extract_stays(self) -> pl.DataFrame: """Extract ICU stay metadata (stay_id, patient_id, times). @@ -259,8 +286,9 @@ def _validate_stays(self, stays: pl.DataFrame) -> None: # Check for negative or invalid LOS invalid_los = stays.filter((pl.col("los_days").is_null()) | (pl.col("los_days") < 0)) if len(invalid_los) > 0: - console.print( - f"[yellow]Warning: Found {len(invalid_los)} stays with invalid LOS[/yellow]" + raise ValueError( + f"Found {len(invalid_los)} stays with invalid LOS. " + "Extraction must fail closed until the upstream stay metadata is fixed." ) def _validate_timeseries( @@ -282,7 +310,10 @@ def _validate_timeseries( missing = expected_stay_ids - timeseries_stay_ids if missing: - console.print(f"[yellow]Warning: {len(missing)} stays have no timeseries data[/yellow]") + raise ValueError( + f"Found {len(missing)} stays with no timeseries data. " + "Extraction must fail closed until upstream timeseries coverage is complete." + ) # Check that expected features are present timeseries_features = { @@ -292,8 +323,10 @@ def _validate_timeseries( } missing_features = set(feature_names) - timeseries_features if missing_features: - console.print( - f"[yellow]Warning: Missing features in timeseries: {missing_features}[/yellow]" + raise ValueError( + "Missing expected timeseries features: " + f"{sorted(missing_features)}. " + "Extraction must fail closed until the upstream export is complete." ) def _validate_labels(self, labels: pl.DataFrame, stay_ids: List[int]) -> None: @@ -315,13 +348,45 @@ def _validate_labels(self, labels: pl.DataFrame, stay_ids: List[int]) -> None: # Check for missing labels missing = expected_stay_ids - label_stay_ids if missing: - console.print(f"[yellow]Warning: {len(missing)} stays have no labels[/yellow]") + raise ValueError( + f"Found {len(missing)} stays with no labels. " + "Extraction must fail closed until label coverage is complete." + ) # Check for extra labels (shouldn't happen, but worth checking) extra = label_stay_ids - expected_stay_ids if extra: console.print(f"[yellow]Warning: {len(extra)} labels for stays not in dataset[/yellow]") + def _emit_label_quality_warnings( + self, task_config: LabelConfig, task_labels: pl.DataFrame + ) -> None: + """Emit non-fatal quality warnings for extracted labels.""" + if not task_config.quality_checks or "label" not in task_labels.columns: + return + + total = len(task_labels) + if total == 0: + return + + max_missing_percentage = task_config.quality_checks.get("max_missing_percentage") + if max_missing_percentage is None: + return + + missing = task_labels["label"].null_count() + missing_percentage = (missing / total) * 100.0 + + if missing_percentage > float(max_missing_percentage): + console.print( + "[yellow]Warning: " + f"Task '{task_config.task_name}' has {missing_percentage:.1f}% missing labels " + f"({missing}/{total}), exceeding the configured maximum of " + f"{float(max_missing_percentage):.1f}%. " + "This usually means the upstream label-source coverage is insufficient " + "(for example, AKI without enough post-observation creatinine data)." + "[/yellow]" + ) + def extract_labels(self, stay_ids: List[int], task_configs: List[LabelConfig]) -> pl.DataFrame: """Extract labels for multiple downstream tasks. @@ -341,6 +406,7 @@ def extract_labels(self, stay_ids: List[int], task_configs: List[LabelConfig]) - DataFrame with stay_id and one column per task (named by task_name). """ # Step 1: Identify all required data sources + self._label_quality_stats = {} required_sources = set() for task_config in task_configs: required_sources.update(task_config.label_sources) @@ -381,6 +447,11 @@ def extract_labels(self, stay_ids: List[int], task_configs: List[LabelConfig]) - # Compute labels task_labels = builder.build_labels(raw_data) + self._emit_label_quality_warnings(task_config, task_labels) + if "label" in task_labels.columns: + self._label_quality_stats[task_config.task_name] = builder.build_quality_stats( + task_labels + ) # For single-label tasks, the builder returns a 'label' column # that we rename to the task name. @@ -572,6 +643,24 @@ def _check_existing_extraction(self) -> Optional[Dict[str, pl.DataFrame]]: "Will overwrite.[/yellow]" ) return None + + current_signature = self._get_upstream_source_signature() + existing_signature = existing_metadata.get("upstream_source_signature") + if current_signature != existing_signature: + console.print( + "[yellow]Warning: Existing extraction was built from different " + "upstream inputs. Will overwrite.[/yellow]" + ) + return None + + current_label_manifest = self._build_label_manifest() + existing_label_manifest = existing_metadata.get("label_manifest") + if existing_label_manifest != current_label_manifest: + console.print( + "[yellow]Warning: Existing extraction was built from a different " + "task config or label builder version. Will overwrite.[/yellow]" + ) + return None except Exception: # If metadata can't be read, assume we should overwrite return None diff --git a/src/slices/data/extractors/ricu.py b/src/slices/data/extractors/ricu.py index 9f00ff6..91f511b 100644 --- a/src/slices/data/extractors/ricu.py +++ b/src/slices/data/extractors/ricu.py @@ -6,7 +6,8 @@ Usage: # Step 1: Run R extraction Rscript scripts/preprocessing/extract_with_ricu.R \ - --dataset miiv --output_dir data/ricu_output/miiv + --dataset miiv --output_dir data/ricu_output/miiv \ + --raw_export_horizon_hours 48 # Step 2: Run Python processing uv run python scripts/preprocessing/extract_ricu.py \ @@ -14,7 +15,8 @@ """ from datetime import datetime -from typing import Dict, List, Optional +from pathlib import Path +from typing import Dict, List, Optional, Union import polars as pl import yaml @@ -25,11 +27,27 @@ TextColumn, ) -from slices.constants import FEATURE_BLOCKLIST +from slices.constants import FEATURE_BLOCKLIST, LABEL_HORIZON_HOURS from slices.data.config_schemas import TimeSeriesConceptConfig +from slices.data.labels import LabelBuilderFactory, LabelConfig from .base import BaseExtractor, ExtractorConfig + +def _resolve_parquet_path(path: Path) -> Union[Path, str]: + """Resolve a parquet path that may be a single file or a directory of parts. + + The R extraction script writes timeseries as a directory of part files + to avoid exceeding memory limits. Polars scan_parquet accepts both a + single file and a glob pattern. + """ + if path.is_file(): + return path + if path.is_dir(): + return str(path / "*.parquet") + raise FileNotFoundError(f"Parquet path not found: {path}") + + console = Console() @@ -37,9 +55,9 @@ class RicuExtractor(BaseExtractor): """Reads pre-extracted RICU parquet output. Expects parquet_root to contain: - ricu_timeseries.parquet, ricu_stays.parquet, - ricu_mortality.parquet, ricu_diagnoses.parquet, - ricu_metadata.yaml + ricu_timeseries.parquet (file or directory of part files), + ricu_stays.parquet, ricu_mortality.parquet, + ricu_diagnoses.parquet, ricu_metadata.yaml Unlike other extractors, this class does NOT use DuckDB or concept YAML files. All extraction and binning is handled by the R script; this class @@ -57,6 +75,7 @@ def __init__(self, config: ExtractorConfig) -> None: self._stays_cache: Optional[pl.DataFrame] = None self._metadata: Optional[dict] = None + self._task_configs_cache: Optional[List[LabelConfig]] = None if not self.parquet_root.exists(): raise ValueError(f"RICU output directory not found: {self.parquet_root}") @@ -67,6 +86,112 @@ def __init__(self, config: ExtractorConfig) -> None: raise ValueError(f"ricu_metadata.yaml not found in {self.parquet_root}") with open(metadata_path) as f: self._metadata = yaml.safe_load(f) + self._validate_ricu_horizon() + + def _get_raw_export_horizon_hours(self) -> int: + """Return the upstream RICU export horizon in hours. + + New exports should store ``raw_export_horizon_hours`` explicitly. Older + exports only have ``seq_length_hours``; keep supporting that field for + backward compatibility. + """ + raw_horizon = self._metadata.get("raw_export_horizon_hours") + if raw_horizon is None: + raw_horizon = self._metadata.get("seq_length_hours") + + if raw_horizon is None: + raise ValueError( + "ricu_metadata.yaml is missing both raw_export_horizon_hours and " + "seq_length_hours; cannot validate the export horizon safely." + ) + + return int(raw_horizon) + + def _get_task_configs_cached(self) -> List[LabelConfig]: + """Load and cache task configs for horizon validation and metadata.""" + if self._task_configs_cache is None: + self._task_configs_cache = self._load_task_configs(self.config.tasks) + return self._task_configs_cache + + def _get_required_raw_export_horizon_hours( + self, task_configs: Optional[List[LabelConfig]] = None + ) -> int: + """Return the minimum upstream raw export horizon required by this extraction.""" + required_horizon = int(self.config.seq_length_hours) + for task_config in task_configs or self._get_task_configs_cached(): + builder = LabelBuilderFactory.create(task_config) + required_horizon = max( + required_horizon, + int(builder.required_raw_timeseries_horizon_hours()), + ) + + return required_horizon + + def _validate_ricu_horizon(self) -> None: + """Validate that the Python extraction horizon does not exceed the R export.""" + raw_export_horizon = self._get_raw_export_horizon_hours() + if self.config.seq_length_hours > raw_export_horizon: + raise ValueError( + "Python extraction requests seq_length_hours=" + f"{self.config.seq_length_hours}, but the upstream RICU export only " + f"contains {raw_export_horizon} hours. Re-run the R export with a longer " + "horizon or lower the Python extraction seq_length_hours." + ) + + task_requirements = [] + for task_config in self._get_task_configs_cached(): + builder = LabelBuilderFactory.create(task_config) + required_horizon = int(builder.required_raw_timeseries_horizon_hours()) + if required_horizon > self.config.seq_length_hours: + task_requirements.append((task_config.task_name, required_horizon)) + + required_raw_export_horizon = self._get_required_raw_export_horizon_hours( + self._get_task_configs_cached() + ) + if raw_export_horizon < required_raw_export_horizon: + requirement_str = ", ".join( + f"{task_name}={required_horizon}h" + for task_name, required_horizon in task_requirements + ) + raise ValueError( + "Active task labels require a longer upstream raw timeseries horizon than " + f"the current RICU export provides. Upstream export: {raw_export_horizon}h; " + f"required: {required_raw_export_horizon}h. " + f"Task-specific requirements: {requirement_str}. " + "Keep the model input at 24h if desired, but re-run the R export with a " + "longer raw horizon so forward-looking labels can be built safely." + ) + + def _iter_upstream_files(self) -> List[Path]: + """Return the upstream files that define the extraction contents.""" + files: List[Path] = [] + for path in sorted(self.parquet_root.glob("ricu_*")): + if path.is_file(): + files.append(path) + elif path.is_dir(): + files.extend(sorted(p for p in path.rglob("*") if p.is_file())) + return files + + def _get_upstream_source_signature(self) -> dict: + """Fingerprint upstream RICU inputs for safe resume behavior.""" + files = [] + for path in self._iter_upstream_files(): + stat = path.stat() + files.append( + { + "path": str(path.relative_to(self.parquet_root)), + "size_bytes": stat.st_size, + "mtime_ns": stat.st_mtime_ns, + } + ) + + return { + "dataset": self._metadata.get("dataset"), + "ricu_seq_length_hours": int(self._metadata["seq_length_hours"]), + "ricu_raw_export_horizon_hours": self._get_raw_export_horizon_hours(), + "feature_blocklist": sorted(FEATURE_BLOCKLIST), + "files": files, + } def _get_dataset_name(self) -> str: return self._metadata["dataset"] @@ -85,7 +210,7 @@ def extract_timeseries(self, stay_ids: List[int]) -> pl.DataFrame: Categorical string columns (e.g. avpu, mech_vent) are mapped to ordinal floats using CATEGORICAL_ENCODINGS. """ - path = self.parquet_root / "ricu_timeseries.parquet" + path = _resolve_parquet_path(self.parquet_root / "ricu_timeseries.parquet") df = pl.scan_parquet(path).filter(pl.col("stay_id").is_in(stay_ids)).collect() return self._encode_categorical_columns(df) @@ -141,10 +266,58 @@ def extract_data_source(self, source_name: str, stay_ids: List[int]) -> pl.DataF raise ValueError( f"Unknown data source '{source_name}'. " f"Available: {list(dispatch.keys())}" ) - path = self.parquet_root / dispatch[source_name] + path = _resolve_parquet_path(self.parquet_root / dispatch[source_name]) df = pl.scan_parquet(path).filter(pl.col("stay_id").is_in(stay_ids)).collect() if source_name == "timeseries": df = self._encode_categorical_columns(df) + elif source_name == "mortality_info": + df = self._migrate_mortality_schema(df) + return df + + @staticmethod + def _migrate_mortality_schema(df: pl.DataFrame) -> pl.DataFrame: + """Migrate legacy mortality parquet (date_of_death only) to new schema. + + If the parquet has the old single-column format, infer precision from + the Polars dtype so the label builder can use precision-aware logic. + """ + if "death_time_precision" in df.columns: + return df # already new schema + + if "date_of_death" not in df.columns: + return df + + dod_dtype = df["date_of_death"].dtype + if dod_dtype == pl.Date: + # Date-only column — treat as date precision + df = df.with_columns( + pl.lit(None).cast(pl.Datetime("us", "UTC")).alias("death_time"), + pl.col("date_of_death").cast(pl.Date).alias("death_date"), + pl.when(pl.col("date_of_death").is_not_null()) + .then(pl.lit("date")) + .otherwise(pl.lit(None)) + .alias("death_time_precision"), + pl.when(pl.col("date_of_death").is_not_null()) + .then(pl.lit("legacy")) + .otherwise(pl.lit(None)) + .alias("death_source"), + ) + else: + # Legacy datetimes are ambiguous: older exports may have stored + # date-only values as midnight-cast timestamps. Treat them as date + # precision unless explicit precision metadata is present. + df = df.with_columns( + pl.lit(None).cast(pl.Datetime("us", "UTC")).alias("death_time"), + pl.col("date_of_death").cast(pl.Date).alias("death_date"), + pl.when(pl.col("date_of_death").is_not_null()) + .then(pl.lit("date")) + .otherwise(pl.lit(None)) + .alias("death_time_precision"), + pl.when(pl.col("date_of_death").is_not_null()) + .then(pl.lit("legacy")) + .otherwise(pl.lit(None)) + .alias("death_source"), + ) return df def run(self) -> None: @@ -183,9 +356,23 @@ def run(self) -> None: stays = self.extract_stays() progress.update(task, completed=True) + # Filter by valid/minimum stay length. RICU exports can include source + # records with missing or negative discharge-derived LOS; those rows + # cannot enter the benchmark cohort, so exclude them before validating + # the extracted stay universe. + invalid_los = stays.filter((pl.col("los_days").is_null()) | (pl.col("los_days") < 0)) + if len(invalid_los) > 0: + console.print( + f"[yellow]Warning: Excluding {len(invalid_los)} stays with invalid LOS " + "before benchmark cohort filtering.[/yellow]" + ) + stays_with_valid_los = stays.filter( + pl.col("los_days").is_not_null() & (pl.col("los_days") >= 0) + ) + # Filter by minimum stay length min_los_days = self.config.min_stay_hours / 24.0 - stays_filtered = stays.filter(pl.col("los_days") >= min_los_days) + stays_filtered = stays_with_valid_los.filter(pl.col("los_days") >= min_los_days) console.print(f"Found {len(stays)} ICU stays") console.print( @@ -243,10 +430,31 @@ def run(self) -> None: ) progress.update(task, completed=True) + obs_counts = dense_timeseries.select( + "stay_id", + pl.col("mask").list.eval(pl.element().list.sum()).list.sum().alias("_n_observed"), + ) + empty_stay_ids = obs_counts.filter(pl.col("_n_observed") == 0)["stay_id"].to_list() + if empty_stay_ids: + console.print( + f"[yellow]Excluding {len(empty_stay_ids)} stays with zero observations " + "in the model input window.[/yellow]" + ) + dense_timeseries = dense_timeseries.filter(~pl.col("stay_id").is_in(empty_stay_ids)) + stays_filtered = stays_filtered.filter(~pl.col("stay_id").is_in(empty_stay_ids)) + self._stays_cache = stays_filtered + stay_ids = stays_filtered["stay_id"].to_list() + + if not stay_ids: + console.print( + "[red]Error: No stays remaining after zero-observation filtering![/red]" + ) + return + # ----------------------------------------------------------------- # Step 3: Extract labels # ----------------------------------------------------------------- - task_configs = self._load_task_configs(self.config.tasks) + task_configs = self._get_task_configs_cached() task_names = [tc.task_name for tc in task_configs] if task_configs: @@ -301,6 +509,8 @@ def run(self) -> None: self._validate_labels(labels, stay_ids) + label_manifest = self._build_label_manifest(task_configs) + metadata = { "dataset": self._get_dataset_name(), "feature_set": self.config.feature_set, @@ -308,9 +518,20 @@ def run(self) -> None: "feature_names": feature_names, "n_features": len(feature_names), "seq_length_hours": self.config.seq_length_hours, + "input_seq_length_hours": self.config.seq_length_hours, + "label_horizon_hours": LABEL_HORIZON_HOURS, + "raw_export_horizon_hours": self._get_raw_export_horizon_hours(), + "required_raw_export_horizon_hours": self._get_required_raw_export_horizon_hours( + task_configs + ), "min_stay_hours": self.config.min_stay_hours, "task_names": task_names, "n_stays": len(stays_filtered), + "zero_observation_stays_excluded": True, + "n_zero_observation_stays_excluded": len(empty_stay_ids), + "label_manifest": label_manifest, + "label_quality_stats": self._label_quality_stats, + "upstream_source_signature": self._get_upstream_source_signature(), "extraction_config": { "parquet_root": str(self.parquet_root), "output_dir": str(self.output_dir), diff --git a/src/slices/data/labels/aki.py b/src/slices/data/labels/aki.py index d497be9..2f46f68 100644 --- a/src/slices/data/labels/aki.py +++ b/src/slices/data/labels/aki.py @@ -1,10 +1,12 @@ """AKI (Acute Kidney Injury) label builder using KDIGO criteria.""" import logging -from typing import Dict +from typing import Any, Dict import polars as pl +from slices.constants import SEQ_LENGTH_HOURS + from .base import LabelBuilder logger = logging.getLogger(__name__) @@ -17,11 +19,42 @@ class AKILabelBuilder(LabelBuilder): - Creatinine rise >= 0.3 mg/dL within any 48h window - Creatinine >= 1.5x baseline within 7 days of baseline measurement - Baseline: minimum creatinine in first 24h (standard clinical proxy). - Only evaluates criteria AFTER the observation window ends. - Label = null if no creatinine measurements available. + Baseline: minimum creatinine in first baseline_window_hours + (default benchmark observation window, 24h). + Detection window: hours [observation_window_hours, observation_window_hours + + prediction_window_hours) — forward-looking to avoid data leakage. + Label = null if no creatinine in the prediction window. """ + SEMANTIC_VERSION = "1.1.0" + + def required_raw_timeseries_horizon_hours(self) -> int: + """AKI needs baseline data plus post-observation creatinine measurements.""" + prediction_hours = self.config.prediction_window_hours + if prediction_hours is None: + raise ValueError( + "AKI task requires prediction_window_hours to define the forward-looking " + "detection window." + ) + + obs_hours = self.config.observation_window_hours or SEQ_LENGTH_HOURS + return int(obs_hours + prediction_hours) + + def build_quality_stats(self, labels: pl.DataFrame) -> Dict[str, Any]: + """Add AKI-specific null-reason counts to the default quality stats.""" + null_reason_counts = { + "no_creatinine_or_baseline": int( + self._last_quality_stats.get("no_creatinine_or_baseline", 0) + ), + "no_post_obs_creatinine": int( + self._last_quality_stats.get("no_post_obs_creatinine", 0) + ), + } + stats = super().build_quality_stats(labels) + stats["null_reason_counts"] = null_reason_counts + self._last_quality_stats = stats + return stats + def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: """Build AKI labels from stays and timeseries data. @@ -36,6 +69,10 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: timeseries = raw_data["timeseries"] if len(stays) == 0: + self._last_quality_stats = { + "no_creatinine_or_baseline": 0, + "no_post_obs_creatinine": 0, + } return pl.DataFrame( { "stay_id": pl.Series([], dtype=pl.Int64), @@ -43,18 +80,32 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: } ) - obs_hours = self.config.observation_window_hours or 48 crea_col = self.config.label_params.get("creatinine_col", "crea") - baseline_hours = self.config.label_params.get("baseline_window_hours", 24) + baseline_hours = self.config.label_params.get("baseline_window_hours", SEQ_LENGTH_HOURS) abs_threshold = self.config.label_params.get("absolute_rise_threshold", 0.3) rel_threshold = self.config.label_params.get("relative_rise_threshold", 1.5) rel_window_hours = self.config.label_params.get("relative_window_hours", 168) + # Forward-looking prediction window + obs_hours = self.config.observation_window_hours + prediction_hours = self.config.prediction_window_hours + if not prediction_hours: + raise ValueError( + "AKI task requires prediction_window_hours to define the forward-looking " + "detection window. Set it in the task YAML (e.g. prediction_window_hours: 24)." + ) + detection_start = obs_hours + detection_end = obs_hours + prediction_hours + stay_ids = stays["stay_id"].to_list() # Validate creatinine column exists if crea_col not in timeseries.columns: logger.warning(f"Creatinine column '{crea_col}' not in timeseries; all labels null") + self._last_quality_stats = { + "no_creatinine_or_baseline": len(stay_ids), + "no_post_obs_creatinine": 0, + } return pl.DataFrame( { "stay_id": pl.Series(stay_ids, dtype=pl.Int64), @@ -76,16 +127,16 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: .agg(pl.col(crea_col).min().alias("baseline")) ) - # Post-observation creatinine values - post_obs = crea_ts.filter(pl.col("hour") >= obs_hours) + # Prediction window creatinine values (detection window) + post_obs = crea_ts.filter( + (pl.col("hour") >= detection_start) & (pl.col("hour") < detection_end) + ) - # Join baselines with post-observation data + # Join baselines with post-baseline data post_with_baseline = post_obs.join(baselines, on="stay_id", how="inner") # --- KDIGO Criterion 1: Relative rise (>= 1.5x baseline within 7 days) --- - # The 7-day window is relative to the baseline measurement period. - # Baseline is measured in first baseline_hours, so the 7-day window extends - # from hour 0 to hour rel_window_hours. + # post_with_baseline already filtered to [detection_start, detection_end) rel_aki = ( post_with_baseline.filter(pl.col("hour") < rel_window_hours) .filter(pl.col(crea_col) >= pl.col("baseline") * rel_threshold) @@ -95,8 +146,11 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: ) # --- KDIGO Criterion 2: Absolute rise (>= 0.3 mg/dL within 48h window) --- - # Self-join to find pairs within 48h windows - abs_aki = self._check_absolute_criterion_vectorized(post_obs, crea_col, abs_threshold) + # Use all creatinine data but only flag rises where the later measurement + # falls in the detection window [detection_start, detection_end). + abs_aki = self._check_absolute_criterion_vectorized( + crea_ts, crea_col, abs_threshold, detection_start, detection_end + ) # Combine criteria: AKI = 1 if either criterion is met stay_df = pl.DataFrame({"stay_id": pl.Series(stay_ids, dtype=pl.Int64)}) @@ -107,7 +161,7 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: # Stays with baseline (needed for evaluation) stays_with_baseline = baselines.select("stay_id").unique() - # Stays with post-obs data + # Stays with post-baseline data stays_with_post_obs = post_obs.select("stay_id").unique() result = ( @@ -130,16 +184,28 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: .join(abs_aki, on="stay_id", how="left") ) + missing_crea_or_baseline = result.filter( + pl.col("has_crea").is_null() | pl.col("has_baseline").is_null() + ).height + missing_post_obs = result.filter( + pl.col("has_crea").is_not_null() + & pl.col("has_baseline").is_not_null() + & pl.col("has_post_obs").is_null() + ).height + self._last_quality_stats = { + "no_creatinine_or_baseline": int(missing_crea_or_baseline), + "no_post_obs_creatinine": int(missing_post_obs), + } + # Build label: - # - null if no creatinine data or no baseline - # - 0 if no post-obs data (survived observation without AKI evidence after) + # - null if no creatinine data, no baseline, or no data in prediction window # - 1 if either AKI criterion met # - 0 otherwise result = result.with_columns( pl.when(pl.col("has_crea").is_null() | pl.col("has_baseline").is_null()) .then(None) .when(pl.col("has_post_obs").is_null()) - .then(0) + .then(None) .when(pl.col("rel_aki").is_not_null() | pl.col("abs_aki").is_not_null()) .then(1) .otherwise(0) @@ -156,26 +222,31 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: n_no_aki = result_df.filter(pl.col("label") == 0).height n_null = result_df["label"].null_count() logger.info( - f"AKI labels: {n_aki} AKI, {n_no_aki} no AKI, {n_null} excluded (no creatinine)" + f"AKI labels: {n_aki} AKI, {n_no_aki} no AKI, " + f"{n_null} excluded ({missing_crea_or_baseline} no creatinine/baseline, " + f"{missing_post_obs} no creatinine in prediction window)" ) return result_df def _check_absolute_criterion_vectorized( self, - post_obs: pl.DataFrame, + crea_ts: pl.DataFrame, crea_col: str, abs_threshold: float, + detection_start_hour: int, + detection_end_hour: int, ) -> pl.DataFrame: """Check KDIGO absolute rise criterion using vectorized operations. Finds any pair of creatinine values within a 48h window where the later - value exceeds the earlier by >= abs_threshold. + value exceeds the earlier by >= abs_threshold and the later measurement + falls in the detection window [detection_start_hour, detection_end_hour). Returns: DataFrame with stay_id column for stays meeting the criterion. """ - if len(post_obs) == 0: + if len(crea_ts) == 0: return pl.DataFrame( { "stay_id": pl.Series([], dtype=pl.Int64), @@ -184,21 +255,24 @@ def _check_absolute_criterion_vectorized( ) # Self-join within each stay: pair each measurement with later ones - left = post_obs.select( + left = crea_ts.select( pl.col("stay_id"), pl.col("hour").alias("hour_i"), pl.col(crea_col).alias("crea_i"), ) - right = post_obs.select( + right = crea_ts.select( pl.col("stay_id"), pl.col("hour").alias("hour_j"), pl.col(crea_col).alias("crea_j"), ) - # Join on stay_id and filter: j > i, within 48h, rise >= threshold + # Join on stay_id and filter: j > i, within 48h, rise >= threshold, + # and the later measurement is in the detection window [start, end) pairs = left.join(right, on="stay_id", how="inner").filter( (pl.col("hour_j") > pl.col("hour_i")) & (pl.col("hour_j") - pl.col("hour_i") <= 48) + & (pl.col("hour_j") >= detection_start_hour) + & (pl.col("hour_j") < detection_end_hour) & (pl.col("crea_j") - pl.col("crea_i") >= abs_threshold) ) diff --git a/src/slices/data/labels/base.py b/src/slices/data/labels/base.py index db40af9..830ea19 100644 --- a/src/slices/data/labels/base.py +++ b/src/slices/data/labels/base.py @@ -1,8 +1,10 @@ """Base classes for label extraction.""" +import hashlib +import json from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import Dict, List, Optional +from typing import Any, Dict, List, Optional import polars as pl @@ -26,6 +28,7 @@ class LabelConfig: # Label definition label_sources: List[str] = field(default_factory=list) # Required data sources label_params: Dict = field(default_factory=dict) # Task-specific parameters + quality_checks: Dict = field(default_factory=dict) # Optional alert thresholds only # Evaluation metrics primary_metric: str = "auroc" @@ -47,6 +50,33 @@ class LabelBuilder(ABC): (e.g., mortality flags, creatinine values) into prediction labels. """ + SEMANTIC_VERSION: str = "1.0.0" + + @staticmethod + def config_hash(config: LabelConfig) -> str: + """Compute a deterministic hash of the label-affecting config fields. + + Returns: + 16-char hex digest of the config's label-relevant fields. + """ + # NOTE: quality_checks intentionally excluded because they only control + # warnings/analysis thresholds, not the label semantics themselves. + supported_datasets = ( + sorted(config.supported_datasets) if config.supported_datasets is not None else None + ) + hashable = { + "task_name": config.task_name, + "task_type": config.task_type, + "prediction_window_hours": config.prediction_window_hours, + "observation_window_hours": config.observation_window_hours, + "gap_hours": config.gap_hours, + "label_sources": sorted(config.label_sources), + "label_params": config.label_params, + "supported_datasets": supported_datasets, + } + content = json.dumps(hashable, sort_keys=True, default=str) + return hashlib.sha256(content.encode()).hexdigest()[:16] + def __init__(self, config: LabelConfig) -> None: """Initialize label builder with configuration. @@ -54,6 +84,64 @@ def __init__(self, config: LabelConfig) -> None: config: Label configuration specifying label definition. """ self.config = config + self._last_quality_stats: Dict[str, Any] = {} + + def required_raw_timeseries_horizon_hours(self) -> int: + """Return the raw timeseries horizon needed to build this task's labels. + + This is independent of the model input sequence length. The extractor uses + it to validate that the upstream export retains enough post-observation + data for forward-looking labels. + + Returns: + Maximum hour offset needed from the raw ``timeseries`` source. + Returns 0 for tasks that do not depend on raw timeseries labels. + """ + if "timeseries" not in self.config.label_sources: + return 0 + + return int(self.config.observation_window_hours or 0) + + def build_quality_stats(self, labels: pl.DataFrame) -> Dict[str, Any]: + """Build serializable task-level quality stats from extracted labels. + + Args: + labels: Builder output with ``stay_id`` and ``label`` columns. + + Returns: + Dictionary of quality stats suitable for persistence in metadata.yaml. + """ + total = len(labels) + if "label" not in labels.columns: + stats: Dict[str, Any] = {"total_stays": total} + self._last_quality_stats = stats + return stats + + null_count = labels["label"].null_count() + non_null = total - null_count + + stats = { + "total_stays": total, + "non_null_labels": non_null, + "null_labels": null_count, + "null_percentage": ((null_count / total) * 100.0) if total > 0 else 0.0, + } + + if self.config.task_type in {"binary", "binary_classification"}: + positives = labels.filter(pl.col("label") == 1).height + negatives = labels.filter(pl.col("label") == 0).height + stats.update( + { + "positive_labels": positives, + "negative_labels": negatives, + "positive_prevalence_non_null": ( + (positives / non_null) if non_null > 0 else None + ), + } + ) + + self._last_quality_stats = stats + return stats @abstractmethod def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: diff --git a/src/slices/data/labels/los.py b/src/slices/data/labels/los.py index 58b0c6e..3af32bc 100644 --- a/src/slices/data/labels/los.py +++ b/src/slices/data/labels/los.py @@ -17,6 +17,8 @@ class LOSLabelBuilder(LabelBuilder): Clips at 0.0 (stays exactly at boundary get 0). """ + SEMANTIC_VERSION = "1.0.0" + def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: self.validate_inputs(raw_data) stays = raw_data["stays"] diff --git a/src/slices/data/labels/mortality.py b/src/slices/data/labels/mortality.py index e24aec4..1712f90 100644 --- a/src/slices/data/labels/mortality.py +++ b/src/slices/data/labels/mortality.py @@ -27,22 +27,35 @@ class MortalityLabelBuilder(LabelBuilder): When observation_window_hours is None (legacy behavior): - Label = 1 if death occurred within prediction_window_hours of admission + + Precision-aware labeling: + When death_time_precision is available (new schema), uses it to determine + comparison logic per row: + - "timestamp": exact datetime comparison using death_time + - "date": interval-based conservative comparison using death_date + (treats as [00:00:00, 23:59:59.999] interval; boundary overlaps → null) + - "unknown" or missing: uses explicit death evidence and preserves + unknown outcomes as null """ + SEMANTIC_VERSION = "2.3.0" + def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: """Build mortality labels from stay and mortality data. Expected raw_data sources: - 'stays': stay_id, intime, outtime - - 'mortality_info': stay_id, date_of_death, hospital_expire_flag, - dischtime, discharge_location + - 'mortality_info': stay_id, death_time, death_date, death_time_precision, + death_source, hospital_expire_flag, dischtime, + discharge_location (and optionally legacy date_of_death) Args: raw_data: Dictionary with 'stays' and 'mortality_info' DataFrames. Returns: DataFrame with stay_id and binary label (1=died, 0=survived). - Label is null for stays where death occurred during observation window. + Label is null for stays where death occurred during observation window, + or where date-only precision cannot resolve boundary cases. """ self.validate_inputs(raw_data) @@ -61,6 +74,10 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: # Join mortality info with stays merged = stays.join(mortality, on="stay_id", how="left") + # Ensure precision-aware columns and nullable outcome evidence exist + # (backward compat with legacy data). + merged = self._ensure_precision_columns(merged) + # Compute label based on prediction window window_hours = self.config.prediction_window_hours obs_hours = self.config.observation_window_hours @@ -71,7 +88,6 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: window_hours is not None and window_hours != -1 and obs_hours is None ) needs_outtime = window_hours == -1 - needs_date_of_death = window_hours is not None if needs_intime and merged["intime"].null_count() == len(merged): raise ValueError( @@ -87,113 +103,426 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: "but all values are null. This dataset may not provide discharge timestamps." ) - if needs_date_of_death and merged["date_of_death"].null_count() == len(merged): - logger.warning( - f"Task '{self.config.task_name}': 'date_of_death' is null for all stays. " - "This dataset (e.g., eICU) may not provide death timestamps. " - "All windowed mortality labels will be null (excluded). " - "Consider using hospital-mortality (hospital_expire_flag) instead." - ) - - # Check dtype of date_of_death to handle DATE vs DATETIME properly - # This avoids edge cases where DATE type is cast to DATETIME with 00:00:00 - date_of_death_dtype = mortality["date_of_death"].dtype - is_date_type = date_of_death_dtype == pl.Date - if window_hours is None and obs_hours is None: # Hospital mortality, no observation window (legacy) - # Null hospital_expire_flag means outcome is unknown — keep as null. - # ICUDataset's handle_missing_labels='filter' will exclude these stays. - labels = merged.select( - [ - "stay_id", - pl.col("hospital_expire_flag").cast(pl.Int32).alias("label"), - ] - ) - - elif window_hours is None and obs_hours is not None: - # Hospital mortality with observation window exclusion - # Patients who died during observation are excluded (label=null) - obs_end_datetime = pl.col("intime") + pl.duration(hours=obs_hours) - left_icu_during_obs = pl.col("outtime") < obs_end_datetime - - if is_date_type: - obs_end = obs_end_datetime.cast(pl.Date) - died_during_obs = ( - pl.col("date_of_death").is_not_null() - & (pl.col("date_of_death") <= obs_end) - & left_icu_during_obs - ) - else: - died_during_obs = ( - pl.col("date_of_death").is_not_null() - & (pl.col("date_of_death") <= obs_end_datetime) - & left_icu_during_obs - ) - labels = merged.select( [ "stay_id", - pl.when(died_during_obs) - .then(None) - .when(pl.col("hospital_expire_flag") == 1) + pl.when(self._death_evidence_expr()) .then(1) - .otherwise(0) + .when(self._known_survivor_expr()) + .then(0) + .otherwise(None) .cast(pl.Int32) .alias("label"), ] ) + elif window_hours is None and obs_hours is not None: + # Hospital mortality with observation window exclusion + labels = self._build_hospital_mortality_with_obs(merged, obs_hours) + elif window_hours == -1 and obs_hours is None: # ICU mortality (died during or at ICU discharge), no observation window - if is_date_type: - # For DATE type, compare dates only (ignores time component) - comparison = pl.col("date_of_death") <= pl.col("outtime").cast(pl.Date) - else: - # For DATETIME type, compare full datetime for precision - comparison = pl.col("date_of_death") <= pl.col("outtime") - - labels = merged.select( - [ - "stay_id", - pl.when(pl.col("date_of_death").is_not_null() & comparison) - .then(1) - .otherwise(0) - .alias("label"), - ] - ) + labels = self._build_icu_mortality(merged) elif obs_hours is not None: - # Time-bounded mortality with observation window (recommended) - # Prediction window starts after observation + gap ends + # Time-bounded mortality with observation window (recommended path) labels = self._build_windowed_mortality_labels( - merged, obs_hours, gap_hours, window_hours, is_date_type + merged, obs_hours, gap_hours, window_hours ) else: - # Legacy: Time-bounded mortality from admission (e.g., 24h, 48h from intime) - if is_date_type: - # For DATE type, compare dates only - comparison = pl.col("date_of_death") <= ( - pl.col("intime") + pl.duration(hours=window_hours) - ).cast(pl.Date) + # Legacy: Time-bounded mortality from admission + labels = self._build_legacy_windowed(merged, window_hours) + + return labels + + @staticmethod + def _ensure_precision_columns(merged: pl.DataFrame) -> pl.DataFrame: + """Ensure precision-aware columns exist, migrating legacy data if needed.""" + merged = MortalityLabelBuilder._normalize_datetime_columns(merged) + merged = MortalityLabelBuilder._ensure_outcome_columns(merged) + + if "death_time_precision" in merged.columns: + # Keep death_time in the repo's canonical tz-naive datetime form. + if "death_time" in merged.columns: + dt_dtype = merged["death_time"].dtype + if dt_dtype not in ( + pl.Datetime, + pl.Datetime("us"), + pl.Datetime("ns"), + pl.Datetime("ms"), + ): + # Strip timezone if present, or cast to Datetime + merged = merged.with_columns( + pl.col("death_time").cast(pl.Datetime("us")).alias("death_time") + ) else: - # For DATETIME type, compare full datetime for precision - comparison = pl.col("date_of_death") <= ( - pl.col("intime") + pl.duration(hours=window_hours) + merged = merged.with_columns( + pl.lit(None).cast(pl.Datetime("us")).alias("death_time") ) + if "death_date" not in merged.columns: + merged = merged.with_columns(pl.lit(None).cast(pl.Date).alias("death_date")) + if "death_source" not in merged.columns: + merged = merged.with_columns(pl.lit(None).cast(pl.Utf8).alias("death_source")) + return merged + + # Legacy schema: only date_of_death exists + if "date_of_death" not in merged.columns: + # No death info at all — add empty precision columns + return merged.with_columns( + pl.lit(None).cast(pl.Datetime("us")).alias("death_time"), + pl.lit(None).cast(pl.Date).alias("death_date"), + pl.lit(None).cast(pl.Utf8).alias("death_time_precision"), + pl.lit(None).cast(pl.Utf8).alias("death_source"), + ) - labels = merged.select( - [ - "stay_id", - pl.when(pl.col("date_of_death").is_not_null() & comparison) - .then(1) - .otherwise(0) - .alias("label"), - ] + # Infer precision from dtype + dod_dtype = merged["date_of_death"].dtype + if dod_dtype == pl.Date: + return merged.with_columns( + pl.lit(None).cast(pl.Datetime("us")).alias("death_time"), + pl.col("date_of_death").cast(pl.Date).alias("death_date"), + pl.when(pl.col("date_of_death").is_not_null()) + .then(pl.lit("date")) + .otherwise(pl.lit(None)) + .alias("death_time_precision"), + pl.lit("legacy").alias("death_source"), + ) + else: + # Legacy datetimes are ambiguous: older exports may have stored + # date-only values as midnight-cast timestamps. Treat them as + # date precision unless the upstream schema explicitly provides + # a death_time_precision column. + if dod_dtype not in ( + pl.Datetime, + pl.Datetime("us"), + pl.Datetime("ns"), + pl.Datetime("ms"), + pl.Datetime("us"), + ): + merged = merged.with_columns( + pl.col("date_of_death").cast(pl.Datetime("us")).alias("date_of_death") + ) + return merged.with_columns( + pl.lit(None).cast(pl.Datetime("us")).alias("death_time"), + pl.col("date_of_death").cast(pl.Date).alias("death_date"), + pl.when(pl.col("date_of_death").is_not_null()) + .then(pl.lit("date")) + .otherwise(pl.lit(None)) + .alias("death_time_precision"), + pl.lit("legacy").alias("death_source"), ) - return labels + @staticmethod + def _ensure_outcome_columns(merged: pl.DataFrame) -> pl.DataFrame: + """Ensure nullable outcome/evidence columns exist with stable dtypes.""" + exprs = [] + + if "hospital_expire_flag" in merged.columns: + exprs.append( + pl.col("hospital_expire_flag").cast(pl.Int32).alias("hospital_expire_flag") + ) + else: + exprs.append(pl.lit(None).cast(pl.Int32).alias("hospital_expire_flag")) + + if "dischtime" in merged.columns: + dtype = merged["dischtime"].dtype + if dtype not in ( + pl.Datetime, + pl.Datetime("us"), + pl.Datetime("ns"), + pl.Datetime("ms"), + ): + exprs.append(pl.col("dischtime").cast(pl.Datetime("us")).alias("dischtime")) + else: + exprs.append(pl.lit(None).cast(pl.Datetime("us")).alias("dischtime")) + + if "discharge_location" in merged.columns: + exprs.append(pl.col("discharge_location").cast(pl.Utf8).alias("discharge_location")) + else: + exprs.append(pl.lit(None).cast(pl.Utf8).alias("discharge_location")) + + if exprs: + merged = merged.with_columns(exprs) + + return merged + + @staticmethod + def _normalize_datetime_columns(merged: pl.DataFrame) -> pl.DataFrame: + """Normalize timestamp columns to tz-naive microsecond datetimes. + + RICU parquet exports use UTC-aware timestamps. Downstream label logic + also compares against date-only fields, so all temporal columns must + share the same tz-naive dtype before any comparisons run. + """ + exprs = [] + for column in ("intime", "outtime", "dischtime", "death_time", "date_of_death"): + if column not in merged.columns: + continue + + dtype = merged[column].dtype + if isinstance(dtype, pl.datatypes.Datetime): + exprs.append(pl.col(column).cast(pl.Datetime("us")).alias(column)) + + if exprs: + return merged.with_columns(exprs) + + return merged + + @staticmethod + def _death_location_evidence_expr() -> pl.Expr: + """Return true for discharge destinations/statuses that directly encode death.""" + normalized_location = pl.col("discharge_location").str.strip_chars().str.to_lowercase() + return normalized_location.is_in(["death", "died", "expired", "deceased", "dead"]) + + @staticmethod + def _death_evidence_expr() -> pl.Expr: + """Return true when any reliable mortality evidence is present.""" + flag_death = (pl.col("hospital_expire_flag") == 1).fill_null(False) + location_death = MortalityLabelBuilder._death_location_evidence_expr().fill_null(False) + timed_death = pl.col("death_time").is_not_null() | pl.col("death_date").is_not_null() + return flag_death | location_death | timed_death + + @staticmethod + def _hospital_death_date_evidence_expr() -> pl.Expr: + """Return true when date-only death evidence supports in-hospital death. + + MIMIC-IV ``patients.dod`` is post-discharge mortality evidence. It must + not by itself turn a hospital survivor into a hospital mortality label. + Date-only evidence is accepted for hospital mortality only when explicit + hospital death evidence exists, or when the hospital outcome flag is + unknown and the death date is not after hospital discharge. + """ + flag_death = (pl.col("hospital_expire_flag") == 1).fill_null(False) + location_death = MortalityLabelBuilder._death_location_evidence_expr().fill_null(False) + flag_survivor = (pl.col("hospital_expire_flag") == 0).fill_null(False) + date_present = pl.col("death_date").is_not_null() + date_not_after_discharge = ( + pl.col("dischtime").is_not_null() + & (pl.col("death_date") <= pl.col("dischtime").cast(pl.Date)) + ).fill_null(False) + + return date_present & ( + flag_death + | location_death + | (pl.col("hospital_expire_flag").is_null() & ~flag_survivor & date_not_after_discharge) + ) + + @staticmethod + def _hospital_death_evidence_expr() -> pl.Expr: + """Return true for mortality evidence that supports hospital death.""" + flag_death = (pl.col("hospital_expire_flag") == 1).fill_null(False) + location_death = MortalityLabelBuilder._death_location_evidence_expr().fill_null(False) + timestamp_death = pl.col("death_time").is_not_null() + date_death = MortalityLabelBuilder._hospital_death_date_evidence_expr() + return flag_death | location_death | timestamp_death | date_death + + @staticmethod + def _known_hospital_survivor_expr() -> pl.Expr: + """Return true for explicit hospital survivors with no hospital-death evidence.""" + return ( + (pl.col("hospital_expire_flag") == 0) + & ~MortalityLabelBuilder._hospital_death_evidence_expr() + ).fill_null(False) + + @staticmethod + def _known_survivor_expr() -> pl.Expr: + """Return true only for explicit survivor outcomes with no conflicting death evidence.""" + return ( + (pl.col("hospital_expire_flag") == 0) & ~MortalityLabelBuilder._death_evidence_expr() + ).fill_null(False) + + @staticmethod + def _effective_death_time_expr() -> pl.Expr: + """Use exact death_time, or death-coded discharge time when no death_date exists.""" + discharge_death_time = ( + pl.when( + MortalityLabelBuilder._death_evidence_expr() + & pl.col("death_date").is_null() + & pl.col("dischtime").is_not_null() + ) + .then(pl.col("dischtime")) + .otherwise(pl.lit(None).cast(pl.Datetime("us"))) + ) + return pl.coalesce([pl.col("death_time"), discharge_death_time]) + + @staticmethod + def _effective_death_precision_expr() -> pl.Expr: + """Return the precision available after discharge-time death evidence fallback.""" + effective_death_time = MortalityLabelBuilder._effective_death_time_expr() + return ( + pl.when(effective_death_time.is_not_null()) + .then(pl.lit("timestamp")) + .when(pl.col("death_date").is_not_null()) + .then(pl.lit("date")) + .when(MortalityLabelBuilder._death_evidence_expr()) + .then(pl.lit("unknown")) + .otherwise(pl.lit(None)) + ) + + @staticmethod + def _hospital_effective_death_time_expr() -> pl.Expr: + """Use exact hospital death time, or death-coded discharge time fallback.""" + discharge_death_time = ( + pl.when( + MortalityLabelBuilder._hospital_death_evidence_expr() + & pl.col("death_date").is_null() + & pl.col("dischtime").is_not_null() + ) + .then(pl.col("dischtime")) + .otherwise(pl.lit(None).cast(pl.Datetime("us"))) + ) + return pl.coalesce([pl.col("death_time"), discharge_death_time]) + + @staticmethod + def _hospital_effective_death_precision_expr() -> pl.Expr: + """Return the precision available for hospital mortality timing.""" + effective_death_time = MortalityLabelBuilder._hospital_effective_death_time_expr() + return ( + pl.when(effective_death_time.is_not_null()) + .then(pl.lit("timestamp")) + .when(MortalityLabelBuilder._hospital_death_date_evidence_expr()) + .then(pl.lit("date")) + .when(MortalityLabelBuilder._hospital_death_evidence_expr()) + .then(pl.lit("unknown")) + .otherwise(pl.lit(None)) + ) + + def _build_hospital_mortality_with_obs( + self, merged: pl.DataFrame, obs_hours: int + ) -> pl.DataFrame: + """Hospital mortality with observation window exclusion.""" + obs_end = pl.col("intime") + pl.duration(hours=obs_hours) + effective_death_time = self._hospital_effective_death_time_expr() + effective_precision = self._hospital_effective_death_precision_expr() + hospital_death_evidence = self._hospital_death_evidence_expr() + + # Observation windows are half-open: [intime, obs_end). Deaths exactly + # at obs_end belong to the prediction period when gap_hours == 0. + ts_died_during_obs = effective_death_time.is_not_null() & (effective_death_time < obs_end) + + # For date precision: conservative interval logic. A date-only death is + # [00:00, 23:59:59]; if it is fully before obs_end or overlaps the + # boundary from the observation side, exclude. + date_start = pl.col("death_date").cast(pl.Datetime("us")) + date_end = date_start + pl.duration(hours=23, minutes=59, seconds=59) + date_died_during_obs = pl.col("death_date").is_not_null() & (date_end < obs_end) + date_obs_boundary = ( + pl.col("death_date").is_not_null() & (date_start < obs_end) & (date_end >= obs_end) + ) + + died_during_obs = ( + pl.when(effective_precision == "timestamp") + .then(ts_died_during_obs) + .when(effective_precision == "date") + .then(date_died_during_obs | date_obs_boundary) + .otherwise(pl.lit(False)) + ) + + return merged.select( + [ + "stay_id", + pl.when(died_during_obs) + .then(None) + .when(hospital_death_evidence) + .then(1) + .when(self._known_hospital_survivor_expr()) + .then(0) + .otherwise(None) + .cast(pl.Int32) + .alias("label"), + ] + ) + + def _build_icu_mortality(self, merged: pl.DataFrame) -> pl.DataFrame: + """ICU mortality (died during or at ICU discharge), no observation window.""" + effective_death_time = self._effective_death_time_expr() + effective_precision = self._effective_death_precision_expr() + + ts_died = effective_death_time.is_not_null() & (effective_death_time <= pl.col("outtime")) + date_died = pl.col("death_date").is_not_null() & ( + pl.col("death_date") <= pl.col("outtime").cast(pl.Date) + ) + + died_in_icu = ( + pl.when(effective_precision == "timestamp") + .then(ts_died) + .when(effective_precision == "date") + .then(date_died) + .otherwise(pl.lit(None)) + ) + known_not_icu_death = self._death_evidence_expr() & died_in_icu.is_not_null() & ~died_in_icu + + return merged.select( + [ + "stay_id", + pl.when(died_in_icu) + .then(1) + .when(self._known_survivor_expr() | known_not_icu_death) + .then(0) + .otherwise(None) + .cast(pl.Int32) + .alias("label"), + ] + ) + + def _build_legacy_windowed(self, merged: pl.DataFrame, window_hours: int) -> pl.DataFrame: + """Legacy: Time-bounded mortality from admission.""" + boundary = pl.col("intime") + pl.duration(hours=window_hours) + effective_death_time = self._effective_death_time_expr() + effective_precision = self._effective_death_precision_expr() + + ts_died = effective_death_time.is_not_null() & (effective_death_time <= boundary) + # Date-only: death_date end of day <= boundary + date_died_definite = pl.col("death_date").is_not_null() & ( + ( + pl.col("death_date").cast(pl.Datetime("us")) + + pl.duration(hours=23, minutes=59, seconds=59) + ) + <= boundary + ) + # Date-only: boundary overlap (death_date start-of-day <= boundary < end-of-day) + date_boundary_overlap = pl.col("death_date").is_not_null() & ( + (pl.col("death_date").cast(pl.Datetime("us")) <= boundary) + & ( + boundary + < ( + pl.col("death_date").cast(pl.Datetime("us")) + + pl.duration(hours=23, minutes=59, seconds=59) + ) + ) + ) + + died = ( + pl.when(effective_precision == "timestamp") + .then(ts_died) + .when(effective_precision == "date") + .then( + pl.when(date_died_definite) + .then(pl.lit(True)) + .when(date_boundary_overlap) + .then(pl.lit(None)) # null = ambiguous + .otherwise(pl.lit(False)) + ) + .otherwise(pl.lit(None)) + ) + + return merged.select( + [ + "stay_id", + pl.when(died.is_null() & self._known_survivor_expr()) + .then(0) + .when(died.is_null()) + .then(None) + .when(died) + .then(1) + .otherwise(0) + .cast(pl.Int32) + .alias("label"), + ] + ) def _build_windowed_mortality_labels( self, @@ -201,7 +530,6 @@ def _build_windowed_mortality_labels( obs_hours: int, gap_hours: int, prediction_hours: int, - is_date_type: bool, ) -> pl.DataFrame: """Build mortality labels with explicit observation and prediction windows. @@ -209,148 +537,129 @@ def _build_windowed_mortality_labels( |---- observation ----|-- gap --|---- prediction ----| intime obs_end gap_end pred_end - Timeline (for prediction_hours == -1, until ICU discharge): - |---- observation ----|-- gap --|---- prediction (until discharge) ----| - intime obs_end gap_end outtime - - Args: - merged: DataFrame with stays and mortality info joined. - obs_hours: Hours of observation window from admission. - gap_hours: Hours of gap between observation and prediction. - prediction_hours: Hours of prediction window, or -1 for "until ICU discharge". - is_date_type: Whether date_of_death is DATE (vs DATETIME). - - Returns: - DataFrame with stay_id and label: - - 1: Death occurred during prediction window - - 0: Survived prediction window (or died after) - - null: Death occurred during observation or gap window (excluded) - - Note: - Uses outtime (ICU discharge time) in addition to date_of_death to determine - if death occurred during observation. This fixes false positives when - date_of_death is a DATE type (day-level precision) - if outtime >= obs_end, - the patient was still in ICU at observation end and could not have died - during observation. + Uses precision-aware per-row logic: + - "timestamp": exact datetime comparison using death_time + - "date": interval-based [00:00:00, 23:59:59.999] comparison using death_date; + boundary overlaps produce null (conservative exclusion) + - "unknown"/missing: true survivors remain 0, unknown outcomes are null """ - # Calculate window boundaries - # obs_end = intime + obs_hours - # gap_end = obs_end + gap_hours = intime + obs_hours + gap_hours prediction_start_hours = obs_hours + gap_hours - - # Check if prediction window extends until ICU discharge until_icu_discharge = prediction_hours == -1 + effective_death_time = self._effective_death_time_expr() + effective_precision = self._effective_death_precision_expr() - # Use outtime to determine if patient was still in ICU at end of observation. - # This is more reliable than date_of_death for DATE types (day-level precision). - # If outtime >= intime + obs_hours, patient was alive at observation end. - obs_end_datetime = pl.col("intime") + pl.duration(hours=obs_hours) - left_icu_during_obs = pl.col("outtime") < obs_end_datetime - - if is_date_type: - # For DATE type, cast boundaries to Date for comparison - obs_end = obs_end_datetime.cast(pl.Date) - pred_start = (pl.col("intime") + pl.duration(hours=prediction_start_hours)).cast( - pl.Date - ) + obs_end = pl.col("intime") + pl.duration(hours=obs_hours) + pred_start = pl.col("intime") + pl.duration(hours=prediction_start_hours) - # Death during observation window (exclude these stays) - # Must check BOTH date_of_death AND outtime because date_of_death has only - # day-level precision which causes false positives (see issue with DATE vs DATETIME). - # If outtime >= obs_end, patient was still in ICU at observation end, so they - # could not have died during observation regardless of what date_of_death says. - died_during_obs = ( - pl.col("date_of_death").is_not_null() - & (pl.col("date_of_death") <= obs_end) - & left_icu_during_obs + if until_icu_discharge: + pred_end = pl.col("outtime") + else: + prediction_end_hours = prediction_start_hours + prediction_hours + pred_end = pl.col("intime") + pl.duration(hours=prediction_end_hours) + + # --- Timestamp precision: exact comparisons --- + ts_died_during_obs = effective_death_time.is_not_null() & (effective_death_time < obs_end) + ts_died_during_gap = ( + ( + effective_death_time.is_not_null() + & (effective_death_time > obs_end) + & (effective_death_time < pred_start) ) + if gap_hours > 0 + else pl.lit(False) + ) + ts_died_during_pred = ( + effective_death_time.is_not_null() + & (effective_death_time >= pred_start) + & (effective_death_time <= pred_end) + ) - # Death during gap period (exclude — ambiguous whether model could predict) - died_during_gap = ( - ( - pl.col("date_of_death").is_not_null() - & (pl.col("date_of_death") > obs_end) - & (pl.col("date_of_death") < pred_start) - ) - if gap_hours > 0 - else pl.lit(False) - ) + # --- Date precision: interval-based conservative logic --- + # A date-only death represents [date 00:00:00, date 23:59:59.999]. + # "Definite" means the entire interval falls within the region. + # "Overlap" means the interval straddles a boundary → null. + date_start = pl.col("death_date").cast(pl.Datetime("us")) + date_end = date_start + pl.duration(hours=23, minutes=59, seconds=59) - if until_icu_discharge: - # Prediction window ends at ICU discharge (outtime) - pred_end = pl.col("outtime").cast(pl.Date) - else: - # Prediction window ends at fixed time after observation - prediction_end_hours = prediction_start_hours + prediction_hours - pred_end = (pl.col("intime") + pl.duration(hours=prediction_end_hours)).cast( - pl.Date - ) + # Died before obs_end? The entire date interval is strictly before the + # half-open observation boundary. + date_died_during_obs = pl.col("death_date").is_not_null() & (date_end < obs_end) - # Death during prediction window - # Use >= pred_start to include deaths exactly at prediction start - died_during_pred = ( - pl.col("date_of_death").is_not_null() - & (pl.col("date_of_death") >= pred_start) - & (pl.col("date_of_death") <= pred_end) - ) - else: - # For DATETIME type, use full precision - obs_end = obs_end_datetime - pred_start = pl.col("intime") + pl.duration(hours=prediction_start_hours) - - # Death during observation window (exclude these stays) - # Also check outtime for consistency with DATE type logic - died_during_obs = ( - pl.col("date_of_death").is_not_null() - & (pl.col("date_of_death") <= obs_end) - & left_icu_during_obs - ) + # Date interval overlaps obs_end boundary (part before, part after) + date_obs_boundary = ( + pl.col("death_date").is_not_null() & (date_start < obs_end) & (date_end >= obs_end) + ) - # Death during gap period (exclude — ambiguous whether model could predict) - died_during_gap = ( - ( - pl.col("date_of_death").is_not_null() - & (pl.col("date_of_death") > obs_end) - & (pl.col("date_of_death") < pred_start) - ) - if gap_hours > 0 - else pl.lit(False) - ) + # Definite during prediction: entire interval within [pred_start, pred_end] + date_died_during_pred_definite = ( + pl.col("death_date").is_not_null() & (date_start >= pred_start) & (date_end <= pred_end) + ) - if until_icu_discharge: - # Prediction window ends at ICU discharge (outtime) - pred_end = pl.col("outtime") - else: - # Prediction window ends at fixed time after observation - prediction_end_hours = prediction_start_hours + prediction_hours - pred_end = pl.col("intime") + pl.duration(hours=prediction_end_hours) - - # Death during prediction window - # Use >= pred_start to include deaths exactly at prediction start - died_during_pred = ( - pl.col("date_of_death").is_not_null() - & (pl.col("date_of_death") >= pred_start) - & (pl.col("date_of_death") <= pred_end) + # Date interval overlaps pred_start boundary + date_pred_start_overlap = ( + pl.col("death_date").is_not_null() + & (date_start < pred_start) + & (date_end >= pred_start) + ) + + # Date interval overlaps pred_end boundary + date_pred_end_overlap = ( + pl.col("death_date").is_not_null() & (date_start <= pred_end) & (date_end > pred_end) + ) + + # Definite during gap: entire interval within (obs_end, pred_start) + date_died_during_gap = ( + (pl.col("death_date").is_not_null() & (date_start > obs_end) & (date_end < pred_start)) + if gap_hours > 0 + else pl.lit(False) + ) + + # Any boundary overlap → ambiguous → null + date_any_boundary = date_obs_boundary | date_pred_start_overlap | date_pred_end_overlap + + # --- Per-row label computation --- + # Timestamp precision path + ts_label = ( + pl.when(ts_died_during_obs) + .then(None) + .when(ts_died_during_gap) + .then(None) + .when(ts_died_during_pred) + .then(1) + .otherwise(0) + ) + + # Date precision path + date_label = ( + pl.when(date_any_boundary) + .then(None) # ambiguous boundary → exclude + .when(date_died_during_obs) + .then(None) # definite during obs → exclude + .when(date_died_during_gap) + .then(None) # definite during gap → exclude + .when(date_died_during_pred_definite) + .then(1) # definite during prediction → positive + .otherwise(0) # definite outside → negative + ) + + # Combine by precision + label_expr = ( + pl.when(effective_precision == "timestamp") + .then(ts_label) + .when(effective_precision == "date") + .then(date_label) + .otherwise( + # Unknown death timing or unknown outcome → exclude. Explicit + # survivors remain negatives. + pl.when(self._known_survivor_expr()) + .then(0) + .otherwise(None) ) + ) - # Build labels: - # - null if died during observation (can't predict the past) - # - null if died during gap period (ambiguous, exclude) - # - 1 if died during prediction window - # - 0 otherwise (survived or died after prediction window) - labels = merged.select( + return merged.select( [ "stay_id", - pl.when(died_during_obs) - .then(None) - .when(died_during_gap) - .then(None) - .when(died_during_pred) - .then(1) - .otherwise(0) - .cast(pl.Int32) - .alias("label"), + label_expr.cast(pl.Int32).alias("label"), ] ) - - return labels diff --git a/src/slices/data/preparation.py b/src/slices/data/preparation.py new file mode 100644 index 0000000..e6d36af --- /dev/null +++ b/src/slices/data/preparation.py @@ -0,0 +1,246 @@ +"""Utilities for preparing processed datasets for training. + +This module owns the canonical split-generation and normalization-stat logic +used by both standalone scripts and combined-dataset assembly. +""" + +import os +from pathlib import Path +from tempfile import NamedTemporaryFile + +import numpy as np +import polars as pl +import yaml + +from slices.constants import TEST_RATIO, TRAIN_RATIO, VAL_RATIO +from slices.data.tensor_cache import ( + get_data_fingerprint, + get_preprocessing_fingerprint, +) + + +def _atomic_yaml_write(path: Path, data: dict) -> None: + """Write YAML atomically using a temp file and rename.""" + with NamedTemporaryFile(dir=path.parent, suffix=".yaml", mode="w", delete=False) as tmp: + tmp_path = tmp.name + yaml.dump(data, tmp, default_flow_style=False) + os.replace(tmp_path, path) + + +def compute_patient_splits( + static_df: pl.DataFrame, + timeseries_df: pl.DataFrame, + train_ratio: float, + val_ratio: float, + test_ratio: float, + seed: int, +) -> dict: + """Compute patient-level train/val/test splits.""" + total = train_ratio + val_ratio + test_ratio + if not np.isclose(total, 1.0): + raise ValueError(f"Split ratios must sum to 1.0, got {total}") + + stay_ids = timeseries_df["stay_id"].to_list() + stay_to_patient = dict(zip(static_df["stay_id"].to_list(), static_df["patient_id"].to_list())) + + unique_patients = sorted(set(stay_to_patient.values())) + n_patients = len(unique_patients) + + print(f" Total patients: {n_patients:,}") + print(f" Total stays: {len(stay_ids):,}") + + if n_patients == len(stay_ids) and n_patients > 0: + patient_set = set(stay_to_patient.values()) + stay_set = set(stay_ids) + if patient_set == stay_set: + print( + " WARNING: patient_id == stay_id for all stays. This dataset likely " + "lacks true patient-level identifiers (e.g. HiRID, SICdb). " + "Patient-level split cannot prevent leakage from repeat ICU admissions." + ) + + rng = np.random.RandomState(seed) + patient_indices = np.arange(n_patients) + rng.shuffle(patient_indices) + shuffled_patients = [unique_patients[i] for i in patient_indices] + + n_train = int(n_patients * train_ratio) + n_val = int(n_patients * val_ratio) + + train_patients = set(shuffled_patients[:n_train]) + val_patients = set(shuffled_patients[n_train : n_train + n_val]) + test_patients = set(shuffled_patients[n_train + n_val :]) + + assert train_patients.isdisjoint(val_patients), "Train/val overlap!" + assert train_patients.isdisjoint(test_patients), "Train/test overlap!" + assert val_patients.isdisjoint(test_patients), "Val/test overlap!" + + train_indices = [] + val_indices = [] + test_indices = [] + + for idx, stay_id in enumerate(stay_ids): + patient_id = stay_to_patient[stay_id] + if patient_id in train_patients: + train_indices.append(idx) + elif patient_id in val_patients: + val_indices.append(idx) + else: + test_indices.append(idx) + + return { + "seed": seed, + "train_ratio": train_ratio, + "val_ratio": val_ratio, + "test_ratio": test_ratio, + "train_patients": sorted(train_patients), + "val_patients": sorted(val_patients), + "test_patients": sorted(test_patients), + "train_indices": train_indices, + "val_indices": val_indices, + "test_indices": test_indices, + "train_stays": len(train_indices), + "val_stays": len(val_indices), + "test_stays": len(test_indices), + } + + +def compute_normalization_stats( + timeseries_df: pl.DataFrame, + train_indices: list, + feature_names: list, + seq_length: int, +) -> dict: + """Compute per-feature normalization stats on the training split only.""" + n_features = len(feature_names) + n_train = len(train_indices) + + raw_timeseries = timeseries_df["timeseries"].to_list() + raw_masks = timeseries_df["mask"].to_list() + + sums = np.zeros(n_features) + sq_sums = np.zeros(n_features) + counts = np.zeros(n_features) + + for progress, idx in enumerate(train_indices): + if (progress + 1) % 5000 == 0: + print(f" Processing {progress + 1:,}/{n_train:,} training samples...") + + ts_arr = np.array(raw_timeseries[idx][:seq_length], dtype=np.float64) + mask_arr = np.array(raw_masks[idx][:seq_length], dtype=bool) + + valid = mask_arr & np.isfinite(ts_arr) + ts_valid = np.where(valid, ts_arr, 0.0) + + counts += valid.sum(axis=0) + sums += ts_valid.sum(axis=0) + sq_sums += (ts_valid**2).sum(axis=0) + + safe_counts = np.maximum(counts, 1) + means = sums / safe_counts + variance = (sq_sums - counts * means**2) / np.maximum(safe_counts - 1, 1) + stds = np.sqrt(np.maximum(variance, 0.0)) + stds = np.where(stds > 1e-6, stds, 1.0) + means = np.where(counts > 0, means, 0.0) + + return { + "feature_means": means.tolist(), + "feature_stds": stds.tolist(), + "feature_names": feature_names, + "train_indices": train_indices, + "observation_counts": counts.tolist(), + } + + +def prepare_processed_dataset( + processed_dir: Path, + seed: int, + dataset_name: str | None = None, +) -> tuple[dict, dict]: + """Generate and save splits plus normalization stats for a processed dataset.""" + dataset_label = dataset_name or processed_dir.name + + print("=" * 70) + print(f"Dataset Preparation — {dataset_label}") + print("=" * 70) + + if not processed_dir.exists(): + raise FileNotFoundError( + f"Processed directory not found: {processed_dir}\nRun extraction before preparation." + ) + + print("\n1. Loading metadata...") + with open(processed_dir / "metadata.yaml") as f: + metadata = yaml.safe_load(f) + + feature_names = metadata["feature_names"] + seq_length = metadata["seq_length_hours"] + print(f" Features: {len(feature_names)}") + print(f" Sequence length: {seq_length} hours") + + print("\n2. Loading parquet files...") + static_df = pl.read_parquet(processed_dir / "static.parquet") + timeseries_df = pl.read_parquet(processed_dir / "timeseries.parquet") + print(f" Loaded {len(timeseries_df):,} stays") + + print("\n3. Computing patient-level splits...") + splits = compute_patient_splits( + static_df=static_df, + timeseries_df=timeseries_df, + train_ratio=TRAIN_RATIO, + val_ratio=VAL_RATIO, + test_ratio=TEST_RATIO, + seed=seed, + ) + + print(f" Train: {splits['train_stays']:,} stays ({len(splits['train_patients']):,} patients)") + print(f" Val: {splits['val_stays']:,} stays ({len(splits['val_patients']):,} patients)") + print(f" Test: {splits['test_stays']:,} stays ({len(splits['test_patients']):,} patients)") + + splits_path = processed_dir / "splits.yaml" + splits_to_save = { + "seed": splits["seed"], + "train_ratio": splits["train_ratio"], + "val_ratio": splits["val_ratio"], + "test_ratio": splits["test_ratio"], + "train_indices": splits["train_indices"], + "val_indices": splits["val_indices"], + "test_indices": splits["test_indices"], + "train_stays": splits["train_stays"], + "val_stays": splits["val_stays"], + "test_stays": splits["test_stays"], + "train_patients": splits["train_patients"], + "val_patients": splits["val_patients"], + "test_patients": splits["test_patients"], + } + _atomic_yaml_write(splits_path, splits_to_save) + print(f"\n Saved: {splits_path}") + + print("\n4. Computing normalization statistics (train set only)...") + stats = compute_normalization_stats( + timeseries_df=timeseries_df, + train_indices=splits["train_indices"], + feature_names=feature_names, + seq_length=seq_length, + ) + stats["normalize"] = True + stats["split_hash"] = None + stats["normalization_index_space"] = "raw_full_cohort" + stats["train_indices_count"] = len(splits["train_indices"]) + stats["data_fingerprint"] = get_data_fingerprint(processed_dir) + stats["preprocessing_fingerprint"] = get_preprocessing_fingerprint() + + stats_path = processed_dir / "normalization_stats.yaml" + _atomic_yaml_write(stats_path, stats) + print(f" Saved: {stats_path}") + + print("\n" + "=" * 70) + print("Preparation Complete") + print("=" * 70) + print(f"\nOutput directory: {processed_dir}") + print(" - splits.yaml") + print(" - normalization_stats.yaml") + print("\nYou can now run training:") + print(f" uv run python scripts/training/pretrain.py dataset={dataset_label} ssl=mae") + + return splits_to_save, stats diff --git a/src/slices/data/sliding_window.py b/src/slices/data/sliding_window.py index e2e7395..754f1de 100644 --- a/src/slices/data/sliding_window.py +++ b/src/slices/data/sliding_window.py @@ -28,18 +28,18 @@ class SlidingWindowDataset(Dataset): learn better representations. Example: - >>> base_dataset = ICUDataset("data/processed/mimic-iv-168h") # 168h sequences + >>> base_dataset = ICUDataset("data/processed/miiv-168h") # 168h sequences >>> windowed = SlidingWindowDataset( ... base_dataset, - ... window_size=48, # 48h windows - ... stride=24, # 24h stride (50% overlap) + ... window_size=24, # 24h benchmark windows + ... stride=12, # 12h stride (50% overlap) ... stay_indices=[0, 1, 2, 10, 20], # Only these stays (e.g., train split) ... ) - >>> len(windowed) # More samples than base dataset - 15 + >>> len(windowed) > len(base_dataset) + True >>> sample = windowed[0] >>> sample["timeseries"].shape - torch.Size([48, 9]) + torch.Size([24, 9]) Attributes: base_dataset: The underlying ICUDataset. diff --git a/src/slices/data/splits.py b/src/slices/data/splits.py new file mode 100644 index 0000000..b3ad2fb --- /dev/null +++ b/src/slices/data/splits.py @@ -0,0 +1,659 @@ +"""Patient-level split logic for ICU datasets. + +Handles loading/computing/caching of train/val/test splits, label filtering, +and label-efficiency subsampling. Extracted from ICUDataModule for modularity. +""" + +import logging +from pathlib import Path +from typing import Any, List, Optional, Set, Tuple + +import numpy as np +import polars as pl +import yaml + +logger = logging.getLogger(__name__) + + +def load_cached_splits( + processed_dir: Path, + static_df: pl.DataFrame, + stay_ids: List[int], + seed: int, + train_ratio: float, + val_ratio: float, + test_ratio: float, + cohort_stay_ids: Optional[List[int]] = None, +) -> Optional[Tuple[List[int], List[int], List[int]]]: + """Load cached splits from splits.yaml if valid. + + Validates that cached splits match current parameters (seed, ratios) and + that the patient lists are consistent with the full cohort. The returned + indices are always mapped over ``stay_ids``, which may be task-filtered. + + Args: + processed_dir: Path to processed data directory containing splits.yaml. + static_df: Static dataframe with stay_id -> patient_id mapping. + stay_ids: List of stay_ids to map into split indices. For supervised + tasks this may already be label-filtered. + seed: Random seed for split computation. + train_ratio: Fraction of patients for training. + val_ratio: Fraction of patients for validation. + test_ratio: Fraction of patients for testing. + cohort_stay_ids: Optional full-cohort stay_ids used to validate that + cached patient lists are global rather than task-filtered. When not + provided, ``stay_ids`` is treated as the full cohort for backward + compatibility. + + Returns: + Tuple of (train_indices, val_indices, test_indices) if cache is valid, + None otherwise. + """ + split_path = processed_dir / "splits.yaml" + if not split_path.exists(): + return None + + try: + with open(split_path) as f: + cached = yaml.safe_load(f) + + # Validate parameters match + if ( + cached.get("seed") != seed + or not np.isclose(cached.get("train_ratio", 0), train_ratio) + or not np.isclose(cached.get("val_ratio", 0), val_ratio) + or not np.isclose(cached.get("test_ratio", 0), test_ratio) + ): + logger.debug("Cached splits have different parameters, recomputing") + return None + + # Get patient lists from cache + train_patients = set(cached.get("train_patients", [])) + val_patients = set(cached.get("val_patients", [])) + test_patients = set(cached.get("test_patients", [])) + + if not train_patients or not val_patients or not test_patients: + logger.debug("Cached splits missing patient lists, recomputing") + return None + + # Validate patient sets are disjoint + if not ( + train_patients.isdisjoint(val_patients) + and train_patients.isdisjoint(test_patients) + and val_patients.isdisjoint(test_patients) + ): + logger.debug("Cached splits have overlapping patients, recomputing") + return None + + # Get stay_id -> patient_id mapping + stay_to_patient = dict( + zip(static_df["stay_id"].to_list(), static_df["patient_id"].to_list()) + ) + + # Validate cached patient lists against the full cohort, not the + # possibly task-filtered stay_ids. Supervised runs should reuse the + # global patient assignment and then filter inside those sets. + full_stay_ids = cohort_stay_ids if cohort_stay_ids is not None else stay_ids + full_patients = set() + for stay_id in full_stay_ids: + patient_id = stay_to_patient.get(stay_id) + if patient_id is None: + logger.debug( + f"Stay {stay_id} missing from static stay->patient mapping, recomputing" + ) + return None + full_patients.add(patient_id) + cached_patients = train_patients | val_patients | test_patients + + if full_patients != cached_patients: + logger.debug( + f"Cached splits have different full-cohort patients " + f"(cached: {len(cached_patients)}, full cohort: {len(full_patients)}), " + f"recomputing" + ) + return None + + # The current mapped cohort may be task-filtered, but every remaining + # patient must come from the cached full-cohort split. + current_patients = set() + for stay_id in stay_ids: + patient_id = stay_to_patient.get(stay_id) + if patient_id is None: + logger.debug( + f"Stay {stay_id} missing from static stay->patient mapping, recomputing" + ) + return None + current_patients.add(patient_id) + if not current_patients.issubset(cached_patients): + logger.debug( + f"Cached splits do not cover current patients " + f"(cached: {len(cached_patients)}, current: {len(current_patients)}), " + f"recomputing" + ) + return None + + # Reconstruct indices from patient lists + train_indices = [] + val_indices = [] + test_indices = [] + + for idx, stay_id in enumerate(stay_ids): + patient_id = stay_to_patient.get(stay_id) + if patient_id in train_patients: + train_indices.append(idx) + elif patient_id in val_patients: + val_indices.append(idx) + elif patient_id in test_patients: + test_indices.append(idx) + else: + logger.debug(f"Stay {stay_id} has unknown patient {patient_id}, recomputing") + return None + + logger.info( + f"Loaded cached splits: " + f"{len(train_indices)} train, {len(val_indices)} val, {len(test_indices)} test" + ) + return train_indices, val_indices, test_indices + + except Exception as e: + logger.debug(f"Failed to load cached splits: {e}, recomputing") + return None + + +def filter_stays_with_missing_labels( + stay_ids: List[int], + labels_df: pl.DataFrame, + task_name: Optional[str], +) -> Tuple[List[int], Set[int]]: + """Filter out stays with missing labels for the configured task. + + This is applied after the global patient split is defined. Split indices are + then mapped onto the filtered stay list so DataModule and Dataset stay + index-consistent without changing patient assignment per task. + + Args: + stay_ids: List of all stay_ids from timeseries. + labels_df: Labels dataframe with task columns. + task_name: Task name to filter for, or None for unsupervised. + + Returns: + Tuple of (filtered_stay_ids, excluded_stay_ids_set). + """ + if task_name is None: + return stay_ids, set() + + logger.debug(f"Checking labels for task '{task_name}'") + + # Detect multi-label tasks: columns prefixed with "{task_name}_" + multilabel_cols = [c for c in labels_df.columns if c.startswith(f"{task_name}_")] + is_multilabel = len(multilabel_cols) > 0 and task_name not in labels_df.columns + + # Get stays with valid labels for this task + if not is_multilabel and task_name not in labels_df.columns: + raise ValueError(f"Task '{task_name}' not found in labels. Available: {labels_df.columns}") + + # Find stays with non-null labels + if is_multilabel: + valid_labels_df = labels_df.filter( + pl.all_horizontal([pl.col(c).is_not_null() for c in multilabel_cols]) + ) + else: + valid_labels_df = labels_df.filter(pl.col(task_name).is_not_null()) + valid_stay_ids = set(valid_labels_df["stay_id"].to_list()) + + # Filter stay_ids maintaining order + filtered_stay_ids = [sid for sid in stay_ids if sid in valid_stay_ids] + excluded_stay_ids = set(stay_ids) - valid_stay_ids + + if excluded_stay_ids: + pct_excluded = len(excluded_stay_ids) / len(stay_ids) * 100 + logger.warning( + f"Excluding {len(excluded_stay_ids):,} stays ({pct_excluded:.1f}%) " + f"with missing '{task_name}' labels. Remaining: {len(filtered_stay_ids):,}" + ) + + return filtered_stay_ids, excluded_stay_ids + + +def _split_patient_sets( + patient_ids: Set[int], + seed: int, + train_ratio: float, + val_ratio: float, +) -> Tuple[Set[int], Set[int], Set[int]]: + """Split patient IDs deterministically into train/val/test sets.""" + unique_patients = sorted(patient_ids) + n_patients = len(unique_patients) + + rng = np.random.RandomState(seed) + patient_indices = np.arange(n_patients) + rng.shuffle(patient_indices) + shuffled_patients = [unique_patients[i] for i in patient_indices] + + n_train = int(n_patients * train_ratio) + n_val = int(n_patients * val_ratio) + + train_patients = set(shuffled_patients[:n_train]) + val_patients = set(shuffled_patients[n_train : n_train + n_val]) + test_patients = set(shuffled_patients[n_train + n_val :]) + + return train_patients, val_patients, test_patients + + +def _indices_from_patient_sets( + stay_ids: List[int], + stay_to_patient: dict[int, int], + train_patients: Set[int], + val_patients: Set[int], + test_patients: Set[int], +) -> Tuple[List[int], List[int], List[int]]: + """Map stay IDs to split indices according to precomputed patient sets.""" + train_indices = [] + val_indices = [] + test_indices = [] + + for idx, stay_id in enumerate(stay_ids): + patient_id = stay_to_patient.get(stay_id) + if patient_id is None: + raise ValueError( + f"Stay {stay_id} has no patient_id mapping. " + "Patient IDs are required for patient-level splits." + ) + if patient_id in train_patients: + train_indices.append(idx) + elif patient_id in val_patients: + val_indices.append(idx) + elif patient_id in test_patients: + test_indices.append(idx) + else: + raise ValueError( + f"Stay {stay_id} has patient_id {patient_id} which is not " + "in any split. This indicates a bug in split computation." + ) + + return train_indices, val_indices, test_indices + + +def save_global_split_info( + processed_dir: Path, + static_df: pl.DataFrame, + stay_ids: List[int], + seed: int, + train_ratio: float, + val_ratio: float, + test_ratio: float, + dataset: Optional[Any] = None, + train_subset_indices: Optional[List[int]] = None, + label_fraction: float = 1.0, +) -> None: + """Save canonical full-cohort patient split information.""" + stay_to_patient = dict(zip(static_df["stay_id"].to_list(), static_df["patient_id"].to_list())) + full_patient_ids = {stay_to_patient[sid] for sid in stay_ids if sid in stay_to_patient} + train_patients, val_patients, test_patients = _split_patient_sets( + full_patient_ids, + seed, + train_ratio, + val_ratio, + ) + train_indices, val_indices, test_indices = _indices_from_patient_sets( + stay_ids, + stay_to_patient, + train_patients, + val_patients, + test_patients, + ) + + split_info = { + "seed": seed, + "train_ratio": train_ratio, + "val_ratio": val_ratio, + "test_ratio": test_ratio, + "label_fraction": label_fraction, + "train_patients": sorted(train_patients), + "val_patients": sorted(val_patients), + "test_patients": sorted(test_patients), + "train_stays": len(train_indices), + "val_stays": len(val_indices), + "test_stays": len(test_indices), + } + + if dataset is not None and train_subset_indices is not None: + train_subset_stay_ids = [dataset.stay_ids[i] for i in train_subset_indices] + train_subset_patients = sorted( + {stay_to_patient[dataset.stay_ids[i]] for i in train_subset_indices} + ) + split_info.update( + { + "train_subset_patients": train_subset_patients, + "train_subset_stays": len(train_subset_indices), + "train_subset_stay_ids": train_subset_stay_ids, + } + ) + + split_path = processed_dir / "splits.yaml" + with open(split_path, "w") as f: + yaml.dump(split_info, f, default_flow_style=False) + + +def compute_patient_level_splits( + processed_dir: Path, + task_name: Optional[str], + seed: int, + train_ratio: float, + val_ratio: float, + test_ratio: float, +) -> Tuple[ + List[int], + List[int], + List[int], + pl.DataFrame, + pl.DataFrame, + List[int], + List[int], + Set[int], + List[int], +]: + """Compute patient-level train/val/test splits from parquet files. + + First checks for cached full-cohort splits in splits.yaml. If found and + parameters match, loads the global patient assignments from cache. Otherwise, + computes global patient assignments from scratch. + + Loads static and timeseries data directly without requiring the full dataset + to be initialized. This is called BEFORE dataset creation to ensure + normalization statistics use only training data (prevents data leakage). + + IMPORTANT: Patient assignment is computed on the full cohort first. Stays + with missing task labels are filtered only after that assignment, so SSL + pretraining and downstream task splits share one global patient partition. + + Uses deterministic shuffling of patient_id to assign patients to splits. + All stays from a patient go to the same split. + + Args: + processed_dir: Path to processed data directory. + task_name: Task name for label filtering, or None for unsupervised. + seed: Random seed for split computation. + train_ratio: Fraction of patients for training. + val_ratio: Fraction of patients for validation. + test_ratio: Fraction of patients for testing. + + Returns: + Tuple of (train_indices, val_indices, test_indices, static_df, labels_df, + all_stay_ids, filtered_stay_ids, excluded_stay_ids, + normalization_train_indices). The first three index lists are + mapped over ``filtered_stay_ids`` for supervised tasks. The + normalization index list is always mapped over ``all_stay_ids`` + so downstream tasks share the same dataset/seed normalizer as + SSL pretraining. + """ + logger.info("Loading data for split computation...") + + # Load only needed columns for efficiency + static_path = processed_dir / "static.parquet" + logger.debug(f"Loading static data from {static_path.name}") + static_df = pl.read_parquet(static_path, columns=["stay_id", "patient_id"]) + + # Load only stay_id column from timeseries (much faster than full load) + timeseries_path = processed_dir / "timeseries.parquet" + logger.debug(f"Loading stay_ids from {timeseries_path.name}") + timeseries_df = pl.read_parquet(timeseries_path, columns=["stay_id"]) + all_stay_ids = timeseries_df["stay_id"].to_list() + + # Load labels to filter out stays with missing labels + labels_path = processed_dir / "labels.parquet" + logger.debug(f"Loading labels from {labels_path.name}") + labels_df = pl.read_parquet(labels_path) + + # Filter task labels before mapping indices, but do not let this change the + # patient assignment. SSL and every downstream task must share the same + # full-cohort patient split. + stay_ids, excluded_stay_ids = filter_stays_with_missing_labels( + all_stay_ids, labels_df, task_name + ) + + # Build the full-cohort stay_id -> patient_id mapping once. Downstream + # optimization indices may be task-filtered, but normalization must remain + # anchored to this unfiltered index space. + logger.debug("Building stay-to-patient mapping") + stay_to_patient = dict(zip(static_df["stay_id"].to_list(), static_df["patient_id"].to_list())) + + # Try to load cached splits first + logger.debug("Checking for cached splits") + cached_splits = load_cached_splits( + processed_dir, + static_df, + stay_ids, + seed, + train_ratio, + val_ratio, + test_ratio, + cohort_stay_ids=all_stay_ids, + ) + if cached_splits is not None: + train_indices, val_indices, test_indices = cached_splits + with open(processed_dir / "splits.yaml") as f: + cached_split_info = yaml.safe_load(f) or {} + train_patients = set(cached_split_info.get("train_patients", [])) + normalization_train_indices = [ + idx + for idx, stay_id in enumerate(all_stay_ids) + if stay_to_patient[stay_id] in train_patients + ] + return ( + train_indices, + val_indices, + test_indices, + static_df, + labels_df, + all_stay_ids, + stay_ids, + excluded_stay_ids, + normalization_train_indices, + ) + + logger.info("Computing full-cohort patient-level splits...") + + # Get unique full-cohort patients (sorted for deterministic ordering across + # Python runs). Task filtering is applied later when mapping stay indices. + full_patient_ids = {stay_to_patient[sid] for sid in all_stay_ids if sid in stay_to_patient} + unique_patients = sorted(full_patient_ids) + n_patients = len(unique_patients) + logger.debug(f"Found {n_patients:,} unique patients") + + # Warn if patient_id == stay_id for all stays (e.g. HiRID, SICdb) + # This means the dataset lacks true patient-level IDs, so multiple ICU + # stays from the same real patient may leak across splits. + if n_patients == len(all_stay_ids) and n_patients > 0: + full_stay_set = set(all_stay_ids) + if full_patient_ids == full_stay_set: + logger.warning( + "patient_id == stay_id for all stays. This dataset likely lacks " + "true patient-level identifiers (e.g. HiRID, SICdb). " + "Patient-level split cannot prevent leakage from repeat ICU admissions." + ) + + # Shuffle full-cohort patients deterministically using seed + logger.debug(f"Shuffling patients (seed={seed})") + train_patients, val_patients, test_patients = _split_patient_sets( + full_patient_ids, + seed, + train_ratio, + val_ratio, + ) + logger.debug( + f"Split: {len(train_patients):,} train, " + f"{len(val_patients):,} val, {len(test_patients):,} test patients" + ) + + # Verify no patient overlap between splits (data leakage check) + assert train_patients.isdisjoint( + val_patients + ), "Patient leakage detected: train/val splits have overlapping patients" + assert train_patients.isdisjoint( + test_patients + ), "Patient leakage detected: train/test splits have overlapping patients" + assert val_patients.isdisjoint( + test_patients + ), "Patient leakage detected: val/test splits have overlapping patients" + + # Verify all patients are accounted for in exactly one split + all_patients_in_splits = train_patients | val_patients | test_patients + all_unique_patients = set(unique_patients) + missing_patients = all_unique_patients - all_patients_in_splits + extra_patients = all_patients_in_splits - all_unique_patients + + assert not missing_patients, ( + f"Patient split validation failed: {len(missing_patients)} patients " + f"not assigned to any split. First 5: {list(missing_patients)[:5]}" + ) + assert not extra_patients, ( + f"Patient split validation failed: {len(extra_patients)} patients " + f"in splits but not in data. First 5: {list(extra_patients)[:5]}" + ) + + # Map global patient assignments back to the current stay list. For + # supervised tasks, stay_ids has already been label-filtered, so this + # applies task filtering inside the global train/val/test patient sets. + logger.debug("Mapping global patient split to current stay indices") + train_indices, val_indices, test_indices = _indices_from_patient_sets( + stay_ids, + stay_to_patient, + train_patients, + val_patients, + test_patients, + ) + normalization_train_indices, _, _ = _indices_from_patient_sets( + all_stay_ids, + stay_to_patient, + train_patients, + val_patients, + test_patients, + ) + + # Final validation: all stays should be accounted for + total_assigned = len(train_indices) + len(val_indices) + len(test_indices) + assert total_assigned == len(stay_ids), ( + f"Split validation failed: {total_assigned} stays assigned but " + f"{len(stay_ids)} total stays in dataset" + ) + + logger.info( + f"Splits computed: {len(train_indices):,} train, " + f"{len(val_indices):,} val, {len(test_indices):,} test stays" + ) + + return ( + train_indices, + val_indices, + test_indices, + static_df, + labels_df, + all_stay_ids, + stay_ids, + excluded_stay_ids, + normalization_train_indices, + ) + + +def subsample_train_indices( + train_indices: List[int], + label_fraction: float, + seed: int, +) -> List[int]: + """Subsample training indices for label-efficiency ablations. + + Uses a separate RNG seeded with the provided seed to ensure deterministic + subsampling. Always selects at least 1 sample. + + Args: + train_indices: Full list of training indices. + label_fraction: Fraction of training data to use (0, 1]. + seed: Random seed for reproducibility. + + Returns: + Subsampled list of training indices. + """ + n_full = len(train_indices) + n_subsample = max(1, int(n_full * label_fraction)) + + rng = np.random.RandomState(seed) + subsample_idx = rng.choice(n_full, size=n_subsample, replace=False) + subsample_idx.sort() # Maintain original order + subsampled = [train_indices[i] for i in subsample_idx] + + logger.info( + f"Label fraction={label_fraction}: using {n_subsample:,}/{n_full:,} " + f"training samples ({label_fraction * 100:.1f}%)" + ) + + return subsampled + + +def save_split_info( + processed_dir: Path, + dataset: Any, + train_indices: List[int], + val_indices: List[int], + test_indices: List[int], + seed: int, + train_ratio: float, + val_ratio: float, + test_ratio: float, + train_subset_indices: Optional[List[int]] = None, + label_fraction: float = 1.0, +) -> None: + """Save split information to file for reproducibility. + + Args: + processed_dir: Path to processed data directory. + dataset: ICUDataset instance with stay_ids and static_df. + train_indices: Training split indices. + val_indices: Validation split indices. + test_indices: Test split indices. + seed: Random seed used. + train_ratio: Training ratio used. + val_ratio: Validation ratio used. + test_ratio: Test ratio used. + train_subset_indices: Optional optimization subset indices for + label-efficiency runs. When provided and different from + train_indices, persisted separately from the full split provenance. + label_fraction: Label fraction used for the optimization subset. + """ + static_df = dataset.static_df + stay_to_patient = dict(zip(static_df["stay_id"].to_list(), static_df["patient_id"].to_list())) + + train_patients = sorted({stay_to_patient[dataset.stay_ids[i]] for i in train_indices}) + val_patients = sorted({stay_to_patient[dataset.stay_ids[i]] for i in val_indices}) + test_patients = sorted({stay_to_patient[dataset.stay_ids[i]] for i in test_indices}) + + split_info = { + "seed": seed, + "train_ratio": train_ratio, + "val_ratio": val_ratio, + "test_ratio": test_ratio, + "label_fraction": label_fraction, + "train_patients": train_patients, + "val_patients": val_patients, + "test_patients": test_patients, + "train_stays": len(train_indices), + "val_stays": len(val_indices), + "test_stays": len(test_indices), + } + + if train_subset_indices is not None and train_subset_indices != train_indices: + train_subset_stay_ids = [dataset.stay_ids[i] for i in train_subset_indices] + train_subset_patients = sorted( + {stay_to_patient[dataset.stay_ids[i]] for i in train_subset_indices} + ) + split_info.update( + { + "train_subset_patients": train_subset_patients, + "train_subset_stays": len(train_subset_indices), + "train_subset_stay_ids": train_subset_stay_ids, + } + ) + + split_path = processed_dir / "splits.yaml" + with open(split_path, "w") as f: + yaml.dump(split_info, f, default_flow_style=False) diff --git a/src/slices/data/tasks/aki_kdigo.yaml b/src/slices/data/tasks/aki_kdigo.yaml index 5edbe2d..57ae24d 100644 --- a/src/slices/data/tasks/aki_kdigo.yaml +++ b/src/slices/data/tasks/aki_kdigo.yaml @@ -1,17 +1,22 @@ task_name: aki_kdigo task_type: binary -observation_window_hours: 48 -prediction_window_hours: null +observation_window_hours: 24 +prediction_window_hours: 24 # Look ahead 24h (hours 24-48) for AKI label gap_hours: 0 label_sources: - stays - timeseries label_params: creatinine_col: crea - baseline_window_hours: 24 + baseline_window_hours: 24 # Use full observation window for baseline absolute_rise_threshold: 0.3 relative_rise_threshold: 1.5 relative_window_hours: 168 # 7 days per KDIGO 2012 specification +quality_checks: + # AKI is evaluated on the creatinine-observed filtered cohort. eICU has + # materially higher missingness than MIMIC, so warn only above the current + # expected cross-dataset envelope. + max_missing_percentage: 35.0 primary_metric: auprc additional_metrics: - auroc diff --git a/src/slices/data/tasks/los_remaining.yaml b/src/slices/data/tasks/los_remaining.yaml index 17ccc2b..cabbeef 100644 --- a/src/slices/data/tasks/los_remaining.yaml +++ b/src/slices/data/tasks/los_remaining.yaml @@ -1,6 +1,6 @@ task_name: los_remaining task_type: regression -observation_window_hours: 48 +observation_window_hours: 24 prediction_window_hours: null gap_hours: 0 label_sources: [stays] diff --git a/src/slices/data/tasks/mortality.yaml b/src/slices/data/tasks/mortality.yaml index c2dc51d..3ba835e 100644 --- a/src/slices/data/tasks/mortality.yaml +++ b/src/slices/data/tasks/mortality.yaml @@ -2,19 +2,19 @@ # Predict if patient will die during the remainder of their ICU stay AFTER observation ends # # Timeline (with default settings): -# |---- observation (48h) ----|-- gap (0h) --|---- prediction (until ICU discharge) ----| -# intime +48h +48h outtime +# |---- observation (24h) ----|-- gap (0h) --|---- prediction (until ICU discharge) ----| +# intime +24h +24h outtime # -# - Model sees: hours 0-48 of ICU stay +# - Model sees: hours 0-24 of ICU stay # - Model predicts: will patient die before ICU discharge? -# - Patients who die during observation (hours 0-48) are excluded (label=null) +# - Patients who die during observation (hours 0-24) are excluded (label=null) task_name: mortality task_type: binary # Prediction parameters prediction_window_hours: -1 # -1 means "until ICU discharge" (entire remaining ICU stay) -observation_window_hours: 48 # Must match seq_length_hours in data config +observation_window_hours: 24 # Must match seq_length_hours in data config gap_hours: 0 # No gap between observation end and prediction start # Data requirements diff --git a/src/slices/data/tasks/mortality_24h.yaml b/src/slices/data/tasks/mortality_24h.yaml index 510494f..d762160 100644 --- a/src/slices/data/tasks/mortality_24h.yaml +++ b/src/slices/data/tasks/mortality_24h.yaml @@ -2,19 +2,19 @@ # Predict if patient will die within 24 hours AFTER the observation window ends # # Timeline (with default settings): -# |---- observation (48h) ----|-- gap (0h) --|---- prediction (24h) ----| -# intime +48h +48h +72h +# |---- observation (24h) ----|-- gap (0h) --|---- prediction (24h) ----| +# intime +24h +24h +48h # -# - Model sees: hours 0-48 of ICU stay -# - Model predicts: will patient die in hours 48-72? -# - Patients who die during observation (hours 0-48) are excluded (label=null) +# - Model sees: hours 0-24 of ICU stay +# - Model predicts: will patient die in hours 24-48? +# - Patients who die during observation (hours 0-24) are excluded (label=null) task_name: mortality_24h task_type: binary # Prediction parameters prediction_window_hours: 24 # Predict mortality in 24h window after observation ends -observation_window_hours: 48 # Must match seq_length_hours in data config +observation_window_hours: 24 # Must match seq_length_hours in data config gap_hours: 0 # No gap between observation end and prediction start # Note: If gap_hours > 0, deaths during the gap period are excluded (label=null) # similar to observation-period deaths, since they are ambiguous for prediction. diff --git a/src/slices/data/tasks/mortality_hospital.yaml b/src/slices/data/tasks/mortality_hospital.yaml index f3dab36..12f07ba 100644 --- a/src/slices/data/tasks/mortality_hospital.yaml +++ b/src/slices/data/tasks/mortality_hospital.yaml @@ -2,12 +2,12 @@ # Predict if patient will die before hospital discharge # # Timeline (with default settings): -# |---- observation (48h) ----|---- prediction (until hospital discharge) ----| -# intime +48h dischtime +# |---- observation (24h) ----|---- prediction (until hospital discharge) ----| +# intime +24h dischtime # -# - Model sees: hours 0-48 of ICU stay +# - Model sees: hours 0-24 of ICU stay # - Model predicts: will patient die before hospital discharge? -# - Patients who die during observation (hours 0-48) are excluded (label=null) +# - Patients who die during observation (hours 0-24) are excluded (label=null) # # Unlike mortality.yaml (ICU mortality using outtime), this uses # hospital_expire_flag which covers the full hospitalization. @@ -17,7 +17,7 @@ task_type: binary # Prediction parameters prediction_window_hours: null # null means hospital mortality (uses hospital_expire_flag) -observation_window_hours: 48 # Must match seq_length_hours in data config +observation_window_hours: 24 # Must match seq_length_hours in data config gap_hours: 0 # Data requirements diff --git a/src/slices/data/tensor_cache.py b/src/slices/data/tensor_cache.py new file mode 100644 index 0000000..7b842f6 --- /dev/null +++ b/src/slices/data/tensor_cache.py @@ -0,0 +1,508 @@ +"""Tensor caching and normalization stats I/O for ICU datasets. + +Handles loading/saving of preprocessed tensor caches and normalization +statistics to avoid recomputation on subsequent runs. +Extracted from ICUDataset for modularity. +""" + +import hashlib +import inspect +import json +import logging +import os +import tempfile +import warnings +from functools import lru_cache +from pathlib import Path +from typing import Any, Dict, List, Optional + +import torch +import yaml + +logger = logging.getLogger(__name__) +_CACHE_FINGERPRINT_VERSION = "2026-04-22" +_NORMALIZATION_STATS_VERSION = "full-cohort-train-v2" + + +def _fingerprint_payload(payload: Dict[str, Any]) -> str: + """Return a short stable fingerprint for a JSON-serializable payload.""" + content = json.dumps(payload, sort_keys=True, default=str) + return hashlib.md5(content.encode()).hexdigest()[:12] + + +def _path_signature(path: Path) -> Dict[str, Any]: + """Return a lightweight fingerprint for a file or directory path.""" + if not path.exists(): + return {"exists": False} + + stat = path.stat() + signature: Dict[str, Any] = { + "exists": True, + "is_dir": path.is_dir(), + "size": stat.st_size, + "mtime_ns": stat.st_mtime_ns, + } + + if path.is_dir(): + signature["children"] = sorted(child.name for child in path.iterdir()) + + return signature + + +def get_data_fingerprint(data_dir: Path) -> str: + """Fingerprint the processed dataset contents used to build caches.""" + payload = { + "version": _CACHE_FINGERPRINT_VERSION, + "metadata": _path_signature(data_dir / "metadata.yaml"), + "static": _path_signature(data_dir / "static.parquet"), + "timeseries": _path_signature(data_dir / "timeseries.parquet"), + "labels": _path_signature(data_dir / "labels.parquet"), + } + return _fingerprint_payload(payload) + + +@lru_cache(maxsize=1) +def get_preprocessing_fingerprint() -> str: + """Fingerprint the tensor-preprocessing code path that builds caches.""" + from slices.data import tensor_preprocessing as tp + + payload = {"version": _CACHE_FINGERPRINT_VERSION} + for name in ( + "extract_tensors_from_dataframe", + "convert_raw_to_tensors", + "compute_normalization_stats", + "apply_normalization_and_imputation", + ): + payload[name] = inspect.getsource(getattr(tp, name)) + return _fingerprint_payload(payload) + + +def _validate_cache_fingerprints( + cached: Dict[str, Any], + *, + artifact_name: str, + expected_data_fingerprint: str, + expected_preprocessing_fingerprint: str, +) -> bool: + """Return True when cached metadata matches current data and code fingerprints.""" + cached_data_fingerprint = cached.get("data_fingerprint") + cached_preprocessing_fingerprint = cached.get("preprocessing_fingerprint") + + if not cached_data_fingerprint or not cached_preprocessing_fingerprint: + warnings.warn( + f"{artifact_name} cache is missing freshness fingerprints. " + "Ignoring cache and recomputing.", + UserWarning, + ) + return False + + if cached_data_fingerprint != expected_data_fingerprint: + logger.debug("%s data fingerprint mismatch, recomputing", artifact_name) + return False + + if cached_preprocessing_fingerprint != expected_preprocessing_fingerprint: + logger.debug("%s preprocessing fingerprint mismatch, recomputing", artifact_name) + return False + + return True + + +def _compute_split_hash(train_indices: List[int], normalize: bool) -> str: + """Compute a stable hash from sorted train indices and normalize flag. + + Used to key normalization stats files so concurrent runs with different + splits write to different files instead of overwriting each other. + """ + content = f"{_NORMALIZATION_STATS_VERSION}|{sorted(train_indices)}|{normalize}" + return hashlib.md5(content.encode()).hexdigest()[:12] + + +def _normalization_stats_path(data_dir: Path, split_hash: str) -> Path: + """Return the hash-keyed normalization stats file path.""" + return data_dir / f"normalization_stats_{split_hash}.yaml" + + +def load_normalization_stats( + data_dir: Path, + current_train_indices: Optional[List[int]], + normalize: bool, +) -> Optional[Dict[str, Any]]: + """Load cached preprocessing statistics from file if they match current split. + + Looks for a hash-keyed file first (normalization_stats_.yaml), then + falls back to the legacy normalization_stats.yaml with set-comparison validation. + + Args: + data_dir: Path to data directory containing normalization stats. + current_train_indices: List of training indices for current split, + or None for unsupervised. + normalize: Whether normalization is enabled. + + Returns: + Dictionary with 'feature_means' and 'feature_stds' if file exists and split matches, + None otherwise. + """ + expected_data_fingerprint = get_data_fingerprint(data_dir) + expected_preprocessing_fingerprint = get_preprocessing_fingerprint() + + # Hash-keyed path (new format) — hash guarantees index match + if current_train_indices is not None: + split_hash = _compute_split_hash(current_train_indices, normalize) + hashed_path = _normalization_stats_path(data_dir, split_hash) + if hashed_path.exists(): + try: + with open(hashed_path) as f: + stats = yaml.safe_load(f) + if not _validate_cache_fingerprints( + stats, + artifact_name="Normalization stats", + expected_data_fingerprint=expected_data_fingerprint, + expected_preprocessing_fingerprint=expected_preprocessing_fingerprint, + ): + return None + logger.debug(f"Loaded normalization stats from {hashed_path}") + return stats + except Exception as e: + warnings.warn( + f"Failed to load normalization stats from {hashed_path}: {e}. " + "Will recompute statistics.", + UserWarning, + ) + return None + + # Legacy fallback — validate by comparing train_indices sets + legacy_path = data_dir / "normalization_stats.yaml" + if not legacy_path.exists(): + return None + + try: + with open(legacy_path) as f: + stats = yaml.safe_load(f) + + cached_train_indices = stats.get("train_indices") + current_train_set = set(current_train_indices) if current_train_indices else None + cached_train_set = set(cached_train_indices) if cached_train_indices else None + + if current_train_set != cached_train_set: + warnings.warn( + f"Cached normalization stats were computed on a different training split. " + f"Cached: {len(cached_train_set) if cached_train_set else 0} samples, " + f"Current: {len(current_train_set) if current_train_set else 0} samples. " + "Recomputing statistics to prevent data leakage.", + UserWarning, + ) + return None + + if not _validate_cache_fingerprints( + stats, + artifact_name="Legacy normalization stats", + expected_data_fingerprint=expected_data_fingerprint, + expected_preprocessing_fingerprint=expected_preprocessing_fingerprint, + ): + return None + + return stats + except Exception as e: + warnings.warn( + f"Failed to load normalization stats from {legacy_path}: {e}. " + "Will recompute statistics.", + UserWarning, + ) + return None + + +def save_normalization_stats( + data_dir: Path, + feature_means: torch.Tensor, + feature_stds: torch.Tensor, + feature_names: List[str], + train_indices: Optional[List[int]], + normalize: bool, +) -> None: + """Save computed normalization statistics to a hash-keyed file atomically. + + Uses tempfile + os.replace for atomic writes, preventing corruption from + concurrent runs. The file is keyed by a hash of train_indices + normalize, + so different splits write to different files. + + Args: + data_dir: Path to data directory. + feature_means: Per-feature means tensor. + feature_stds: Per-feature stds tensor. + feature_names: List of feature names. + train_indices: Optional list of training indices used to compute stats. + normalize: Whether normalization is enabled. + """ + if train_indices is not None: + split_hash = _compute_split_hash(train_indices, normalize) + stats_path = _normalization_stats_path(data_dir, split_hash) + else: + split_hash = "unsupervised" + stats_path = data_dir / "normalization_stats.yaml" + + stats = { + "feature_means": feature_means.tolist(), + "feature_stds": feature_stds.tolist(), + "feature_names": feature_names, + "normalization_stats_version": _NORMALIZATION_STATS_VERSION, + "normalization_index_space": "raw_full_cohort", + "split_hash": split_hash, + "train_indices_count": len(train_indices) if train_indices else 0, + "train_indices": train_indices, + "normalize": normalize, + "data_fingerprint": get_data_fingerprint(data_dir), + "preprocessing_fingerprint": get_preprocessing_fingerprint(), + } + + try: + fd, tmp_path = tempfile.mkstemp(dir=data_dir, suffix=".yaml.tmp") + try: + with os.fdopen(fd, "w") as f: + yaml.dump(stats, f, default_flow_style=False) + os.replace(tmp_path, stats_path) + except Exception: + try: + os.unlink(tmp_path) + except OSError: + pass + raise + except Exception as e: + warnings.warn( + f"Failed to save normalization stats to {stats_path}: {e}", + UserWarning, + ) + + +def get_tensor_cache_key( + seq_length: int, + n_features: int, +) -> str: + """Generate a hash key for raw tensor caching based on shape parameters. + + The cache stores raw (unnormalized) tensors for the full dataset. + Normalization is applied at runtime using run-specific train_indices, + so the cache key only depends on tensor shape parameters. + + Args: + seq_length: Sequence length. + n_features: Number of features. + + Returns: + Hash string to use as cache identifier. + """ + params = { + "seq_length": seq_length, + "n_features": n_features, + } + params_str = str(sorted(params.items())) + return hashlib.md5(params_str.encode()).hexdigest()[:12] + + +def get_tensor_cache_path( + data_dir: Path, + seq_length: int, + n_features: int, +) -> Path: + """Get the path to the raw tensor cache file. + + Args: + data_dir: Path to data directory. + seq_length: Sequence length. + n_features: Number of features. + + Returns: + Path to the tensor cache file. + """ + cache_key = get_tensor_cache_key(seq_length, n_features) + cache_dir = data_dir / ".tensor_cache" + return cache_dir / f"tensors_{cache_key}.pt" + + +def _try_load_cache( + cache_path: Path, + n_features: int, + seq_length: int, + *, + expected_data_fingerprint: str, + expected_preprocessing_fingerprint: str, +) -> Optional[Dict[str, Any]]: + """Try to load and validate a raw tensor cache file. + + Args: + cache_path: Path to the cache file. + n_features: Expected number of features. + seq_length: Expected sequence length. + + Returns: + Dictionary with cached raw tensors if valid, None otherwise. + """ + if not cache_path.exists(): + return None + + try: + logger.debug(f"Loading cached tensors from {cache_path.name}") + cached = torch.load(cache_path, weights_only=True) + + # Validate cache metadata + if cached.get("n_features") != n_features: + logger.debug("Cached tensors have different feature count, recomputing") + return None + if cached.get("seq_length") != seq_length: + logger.debug("Cached tensors have different sequence length, recomputing") + return None + if not _validate_cache_fingerprints( + cached, + artifact_name="Tensor", + expected_data_fingerprint=expected_data_fingerprint, + expected_preprocessing_fingerprint=expected_preprocessing_fingerprint, + ): + return None + + # Validate tensor shapes (support both old list and new stacked format) + # NOTE: cannot use `x or y` here — `or` calls bool() on tensors, which + # raises RuntimeError for multi-element tensors. + timeseries = cached.get("timeseries_tensor") + if timeseries is None: + timeseries = cached.get("timeseries_tensors") + masks = cached.get("mask_tensor") + if masks is None: + masks = cached.get("mask_tensors") + if timeseries is None or masks is None: + logger.debug("Cached tensors missing data, recomputing") + return None + + # Convert old list format to stacked if needed + if isinstance(timeseries, list): + logger.debug("Converting cached tensors from list to stacked format") + timeseries = torch.stack(timeseries) + masks = torch.stack(masks) + cached["timeseries_tensor"] = timeseries + cached["mask_tensor"] = masks + + n_samples = timeseries.shape[0] if hasattr(timeseries, "shape") else len(timeseries) + logger.info(f"Loaded {n_samples:,} cached raw samples") + return cached + + except Exception as e: + logger.debug(f"Failed to load cached tensors: {e}, recomputing") + return None + + +def load_cached_tensors( + data_dir: Path, + seq_length: int, + n_features: int, +) -> Optional[Dict[str, Any]]: + """Load cached raw tensors if they exist and are valid. + + The cache stores raw (unnormalized) tensors for the full dataset. + Normalization and stay filtering are applied after loading. + + Args: + data_dir: Path to data directory. + seq_length: Sequence length. + n_features: Number of features. + + Returns: + Dictionary with cached raw tensors and metadata if valid, None otherwise. + """ + cache_path = get_tensor_cache_path(data_dir, seq_length, n_features) + return _try_load_cache( + cache_path, + n_features, + seq_length, + expected_data_fingerprint=get_data_fingerprint(data_dir), + expected_preprocessing_fingerprint=get_preprocessing_fingerprint(), + ) + + +def save_cached_tensors( + data_dir: Path, + timeseries_tensor: torch.Tensor, + mask_tensor: torch.Tensor, + seq_length: int, + n_features: int, +) -> None: + """Save raw tensors to cache file atomically. + + Stores raw (unnormalized) tensors for the full dataset. Normalization + and stay filtering are applied at runtime after loading. + + Args: + data_dir: Path to data directory. + timeseries_tensor: Raw timeseries tensor (unnormalized). + mask_tensor: Observation mask tensor. + seq_length: Sequence length. + n_features: Number of features. + """ + cache_path = get_tensor_cache_path(data_dir, seq_length, n_features) + cache_dir = cache_path.parent + cache_dir.mkdir(parents=True, exist_ok=True) + + cache_data = { + "timeseries_tensor": timeseries_tensor, + "mask_tensor": mask_tensor, + "n_features": n_features, + "seq_length": seq_length, + "data_fingerprint": get_data_fingerprint(data_dir), + "preprocessing_fingerprint": get_preprocessing_fingerprint(), + } + + try: + logger.debug(f"Saving raw tensors to cache: {cache_path.name}") + fd, tmp_path = tempfile.mkstemp(dir=cache_dir, suffix=".pt.tmp") + os.close(fd) + try: + torch.save(cache_data, tmp_path) + os.replace(tmp_path, cache_path) + except Exception: + try: + os.unlink(tmp_path) + except OSError: + pass + raise + except Exception as e: + logger.warning(f"Failed to save tensor cache to {cache_path}: {e}") + + +def save_dataset_metadata( + data_dir: Path, + task_name: Optional[str], + handle_missing_labels: str, + removed_samples: List[tuple], +) -> None: + """Save dataset metadata including removed samples for reproducibility. + + Creates or updates dataset_metadata.yaml with information about: + - Number of samples used (original vs. final after filtering) + - Removed samples and reasons for removal + - Task name and label handling strategy + + Args: + data_dir: Path to data directory. + task_name: Name of the task. + handle_missing_labels: Label handling strategy. + removed_samples: List of (stay_id, reason) tuples. + """ + if not removed_samples: + return # Nothing to save if no samples were removed + + metadata_path = data_dir / "dataset_metadata.yaml" + + metadata = { + "task_name": task_name, + "handle_missing_labels": handle_missing_labels, + "removed_samples_count": len(removed_samples), + "removed_samples": [ + {"stay_id": stay_id, "reason": reason} for stay_id, reason in removed_samples + ], + } + + try: + with open(metadata_path, "w") as f: + yaml.dump(metadata, f, default_flow_style=False) + except Exception as e: + warnings.warn( + f"Failed to save dataset metadata to {metadata_path}: {e}", + UserWarning, + ) diff --git a/src/slices/data/tensor_preprocessing.py b/src/slices/data/tensor_preprocessing.py new file mode 100644 index 0000000..9b12693 --- /dev/null +++ b/src/slices/data/tensor_preprocessing.py @@ -0,0 +1,334 @@ +"""Tensor preprocessing functions for ICU time-series data. + +Pure functions for converting raw nested-list data to padded tensors, +computing normalization statistics, and applying normalization + imputation. +Extracted from ICUDataset._precompute_tensors for modularity. +""" + +import logging +from typing import TYPE_CHECKING, Dict, List + +import numpy as np + +if TYPE_CHECKING: + import polars as pl +import torch +from tqdm import tqdm + +logger = logging.getLogger(__name__) + +# Constants +MIN_STD_THRESHOLD = 1e-6 # Minimum standard deviation to avoid division by zero + + +def extract_tensors_from_dataframe( + timeseries_df: "pl.DataFrame", + seq_length: int, + n_features: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Extract tensors directly from a Polars DataFrame using zero-copy operations. + + Uses Polars explode() to flatten nested List(List(...)) columns into flat + arrays, then reshapes to 3D tensors. This avoids to_list() which creates + Python objects (~28 bytes per float) and can use 3-4x more memory than the + actual data, causing OOM on large datasets. + + Requires uniform sequence lengths (all rows must have exactly seq_length + timesteps with n_features features each). Falls back to convert_raw_to_tensors + if sequences are non-uniform. + + Args: + timeseries_df: Polars DataFrame with 'timeseries' and 'mask' columns, + where each is List(List(Float64|Boolean)). + seq_length: Expected sequence length per sample. + n_features: Expected number of features per timestep. + + Returns: + Tuple of (timeseries_tensor, masks_tensor), each shaped + (n_samples, seq_length, n_features). + """ + import polars as pl + + n_samples = len(timeseries_df) + logger.debug("[1/3] Extracting tensors from DataFrame...") + + # Verify uniform sequence lengths for reshape safety + ts_lens = timeseries_df["timeseries"].list.len() + if ts_lens.min() != ts_lens.max(): + logger.info("Non-uniform sequence lengths detected, falling back to list conversion") + raw_timeseries = timeseries_df["timeseries"].to_list() + raw_masks = timeseries_df["mask"].to_list() + return convert_raw_to_tensors(raw_timeseries, raw_masks, seq_length, n_features) + + # Check if stored sequence length differs from requested — fall back to + # convert_raw_to_tensors which handles truncation/padding correctly + stored_seq_length = int(ts_lens.min()) # type: ignore[arg-type] + if stored_seq_length != seq_length: + logger.info( + "Stored sequence length (%d) differs from requested (%d)," + " falling back to list conversion for truncation/padding", + stored_seq_length, + seq_length, + ) + raw_timeseries = timeseries_df["timeseries"].to_list() + raw_masks = timeseries_df["mask"].to_list() + return convert_raw_to_tensors(raw_timeseries, raw_masks, seq_length, n_features) + + # Fast path: explode nested lists to flat array, then reshape. + # This stays in Arrow/numpy memory without creating Python objects. + # Extract columns as Series first, then process sequentially to limit + # peak memory (avoids holding both float64 intermediate + final tensor). + ts_series = timeseries_df["timeseries"] + mask_series = timeseries_df["mask"] + + # Cast to Float32 in Polars to avoid a float64 intermediate array (~5 GB saved) + timeseries_np = ( + ts_series.explode() + .explode() + .cast(pl.Float32) + .to_numpy(writable=True) + .reshape(n_samples, seq_length, n_features) + ) + del ts_series + timeseries_tensor = torch.from_numpy(timeseries_np) + + masks_np = mask_series.explode().explode().to_numpy().reshape(n_samples, seq_length, n_features) + del mask_series + masks_tensor = torch.from_numpy(masks_np.copy()) # copy needed: bool dtype + del masks_np + + return timeseries_tensor, masks_tensor + + +def convert_raw_to_tensors( + raw_timeseries: List[List[List[float]]], + raw_masks: List[List[List[bool]]], + seq_length: int, + n_features: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Convert nested lists to padded tensors (Step 1 of preprocessing). + + Converts all samples to numpy arrays in batches for memory efficiency, + then stacks into single tensors. Samples shorter than seq_length are + padded with NaN/False; longer samples are truncated. + + Note: For large datasets, prefer extract_tensors_from_dataframe() which + avoids the expensive to_list() conversion. + + Args: + raw_timeseries: List of timeseries arrays (n_samples x seq_len x n_features). + raw_masks: List of mask arrays (n_samples x seq_len x n_features). + seq_length: Target sequence length (pad/truncate to this). + n_features: Number of features per timestep. + + Returns: + Tuple of (timeseries_tensor, masks_tensor), each shaped + (n_samples, seq_length, n_features). + """ + n_samples = len(raw_timeseries) + logger.debug("[1/3] Converting to tensors...") + + # Pre-allocate output arrays to avoid 2x peak memory from list + np.stack + timeseries_np = np.full((n_samples, seq_length, n_features), np.nan, dtype=np.float32) + masks_np = np.zeros((n_samples, seq_length, n_features), dtype=bool) + + sample_iter = range(n_samples) + if n_samples >= 50000: # Only show progress for large datasets + sample_iter = tqdm(sample_iter, desc=" Converting samples", unit="sample") + + for i in sample_iter: + ts_data = raw_timeseries[i] + mask_data = raw_masks[i] + actual_len = min(len(ts_data), seq_length) + + timeseries_np[i, :actual_len] = np.array(ts_data[:actual_len], dtype=np.float32) + masks_np[i, :actual_len] = np.array(mask_data[:actual_len], dtype=bool) + + # Convert to tensors (torch.from_numpy shares memory, no copy) + timeseries_tensor = torch.from_numpy(timeseries_np) # (n_samples, seq_len, n_features) + masks_tensor = torch.from_numpy(masks_np) # (n_samples, seq_len, n_features) + + return timeseries_tensor, masks_tensor + + +def compute_normalization_stats( + timeseries_tensor: torch.Tensor, + masks_tensor: torch.Tensor, + train_indices: List[int], + n_features: int, + normalize: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute per-feature mean and std from training samples only (Step 2). + + CRITICAL: Uses only training samples to prevent data leakage from + validation/test sets. + + Args: + timeseries_tensor: Full timeseries tensor (n_samples, seq_len, n_features). + masks_tensor: Full mask tensor (n_samples, seq_len, n_features). + train_indices: Indices of training samples. + n_features: Number of features. + normalize: Whether normalization is enabled (affects std computation). + + Returns: + Tuple of (feature_means, feature_stds), each shaped (n_features,). + """ + logger.debug("[2/3] Computing normalization statistics...") + + feature_means = torch.zeros(n_features) + feature_stds = torch.ones(n_features) + + # CRITICAL: Use only training samples to prevent data leakage + train_ts = timeseries_tensor[train_indices] # (n_train, seq_len, n_features) + train_masks = masks_tensor[train_indices] + + # Vectorized computation of mean and std per feature + # Reshape to (n_samples * seq_len, n_features) for easier computation + flat_ts = train_ts.reshape(-1, n_features) # (n_train * seq_len, n_features) + flat_masks = train_masks.reshape(-1, n_features) + + # Create combined mask: observed AND not NaN + valid_mask = flat_masks & ~torch.isnan(flat_ts) + + # Vectorized mean computation using masked tensor operations + # Replace invalid values with 0 for sum, count valid entries + masked_ts = torch.where(valid_mask, flat_ts, torch.zeros_like(flat_ts)) + valid_counts = valid_mask.sum(dim=0).float() # (n_features,) + feature_sums = masked_ts.sum(dim=0) # (n_features,) + + # Compute means (avoid div by zero) + feature_means = torch.where( + valid_counts > 0, + feature_sums / valid_counts, + torch.zeros(n_features), + ) + + if normalize: + # Vectorized std computation + # Compute squared deviations from mean + deviations = torch.where( + valid_mask, + (flat_ts - feature_means.unsqueeze(0)) ** 2, + torch.zeros_like(flat_ts), + ) + variance = torch.where( + valid_counts > 1, + deviations.sum(dim=0) / (valid_counts - 1), # Bessel's correction + torch.ones(n_features), + ) + feature_stds = torch.sqrt(variance) + + # Clamp minimum std to avoid division by zero + feature_stds = torch.clamp(feature_stds, min=MIN_STD_THRESHOLD) + + logger.debug(f"Computed normalization stats for {n_features} features") + return feature_means, feature_stds + + +def apply_normalization_and_imputation( + timeseries_tensor: torch.Tensor, + feature_means: torch.Tensor, + feature_stds: torch.Tensor, + normalize: bool, + n_features: int, +) -> torch.Tensor: + """Normalize then impute missing values (Step 3 of preprocessing). + + When normalize=True: z-score normalize, then zero-fill (0 = population mean). + When normalize=False: impute with feature means when available. If callers + provide zero means, this degenerates to an explicit zero-fill fallback. + + Args: + timeseries_tensor: Tensor to normalize (n_samples, seq_len, n_features). + Modified in-place and returned. + feature_means: Per-feature means (n_features,). + feature_stds: Per-feature stds (n_features,). + normalize: Whether to apply z-score normalization. + n_features: Number of features. + + Returns: + The normalized/imputed timeseries tensor. + """ + logger.debug("[3/3] Applying normalization and imputation...") + + if normalize: + # Normalize first (NaN positions stay NaN through arithmetic) + timeseries_tensor = (timeseries_tensor - feature_means) / feature_stds + # After z-score normalization, 0 = feature mean in original space. + # This is the most neutral default: "no information, assume population average." + timeseries_tensor = torch.nan_to_num(timeseries_tensor, nan=0.0) + else: + # Without normalization, keep data in original units and impute each + # missing value with the corresponding feature mean. + for f in range(n_features): + nan_mask = torch.isnan(timeseries_tensor[:, :, f]) + timeseries_tensor[:, :, f][nan_mask] = feature_means[f] + + return timeseries_tensor + + +def compute_single_sample_stages( + raw_ts: List[List[float]], + raw_mask: List[List[bool]], + seq_length: int, + n_features: int, + feature_means: torch.Tensor, + feature_stds: torch.Tensor, + normalize: bool, +) -> Dict[str, Dict[str, torch.Tensor]]: + """Re-run the preprocessing pipeline for a single sample, capturing each stage. + + Useful for debugging to see exactly what transformations are applied. + + The stages are: + - grid: Raw 2D tensor (seq_length, n_features) with NaN for missing + - normalized: After z-score normalization + zero-fill (what model sees) + + Args: + raw_ts: Raw timeseries nested list for one sample. + raw_mask: Raw mask nested list for one sample. + seq_length: Target sequence length. + n_features: Number of features. + feature_means: Per-feature means from training set. + feature_stds: Per-feature stds from training set. + normalize: Whether normalization is enabled. + + Returns: + Dict with keys 'grid', 'normalized', each containing: + - 'timeseries': The tensor at that stage + - 'mask': The observation mask (same for all stages) + """ + # Stage 1: GRID - Convert to tensor with padding/truncation + actual_len = len(raw_ts) + if actual_len >= seq_length: + grid_ts = np.array(raw_ts[:seq_length], dtype=np.float32) + mask_arr = np.array(raw_mask[:seq_length], dtype=bool) + else: + grid_ts = np.full((seq_length, n_features), np.nan, dtype=np.float32) + mask_arr = np.zeros((seq_length, n_features), dtype=bool) + grid_ts[:actual_len] = np.array(raw_ts, dtype=np.float32) + mask_arr[:actual_len] = np.array(raw_mask, dtype=bool) + + grid_tensor = torch.from_numpy(grid_ts) + mask_tensor = torch.from_numpy(mask_arr) + + # Stage 2: NORMALIZED - z-score normalize then impute + if normalize: + normalized_tensor = (grid_tensor - feature_means) / feature_stds + normalized_tensor = torch.nan_to_num(normalized_tensor, nan=0.0) + else: + normalized_tensor = grid_tensor.clone() + for f in range(n_features): + nan_mask = torch.isnan(normalized_tensor[:, f]) + normalized_tensor[:, f][nan_mask] = feature_means[f] + + return { + "grid": { + "timeseries": grid_tensor, + "mask": mask_tensor, + }, + "normalized": { + "timeseries": normalized_tensor, + "mask": mask_tensor, + }, + } diff --git a/src/slices/debug/embeddings.py b/src/slices/debug/embeddings.py index ea24cc6..e2b9042 100644 --- a/src/slices/debug/embeddings.py +++ b/src/slices/debug/embeddings.py @@ -604,7 +604,7 @@ def load_embeddings_from_file( raise ImportError("PyTorch is required to load .pt files.") import torch - data = torch.load(path, map_location="cpu") + data = torch.load(path, map_location="cpu", weights_only=True) if isinstance(data, dict): embeddings = _to_numpy(data["embeddings"]) stay_ids = _to_numpy(data.get("stay_ids")) if "stay_ids" in data else None diff --git a/src/slices/debug/plots.py b/src/slices/debug/plots.py index 6afd7eb..efdb088 100644 --- a/src/slices/debug/plots.py +++ b/src/slices/debug/plots.py @@ -348,9 +348,11 @@ def plot_missingness_heatmap( # Create figure with subplots for each patient fig, axes = plt.subplots(n_stays_actual, 1, figsize=(12, 2 * n_stays_actual), sharex=True) if n_stays_actual == 1: - axes = [axes] + axes_list: List[Any] = [axes] + else: + axes_list = list(axes) - for idx, ax in enumerate(axes): + for idx, ax in enumerate(axes_list): ax.imshow( mask[idx].T, aspect="auto", cmap="Greens", vmin=0, vmax=1, interpolation="nearest" ) @@ -365,7 +367,7 @@ def plot_missingness_heatmap( ax.set_yticks(range(n_features)) ax.set_yticklabels(feature_names, fontsize=6) - axes[-1].set_xlabel("Hour") + axes_list[-1].set_xlabel("Hour") fig.suptitle(f"{title}\n(Green = Observed)") fig.tight_layout() diff --git a/src/slices/debug/sampling.py b/src/slices/debug/sampling.py index 20689bb..c3ab14e 100644 --- a/src/slices/debug/sampling.py +++ b/src/slices/debug/sampling.py @@ -151,7 +151,7 @@ def get_default_sentinel_slots() -> List[SentinelSlot]: 2. Short stay + died - label alignment edge case 3. Medium stay + low missingness - "clean" baseline (1-3 days) 4. Medium stay + high missingness - test imputation (1-3 days) - 5. Long stay (>5 days) - test truncation to 48h window + 5. Long stay (>5 days) - test truncation to the 24h benchmark window 6. Young patient (<40) - demographic edge 7. Old patient (>80) - demographic edge 8. Random sample - catch unexpected issues diff --git a/src/slices/debug/snapshots.py b/src/slices/debug/snapshots.py index 70d9903..a627c1f 100644 --- a/src/slices/debug/snapshots.py +++ b/src/slices/debug/snapshots.py @@ -14,6 +14,8 @@ import polars as pl import yaml +from slices.constants import SEQ_LENGTH_HOURS + # Import the canonical PipelineStage from staged_snapshots # This module previously had its own PipelineStage with different stage names. # For backwards compatibility, we keep the old stage names as aliases. @@ -43,6 +45,9 @@ class LegacyPipelineStage(str, Enum): FINAL = "final" +SnapshotStage = Union[PipelineStage, LegacyPipelineStage] + + @dataclass class SnapshotConfig: """Configuration for pipeline snapshots. @@ -57,9 +62,11 @@ class SnapshotConfig: """ output_dir: Union[str, Path] = "debug_snapshots" - stages: List[PipelineStage] = field(default_factory=lambda: list(PipelineStage)) + stages: List[SnapshotStage] = field( + default_factory=lambda: list(PipelineStage) + list(LegacyPipelineStage) + ) stay_ids: Optional[List[int]] = None - max_hours: int = 48 + max_hours: int = SEQ_LENGTH_HOURS include_masks: bool = True flatten_arrays: bool = True @@ -75,7 +82,7 @@ class PipelineSnapshot: timestamp: When snapshot was captured. """ - stage: PipelineStage + stage: SnapshotStage data: pl.DataFrame metadata: Dict[str, Any] = field(default_factory=dict) timestamp: Optional[str] = None @@ -109,7 +116,7 @@ def capture_stays_snapshot( df = df.filter(pl.col("stay_id").is_in(stay_ids)) return PipelineSnapshot( - stage=PipelineStage.STAYS, + stage=LegacyPipelineStage.STAYS, data=df, metadata={ "n_stays": len(df), @@ -147,7 +154,7 @@ def capture_labels_snapshot( } return PipelineSnapshot( - stage=PipelineStage.LABELS, + stage=LegacyPipelineStage.LABELS, data=df, metadata={ "n_stays": len(df), @@ -182,7 +189,7 @@ def capture_dense_snapshot( df = flatten_dense_timeseries(df, feature_names) return PipelineSnapshot( - stage=PipelineStage.DENSE_TIMESERIES, + stage=LegacyPipelineStage.DENSE_TIMESERIES, data=df, metadata={ "n_stays": ( @@ -305,7 +312,7 @@ def flatten_dense_timeseries( def unflatten_timeseries( flat_df: pl.DataFrame, feature_names: List[str], - seq_length: int = 48, + seq_length: int = SEQ_LENGTH_HOURS, ) -> pl.DataFrame: """Convert flattened timeseries back to nested format. @@ -393,9 +400,9 @@ def export_snapshot( def export_all_snapshots( - snapshots: Dict[PipelineStage, PipelineSnapshot], + snapshots: Dict[SnapshotStage, PipelineSnapshot], config: SnapshotConfig, -) -> Dict[PipelineStage, Path]: +) -> Dict[SnapshotStage, Path]: """Export all captured snapshots to CSV. Args: @@ -524,7 +531,7 @@ class DebuggableExtractor(SnapshotMixin, RicuExtractor): """ _snapshot_config: Optional[SnapshotConfig] = None - _snapshots: Dict[PipelineStage, PipelineSnapshot] = {} + _snapshots: Dict[SnapshotStage, PipelineSnapshot] = {} def enable_snapshots(self, config: SnapshotConfig) -> None: """Enable snapshot capture for this extractor run. @@ -537,7 +544,7 @@ def enable_snapshots(self, config: SnapshotConfig) -> None: def capture_snapshot( self, - stage: PipelineStage, + stage: SnapshotStage, data: pl.DataFrame, metadata: Optional[Dict[str, Any]] = None, ) -> None: @@ -566,7 +573,7 @@ def capture_snapshot( metadata=metadata or {}, ) - def get_snapshots(self) -> Dict[PipelineStage, PipelineSnapshot]: + def get_snapshots(self) -> Dict[SnapshotStage, PipelineSnapshot]: """Get all captured snapshots. Returns: @@ -574,7 +581,7 @@ def get_snapshots(self) -> Dict[PipelineStage, PipelineSnapshot]: """ return self._snapshots.copy() - def export_snapshots(self) -> Dict[PipelineStage, Path]: + def export_snapshots(self) -> Dict[SnapshotStage, Path]: """Export all captured snapshots to disk. Returns: diff --git a/src/slices/eval/README.md b/src/slices/eval/README.md index aff5355..fc8c19f 100644 --- a/src/slices/eval/README.md +++ b/src/slices/eval/README.md @@ -1,275 +1,118 @@ # Evaluation Module -This module provides configurable metrics and evaluation utilities for clinical prediction tasks. +This package covers downstream metrics, fairness analysis, statistical testing, +inference helpers, and SSL imputation checks. -## Overview +## What Is Implemented -The `eval` module separates evaluation concerns from training logic: -- **Configurable metrics**: Specify which metrics to compute via config -- **Task-specific defaults**: Sensible minimal metrics per task type -- **Extensible design**: Easy to add new metrics and analysis functions -- **Fairness placeholders**: Interface designed for future fairness analysis +- Configurable metric collections in `metrics.py` +- Binary, multiclass, multilabel, and regression metrics +- Fairness primitives in `fairness.py` +- Structured subgroup evaluation in `fairness_evaluator.py` +- Statistical summaries and pairwise tests in `statistical.py` +- Batched inference helper in `inference.py` +- Reconstruction/imputation evaluation for SSL encoders in `imputation.py` -## Quick Start +## Metrics -### Basic Usage +`MetricConfig` and `build_metrics()` are the main entry points. ```python from slices.eval import MetricConfig, build_metrics -# Create metric configuration -config = MetricConfig( - task_type="binary", - n_classes=2, - metrics=["auroc", "auprc"], -) +config = MetricConfig(task_type="binary", threshold=0.5) -# Build metrics with prefix for logging val_metrics = build_metrics(config, prefix="val") -test_metrics = build_metrics(config, prefix="test") - -# Use in training loop -val_metrics.update(predictions, labels) -results = val_metrics.compute() -# Results: {'val/auroc': 0.85, 'val/auprc': 0.78} -``` - -### Using Defaults - -```python -from slices.eval import get_default_metrics - -# Get default metrics for task type -val_metrics = get_default_metrics("binary", prefix="val") -# Uses default: [auroc, auprc] ``` -### In Configuration Files - -```yaml -# configs/finetune.yaml or configs/eval/default.yaml -eval: - metrics: - # Specify metrics to compute (null = use defaults) - names: [auroc, auprc, accuracy, f1] - - # Decision threshold for binary classification - threshold: 0.5 -``` - -## Available Metrics - -### Binary Classification -- `auroc` - Area Under ROC Curve (default) -- `auprc` - Area Under Precision-Recall Curve (default) -- `accuracy` - Accuracy at 0.5 threshold -- `f1` - F1 Score - -### Multiclass Classification -- `auroc` - One-vs-Rest AUROC (default) -- `auprc` - One-vs-Rest AUPRC -- `accuracy` - Multiclass accuracy (default) -- `f1` - Macro F1 Score +Supported task types: -### Multilabel Classification -- `auroc` - Per-label AUROC (default) -- `auprc` - Per-label AUPRC -- `accuracy` - Subset accuracy -- `f1` - Macro F1 Score +- `binary`: `auroc`, `auprc`, `accuracy`, `f1`, `precision`, `recall`, `specificity`, `brier_score`, `ece` +- `multiclass`: `auroc`, `auprc`, `accuracy`, `f1`, `precision`, `recall` +- `multilabel`: `auroc`, `auprc`, `accuracy`, `f1` +- `regression`: `mse`, `mae`, `rmse`, `r2` -### Regression -- Not yet implemented (placeholder) -- Planned: MSE, MAE, R², RMSE +Default metric sets match the public benchmark export surface: -## Integration with Training +- Binary: `auroc`, `auprc`, `accuracy`, `f1`, `precision`, `recall`, `specificity`, `brier_score`, `ece` +- Multiclass: `auroc`, `accuracy` +- Multilabel: `auroc` +- Regression: `mse`, `mae`, `r2` -The `FineTuneModule` automatically uses eval metrics: +## Fairness -```python -from slices.training import FineTuneModule - -# Metrics are built from config -module = FineTuneModule(config) +There are two layers: -# During training, metrics are logged automatically -# val/auroc, val/auprc, test/auroc, test/auprc, etc. -``` +1. `fairness.py` provides reusable metric primitives such as demographic parity + difference, equalized odds difference, and disparate impact ratio. +2. `FairnessEvaluator` in `fairness_evaluator.py` builds full subgroup reports + from predictions, labels, stay IDs, and `static.parquet` demographics. -Metrics are read from `config.eval.metrics.names` if specified, otherwise defaults are used. +Current evaluator behavior matches the benchmark pipeline: -## Extending with New Metrics +- sex and age-group analysis for all supported datasets when columns exist +- race/ethnicity fairness only on MIMIC-IV rows +- canonical race bins: `White`, `Black`, `Hispanic`, `Asian`, `Other` +- subgroup thresholds enforced on unique patients rather than stays +- flattened `fairness/*` output via `flatten_fairness_report()` -### Adding a New Metric +Typical usage: -1. **Update `AVAILABLE_METRICS` dict**: ```python -AVAILABLE_METRICS = { - "binary": ["auroc", "auprc", "accuracy", "f1", "sensitivity"], # Added - ... -} -``` - -2. **Add to `_build_metric()` function**: -```python -def _build_metric(name: str, task_type: TaskType, n_classes: int): - # ... - elif name == "sensitivity": - return Recall(task=task, **kwargs) - # ... -``` +from slices.eval import FairnessEvaluator, flatten_fairness_report -3. **Update docstring** to document the new metric - -### Adding a New Task Type - -1. Add to `TaskType` literal -2. Add to `AVAILABLE_METRICS` and `DEFAULT_METRICS` -3. Handle in `_build_metric()` function - -## Fairness Analysis (Planned) - -The `fairness.py` module provides placeholders for future fairness analysis: - -```python -from slices.eval.fairness import ( - compute_subgroup_metrics, - compute_demographic_parity, - compute_equalized_odds, +evaluator = FairnessEvaluator( + static_df=static_df, + protected_attributes=["gender", "age_group", "race"], + min_subgroup_size=50, + task_type="binary", + dataset_name="miiv", ) -# Not yet implemented - raises NotImplementedError -# Interface designed for future expansion +report = evaluator.evaluate(predictions, labels, stay_ids) +flat_report = flatten_fairness_report(report) ``` -### Planned Features +For batch benchmark sweeps, the repository uses +`scripts/eval/evaluate_fairness.py`, which reruns inference from the recorded +evaluation checkpoint and writes `fairness/*` keys back to W&B. -- **Subgroup performance**: Metrics per demographic group (age, gender, race, insurance) -- **Demographic parity**: Equal positive prediction rates across groups -- **Equalized odds**: Equal TPR/FPR across groups -- **Calibration analysis**: Calibration curves per subgroup +## Statistical Utilities -### Available Demographics (MIMIC-IV) +`statistical.py` provides the utilities used by the export pipeline: -The extractor provides these demographic attributes in `static.parquet`: -- `age` - Patient age at admission -- `gender` - Patient gender (M/F) -- `race` - Patient race/ethnicity -- `insurance` - Insurance type -- `admission_type` - Emergency vs elective -- `admission_location` - Where patient was admitted from +- `bootstrap_ci` +- `paired_bootstrap_test` +- `paired_wilcoxon_signed_rank` +- `bonferroni_correction` +- `cohens_d` -## Design Principles +These are exported from `slices.eval` for direct reuse. -### Minimal but Extendable -- Start with essential metrics (AUROC, AUPRC) -- Easy to add more without breaking existing code -- No unnecessary complexity +## Inference Helper -### Configuration-Driven -- Metrics specified in YAML config files -- Sensible defaults that "just work" -- Override when needed for experiments +`run_inference()` standardizes batched evaluation over a dataloader and returns +predictions, labels, and stay IDs in a format that downstream fairness and +export code can consume. -### Separation of Concerns -- Evaluation logic separate from training logic -- Metrics independent of Lightning module -- Can be used standalone or in training +## Imputation Evaluation -### Clinical Focus -- Metrics relevant to clinical prediction (AUROC, calibration) -- Demographics available for fairness analysis -- Designed for high-stakes decision support +`ImputationEvaluator` measures how well an SSL encoder reconstructs masked +values under controlled masking schemes. This is separate from downstream task +metrics and is intended for representation diagnostics rather than benchmark +headline results. -## Examples - -### Different Task Types - -```python -# Binary classification (mortality) -binary_config = MetricConfig(task_type="binary", metrics=["auroc", "auprc"]) -metrics = build_metrics(binary_config, prefix="val") - -# Multiclass classification (diagnosis) -multiclass_config = MetricConfig( - task_type="multiclass", - n_classes=5, - metrics=["auroc", "accuracy"], -) -metrics = build_metrics(multiclass_config, prefix="test") - -# Multilabel classification (complications) -multilabel_config = MetricConfig( - task_type="multilabel", - n_classes=10, - metrics=["auroc", "f1"], -) -metrics = build_metrics(multilabel_config, prefix="val") -``` - -### Custom Metric Sets - -```python -# Comprehensive evaluation -full_config = MetricConfig( - task_type="binary", - metrics=["auroc", "auprc", "accuracy", "f1"], -) - -# Minimal evaluation (faster) -minimal_config = MetricConfig( - task_type="binary", - metrics=["auroc"], -) - -# Task-specific focus -clinical_config = MetricConfig( - task_type="binary", - metrics=["auroc", "auprc"], # Focus on ranking quality -) -``` - -### Integration with Hydra Config - -```python -# configs/finetune.yaml -defaults: - - eval: default # Use configs/eval/default.yaml - - _self_ - -task: - task_name: mortality_24h - task_type: binary_classification - -eval: - metrics: - names: [auroc, auprc, f1] # Override defaults -``` - -## Testing - -See `tests/test_metrics.py` for comprehensive tests covering: -- Metric configuration validation -- Metric building for different task types -- Integration with torchmetrics -- Default metric selection -- Error handling - -Run tests: -```bash -uv run pytest tests/test_metrics.py -v -``` +## Integration Points -## Related Modules +- `FineTuneModule` builds metric collections from `config.eval.metrics.*` +- `scripts/eval/evaluate_fairness.py` uses `run_inference()` and `FairnessEvaluator` +- `scripts/export_results.py` consumes both task metrics and `fairness/*` keys -- `slices.training.finetune_module` - Uses eval metrics during training -- `slices.data.labels.base` - Defines `LabelConfig.primary_metric` -- `torchmetrics` - Underlying metric implementations +## Extending -## Future Work +When adding a new metric: -- [ ] Implement regression metrics (MSE, MAE, R²) -- [ ] Implement fairness analysis functions -- [ ] Add calibration metrics (ECE, Brier score) -- [ ] Add confidence intervals for metrics -- [ ] Add statistical significance testing -- [ ] Create visualization utilities (ROC curves, calibration plots) +1. Register it in `AVAILABLE_METRICS` and `DEFAULT_METRICS` in `metrics.py` +2. Implement it in `_build_metric()` +3. Add regression or classification coverage in `tests/test_metrics.py` +4. Update this README if the public surface changes diff --git a/src/slices/eval/__init__.py b/src/slices/eval/__init__.py index ec8cc6a..29293d8 100644 --- a/src/slices/eval/__init__.py +++ b/src/slices/eval/__init__.py @@ -7,21 +7,33 @@ - Imputation evaluation for SSL encoder quality assessment """ -from slices.eval.fairness_evaluator import FairnessEvaluator +from slices.eval.fairness_evaluator import FairnessEvaluator, flatten_fairness_report from slices.eval.imputation import ImputationEvaluator +from slices.eval.inference import run_inference from slices.eval.metrics import ( MetricConfig, build_metrics, get_default_metrics, ) -from slices.eval.statistical import bootstrap_ci, paired_bootstrap_test +from slices.eval.statistical import ( + bonferroni_correction, + bootstrap_ci, + cohens_d, + paired_bootstrap_test, + paired_wilcoxon_signed_rank, +) __all__ = [ "MetricConfig", "build_metrics", "get_default_metrics", "FairnessEvaluator", + "flatten_fairness_report", "ImputationEvaluator", + "bonferroni_correction", "bootstrap_ci", + "cohens_d", "paired_bootstrap_test", + "paired_wilcoxon_signed_rank", + "run_inference", ] diff --git a/src/slices/eval/fairness_evaluator.py b/src/slices/eval/fairness_evaluator.py index 09b5d9a..84ca485 100644 --- a/src/slices/eval/fairness_evaluator.py +++ b/src/slices/eval/fairness_evaluator.py @@ -3,17 +3,16 @@ Computes per-group AUROC, worst-group AUROC, demographic parity, and equalized odds across protected attributes (gender, age group, race). -This module builds on the lower-level fairness utilities in fairness.py -to provide a high-level evaluator that works with static patient data -and produces structured reports suitable for W&B logging. - -Protected attributes: -- gender (M/F) — available in all datasets -- age_group (18-44, 45-64, 65-79, 80+) — available in all datasets -- race — available in MIMIC + eICU only +Race handling follows the benchmark protocol: +- race/ethnicity is evaluated only on MIMIC-IV rows +- raw race strings are mapped into the canonical five-bin schema + (White, Black, Hispanic, Asian, Other) +- subgroup inclusion thresholds are enforced on unique patients, not stays """ import logging +import re +from numbers import Real from typing import Any, Dict, List, Optional, Tuple import polars as pl @@ -30,6 +29,24 @@ logger = logging.getLogger(__name__) +def flatten_fairness_report(report: dict[str, Any]) -> dict[str, Any]: + """Flatten a nested fairness report into W&B-safe ``fairness/*`` keys.""" + + flat: dict[str, Any] = {} + + def _flatten(prefix: str, value: Any) -> None: + if isinstance(value, Real) and not isinstance(value, bool): + flat[prefix] = value + elif isinstance(value, dict): + for sub_key, sub_value in value.items(): + _flatten(f"{prefix}/{sub_key}", sub_value) + + for attr, metrics in report.items(): + _flatten(f"fairness/{attr}", metrics) + + return flat + + class FairnessEvaluator: """Compute fairness metrics from model predictions + demographics. @@ -47,6 +64,7 @@ class FairnessEvaluator: AGE_BINS = [(18, 44), (45, 64), (65, 79), (80, float("inf"))] AGE_LABELS = ["18-44", "45-64", "65-79", "80+"] + RACE_LABELS = ["White", "Black", "Hispanic", "Asian", "Other"] def __init__( self, @@ -54,6 +72,7 @@ def __init__( protected_attributes: Optional[List[str]] = None, min_subgroup_size: int = 50, task_type: str = "binary", + dataset_name: Optional[str] = None, ) -> None: """Initialize FairnessEvaluator. @@ -62,14 +81,20 @@ def __init__( Must contain 'stay_id' column. May contain 'gender', 'age', 'race'. protected_attributes: List of attributes to evaluate. Defaults to ["gender", "age_group"]. - min_subgroup_size: Minimum samples for a subgroup to be included. + min_subgroup_size: Minimum patients for a subgroup to be included. task_type: Task type ("binary", "multiclass", "regression"). Determines which per-group metrics are computed. + dataset_name: Dataset identifier ("miiv", "eicu", "combined", ...). + Used for dataset-specific subgroup rules such as MIMIC-only race + analysis. When omitted, the evaluator falls back to columns + available in ``static_df`` for backwards compatibility. """ self.static_df = static_df self.protected_attributes = protected_attributes or ["gender", "age_group"] self.min_subgroup_size = min_subgroup_size self.task_type = task_type + self.dataset_name = dataset_name.lower() if dataset_name is not None else None + self._static_dict = {row["stay_id"]: row for row in self.static_df.to_dicts()} self._available_attributes = self._detect_available_attributes() def _detect_available_attributes(self) -> List[str]: @@ -83,10 +108,24 @@ def _detect_available_attributes(self) -> List[str]: available.append("gender") if "age" in self.static_df.columns: available.append("age_group") # Derived from age - if "race" in self.static_df.columns: + if "race" in self.static_df.columns and self._race_analysis_available(): available.append("race") return [a for a in self.protected_attributes if a in available] + def _race_analysis_available(self) -> bool: + """Return whether race fairness should be evaluated for this dataset.""" + if self.dataset_name == "eicu": + return False + if self.dataset_name == "combined": + if "source_dataset" not in self.static_df.columns: + return False + source_values = { + str(value).strip().lower() + for value in self.static_df["source_dataset"].drop_nulls().to_list() + } + return any("miiv" in value or "mimic" in value for value in source_values) + return True + def _bin_age(self, ages: torch.Tensor) -> torch.Tensor: """Bin continuous age into groups. @@ -96,16 +135,46 @@ def _bin_age(self, ages: torch.Tensor) -> torch.Tensor: Returns: Tensor of group indices (0-3). """ - groups = torch.zeros_like(ages, dtype=torch.long) + groups = torch.full_like(ages, fill_value=-1, dtype=torch.long) + finite_ages = torch.isfinite(ages) for i, (low, high) in enumerate(self.AGE_BINS): - groups[(ages >= low) & (ages <= high)] = i + groups[finite_ages & (ages >= low) & (ages <= high)] = i return groups + @staticmethod + def _canonicalize_gender(value: Any) -> Optional[str]: + """Normalize gender labels and exclude missing/unknown values.""" + if value is None: + return None + + text = str(value).strip() + if not text: + return None + + normalized = re.sub(r"\s+", " ", text).strip().lower() + if normalized in { + "unknown", + "unk", + "missing", + "not specified", + "not available", + "na", + "n/a", + "none", + "null", + }: + return None + if normalized in {"m", "male"}: + return "M" + if normalized in {"f", "female"}: + return "F" + return text + def _encode_attribute( self, stay_ids: List[int], attribute: str, - ) -> Tuple[torch.Tensor, Dict[int, str]]: + ) -> Tuple[torch.Tensor, Dict[int, str], List[Any]]: """Encode attribute as integer group IDs. Args: @@ -113,33 +182,38 @@ def _encode_attribute( attribute: Attribute name ("gender", "age_group", "race"). Returns: - Tuple of (group_ids tensor, mapping from int -> group name). + Tuple of: + - group_ids tensor aligned to ``stay_ids`` + - mapping from int -> group name + - patient_ids aligned to ``stay_ids`` for patient-threshold logic """ - # Build stay_id -> row lookup - static_dict = {row["stay_id"]: row for row in self.static_df.to_dicts()} + patient_ids = [] + rows = [] + for sid in stay_ids: + row = self._static_dict.get(sid, {}) + rows.append(row) + patient_ids.append(row.get("patient_id", sid)) if attribute == "age_group": ages = [] - for sid in stay_ids: - row = static_dict.get(sid, {}) + for row in rows: age = row.get("age") ages.append(float(age) if age is not None else -1.0) ages_tensor = torch.tensor(ages) group_ids = self._bin_age(ages_tensor) # Mark missing ages as -1 - group_ids[ages_tensor < 0] = -1 + group_ids[(ages_tensor < 0) | ~torch.isfinite(ages_tensor)] = -1 group_names = {i: label for i, label in enumerate(self.AGE_LABELS)} group_names[-1] = "unknown" - return group_ids, group_names + return group_ids, group_names, patient_ids elif attribute == "gender": # Map gender values to integers unique_vals = set() raw_vals = [] - for sid in stay_ids: - row = static_dict.get(sid, {}) - val = row.get("gender") + for row in rows: + val = self._canonicalize_gender(row.get("gender")) raw_vals.append(val) if val is not None: unique_vals.add(val) @@ -150,29 +224,69 @@ def _encode_attribute( ) group_names = {i: str(v) for v, i in val_to_id.items()} group_names[-1] = "unknown" - return group_ids, group_names + return group_ids, group_names, patient_ids elif attribute == "race": - unique_vals = set() - raw_vals = [] - for sid in stay_ids: - row = static_dict.get(sid, {}) - val = row.get("race") - raw_vals.append(val) - if val is not None: - unique_vals.add(val) + canonical_to_id = {label: i for i, label in enumerate(self.RACE_LABELS)} + canonical_vals = [] + for row in rows: + canonical_vals.append(self._canonicalize_race(row)) - val_to_id = {v: i for i, v in enumerate(sorted(unique_vals))} group_ids = torch.tensor( - [val_to_id.get(v, -1) if v is not None else -1 for v in raw_vals], dtype=torch.long + [ + canonical_to_id.get(value, -1) if value is not None else -1 + for value in canonical_vals + ], + dtype=torch.long, ) - group_names = {i: str(v) for v, i in val_to_id.items()} + group_names = {i: label for label, i in canonical_to_id.items()} group_names[-1] = "unknown" - return group_ids, group_names + return group_ids, group_names, patient_ids else: raise ValueError(f"Unknown attribute: {attribute}") + def _canonicalize_race(self, row: Dict[str, Any]) -> Optional[str]: + """Map raw race strings into the planned five-bin schema.""" + if not self._row_is_miiv_for_race(row): + return None + + raw_value = row.get("race") + if raw_value is None: + return None + + text = str(raw_value).strip() + if not text: + return None + + normalized = re.sub(r"[/_\\-]+", " ", text.upper()) + normalized = re.sub(r"\s+", " ", normalized).strip() + + missing_markers = ("UNKNOWN", "DECLIN", "UNABLE", "PATIENT REFUSED", "NOT SPECIFIED") + if any(marker in normalized for marker in missing_markers): + return None + if "HISPANIC" in normalized or "LATINO" in normalized: + return "Hispanic" + if "BLACK" in normalized or "AFRICAN" in normalized: + return "Black" + if "ASIAN" in normalized: + return "Asian" + if "WHITE" in normalized: + return "White" + return "Other" + + def _row_is_miiv_for_race(self, row: Dict[str, Any]) -> bool: + """Return whether a row should participate in race fairness analysis.""" + if self.dataset_name == "eicu": + return False + source_dataset = row.get("source_dataset") + if source_dataset is not None: + source_value = str(source_dataset).strip().lower() + return "miiv" in source_value or "mimic" in source_value + if self.dataset_name == "combined": + return False + return True + def evaluate( self, predictions: torch.Tensor, @@ -193,38 +307,70 @@ def evaluate( Returns: Structured dict with per-attribute results, loggable to W&B. """ + if predictions.ndim > 1 and predictions.shape[-1] == 1: + predictions = predictions.squeeze(-1) + if labels.ndim > 1 and labels.shape[-1] == 1: + labels = labels.squeeze(-1) + report: Dict[str, Any] = {} for attr in self._available_attributes: - group_ids, group_names = self._encode_attribute(stay_ids, attr) + group_ids, group_names, patient_ids = self._encode_attribute(stay_ids, attr) # Get unique valid groups (exclude -1 = unknown) unique_groups = [g for g in group_ids.unique().tolist() if g >= 0] - # Filter groups below min_subgroup_size + # Filter groups below the benchmark patient-count threshold. valid_groups = [] + group_sizes = {} + group_sample_sizes = {} for g in unique_groups: group_mask = group_ids == g - if group_mask.sum().item() >= self.min_subgroup_size: + patient_count = len( + { + patient_ids[i] + for i, is_member in enumerate(group_mask.tolist()) + if is_member and patient_ids[i] is not None + } + ) + sample_count = int(group_mask.sum().item()) + group_sizes[group_names[g]] = patient_count + group_sample_sizes[group_names[g]] = sample_count + if patient_count >= self.min_subgroup_size: valid_groups.append(g) if len(valid_groups) < 2: logger.warning( - "Attribute '%s': fewer than 2 groups with >= %d samples, skipping", + "Attribute '%s': fewer than 2 groups with >= %d patients, skipping", attr, self.min_subgroup_size, ) continue - group_sizes = {group_names[g]: int((group_ids == g).sum().item()) for g in valid_groups} + valid_group_sizes = {group_names[g]: group_sizes[group_names[g]] for g in valid_groups} + valid_group_sample_sizes = { + group_names[g]: group_sample_sizes[group_names[g]] for g in valid_groups + } if self.task_type == "regression": report[attr] = self._evaluate_regression( - predictions, labels, group_ids, group_names, valid_groups, group_sizes + predictions, + labels, + group_ids, + group_names, + valid_groups, + valid_group_sizes, + valid_group_sample_sizes, ) else: report[attr] = self._evaluate_binary( - predictions, labels, group_ids, group_names, valid_groups, group_sizes + predictions, + labels, + group_ids, + group_names, + valid_groups, + valid_group_sizes, + valid_group_sample_sizes, ) return report @@ -237,41 +383,56 @@ def _evaluate_binary( group_names: Dict[int, str], valid_groups: List[int], group_sizes: Dict[str, int], + group_sample_sizes: Dict[str, int], ) -> Dict[str, Any]: """Compute binary classification fairness metrics.""" per_group_auroc: Dict[str, float] = {} per_group_auprc: Dict[str, float] = {} auroc_values = [] auprc_values = [] + metric_valid_groups = [] + single_class_groups = [] for g in valid_groups: group_mask = group_ids == g g_preds = predictions[group_mask] g_labels = labels[group_mask].long() + group_name = group_names[g] if g_labels.unique().numel() < 2: - per_group_auroc[group_names[g]] = float("nan") - per_group_auprc[group_names[g]] = float("nan") + per_group_auroc[group_name] = float("nan") + per_group_auprc[group_name] = float("nan") + single_class_groups.append(group_name) continue auroc_metric = AUROC(task="binary") auroc_val = auroc_metric(g_preds, g_labels).item() - per_group_auroc[group_names[g]] = auroc_val + per_group_auroc[group_name] = auroc_val auroc_values.append(auroc_val) auprc_metric = BinaryAveragePrecision() auprc_val = auprc_metric(g_preds, g_labels).item() - per_group_auprc[group_names[g]] = auprc_val + per_group_auprc[group_name] = auprc_val auprc_values.append(auprc_val) - - worst_group_auroc = min(auroc_values) if auroc_values else float("nan") - worst_group_auprc = min(auprc_values) if auprc_values else float("nan") - auroc_gap = ( - (max(auroc_values) - min(auroc_values)) if len(auroc_values) >= 2 else float("nan") - ) - auprc_gap = ( - (max(auprc_values) - min(auprc_values)) if len(auprc_values) >= 2 else float("nan") - ) + metric_valid_groups.append(group_name) + + metrics_comparable = len(metric_valid_groups) >= 2 + if metrics_comparable: + worst_group_auroc = min(auroc_values) + worst_group_auprc = min(auprc_values) + auroc_gap = max(auroc_values) - min(auroc_values) + auprc_gap = max(auprc_values) - min(auprc_values) + else: + logger.warning( + "Binary fairness discrimination metrics are not comparable: " + "%d/%d size-valid groups have both outcome classes.", + len(metric_valid_groups), + len(valid_groups), + ) + worst_group_auroc = float("nan") + worst_group_auprc = float("nan") + auroc_gap = float("nan") + auprc_gap = float("nan") valid_mask = torch.zeros_like(group_ids, dtype=torch.bool) for g in valid_groups: @@ -294,7 +455,10 @@ def _evaluate_binary( "equalized_odds_diff": eo_diff, "disparate_impact_ratio": di_ratio, "n_valid_groups": len(valid_groups), + "n_metric_valid_groups": len(metric_valid_groups), + "n_single_class_groups": len(single_class_groups), "group_sizes": group_sizes, + "group_sample_sizes": group_sample_sizes, } def _evaluate_regression( @@ -305,17 +469,19 @@ def _evaluate_regression( group_names: Dict[int, str], valid_groups: List[int], group_sizes: Dict[str, int], + group_sample_sizes: Dict[str, int], ) -> Dict[str, Any]: """Compute regression fairness metrics (per-group MSE, MAE, R2).""" per_group_mse: Dict[str, float] = {} per_group_mae: Dict[str, float] = {} per_group_r2: Dict[str, float] = {} mse_values = [] + mae_values = [] for g in valid_groups: group_mask = group_ids == g - g_preds = predictions[group_mask].float() - g_labels = labels[group_mask].float() + g_preds = predictions[group_mask].float().reshape(-1) + g_labels = labels[group_mask].float().reshape(-1) residuals = g_preds - g_labels mse = (residuals**2).mean().item() @@ -331,16 +497,24 @@ def _evaluate_regression( per_group_mae[name] = mae per_group_r2[name] = r2 mse_values.append(mse) + mae_values.append(mae) worst_group_mse = max(mse_values) if mse_values else float("nan") + worst_group_mae = max(mae_values) if mae_values else float("nan") + mse_gap = (max(mse_values) - min(mse_values)) if len(mse_values) >= 2 else float("nan") + mae_gap = (max(mae_values) - min(mae_values)) if len(mae_values) >= 2 else float("nan") return { "per_group_mse": per_group_mse, "per_group_mae": per_group_mae, "per_group_r2": per_group_r2, "worst_group_mse": worst_group_mse, + "worst_group_mae": worst_group_mae, + "mse_gap": mse_gap, + "mae_gap": mae_gap, "n_valid_groups": len(valid_groups), "group_sizes": group_sizes, + "group_sample_sizes": group_sample_sizes, } def print_report(self, report: Dict[str, Any]) -> None: @@ -361,19 +535,33 @@ def print_report(self, report: Dict[str, Any]) -> None: # Binary classification report print(" Per-group AUROC:") for group, auroc in metrics["per_group_auroc"].items(): - size = metrics["group_sizes"].get(group, "?") + patient_size = metrics["group_sizes"].get(group, "?") + sample_size = metrics.get("group_sample_sizes", {}).get(group, "?") if isinstance(auroc, float) and auroc != auroc: # NaN check - print(f" {group} (n={size}): N/A (single class)") + print( + f" {group} (patients={patient_size}, stays={sample_size}): " + "N/A (single class)" + ) else: - print(f" {group} (n={size}): {auroc:.4f}") + print( + f" {group} (patients={patient_size}, stays={sample_size}): " + f"{auroc:.4f}" + ) print(" Per-group AUPRC:") for group, auprc in metrics.get("per_group_auprc", {}).items(): - size = metrics["group_sizes"].get(group, "?") + patient_size = metrics["group_sizes"].get(group, "?") + sample_size = metrics.get("group_sample_sizes", {}).get(group, "?") if isinstance(auprc, float) and auprc != auprc: - print(f" {group} (n={size}): N/A (single class)") + print( + f" {group} (patients={patient_size}, stays={sample_size}): " + "N/A (single class)" + ) else: - print(f" {group} (n={size}): {auprc:.4f}") + print( + f" {group} (patients={patient_size}, stays={sample_size}): " + f"{auprc:.4f}" + ) wg = metrics["worst_group_auroc"] if isinstance(wg, float) and wg == wg: # Not NaN @@ -409,11 +597,15 @@ def print_report(self, report: Dict[str, Any]) -> None: # Regression report print(" Per-group MSE / MAE / R2:") for group in metrics["per_group_mse"]: - size = metrics["group_sizes"].get(group, "?") + patient_size = metrics["group_sizes"].get(group, "?") + sample_size = metrics.get("group_sample_sizes", {}).get(group, "?") mse = metrics["per_group_mse"][group] mae = metrics["per_group_mae"][group] r2 = metrics["per_group_r2"][group] - print(f" {group} (n={size}): MSE={mse:.4f} MAE={mae:.4f} R2={r2:.4f}") + print( + f" {group} (patients={patient_size}, stays={sample_size}): " + f"MSE={mse:.4f} MAE={mae:.4f} R2={r2:.4f}" + ) wg = metrics["worst_group_mse"] if isinstance(wg, float) and wg == wg: diff --git a/src/slices/eval/fairness_metadata.py b/src/slices/eval/fairness_metadata.py new file mode 100644 index 0000000..eb45c11 --- /dev/null +++ b/src/slices/eval/fairness_metadata.py @@ -0,0 +1,95 @@ +"""Shared metadata contract for post-run fairness summaries.""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +from typing import Any + +FAIRNESS_SUMMARY_SCHEMA_VERSION = "2026-04-23.v2" +FAIRNESS_SCRIPT_VERSION = "evaluate_fairness.py:2026-04-23.v2" +FAIRNESS_DEFAULT_PROTECTED_ATTRIBUTES = ["gender", "age_group", "race"] +FAIRNESS_DEFAULT_MIN_SUBGROUP_SIZE = 50 + +EVAL_ARTIFACT_PATH_KEY = "_eval_artifact_path" +EVAL_ARTIFACT_SHA256_KEY = "_eval_artifact_sha256" +FAIRNESS_SCHEMA_VERSION_KEY = "_fairness_summary_schema_version" +FAIRNESS_SCRIPT_VERSION_KEY = "_fairness_script_version" +FAIRNESS_ARTIFACT_PATH_KEY = "_fairness_artifact_path" +FAIRNESS_ARTIFACT_SHA256_KEY = "_fairness_artifact_sha256" +FAIRNESS_ARTIFACT_SOURCE_KEY = "_fairness_artifact_source" +FAIRNESS_CHECKPOINT_SOURCE_KEY = "_fairness_checkpoint_source" +FAIRNESS_PROTECTED_ATTRIBUTES_KEY = "_fairness_protected_attributes" +FAIRNESS_MIN_SUBGROUP_SIZE_KEY = "_fairness_min_subgroup_size" + +FAIRNESS_METADATA_COLUMNS = [ + FAIRNESS_SCHEMA_VERSION_KEY, + FAIRNESS_SCRIPT_VERSION_KEY, + FAIRNESS_ARTIFACT_PATH_KEY, + FAIRNESS_ARTIFACT_SHA256_KEY, + FAIRNESS_ARTIFACT_SOURCE_KEY, + FAIRNESS_CHECKPOINT_SOURCE_KEY, + FAIRNESS_PROTECTED_ATTRIBUTES_KEY, + FAIRNESS_MIN_SUBGROUP_SIZE_KEY, +] + +FAIRNESS_CLEAR_PREFIXES = ("fairness/", "_fairness_") + + +def normalize_protected_attributes(protected_attributes: list[str]) -> list[str]: + """Return de-duplicated protected attributes while preserving order.""" + return list(dict.fromkeys(str(attr) for attr in protected_attributes)) + + +def encode_protected_attributes(protected_attributes: list[str]) -> str: + """Encode protected attributes into a stable W&B summary scalar.""" + return json.dumps(normalize_protected_attributes(protected_attributes)) + + +def decode_protected_attributes(value: Any) -> list[str] | None: + """Decode protected-attribute metadata from W&B/export representations.""" + if value is None: + return None + if isinstance(value, str): + text = value.strip() + if not text: + return None + try: + parsed = json.loads(text) + except json.JSONDecodeError: + parsed = [part.strip() for part in text.split(",") if part.strip()] + elif isinstance(value, (list, tuple, set)): + parsed = list(value) + else: + return None + return normalize_protected_attributes([str(attr) for attr in parsed]) + + +def canonical_artifact_id(path_value: Any) -> str: + """Return a path identifier stable across output-root rebases.""" + if path_value is None: + return "" + + text = str(path_value).strip().replace("\\", "/") + if not text: + return "" + + marker = "/outputs/" + if marker in text: + return "outputs/" + text.split(marker, 1)[1].strip("/") + if text.startswith("outputs/"): + return text.strip("/") + index = text.find("outputs/") + if index >= 0: + return text[index:].strip("/") + return text.strip("/") + + +def file_sha256(path: str | Path) -> str: + """Return the SHA256 digest for a local artifact file.""" + digest = hashlib.sha256() + with Path(path).open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() diff --git a/src/slices/eval/imputation.py b/src/slices/eval/imputation.py index 1cdf756..2534873 100644 --- a/src/slices/eval/imputation.py +++ b/src/slices/eval/imputation.py @@ -10,8 +10,8 @@ - feature_block: Mask entire features for the full window - temporal_block: Mask contiguous hour blocks across all features -For MAE models: extracts encoder weights, uses a linear decoder for probing. -For non-MAE models: train lightweight linear decoder (d_model -> d_input). +Checkpoint loaders create a lightweight probe decoder (d_model -> d_input). +This probe should be trained on the training split before evaluation. Metrics: NRMSE per feature, MAE overall. """ @@ -23,6 +23,8 @@ import torch.nn as nn from torch.utils.data import DataLoader +from slices.models.encoders import ObservationTransformerEncoder, SMARTEncoder + logger = logging.getLogger(__name__) @@ -86,6 +88,30 @@ def forward( return output +class _SmartEncoderTimestepAdapter(nn.Module): + """Wrap SMART pooling=none output into timestep-level representations. + + SMART returns ``(B, V, T, d_model)`` for SSL. The generic imputation probe + expects ``(B, T, d_model)`` so one decoder can reconstruct all features at + each timestep. + """ + + def __init__(self, encoder: SMARTEncoder) -> None: + super().__init__() + self.encoder = encoder + + def get_output_dim(self) -> int: + return self.encoder.get_output_dim() + + def forward( + self, x: torch.Tensor, mask: Optional[torch.Tensor] = None, **kwargs + ) -> torch.Tensor: + encoded = self.encoder(x, mask=mask, **kwargs) + if encoded.dim() == 4: + return encoded.mean(dim=1) + return encoded + + class ImputationEvaluator: """Evaluate reconstruction quality of SSL embeddings. @@ -150,10 +176,10 @@ def from_mae_checkpoint( device: str = "cpu", feature_names: Optional[List[str]] = None, ) -> "ImputationEvaluator": - """Load MAE encoder for reconstruction evaluation. + """Load a MAE encoder and initialize a probe decoder for evaluation. Extracts the encoder from a full MAE pretraining checkpoint and - creates a linear decoder for probing reconstruction quality. + creates a lightweight probe decoder for reconstruction quality. The MAE's native decoder operates in observation-token space which is incompatible with the timestep-level masking used by evaluate(). @@ -166,7 +192,7 @@ def from_mae_checkpoint( feature_names: Optional feature names for per-feature reporting. Returns: - ImputationEvaluator with MAE encoder (adapted) and linear decoder. + ImputationEvaluator with MAE encoder (adapted) and probe decoder. """ from slices.models.encoders import build_encoder @@ -189,8 +215,12 @@ def from_mae_checkpoint( seq_length = encoder.config.max_seq_length # Wrap observation-level encoder for timestep-level output - if hasattr(encoder, "tokenize"): + # (ObservationTransformerEncoder produces per-observation tokens that + # need scatter-back; obs_aware TransformerEncoder already outputs per-timestep) + if isinstance(encoder, ObservationTransformerEncoder): encoder = _ObsEncoderTimestepAdapter(encoder, seq_length) + elif isinstance(encoder, SMARTEncoder): + encoder = _SmartEncoderTimestepAdapter(encoder) encoder = encoder.to(device).eval() @@ -239,9 +269,13 @@ def from_encoder_checkpoint( encoder.load_state_dict(checkpoint["encoder_state_dict"]) # Wrap observation-level encoder for timestep-level output + # (ObservationTransformerEncoder produces per-observation tokens; + # obs_aware TransformerEncoder already outputs per-timestep) seq_length = encoder_config.get("max_seq_length", 168) - if hasattr(encoder, "tokenize"): + if isinstance(encoder, ObservationTransformerEncoder): encoder = _ObsEncoderTimestepAdapter(encoder, seq_length) + elif isinstance(encoder, SMARTEncoder): + encoder = _SmartEncoderTimestepAdapter(encoder) else: raise ValueError( "Checkpoint does not contain encoder_config. " @@ -322,10 +356,11 @@ def train_decoder( max_epochs: int = 10, lr: float = 1e-3, ) -> Dict[str, Any]: - """Train lightweight decoder for non-MAE models. + """Train a lightweight reconstruction probe on frozen encoder features. - Freezes encoder and trains only the decoder to reconstruct - observed values from encoder representations. + Freezes the encoder and trains only the decoder to reconstruct + observed values from encoder representations. This is the intended + evaluation path for both MAE checkpoints and saved encoder checkpoints. Args: dataloader: DataLoader providing batches with 'timeseries' and 'mask'. diff --git a/src/slices/eval/inference.py b/src/slices/eval/inference.py new file mode 100644 index 0000000..0adb290 --- /dev/null +++ b/src/slices/eval/inference.py @@ -0,0 +1,152 @@ +"""Shared inference utilities for model evaluation. + +Provides a reusable function for collecting predictions, labels, and stay IDs +from a test dataloader — used by both in-training fairness evaluation and +the standalone post-run fairness script. +""" + +from pathlib import Path +from typing import List, Tuple, Union + +import numpy as np +import torch +import torch.nn as nn +from torch.utils.data import DataLoader + + +def run_inference( + model: nn.Module, + dataloader: DataLoader, + device: Union[torch.device, str] = "cpu", +) -> Tuple[torch.Tensor, torch.Tensor, List[int]]: + """Run inference on a dataloader, collecting predictions, labels, and stay IDs. + + Args: + model: Model that accepts (timeseries, mask) and returns a dict + with a 'probs' key. Must already be in eval mode. + dataloader: DataLoader yielding batches with 'timeseries', 'mask', + 'label', and 'stay_id' keys. + device: Device to run inference on. + + Returns: + Tuple of (predictions, labels, stay_ids) where: + - predictions: (N,) tensor of probabilities (positive class for binary) + - labels: (N,) tensor of ground truth labels + - stay_ids: list of N stay ID integers + """ + all_preds: list[torch.Tensor] = [] + all_labels: list[torch.Tensor] = [] + all_stay_ids: list[int] = [] + + model.eval() + for batch in dataloader: + with torch.no_grad(): + outputs = model( + batch["timeseries"].to(device), + batch["mask"].to(device), + ) + probs = outputs["probs"] + # For binary classification with 2-class output, take positive class + if probs.dim() > 1 and probs.shape[1] == 2: + all_preds.append(probs[:, 1].cpu()) + else: + all_preds.append(probs.reshape(-1).cpu()) + all_labels.append(batch["label"].reshape(-1).cpu()) + all_stay_ids.extend( + batch["stay_id"].tolist() + if isinstance(batch["stay_id"], torch.Tensor) + else batch["stay_id"] + ) + + predictions = torch.cat(all_preds, dim=0) + labels = torch.cat(all_labels, dim=0) + return predictions, labels, all_stay_ids + + +def extract_tabular_features(dataset, indices: list[int]) -> tuple[np.ndarray, np.ndarray]: + """Extract the tabular summary features used by the XGBoost baseline.""" + idx_tensor = torch.tensor(indices, dtype=torch.long) + ts = dataset._timeseries_tensor[idx_tensor] # (N, T, D) + mask = dataset._mask_tensor[idx_tensor] # (N, T, D) bool + + n_samples, n_timesteps, n_features = ts.shape + mask_float = mask.float() + + obs_count = mask_float.sum(dim=1) + obs_frac = obs_count / n_timesteps + + masked_ts = ts * mask_float + feat_sum = masked_ts.sum(dim=1) + safe_count = obs_count.clamp(min=1) + feat_mean = feat_sum / safe_count + + diff_sq = ((ts - feat_mean.unsqueeze(1)) ** 2) * mask_float + feat_var = diff_sq.sum(dim=1) / safe_count + feat_std = torch.sqrt(feat_var) + + zeros = ts.new_zeros(n_samples, n_features) + ts_for_min = ts.clone() + ts_for_min[~mask] = float("inf") + raw_min = ts_for_min.min(dim=1).values + feat_min = torch.where(obs_count > 0, raw_min, zeros) + + ts_for_max = ts.clone() + ts_for_max[~mask] = float("-inf") + raw_max = ts_for_max.max(dim=1).values + feat_max = torch.where(obs_count > 0, raw_max, zeros) + + first_idx = mask.float().argmax(dim=1) + last_idx = n_timesteps - 1 - mask.flip(dims=[1]).float().argmax(dim=1) + + feat_first = ts.gather(1, first_idx.unsqueeze(1)).squeeze(1) + feat_last = ts.gather(1, last_idx.unsqueeze(1)).squeeze(1) + + no_obs = obs_count == 0 + feat_mean = torch.nan_to_num(feat_mean, nan=0.0) + feat_first[no_obs] = 0.0 + feat_last[no_obs] = 0.0 + + features = torch.stack( + [feat_mean, feat_std, feat_min, feat_max, feat_first, feat_last, obs_count, obs_frac], + dim=-1, + ) + x = features.reshape(n_samples, n_features * 8).numpy() + + if dataset._labels_tensor is not None: + y = dataset._labels_tensor[idx_tensor].numpy() + else: + y = np.zeros(n_samples, dtype=np.float32) + + return x, y + + +def load_xgboost_model(task_type: str, model_path: str | Path): + """Load a saved XGBoost baseline model from disk.""" + from xgboost import XGBClassifier, XGBRegressor + + model = XGBRegressor() if task_type == "regression" else XGBClassifier() + model.load_model(str(model_path)) + return model + + +def run_xgboost_inference( + model_path: str | Path, + task_type: str, + dataset, + indices: list[int], +) -> Tuple[torch.Tensor, torch.Tensor, List[int]]: + """Run inference with a saved XGBoost baseline over a dataset subset.""" + model = load_xgboost_model(task_type, model_path) + features, labels = extract_tabular_features(dataset, indices) + + if task_type == "regression": + predictions = model.predict(features) + else: + predictions = model.predict_proba(features)[:, 1] + + stay_ids = [dataset.stay_ids[i] for i in indices] + return ( + torch.tensor(predictions, dtype=torch.float32), + torch.tensor(labels, dtype=torch.float32), + stay_ids, + ) diff --git a/src/slices/eval/metrics.py b/src/slices/eval/metrics.py index d55e6cf..48b4f76 100644 --- a/src/slices/eval/metrics.py +++ b/src/slices/eval/metrics.py @@ -52,9 +52,19 @@ "regression": ["mse", "mae", "rmse", "r2"], } -# Default metrics per task type (minimal set for clinical tasks) +# Default metrics per task type for benchmark exports. DEFAULT_METRICS = { - "binary": ["auroc", "auprc", "brier_score", "ece"], + "binary": [ + "auroc", + "auprc", + "accuracy", + "f1", + "precision", + "recall", + "specificity", + "brier_score", + "ece", + ], "multiclass": ["auroc", "accuracy"], "multilabel": ["auroc"], "regression": ["mse", "mae", "r2"], diff --git a/src/slices/eval/statistical.py b/src/slices/eval/statistical.py index b4056c2..1be84e8 100644 --- a/src/slices/eval/statistical.py +++ b/src/slices/eval/statistical.py @@ -1,20 +1,15 @@ -"""Bootstrap confidence intervals and statistical tests for metric comparison. - -Provides non-parametric bootstrap CIs for any torchmetrics-compatible metric, -and paired bootstrap tests for comparing two models on the same test set. - -Example: - >>> from torchmetrics import AUROC - >>> ci = bootstrap_ci(AUROC(task="binary"), preds, labels) - >>> print(f"AUROC: {ci['point']:.3f} ({ci['ci_lower']:.3f}-{ci['ci_upper']:.3f})") - >>> - >>> p = paired_bootstrap_test( - ... AUROC(task="binary"), preds_a, preds_b, labels - ... ) - >>> print(f"p-value: {p['p_value']:.4f}") +"""Bootstrap confidence intervals and paired statistical tests. + +Provides: +- non-parametric bootstrap CIs for torchmetrics-compatible metrics +- paired bootstrap tests on shared test sets +- paired Wilcoxon signed-rank tests for per-seed / per-task comparisons +- Bonferroni correction for multiple comparisons +- paired Cohen's d effect sizes """ -from typing import Any, Callable, Dict, Optional, Union +import math +from typing import Any, Callable, Dict, Iterable, Optional, Sequence, Union import torch from torchmetrics import Metric @@ -50,7 +45,12 @@ def bootstrap_ci( generator.manual_seed(seed) n = len(targets) - point = _compute_metric(metric_fn, preds, targets) + requires_class_diversity = _metric_requires_class_diversity(metric_fn) + point = ( + float("nan") + if requires_class_diversity and not _has_class_diversity(targets) + else _compute_metric(metric_fn, preds, targets) + ) scores = [] for _ in range(n_bootstraps): @@ -58,8 +58,10 @@ def bootstrap_ci( boot_preds = preds[indices] boot_targets = targets[indices] + if requires_class_diversity and not _has_class_diversity(boot_targets): + continue score = _compute_metric(metric_fn, boot_preds, boot_targets) - if score == score: # skip NaN + if math.isfinite(score): scores.append(score) if not scores: @@ -120,8 +122,13 @@ def paired_bootstrap_test( generator.manual_seed(seed) n = len(targets) - score_a = _compute_metric(metric_fn, preds_a, targets) - score_b = _compute_metric(metric_fn, preds_b, targets) + requires_class_diversity = _metric_requires_class_diversity(metric_fn) + if requires_class_diversity and not _has_class_diversity(targets): + score_a = float("nan") + score_b = float("nan") + else: + score_a = _compute_metric(metric_fn, preds_a, targets) + score_b = _compute_metric(metric_fn, preds_b, targets) # Observed delta: positive means A is better when higher_is_better=True observed_delta = score_a - score_b @@ -135,10 +142,13 @@ def paired_bootstrap_test( for _ in range(n_bootstraps): indices = torch.randint(0, n, (n,), generator=generator) - boot_a = _compute_metric(metric_fn, preds_a[indices], targets[indices]) - boot_b = _compute_metric(metric_fn, preds_b[indices], targets[indices]) + boot_targets = targets[indices] + if requires_class_diversity and not _has_class_diversity(boot_targets): + continue + boot_a = _compute_metric(metric_fn, preds_a[indices], boot_targets) + boot_b = _compute_metric(metric_fn, preds_b[indices], boot_targets) - if boot_a != boot_a or boot_b != boot_b: # skip NaN + if not (math.isfinite(boot_a) and math.isfinite(boot_b)): continue valid += 1 @@ -154,7 +164,7 @@ def paired_bootstrap_test( if boot_delta <= 2 * observed_delta: count_extreme += 1 - p_value = count_extreme / max(valid, 1) + p_value = count_extreme / valid if valid else float("nan") return { "score_a": score_a, @@ -165,6 +175,137 @@ def paired_bootstrap_test( } +def paired_wilcoxon_signed_rank( + values_a: Sequence[float], + values_b: Sequence[float], + correction: bool = True, +) -> Dict[str, float]: + """Paired Wilcoxon signed-rank test. + + Delegates to SciPy's implementation so small-sample exact handling and + tied absolute differences match the statistical reference implementation. + + Args: + values_a: First paired sample. + values_b: Second paired sample. + correction: Whether to apply a 0.5 continuity correction. + + Returns: + Dictionary with Wilcoxon statistic, z-score, p-value, and pair counts. + """ + pairs = _finite_pairs(values_a, values_b) + n_pairs = len(pairs) + if n_pairs == 0: + return { + "statistic": 0.0, + "z_score": 0.0, + "p_value": float("nan"), + "n_pairs": 0.0, + "n_nonzero_pairs": 0.0, + } + + diffs = [a - b for a, b in pairs] + nonzero_diffs = [diff for diff in diffs if diff != 0.0] + n_nonzero = len(nonzero_diffs) + + if n_nonzero == 0: + return { + "statistic": 0.0, + "z_score": 0.0, + "p_value": 1.0, + "n_pairs": float(n_pairs), + "n_nonzero_pairs": 0.0, + } + + try: + from scipy.stats import wilcoxon + except ImportError as exc: + raise RuntimeError( + "SciPy is required for Wilcoxon signed-rank tests. " + "Install the project dependencies with `uv sync`." + ) from exc + + # Exact enumeration is cheap for the benchmark's usual 5-seed comparisons + # and avoids anti-conservative normal approximations when all ranks tie. + method = "exact" if n_nonzero <= 50 else "auto" + result = wilcoxon( + nonzero_diffs, + zero_method="wilcox", + correction=correction, + alternative="two-sided", + method=method, + ) + z_score = getattr(result, "zstatistic", float("nan")) + + return { + "statistic": float(result.statistic), + "z_score": float(z_score), + "p_value": float(result.pvalue), + "n_pairs": float(n_pairs), + "n_nonzero_pairs": float(n_nonzero), + } + + +def bonferroni_correction(p_values: Sequence[float]) -> list[float]: + """Apply Bonferroni correction, preserving NaNs.""" + finite_count = sum(0 if _is_nan(p) else 1 for p in p_values) + if finite_count == 0: + return [float("nan") for _ in p_values] + + corrected = [] + for p in p_values: + if _is_nan(p): + corrected.append(float("nan")) + else: + corrected.append(min(float(p) * finite_count, 1.0)) + return corrected + + +def cohens_d( + values_a: Sequence[float], + values_b: Sequence[float], + paired: bool = False, +) -> float: + """Compute Cohen's d effect size. + + Args: + values_a: First sample. + values_b: Second sample. + paired: When True, compute the paired effect size using the standard + deviation of paired differences. + """ + pairs = _finite_pairs(values_a, values_b) + if not pairs: + return float("nan") + + sample_a = [a for a, _ in pairs] + sample_b = [b for _, b in pairs] + + if paired: + diffs = [a - b for a, b in pairs] + return _cohens_d_from_differences(diffs) + + n_a = len(sample_a) + n_b = len(sample_b) + if n_a < 2 or n_b < 2: + return float("nan") + + mean_a = sum(sample_a) / n_a + mean_b = sum(sample_b) / n_b + var_a = _sample_variance(sample_a) + var_b = _sample_variance(sample_b) + + pooled_var = ((n_a - 1) * var_a + (n_b - 1) * var_b) / max(n_a + n_b - 2, 1) + pooled_std = math.sqrt(max(pooled_var, 0.0)) + if pooled_std == 0.0: + diff = mean_a - mean_b + if diff == 0.0: + return 0.0 + return math.copysign(float("inf"), diff) + + return (mean_a - mean_b) / pooled_std + + def _compute_metric( metric_fn: Union[Metric, Callable[[torch.Tensor, torch.Tensor], torch.Tensor]], preds: torch.Tensor, @@ -177,3 +318,57 @@ def _compute_metric( return metric_fn.compute().item() else: return metric_fn(preds, targets).item() + + +def _metric_requires_class_diversity( + metric_fn: Union[Metric, Callable[[torch.Tensor, torch.Tensor], torch.Tensor]], +) -> bool: + """Return whether a metric is undefined for single-class target samples.""" + if not isinstance(metric_fn, Metric): + return False + metric_name = metric_fn.__class__.__name__.lower() + return "auroc" in metric_name or "averageprecision" in metric_name + + +def _has_class_diversity(targets: torch.Tensor) -> bool: + flat = targets.detach().reshape(-1) + if flat.numel() < 2: + return False + return torch.unique(flat).numel() >= 2 + + +def _finite_pairs( + values_a: Iterable[float], + values_b: Iterable[float], +) -> list[tuple[float, float]]: + pairs = [] + for a, b in zip(values_a, values_b, strict=True): + a_val = float(a) + b_val = float(b) + if math.isfinite(a_val) and math.isfinite(b_val): + pairs.append((a_val, b_val)) + return pairs + + +def _sample_variance(values: Sequence[float]) -> float: + if len(values) < 2: + return float("nan") + mean = sum(values) / len(values) + return sum((value - mean) ** 2 for value in values) / (len(values) - 1) + + +def _cohens_d_from_differences(differences: Sequence[float]) -> float: + if len(differences) < 2: + return float("nan") + + mean_diff = sum(differences) / len(differences) + std_diff = math.sqrt(max(_sample_variance(differences), 0.0)) + if std_diff == 0.0: + if mean_diff == 0.0: + return 0.0 + return math.copysign(float("inf"), mean_diff) + return mean_diff / std_diff + + +def _is_nan(value: float) -> bool: + return isinstance(value, float) and math.isnan(value) diff --git a/src/slices/models/common.py b/src/slices/models/common.py index 0c97640..37652e5 100644 --- a/src/slices/models/common.py +++ b/src/slices/models/common.py @@ -66,7 +66,8 @@ def apply_pooling( lengths = padding_mask.sum(dim=1) batch_idx = torch.arange(x.size(0), device=x.device) last_idx = (lengths - 1).clamp(min=0) - return x[batch_idx, last_idx, :] + pooled = x[batch_idx, last_idx, :] + return torch.where(lengths.unsqueeze(-1) > 0, pooled, torch.zeros_like(pooled)) else: return x[:, -1, :] @@ -84,7 +85,12 @@ def apply_pooling( if padding_mask is not None: mask_expanded = padding_mask.unsqueeze(-1) x_masked = x.masked_fill(~mask_expanded, float("-inf")) - return x_masked.max(dim=1)[0] + pooled = x_masked.max(dim=1)[0] + all_empty = ~padding_mask.any(dim=1) + if all_empty.any(): + pooled = pooled.clone() + pooled[all_empty] = 0 + return pooled else: return x.max(dim=1)[0] @@ -92,6 +98,24 @@ def apply_pooling( raise ValueError(f"Unknown pooling strategy: {pooling}") +def build_sinusoidal_pe(max_len: int, d_model: int) -> torch.Tensor: + """Build sinusoidal positional encoding tensor. + + Args: + max_len: Maximum sequence length. + d_model: Model dimension. + + Returns: + Tensor of shape (max_len, d_model) with sinusoidal encodings. + """ + position = torch.arange(max_len).unsqueeze(1).float() + div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) + pe = torch.zeros(max_len, d_model) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term[: pe[:, 1::2].shape[1]]) + return pe + + class PositionalEncoding(nn.Module): """Sinusoidal positional encoding. @@ -109,15 +133,7 @@ def __init__(self, d_model: int, max_seq_length: int = 5000, dropout: float = 0. """ super().__init__() self.dropout = nn.Dropout(p=dropout) - - position = torch.arange(max_seq_length).unsqueeze(1) - div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) - - pe = torch.zeros(max_seq_length, d_model) - pe[:, 0::2] = torch.sin(position * div_term) - pe[:, 1::2] = torch.cos(position * div_term[: pe[:, 1::2].shape[1]]) - - self.register_buffer("pe", pe) + self.register_buffer("pe", build_sinusoidal_pe(max_seq_length, d_model)) def forward(self, x: torch.Tensor) -> torch.Tensor: """Add positional encoding to input. diff --git a/src/slices/models/encoders/README.md b/src/slices/models/encoders/README.md index 3f41368..b250ec9 100644 --- a/src/slices/models/encoders/README.md +++ b/src/slices/models/encoders/README.md @@ -1,288 +1,134 @@ -# Transformer Encoder for ICU Time-Series +# Encoder Architectures -A modular, configurable transformer architecture designed for self-supervised learning on ICU time-series data with missing values and variable-length sequences. +This package contains the encoder backbones used across SSL pretraining, +downstream finetuning, and baseline experiments. -## Features +## Canonical Benchmark Encoder -- **Modular Architecture**: Clean separation between input projection, positional encoding, transformer layers, and pooling -- **Missing Value Support**: Handles observation masks for feature-level missingness -- **Variable-Length Sequences**: Proper padding mask handling for variable-length ICU stays -- **Multiple Pooling Strategies**: Mean, max, CLS token, last timestep, or no pooling -- **Modern Best Practices**: Pre-LN by default (more stable training) -- **Configurable**: All hyperparameters exposed via dataclass config +The main benchmark uses `TransformerEncoder` from `transformer.py`. -## Quick Start +Current canonical config: -```python -from slices.models.encoders import TransformerConfig, TransformerEncoder - -# Create config -config = TransformerConfig( - d_input=35, # Number of input features - d_model=128, # Model dimension - n_layers=4, # Number of transformer layers - n_heads=8, # Number of attention heads - d_ff=512, # Feedforward dimension - pooling="mean", # Pooling strategy -) - -# Create encoder -encoder = TransformerEncoder(config) - -# Process ICU time-series -x = torch.randn(32, 48, 35) # (batch, seq_len, features) -embeddings = encoder(x) # (batch, d_model) -``` +- `d_model=64` +- `n_layers=2` +- `n_heads=4` +- `d_ff=256` +- `obs_aware=true` +- `pooling=mean` for downstream tasks +- `pooling=none` for SSL pretraining -## Architecture +The canonical source of truth for these defaults is +`configs/model/transformer.yaml`. -### Input Pipeline +## Obs-Aware Timestep Tokenization -1. **Input Projection**: Linear layer maps raw features to model dimension - - Input: `(B, T, d_input)` → Output: `(B, T, d_model)` +The benchmark transformer is not a plain linear projection over imputed values. +When `obs_aware=True`, each timestep token is built from: -2. **Positional Encoding**: Sinusoidal position embeddings (optional) - - Provides temporal ordering information - - Can be disabled with `use_positional_encoding=False` +```text +concat(values_t, obs_mask_t) -> MLP -> d_model +``` -3. **Optional CLS Token**: Prepended for BERT-style pooling - - Only added when `pooling="cls"` +This gives the encoder direct access to both observed values and missingness +patterns. In the SSL path, `tokenize()` also marks fully unobserved hours as +invalid timestep tokens so they can be excluded from masking and loss logic. -### Transformer Layers +## Core APIs -Each layer consists of: -- **Multi-head Self-Attention**: Captures dependencies across timesteps -- **Feedforward Network**: Two-layer MLP with nonlinearity -- **Layer Normalization**: Pre-LN (default) or Post-LN -- **Residual Connections**: Around attention and feedforward +`TransformerEncoder` exposes three levels of interface: -``` -┌──────────────────────────┐ -│ Input (B, T, d_model) │ -└──────────┬───────────────┘ - │ - ┌──────▼──────┐ - │ LayerNorm │ (Pre-LN only) - └──────┬──────┘ - │ - ┌──────▼──────────┐ - │ Multi-head │ - │ Self-Attention │ - └──────┬──────────┘ - │ - ┌──────▼──────┐ - │ + Residual │ - └──────┬──────┘ - │ - ┌──────▼──────┐ - │ LayerNorm │ - └──────┬──────┘ - │ - ┌──────▼──────┐ - │ Feedforward│ - │ (d_ff) │ - └──────┬──────┘ - │ - ┌──────▼──────┐ - │ + Residual │ - └──────┬──────┘ - │ - ┌──────▼──────┐ - │ LayerNorm │ (Pre-LN only) - └──────┬──────┘ - │ - Output (B, T, d_model) -``` +1. `forward(x, mask, padding_mask)` for normal downstream use +2. `tokenize(x, obs_mask)` for timestep-token creation in SSL objectives +3. `encode(tokens, padding_mask)` for running transformer layers on prebuilt + visible-token subsets -### Pooling Strategies +The SSL objectives use `tokenize()` and `encode()` directly so they can apply +their own masking strategies while sharing the same encoder body. -| Strategy | Output Shape | Use Case | -|----------|--------------|----------| -| `mean` | `(B, d_model)` | Average over valid timesteps (recommended for SSL) | -| `max` | `(B, d_model)` | Max pooling over valid timesteps | -| `cls` | `(B, d_model)` | Use [CLS] token representation (BERT-style) | -| `last` | `(B, d_model)` | Use last valid timestep | -| `none` | `(B, T, d_model)` | Return per-timestep embeddings | +## Pooling -## Configuration +Supported pooling strategies for `TransformerEncoder`: -### Basic Parameters +- `mean` +- `max` +- `cls` +- `last` +- `none` -```python -@dataclass -class TransformerConfig(EncoderConfig): - # Input/Output - d_input: int = 35 # Number of input features - d_model: int = 128 # Model dimension - max_seq_length: int = 168 # Maximum sequence length - - # Architecture - n_layers: int = 4 # Number of transformer layers - n_heads: int = 8 # Number of attention heads - d_ff: int = 512 # Feedforward dimension - - # Regularization - dropout: float = 0.1 # Dropout probability - - # Options - activation: str = "gelu" # Activation: gelu, relu, silu - prenorm: bool = True # Pre-LN (True) or Post-LN (False) - use_positional_encoding: bool = True - pooling: str = "mean" # Pooling strategy -``` +Use `pooling="none"` only when the caller needs per-token outputs, which is the +case for MAE, JEPA, contrastive, TS2Vec, and SMART pretraining flows. -### Common Configurations +Example: -**Small (for testing/debugging)** ```python -config = TransformerConfig( - d_model=64, n_layers=2, n_heads=4, d_ff=256 -) -# ~200K parameters -``` +from slices.models.encoders import TransformerConfig, TransformerEncoder -**Medium (default)** -```python config = TransformerConfig( - d_model=128, n_layers=4, n_heads=8, d_ff=512 + d_input=35, + d_model=64, + n_layers=2, + n_heads=4, + d_ff=256, + obs_aware=True, + pooling="mean", ) -# ~1.5M parameters -``` -**Large (research-scale)** -```python -config = TransformerConfig( - d_model=512, n_layers=8, n_heads=16, d_ff=2048 -) -# ~25M parameters -``` - -## Handling Missing Values - -ICU time-series data often has missing values. The encoder accepts an observation mask: - -```python -# Create observation mask (True = observed, False = missing) -obs_mask = torch.rand(B, T, D) > 0.3 # 30% missing - -# Important: Impute missing values BEFORE passing to encoder -x_imputed = impute(x, obs_mask) # forward fill, mean, etc. - -# Pass both data and mask -embeddings = encoder(x_imputed, mask=obs_mask) -``` - -**Note**: The observation mask is currently used for logging/analysis. Missing values should be imputed before encoding. Future SSL objectives may incorporate the mask explicitly (e.g., masked prediction). - -## Handling Variable-Length Sequences - -ICU stays have different lengths. Use padding masks to handle this: - -```python -# Create padding mask (True = valid, False = padding) -padding_mask = torch.ones(B, T, dtype=torch.bool) -for i, length in enumerate(sequence_lengths): - padding_mask[i, length:] = False - -# Forward pass -embeddings = encoder(x, padding_mask=padding_mask) -``` - -**Important**: Pooling strategies (`mean`, `max`, `last`) correctly handle padding: -- `mean`: Averages only over valid timesteps -- `max`: Max only over valid timesteps -- `last`: Uses last valid timestep (not last position) - -## Design Decisions - -### Why Pre-LN (Pre-LayerNorm)? - -Modern transformers use Pre-LN instead of the original Post-LN: -- **More stable training**: Gradients flow better, especially for deep models -- **Less sensitive to learning rate**: Easier to tune -- **Standard in recent work**: BERT variants, GPT-3, ViT all use Pre-LN - -Set `prenorm=False` if you need original transformer behavior. - -### Why Sinusoidal Positional Encoding? - -- **Extrapolation**: Can handle sequences longer than seen during training -- **Deterministic**: No learned parameters, more interpretable -- **Standard**: Works well for time-series (timestamps are ordered) - -Alternative: Learned positional embeddings (not yet implemented). - -### Why Mean Pooling by Default? - -For SSL pretraining, mean pooling is recommended: -- **Robust**: Less sensitive to outliers than max pooling -- **Uses all information**: Unlike CLS or last timestep -- **Simple**: No extra parameters - -For supervised tasks, CLS pooling may work better (requires fine-tuning). - -## Integration with SSL Objectives - -The encoder is designed to work with SSL objectives: - -```python -from slices.models.encoders import TransformerEncoder -from slices.models.pretraining import MaskedAutoencoderSSL # TODO - -# Create encoder encoder = TransformerEncoder(config) - -# Wrap in SSL objective -ssl_model = MaskedAutoencoderSSL(encoder, ssl_config) - -# Pretraining -x, obs_mask = batch -loss, metrics = ssl_model(x, obs_mask) -loss.backward() +embeddings = encoder(x, mask=obs_mask) ``` -After pretraining, extract the encoder: +For SSL: + ```python -pretrained_encoder = ssl_model.get_encoder() +ssl_config = TransformerConfig( + d_input=35, + d_model=64, + n_layers=2, + n_heads=4, + d_ff=256, + obs_aware=True, + pooling="none", +) ``` -## Performance Tips +## Other Encoders In This Package -### Memory Efficiency +- `ObservationTransformerEncoder` in `observation.py` + - observation-level tokenization + - kept as an alternate architecture, not the canonical benchmark encoder +- `SMARTEncoder` in `smart.py` + - MART-style architecture required by the appendix SMART objective +- `GRUDEncoder` in `gru_d.py` + - GRU-D baseline for contextual comparisons +- `LinearEncoder` in `linear.py` + - simple baseline / utility encoder -- Use gradient checkpointing for very deep models (not yet implemented) -- Use `torch.no_grad()` for inference -- Reduce batch size if OOM +## Which Encoder Goes With Which Objective -### Training Speed +- `mae`, `jepa`, `contrastive`, `ts2vec` + - use `TransformerEncoder` + - require `obs_aware=True` + - require `pooling="none"` during pretraining +- `smart` + - requires `SMARTEncoder` + - requires `pooling="none"` -- Use mixed precision training (AMP) -- Smaller models train faster (d_model=64-128 for initial experiments) -- Use `pooling="none"` only if you need per-timestep outputs +The compatibility checks are enforced in `scripts/training/pretrain.py`. -### Model Size vs. Performance +## Notes On Missingness And Padding -For ICU data (typical: 35 features, 48-hour windows): -- **Small models (d_model=64-128)**: Usually sufficient, faster training -- **Large models (d_model≥256)**: May overfit on small datasets -- **Deep models (n_layers≥6)**: Use Pre-LN and careful tuning +- Observation masks are used directly by the obs-aware transformer path. +- Padding masks are handled at the sequence level and inverted internally to + match PyTorch attention conventions. +- Fully unobserved hours are treated as invalid SSL timesteps in + `TransformerEncoder.tokenize()`. ## Testing -Run the comprehensive test suite: -```bash -uv run pytest tests/test_transformer_encoder.py -v -``` - -Run the example script: -```bash -uv run python examples/transformer_encoder_example.py -``` +Relevant test coverage lives in: -## References +- `tests/test_transformer_encoder.py` +- `tests/test_smart_encoder.py` +- `tests/test_factories.py` -- **Attention Is All You Need** (Vaswani et al., 2017): Original transformer -- **On Layer Normalization in the Transformer Architecture** (Xiong et al., 2020): Pre-LN vs Post-LN -- **BERT** (Devlin et al., 2019): CLS token pooling -- **ricu** (GigaScience, 2023): ICU data concept dictionaries -- **YAIB** (ICLR, 2024): ICU benchmark tasks +Small executable examples live in the focused encoder tests above. diff --git a/src/slices/models/encoders/__init__.py b/src/slices/models/encoders/__init__.py index 26802a4..a39de3f 100644 --- a/src/slices/models/encoders/__init__.py +++ b/src/slices/models/encoders/__init__.py @@ -1,7 +1,8 @@ """Encoder architectures for time-series data.""" -from .base import BaseEncoder, EncoderConfig +from .base import BaseEncoder, EncoderConfig, SSLTokenizingEncoder from .factory import build_encoder, get_encoder_config_class +from .gru_d import GRUDConfig, GRUDEncoder from .linear import LinearConfig, LinearEncoder from .observation import ObservationTransformerConfig, ObservationTransformerEncoder from .smart import SMARTEncoder, SMARTEncoderConfig @@ -12,12 +13,15 @@ "BaseEncoder", "EncoderConfig", "EncoderWithMissingToken", + "GRUDConfig", + "GRUDEncoder", "LinearConfig", "LinearEncoder", "ObservationTransformerConfig", "ObservationTransformerEncoder", "SMARTEncoder", "SMARTEncoderConfig", + "SSLTokenizingEncoder", "TransformerConfig", "TransformerEncoder", "build_encoder", diff --git a/src/slices/models/encoders/base.py b/src/slices/models/encoders/base.py index 9fb9b02..7e31986 100644 --- a/src/slices/models/encoders/base.py +++ b/src/slices/models/encoders/base.py @@ -2,7 +2,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Optional +from typing import Any, Dict, Optional, Protocol, Tuple import torch import torch.nn as nn @@ -61,3 +61,39 @@ def get_output_dim(self) -> int: Output dimension (typically d_model). """ return self.config.d_model + + def handles_missingness_intrinsically(self) -> bool: + """Return whether the encoder consumes observation masks natively. + + Downstream checkpoint loading uses this capability to decide whether an + encoder still needs EncoderWithMissingToken wrapping. + """ + return False + + +class SSLTokenizingEncoder(Protocol): + """Protocol for encoders that expose the SSL timestep-token interface. + + Controlled SSL objectives operate on timestep tokens directly. Keeping this + as a protocol makes the required encoder capability explicit without forcing + every baseline encoder to inherit a second abstract base class. + """ + + config: EncoderConfig + + def tokenize( + self, + x: torch.Tensor, + obs_mask: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, Any]]: + """Convert dense values and observation masks into timestep tokens.""" + + def encode( + self, + tokens: torch.Tensor, + padding_mask: torch.Tensor, + ) -> torch.Tensor: + """Encode prebuilt timestep tokens.""" + + def get_output_dim(self) -> int: + """Return the encoder output dimension.""" diff --git a/src/slices/models/encoders/factory.py b/src/slices/models/encoders/factory.py index bc00e6f..5e71d1d 100644 --- a/src/slices/models/encoders/factory.py +++ b/src/slices/models/encoders/factory.py @@ -7,6 +7,7 @@ from typing import Any, Dict, Type from .base import BaseEncoder, EncoderConfig +from .gru_d import GRUDConfig, GRUDEncoder from .linear import LinearConfig, LinearEncoder from .observation import ObservationTransformerConfig, ObservationTransformerEncoder from .smart import SMARTEncoder, SMARTEncoderConfig @@ -18,6 +19,7 @@ "linear": LinearEncoder, "smart": SMARTEncoder, "observation_transformer": ObservationTransformerEncoder, + "gru_d": GRUDEncoder, } # Registry of encoder configs @@ -26,6 +28,7 @@ "linear": LinearConfig, "smart": SMARTEncoderConfig, "observation_transformer": ObservationTransformerConfig, + "gru_d": GRUDConfig, } diff --git a/src/slices/models/encoders/gru_d.py b/src/slices/models/encoders/gru_d.py new file mode 100644 index 0000000..2228a77 --- /dev/null +++ b/src/slices/models/encoders/gru_d.py @@ -0,0 +1,186 @@ +"""GRU-D encoder for ICU time-series data with learnable decay. + +Implements GRU-D (Che et al. 2018) which handles missing data natively through +input decay (toward empirical mean) and hidden state decay. Compatible with the +existing z-normalized, zero-filled data pipeline where empirical mean = 0. + +This encoder does NOT use EncoderWithMissingToken — missingness is handled +intrinsically via the decay mechanism. +""" + +from dataclasses import dataclass +from typing import Optional + +import torch +import torch.nn as nn + +from .base import BaseEncoder, EncoderConfig + + +@dataclass +class GRUDConfig(EncoderConfig): + """Configuration for GRU-D encoder.""" + + n_gru_layers: int = 1 + + +class GRUDEncoder(BaseEncoder): + """GRU-D encoder with learnable input and hidden state decay. + + Input at each timestep is concat(x_decayed, obs_mask) of size 2*d_input. + Output is the final hidden state of shape (B, d_model). + + Args: + config: GRU-D configuration. + + Example: + >>> config = GRUDConfig(d_input=35, d_model=64, n_gru_layers=1) + >>> encoder = GRUDEncoder(config) + >>> x = torch.randn(4, 48, 35) + >>> mask = torch.rand(4, 48, 35) > 0.3 + >>> out = encoder(x, mask=mask) # (4, 64) + """ + + def __init__(self, config: GRUDConfig) -> None: + super().__init__(config) + self.config: GRUDConfig = config + + d_input = config.d_input + d_model = config.d_model + + # GRU cell: input is concat(x_decayed, mask) = 2 * d_input + self.gru_cell = nn.GRUCell(input_size=2 * d_input, hidden_size=d_model) + + # Input decay parameters + self.W_gamma_x = nn.Parameter(torch.Tensor(d_input)) + self.b_gamma_x = nn.Parameter(torch.Tensor(d_input)) + + # Hidden state decay parameters: feature-wise deltas -> hidden units. + self.W_gamma_h = nn.Parameter(torch.Tensor(d_model, d_input)) + self.b_gamma_h = nn.Parameter(torch.Tensor(d_model)) + + # Empirical mean in z-space is 0 (after z-normalization) + self.register_buffer("x_mean", torch.zeros(d_input)) + + self.dropout = nn.Dropout(config.dropout) + + self._reset_parameters() + + def handles_missingness_intrinsically(self) -> bool: + """GRU-D models missingness via input and hidden-state decay.""" + return True + + def _reset_parameters(self) -> None: + nn.init.uniform_(self.W_gamma_x, -0.1, 0.1) + nn.init.zeros_(self.b_gamma_x) + nn.init.uniform_(self.W_gamma_h, -0.1, 0.1) + nn.init.zeros_(self.b_gamma_h) + + def _compute_input_decay(self, delta_t: torch.Tensor) -> torch.Tensor: + """Compute feature-wise input decay gamma_x for one timestep.""" + return torch.exp(-torch.relu(self.W_gamma_x * delta_t + self.b_gamma_x)) + + def _compute_hidden_decay(self, delta_t: torch.Tensor) -> torch.Tensor: + """Compute hidden-state decay gamma_h from feature-wise deltas.""" + return torch.exp(-torch.relu(delta_t @ self.W_gamma_h.t() + self.b_gamma_h)) + + def _impute_inputs( + self, + x_t: torch.Tensor, + mask_t: torch.Tensor, + gamma_x: torch.Tensor, + x_last_observed: torch.Tensor, + ) -> torch.Tensor: + """Apply GRU-D input decay toward the empirical mean for missing values.""" + mask_float = mask_t.float() + return mask_float * x_t + (1 - mask_float) * ( + gamma_x * x_last_observed + (1 - gamma_x) * self.x_mean + ) + + def _update_last_observed( + self, + x_t: torch.Tensor, + mask_t: torch.Tensor, + x_last_observed: torch.Tensor, + ) -> torch.Tensor: + """Update x_last only with actually observed raw values.""" + return torch.where(mask_t, x_t, x_last_observed) + + def _compute_time_deltas(self, mask: torch.Tensor) -> torch.Tensor: + """Compute time since last observation per feature. + + Args: + mask: Boolean observation mask (B, T, D). True = observed. + + Returns: + Time deltas (B, T, D) in float. delta[t=0] = 0. + For t > 0: delta[t] = 1 if observed at t-1, else delta[t-1] + 1. + """ + B, T, D = mask.shape + delta = torch.zeros(B, T, D, device=mask.device) + for t in range(1, T): + delta[:, t, :] = torch.where( + mask[:, t - 1, :], + torch.ones(B, D, device=mask.device), + delta[:, t - 1, :] + 1, + ) + return delta + + def forward( + self, + x: torch.Tensor, + mask: Optional[torch.Tensor] = None, + padding_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Forward pass with learnable decay for missing values. + + Args: + x: Input tensor (B, T, D). + mask: Observation mask (B, T, D), True = observed. + If None, all values treated as observed. + padding_mask: Unused, kept for API compatibility. + + Returns: + Final hidden state (B, d_model). + """ + B, T, D = x.shape + device = x.device + + if mask is None: + mask = torch.ones(B, T, D, dtype=torch.bool, device=device) + + mask_float = mask.float() + delta = self._compute_time_deltas(mask) + + h = torch.zeros(B, self.config.d_model, device=device) + x_last_observed = self.x_mean.expand(B, -1).clone() + + for t in range(T): + # Input decay + gamma_x = self._compute_input_decay(delta[:, t, :]) + x_decayed = self._impute_inputs( + x[:, t, :], + mask[:, t, :], + gamma_x, + x_last_observed, + ) + + # Hidden state decay + gamma_h = self._compute_hidden_decay(delta[:, t, :]) + h = gamma_h * h + + # GRU update + gru_input = torch.cat([x_decayed, mask_float[:, t, :]], dim=-1) + h = self.gru_cell(gru_input, h) + + x_last_observed = self._update_last_observed( + x[:, t, :], + mask[:, t, :], + x_last_observed, + ) + + h = self.dropout(h) + return h + + def get_output_dim(self) -> int: + return self.config.d_model diff --git a/src/slices/models/encoders/observation.py b/src/slices/models/encoders/observation.py index 94145fb..936548f 100644 --- a/src/slices/models/encoders/observation.py +++ b/src/slices/models/encoders/observation.py @@ -97,6 +97,10 @@ def __init__(self, config: ObservationTransformerConfig) -> None: else: self.final_norm = nn.Identity() + def handles_missingness_intrinsically(self) -> bool: + """Observation tokenization only materializes observed measurements.""" + return True + def _validate_config(self) -> None: if self.config.d_model % self.config.n_heads != 0: raise ValueError( @@ -210,7 +214,11 @@ def encode( (B, N, d_model) encoded tokens. """ # Convert to PyTorch convention: True = ignore - key_padding_mask = ~padding_mask + key_padding_mask = ~padding_mask.to(dtype=torch.bool) + all_masked = key_padding_mask.all(dim=1) + if all_masked.any(): + key_padding_mask = key_padding_mask.clone() + key_padding_mask[all_masked, 0] = False x = tokens for layer in self.layers: diff --git a/src/slices/models/encoders/smart.py b/src/slices/models/encoders/smart.py index 0ed4a34..0854015 100644 --- a/src/slices/models/encoders/smart.py +++ b/src/slices/models/encoders/smart.py @@ -84,8 +84,9 @@ class SeqAttentionBlock(nn.Module): using additive observation-based attention biases (SMART's key innovation). The attention mask uses an additive scheme where mask[i] + mask[j] creates - graded attention: 0 (neither observed) → blocked, 1 (one observed) → partial, - 2 (both observed) → full attention. + graded finite attention bias: 0 (neither observed) is the lowest bias, + 1 (one observed) is intermediate, and 2 (both observed) is highest. + Padding remains hard-blocked, but missing/missing pairs are not. """ def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1) -> None: @@ -141,8 +142,9 @@ def forward( obs_mask_extended = torch.cat([query_observed, obs_mask], dim=2) # (B, V, T+1) obs_mask_flat = obs_mask_extended.reshape(B * V, T_plus_1).float() # (B*V, T+1) - # Additive mask: mask[i] + mask[j] gives 0, 1, or 2 - # This creates graded attention bias (higher = more attention) + # Additive mask: mask[i] + mask[j] gives a finite 0, 1, or 2 bias. + # This is graded attention bias (higher = more attention), not a hard + # block for neither-observed timestep pairs. attn_bias = obs_mask_flat.unsqueeze(-1) + obs_mask_flat.unsqueeze(-2) # (B*V, T+1, T+1) # Handle padding mask if provided (add -inf for padded positions) @@ -266,7 +268,6 @@ def forward(self, x: torch.Tensor, obs_mask: torch.Tensor) -> torch.Tensor: scale = self.d_head**-0.5 attn_weights = torch.matmul(q, k.transpose(-2, -1)) * scale # (B, n_heads, V, V) attn_weights = F.softmax(attn_weights, dim=-1) - attn_weights = self.dropout(attn_weights) attn_out = torch.matmul(attn_weights, v) # (B, n_heads, V, (T+1)*d_head) # Unpack temporal dimension from output @@ -412,10 +413,11 @@ def __init__(self, config: SMARTEncoderConfig) -> None: nn.init.normal_(self.query, std=0.02) # Positional encoding for temporal dimension + # Original SMART has no dropout in positional encoding self.pos_encoder = PositionalEncoding( d_model=config.d_model, max_seq_length=config.max_seq_length + 1, # +1 for query token - dropout=config.dropout, + dropout=0.0, ) # MART blocks @@ -431,8 +433,9 @@ def __init__(self, config: SMARTEncoderConfig) -> None: ] ) - # Final layer norm - self.final_norm = nn.LayerNorm(config.d_model) + def handles_missingness_intrinsically(self) -> bool: + """SMART embeds value and mask pairs jointly at the input layer.""" + return True def _validate_config(self) -> None: """Validate configuration parameters.""" @@ -493,9 +496,6 @@ def forward( for block in self.blocks: x = block(x, mask, padding_mask) - # Final layer norm - x = self.final_norm(x) - # Apply pooling return self._apply_pooling(x, mask) diff --git a/src/slices/models/encoders/transformer.py b/src/slices/models/encoders/transformer.py index 7f06224..be5ac00 100644 --- a/src/slices/models/encoders/transformer.py +++ b/src/slices/models/encoders/transformer.py @@ -5,12 +5,17 @@ """ from dataclasses import dataclass -from typing import Optional +from typing import Any, Dict, Optional, Tuple import torch import torch.nn as nn -from slices.models.common import PositionalEncoding, apply_pooling, get_activation +from slices.models.common import ( + PositionalEncoding, + apply_pooling, + build_sinusoidal_pe, + get_activation, +) from .base import BaseEncoder, EncoderConfig @@ -30,6 +35,7 @@ class TransformerConfig(EncoderConfig): use_positional_encoding: bool = True pooling: str = "mean" # Pooling strategy: mean, max, cls, last, none prenorm: bool = True # Pre-LN vs Post-LN transformer + obs_aware: bool = False # Enable observation-aware MLP projection for SSL class TransformerEncoderLayer(nn.Module): @@ -176,8 +182,21 @@ def __init__(self, config: TransformerConfig) -> None: # Validate configuration before creating modules self._validate_config() - # Input projection: (B, T, d_input) -> (B, T, d_model) - self.input_proj = nn.Linear(config.d_input, config.d_model) + # Input projection — mutually exclusive paths: + # obs_aware: concat (values, obs_mask) → MLP → d_model (for SSL pretraining) + # standard: values → Linear → d_model (for finetuning / supervised) + if config.obs_aware: + self.obs_proj = nn.Sequential( + nn.Linear(2 * config.d_input, config.d_ff), + nn.GELU(), + nn.Dropout(config.dropout), + nn.Linear(config.d_ff, config.d_model), + ) + self.register_buffer( + "time_pe", build_sinusoidal_pe(config.max_seq_length, config.d_model) + ) + else: + self.input_proj = nn.Linear(config.d_input, config.d_model) # Optional CLS token for pooling if config.pooling == "cls": @@ -215,6 +234,10 @@ def __init__(self, config: TransformerConfig) -> None: else: self.final_norm = nn.Identity() + def handles_missingness_intrinsically(self) -> bool: + """Observation-aware transformers consume masks in the input projection.""" + return self.config.obs_aware + def _validate_config(self) -> None: """Validate configuration parameters.""" if self.config.d_model % self.config.n_heads != 0: @@ -229,6 +252,80 @@ def _validate_config(self) -> None: f"Invalid pooling '{self.config.pooling}'. " f"Choose from: {valid_pooling}" ) + def tokenize( + self, + x: torch.Tensor, + obs_mask: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, Any]]: + """Convert timestep data into obs-aware tokens. + + Concatenates values and observation mask, projects through obs_proj MLP, + and adds sinusoidal time positional encoding. + + Requires obs_aware=True in the encoder config. + + Args: + x: Input values of shape (B, T, D). + obs_mask: Boolean mask (B, T, D), True = observed. + + Returns: + tokens: (B, T, d_model) embedded tokens. + padding_mask: (B, T) all True (fixed T, no padding). + token_info: dict with timestep_idx, values, obs_mask. + + Raises: + RuntimeError: If encoder was not configured with obs_aware=True. + """ + if not self.config.obs_aware: + raise RuntimeError( + "tokenize() requires obs_aware=True. " + "Create the encoder with TransformerConfig(obs_aware=True)." + ) + B, T, D = x.shape + device = x.device + + # Concat values + mask -> (B, T, 2D) -> obs_proj -> (B, T, d_model) + combined = torch.cat([x, obs_mask.float()], dim=-1) # (B, T, 2D) + tokens = self.obs_proj(combined) # (B, T, d_model) + + # Add sinusoidal time PE + tokens = tokens + self.time_pe[:T].unsqueeze(0) # (B, T, d_model) + + # Fully unobserved hours are not eligible SSL tokens. + padding_mask = obs_mask.any(dim=-1) + + token_info = { + "timestep_idx": torch.arange(T, device=device).unsqueeze(0).expand(B, -1), + "values": x, + "obs_mask": obs_mask, + "valid_timestep_mask": padding_mask, + } + + return tokens, padding_mask, token_info + + def encode( + self, + tokens: torch.Tensor, + padding_mask: torch.Tensor, + ) -> torch.Tensor: + """Run transformer layers on tokens. + + Args: + tokens: (B, N, d_model) token embeddings. + padding_mask: (B, N) True = valid, False = padding. + + Returns: + (B, N, d_model) encoded tokens. + """ + key_padding_mask = self._key_padding_mask_for_attention(padding_mask) + + x = tokens + for layer in self.layers: + x = layer(x, key_padding_mask=key_padding_mask) + + x = self.final_norm(x) + return x + def forward( self, x: torch.Tensor, @@ -256,8 +353,15 @@ def forward( """ B, T, D = x.shape - # Input projection - x = self.input_proj(x) # (B, T, d_model) + # Input projection (mutually exclusive paths) + if self.config.obs_aware: + if mask is None: + mask = torch.ones(B, T, D, dtype=torch.bool, device=x.device) + combined = torch.cat([x, mask.float()], dim=-1) # (B, T, 2D) + x = self.obs_proj(combined) # (B, T, d_model) + x = x + self.time_pe[:T].unsqueeze(0) + else: + x = self.input_proj(x) # (B, T, d_model) # Add CLS token if using CLS pooling if self.config.pooling == "cls": @@ -269,14 +373,15 @@ def forward( cls_mask = torch.ones(B, 1, dtype=torch.bool, device=x.device) padding_mask = torch.cat([cls_mask, padding_mask], dim=1) # (B, T+1) - # Add positional encoding - x = self.pos_encoder(x) # (B, T, d_model) or (B, T+1, d_model) + # Add positional encoding (skip for obs_aware — time_pe already added) + if not self.config.obs_aware: + x = self.pos_encoder(x) # (B, T, d_model) or (B, T+1, d_model) # Convert padding mask to PyTorch convention (True = ignore) # Our convention: True = valid, False = padding # PyTorch convention: True = padding (ignore), False = valid if padding_mask is not None: - key_padding_mask = ~padding_mask # Invert + key_padding_mask = self._key_padding_mask_for_attention(padding_mask) else: key_padding_mask = None @@ -307,6 +412,22 @@ def _apply_pooling( """ return apply_pooling(x, self.config.pooling, padding_mask) + @staticmethod + def _key_padding_mask_for_attention(padding_mask: torch.Tensor) -> torch.Tensor: + """Convert valid-token masks to PyTorch attention masks without all-masked rows. + + PyTorch attention returns NaNs when a sample has every key masked. Fully + unobserved ICU stays can create that case. We expose one finite dummy token + to attention, while downstream pooling/loss code still receives the original + padding mask and therefore treats the sample as having no valid timesteps. + """ + key_padding_mask = ~padding_mask.to(dtype=torch.bool) + all_masked = key_padding_mask.all(dim=1) + if all_masked.any(): + key_padding_mask = key_padding_mask.clone() + key_padding_mask[all_masked, 0] = False + return key_padding_mask + def get_output_dim(self) -> int: """Return the output dimension of the encoder. diff --git a/src/slices/models/encoders/wrapper.py b/src/slices/models/encoders/wrapper.py index 0f47f9c..b808c2d 100644 --- a/src/slices/models/encoders/wrapper.py +++ b/src/slices/models/encoders/wrapper.py @@ -5,8 +5,8 @@ and finetuning stages. During MAE pretraining, the encoder sees MISSING_TOKEN at positions where -obs_mask=False. Without this wrapper, finetuning would use forward-filled -imputed values instead, creating a mismatch that may degrade performance. +obs_mask=False. Without this wrapper, finetuning would use normalized +zero-filled values instead, creating a mismatch that may degrade performance. """ from typing import Optional diff --git a/src/slices/models/pretraining/README.md b/src/slices/models/pretraining/README.md index 619b362..e3ec45b 100644 --- a/src/slices/models/pretraining/README.md +++ b/src/slices/models/pretraining/README.md @@ -1,83 +1,135 @@ # Self-Supervised Learning Objectives -This module implements various SSL objectives for pretraining on unlabeled ICU time-series data. +This package contains the SSL objectives used by SLICES pretraining. -## Available Objectives +## Implemented Objectives -### MAE (Masked Autoencoder) +- `mae` + - timestep-level masked autoencoding + - reconstructs masked timestep features in input space +- `jepa` + - timestep-level latent prediction with an EMA target encoder +- `contrastive` + - instance-level or temporal NT-Xent over masked views +- `ts2vec` + - temporal contrastive extension with noise augmentation and hierarchical loss +- `smart` + - appendix-only external reference built around `SMARTEncoder` -Learns representations by masking portions of the input and reconstructing them. +The registry lives in `factory.py`. -**Key Features**: -- 4 masking strategies: random, block, timestep, feature -- Respects observation mask (ICU data-aware) -- Lightweight transformer decoder -- Configurable mask ratio and decoder architecture +## Encoder Requirements -**Usage**: -```python -from slices.models.pretraining import MAEConfig, MAEObjective - -mae_config = MAEConfig( - mask_ratio=0.15, - mask_strategy="block", - decoder_d_model=64, - decoder_n_layers=2, -) -mae = MAEObjective(encoder, mae_config) +The objectives do not all share the same encoder contract. -loss, metrics = mae(x, obs_mask) -``` +| Objective | Required Encoder | Key Requirements | +|---|---|---| +| `mae` | `TransformerEncoder` | `obs_aware=True`, `pooling="none"` | +| `jepa` | `TransformerEncoder` | `obs_aware=True`, `pooling="none"` | +| `contrastive` | `TransformerEncoder` | `obs_aware=True`, `pooling="none"` | +| `ts2vec` | `TransformerEncoder` | `obs_aware=True`, `pooling="none"` | +| `smart` | `SMARTEncoder` | `pooling="none"` | -**See**: `docs/MAE_IMPLEMENTATION.md` for detailed documentation. +These constraints are enforced in the objective constructors and by +`scripts/training/pretrain.py`. -## Factory Pattern +## Benchmark Design -Use the factory for easy switching between objectives: +The core controlled benchmark objectives are `mae`, `jepa`, and `contrastive`. +They share the same timestep-level obs-aware transformer encoder and differ only +in the SSL objective and masking logic. -```python -from slices.models.pretraining import build_ssl_objective +Highlights: -# Build any SSL objective -ssl_obj = build_ssl_objective(encoder, config) -``` +- MAE reconstructs masked timestep features +- JEPA predicts masked latent targets from an EMA teacher +- Contrastive aligns sequence or temporal representations across masked views +- TS2Vec gives the contrastive family a temporal extension, but it is not a + core thesis vertex +- SMART remains outside the controlled comparison because it swaps in MART -## Adding New Objectives +## Example -1. Create objective class extending `BaseSSLObjective` -2. Create config extending `SSLConfig` -3. Register in `factory.py` -4. Add tests in `tests/test_.py` -5. Add documentation in `docs/_IMPLEMENTATION.md` +```python +from slices.models.pretraining import MAEConfig, build_ssl_objective -## File Structure +ssl_config = MAEConfig(mask_ratio=0.5) +ssl_objective = build_ssl_objective(encoder, ssl_config) +loss, metrics = ssl_objective(x, obs_mask) ``` + +## File Layout + +```text pretraining/ -├── __init__.py # Exports -├── base.py # Abstract base classes -├── factory.py # Factory functions and registry -├── mae.py # MAE implementation -└── README.md # This file +├── base.py +├── factory.py +├── masking.py +├── mae.py +├── jepa.py +├── contrastive.py +├── ts2vec.py +├── smart.py +└── README.md ``` -## Requirements +## Objective Notes + +### MAE + +- operates on timestep tokens +- encoder sees only visible tokens +- decoder reconstructs full `(B, T, D)` values +- loss is computed only on observed features at masked timesteps + +### JEPA -All SSL objectives require: -- Encoder with `pooling="none"` (for per-timestep outputs) -- Input shape: `(B, T, D)` -- Observation mask shape: `(B, T, D)` +- uses the same timestep-token interface as MAE +- target encoder is an EMA copy kept in eval mode +- supports `mse` and `cosine` latent losses -## Future Objectives +### Contrastive + +- supports `instance` and `temporal` modes +- standard benchmark setting uses complementary masked views in instance mode + +### TS2Vec + +- adds input-level Gaussian noise and optional cropping +- uses independent timestep masks +- applies a hierarchical temporal contrastive loss +- treated as a temporal-contrastive extension, not one of the three controlled + thesis objectives + +### SMART + +- uses element-wise masking, not timestep masking +- requires the MART-style `SMARTEncoder` +- included as an external appendix reference rather than the main benchmark + +## Configuration + +Hydra objective configs live in `configs/ssl/`: + +- `configs/ssl/mae.yaml` +- `configs/ssl/jepa.yaml` +- `configs/ssl/contrastive.yaml` +- `configs/ssl/ts2vec.yaml` +- `configs/ssl/smart.yaml` + +Switch objectives at the CLI with: + +```bash +uv run python scripts/training/pretrain.py dataset=miiv ssl=jepa +``` -Planned implementations: -- **SimCLR**: Contrastive learning with augmentations -- **MoCo**: Momentum contrast -- **JEPA**: Joint-Embedding Predictive Architecture -- **TimeMAE**: Time-specific MAE variant -- **TS-TCC**: Time-series contrastive coding +## Extending -## References +To add a new objective: -- **MAE**: He et al. (2022) "Masked Autoencoders Are Scalable Vision Learners" -- **BERT**: Devlin et al. (2019) "BERT: Pre-training of Deep Bidirectional Transformers" +1. Add a config dataclass extending `SSLConfig` +2. Implement a `BaseSSLObjective` subclass +3. Register both in `factory.py` +4. Add objective-specific tests under `tests/` +5. Add or update the relevant Hydra config in `configs/ssl/` diff --git a/src/slices/models/pretraining/__init__.py b/src/slices/models/pretraining/__init__.py index 3f7af80..5da69ca 100644 --- a/src/slices/models/pretraining/__init__.py +++ b/src/slices/models/pretraining/__init__.py @@ -6,6 +6,7 @@ from .jepa import JEPAConfig, JEPAObjective from .mae import MAEConfig, MAEObjective from .smart import SMARTObjective, SMARTSSLConfig +from .ts2vec import TS2VecConfig, TS2VecObjective __all__ = [ "BaseSSLObjective", @@ -18,6 +19,8 @@ "MAEObjective", "SMARTObjective", "SMARTSSLConfig", + "TS2VecConfig", + "TS2VecObjective", "build_ssl_objective", "get_ssl_config_class", ] diff --git a/src/slices/models/pretraining/base.py b/src/slices/models/pretraining/base.py index f62d712..7929a6b 100644 --- a/src/slices/models/pretraining/base.py +++ b/src/slices/models/pretraining/base.py @@ -2,11 +2,13 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Dict, Tuple +from typing import Dict, Tuple, cast import torch import torch.nn as nn +from slices.models.encoders.base import SSLTokenizingEncoder + @dataclass class SSLConfig: @@ -57,3 +59,44 @@ def get_encoder(self) -> nn.Module: The encoder module. """ return self.encoder + + +def require_ssl_tokenizing_encoder( + encoder: nn.Module, + objective_name: str, +) -> SSLTokenizingEncoder: + """Validate and type-narrow the encoder contract used by controlled SSL. + + MAE, JEPA, contrastive, and TS2Vec share a timestep-token interface instead + of using the pooled downstream encoder forward path. This helper keeps that + contract in one place while preserving each objective's public behavior. + """ + + missing_methods = [ + method_name + for method_name in ("tokenize", "encode") + if not callable(getattr(encoder, method_name, None)) + ] + if missing_methods: + missing = ", ".join(missing_methods) + raise ValueError( + f"{objective_name} requires an encoder with tokenize()/encode() SSL " + f"tokenization support, but {type(encoder).__name__} is missing: {missing}" + ) + + encoder_config = getattr(encoder, "config", None) + if not getattr(encoder_config, "obs_aware", False): + raise ValueError( + f"{objective_name} requires an encoder with obs_aware=True " + "(e.g., TransformerEncoder with obs_aware=True). Got: " + f"{type(encoder).__name__}" + ) + + encoder_pooling = getattr(encoder_config, "pooling", "none") + if encoder_pooling != "none": + raise ValueError( + f"{objective_name} requires encoder with pooling='none' to get " + f"per-token representations, but got pooling='{encoder_pooling}'" + ) + + return cast(SSLTokenizingEncoder, encoder) diff --git a/src/slices/models/pretraining/contrastive.py b/src/slices/models/pretraining/contrastive.py index b67e010..38144ce 100644 --- a/src/slices/models/pretraining/contrastive.py +++ b/src/slices/models/pretraining/contrastive.py @@ -1,19 +1,33 @@ """Contrastive (SimCLR-style) SSL objective for ICU time-series. -Observation-level tokenization variant: uses two different random masks as -two augmented "views" of the same sample, then applies NT-Xent contrastive -loss on the pooled representations. +Timestep-level tokenization variant: creates two "views" of the same sample +via timestep masks, then applies NT-Xent contrastive loss. + +By default, views use **complementary masks** (view 2 = ~view 1), ensuring +zero overlap and forcing the encoder to learn abstract temporal semantics +from non-overlapping windows. Independent random masks are available as a +fallback (complementary_masks=False). + +Supports two modes: +- **instance** (default): SimCLR-style — mean-pool each view to a single + sequence-level embedding, then NT-Xent on a (2B x 2B) matrix. + Compatible with both complementary and independent masks. +- **temporal**: Positive pairs are formed from timesteps visible in + both views (~25% at mask_ratio=0.5). Each overlapping timestep has two + independently contextualized representations (different self-attention + contexts) that form a natural positive pair. All other tokens across the + batch are negatives. NT-Xent operates on a (2N x 2N) matrix where + N = number of overlap tokens. Requires complementary_masks=False. Architecture: -1. ObservationTransformerEncoder.tokenize() → one token per observed measurement -2. Two independent random masks → two different subsets of tokens (views) -3. Encoder processes each view separately → per-token representations -4. Mean-pool over visible tokens → sequence-level embeddings -5. Projection head → low-dimensional normalized embeddings -6. NT-Xent loss: positive pairs = same sample's two views - -Key difference from MAE: discriminative (not reconstructive), global (not local). -Key difference from JEPA: global invariance (not local positional prediction). +1. TransformerEncoder.tokenize() -> one token per timestep (B, T, d_model) +2. Complementary masks (default) or two independent random masks -> two views +3. Encoder processes each view separately -> per-token representations +4. instance: mean-pool -> project -> NT-Xent + temporal: scatter to (B,T,d) -> gather overlap tokens -> project -> NT-Xent + +Key difference from MAE: discriminative (not reconstructive). +Key difference from JEPA: invariance (not positional prediction). """ from dataclasses import dataclass @@ -23,16 +37,24 @@ import torch.nn as nn import torch.nn.functional as F -from .base import BaseSSLObjective, SSLConfig -from .masking import create_observation_mask, extract_visible +from .base import BaseSSLObjective, SSLConfig, require_ssl_tokenizing_encoder +from .masking import ( + create_complementary_timestep_masks, + create_timestep_mask, + extract_visible_timesteps, + scatter_visible_timesteps, +) @dataclass class ContrastiveConfig(SSLConfig): - """Configuration for observation-level contrastive objective.""" + """Configuration for timestep-level contrastive objective.""" name: str = "contrastive" + # Mode: "temporal" (per-timestep overlap pairs) or "instance" (mean-pool) + mode: str = "instance" + # Masking mask_ratio: float = 0.5 @@ -41,7 +63,23 @@ class ContrastiveConfig(SSLConfig): proj_output_dim: int = 128 # Temperature for NT-Xent - temperature: float = 0.1 + temperature: float = 0.07 + + # Use complementary masks for views. Extremely sparse samples may fall back + # to a shared timestep so neither view becomes empty after eligibility + # filtering. + complementary_masks: bool = True + + def __post_init__(self) -> None: + if self.mode not in ("temporal", "instance"): + raise ValueError( + f"ContrastiveConfig.mode must be 'temporal' or 'instance', " f"got '{self.mode}'" + ) + if self.complementary_masks and self.mode == "temporal": + raise ValueError( + "complementary_masks=True is incompatible with mode='temporal'. " + "Temporal mode requires overlapping masks to form positive pairs." + ) class ProjectionHead(nn.Module): @@ -78,38 +116,19 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class ContrastiveObjective(BaseSSLObjective): - """Observation-level contrastive (SimCLR-style) SSL for ICU time-series. + """Timestep-level contrastive SSL for ICU time-series. - Flow: - 1. encoder.tokenize(x, obs_mask) → observation tokens - 2. Two independent random masks → two views - 3. For each view: extract_visible → encode → mean_pool → (B, d_model) - 4. projection_head(pooled) → z1, z2 (B, proj_dim), L2-normalized - 5. NT-Xent loss: match positive pairs across views + Supports temporal mode (per-timestep overlap pairs) and instance mode + (mean-pool, SimCLR-style). See module docstring for details. - Requires ObservationTransformerEncoder with pooling='none'. + Requires encoder with tokenize()/encode() and pooling='none'. """ def __init__(self, encoder: nn.Module, config: ContrastiveConfig) -> None: super().__init__(encoder, config) self.config: ContrastiveConfig = config - # Validate encoder type - if not hasattr(encoder, "tokenize") or not hasattr(encoder, "encode"): - raise ValueError( - "Contrastive requires an encoder with tokenize() and encode() " - "methods (e.g., ObservationTransformerEncoder). Got: " - f"{type(encoder).__name__}" - ) - - # Validate pooling - encoder_pooling = getattr(encoder.config, "pooling", "none") - if encoder_pooling != "none": - raise ValueError( - "Contrastive requires encoder with pooling='none' to get " - "per-token representations for mean-pooling, but got " - f"pooling='{encoder_pooling}'" - ) + require_ssl_tokenizing_encoder(encoder, "Contrastive") d_encoder = encoder.get_output_dim() @@ -140,26 +159,76 @@ def forward( # 1. Tokenize (shared between both views) tokens, padding_mask, token_info = self.encoder.tokenize(x, obs_mask) - # 2. Two independent random masks - ssl_mask_1 = create_observation_mask(padding_mask, self.config.mask_ratio, device) - ssl_mask_2 = create_observation_mask(padding_mask, self.config.mask_ratio, device) + valid_timestep_mask = token_info["valid_timestep_mask"] + + # 2. Create two views via timestep masks + ssl_mask_1 = create_timestep_mask( + B, + T, + self.config.mask_ratio, + device, + valid_timestep_mask=valid_timestep_mask, + ) + if self.config.complementary_masks: + ssl_mask_1, ssl_mask_2 = create_complementary_timestep_masks( + ssl_mask_1, + valid_timestep_mask, + ) + else: + ssl_mask_2 = create_timestep_mask( + B, + T, + self.config.mask_ratio, + device, + valid_timestep_mask=valid_timestep_mask, + ) - # 3. View 1: extract → encode → mean pool - vis_tokens_1, vis_padding_1 = extract_visible(tokens, ssl_mask_1, padding_mask) + effective_mask_1 = ssl_mask_1 & valid_timestep_mask + effective_mask_2 = ssl_mask_2 & valid_timestep_mask + + # 3. Encode both views + vis_tokens_1, vis_padding_1 = extract_visible_timesteps( + tokens, + ssl_mask_1, + valid_timestep_mask=valid_timestep_mask, + ) encoded_1 = self.encoder.encode(vis_tokens_1, vis_padding_1) - pooled_1 = self._mean_pool(encoded_1, vis_padding_1) # (B, d_model) - # 4. View 2: extract → encode → mean pool - vis_tokens_2, vis_padding_2 = extract_visible(tokens, ssl_mask_2, padding_mask) + vis_tokens_2, vis_padding_2 = extract_visible_timesteps( + tokens, + ssl_mask_2, + valid_timestep_mask=valid_timestep_mask, + ) encoded_2 = self.encoder.encode(vis_tokens_2, vis_padding_2) - pooled_2 = self._mean_pool(encoded_2, vis_padding_2) # (B, d_model) - - # 5. Project to contrastive space - z1 = self.projection_head(pooled_1) # (B, proj_dim) - z2 = self.projection_head(pooled_2) # (B, proj_dim) - # 6. NT-Xent loss - loss, metrics = self._nt_xent_loss(z1, z2, ssl_mask_1, ssl_mask_2, padding_mask) + if self.config.mode == "temporal": + # 4a. Scatter encoded tokens back to full (B, T, d) grid + full_1 = self._scatter_to_full(encoded_1, effective_mask_1, T) + full_2 = self._scatter_to_full(encoded_2, effective_mask_2, T) + + # 5a. Temporal NT-Xent on overlapping timesteps + loss, metrics = self._temporal_nt_xent_loss( + full_1, + full_2, + effective_mask_1, + effective_mask_2, + valid_timestep_mask, + ) + else: + # 4b. Instance mode: mean-pool -> project -> instance NT-Xent + pooled_1 = self._mean_pool(encoded_1, vis_padding_1) + pooled_2 = self._mean_pool(encoded_2, vis_padding_2) + + z1 = self.projection_head(pooled_1) # (B, proj_dim) + z2 = self.projection_head(pooled_2) # (B, proj_dim) + + loss, metrics = self._nt_xent_loss( + z1, + z2, + effective_mask_1, + effective_mask_2, + valid_timestep_mask, + ) return loss, metrics @@ -182,26 +251,209 @@ def _mean_pool( counts = padding_mask.sum(dim=1, keepdim=True).clamp(min=1).float() # (B, 1) return summed / counts + @staticmethod + def _timestep_count_metrics( + ssl_mask_1: torch.Tensor, + ssl_mask_2: torch.Tensor, + eligible_mask: torch.Tensor, + ) -> Dict[str, float]: + """Count visible/masked timesteps over eligible non-empty timesteps only.""" + eligible_mask = eligible_mask.to(device=ssl_mask_1.device, dtype=torch.bool) + ssl_mask_1 = ssl_mask_1.to(dtype=torch.bool) & eligible_mask + ssl_mask_2 = ssl_mask_2.to(dtype=torch.bool) & eligible_mask + denom = max(int(eligible_mask.shape[0]), 1) + + n_vis_1 = ssl_mask_1.sum().item() + n_vis_2 = ssl_mask_2.sum().item() + n_masked_1 = ((~ssl_mask_1) & eligible_mask).sum().item() + n_masked_2 = ((~ssl_mask_2) & eligible_mask).sum().item() + + return { + "contrastive_n_visible_view1": n_vis_1 / denom, + "contrastive_n_visible_view2": n_vis_2 / denom, + "contrastive_n_masked_view1": n_masked_1 / denom, + "contrastive_n_masked_view2": n_masked_2 / denom, + } + + @staticmethod + def _scatter_to_full( + encoded: torch.Tensor, + ssl_mask: torch.Tensor, + n_timesteps: int, + ) -> torch.Tensor: + """Scatter visible encoded tokens back to full (B, T, d) tensor. + + Places encoded visible tokens at their original temporal positions, + with zeros at masked positions. + + Args: + encoded: (B, n_vis, d_enc) encoded visible tokens. + ssl_mask: (B, T) True = visible. + n_timesteps: Total number of timesteps T. + + Returns: + (B, T, d_enc) with encoded tokens at visible positions, zeros elsewhere. + """ + return scatter_visible_timesteps(encoded, ssl_mask, n_timesteps) + + def _temporal_nt_xent_loss( + self, + full_1: torch.Tensor, + full_2: torch.Tensor, + ssl_mask_1: torch.Tensor, + ssl_mask_2: torch.Tensor, + eligible_mask: torch.Tensor, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: + """Compute temporal NT-Xent loss on overlapping timestep tokens. + + Timesteps visible in both views have two independently contextualized + representations (from different self-attention contexts). These form + natural positive pairs. All other tokens in the batch are negatives. + + Args: + full_1: (B, T, d_enc) scattered encoded tokens from view 1. + full_2: (B, T, d_enc) scattered encoded tokens from view 2. + ssl_mask_1: (B, T) True = visible in view 1. + ssl_mask_2: (B, T) True = visible in view 2. + eligible_mask: (B, T) True for non-empty timesteps eligible for masking. + + Returns: + (loss, metrics_dict) + """ + B = ssl_mask_1.shape[0] + T = ssl_mask_1.shape[1] + temperature = self.config.temperature + + overlap = ssl_mask_1 & ssl_mask_2 # (B, T) + N = int(overlap.sum().item()) + + # Edge case: not enough overlap tokens for contrastive learning + if N < 2: + loss = full_1.sum() * 0.0 # zero with grad connectivity + with torch.no_grad(): + count_metrics = self._timestep_count_metrics( + ssl_mask_1, + ssl_mask_2, + eligible_mask, + ) + metrics = { + "contrastive_loss": loss.detach(), + "ssl_loss": loss.detach(), + "contrastive_accuracy": torch.tensor(0.0), + "contrastive_pos_similarity": torch.tensor(0.0), + "contrastive_temperature": temperature, + "contrastive_n_timesteps": T, + "contrastive_n_overlap_tokens": 0, + "contrastive_n_overlap_per_sample": 0.0, + } + metrics.update(count_metrics) + return loss, metrics + + # Gather overlap tokens — boolean indexing in row-major order ensures + # full_1[overlap] and full_2[overlap] are aligned (same batch, time pairs) + tokens_1 = full_1[overlap] # (N, d_enc) + tokens_2 = full_2[overlap] # (N, d_enc) + + # Project per-token through shared projection head + z1 = self.projection_head(tokens_1) # (N, proj_dim) + z2 = self.projection_head(tokens_2) # (N, proj_dim) + + # NT-Xent on (2N, 2N) similarity matrix + z = torch.cat([z1, z2], dim=0) # (2N, proj_dim) + sim_matrix = torch.mm(z, z.t()) / temperature # (2N, 2N) + + labels = torch.cat( + [ + torch.arange(N, 2 * N, device=z.device), + torch.arange(N, device=z.device), + ] + ) + + mask = torch.eye(2 * N, dtype=torch.bool, device=z.device) + sim_matrix = sim_matrix.masked_fill(mask, float("-inf")) + + loss = F.cross_entropy(sim_matrix, labels) + + with torch.no_grad(): + preds = sim_matrix.argmax(dim=1) + accuracy = (preds == labels).float().mean() + pos_sim = F.cosine_similarity(z1, z2, dim=-1).mean() + count_metrics = self._timestep_count_metrics( + ssl_mask_1, + ssl_mask_2, + eligible_mask, + ) + + metrics = { + "contrastive_loss": loss.detach(), + "ssl_loss": loss.detach(), + "contrastive_accuracy": accuracy, + "contrastive_pos_similarity": pos_sim, + "contrastive_temperature": temperature, + "contrastive_n_timesteps": T, + "contrastive_n_overlap_tokens": N, + "contrastive_n_overlap_per_sample": N / B, + } + metrics.update(count_metrics) + + return loss, metrics + def _nt_xent_loss( self, z1: torch.Tensor, z2: torch.Tensor, ssl_mask_1: torch.Tensor, ssl_mask_2: torch.Tensor, - padding_mask: torch.Tensor, + eligible_mask: torch.Tensor, ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: """Compute NT-Xent (Normalized Temperature-scaled Cross Entropy) loss. Args: z1: (B, proj_dim) L2-normalized projections from view 1. z2: (B, proj_dim) L2-normalized projections from view 2. - ssl_mask_1: (B, max_obs) mask for view 1 (for metrics). - ssl_mask_2: (B, max_obs) mask for view 2 (for metrics). - padding_mask: (B, max_obs) padding mask (for metrics). + ssl_mask_1: (B, T) mask for view 1 (for metrics). + ssl_mask_2: (B, T) mask for view 2 (for metrics). + eligible_mask: (B, T) True for non-empty timesteps eligible for masking. Returns: (loss, metrics_dict) """ + valid_samples = ssl_mask_1.any(dim=1) & ssl_mask_2.any(dim=1) + n_valid_samples = int(valid_samples.sum().item()) + if n_valid_samples < 2: + loss = (z1.sum() + z2.sum()) * 0.0 + with torch.no_grad(): + T = ssl_mask_1.shape[1] + valid_mask_1 = ssl_mask_1[valid_samples] + valid_mask_2 = ssl_mask_2[valid_samples] + valid_eligible_mask = eligible_mask[valid_samples] + count_metrics = self._timestep_count_metrics( + valid_mask_1, + valid_mask_2, + valid_eligible_mask, + ) + metrics = { + "contrastive_loss": loss.detach(), + "ssl_loss": loss.detach(), + "contrastive_accuracy": torch.tensor(0.0, device=z1.device), + "contrastive_pos_similarity": torch.tensor(0.0, device=z1.device), + "contrastive_alignment": torch.tensor(0.0, device=z1.device), + "contrastive_uniformity": torch.tensor(0.0, device=z1.device), + "contrastive_effective_rank": torch.tensor(0.0, device=z1.device), + "contrastive_temperature": self.config.temperature, + "contrastive_n_timesteps": T, + "contrastive_n_samples_used": n_valid_samples, + "contrastive_n_samples_skipped": z1.shape[0] - n_valid_samples, + } + metrics.update(count_metrics) + return loss, metrics + + z1 = z1[valid_samples] + z2 = z2[valid_samples] + ssl_mask_1 = ssl_mask_1[valid_samples] + ssl_mask_2 = ssl_mask_2[valid_samples] + eligible_mask = eligible_mask[valid_samples] + B = z1.shape[0] temperature = self.config.temperature @@ -231,24 +483,45 @@ def _nt_xent_loss( # Positive pair similarities (before temperature scaling) pos_sim = F.cosine_similarity(z1, z2, dim=-1).mean() - # Token statistics - n_total = padding_mask.sum().item() - n_vis_1 = (ssl_mask_1 & padding_mask).sum().item() - n_vis_2 = (ssl_mask_2 & padding_mask).sum().item() - n_masked_1 = ((~ssl_mask_1) & padding_mask).sum().item() - n_masked_2 = ((~ssl_mask_2) & padding_mask).sum().item() + # --- Collapse monitoring (Wang & Isola 2020) --- + # Alignment: mean L2 distance between positive pairs (lower = better) + alignment = (z1 - z2).norm(dim=-1).pow(2).mean() + + # Uniformity: log avg pairwise Gaussian potential on hypersphere + # Lower = more uniform = better spread of representations + # Use z1 only (one view per sample) to avoid inflating with positives + sq_pdist = torch.cdist(z1, z1, p=2).pow(2) # (B, B) + uniformity = sq_pdist.mul(-2).exp().mean().log() + + # Effective rank via singular value entropy (Roy & Vetterli 2007) + # High effective rank = diverse representation dimensions + # Low effective rank → collapse to a low-dim subspace + _, s, _ = torch.svd_lowrank(z1 - z1.mean(dim=0), q=min(B, z1.shape[1])) + p = s.clamp(min=0).pow(2) # eigenvalues (squared singular values) + p = p / p.sum().clamp(min=1e-12) + eff_rank = (-p * p.clamp(min=1e-7).log()).sum().exp() + + # Timestep statistics over eligible non-empty timesteps only. + T = ssl_mask_1.shape[1] + count_metrics = self._timestep_count_metrics( + ssl_mask_1, + ssl_mask_2, + eligible_mask, + ) metrics = { "contrastive_loss": loss.detach(), "ssl_loss": loss.detach(), "contrastive_accuracy": accuracy, "contrastive_pos_similarity": pos_sim, + "contrastive_alignment": alignment, + "contrastive_uniformity": uniformity, + "contrastive_effective_rank": eff_rank, "contrastive_temperature": temperature, - "contrastive_n_tokens_per_sample": n_total / B, - "contrastive_n_visible_view1": n_vis_1 / B, - "contrastive_n_visible_view2": n_vis_2 / B, - "contrastive_n_masked_view1": n_masked_1 / B, - "contrastive_n_masked_view2": n_masked_2 / B, + "contrastive_n_timesteps": T, + "contrastive_n_samples_used": B, + "contrastive_n_samples_skipped": valid_samples.numel() - B, } + metrics.update(count_metrics) return loss, metrics diff --git a/src/slices/models/pretraining/factory.py b/src/slices/models/pretraining/factory.py index c891149..470f52b 100644 --- a/src/slices/models/pretraining/factory.py +++ b/src/slices/models/pretraining/factory.py @@ -13,6 +13,7 @@ from .jepa import JEPAConfig, JEPAObjective from .mae import MAEConfig, MAEObjective from .smart import SMARTObjective, SMARTSSLConfig +from .ts2vec import TS2VecConfig, TS2VecObjective # Registry of available SSL objectives SSL_REGISTRY: Dict[str, Type[BaseSSLObjective]] = { @@ -20,6 +21,7 @@ "smart": SMARTObjective, "jepa": JEPAObjective, "contrastive": ContrastiveObjective, + "ts2vec": TS2VecObjective, } # Registry of SSL configs @@ -28,6 +30,7 @@ "smart": SMARTSSLConfig, "jepa": JEPAConfig, "contrastive": ContrastiveConfig, + "ts2vec": TS2VecConfig, } diff --git a/src/slices/models/pretraining/jepa.py b/src/slices/models/pretraining/jepa.py index 53b676a..81489bb 100644 --- a/src/slices/models/pretraining/jepa.py +++ b/src/slices/models/pretraining/jepa.py @@ -1,13 +1,13 @@ """Joint-Embedding Predictive Architecture (JEPA) for ICU time-series SSL. -Observation-level tokenization variant: same masking as MAE, but predicts +Timestep-level tokenization variant: same masking as MAE, but predicts latent representations of masked tokens instead of raw input values. Architecture: -1. ObservationTransformerEncoder.tokenize() → one token per observed measurement -2. Random mask: 50% of observation tokens are masked (configurable via mask_ratio) -3. Online encoder processes only visible tokens (50%) -4. EMA target encoder processes ALL tokens → target representations +1. TransformerEncoder.tokenize() -> one token per timestep (B, T, d_model) +2. Random mask: 50% of timestep tokens are masked (configurable via mask_ratio) +3. Online encoder processes only visible tokens +4. EMA target encoder processes ALL tokens -> target representations 5. Predictor reassembles (visible encoded + mask tokens), predicts target repr 6. MSE/cosine loss on masked positions in latent space @@ -16,7 +16,6 @@ """ import copy -import math from dataclasses import dataclass from typing import Dict, Tuple @@ -24,18 +23,27 @@ import torch.nn as nn import torch.nn.functional as F -from .base import BaseSSLObjective, SSLConfig -from .masking import create_observation_mask, extract_visible +from slices.models.common import build_sinusoidal_pe + +from .base import BaseSSLObjective, SSLConfig, require_ssl_tokenizing_encoder +from .masking import ( + create_block_timestep_mask, + create_timestep_mask, + extract_visible_timesteps, + scatter_visible_timesteps, +) @dataclass class JEPAConfig(SSLConfig): - """Configuration for observation-level JEPA objective.""" + """Configuration for timestep-level JEPA objective.""" name: str = "jepa" # Masking mask_ratio: float = 0.5 + mask_strategy: str = "block" # "random" or "block" + mask_n_blocks: int = 3 # Number of contiguous blocks (only for block strategy) # Predictor parameters (mirrors MAE decoder for fairness) predictor_d_model: int = 128 @@ -45,7 +53,7 @@ class JEPAConfig(SSLConfig): predictor_dropout: float = 0.1 # Momentum encoder - momentum_base: float = 0.996 + momentum_base: float = 0.999 momentum_final: float = 1.0 # Loss @@ -56,13 +64,12 @@ class JEPAPredictor(nn.Module): """Lightweight transformer predictor for JEPA. Same architecture as MAEDecoder but output_proj maps to d_encoder - (representation vectors) instead of scalar values. + (representation vectors) instead of D feature values. """ def __init__( self, d_encoder: int, - n_features: int, max_seq_length: int, config: JEPAConfig, ) -> None: @@ -75,10 +82,7 @@ def __init__( self.mask_token = nn.Parameter(torch.zeros(1, 1, d_pred)) nn.init.normal_(self.mask_token, std=0.02) - self.feature_embed = nn.Embedding(n_features, d_pred) - - pe = self._build_sinusoidal_pe(max_seq_length, d_pred) - self.register_buffer("time_pe", pe) + self.register_buffer("time_pe", build_sinusoidal_pe(max_seq_length, d_pred)) predictor_layer = nn.TransformerEncoderLayer( d_model=d_pred, @@ -96,106 +100,94 @@ def __init__( self.embed_dropout = nn.Dropout(config.predictor_dropout) - # Output projects to d_encoder (representation space), not scalar + # Output projects to d_encoder (representation space) self.output_proj = nn.Linear(d_pred, d_encoder) - @staticmethod - def _build_sinusoidal_pe(max_len: int, d_model: int) -> torch.Tensor: - position = torch.arange(max_len).unsqueeze(1).float() - div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) - pe = torch.zeros(max_len, d_model) - pe[:, 0::2] = torch.sin(position * div_term) - pe[:, 1::2] = torch.cos(position * div_term[: pe[:, 1::2].shape[1]]) - return pe - def forward( self, encoded_visible: torch.Tensor, ssl_mask: torch.Tensor, token_info: Dict[str, torch.Tensor], - max_tokens: int, - token_padding_mask: torch.Tensor, + n_timesteps: int, ) -> torch.Tensor: - """Predict target representations at all positions. + """Predict target representations at all timestep positions. Args: encoded_visible: (B, n_vis, d_encoder) encoded visible tokens. - ssl_mask: (B, max_obs) bool, True = visible, False = masked. - token_info: dict with timestep_idx, feature_idx (B, max_obs). - max_tokens: total number of token positions (max_obs). - token_padding_mask: (B, max_obs) True = valid token, False = padding. + ssl_mask: (B, T) bool, True = visible, False = masked. + token_info: dict with timestep_idx (B, T) and optional + valid_timestep_mask (B, T) marking timesteps that had at + least one observed variable. + n_timesteps: total number of timesteps T. Returns: - (B, max_tokens, d_encoder) predicted representations per token. + (B, T, d_encoder) predicted representations per timestep. """ - B = encoded_visible.shape[0] - d_pred = self.config.predictor_d_model - vis_proj = self.encoder_proj(encoded_visible) # (B, n_vis, d_pred) - full_tokens = self.mask_token.expand(B, max_tokens, d_pred).clone() + valid_timestep_mask = token_info.get("valid_timestep_mask") + if valid_timestep_mask is None: + visible_mask = ssl_mask + else: + visible_mask = ssl_mask & valid_timestep_mask.to( + device=ssl_mask.device, + dtype=torch.bool, + ) - # Scatter visible tokens to original positions (same logic as MAEDecoder) - vis_indices = ssl_mask.float().argsort(dim=1, descending=True, stable=True) - n_vis = vis_proj.shape[1] - scatter_idx = vis_indices[:, :n_vis] - scatter_idx_expanded = scatter_idx.unsqueeze(-1).expand(-1, -1, d_pred) - full_tokens.scatter_(1, scatter_idx_expanded, vis_proj) + full_tokens = scatter_visible_timesteps( + vis_proj, + visible_mask, + n_timesteps, + fill_value=self.mask_token, + ) - # Add positional information + # Add time PE timestep_idx = token_info["timestep_idx"] - feature_idx = token_info["feature_idx"] - - full_tokens = full_tokens + self.feature_embed(feature_idx) full_tokens = full_tokens + self.time_pe[timestep_idx] full_tokens = self.embed_dropout(full_tokens) - key_padding_mask = ~token_padding_mask + predictor_padding_mask = None + if valid_timestep_mask is not None: + predictor_padding_mask = ~valid_timestep_mask.to( + device=full_tokens.device, + dtype=torch.bool, + ) + all_masked = predictor_padding_mask.all(dim=1) + if all_masked.any(): + predictor_padding_mask = predictor_padding_mask.clone() + predictor_padding_mask[all_masked, 0] = False - decoded = self.predictor(full_tokens, src_key_padding_mask=key_padding_mask) + # Fully unobserved hours should not be available as predictor context. + decoded = self.predictor(full_tokens, src_key_padding_mask=predictor_padding_mask) # Project to encoder representation space - predictions = self.output_proj(decoded) # (B, max_tokens, d_encoder) + predictions = self.output_proj(decoded) # (B, T, d_encoder) return predictions class JEPAObjective(BaseSSLObjective): - """Observation-level JEPA for ICU time-series. + """Timestep-level JEPA for ICU time-series. Flow: - 1. encoder.tokenize(x, obs_mask) → observation tokens + padding + info - 2. Random mask: 75% tokens masked, 25% visible - 3. Online encoder.encode(visible_tokens) → context representations - 4. EMA target encoder tokenize+encode(ALL tokens) → target representations - 5. Predictor(context, mask, info) → predicted representations + 1. encoder.tokenize(x, obs_mask) -> timestep tokens + padding + info + 2. Random mask: mask_ratio timesteps masked + 3. Online encoder.encode(visible_tokens) -> context representations + 4. EMA target encoder tokenize+encode(ALL tokens) -> target representations + 5. Predictor(context, mask, info) -> predicted representations 6. MSE/cosine loss on masked positions in latent space - Requires ObservationTransformerEncoder with pooling='none'. + Requires encoder with tokenize()/encode() and pooling='none'. """ def __init__(self, encoder: nn.Module, config: JEPAConfig) -> None: super().__init__(encoder, config) self.config: JEPAConfig = config - # Validate encoder type - if not hasattr(encoder, "tokenize") or not hasattr(encoder, "encode"): - raise ValueError( - "JEPA requires an encoder with tokenize() and encode() methods " - "(e.g., ObservationTransformerEncoder). Got: " - f"{type(encoder).__name__}" - ) - - # Validate pooling - encoder_pooling = getattr(encoder.config, "pooling", "none") - if encoder_pooling != "none": - raise ValueError( - "JEPA requires encoder with pooling='none' to get per-token " - f"representations, but got pooling='{encoder_pooling}'" - ) + require_ssl_tokenizing_encoder(encoder, "JEPA") d_encoder = encoder.get_output_dim() - n_features = encoder.config.d_input + self.d_encoder = d_encoder max_seq_length = encoder.config.max_seq_length self.missing_token = None @@ -209,7 +201,6 @@ def __init__(self, encoder: nn.Module, config: JEPAConfig) -> None: # Create predictor self.predictor = JEPAPredictor( d_encoder=d_encoder, - n_features=n_features, max_seq_length=max_seq_length, config=config, ) @@ -233,41 +224,63 @@ def forward( B, T, D = x.shape device = x.device - # 1. Tokenize observed measurements (online encoder) + # 1. Tokenize timesteps (online encoder) tokens, padding_mask, token_info = self.encoder.tokenize(x, obs_mask) - max_obs = tokens.shape[1] - - # 2. Create SSL mask on observation tokens - ssl_mask = create_observation_mask(padding_mask, self.config.mask_ratio, device) + valid_timestep_mask = token_info["valid_timestep_mask"] + + # 2. Create SSL mask on timesteps + if self.config.mask_strategy == "block": + ssl_mask = create_block_timestep_mask( + B, + T, + self.config.mask_ratio, + device, + n_blocks=self.config.mask_n_blocks, + valid_timestep_mask=valid_timestep_mask, + ) + else: + ssl_mask = create_timestep_mask( + B, + T, + self.config.mask_ratio, + device, + valid_timestep_mask=valid_timestep_mask, + ) + effective_visible_mask = ssl_mask & valid_timestep_mask # 3. Extract visible tokens - visible_tokens, vis_padding = extract_visible(tokens, ssl_mask, padding_mask) + visible_tokens, vis_padding = extract_visible_timesteps( + tokens, + ssl_mask, + valid_timestep_mask=valid_timestep_mask, + ) # 4. Encode visible tokens only (online encoder) encoded_visible = self.encoder.encode(visible_tokens, vis_padding) # 5. Target encoder processes ALL tokens (no masking, no grad) - # Keep target encoder in eval mode to disable dropout — Lightning's - # model.train() call at epoch start would otherwise re-enable it, - # corrupting target representations with stochastic noise. self.target_encoder.eval() with torch.no_grad(): target_tokens, target_padding, _ = self.target_encoder.tokenize(x, obs_mask) target_repr = self.target_encoder.encode( target_tokens, target_padding - ) # (B, max_obs, d_model) + ) # (B, T, d_model) # 6. Predictor predicts target representations predicted_repr = self.predictor( encoded_visible=encoded_visible, - ssl_mask=ssl_mask, + ssl_mask=effective_visible_mask, token_info=token_info, - max_tokens=max_obs, - token_padding_mask=padding_mask, - ) # (B, max_obs, d_encoder) + n_timesteps=T, + ) # (B, T, d_encoder) # 7. Compute loss on masked positions - loss, metrics = self._compute_loss(predicted_repr, target_repr, ssl_mask, padding_mask) + loss, metrics = self._compute_loss( + predicted_repr, + target_repr, + ssl_mask, + valid_timestep_mask, + ) return loss, metrics @@ -276,29 +289,56 @@ def _compute_loss( predicted: torch.Tensor, target: torch.Tensor, ssl_mask: torch.Tensor, - padding_mask: torch.Tensor, + valid_timestep_mask: torch.Tensor, ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: - """Compute loss on masked token representations. + """Compute loss on masked timestep representations. Args: - predicted: (B, max_obs, d_encoder) predicted representations. - target: (B, max_obs, d_encoder) target representations. - ssl_mask: (B, max_obs) True = visible, False = masked. - padding_mask: (B, max_obs) True = valid token. + predicted: (B, T, d_encoder) predicted representations. + target: (B, T, d_encoder) target representations. + ssl_mask: (B, T) True = visible, False = masked. Returns: (loss, metrics_dict) """ - # Loss only on masked AND valid positions - loss_mask = (~ssl_mask) & padding_mask # (B, max_obs) + # Loss only on masked timesteps that had at least one observation. + loss_mask = (~ssl_mask) & valid_timestep_mask # (B, T) + + # Collapse monitoring metrics (computed on raw target before normalization) + with torch.no_grad(): + B, T = ssl_mask.shape + n_masked = loss_mask.sum().item() + n_visible = (ssl_mask & valid_timestep_mask).sum().item() + + # Flatten to (B*T, d_encoder) for batch-level statistics + target_flat = target.reshape(-1, self.d_encoder) + + # Std of each feature dim across all positions — collapse → 0 + target_repr_std = target_flat.std(dim=0).mean().item() + + # Mean pairwise cosine similarity (sampled for efficiency) + # Sample up to 256 vectors to keep cost O(1) + n_vecs = target_flat.shape[0] + if n_vecs > 256: + idx = torch.randperm(n_vecs, device=target_flat.device)[:256] + sampled = target_flat[idx] + else: + sampled = target_flat + sampled_norm = F.normalize(sampled, dim=-1) + cos_matrix = sampled_norm @ sampled_norm.T + # Exclude diagonal (self-similarity = 1) + n_sampled = sampled_norm.shape[0] + mask_diag = ~torch.eye(n_sampled, dtype=torch.bool, device=cos_matrix.device) + target_repr_mean_cos = cos_matrix[mask_diag].mean().item() + + # Normalize targets to prevent scale/shift collapse + target = F.layer_norm(target, [self.d_encoder]) if self.config.loss_type == "mse": element_loss = F.mse_loss(predicted, target, reduction="none") - # Average over d_encoder dimension - element_loss = element_loss.mean(dim=-1) # (B, max_obs) + element_loss = element_loss.mean(dim=-1) # (B, T) elif self.config.loss_type == "cosine": - # Cosine distance: 1 - cosine_similarity - cos_sim = F.cosine_similarity(predicted, target, dim=-1) # (B, max_obs) + cos_sim = F.cosine_similarity(predicted, target, dim=-1) # (B, T) element_loss = 1.0 - cos_sim else: raise ValueError(f"Unknown loss type: {self.config.loss_type}") @@ -306,19 +346,16 @@ def _compute_loss( loss = (element_loss * loss_mask.float()).sum() / loss_mask.float().sum().clamp(min=1) with torch.no_grad(): - B = padding_mask.shape[0] - n_total_tokens = padding_mask.sum().item() - n_masked = loss_mask.sum().item() - n_visible = (ssl_mask & padding_mask).sum().item() - metrics = { "jepa_loss": loss.detach(), "ssl_loss": loss.detach(), - "jepa_mask_ratio_actual": n_masked / max(n_total_tokens, 1), - "jepa_n_tokens_per_sample": n_total_tokens / B, + "jepa_mask_ratio_actual": n_masked / max(valid_timestep_mask.sum().item(), 1), + "jepa_n_timesteps": T, "jepa_n_visible_per_sample": n_visible / B, "jepa_n_masked_per_sample": n_masked / B, "jepa_momentum": self._current_momentum, + "jepa_target_repr_std": target_repr_std, + "jepa_target_repr_mean_cos": target_repr_mean_cos, } return loss, metrics diff --git a/src/slices/models/pretraining/mae.py b/src/slices/models/pretraining/mae.py index b1b04c6..aec636f 100644 --- a/src/slices/models/pretraining/mae.py +++ b/src/slices/models/pretraining/mae.py @@ -1,39 +1,40 @@ """Masked Autoencoder (MAE) for self-supervised learning on ICU time-series. -Observation-level tokenization variant: each token = one observed (timestep, feature, value) -triplet. The encoder only sees visible (unmasked) tokens, matching the original MAE design. +Timestep-level tokenization variant: each token = one timestep with all D features. +The encoder only sees visible (unmasked) timestep tokens, matching the original MAE design. Architecture: -1. ObservationTransformerEncoder.tokenize() → one token per observed measurement -2. Random mask: 50% of observation tokens are masked (configurable via mask_ratio) -3. Encoder processes only visible tokens (50%) -4. Decoder reassembles full sequence (visible + mask tokens), predicts scalar values -5. MSE loss on masked token values only +1. TransformerEncoder.tokenize() -> one token per timestep (B, T, d_model) +2. Random mask: 50% of timestep tokens are masked (configurable via mask_ratio) +3. Encoder processes only visible tokens +4. Decoder reassembles full sequence (visible + mask tokens), predicts D features per timestep +5. MSE loss on observed features at masked timesteps only """ -import math from dataclasses import dataclass from typing import Dict, Tuple import torch import torch.nn as nn -from .base import BaseSSLObjective, SSLConfig -from .masking import create_observation_mask, extract_visible +from slices.models.common import build_sinusoidal_pe + +from .base import BaseSSLObjective, SSLConfig, require_ssl_tokenizing_encoder +from .masking import create_timestep_mask, extract_visible_timesteps, scatter_visible_timesteps @dataclass class MAEConfig(SSLConfig): - """Configuration for observation-level MAE objective. + """Configuration for timestep-level MAE objective. - The encoder only sees visible (unmasked) observation tokens. - Mask ratio is applied to observation tokens (not timesteps). + The encoder only sees visible (unmasked) timestep tokens. + Mask ratio is applied to timesteps (not individual observations). """ name: str = "mae" # Masking - mask_ratio: float = 0.5 # Fraction of observation tokens to mask + mask_ratio: float = 0.5 # Fraction of timesteps to mask # Decoder parameters decoder_d_model: int = 128 @@ -44,12 +45,12 @@ class MAEConfig(SSLConfig): class MAEDecoder(nn.Module): - """Lightweight decoder for observation-level MAE. + """Lightweight decoder for timestep-level MAE. Reassembles encoded visible tokens with learnable mask tokens at masked - positions. Adds positional information (feature + temporal) to all tokens - before running through lightweight transformer layers. Predicts a scalar - value per token. + timestep positions. Adds temporal positional information to all tokens + before running through lightweight transformer layers. Predicts D features + per timestep. """ def __init__( @@ -70,17 +71,10 @@ def __init__( self.mask_token = nn.Parameter(torch.zeros(1, 1, d_dec)) nn.init.normal_(self.mask_token, std=0.02) - # Positional information for decoder - self.feature_embed = nn.Embedding(n_features, d_dec) - # Sinusoidal time PE for decoder - pe = self._build_sinusoidal_pe(max_seq_length, d_dec) - self.register_buffer("time_pe", pe) + self.register_buffer("time_pe", build_sinusoidal_pe(max_seq_length, d_dec)) - # Decoder transformer layers. - # PyTorch's TransformerEncoderLayer is used here (not DecoderLayer) because - # MAE's decoder uses self-attention over all tokens (visible + mask), not - # cross-attention between encoder output and mask tokens. + # Decoder transformer layers decoder_layer = nn.TransformerEncoderLayer( d_model=d_dec, nhead=config.decoder_n_heads, @@ -98,124 +92,97 @@ def __init__( # Dropout on input embeddings before decoder transformer self.embed_dropout = nn.Dropout(config.decoder_dropout) - # Output: predict scalar value per token - self.output_proj = nn.Linear(d_dec, 1) - - @staticmethod - def _build_sinusoidal_pe(max_len: int, d_model: int) -> torch.Tensor: - position = torch.arange(max_len).unsqueeze(1).float() - div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) - pe = torch.zeros(max_len, d_model) - pe[:, 0::2] = torch.sin(position * div_term) - pe[:, 1::2] = torch.cos(position * div_term[: pe[:, 1::2].shape[1]]) - return pe + # Output: predict D features per timestep + self.output_proj = nn.Linear(d_dec, n_features) def forward( self, encoded_visible: torch.Tensor, ssl_mask: torch.Tensor, token_info: Dict[str, torch.Tensor], - max_tokens: int, - token_padding_mask: torch.Tensor, + n_timesteps: int, ) -> torch.Tensor: - """Decode: reassemble visible + mask tokens, predict values. + """Decode: reassemble visible + mask tokens, predict feature values. Args: encoded_visible: (B, n_vis, d_encoder) encoded visible tokens. - ssl_mask: (B, max_obs) bool, True = visible, False = masked. - token_info: dict with timestep_idx, feature_idx (B, max_obs). - max_tokens: total number of token positions (max_obs). - token_padding_mask: (B, max_obs) True = valid token, False = padding. + ssl_mask: (B, T) bool, True = visible, False = masked. + token_info: dict with timestep_idx (B, T). + n_timesteps: total number of timesteps T. Returns: - (B, max_tokens) predicted scalar value per token position. + (B, T, D) predicted feature values per timestep. """ - B = encoded_visible.shape[0] - d_dec = self.config.decoder_d_model - # Project visible tokens to decoder space vis_proj = self.encoder_proj(encoded_visible) # (B, n_vis, d_dec) - # Build full decoder input: place visible tokens and mask tokens - full_tokens = self.mask_token.expand(B, max_tokens, d_dec).clone() - - # Place visible tokens at their original positions. - # ssl_mask is (B, max_obs), True at visible positions. - # Stable descending argsort on ssl_mask.float() puts visible (1.0) positions - # first. This mirrors extract_visible(), so vis_indices[:, :n_vis] gives - # the original positions of the n_vis visible tokens in the same order - # they were fed to the encoder. - # Padding invariant: padding positions have ssl_mask=True (see - # create_observation_mask), so they sort among visible positions. - # But because tokenization places valid tokens before padding and argsort - # is stable, valid-visible tokens always precede padding tokens in the - # sorted order. We only scatter the first n_vis entries, which are the - # actual visible tokens, not padding. - vis_indices = ssl_mask.float().argsort(dim=1, descending=True, stable=True) - # vis_indices[:, :n_vis] gives the original positions of visible tokens - n_vis = vis_proj.shape[1] - scatter_idx = vis_indices[:, :n_vis] # (B, n_vis) - scatter_idx_expanded = scatter_idx.unsqueeze(-1).expand(-1, -1, d_dec) # (B, n_vis, d_dec) - full_tokens.scatter_(1, scatter_idx_expanded, vis_proj) - - # Add positional information to ALL token positions - timestep_idx = token_info["timestep_idx"] # (B, max_tokens) - feature_idx = token_info["feature_idx"] # (B, max_tokens) - - full_tokens = full_tokens + self.feature_embed(feature_idx) + valid_timestep_mask = token_info.get("valid_timestep_mask") + if valid_timestep_mask is None: + visible_mask = ssl_mask + else: + visible_mask = ssl_mask & valid_timestep_mask.to( + device=ssl_mask.device, + dtype=torch.bool, + ) + + full_tokens = scatter_visible_timesteps( + vis_proj, + visible_mask, + n_timesteps, + fill_value=self.mask_token, + ) + + # Add time PE to all positions + timestep_idx = token_info["timestep_idx"] # (B, T) full_tokens = full_tokens + self.time_pe[timestep_idx] full_tokens = self.embed_dropout(full_tokens) - # Build key_padding_mask for decoder (True = ignore in PyTorch convention) - key_padding_mask = ~token_padding_mask # (B, max_tokens) + decoder_padding_mask = None + if valid_timestep_mask is not None: + decoder_padding_mask = ~valid_timestep_mask.to( + device=full_tokens.device, + dtype=torch.bool, + ) + all_masked = decoder_padding_mask.all(dim=1) + if all_masked.any(): + decoder_padding_mask = decoder_padding_mask.clone() + decoder_padding_mask[all_masked, 0] = False - # Run decoder transformer - decoded = self.decoder(full_tokens, src_key_padding_mask=key_padding_mask) + # Fully unobserved hours should not be available as decoder context. + if decoder_padding_mask is None: + decoded = self.decoder(full_tokens) + else: + decoded = self.decoder(full_tokens, src_key_padding_mask=decoder_padding_mask) - # Predict scalar value per token - predictions = self.output_proj(decoded).squeeze(-1) # (B, max_tokens) + # Predict D features per timestep + predictions = self.output_proj(decoded) # (B, T, D) return predictions class MAEObjective(BaseSSLObjective): - """Observation-level Masked Autoencoder for ICU time-series. + """Timestep-level Masked Autoencoder for ICU time-series. Flow: - 1. encoder.tokenize(x, obs_mask) → observation tokens + padding + info - 2. Random mask: 75% of tokens masked, 25% visible - 3. encoder.encode(visible_tokens) → encoded visible - 4. decoder(encoded_visible, mask, info) → predicted values per token - 5. MSE loss on masked token values + 1. encoder.tokenize(x, obs_mask) -> timestep tokens + padding + info + 2. Random mask: mask_ratio of timesteps masked + 3. encoder.encode(visible_tokens) -> encoded visible + 4. decoder(encoded_visible, mask, info) -> predicted values (B, T, D) + 5. MSE loss on observed features at masked timesteps - Requires ObservationTransformerEncoder with pooling='none'. + Requires encoder with tokenize()/encode() and pooling='none'. """ def __init__(self, encoder: nn.Module, config: MAEConfig) -> None: super().__init__(encoder, config) self.config: MAEConfig = config - # Validate encoder type - if not hasattr(encoder, "tokenize") or not hasattr(encoder, "encode"): - raise ValueError( - "MAE requires an encoder with tokenize() and encode() methods " - "(e.g., ObservationTransformerEncoder). Got: " - f"{type(encoder).__name__}" - ) - - # Validate pooling - encoder_pooling = getattr(encoder.config, "pooling", "none") - if encoder_pooling != "none": - raise ValueError( - "MAE requires encoder with pooling='none' to get per-token " - f"representations, but got pooling='{encoder_pooling}'" - ) + require_ssl_tokenizing_encoder(encoder, "MAE") d_encoder = encoder.get_output_dim() n_features = encoder.config.d_input max_seq_length = encoder.config.max_seq_length - # No missing_token needed -- observation encoder handles this intrinsically self.missing_token = None # Create decoder @@ -231,7 +198,7 @@ def forward( x: torch.Tensor, obs_mask: torch.Tensor, ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: - """Compute observation-level MAE loss. + """Compute timestep-level MAE loss. Args: x: Input tensor (B, T, D). @@ -243,18 +210,26 @@ def forward( B, T, D = x.shape device = x.device - # 1. Tokenize observed measurements + # 1. Tokenize timesteps tokens, padding_mask, token_info = self.encoder.tokenize(x, obs_mask) - # tokens: (B, max_obs, d_model), padding_mask: (B, max_obs) - max_obs = tokens.shape[1] - true_values = token_info["values"] # (B, max_obs) - - # 2. Create SSL mask on observation tokens - # ssl_mask: True = visible, False = masked - ssl_mask = create_observation_mask(padding_mask, self.config.mask_ratio, device) + # tokens: (B, T, d_model) + + # 2. Create SSL mask on timesteps + ssl_mask = create_timestep_mask( + B, + T, + self.config.mask_ratio, + device, + valid_timestep_mask=token_info["valid_timestep_mask"], + ) # 3. Extract visible tokens - visible_tokens, vis_padding = extract_visible(tokens, ssl_mask, padding_mask) + valid_timestep_mask = token_info["valid_timestep_mask"] + visible_tokens, vis_padding = extract_visible_timesteps( + tokens, + ssl_mask, + valid_timestep_mask=valid_timestep_mask, + ) # 4. Encode visible tokens only encoded_visible = self.encoder.encode(visible_tokens, vis_padding) @@ -264,13 +239,18 @@ def forward( encoded_visible=encoded_visible, ssl_mask=ssl_mask, token_info=token_info, - max_tokens=max_obs, - token_padding_mask=padding_mask, + n_timesteps=T, + ) # (B, T, D) + + # 6. Compute loss on observed features at masked timesteps + loss, metrics = self._compute_loss( + predictions, + x, + ssl_mask, + obs_mask, + valid_timestep_mask, ) - # 6. Compute loss on masked token values - loss, metrics = self._compute_loss(predictions, true_values, ssl_mask, padding_mask) - return loss, metrics def _compute_loss( @@ -278,37 +258,41 @@ def _compute_loss( predictions: torch.Tensor, true_values: torch.Tensor, ssl_mask: torch.Tensor, - padding_mask: torch.Tensor, + obs_mask: torch.Tensor, + valid_timestep_mask: torch.Tensor, ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: - """Compute MSE loss on masked token predictions. + """Compute MSE loss on observed features at masked timesteps. Args: - predictions: (B, max_obs) predicted values. - true_values: (B, max_obs) original observed values. - ssl_mask: (B, max_obs) True = visible, False = masked. - padding_mask: (B, max_obs) True = valid token. + predictions: (B, T, D) predicted feature values. + true_values: (B, T, D) original input values. + ssl_mask: (B, T) True = visible, False = masked. + obs_mask: (B, T, D) True = observed. + valid_timestep_mask: (B, T) True for timesteps with observations. Returns: (loss, metrics_dict) """ - # Loss only on masked AND valid positions - loss_mask = (~ssl_mask) & padding_mask # (B, max_obs) + # Loss mask: masked timesteps AND observed features -> (B, T, D) + loss_mask = (~ssl_mask).unsqueeze(-1) & obs_mask squared_error = (predictions - true_values) ** 2 loss = (squared_error * loss_mask.float()).sum() / loss_mask.float().sum().clamp(min=1) with torch.no_grad(): - B = padding_mask.shape[0] - n_total_tokens = padding_mask.sum().item() - n_masked = loss_mask.sum().item() - n_visible = (ssl_mask & padding_mask).sum().item() - - # Visible (unmasked) reconstruction for monitoring - visible_mask = ssl_mask & padding_mask - if visible_mask.sum() > 0: + B, T, D = true_values.shape + n_timesteps = T + n_valid_timesteps = valid_timestep_mask.sum().item() + n_masked_timesteps = ((~ssl_mask) & valid_timestep_mask).sum().item() + n_visible_timesteps = (ssl_mask & valid_timestep_mask).sum().item() + n_loss_positions = loss_mask.sum().item() + + # Visible reconstruction for monitoring + visible_loss_mask = ssl_mask.unsqueeze(-1) & obs_mask + if visible_loss_mask.sum() > 0: visible_loss = ( - squared_error * visible_mask.float() - ).sum() / visible_mask.float().sum() + squared_error * visible_loss_mask.float() + ).sum() / visible_loss_mask.float().sum() else: visible_loss = torch.tensor(0.0, device=loss.device) @@ -317,12 +301,11 @@ def _compute_loss( "ssl_loss": loss.detach(), "mae_recon_loss_masked": loss.detach(), "mae_recon_loss_visible": visible_loss, - "mae_mask_ratio_actual": n_masked / max(n_total_tokens, 1), - "mae_obs_ratio": n_total_tokens - / max(padding_mask.shape[0] * padding_mask.shape[1], 1), - "mae_n_tokens_per_sample": n_total_tokens / B, - "mae_n_visible_per_sample": n_visible / B, - "mae_n_masked_per_sample": n_masked / B, + "mae_mask_ratio_actual": n_masked_timesteps / max(n_valid_timesteps, 1), + "mae_n_timesteps": n_timesteps, + "mae_n_visible_per_sample": n_visible_timesteps / B, + "mae_n_masked_per_sample": n_masked_timesteps / B, + "mae_n_loss_positions": n_loss_positions, } return loss, metrics diff --git a/src/slices/models/pretraining/masking.py b/src/slices/models/pretraining/masking.py index 101cf72..b273597 100644 --- a/src/slices/models/pretraining/masking.py +++ b/src/slices/models/pretraining/masking.py @@ -79,3 +79,285 @@ def extract_visible( vis_padding = vis_positions < n_visible.unsqueeze(1) # (B, max_vis) return visible_tokens, vis_padding + + +def create_timestep_mask( + batch_size: int, + n_timesteps: int, + mask_ratio: float, + device: torch.device, + valid_timestep_mask: torch.Tensor | None = None, +) -> torch.Tensor: + """Create random mask at timestep level. + + Args: + batch_size: Batch size B. + n_timesteps: Number of timesteps T. + mask_ratio: Fraction of timesteps to mask. + device: Device. + valid_timestep_mask: Optional (B, T) bool mask marking timesteps with + at least one observed variable. When provided, fully unobserved + timesteps are always treated as visible/excluded from SSL masking. + + Returns: + ssl_mask: (B, T) bool mask, True = visible, False = masked. + """ + if valid_timestep_mask is None: + valid_timestep_mask = torch.ones(batch_size, n_timesteps, dtype=torch.bool, device=device) + else: + valid_timestep_mask = valid_timestep_mask.to(device=device, dtype=torch.bool) + + ssl_mask = torch.ones(batch_size, n_timesteps, dtype=torch.bool, device=device) + for b in range(batch_size): + valid_idx = valid_timestep_mask[b].nonzero(as_tuple=True)[0] + n_valid = int(valid_idx.numel()) + n_masked = _masked_timestep_budget(n_valid, mask_ratio) + if n_masked == 0: + continue + + masked_ordinals = torch.randperm(n_valid, device=device)[:n_masked] + ssl_mask[b, valid_idx[masked_ordinals]] = False + + return ssl_mask + + +def create_complementary_timestep_masks( + primary_mask: torch.Tensor, + valid_timestep_mask: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Create paired timestep masks without allowing an empty eligible view. + + The secondary mask is complementary over eligible timesteps whenever + possible. Sparse samples need a fallback: + + - 0 eligible timesteps: both masks remain effectively empty + - 1 eligible timestep: expose it in both views + - >=2 eligible timesteps: keep views complementary, but if the primary view + happened to keep every eligible timestep visible, move one timestep into + the secondary view so both remain encodable + """ + primary = primary_mask.to(dtype=torch.bool).clone() + valid = valid_timestep_mask.to(device=primary.device, dtype=torch.bool) + secondary = (~primary) | (~valid) + + batch_size = primary.shape[0] + for b in range(batch_size): + valid_idx = valid[b].nonzero(as_tuple=True)[0] + n_valid = int(valid_idx.numel()) + + if n_valid == 0: + secondary[b] = True + continue + + if n_valid == 1: + idx = valid_idx[0] + primary[b, idx] = True + secondary[b, idx] = True + continue + + secondary_visible = (secondary[b] & valid[b]).nonzero(as_tuple=True)[0] + if secondary_visible.numel() == 0: + primary_visible = (primary[b] & valid[b]).nonzero(as_tuple=True)[0] + idx = primary_visible[0] + primary[b, idx] = False + secondary[b, idx] = True + + return primary, secondary + + +def _masked_timestep_budget(n_valid: int, mask_ratio: float) -> int: + """Return the integer masked-token budget while keeping one valid token visible.""" + if n_valid <= 1 or mask_ratio <= 0: + return 0 + + requested = int(n_valid * mask_ratio + 0.5) + return min(max(requested, 1), n_valid - 1) + + +def _random_positive_composition(total: int, parts: int) -> torch.Tensor: + """Split ``total`` into ``parts`` positive integer pieces.""" + if parts <= 1: + return torch.tensor([total], dtype=torch.long) + if total == parts: + return torch.ones(parts, dtype=torch.long) + + cuts = torch.randperm(total - 1)[: parts - 1] + 1 + cuts = cuts.sort().values + boundaries = torch.cat( + [ + torch.zeros(1, dtype=torch.long), + cuts.to(dtype=torch.long), + torch.tensor([total], dtype=torch.long), + ] + ) + return boundaries[1:] - boundaries[:-1] + + +def _random_nonnegative_composition(total: int, parts: int) -> torch.Tensor: + """Split ``total`` into ``parts`` non-negative integer pieces.""" + if parts <= 0: + return torch.empty(0, dtype=torch.long) + if total <= 0: + return torch.zeros(parts, dtype=torch.long) + + assignments = torch.randint(parts, (total,)) + return torch.bincount(assignments, minlength=parts).to(dtype=torch.long) + + +def create_block_timestep_mask( + batch_size: int, + n_timesteps: int, + mask_ratio: float, + device: torch.device, + n_blocks: int = 3, + valid_timestep_mask: torch.Tensor | None = None, +) -> torch.Tensor: + """Create contiguous block mask at timestep level. + + Masks up to ``n_blocks`` non-overlapping contiguous spans that hit the + requested masked-token budget after taking the union over valid timesteps. + Fully unobserved timesteps are never counted in the SSL budget and are + marked visible so losses ignore them. + + Args: + batch_size: Batch size B. + n_timesteps: Number of timesteps T. + mask_ratio: Fraction of timesteps to mask. + device: Device. + n_blocks: Number of contiguous blocks to mask (default 3). + valid_timestep_mask: Optional (B, T) bool mask marking timesteps with + at least one observed variable. + + Returns: + ssl_mask: (B, T) bool mask, True = visible, False = masked. + """ + if valid_timestep_mask is None: + valid_timestep_mask = torch.ones(batch_size, n_timesteps, dtype=torch.bool, device=device) + else: + valid_timestep_mask = valid_timestep_mask.to(device=device, dtype=torch.bool) + + ssl_mask = torch.ones(batch_size, n_timesteps, dtype=torch.bool, device=device) + requested_blocks = max(int(n_blocks), 1) + + for b in range(batch_size): + valid_idx = valid_timestep_mask[b].nonzero(as_tuple=True)[0] + n_valid = int(valid_idx.numel()) + n_masked = _masked_timestep_budget(n_valid, mask_ratio) + if n_masked == 0: + continue + + # Distinct blocks need at least one visible eligible timestep between + # them. Reduce the block count when the requested budget leaves too few + # visible timesteps to separate every block. + max_separated_blocks = n_valid - n_masked + 1 + block_count = min(requested_blocks, n_masked, max_separated_blocks) + + lengths = _random_positive_composition(n_masked, block_count) + total_gap = n_valid - n_masked + extra_gap = total_gap - (block_count - 1) + gap_extras = _random_nonnegative_composition(extra_gap, block_count + 1) + gaps = gap_extras.clone() + if block_count > 1: + gaps[1:block_count] += 1 + + ordinal_mask = torch.zeros(n_valid, dtype=torch.bool, device=device) + cursor = int(gaps[0].item()) + for block_idx in range(block_count): + length = int(lengths[block_idx].item()) + ordinal_mask[cursor : cursor + length] = True + cursor += length + int(gaps[block_idx + 1].item()) + + ssl_mask[b, valid_idx[ordinal_mask]] = False + + return ssl_mask + + +def extract_visible_timesteps( + tokens: torch.Tensor, + ssl_mask: torch.Tensor, + valid_timestep_mask: torch.Tensor | None = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Extract visible timestep tokens from full sequence. + + Args: + tokens: (B, T, d_model) + ssl_mask: (B, T) True = visible, False = masked. + valid_timestep_mask: Optional (B, T) bool mask for timesteps that + should participate in SSL tokenization. + + Returns: + visible_tokens: (B, max_vis, d_model) + vis_padding: (B, max_vis) True = valid visible token + """ + B, T, d_model = tokens.shape + + if valid_timestep_mask is None: + visible_mask = ssl_mask + else: + visible_mask = ssl_mask & valid_timestep_mask.to(device=tokens.device, dtype=torch.bool) + + n_visible = visible_mask.sum(dim=1) # (B,) + max_vis = max(int(n_visible.max().item()), 1) + + # Argsort: visible (True=1) first + sort_idx = visible_mask.float().argsort(dim=1, descending=True, stable=True) + sort_idx_expanded = sort_idx.unsqueeze(-1).expand(-1, -1, d_model) + + sorted_tokens = tokens.gather(1, sort_idx_expanded) + visible_tokens = sorted_tokens[:, :max_vis, :] # (B, max_vis, d_model) + + vis_positions = torch.arange(max_vis, device=tokens.device).unsqueeze(0) + vis_padding = vis_positions < n_visible.unsqueeze(1) # (B, max_vis) + + return visible_tokens, vis_padding + + +def scatter_visible_timesteps( + visible_tokens: torch.Tensor, + visible_mask: torch.Tensor, + n_timesteps: int, + fill_value: torch.Tensor | None = None, +) -> torch.Tensor: + """Place variable-count visible tokens back on a full timestep grid. + + `extract_visible_timesteps` pads each batch row to the maximum visible-token + count. This helper scatters only the real visible tokens for each sample, so + padded tokens cannot overwrite masked timestep positions. + + Args: + visible_tokens: (B, max_vis, d_model) tokens returned by an encoder. + visible_mask: (B, T) True at the original visible timestep positions. + n_timesteps: Total number of timesteps T. + fill_value: Optional (1, 1, d_model) or broadcastable fill tensor. When + omitted, masked positions are filled with zeros. + + Returns: + (B, T, d_model) tensor with visible tokens restored to their original + timestep positions. + """ + B, max_vis, d_model = visible_tokens.shape + device = visible_tokens.device + visible_mask = visible_mask.to(device=device, dtype=torch.bool) + + if visible_mask.shape != (B, n_timesteps): + raise ValueError( + f"visible_mask must have shape ({B}, {n_timesteps}), " + f"got {tuple(visible_mask.shape)}" + ) + + if fill_value is None: + full = visible_tokens.new_zeros((B, n_timesteps, d_model)) + else: + fill = fill_value.to(device=device, dtype=visible_tokens.dtype) + full = fill.expand(B, n_timesteps, d_model).clone() + + visible_indices = visible_mask.float().argsort(dim=1, descending=True, stable=True) + visible_counts = visible_mask.sum(dim=1).clamp(max=max_vis) + + for b in range(B): + n_visible = int(visible_counts[b].item()) + if n_visible == 0: + continue + full[b, visible_indices[b, :n_visible]] = visible_tokens[b, :n_visible].to(full.dtype) + + return full diff --git a/src/slices/models/pretraining/smart.py b/src/slices/models/pretraining/smart.py index a1692d9..c4efe57 100644 --- a/src/slices/models/pretraining/smart.py +++ b/src/slices/models/pretraining/smart.py @@ -173,6 +173,8 @@ def __init__(self, encoder: nn.Module, config: SMARTSSLConfig) -> None: # Freeze target encoder - only updated via momentum for param in self.target_encoder.parameters(): param.requires_grad = False + # Keep the EMA teacher deterministic even while the online model trains. + self.target_encoder.eval() # Create predictor (simple MLP like original) self.predictor = SMARTPredictor( @@ -184,6 +186,12 @@ def __init__(self, encoder: nn.Module, config: SMARTSSLConfig) -> None: # Track momentum for logging self._current_momentum = config.momentum_base + def train(self, mode: bool = True) -> "SMARTObjective": + """Keep the EMA teacher in eval mode even when the objective trains.""" + super().train(mode) + self.target_encoder.eval() + return self + def forward( self, x: torch.Tensor, diff --git a/src/slices/models/pretraining/ts2vec.py b/src/slices/models/pretraining/ts2vec.py new file mode 100644 index 0000000..74844f2 --- /dev/null +++ b/src/slices/models/pretraining/ts2vec.py @@ -0,0 +1,498 @@ +"""TS2Vec-style temporal contrastive SSL objective for ICU time-series. + +Addresses the augmentation limitation of the standard contrastive objective +(masking-only) by combining timestamp masking with input-level Gaussian noise +and optional random cropping — augmentations natural to contrastive learning. + +Key differences from the standard contrastive objective: +- **Input-level augmentation**: Gaussian noise applied to raw values before + tokenization, so the obs-aware MLP sees genuinely different inputs per view + (not just different attention contexts from masking). +- **Independent masks**: Two independent random masks (not complementary), + producing ~25% temporal overlap at mask_ratio=0.5. +- **Temporal contrastive loss**: Per-timestep positive pairs (same timestep + across two views). Within-sample timesteps are hard negatives, cross-sample + timesteps are additional negatives. +- **Hierarchical pooling**: Max-pool over increasing temporal scales, computing + the temporal loss at each scale. Captures multi-scale temporal structure. + +Architecture: +1. Add independent Gaussian noise to create two input views +2. Tokenize each view independently via encoder.tokenize() +3. Apply independent random timestep masks to each view +4. Encode visible tokens for each view via encoder.encode() +5. Scatter back to full (B, T, d) grids +6. Project per-timestep through shared projection head +7. Compute hierarchical temporal contrastive loss on overlapping timesteps + +Reference: Yue et al., "TS2Vec: Towards Universal Representation of Time +Series", AAAI 2022. +""" + +from dataclasses import dataclass +from typing import Dict, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .base import BaseSSLObjective, SSLConfig, require_ssl_tokenizing_encoder +from .masking import create_timestep_mask, extract_visible_timesteps, scatter_visible_timesteps + + +@dataclass +class TS2VecConfig(SSLConfig): + """Configuration for TS2Vec-style temporal contrastive objective.""" + + name: str = "ts2vec" + + # Masking (same default as other objectives for comparability) + mask_ratio: float = 0.5 + + # Input-level augmentation + noise_scale: float = 0.01 # Gaussian noise std (0 = masking only) + crop_ratio: float = 1.0 # Sub-sequence crop ratio (1.0 = no crop) + + # Projection head + proj_hidden_dim: int = 256 + proj_output_dim: int = 64 + + # Temperature for temporal NT-Xent + temperature: float = 0.05 + + # Hierarchical temporal pooling + n_hierarchical_scales: int = 4 # Number of max-pool scales (1 = no hierarchy) + + def __post_init__(self) -> None: + if self.noise_scale < 0: + raise ValueError(f"noise_scale must be >= 0, got {self.noise_scale}") + if not 0 < self.crop_ratio <= 1.0: + raise ValueError(f"crop_ratio must be in (0, 1], got {self.crop_ratio}") + if self.n_hierarchical_scales < 1: + raise ValueError( + f"n_hierarchical_scales must be >= 1, got {self.n_hierarchical_scales}" + ) + + +class TemporalProjectionHead(nn.Module): + """Per-timestep MLP projection head with L2 normalization.""" + + def __init__(self, d_input: int, hidden_dim: int, output_dim: int) -> None: + super().__init__() + self.net = nn.Sequential( + nn.Linear(d_input, hidden_dim), + nn.LayerNorm(hidden_dim), + nn.ReLU(inplace=True), + nn.Linear(hidden_dim, output_dim), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Project and L2-normalize per-timestep representations. + + Args: + x: (B, T, d_input) or (N, d_input) + + Returns: + L2-normalized projections, same leading dims. + """ + z = self.net(x) + return F.normalize(z, dim=-1) + + +class TS2VecObjective(BaseSSLObjective): + """TS2Vec-style temporal contrastive SSL for ICU time-series. + + Uses input-level augmentation (noise + masking) and hierarchical temporal + contrastive loss. See module docstring for details. + """ + + def __init__(self, encoder: nn.Module, config: TS2VecConfig) -> None: + super().__init__(encoder, config) + self.config: TS2VecConfig = config + + require_ssl_tokenizing_encoder(encoder, "TS2Vec") + + d_encoder = encoder.get_output_dim() + self.missing_token = None + + self.projection_head = TemporalProjectionHead( + d_input=d_encoder, + hidden_dim=config.proj_hidden_dim, + output_dim=config.proj_output_dim, + ) + + def forward( + self, + x: torch.Tensor, + obs_mask: torch.Tensor, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: + """Compute hierarchical temporal contrastive loss. + + Args: + x: Input tensor (B, T, D). + obs_mask: Observation mask (B, T, D), True = observed. + + Returns: + (loss, metrics_dict) + """ + B, T, D = x.shape + device = x.device + + # 1. Create two augmented input views + x1, x2 = self._create_augmented_views(x, obs_mask) + + # 2. Optional random cropping to shared sub-interval + if self.config.crop_ratio < 1.0: + crop_len = max(2, int(T * self.config.crop_ratio)) + start = torch.randint(0, T - crop_len + 1, (1,)).item() + x1 = x1[:, start : start + crop_len] + x2 = x2[:, start : start + crop_len] + obs_mask = obs_mask[:, start : start + crop_len] + T = crop_len + + # 3. Tokenize each view independently (noise means different MLP inputs) + tokens_1, _, token_info_1 = self.encoder.tokenize(x1, obs_mask) + tokens_2, _, _ = self.encoder.tokenize(x2, obs_mask) + valid_timestep_mask = token_info_1["valid_timestep_mask"] + + # 4. Independent timestep masks + ssl_mask_1 = create_timestep_mask( + B, + T, + self.config.mask_ratio, + device, + valid_timestep_mask=valid_timestep_mask, + ) + ssl_mask_2 = create_timestep_mask( + B, + T, + self.config.mask_ratio, + device, + valid_timestep_mask=valid_timestep_mask, + ) + effective_mask_1 = ssl_mask_1 & valid_timestep_mask + effective_mask_2 = ssl_mask_2 & valid_timestep_mask + + # 5. Encode visible tokens for each view + vis_1, vp_1 = extract_visible_timesteps( + tokens_1, + ssl_mask_1, + valid_timestep_mask=valid_timestep_mask, + ) + enc_1 = self.encoder.encode(vis_1, vp_1) + + vis_2, vp_2 = extract_visible_timesteps( + tokens_2, + ssl_mask_2, + valid_timestep_mask=valid_timestep_mask, + ) + enc_2 = self.encoder.encode(vis_2, vp_2) + + # 6. Scatter back to full (B, T, d) grids + full_1 = self._scatter_to_full(enc_1, effective_mask_1, T) + full_2 = self._scatter_to_full(enc_2, effective_mask_2, T) + + # 7. Project per-timestep + z1 = self.projection_head(full_1) # (B, T, proj_dim) + z2 = self.projection_head(full_2) # (B, T, proj_dim) + + # 8. Find overlap and valid masks + overlap = effective_mask_1 & effective_mask_2 # (B, T) + + # 9. Hierarchical temporal contrastive loss + # Pass per-view masks so max-pool can ignore masked positions + loss, metrics = self._hierarchical_temporal_loss( + z1, + z2, + overlap, + effective_mask_1, + effective_mask_2, + ) + + # Add masking statistics + with torch.no_grad(): + metrics.update( + { + "ts2vec_n_timesteps": torch.tensor(T), + "ts2vec_n_visible_view1": torch.tensor(effective_mask_1.sum().item() / B), + "ts2vec_n_visible_view2": torch.tensor(effective_mask_2.sum().item() / B), + "ts2vec_n_overlap_per_sample": torch.tensor(overlap.sum().item() / B), + "ts2vec_overlap_ratio": overlap.sum().float() + / valid_timestep_mask.sum().clamp(min=1).float(), + } + ) + + return loss, metrics + + def _create_augmented_views( + self, + x: torch.Tensor, + obs_mask: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Create two input views with independent Gaussian noise. + + Noise is applied only to observed values (via obs_mask) to avoid + injecting signal into missing positions. + + Args: + x: (B, T, D) + obs_mask: (B, T, D) True = observed + + Returns: + (x1, x2) each (B, T, D) + """ + if self.config.noise_scale > 0 and self.training: + noise_mask = obs_mask.float() + x1 = x + self.config.noise_scale * torch.randn_like(x) * noise_mask + x2 = x + self.config.noise_scale * torch.randn_like(x) * noise_mask + else: + x1 = x + x2 = x + return x1, x2 + + @staticmethod + def _scatter_to_full( + encoded: torch.Tensor, + ssl_mask: torch.Tensor, + n_timesteps: int, + ) -> torch.Tensor: + """Scatter visible encoded tokens back to full (B, T, d) tensor. + + Args: + encoded: (B, n_vis, d_enc) + ssl_mask: (B, T) True = visible + n_timesteps: Total T + + Returns: + (B, T, d_enc) with encoded tokens at visible positions, zeros elsewhere. + """ + return scatter_visible_timesteps(encoded, ssl_mask, n_timesteps) + + @staticmethod + def _masked_max_pool1d( + z: torch.Tensor, + valid_mask: torch.Tensor, + pool_size: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Max-pool along temporal dim, ignoring masked (invalid) positions. + + Fills invalid positions with -inf before pooling so they are never + selected. A pooled position is valid if ANY timestep in its window + was valid. + + Args: + z: (B, T, d) representations + valid_mask: (B, T) True = valid position + pool_size: Pooling kernel size + + Returns: + (z_pooled, valid_pooled) where z_pooled is (B, T', d) and + valid_pooled is (B, T') bool. + """ + B, T, d = z.shape + + # Truncate to multiple of pool_size + T_trunc = (T // pool_size) * pool_size + z_t = z[:, :T_trunc].permute(0, 2, 1) # (B, d, T_trunc) + mask_trunc = valid_mask[:, :T_trunc] # (B, T_trunc) + + # Fill invalid positions with -inf so max-pool ignores them + inv_mask = ~mask_trunc.unsqueeze(1).expand_as(z_t) # (B, d, T_trunc) + z_t = z_t.masked_fill(inv_mask, float("-inf")) + + z_pooled = F.max_pool1d(z_t, kernel_size=pool_size).permute(0, 2, 1) + + # A pooled position is valid if any input position was valid + valid_pooled = ( + F.max_pool1d(mask_trunc.float().unsqueeze(1), kernel_size=pool_size).squeeze(1).bool() + ) + + # Replace any remaining -inf (all-masked windows) with 0 + z_pooled = z_pooled.masked_fill(z_pooled == float("-inf"), 0.0) + + return z_pooled, valid_pooled + + def _hierarchical_temporal_loss( + self, + z1: torch.Tensor, + z2: torch.Tensor, + overlap: torch.Tensor, + mask_1: torch.Tensor, + mask_2: torch.Tensor, + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: + """Compute temporal contrastive loss at multiple temporal scales. + + At each scale, max-pool over non-overlapping windows (ignoring masked + positions via -inf fill), then compute the temporal contrastive loss + on the pooled representations. + + Args: + z1: (B, T, proj_dim) projected representations from view 1 + z2: (B, T, proj_dim) projected representations from view 2 + overlap: (B, T) True = timestep visible in both views + mask_1: (B, T) True = visible in view 1 + mask_2: (B, T) True = visible in view 2 + + Returns: + (loss, metrics_dict) + """ + B, T, proj_dim = z1.shape + device = z1.device + + # Compute pool sizes: [1, 2, 4, ...] up to n_hierarchical_scales + pool_sizes = [2**i for i in range(self.config.n_hierarchical_scales)] + # Filter out scales larger than T + pool_sizes = [p for p in pool_sizes if p <= T] + if not pool_sizes: + pool_sizes = [1] + + total_loss = torch.tensor(0.0, device=device) + n_scales = 0 + total_overlap_tokens = 0 + + for pool_size in pool_sizes: + if pool_size == 1: + z1_pooled = z1 + z2_pooled = z2 + overlap_pooled = overlap + else: + z1_pooled, valid_1 = self._masked_max_pool1d(z1, mask_1, pool_size) + z2_pooled, valid_2 = self._masked_max_pool1d(z2, mask_2, pool_size) + overlap_pooled = valid_1 & valid_2 + + # Compute temporal contrastive loss at this scale + scale_loss, n_tokens = self._temporal_contrastive_loss( + z1_pooled, z2_pooled, overlap_pooled + ) + + if n_tokens > 0: + total_loss = total_loss + scale_loss + n_scales += 1 + total_overlap_tokens += n_tokens + + if n_scales > 0: + total_loss = total_loss / n_scales + else: + # No overlap at any scale — zero loss with grad connectivity + total_loss = z1.sum() * 0.0 + + with torch.no_grad(): + # Collapse monitoring (Wang & Isola 2020) — computed at scale 1 + # using overlap tokens only for consistency + collapse_metrics = self._compute_collapse_metrics(z1, z2, overlap) + + metrics = { + "ts2vec_loss": total_loss.detach(), + "ssl_loss": total_loss.detach(), + "ts2vec_n_scales_active": torch.tensor(n_scales), + "ts2vec_total_overlap_tokens": torch.tensor(total_overlap_tokens), + } + metrics.update(collapse_metrics) + + return total_loss, metrics + + def _temporal_contrastive_loss( + self, + z1: torch.Tensor, + z2: torch.Tensor, + overlap: torch.Tensor, + ) -> Tuple[torch.Tensor, int]: + """Compute temporal contrastive loss at a single scale. + + For each overlapping timestep, the two views' representations form a + positive pair. All other overlap tokens across the batch (both temporal + neighbors within the same sample and tokens from other samples) serve + as negatives in a standard NT-Xent formulation. + + Args: + z1: (B, T', proj_dim) projected representations from view 1 + z2: (B, T', proj_dim) projected representations from view 2 + overlap: (B, T') True = valid overlap position + + Returns: + (loss, n_overlap_tokens) + """ + temperature = self.config.temperature + + N = int(overlap.sum().item()) + if N < 2: + return torch.tensor(0.0, device=z1.device), 0 + + # Gather overlap tokens from both views + tokens_1 = z1[overlap] # (N, proj_dim) + tokens_2 = z2[overlap] # (N, proj_dim) + + # Re-normalize after potential max-pool + tokens_1 = F.normalize(tokens_1, dim=-1) + tokens_2 = F.normalize(tokens_2, dim=-1) + + # Standard NT-Xent on (2N, 2N) — combines temporal and cross-sample + z = torch.cat([tokens_1, tokens_2], dim=0) # (2N, proj_dim) + sim_matrix = torch.mm(z, z.t()) / temperature # (2N, 2N) + + # Positive pair labels: token i pairs with i+N + labels = torch.cat( + [ + torch.arange(N, 2 * N, device=z.device), + torch.arange(N, device=z.device), + ] + ) + + # Mask self-similarity + mask = torch.eye(2 * N, dtype=torch.bool, device=z.device) + sim_matrix = sim_matrix.masked_fill(mask, float("-inf")) + + loss = F.cross_entropy(sim_matrix, labels) + + return loss, N + + @staticmethod + def _compute_collapse_metrics( + z1: torch.Tensor, + z2: torch.Tensor, + overlap: torch.Tensor, + ) -> Dict[str, torch.Tensor]: + """Compute collapse monitoring metrics on scale-1 overlap tokens. + + Mirrors the standard contrastive objective's collapse monitoring + (Wang & Isola 2020, Roy & Vetterli 2007) for consistent analysis. + + Args: + z1: (B, T, proj_dim) from view 1 + z2: (B, T, proj_dim) from view 2 + overlap: (B, T) True = overlap position + + Returns: + Dict of collapse metrics. + """ + N = int(overlap.sum().item()) + if N < 2: + return { + "ts2vec_alignment": torch.tensor(0.0), + "ts2vec_uniformity": torch.tensor(0.0), + "ts2vec_effective_rank": torch.tensor(0.0), + } + + t1 = F.normalize(z1[overlap], dim=-1) # (N, proj_dim) + t2 = F.normalize(z2[overlap], dim=-1) + + # Alignment: mean squared L2 distance of positive pairs (lower = better) + alignment = (t1 - t2).norm(dim=-1).pow(2).mean() + + # Uniformity: log avg Gaussian potential on hypersphere (lower = better) + # Use view 1 tokens to avoid inflating with positives + sq_pdist = torch.cdist(t1, t1, p=2).pow(2) + uniformity = sq_pdist.mul(-2).exp().mean().log() + + # Effective rank via singular value entropy + centered = t1 - t1.mean(dim=0) + _, s, _ = torch.svd_lowrank(centered, q=min(N, t1.shape[1])) + p = s.clamp(min=0).pow(2) + p = p / p.sum().clamp(min=1e-12) + eff_rank = (-p * p.clamp(min=1e-7).log()).sum().exp() + + return { + "ts2vec_alignment": alignment, + "ts2vec_uniformity": uniformity, + "ts2vec_effective_rank": eff_rank, + } diff --git a/src/slices/training/README.md b/src/slices/training/README.md index 710faa3..25eed9a 100644 --- a/src/slices/training/README.md +++ b/src/slices/training/README.md @@ -1,430 +1,157 @@ # Training Module -This module contains Lightning modules for training SSL and supervised models. +This package contains the Lightning modules and helper utilities used by the +training scripts under `scripts/training/`. -## Components +## Main Components ### `SSLPretrainModule` -Lightning module for SSL pretraining. This module is fully agnostic to: -- **Encoder architecture**: Works with any encoder (transformer, RNN, CNN, etc.) -- **SSL objective**: Works with any SSL objective (MAE, contrastive, JEPA, etc.) +Builds: -All components are built from configuration using factory patterns. +- an encoder from `config.encoder` +- an SSL objective from `config.ssl` +- optimizer and scheduler state from `config.optimizer` and `config.scheduler` -#### Key Features +What it does: -- **Configuration-driven**: All models built from YAML configs via Hydra -- **Flexible optimizers**: Supports Adam, AdamW, SGD with configurable parameters -- **Learning rate schedulers**: Cosine annealing, step decay, plateau, warmup+cosine -- **Automatic logging**: Logs all metrics from SSL objectives to tensorboard/wandb -- **Checkpoint management**: Saves encoder weights separately for downstream tasks -- **Distributed training**: Works with multi-GPU via Lightning +- logs `train/loss` and `val/loss` +- logs objective-specific metrics under `train/*` and `val/*` +- tracks `train/gradient_steps` +- tracks `train/wall_clock_seconds` +- exposes `save_encoder()` for downstream reuse -#### Usage +The actual pretraining entrypoint is `scripts/training/pretrain.py`, which also +enforces SSL/encoder compatibility and saves both last-epoch and best-val +encoder checkpoints. -```python -from omegaconf import OmegaConf -from slices.training import SSLPretrainModule +Example: -# Load config (typically from Hydra) -config = OmegaConf.load("configs/pretrain.yaml") - -# Create module -module = SSLPretrainModule(config) - -# Train with Lightning Trainer -import lightning.pytorch as L -trainer = L.Trainer(max_epochs=100, devices=1) -trainer.fit(module, datamodule=datamodule) - -# Save encoder for downstream tasks -module.save_encoder("pretrained_encoder.pt") -``` - -#### Config Structure - -The module expects a config with three main sections: - -```yaml -encoder: - name: transformer # Encoder architecture name - d_input: 35 # Set automatically from data - d_model: 128 - # ... encoder-specific params - -ssl: - name: mae # SSL objective name - mask_ratio: 0.15 - # ... objective-specific params - -optimizer: - name: adamw - lr: 1.0e-3 - weight_decay: 0.01 - -scheduler: # Optional - name: warmup_cosine - warmup_epochs: 10 - max_epochs: 100 -``` - -#### Methods - -- **`forward(timeseries, mask)`**: Forward pass through SSL objective -- **`training_step(batch, batch_idx)`**: Training step (called by Lightning) -- **`validation_step(batch, batch_idx)`**: Validation step (called by Lightning) -- **`configure_optimizers()`**: Configure optimizer and optional scheduler -- **`get_encoder()`**: Get encoder module (e.g., for inspection) -- **`save_encoder(path)`**: Save encoder weights to file - -#### Logged Metrics - -The module automatically logs: -- `train/loss`: Training loss (per step and per epoch) -- `val/loss`: Validation loss (per epoch) -- `train/{metric}`: All metrics returned by SSL objective (per epoch) -- `val/{metric}`: All validation metrics (per epoch) -- Learning rate (when using scheduler) - -#### Supported Optimizers - -- **adam**: Classic Adam optimizer -- **adamw**: Adam with decoupled weight decay (recommended) -- **sgd**: Stochastic gradient descent with momentum - -#### Supported Schedulers - -- **cosine**: Cosine annealing learning rate decay -- **step**: Step decay (reduce LR every N epochs) -- **plateau**: Reduce LR when validation loss plateaus -- **warmup_cosine**: Linear warmup followed by cosine decay (recommended for SSL) - -#### Example: Different Optimizers - -```yaml -# AdamW (recommended for transformers) -optimizer: - name: adamw - lr: 1.0e-3 - weight_decay: 0.01 - -# Adam (no weight decay) -optimizer: - name: adam - lr: 1.0e-3 - weight_decay: 0.0 - -# SGD with momentum -optimizer: - name: sgd - lr: 1.0e-2 - weight_decay: 1.0e-4 - momentum: 0.9 -``` - -#### Example: Different Schedulers - -```yaml -# Warmup + Cosine (recommended for SSL) -scheduler: - name: warmup_cosine - warmup_epochs: 10 - max_epochs: 100 - eta_min: 1.0e-6 - -# Cosine annealing only -scheduler: - name: cosine - T_max: 100 - eta_min: 0.0 - -# Step decay -scheduler: - name: step - step_size: 30 - gamma: 0.1 - -# Reduce on plateau -scheduler: - name: plateau - mode: min - factor: 0.1 - patience: 10 +```bash +uv run python scripts/training/pretrain.py dataset=miiv ssl=mae +uv run python scripts/training/pretrain.py dataset=miiv ssl=jepa +uv run python scripts/training/pretrain.py dataset=miiv ssl=ts2vec ``` -## Adding New Components +### `FineTuneModule` -### Adding a New Encoder +Composes: -1. Create encoder class in `src/slices/models/encoders/`: +- an encoder +- an optional missing-token wrapper +- a task head +- metric collections for train, val, and test -```python -from .base import BaseEncoder, EncoderConfig +Supported task types: -@dataclass -class MyEncoderConfig(EncoderConfig): - # Add encoder-specific params - hidden_size: int = 256 +- `binary` +- `multiclass` +- `multilabel` +- `regression` -class MyEncoder(BaseEncoder): - def __init__(self, config: MyEncoderConfig): - super().__init__(config) - # ... build encoder +Checkpoint inputs: - def forward(self, x, mask=None, padding_mask=None): - # ... encode input - return encoded -``` +- encoder checkpoint from `encoder.pt` +- full Lightning pretrain checkpoint via `pretrain_checkpoint` +- no checkpoint for supervised-from-scratch runs -2. Register in `src/slices/models/encoders/factory.py`: +The benchmark matrix uses `encoder.pt`, the last encoder from the fixed +pretraining schedule, for all SSL downstream runs. `encoder_best_val.pt` is +saved for diagnostics but is not the thesis downstream checkpoint, because SSL +validation loss is not comparable as an early-stopping signal across objectives. -```python -from .my_encoder import MyEncoder, MyEncoderConfig +What it handles: -ENCODER_REGISTRY["my_encoder"] = MyEncoder -ENCODER_CONFIG_REGISTRY["my_encoder"] = MyEncoderConfig -``` +- linear probing via `training.freeze_encoder=true` +- full finetuning via `training.freeze_encoder=false` +- gradual unfreezing via `training.unfreeze_epoch` +- task-specific losses and metrics +- optional projection heads for dimension-controlled evaluation -3. Create config file `configs/encoder/my_encoder.yaml`: +The main entrypoints are: -```yaml -name: my_encoder -hidden_size: 256 -# ... other params -``` +- `scripts/training/finetune.py` +- `scripts/training/supervised.py` -4. Use in pretraining: +Examples: ```bash -uv run python scripts/pretrain.py encoder=my_encoder -``` - -### Adding a New SSL Objective - -1. Create objective class in `src/slices/models/pretraining/`: - -```python -from .base import BaseSSLObjective, SSLConfig - -@dataclass -class MySSLConfig(SSLConfig): - name: str = "my_ssl" - # Add objective-specific params - -class MySSLObjective(BaseSSLObjective): - def __init__(self, encoder, config: MySSLConfig): - super().__init__(encoder, config) - # ... build objective-specific components - - def forward(self, x, mask): - # ... compute SSL loss - return loss, metrics +uv run python scripts/training/finetune.py dataset=miiv checkpoint=outputs/.../encoder.pt +uv run python scripts/training/finetune.py dataset=miiv checkpoint=outputs/.../encoder.pt training.freeze_encoder=false +uv run python scripts/training/supervised.py dataset=miiv ``` -2. Register in `src/slices/models/pretraining/factory.py`: +## Configuration Layout -```python -from .my_ssl import MySSLObjective, MySSLConfig +The training scripts use Hydra config groups: -SSL_REGISTRY["my_ssl"] = MySSLObjective -CONFIG_REGISTRY["my_ssl"] = MySSLConfig -``` +- `configs/model/` for encoder architecture +- `configs/ssl/` for SSL objective selection +- `configs/tasks/` for downstream task heads and metrics +- `configs/pretrain.yaml`, `configs/finetune.yaml`, `configs/supervised.yaml`, and `configs/xgboost.yaml` for phase-specific orchestration -3. Create config file `configs/ssl/my_ssl.yaml`: +Note that custom encoders belong in `configs/model/`, not `configs/encoder/`. -```yaml -name: my_ssl -# ... ssl-specific params -``` +## Checkpoint Formats -4. Use in pretraining: +The downstream path supports two checkpoint styles: -```bash -uv run python scripts/pretrain.py ssl=my_ssl -``` +1. `encoder.pt` + - compact encoder export for reuse across downstream runs + - includes encoder config and optional learned missing token +2. full Lightning `.ckpt` + - used when reconstructing the entire pretraining module state -### `FineTuneModule` +`scripts/training/finetune.py` records which checkpoint source was actually used +for final test evaluation so post-run fairness evaluation can reuse the same +provenance. -Lightning module for downstream task finetuning. This module composes a pretrained encoder with a task head for supervised learning on clinical prediction tasks. +## Utilities -#### Key Features +`utils.py` contains shared helpers for: -- **Flexible checkpoint loading**: Supports loading from encoder weights (`.pt`) or full pretrain checkpoints (`.ckpt`) -- **Freezing strategies**: Linear probing (frozen encoder), full finetuning, or gradual unfreezing -- **Task support**: Binary classification, multiclass, multilabel, and regression -- **Comprehensive metrics**: AUROC, AUPRC, accuracy, F1 score for classification tasks -- **Class imbalance handling**: Configurable loss functions and class weights +- optimizer and scheduler construction +- callback setup +- W&B logger setup +- checkpoint export +- label-support validation +- fairness evaluation hooks +- data prerequisite checks -#### Usage +## Extending -```python -from omegaconf import OmegaConf -from slices.training import FineTuneModule +To add a new encoder: -# Load config (typically from Hydra) -config = OmegaConf.load("configs/finetune.yaml") +1. implement it in `src/slices/models/encoders/` +2. register it in `src/slices/models/encoders/factory.py` +3. add a Hydra config under `configs/model/` +4. add tests -# Create module with pretrained encoder -module = FineTuneModule( - config=config, - checkpoint_path="outputs/encoder.pt", # From pretraining -) +To add a new SSL objective: -# Train with Lightning Trainer -import lightning.pytorch as L -trainer = L.Trainer(max_epochs=50, devices=1) -trainer.fit(module, datamodule=datamodule) -``` +1. implement it in `src/slices/models/pretraining/` +2. register it in `src/slices/models/pretraining/factory.py` +3. add a Hydra config under `configs/ssl/` +4. add tests -#### Config Structure - -The module expects a config with encoder, task, and training sections: - -```yaml -encoder: - name: transformer - d_model: 128 - # ... encoder config (should match pretrained encoder) - -task: - task_name: mortality_24h - task_type: binary # binary | multiclass | multilabel | regression - head_type: mlp # mlp | linear - hidden_dims: [64] - dropout: 0.1 - -training: - freeze_encoder: true # Linear probing (default) - unfreeze_epoch: null # Set to N for gradual unfreezing - max_epochs: 50 - batch_size: 64 - -optimizer: - name: adamw - lr: 1.0e-4 # Lower LR for finetuning - weight_decay: 0.01 -``` - -#### Finetuning Strategies - -1. **Linear Probing** (default): Freeze encoder, train only task head - ```yaml - training: - freeze_encoder: true - ``` - -2. **Full Finetuning**: Train both encoder and task head - ```yaml - training: - freeze_encoder: false - ``` - -3. **Gradual Unfreezing**: Start frozen, unfreeze after N epochs - ```yaml - training: - freeze_encoder: true - unfreeze_epoch: 5 # Unfreeze at epoch 5 - ``` - -#### Task Types - -- **binary**: Binary classification (2 outputs for CrossEntropyLoss) -- **multiclass**: Multi-class classification (requires `n_classes`) -- **multilabel**: Multi-label classification (requires `n_classes`) -- **regression**: Continuous prediction (1 output, MSELoss) - -#### Logged Metrics - -For classification tasks: -- `train/loss`, `val/loss`, `test/loss`: Cross-entropy loss -- `train/accuracy`, `val/accuracy`, `test/accuracy`: Classification accuracy -- `val/auroc`, `test/auroc`: Area under ROC curve -- `val/auprc`, `test/auprc`: Area under PR curve -- `val/f1`, `test/f1`: F1 score - -For regression tasks: -- `train/loss`, `val/loss`, `test/loss`: MSE loss - -#### Methods - -- **`forward(timeseries, mask)`**: Forward pass through encoder and task head -- **`training_step(batch, batch_idx)`**: Training step (called by Lightning) -- **`validation_step(batch, batch_idx)`**: Validation step (called by Lightning) -- **`test_step(batch, batch_idx)`**: Test step (called by Lightning) -- **`on_train_epoch_start()`**: Called at start of each epoch (handles gradual unfreezing) -- **`configure_optimizers()`**: Configure optimizer and optional scheduler -- **`get_encoder()`**: Get encoder module -- **`get_task_head()`**: Get task head module - -#### Loading Checkpoints - -The module supports two checkpoint formats: - -1. **Encoder weights** (`.pt` file from `pretrain.py`): - ```python - module = FineTuneModule( - config=config, - checkpoint_path="outputs/encoder.pt", - ) - ``` - -2. **Full pretrain checkpoint** (`.ckpt` file from Lightning): - ```python - module = FineTuneModule( - config=config, - pretrain_checkpoint_path="outputs/ssl-last.ckpt", - ) - ``` - -#### Adding a New Task Head - -1. Create task head class in `src/slices/models/heads/`: - -```python -from .base import BaseTaskHead, TaskHeadConfig - -class MyTaskHead(BaseTaskHead): - def __init__(self, config: TaskHeadConfig): - super().__init__(config) - # ... build head architecture - - def forward(self, encoder_output): - # ... compute predictions - return {"logits": logits, "probs": probs} -``` - -2. Register in `src/slices/models/heads/factory.py`: - -```python -from .my_head import MyTaskHead +To add a new task head: -TASK_HEAD_REGISTRY["my_head"] = MyTaskHead -``` - -3. Use in finetuning: - -```yaml -task: - head_type: my_head - # ... other config -``` +1. implement it in `src/slices/models/heads/` +2. register it in `src/slices/models/heads/factory.py` +3. reference it from the relevant task config ## Testing -Test the modules with pytest: +Useful focused test targets: ```bash -# Test pretraining module uv run pytest tests/test_pretrain_module.py -v - -# Test finetuning module uv run pytest tests/test_finetune_module.py -v +uv run pytest tests/test_training_utils.py -v ``` -## See Also +## Related Docs -- [Pretraining Guide](../../../docs/PRETRAINING_GUIDE.md) - Full guide to using the pretraining pipeline -- [Encoder README](../models/encoders/README.md) - Details on encoder architectures -- [SSL Objectives README](../models/pretraining/README.md) - Details on SSL objectives -- [Task Heads README](../models/heads/README.md) - Details on task head architectures +- `src/slices/models/encoders/README.md` +- `src/slices/models/pretraining/README.md` +- `src/slices/eval/README.md` diff --git a/src/slices/training/checkpoint_loading.py b/src/slices/training/checkpoint_loading.py new file mode 100644 index 0000000..6893248 --- /dev/null +++ b/src/slices/training/checkpoint_loading.py @@ -0,0 +1,326 @@ +"""Checkpoint loading utilities for encoder weights. + +Handles loading pretrained encoder weights from various checkpoint formats +(.pt v1/v2/v3, .ckpt Lightning checkpoints), encoder type inference, +and conditional wrapping with EncoderWithMissingToken. +Extracted from FineTuneModule for modularity. +""" + +import logging +import warnings +from typing import Any, Dict, Optional + +import torch +import torch.nn as nn +from omegaconf import DictConfig + +from slices.models.encoders.factory import build_encoder +from slices.models.encoders.wrapper import EncoderWithMissingToken + +logger = logging.getLogger(__name__) + + +def infer_encoder_type(state_dict: Dict[str, Any]) -> Optional[str]: + """Infer encoder type from state_dict keys. + + Used for v2 checkpoints that don't include encoder_config. + + Args: + state_dict: Encoder state dictionary. + + Returns: + Inferred encoder name or None if unknown. + """ + keys = set(state_dict.keys()) + + # GRU-D encoder has gru_cell and decay parameters + if any("gru_cell" in k for k in keys) and any("W_gamma_x" in k for k in keys): + return "gru_d" + + # SMART encoder has distinctive keys + if any("embedder" in k or "blocks" in k or "seq_att" in k for k in keys): + return "smart" + + # Observation transformer has value_proj and feature_embed (not input_proj) + if any("value_proj" in k for k in keys) and any("feature_embed" in k for k in keys): + return "observation_transformer" + + # Standard transformer has input_proj and layers + if any("input_proj" in k or "layers" in k for k in keys): + return "transformer" + + return None + + +def wrap_encoder_with_missing_token( + encoder: nn.Module, + missing_token: Optional[torch.Tensor], +) -> nn.Module: + """Conditionally wrap encoder with EncoderWithMissingToken. + + Some encoders handle missingness intrinsically and should not be wrapped: + - ObservationTransformerEncoder: only tokenizes observed values + - SMARTEncoder: MLPEmbedder jointly embeds (value, mask) pairs + - TransformerEncoder with obs_aware=True: obs_proj handles missingness + + Args: + encoder: The encoder module to potentially wrap. + missing_token: Optional pretrained missing token. If None, + a random token will be initialized. + + Returns: + The encoder, possibly wrapped with EncoderWithMissingToken. + """ + handles_missingness = getattr(encoder, "handles_missingness_intrinsically", None) + if callable(handles_missingness) and handles_missingness(): + logger.info( + "Skipping EncoderWithMissingToken wrapper for %s " + "(handles missingness intrinsically)", + type(encoder).__name__, + ) + return encoder + + d_input = encoder.config.d_input + + if missing_token is not None: + encoder = EncoderWithMissingToken( + encoder=encoder, + d_input=d_input, + missing_token=missing_token, + init_missing_token=False, + ) + logger.info("Wrapped encoder with pretrained missing token") + else: + encoder = EncoderWithMissingToken( + encoder=encoder, + d_input=d_input, + missing_token=None, + init_missing_token=True, + ) + logger.info("Wrapped encoder with randomly initialized missing token") + + return encoder + + +def _apply_finetune_pooling_override( + encoder_config: dict, + config: DictConfig, +) -> dict: + """Return encoder config with a safe downstream pooling override applied.""" + encoder_config = dict(encoder_config) + finetune_pooling = config.encoder.get("pooling", "mean") + ckpt_pooling = encoder_config.get("pooling", "none") + if ckpt_pooling == finetune_pooling: + return encoder_config + + if "cls" in {ckpt_pooling, finetune_pooling}: + raise RuntimeError( + "Cannot override encoder pooling between checkpoint " + f"pooling='{ckpt_pooling}' and finetune pooling='{finetune_pooling}'. " + "CLS pooling adds learned parameters, so use a checkpoint trained with " + "pooling='cls' or choose a non-CLS downstream pooling mode." + ) + + encoder_config["pooling"] = finetune_pooling + logger.info( + "Overriding pooling: %s -> %s (finetuning requires aggregated output)", + ckpt_pooling, + finetune_pooling, + ) + return encoder_config + + +def load_encoder_weights( + encoder: nn.Module, + path: str, + config: DictConfig, + use_missing_token: bool, +) -> nn.Module: + """Load encoder weights from .pt file. + + Handles multiple checkpoint formats: + - v3+: Contains encoder_config for automatic architecture detection + - v2: Contains encoder_state_dict, missing_token but no config + - v1/legacy: Raw state_dict + + For v3+, the encoder is rebuilt from saved config, ensuring architecture + matches between pretraining and finetuning. + + Args: + encoder: The encoder module to load weights into (may be rebuilt). + path: Path to encoder checkpoint. + config: Full Hydra configuration (for pooling override). + use_missing_token: Whether to wrap with EncoderWithMissingToken. + + Returns: + The encoder with loaded weights, possibly wrapped with missing token. + + Raises: + FileNotFoundError: If checkpoint file doesn't exist. + RuntimeError: If state dict keys don't match encoder architecture. + """ + checkpoint = torch.load(path, map_location="cpu", weights_only=True) + + # Detect checkpoint format + if isinstance(checkpoint, dict) and "version" in checkpoint: + version = checkpoint["version"] + state_dict = checkpoint["encoder_state_dict"] + missing_token = checkpoint.get("missing_token", None) + + # Version 3+: Rebuild encoder from saved config + if version >= 3 and "encoder_config" in checkpoint: + encoder_config = dict(checkpoint["encoder_config"]) + encoder_name = encoder_config.pop("name") + encoder_config = _apply_finetune_pooling_override(encoder_config, config) + encoder = build_encoder(encoder_name, encoder_config) + logger.info( + "Rebuilt encoder from checkpoint config: %s (d_model=%s)", + encoder_name, + encoder_config.get("d_model", "N/A"), + ) + elif version == 2: + # v2 checkpoint: try to infer encoder type from state_dict keys + inferred_encoder = infer_encoder_type(state_dict) + config_encoder = config.encoder.name + + if inferred_encoder and inferred_encoder != config_encoder: + raise RuntimeError( + f"Encoder architecture mismatch!\n" + f" Checkpoint appears to be: {inferred_encoder}\n" + f" Config specifies: {config_encoder}\n\n" + f"Fix: Add 'encoder={inferred_encoder}' to your command:\n" + f" uv run python scripts/training/finetune.py \\\n" + f" checkpoint={path} \\\n" + f" encoder={inferred_encoder} \\\n" + f" task.task_name=...\n\n" + f"Or re-run pretraining to create a v3 checkpoint with embedded config." + ) + + encoder.load_state_dict(state_dict) + logger.info("Loaded encoder weights from: %s (format v%d)", path, version) + + # Wrap encoder with missing token if available and enabled + if use_missing_token: + encoder = wrap_encoder_with_missing_token(encoder, missing_token) + else: + # Old format (raw state_dict) + warnings.warn( + f"Loading old-format checkpoint from {path}. " + "Missing token will be randomly initialized. " + "Re-save the encoder using SSLPretrainModule.save_encoder() " + "to include the learned missing_token.", + UserWarning, + ) + encoder.load_state_dict(checkpoint) + logger.info("Loaded encoder weights from: %s (legacy format)", path) + + # Wrap encoder with random missing token if enabled + if use_missing_token: + encoder = wrap_encoder_with_missing_token(encoder, missing_token=None) + + return encoder + + +def load_from_pretrain_checkpoint( + encoder: nn.Module, + path: str, + config: DictConfig, + use_missing_token: bool, +) -> nn.Module: + """Load encoder from full pretraining checkpoint (.ckpt). + + Also extracts the missing_token from the SSL objective if present + and wraps the encoder with EncoderWithMissingToken. + + The encoder architecture is auto-detected from the checkpoint's + hyperparameters (saved by save_hyperparameters() during pretraining), + and the encoder is rebuilt with the correct architecture before + loading weights. + + Args: + encoder: The encoder module to load weights into (may be rebuilt). + path: Path to Lightning checkpoint. + config: Full Hydra configuration (for pooling override). + use_missing_token: Whether to wrap with EncoderWithMissingToken. + + Returns: + The encoder with loaded weights, possibly wrapped with missing token. + + Raises: + FileNotFoundError: If checkpoint file doesn't exist. + KeyError: If checkpoint doesn't contain 'state_dict'. + RuntimeError: If no encoder weights found in checkpoint. + """ + checkpoint = torch.load(path, map_location="cpu", weights_only=False) + + if "state_dict" not in checkpoint: + raise KeyError( + f"Checkpoint at {path} does not contain 'state_dict' key. " + "Is this a valid PyTorch Lightning checkpoint?" + ) + + state_dict = checkpoint["state_dict"] + + # Extract encoder weights (prefixed with "encoder.") + encoder_state_dict = {} + for key, value in state_dict.items(): + if key.startswith("encoder."): + encoder_state_dict[key[8:]] = value # Remove "encoder." prefix + + if not encoder_state_dict: + raise RuntimeError( + f"No encoder weights found in checkpoint {path}. " + "Expected keys prefixed with 'encoder.' in state_dict." + ) + + # Auto-detect encoder architecture from checkpoint hyperparameters + if "hyper_parameters" in checkpoint: + hyper_params = checkpoint["hyper_parameters"] + if "config" in hyper_params and "encoder" in hyper_params["config"]: + ckpt_encoder_cfg = hyper_params["config"]["encoder"] + ckpt_encoder_name = ckpt_encoder_cfg.get("name") + + # Always rebuild encoder from checkpoint config to ensure dimensions match. + # This handles both encoder name mismatches (e.g., transformer vs smart) + # and parameter mismatches (e.g., d_model=64 vs d_model=32). + encoder_config_dict = {k: v for k, v in ckpt_encoder_cfg.items() if k != "name"} + encoder_config_dict = _apply_finetune_pooling_override( + encoder_config_dict, + config, + ) + encoder = build_encoder(ckpt_encoder_name, encoder_config_dict) + logger.info( + "Rebuilt encoder from checkpoint config: %s (d_model=%s)", + ckpt_encoder_name, + encoder_config_dict.get("d_model", "N/A"), + ) + else: + # No hyper_parameters - infer encoder type from state_dict keys + inferred_encoder = infer_encoder_type(encoder_state_dict) + config_encoder = config.encoder.name + + if inferred_encoder and inferred_encoder != config_encoder: + raise RuntimeError( + f"Encoder architecture mismatch!\n" + f" Checkpoint appears to be: {inferred_encoder}\n" + f" Config specifies: {config_encoder}\n\n" + f"Fix: Add 'encoder={inferred_encoder}' to your command:\n" + f" uv run python scripts/training/finetune.py \\\n" + f" pretrain_checkpoint={path} \\\n" + f" encoder={inferred_encoder} \\\n" + f" task.task_name=...\n" + ) + + encoder.load_state_dict(encoder_state_dict) + logger.info("Loaded encoder from pretrain checkpoint: %s", path) + + # Extract missing_token from SSL objective if present + missing_token = None + if "ssl_objective.missing_token" in state_dict: + missing_token = state_dict["ssl_objective.missing_token"] + + # Wrap encoder with missing token if enabled + if use_missing_token: + encoder = wrap_encoder_with_missing_token(encoder, missing_token) + + return encoder diff --git a/src/slices/training/config_schemas.py b/src/slices/training/config_schemas.py index 0a88f59..4be294d 100644 --- a/src/slices/training/config_schemas.py +++ b/src/slices/training/config_schemas.py @@ -16,7 +16,7 @@ from typing import Any, List, Literal, Optional, Union -from pydantic import BaseModel, field_validator +from pydantic import BaseModel, Field, field_validator class TaskConfig(BaseModel): @@ -30,11 +30,22 @@ class TaskConfig(BaseModel): task_name: str task_type: str = "binary" + prediction_window_hours: Optional[int] = None + observation_window_hours: Optional[int] = None + gap_hours: int = 0 + label_sources: List[str] = Field(default_factory=list) + label_params: dict[str, Any] = Field(default_factory=dict) + quality_checks: dict[str, Any] = Field(default_factory=dict) + primary_metric: Optional[str] = None + additional_metrics: List[str] = Field(default_factory=list) head_type: str = "mlp" hidden_dims: List[int] = [64] dropout: float = 0.1 activation: str = "relu" n_classes: Optional[int] = None + class_names: Optional[List[str]] = None + positive_class: Optional[str] = None + supported_datasets: Optional[List[str]] = None use_layer_norm: bool = False projection_dim: Optional[int] = None @@ -86,6 +97,24 @@ class TrainingConfig(BaseModel): early_stopping_patience: Optional[int] = 10 early_stopping_monitor: Optional[str] = None early_stopping_mode: Optional[str] = None + label_smoothing: float = 0.0 + overfit_batches: Union[int, float] = 0 + allow_best_ckpt_fallback: bool = False + + +class PretrainTrainingConfig(BaseModel): + """Validated training configuration for SSL pretraining.""" + + model_config = {"extra": "forbid"} + + max_epochs: int = 100 + batch_size: int = 256 + accelerator: str = "auto" + devices: Any = "auto" + precision: Any = 32 + gradient_clip_val: Optional[float] = 1.0 + accumulate_grad_batches: int = 1 + early_stopping_patience: Optional[int] = None overfit_batches: Union[int, float] = 0 diff --git a/src/slices/training/finetune_module.py b/src/slices/training/finetune_module.py index f1e7aae..66e4f84 100644 --- a/src/slices/training/finetune_module.py +++ b/src/slices/training/finetune_module.py @@ -11,7 +11,6 @@ """ import logging -import warnings from typing import Any, Dict, Literal, Optional import lightning.pytorch as pl @@ -20,13 +19,13 @@ from omegaconf import DictConfig, OmegaConf from slices.eval import MetricConfig, build_metrics -from slices.models.encoders import ( - EncoderWithMissingToken, - ObservationTransformerEncoder, - SMARTEncoder, - build_encoder, -) +from slices.models.encoders import TransformerEncoder, build_encoder from slices.models.heads import TaskHeadConfig, build_task_head +from slices.training.checkpoint_loading import ( + load_encoder_weights, + load_from_pretrain_checkpoint, + wrap_encoder_with_missing_token, +) from slices.training.config_schemas import OptimizerConfig as OptimizerConfigSchema from slices.training.config_schemas import SchedulerConfig as SchedulerConfigSchema from slices.training.config_schemas import TaskConfig as TaskConfigSchema @@ -75,6 +74,11 @@ def __init__( super().__init__() self.save_hyperparameters() self.config = config + if checkpoint_path and pretrain_checkpoint_path: + raise ValueError( + "Provide exactly one checkpoint source: checkpoint_path for encoder.pt " + "or pretrain_checkpoint_path for a full Lightning pretrain .ckpt, not both." + ) # Validate task and training configs (catches typos via extra="forbid") task_dict = OmegaConf.to_container(config.task, resolve=True) @@ -100,12 +104,16 @@ def __init__( # Load pretrained weights (may wrap encoder with EncoderWithMissingToken) if checkpoint_path: - self._load_encoder_weights(checkpoint_path) + self.encoder = load_encoder_weights( + self.encoder, checkpoint_path, config, self.use_missing_token + ) elif pretrain_checkpoint_path: - self._load_from_pretrain_checkpoint(pretrain_checkpoint_path) + self.encoder = load_from_pretrain_checkpoint( + self.encoder, pretrain_checkpoint_path, config, self.use_missing_token + ) elif self.use_missing_token: # No pretrained weights, but still wrap encoder for consistency - self._wrap_encoder_with_missing_token(missing_token=None) + self.encoder = wrap_encoder_with_missing_token(self.encoder, missing_token=None) # Optional projection layer for dimensionality-controlled evaluation. # When comparing encoders with different output dims (e.g., 64 vs 1120), @@ -188,260 +196,6 @@ def _build_task_config(self, config: DictConfig) -> TaskHeadConfig: use_layer_norm=task_cfg.get("use_layer_norm", False), ) - def _load_encoder_weights(self, path: str) -> None: - """Load encoder weights from .pt file. - - Handles multiple checkpoint formats: - - v3+: Contains encoder_config for automatic architecture detection - - v2: Contains encoder_state_dict, missing_token but no config - - v1/legacy: Raw state_dict - - For v3+, the encoder is rebuilt from saved config, ensuring architecture - matches between pretraining and finetuning. - - Args: - path: Path to encoder checkpoint. - - Raises: - FileNotFoundError: If checkpoint file doesn't exist. - RuntimeError: If state dict keys don't match encoder architecture. - """ - checkpoint = torch.load(path, map_location="cpu", weights_only=True) - - # Detect checkpoint format - if isinstance(checkpoint, dict) and "version" in checkpoint: - version = checkpoint["version"] - state_dict = checkpoint["encoder_state_dict"] - missing_token = checkpoint.get("missing_token", None) - - # Version 3+: Rebuild encoder from saved config - if version >= 3 and "encoder_config" in checkpoint: - encoder_config = dict(checkpoint["encoder_config"]) - encoder_name = encoder_config.pop("name") - # Override pooling with finetuning config's value. - # SSL pretraining uses pooling='none' but finetuning needs - # pooling='mean' or 'query' to get a single representation. - # Pooling doesn't affect learned weights, just output aggregation. - finetune_pooling = self.config.encoder.get("pooling", "mean") - ckpt_pooling = encoder_config.get("pooling", "none") - if ckpt_pooling != finetune_pooling: - encoder_config["pooling"] = finetune_pooling - logger.info( - "Overriding pooling: %s -> %s (finetuning requires aggregated output)", - ckpt_pooling, - finetune_pooling, - ) - self.encoder = build_encoder(encoder_name, encoder_config) - logger.info( - "Rebuilt encoder from checkpoint config: %s (d_model=%s)", - encoder_name, - encoder_config.get("d_model", "N/A"), - ) - elif version == 2: - # v2 checkpoint: try to infer encoder type from state_dict keys - inferred_encoder = self._infer_encoder_type(state_dict) - config_encoder = self.config.encoder.name - - if inferred_encoder and inferred_encoder != config_encoder: - raise RuntimeError( - f"Encoder architecture mismatch!\n" - f" Checkpoint appears to be: {inferred_encoder}\n" - f" Config specifies: {config_encoder}\n\n" - f"Fix: Add 'encoder={inferred_encoder}' to your command:\n" - f" uv run python scripts/training/finetune.py \\\n" - f" checkpoint={path} \\\n" - f" encoder={inferred_encoder} \\\n" - f" task.task_name=...\n\n" - f"Or re-run pretraining to create a v3 checkpoint with embedded config." - ) - - self.encoder.load_state_dict(state_dict) - logger.info("Loaded encoder weights from: %s (format v%d)", path, version) - - # Wrap encoder with missing token if available and enabled - if self.use_missing_token: - self._wrap_encoder_with_missing_token(missing_token) - else: - # Old format (raw state_dict) - warnings.warn( - f"Loading old-format checkpoint from {path}. " - "Missing token will be randomly initialized. " - "Re-save the encoder using SSLPretrainModule.save_encoder() " - "to include the learned missing_token.", - UserWarning, - ) - self.encoder.load_state_dict(checkpoint) - logger.info("Loaded encoder weights from: %s (legacy format)", path) - - # Wrap encoder with random missing token if enabled - if self.use_missing_token: - self._wrap_encoder_with_missing_token(missing_token=None) - - def _load_from_pretrain_checkpoint(self, path: str) -> None: - """Load encoder from full pretraining checkpoint (.ckpt). - - Also extracts the missing_token from the SSL objective if present - and wraps the encoder with EncoderWithMissingToken. - - The encoder architecture is auto-detected from the checkpoint's - hyperparameters (saved by save_hyperparameters() during pretraining), - and the encoder is rebuilt with the correct architecture before - loading weights. - - Args: - path: Path to Lightning checkpoint. - - Raises: - FileNotFoundError: If checkpoint file doesn't exist. - KeyError: If checkpoint doesn't contain 'state_dict'. - RuntimeError: If no encoder weights found in checkpoint. - """ - checkpoint = torch.load(path, map_location="cpu", weights_only=False) - - if "state_dict" not in checkpoint: - raise KeyError( - f"Checkpoint at {path} does not contain 'state_dict' key. " - "Is this a valid PyTorch Lightning checkpoint?" - ) - - state_dict = checkpoint["state_dict"] - - # Extract encoder weights (prefixed with "encoder.") - encoder_state_dict = {} - for key, value in state_dict.items(): - if key.startswith("encoder."): - encoder_state_dict[key[8:]] = value # Remove "encoder." prefix - - if not encoder_state_dict: - raise RuntimeError( - f"No encoder weights found in checkpoint {path}. " - "Expected keys prefixed with 'encoder.' in state_dict." - ) - - # Auto-detect encoder architecture from checkpoint hyperparameters - if "hyper_parameters" in checkpoint: - hyper_params = checkpoint["hyper_parameters"] - if "config" in hyper_params and "encoder" in hyper_params["config"]: - ckpt_encoder_cfg = hyper_params["config"]["encoder"] - ckpt_encoder_name = ckpt_encoder_cfg.get("name") - - # Always rebuild encoder from checkpoint config to ensure dimensions match. - # This handles both encoder name mismatches (e.g., transformer vs smart) - # and parameter mismatches (e.g., d_model=64 vs d_model=32). - encoder_config_dict = {k: v for k, v in ckpt_encoder_cfg.items() if k != "name"} - # Override pooling with finetuning config's value. - # SSL pretraining uses pooling='none' but finetuning needs - # pooling='mean' or 'query' to get a single representation. - finetune_pooling = self.config.encoder.get("pooling", "mean") - ckpt_pooling = encoder_config_dict.get("pooling", "none") - if ckpt_pooling != finetune_pooling: - encoder_config_dict["pooling"] = finetune_pooling - logger.info( - "Overriding pooling: %s -> %s (finetuning requires aggregated output)", - ckpt_pooling, - finetune_pooling, - ) - self.encoder = build_encoder(ckpt_encoder_name, encoder_config_dict) - logger.info( - "Rebuilt encoder from checkpoint config: %s (d_model=%s)", - ckpt_encoder_name, - encoder_config_dict.get("d_model", "N/A"), - ) - else: - # No hyper_parameters - infer encoder type from state_dict keys - inferred_encoder = self._infer_encoder_type(encoder_state_dict) - config_encoder = self.config.encoder.name - - if inferred_encoder and inferred_encoder != config_encoder: - raise RuntimeError( - f"Encoder architecture mismatch!\n" - f" Checkpoint appears to be: {inferred_encoder}\n" - f" Config specifies: {config_encoder}\n\n" - f"Fix: Add 'encoder={inferred_encoder}' to your command:\n" - f" uv run python scripts/training/finetune.py \\\n" - f" pretrain_checkpoint={path} \\\n" - f" encoder={inferred_encoder} \\\n" - f" task.task_name=...\n" - ) - - self.encoder.load_state_dict(encoder_state_dict) - logger.info("Loaded encoder from pretrain checkpoint: %s", path) - - # Extract missing_token from SSL objective if present - missing_token = None - if "ssl_objective.missing_token" in state_dict: - missing_token = state_dict["ssl_objective.missing_token"] - - # Wrap encoder with missing token if enabled - if self.use_missing_token: - self._wrap_encoder_with_missing_token(missing_token) - - def _infer_encoder_type(self, state_dict: Dict[str, Any]) -> Optional[str]: - """Infer encoder type from state_dict keys. - - Used for v2 checkpoints that don't include encoder_config. - - Args: - state_dict: Encoder state dictionary. - - Returns: - Inferred encoder name or None if unknown. - """ - keys = set(state_dict.keys()) - - # SMART encoder has distinctive keys - if any("embedder" in k or "blocks" in k or "seq_att" in k for k in keys): - return "smart" - - # Observation transformer has value_proj and feature_embed (not input_proj) - if any("value_proj" in k for k in keys) and any("feature_embed" in k for k in keys): - return "observation_transformer" - - # Standard transformer has input_proj and layers - if any("input_proj" in k or "layers" in k for k in keys): - return "transformer" - - return None - - def _wrap_encoder_with_missing_token(self, missing_token: Optional[torch.Tensor]) -> None: - """Wrap encoder with EncoderWithMissingToken. - - Args: - missing_token: Optional pretrained missing token. If None, - a random token will be initialized. - """ - # Some encoders handle missingness intrinsically and should not be - # wrapped with EncoderWithMissingToken: - # - ObservationTransformerEncoder: only tokenizes observed values - # - SMARTEncoder: MLPEmbedder jointly embeds (value, mask) pairs, - # so it needs to see original values with the mask bit - if isinstance(self.encoder, (ObservationTransformerEncoder, SMARTEncoder)): - logger.info( - "Skipping EncoderWithMissingToken wrapper for %s " - "(handles missingness intrinsically)", - type(self.encoder).__name__, - ) - return - - d_input = self.encoder.config.d_input - - if missing_token is not None: - self.encoder = EncoderWithMissingToken( - encoder=self.encoder, - d_input=d_input, - missing_token=missing_token, - init_missing_token=False, - ) - logger.info("Wrapped encoder with pretrained missing token") - else: - self.encoder = EncoderWithMissingToken( - encoder=self.encoder, - d_input=d_input, - missing_token=None, - init_missing_token=True, - ) - logger.info("Wrapped encoder with randomly initialized missing token") - def _apply_freeze_strategy(self) -> None: """Apply freezing strategy to encoder.""" if self.freeze_strategy: @@ -470,7 +224,8 @@ def _get_criterion(self) -> nn.Module: Loss module. """ if self.task_type in ("binary", "multiclass"): - return nn.CrossEntropyLoss(weight=self._class_weights) + label_smoothing = self.config.training.get("label_smoothing", 0.0) + return nn.CrossEntropyLoss(weight=self._class_weights, label_smoothing=label_smoothing) elif self.task_type == "multilabel": return nn.BCEWithLogitsLoss() elif self.task_type == "regression": @@ -487,11 +242,13 @@ def _setup_metrics(self) -> None: eval_cfg = self.config.get("eval", {}) metrics_cfg = eval_cfg.get("metrics", {}) metric_names = metrics_cfg.get("names", None) + threshold = metrics_cfg.get("threshold", 0.5) metric_config = MetricConfig( task_type=self.task_type, n_classes=output_dim, metrics=metric_names, + threshold=threshold, ) # Build metrics for each stage @@ -514,7 +271,12 @@ def forward( Dictionary with 'logits' and 'probs'. """ # Encoder forward (produces pooled representation) - encoder_out = self.encoder(timeseries, mask=mask) # (B, d_model) + padding_mask = self._downstream_padding_mask(mask) + encoder_out = self.encoder( + timeseries, + mask=mask, + padding_mask=padding_mask, + ) # (B, d_model) # Optional projection to shared dimensionality if self.projection is not None: @@ -523,6 +285,14 @@ def forward( # Task head forward return self.task_head(encoder_out) + def _downstream_padding_mask(self, mask: torch.Tensor) -> Optional[torch.Tensor]: + """Exclude fully unobserved timesteps for obs-aware Transformer pooling.""" + if isinstance(self.encoder, TransformerEncoder) and getattr( + self.encoder.config, "obs_aware", False + ): + return mask.any(dim=-1) + return None + def _validate_labels(self, labels: torch.Tensor) -> None: """One-time check that labels are compatible with the task type. diff --git a/src/slices/training/pretrain_module.py b/src/slices/training/pretrain_module.py index 80ed825..f168a17 100644 --- a/src/slices/training/pretrain_module.py +++ b/src/slices/training/pretrain_module.py @@ -15,6 +15,7 @@ from slices.models.encoders import build_encoder from slices.models.pretraining import build_ssl_objective, get_ssl_config_class from slices.training.config_schemas import OptimizerConfig as OptimizerConfigSchema +from slices.training.config_schemas import PretrainTrainingConfig as PretrainTrainingConfigSchema from slices.training.config_schemas import SchedulerConfig as SchedulerConfigSchema from slices.training.utils import build_optimizer, build_scheduler, save_encoder_checkpoint @@ -77,6 +78,10 @@ def __init__( self.ssl_objective = build_ssl_objective(self.encoder, ssl_config) # Validate and store config for optimizer + if config.get("training") is not None: + training_dict = OmegaConf.to_container(config.training, resolve=True) + PretrainTrainingConfigSchema(**training_dict) + optimizer_dict = OmegaConf.to_container(config.optimizer, resolve=True) OptimizerConfigSchema(**optimizer_dict) @@ -175,13 +180,10 @@ def on_train_batch_end( """ # Update momentum encoder if the SSL objective supports it if hasattr(self.ssl_objective, "momentum_update"): - # Calculate training progress as fraction [0, 1] - if self.trainer.max_steps is not None and self.trainer.max_steps > 0: - progress = self.trainer.global_step / self.trainer.max_steps - else: - # Fallback: use epoch-based progress - max_epochs = self.trainer.max_epochs if self.trainer.max_epochs is not None else 1 - progress = self.trainer.current_epoch / max(1, max_epochs) + # Per-step linear progress [0, 1] — matches original SMART which + # updates momentum after every batch, not every epoch. + total_steps = self.trainer.estimated_stepping_batches + progress = min(self.trainer.global_step / max(1, total_steps), 1.0) self.ssl_objective.momentum_update(progress=progress) def validation_step( diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index b5576aa..9ada5f1 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -5,11 +5,14 @@ """ import math +from dataclasses import fields from pathlib import Path from typing import Any, Dict, Iterator, List, Optional, Union import torch import torch.nn as nn +import wandb +import yaml from lightning.pytorch.callbacks import ( EarlyStopping, LearningRateMonitor, @@ -18,6 +21,21 @@ from lightning.pytorch.loggers import WandbLogger from omegaconf import DictConfig, OmegaConf +from slices.constants import ( + FEATURE_BLOCKLIST, + LABEL_HORIZON_HOURS, + MIN_STAY_HOURS, + SEQ_LENGTH_HOURS, + canonical_downstream_protocol, + downstream_protocol_from_freeze, +) +from slices.data.labels import LabelBuilder, LabelBuilderFactory, LabelConfig + + +class WandbEntityNotFoundError(ValueError): + """Raised when an explicit W&B entity does not exist.""" + + # ============================================================================= # Optimizer / Scheduler # ============================================================================= @@ -160,12 +178,13 @@ def setup_pretrain_callbacks(cfg: DictConfig) -> list: checkpoint_callback = ModelCheckpoint( dirpath=cfg.get("checkpoint_dir", "checkpoints"), - filename="ssl-{epoch:03d}-{val_loss:.4f}", + filename="ssl-{epoch:03d}", monitor="val/loss", mode="min", save_top_k=3, save_last=True, verbose=True, + auto_insert_metric_name=False, ) callbacks.append(checkpoint_callback) @@ -187,29 +206,55 @@ def setup_pretrain_callbacks(cfg: DictConfig) -> list: def setup_finetune_callbacks(cfg: DictConfig, checkpoint_prefix: str = "finetune") -> list: """Set up training callbacks for finetuning / supervised training. - Task-type-aware: monitors val/auprc (max) for classification, - val/mse (min) for regression. + Defaults to the task's declared primary metric when available; otherwise + falls back to val/auprc for classification and val/mse for regression. """ callbacks = [] + lower_is_better = {"loss", "mse", "mae", "rmse", "brier_score", "ece"} + higher_is_better = { + "auroc", + "auprc", + "accuracy", + "f1", + "precision", + "recall", + "specificity", + "r2", + } + task_type = cfg.task.get("task_type", "binary") - if task_type == "regression": - default_monitor, default_mode = "val/mse", "min" + default_metric = cfg.task.get("primary_metric", None) + if default_metric: + default_monitor = default_metric if "/" in default_metric else f"val/{default_metric}" + elif task_type == "regression": + default_monitor = "val/mse" else: - default_monitor, default_mode = "val/auprc", "max" + default_monitor = "val/auprc" + monitor = cfg.training.get("early_stopping_monitor", default_monitor) - mode = cfg.training.get("early_stopping_mode", default_mode) + metric_name = monitor.split("/", 1)[-1] + if metric_name in lower_is_better: + default_mode = "min" + elif metric_name in higher_is_better: + default_mode = "max" + else: + raise ValueError( + f"Cannot infer checkpoint mode for finetune monitor '{monitor}'. " + "Set training.early_stopping_mode explicitly." + ) - metric_filename = monitor.replace("/", "_") + mode = cfg.training.get("early_stopping_mode", default_mode) checkpoint_callback = ModelCheckpoint( dirpath=cfg.get("checkpoint_dir", "checkpoints"), - filename=f"{checkpoint_prefix}-{{epoch:03d}}-{{{metric_filename}:.4f}}", + filename=f"{checkpoint_prefix}-{{epoch:03d}}", monitor=monitor, mode=mode, save_top_k=3, save_last=True, verbose=True, + auto_insert_metric_name=False, ) callbacks.append(checkpoint_callback) @@ -233,6 +278,70 @@ def setup_finetune_callbacks(cfg: DictConfig, checkpoint_prefix: str = "finetune # ============================================================================= +def _add_wandb_tag(tags: list[str], tag: str | None) -> None: + """Append a W&B tag once, ignoring unset values.""" + if not tag: + return + if tag not in tags: + tags.append(tag) + + +def _short_wandb_value(value: Any) -> str: + """Compact a numeric/string config value for stable W&B display names.""" + return str(value).replace(".", "").replace("-", "m") + + +def append_wandb_identity_suffixes(name: str | None, cfg: DictConfig) -> str | None: + """Append scientific identity fields omitted by the base Hydra run name.""" + if not name: + return name + + suffixes: list[str] = [] + experiment_subtype = cfg.get("experiment_subtype") + if experiment_subtype is not None: + suffixes.append(str(experiment_subtype)) + + source_dataset = cfg.get("source_dataset") + if source_dataset is not None: + suffixes.append(f"from_{source_dataset}") + + upstream_lr = cfg.get("upstream_pretrain_lr") + if upstream_lr is not None: + suffixes.append(f"uplr{_short_wandb_value(upstream_lr)}") + elif experiment_subtype == "lr_sensitivity": + optimizer_cfg = cfg.get("optimizer", {}) + lr = optimizer_cfg.get("lr") if optimizer_cfg else None + if lr is not None: + suffixes.append(f"lr{_short_wandb_value(lr)}") + + upstream_mask_ratio = cfg.get("upstream_pretrain_mask_ratio") + if upstream_mask_ratio is not None: + suffixes.append(f"upmr{_short_wandb_value(upstream_mask_ratio)}") + elif experiment_subtype in {"mask_ratio_sensitivity", "view_mask_sensitivity"}: + ssl_cfg = cfg.get("ssl", {}) + mask_ratio = ssl_cfg.get("mask_ratio") if ssl_cfg else None + if mask_ratio is not None: + suffixes.append(f"mr{_short_wandb_value(mask_ratio)}") + + if not suffixes: + return name + return f"{name}_{'_'.join(suffixes)}" + + +def train_label_support_summary(stats: Optional[dict[str, Any]]) -> dict[str, Any]: + """Format train-label support stats for W&B summary logging.""" + if not stats: + return {} + + summary: dict[str, Any] = {} + for key, value in stats.items(): + if value is None: + continue + if isinstance(value, (str, bool, int, float)): + summary[f"train_label_support/{key}"] = value + return summary + + def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: """Set up W&B experiment logger. @@ -242,21 +351,114 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: return None tags = list(cfg.logging.get("wandb_tags", [])) - if cfg.get("sprint") is not None: - tags.append(f"sprint:{cfg.sprint}") + if cfg.get("experiment_class") is not None: + _add_wandb_tag(tags, f"experiment_class:{cfg.experiment_class}") + if cfg.get("experiment_subtype") is not None: + _add_wandb_tag(tags, f"experiment_subtype:{cfg.experiment_subtype}") + if cfg.get("revision") is not None: + _add_wandb_tag(tags, f"revision:{cfg.revision}") + if cfg.get("rerun_reason") is not None: + tag = f"rerun-reason:{cfg.rerun_reason}" + if len(tag) > 64: + tag = tag[:61] + "..." + _add_wandb_tag(tags, tag) + if cfg.get("launch_commit") is not None: + _add_wandb_tag(tags, f"commit:{str(cfg.launch_commit)[:12]}") + if cfg.get("phase") is not None: + _add_wandb_tag(tags, f"phase:{cfg.phase}") + if cfg.get("dataset") is not None: + _add_wandb_tag(tags, f"dataset:{cfg.dataset}") + if cfg.get("paradigm") is not None: + _add_wandb_tag(tags, f"paradigm:{cfg.paradigm}") + elif cfg.get("ssl") and cfg.ssl.get("name") is not None: + _add_wandb_tag(tags, f"paradigm:{cfg.ssl.name}") + task_name = cfg.get("task", {}).get("task_name") + if task_name is not None: + _add_wandb_tag(tags, f"task:{task_name}") + if cfg.get("seed") is not None: + _add_wandb_tag(tags, f"seed:{cfg.seed}") + model_size = cfg.get("model_size") + if model_size is not None: + _add_wandb_tag(tags, f"model_size:{model_size}") + + # Add downstream protocol family tag. Supervised-from-scratch shares the + # full-finetune optimization budget; the phase tag distinguishes it from SSL. + freeze_encoder = cfg.get("training", {}).get("freeze_encoder") + if cfg.get("protocol") is not None: + protocol = canonical_downstream_protocol(cfg.protocol) + _add_wandb_tag(tags, f"protocol:{protocol}") + elif freeze_encoder is not None: + protocol = downstream_protocol_from_freeze(freeze_encoder) + _add_wandb_tag(tags, f"protocol:{protocol}") + + # Add mask_ratio tag for pretrain runs (useful for ablation filtering) + ssl_cfg = cfg.get("ssl", {}) + if ssl_cfg and ssl_cfg.get("mask_ratio") is not None: + _add_wandb_tag(tags, f"mask_ratio:{ssl_cfg.mask_ratio}") + + # Add label_fraction tag when subsampling training data + label_fraction = cfg.get("label_fraction") + if label_fraction is not None and label_fraction < 1.0: + _add_wandb_tag(tags, f"label_fraction:{label_fraction}") + _add_wandb_tag(tags, "ablation:label-efficiency") + + if cfg.get("source_dataset") is not None: + _add_wandb_tag(tags, f"source_dataset:{cfg.source_dataset}") + _add_wandb_tag(tags, "ablation:transfer") + if cfg.get("upstream_pretrain_lr") is not None: + _add_wandb_tag(tags, f"upstream_pretrain_lr:{cfg.upstream_pretrain_lr}") + if cfg.get("upstream_pretrain_mask_ratio") is not None: + _add_wandb_tag(tags, f"upstream_pretrain_mask_ratio:{cfg.upstream_pretrain_mask_ratio}") + tags = tags or None + # Adjust run name: use "probe" prefix instead of "finetune" for frozen encoder + run_name = cfg.logging.get("run_name", None) + if run_name and freeze_encoder is True: + run_name = run_name.replace("_finetune_", "_probe_", 1) + run_name = append_wandb_identity_suffixes(run_name, cfg) + if run_name and model_size is not None: + run_name += f"_{model_size}" + if run_name and label_fraction is not None and label_fraction < 1.0: + frac_str = str(label_fraction).replace(".", "") + run_name += f"_frac{frac_str}" + + # Adjust group to include protocol and label_fraction so that W&B "Group" view + # aggregates exactly the runs that differ only by seed. + group = cfg.logging.get("wandb_group", None) + if group: + if freeze_encoder is True: + group = group.replace("finetune_", "probe_", 1) + group = append_wandb_identity_suffixes(group, cfg) + if model_size is not None: + group += f"_{model_size}" + if label_fraction is not None and label_fraction < 1.0: + frac_str = str(label_fraction).replace(".", "") + group += f"_frac{frac_str}" + + wandb_entity = cfg.logging.get("wandb_entity", None) + logger = WandbLogger( project=cfg.logging.wandb_project, - entity=cfg.logging.get("wandb_entity", None), - name=cfg.logging.get("run_name", None), - group=cfg.logging.get("wandb_group", None), + entity=wandb_entity, + name=run_name, + group=group, tags=tags, save_dir=cfg.output_dir, log_model=False, ) - logger.experiment.config.update(OmegaConf.to_container(cfg, resolve=True)) + try: + logger.experiment.config.update(OmegaConf.to_container(cfg, resolve=True)) + except wandb.errors.CommError as exc: + message = str(exc) + if wandb_entity and "entity" in message and "not found" in message: + raise WandbEntityNotFoundError( + f"W&B entity '{wandb_entity}' could not be found. " + "Remove logging.wandb_entity to use the signed-in account, " + "or provide a valid W&B username/team." + ) from None + raise return logger @@ -284,30 +486,15 @@ def run_fairness_evaluation( print("Fairness Evaluation") print("=" * 80) - from slices.eval.fairness_evaluator import FairnessEvaluator + from slices.eval.fairness_evaluator import FairnessEvaluator, flatten_fairness_report + from slices.eval.inference import run_inference - model.eval() - all_preds, all_labels, all_stay_ids = [], [], [] - for batch in datamodule.test_dataloader(): - with torch.no_grad(): - outputs = model( - batch["timeseries"].to(model.device), - batch["mask"].to(model.device), - ) - probs = outputs["probs"] - if probs.dim() > 1 and probs.shape[1] == 2: - all_preds.append(probs[:, 1].cpu()) - else: - all_preds.append(probs.cpu()) - all_labels.append(batch["label"].cpu()) - all_stay_ids.extend( - batch["stay_id"].tolist() - if isinstance(batch["stay_id"], torch.Tensor) - else batch["stay_id"] - ) - - predictions = torch.cat(all_preds) - labels_tensor = torch.cat(all_labels) + predictions, labels_tensor, all_stay_ids = run_inference( + model, + datamodule.test_dataloader(), + device=model.device, + ) + task_type = cfg.get("task", {}).get("task_type", "binary") evaluator = FairnessEvaluator( static_df=datamodule.dataset.static_df, @@ -315,21 +502,15 @@ def run_fairness_evaluation( fairness_cfg.get("protected_attributes", ["gender", "age_group"]) ), min_subgroup_size=fairness_cfg.get("min_subgroup_size", 50), + task_type=task_type, + dataset_name=getattr(getattr(datamodule, "processed_dir", None), "name", None), ) fairness_report = evaluator.evaluate(predictions, labels_tensor, all_stay_ids) evaluator.print_report(fairness_report) if logger: - for attr, metrics in fairness_report.items(): - for metric_name, value in metrics.items(): - if isinstance(value, (int, float)): - logger.experiment.summary[f"fairness/{attr}/{metric_name}"] = value - elif isinstance(value, dict): - for sub_key, sub_val in value.items(): - if isinstance(sub_val, (int, float)): - logger.experiment.summary[ - f"fairness/{attr}/{metric_name}/{sub_key}" - ] = sub_val + for key, value in flatten_fairness_report(fairness_report).items(): + logger.experiment.summary[key] = value return fairness_report @@ -339,12 +520,199 @@ def run_fairness_evaluation( # ============================================================================= -def validate_data_prerequisites(processed_dir: str, dataset: str) -> None: +def report_and_validate_train_label_support( + *, + datamodule: Any, + task_name: str, + task_type: str, + dataset: str, + seed: int, + label_fraction: float, + min_train_positives: int = 3, +) -> Dict[str, Any]: + """Report effective train-set class support and fail on degenerate subsets. + + This is primarily aimed at label-efficiency runs where aggressive + subsampling can leave the optimization subset with too few positive + examples to produce meaningful metrics. + """ + if task_type != "binary": + return {} + + subset_stats = datamodule.get_train_label_statistics() + full_stats = datamodule.get_train_label_statistics(use_full_train=True) + + if task_name not in subset_stats: + return {} + + subset = subset_stats[task_name] + full = full_stats.get(task_name, subset) + + subset_total = int(subset.get("total", 0)) + subset_pos = int(subset.get("positive", 0)) + subset_neg = int(subset.get("negative", 0)) + subset_prev = float(subset.get("prevalence", 0.0)) + + full_total = int(full.get("total", 0)) + full_pos = int(full.get("positive", 0)) + full_neg = int(full.get("negative", 0)) + full_prev = float(full.get("prevalence", 0.0)) + + print("\n Train label support:") + print( + f" - Dataset / task / seed / fraction: {dataset} / {task_name} / {seed} / {label_fraction}" + ) + print( + f" - Full train split: {full_pos} positive, {full_neg} negative, " + f"{full_total} total ({full_prev * 100:.2f}% positive)" + ) + print( + f" - Optimization subset: {subset_pos} positive, {subset_neg} negative, " + f"{subset_total} total ({subset_prev * 100:.2f}% positive)" + ) + + if subset_pos == 0 or subset_neg == 0: + raise ValueError( + "Binary training subset lost a class: " + f"dataset={dataset}, task={task_name}, seed={seed}, " + f"label_fraction={label_fraction}, positives={subset_pos}, negatives={subset_neg}. " + "Adjust the fraction or seed before training." + ) + + if label_fraction < 1.0 and subset_pos < min_train_positives: + raise ValueError( + "Too few positive training examples in the label-efficiency subset: " + f"dataset={dataset}, task={task_name}, seed={seed}, " + f"label_fraction={label_fraction}, positives={subset_pos} < {min_train_positives}. " + "Increase the fraction or drop this run." + ) + + return { + "task_name": task_name, + "dataset": dataset, + "seed": seed, + "label_fraction": label_fraction, + "full_train_total": full_total, + "full_train_positive": full_pos, + "full_train_negative": full_neg, + "full_train_prevalence": full_prev, + "train_subset_total": subset_total, + "train_subset_positive": subset_pos, + "train_subset_negative": subset_neg, + "train_subset_prevalence": subset_prev, + } + + +def resolve_balanced_class_weights( + task_name: str, + task_type: str, + train_label_stats: Dict[str, Dict[str, Any]], +) -> Optional[List[float]]: + """Resolve ``class_weight='balanced'`` for supported task types. + + Only binary classification is supported. Regression returns ``None`` because + class weights do not apply. Other task types fail closed instead of + constructing incorrect weights. + """ + normalized_task_type = { + "binary_classification": "binary", + "multiclass_classification": "multiclass", + "multilabel_classification": "multilabel", + }.get(task_type, task_type) + + if normalized_task_type == "regression": + return None + + if normalized_task_type != "binary": + raise ValueError( + f"class_weight='balanced' is only supported for binary tasks, got " + f"task_type='{task_type}'. Set class_weight=null or provide explicit weights." + ) + + if task_name not in train_label_stats: + return None + + stats = train_label_stats[task_name] + n_pos = int(stats.get("positive", 0)) + n_neg = int(stats.get("negative", 0)) + n_total = n_pos + n_neg + + if n_pos == 0 or n_neg == 0: + raise ValueError( + f"Cannot compute balanced class weights for '{task_name}': " + f"{n_pos} positive, {n_neg} negative. Check label extraction." + ) + + raw = [n_total / (2 * n_neg), n_total / (2 * n_pos)] + return [w**0.5 for w in raw] + + +def _validate_benchmark_metadata(metadata: dict[str, Any], path: Path) -> None: + """Validate fixed benchmark invariants declared by a processed artifact.""" + declared_values = { + "seq_length_hours": (metadata.get("seq_length_hours"), SEQ_LENGTH_HOURS), + "input_seq_length_hours": (metadata.get("input_seq_length_hours"), SEQ_LENGTH_HOURS), + "min_stay_hours": (metadata.get("min_stay_hours"), MIN_STAY_HOURS), + "label_horizon_hours": (metadata.get("label_horizon_hours"), LABEL_HORIZON_HOURS), + } + extraction_config = metadata.get("extraction_config") or {} + if isinstance(extraction_config, dict): + declared_values["extraction_config.seq_length_hours"] = ( + extraction_config.get("seq_length_hours"), + SEQ_LENGTH_HOURS, + ) + declared_values["extraction_config.min_stay_hours"] = ( + extraction_config.get("min_stay_hours"), + MIN_STAY_HOURS, + ) + + for field_name, (observed, expected) in declared_values.items(): + if observed is None: + continue + try: + observed_int = int(observed) + except (TypeError, ValueError) as exc: + raise RuntimeError( + f"Processed data in {path} has non-integer benchmark invariant " + f"{field_name}: observed={observed!r}, expected={expected}." + ) from exc + if observed_int != expected: + raise RuntimeError( + f"Processed data in {path} violates benchmark invariant " + f"{field_name}: observed={observed}, expected={expected}." + ) + + feature_names = metadata.get("feature_names") or [] + blocked = sorted(set(feature_names) & set(FEATURE_BLOCKLIST)) + if blocked: + raise RuntimeError( + f"Processed data in {path} contains blocked leakage/future-derived " + f"feature(s): {blocked}." + ) + + +def validate_data_prerequisites( + processed_dir: str, + dataset: str, + task_names: Optional[List[str]] = None, + task_configs: Optional[List[Union[LabelConfig, DictConfig, Dict[str, Any]]]] = None, +) -> None: """Validate that required data files exist before training. + Checks file existence and, if task definitions are provided, validates the + label manifest in metadata.yaml to ensure labels were built with the + current builder version and task config. When metadata declares benchmark + window or feature fields, those declarations must match fixed invariants. + + When ``task_configs`` are supplied, they are treated as the source of truth + because they represent the active Hydra-composed task configuration for the + run. ``task_names`` remain as a fallback for callers that only know names. + Raises: FileNotFoundError: If required files are missing. + RuntimeError: If label manifest indicates stale labels. """ + path = Path(processed_dir) if not path.exists(): @@ -360,12 +728,119 @@ def validate_data_prerequisites(processed_dir: str, dataset: str) -> None: f"Run first: uv run python scripts/preprocessing/prepare_dataset.py dataset={dataset}" ) - stats_path = path / "normalization_stats.yaml" - if not stats_path.exists(): - raise FileNotFoundError( - f"normalization_stats.yaml not found in {path}\n" - f"Run first: uv run python scripts/preprocessing/prepare_dataset.py dataset={dataset}" - ) + metadata_path = path / "metadata.yaml" + metadata: dict[str, Any] | None = None + if metadata_path.exists(): + with open(metadata_path) as f: + metadata = yaml.safe_load(f) or {} + _validate_benchmark_metadata(metadata, path) + + label_config_fields = {field.name for field in fields(LabelConfig)} + + def coerce_label_config( + task_config: Union[LabelConfig, DictConfig, Dict[str, Any]], + ) -> LabelConfig: + if isinstance(task_config, LabelConfig): + return task_config + + if isinstance(task_config, DictConfig): + raw_config = OmegaConf.to_container(task_config, resolve=True) + else: + raw_config = dict(task_config) + + if not isinstance(raw_config, dict): + raise TypeError("Task configuration must resolve to a mapping.") + + label_config_dict = { + key: value for key, value in raw_config.items() if key in label_config_fields + } + return LabelConfig(**label_config_dict) + + def get_training_tasks_path() -> Path: + repo_root = Path(__file__).resolve().parents[3] + hydra_tasks_path = repo_root / "configs" / "tasks" + if hydra_tasks_path.exists(): + return hydra_tasks_path + return Path(__file__).resolve().parents[1] / "data" / "tasks" + + resolved_task_configs: List[LabelConfig] = [] + if task_configs is not None: + resolved_task_configs = [coerce_label_config(task_config) for task_config in task_configs] + if task_names is not None: + resolved_names = {task_config.task_name for task_config in resolved_task_configs} + requested_names = set(task_names) + if resolved_names != requested_names: + raise ValueError( + "task_names does not match task_configs: " + f"task_names={sorted(requested_names)}, " + f"task_configs={sorted(resolved_names)}" + ) + elif task_names: + tasks_path = get_training_tasks_path() + for task_name in task_names: + config_file = tasks_path / f"{task_name}.yaml" + if not config_file.exists(): + raise FileNotFoundError( + f"Task config not found for '{task_name}': {config_file}. " + "Cannot validate label freshness safely." + ) + + with open(config_file) as f: + config_dict = yaml.safe_load(f) + resolved_task_configs.append(coerce_label_config(config_dict)) + + # Validate label manifest if task configs are provided or can be resolved. + if resolved_task_configs: + if metadata is None: + raise FileNotFoundError( + f"metadata.yaml not found in {path} — cannot validate label freshness.\n" + f"Re-run extraction: uv run python scripts/preprocessing/" + f"extract_ricu.py dataset={dataset}" + ) + + label_manifest = metadata.get("label_manifest") + if label_manifest is None: + raise RuntimeError( + f"metadata.yaml in {path} has no label_manifest. " + "Labels were extracted before manifest support was added.\n" + f"Re-run extraction: uv run python scripts/preprocessing/" + f"extract_ricu.py dataset={dataset}" + ) + + for current_config in resolved_task_configs: + task_name = current_config.task_name + current_hash = LabelBuilder.config_hash(current_config) + + builder = LabelBuilderFactory.create(current_config) + current_version = builder.SEMANTIC_VERSION + + manifest_entry = label_manifest.get(task_name) + if manifest_entry is None: + raise RuntimeError( + f"Task '{task_name}' not found in label manifest. " + "Labels were extracted without this task.\n" + f"Re-run extraction: uv run python scripts/preprocessing/" + f"extract_ricu.py dataset={dataset}" + ) + + stored_version = manifest_entry.get("builder_version") + stored_hash = manifest_entry.get("config_hash") + + if stored_version != current_version: + raise RuntimeError( + f"Label builder version mismatch for task '{task_name}': " + f"stored={stored_version}, current={current_version}. " + f"Re-run extraction: uv run python scripts/preprocessing/" + f"extract_ricu.py dataset={dataset}" + ) + + if stored_hash != current_hash: + raise RuntimeError( + f"Task config changed for '{task_name}': " + f"stored_hash={stored_hash}, current_hash={current_hash}. " + f"Re-run extraction: uv run python scripts/preprocessing/" + f"extract_ricu.py dataset={dataset}" + ) # ============================================================================= diff --git a/tests/README.md b/tests/README.md index 3c5f913..8362c6a 100644 --- a/tests/README.md +++ b/tests/README.md @@ -1,158 +1,77 @@ # SLICES Test Suite -Comprehensive test suite for the SLICES project, covering data extraction, processing, and task building. +This directory contains the automated regression suite for data preparation, +task labeling, model components, training flows, evaluation, and export logic. -## Running Tests +## What The Suite Covers + +At a high level, the tests exercise: + +- preprocessing and extractor behavior +- dataset loading, caching, normalization, and split logic +- task-label builders for mortality, AKI, and LOS +- encoder implementations and factory registration +- SSL objectives including MAE, JEPA, contrastive, TS2Vec, and SMART +- finetuning and supervised training modules +- metrics, fairness evaluation, and statistical utilities +- experiment-integrity regressions and result export paths + +The suite size changes frequently. For the current inventory and collected test +count, use `pytest --collect-only` instead of relying on this README. + +## Common Commands + +Run the full suite: -### Run all tests ```bash uv run pytest tests/ -v ``` -### Run with coverage report +Run with coverage: + ```bash uv run pytest tests/ --cov=slices --cov-report=html --cov-report=term ``` -### Run specific test file -```bash -uv run pytest tests/test_data_io.py -v -``` +Inspect the current collected tests: -### Run specific test class ```bash -uv run pytest tests/test_data_io.py::TestCsvToParquetAll -v +uv run pytest --collect-only -q tests ``` -### Run specific test function -```bash -uv run pytest tests/test_data_io.py::TestCsvToParquetAll::test_successful_conversion -v -``` +Run a focused file: -### Run tests matching a pattern ```bash -uv run pytest tests/ -k "mortality" -v +uv run pytest tests/test_fairness_evaluator.py -v ``` -### Run with verbose output +Run a focused subset: + ```bash -uv run pytest tests/ -vv +uv run pytest tests/ -k "jepa or ts2vec" -v ``` -## Test Structure +## Testing Philosophy -| File | Purpose | Coverage | -|------|---------|----------| -| `conftest.py` | Shared fixtures and pytest configuration | - | -| `test_package.py` | Package structure, imports, and dependencies | Package-level validation | -| `test_extractor_config.py` | ExtractorConfig validation | Input validation, defaults | -| `test_base_extractor.py` | BaseExtractor abstract class | Core extraction logic, `run()` method, dense conversion | -| `test_data_io.py` | CSV-to-Parquet conversion | File I/O, data integrity | -| `test_dataset_datamodule.py` | ICUDataset and ICUDataModule | Data loading, imputation, normalization, patient-level splits | -| `test_extractor_integration.py` | Extractor + task system integration | Multi-task extraction, label computation | -| `test_task_builders.py` | Task label extraction | Mortality tasks, boundary conditions | -| `test_timeseries_extraction.py` | Time-series extraction pipeline | Hourly binning, feature mapping, edge cases | +The suite is biased toward behavior and regression protection rather than +snapshotting implementation details. -## Test Categories +Typical expectations for new work: -### Unit Tests -- `test_extractor_config.py` - Configuration validation -- `test_base_extractor.py` - Core extractor methods -- `test_data_io.py` - Data I/O utilities -- `test_task_builders.py` - Task builders - -### Integration Tests -- `test_extractor_integration.py` - Full extraction pipeline -- `test_dataset_datamodule.py` - Data loading pipeline - -### Edge Case Tests -- `test_timeseries_extraction.py::TestTimeSeriesEdgeCases` - Empty data, extreme values -- `test_task_builders.py::TestMortalityBoundaryConditions` - Boundary conditions +- add focused unit coverage for new logic +- add an integration or regression test when behavior spans modules +- prefer explicit failure modes for experiment-integrity issues +- keep tests deterministic across seeds and platforms where practical ## Writing New Tests -### Test File Naming -- Test files must start with `test_` or end with `_test.py` -- Test functions must start with `test_` -- Test classes must start with `Test` - -### Using Fixtures -Fixtures from `conftest.py` are automatically available: +General conventions: -```python -def test_something(sample_batch): - # sample_batch fixture is automatically injected - assert sample_batch["timeseries"].shape[0] == 4 -``` - -### Example Test Structure -```python -class TestMyFeature: - """Test suite for MyFeature.""" - - @pytest.fixture - def setup_data(self, tmp_path): - """Create test data.""" - # Setup code - return data - - def test_basic_functionality(self, setup_data): - """Test basic case.""" - result = my_function(setup_data) - assert result == expected - - def test_edge_case(self): - """Test edge case.""" - with pytest.raises(ValueError, match="expected error"): - my_function(invalid_input) - - def test_boundary_condition(self, setup_data): - """Test boundary condition.""" - result = my_function(setup_data, boundary_value) - assert result == expected_boundary_result -``` - -### Best Practices -1. **Test one thing per test** - Each test should verify a single behavior -2. **Use descriptive names** - Test names should describe what they test -3. **Include edge cases** - Test boundary conditions, empty inputs, extreme values -4. **Use fixtures** - Share setup code via fixtures -5. **Mock external dependencies** - Use `unittest.mock.patch` for external calls -6. **Test error handling** - Verify exceptions are raised correctly - -## Test Coverage - -To view test coverage in browser: -```bash -uv run pytest tests/ --cov=slices --cov-report=html -open htmlcov/index.html # macOS -``` +- test files start with `test_` +- test functions start with `test_` +- shared fixtures live in `conftest.py` +- behavior-specific assertions are preferred over broad smoke tests -**Target: >80% coverage on core modules** - -## Key Test Scenarios - -### Data Extraction -- Valid/invalid parquet paths -- Feature mapping loading -- Hourly binning and aggregation -- Negative hour filtering -- Multiple itemids per feature - -### Data Processing -- Imputation strategies (forward_fill, zero, mean, none) -- Normalization -- Dense timeseries conversion -- Observation mask handling - -### Task Building -- Mortality prediction windows (24h, 48h, hospital, ICU) -- Boundary conditions (exact boundaries) -- Empty data handling -- Multiple tasks extraction - -### Data Loading -- Patient-level splits (no leakage) -- Reproducible splits with seeds -- Batch collation -- DataLoader configuration +When a change affects experiment correctness, add a regression test even if unit +coverage already exists elsewhere. This repo depends heavily on config-driven +orchestration, so the failure mode often matters as much as the happy path. diff --git a/tests/conftest.py b/tests/conftest.py index 88dc8bf..97a9b77 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,6 +10,8 @@ import pytest import torch +from slices.constants import SEQ_LENGTH_HOURS + @pytest.fixture def sample_batch() -> dict: @@ -19,7 +21,7 @@ def sample_batch() -> dict: Dictionary with sample batch data including timeseries and labels. """ batch_size = 4 - seq_length = 48 + seq_length = SEQ_LENGTH_HOURS n_features = 35 return { @@ -40,7 +42,7 @@ def sample_config() -> dict: return { "parquet_root": "/tmp/test_data", "output_dir": "/tmp/test_output", - "seq_length_hours": 48, + "seq_length_hours": SEQ_LENGTH_HOURS, "feature_set": "core", } diff --git a/tests/test_aki_builder.py b/tests/test_aki_builder.py index 185c550..4be7013 100644 --- a/tests/test_aki_builder.py +++ b/tests/test_aki_builder.py @@ -2,21 +2,23 @@ import polars as pl import pytest + from slices.data.labels import AKILabelBuilder, LabelBuilderFactory, LabelConfig class TestAKILabelBuilder: """Tests for AKILabelBuilder.""" - def _make_config(self, obs_hours: int = 48) -> LabelConfig: + def _make_config(self, obs_hours: int = 48, prediction_hours: int = 24) -> LabelConfig: return LabelConfig( task_name="aki_kdigo", task_type="binary", observation_window_hours=obs_hours, + prediction_window_hours=prediction_hours, label_sources=["stays", "timeseries"], label_params={ "creatinine_col": "crea", - "baseline_window_hours": 24, + "baseline_window_hours": 48, "absolute_rise_threshold": 0.3, "relative_rise_threshold": 1.5, }, @@ -37,65 +39,69 @@ def _make_timeseries( ) def test_stable_creatinine_no_aki(self): - """Stable creatinine trajectory -> label=0.""" + """Stable creatinine in prediction window (48-72) -> label=0.""" stays = self._make_stays([1]) - # Baseline creatinine ~1.0, stable after obs window + # Baseline creatinine ~1.0 in hours 0-48, stable in prediction window 48-72 ts = self._make_timeseries( stay_id=1, - hours=[0, 6, 12, 18, 24, 30, 48, 54, 60, 72], - crea_values=[1.0, 1.0, 1.1, 1.0, 1.05, 1.0, 1.1, 1.05, 1.1, 1.0], + hours=[0, 6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 66], + crea_values=[1.0, 1.0, 1.1, 1.0, 1.05, 1.0, 1.0, 1.05, 1.1, 1.05, 1.1, 1.0], ) - config = self._make_config(obs_hours=48) + config = self._make_config() builder = AKILabelBuilder(config) labels = builder.build_labels({"stays": stays, "timeseries": ts}) assert labels.filter(pl.col("stay_id") == 1)["label"][0] == 0 def test_absolute_criterion_triggers_aki(self): - """Creatinine rise > 0.3 mg/dL within 48h after obs -> label=1.""" + """Creatinine rise > 0.3 mg/dL within 48h in prediction window -> label=1.""" stays = self._make_stays([1]) + # Baseline in 0-48, rise in prediction window 48-72 ts = self._make_timeseries( stay_id=1, - hours=[0, 12, 48, 54, 60], - crea_values=[1.0, 1.0, 1.0, 1.1, 1.4], # 1.4 - 1.0 = 0.4 >= 0.3 within 12h + hours=[0, 12, 36, 48, 54, 60], + crea_values=[1.0, 1.0, 1.0, 1.0, 1.1, 1.4], # 1.4 - 1.0 = 0.4 >= 0.3 within 12h ) - config = self._make_config(obs_hours=48) + config = self._make_config() builder = AKILabelBuilder(config) labels = builder.build_labels({"stays": stays, "timeseries": ts}) assert labels.filter(pl.col("stay_id") == 1)["label"][0] == 1 def test_relative_criterion_triggers_aki(self): - """Creatinine >= 1.5x baseline after obs -> label=1.""" + """Creatinine >= 1.5x baseline in prediction window -> label=1.""" stays = self._make_stays([1]) - # Baseline min in first 24h: 1.0 - # Post-obs value: 1.6 >= 1.5 * 1.0 = 1.5 -> AKI + # Baseline min in first 48h: 1.0 + # Prediction window value: 1.6 >= 1.5 * 1.0 = 1.5 -> AKI ts = self._make_timeseries( stay_id=1, - hours=[0, 12, 48, 60], - crea_values=[1.0, 1.1, 0.9, 1.6], + hours=[0, 12, 36, 48, 60], + crea_values=[1.0, 1.1, 0.9, 0.9, 1.6], ) - config = self._make_config(obs_hours=48) + config = self._make_config() builder = AKILabelBuilder(config) labels = builder.build_labels({"stays": stays, "timeseries": ts}) assert labels.filter(pl.col("stay_id") == 1)["label"][0] == 1 def test_aki_during_obs_only_no_aki_after(self): - """AKI criteria met only during observation period -> label=0.""" + """AKI criteria met only during observation window (0-48) -> label=0. + + Creatinine spike is in the baseline/observation period, not the prediction window. + """ stays = self._make_stays([1]) - # Spike during obs (hour 24-30), but stable after obs (hour 48+) - # Baseline min in first 24h: 1.0 + # Spike during obs (hour 24-30), but stable in prediction window (hour 48+) + # Baseline min in first 48h: 1.0 ts = self._make_timeseries( stay_id=1, - hours=[0, 12, 24, 30, 48, 60, 72], - crea_values=[1.0, 1.0, 1.5, 1.0, 1.0, 1.1, 1.0], + hours=[0, 12, 24, 30, 36, 48, 54, 60, 66], + crea_values=[1.0, 1.0, 1.5, 1.0, 1.0, 1.0, 1.1, 1.0, 1.05], ) - config = self._make_config(obs_hours=48) + config = self._make_config() builder = AKILabelBuilder(config) labels = builder.build_labels({"stays": stays, "timeseries": ts}) @@ -113,7 +119,7 @@ def test_no_creatinine_measurements_null_label(self): } ) - config = self._make_config(obs_hours=48) + config = self._make_config() builder = AKILabelBuilder(config) labels = builder.build_labels({"stays": stays, "timeseries": ts}) @@ -130,23 +136,23 @@ def test_all_creatinine_null_null_label(self): } ).cast({"crea": pl.Float64}) - config = self._make_config(obs_hours=48) + config = self._make_config() builder = AKILabelBuilder(config) labels = builder.build_labels({"stays": stays, "timeseries": ts}) assert labels.filter(pl.col("stay_id") == 1)["label"][0] is None - def test_no_baseline_no_measurements_in_first_24h(self): - """No creatinine in first 24h (no baseline) -> label=null.""" + def test_no_baseline_no_measurements_in_first_48h(self): + """No creatinine in first 48h (no baseline) -> label=null.""" stays = self._make_stays([1]) - # First creatinine at hour 30 (after baseline window) + # First creatinine at hour 50 (after baseline window) ts = self._make_timeseries( stay_id=1, - hours=[30, 48, 60], + hours=[50, 60, 66], crea_values=[1.0, 1.5, 2.0], ) - config = self._make_config(obs_hours=48) + config = self._make_config() builder = AKILabelBuilder(config) labels = builder.build_labels({"stays": stays, "timeseries": ts}) @@ -156,25 +162,25 @@ def test_both_criteria_independently_trigger(self): """Both absolute and relative criteria should independently trigger AKI.""" stays = self._make_stays([1, 2]) - # Stay 1: only absolute criterion (baseline=1.0, rise of 0.35 in 48h) + # Stay 1: only absolute criterion (baseline=1.0, rise of 0.35 in prediction window) ts1 = self._make_timeseries( stay_id=1, - hours=[0, 12, 48, 60], - crea_values=[1.0, 1.0, 1.0, 1.35], + hours=[0, 12, 36, 48, 60], + crea_values=[1.0, 1.0, 1.0, 1.0, 1.35], ) - # Stay 2: only relative criterion (baseline=0.6, post-obs=0.91 >= 0.9) + # Stay 2: only relative criterion (baseline=0.6, prediction window=0.91 >= 0.9) # 0.91 >= 1.5 * 0.6 = 0.9 -> AKI (relative) # But 0.91 - 0.7 = 0.21 < 0.3 (absolute not met) ts2 = self._make_timeseries( stay_id=2, - hours=[0, 12, 48, 60], - crea_values=[0.6, 0.7, 0.7, 0.91], + hours=[0, 12, 36, 48, 60], + crea_values=[0.6, 0.7, 0.7, 0.7, 0.91], ) ts = pl.concat([ts1, ts2]) - config = self._make_config(obs_hours=48) + config = self._make_config() builder = AKILabelBuilder(config) labels = builder.build_labels({"stays": stays, "timeseries": ts}) @@ -182,20 +188,20 @@ def test_both_criteria_independently_trigger(self): assert labels_dict[1] == 1 # Absolute criterion assert labels_dict[2] == 1 # Relative criterion - def test_no_post_obs_creatinine_label_zero(self): - """Creatinine only during obs window, none after -> label=0.""" + def test_no_prediction_window_creatinine_null_label(self): + """Creatinine only in observation window (0-48), none in prediction window -> label=null.""" stays = self._make_stays([1]) ts = self._make_timeseries( stay_id=1, - hours=[0, 12, 24], - crea_values=[1.0, 1.5, 2.0], + hours=[0, 12, 24, 36], + crea_values=[1.0, 1.5, 2.0, 1.8], ) - config = self._make_config(obs_hours=48) + config = self._make_config() builder = AKILabelBuilder(config) labels = builder.build_labels({"stays": stays, "timeseries": ts}) - assert labels.filter(pl.col("stay_id") == 1)["label"][0] == 0 + assert labels.filter(pl.col("stay_id") == 1)["label"][0] is None def test_empty_stays(self): """Empty stays returns empty DataFrame with correct schema.""" @@ -220,10 +226,10 @@ def test_multiple_stays_mixed_outcomes(self): """Multiple stays with different outcomes.""" stays = self._make_stays([1, 2, 3]) - # Stay 1: AKI (relative criterion) - ts1 = self._make_timeseries(1, [0, 12, 48, 60], [0.8, 0.9, 0.9, 1.5]) - # Stay 2: No AKI (stable) - ts2 = self._make_timeseries(2, [0, 12, 48, 60], [1.0, 1.0, 1.0, 1.1]) + # Stay 1: AKI (relative criterion in prediction window) + ts1 = self._make_timeseries(1, [0, 12, 36, 48, 60], [0.8, 0.9, 0.9, 0.9, 1.5]) + # Stay 2: No AKI (stable in prediction window) + ts2 = self._make_timeseries(2, [0, 12, 36, 48, 60], [1.0, 1.0, 1.0, 1.0, 1.1]) # Stay 3: No creatinine data -> null ts3 = pl.DataFrame( { @@ -250,16 +256,32 @@ def test_validate_inputs_missing_source(self): with pytest.raises(ValueError, match="missing"): builder.build_labels({"stays": self._make_stays([1])}) + def test_missing_prediction_window_raises(self): + """AKI builder must raise if prediction_window_hours is not set.""" + config = LabelConfig( + task_name="aki_kdigo", + task_type="binary", + observation_window_hours=48, + prediction_window_hours=None, + label_sources=["stays", "timeseries"], + label_params={"creatinine_col": "crea", "baseline_window_hours": 48}, + ) + stays = self._make_stays([1]) + ts = self._make_timeseries(1, [0, 12, 48, 60], [1.0, 1.0, 1.0, 1.5]) + builder = AKILabelBuilder(config) + with pytest.raises(ValueError, match="prediction_window_hours"): + builder.build_labels({"stays": stays, "timeseries": ts}) + def test_absolute_rise_exactly_at_threshold(self): """Rise of exactly 0.3 should trigger AKI (>= threshold).""" stays = self._make_stays([1]) ts = self._make_timeseries( stay_id=1, - hours=[0, 12, 48, 54], - crea_values=[1.0, 1.0, 1.0, 1.3], # Rise = exactly 0.3 + hours=[0, 12, 36, 48, 54], + crea_values=[1.0, 1.0, 1.0, 1.0, 1.3], # Rise = exactly 0.3 ) - config = self._make_config(obs_hours=48) + config = self._make_config() builder = AKILabelBuilder(config) labels = builder.build_labels({"stays": stays, "timeseries": ts}) @@ -271,16 +293,31 @@ def test_relative_rise_exactly_at_threshold(self): # Baseline = 1.0, threshold = 1.5 * 1.0 = 1.5 ts = self._make_timeseries( stay_id=1, - hours=[0, 12, 48, 60], - crea_values=[1.0, 1.0, 1.0, 1.5], + hours=[0, 12, 36, 48, 60], + crea_values=[1.0, 1.0, 1.0, 1.0, 1.5], ) - config = self._make_config(obs_hours=48) + config = self._make_config() builder = AKILabelBuilder(config) labels = builder.build_labels({"stays": stays, "timeseries": ts}) assert labels.filter(pl.col("stay_id") == 1)["label"][0] == 1 + def test_short_stay_only_baseline_data_null_label(self): + """Creatinine only in hours 0-48 (baseline), none in 48-72 -> label=null.""" + stays = self._make_stays([1]) + ts = self._make_timeseries( + stay_id=1, + hours=[0, 6, 12, 24, 36, 42], + crea_values=[1.0, 1.1, 1.0, 1.2, 1.0, 1.1], + ) + + config = self._make_config() + builder = AKILabelBuilder(config) + labels = builder.build_labels({"stays": stays, "timeseries": ts}) + + assert labels.filter(pl.col("stay_id") == 1)["label"][0] is None + class TestAKIFactory: """Tests for factory integration with AKILabelBuilder.""" diff --git a/tests/test_analyze_labels.py b/tests/test_analyze_labels.py new file mode 100644 index 0000000..a7bcf03 --- /dev/null +++ b/tests/test_analyze_labels.py @@ -0,0 +1,60 @@ +"""Tests for label analysis script helpers.""" + +import polars as pl + +from scripts.analyze_labels import analyze_task + + +def test_analyze_task_warns_when_missing_rate_exceeds_quality_threshold(): + labels_df = pl.DataFrame( + { + "stay_id": [1, 2, 3, 4, 5], + "aki_kdigo": [1, 0, None, None, None], + } + ) + + stats = analyze_task( + labels_df, + "aki_kdigo", + quality_checks={"max_missing_percentage": 50.0}, + ) + + assert stats["missing_percentage"] == 60.0 + assert len(stats["quality_warnings"]) == 1 + assert "60.0%" in stats["quality_warnings"][0] + assert "50.0%" in stats["quality_warnings"][0] + + +def test_analyze_task_has_no_quality_warning_below_threshold(): + labels_df = pl.DataFrame( + { + "stay_id": [1, 2, 3, 4, 5], + "aki_kdigo": [1, 0, 0, None, 1], + } + ) + + stats = analyze_task( + labels_df, + "aki_kdigo", + quality_checks={"max_missing_percentage": 28.0}, + ) + + assert stats["missing_percentage"] == 20.0 + assert stats["quality_warnings"] == [] + + +def test_analyze_task_summarizes_regression_labels_without_class_casting(): + labels_df = pl.DataFrame( + { + "stay_id": [1, 2, 3, 4], + "los_remaining": [0.5, 2.0, None, 7.5], + } + ) + + stats = analyze_task(labels_df, "los_remaining", task_type="regression") + + assert stats["task_type"] == "regression" + assert stats["n_classes"] is None + assert stats["class_distribution"] == {} + assert stats["label_summary"]["mean"] == 3.3333 + assert stats["label_summary"]["median"] == 2.0 diff --git a/tests/test_base_extractor.py b/tests/test_base_extractor.py index 3392a11..793437d 100644 --- a/tests/test_base_extractor.py +++ b/tests/test_base_extractor.py @@ -8,6 +8,7 @@ import polars as pl import pytest import yaml + from slices.data.extractors.base import BaseExtractor, ExtractorConfig @@ -699,8 +700,8 @@ def test_load_valid_task_config(self, tmp_path): assert loaded[0].task_type == "binary_classification" assert loaded[0].prediction_window_hours == 24 - def test_load_missing_task_config_skipped(self, tmp_path): - """Test that missing task configs are skipped with warning.""" + def test_load_missing_task_config_raises(self, tmp_path): + """Missing task configs should fail closed.""" parquet_root = tmp_path / "parquet" tasks_dir = tmp_path / "tasks" parquet_root.mkdir(parents=True) @@ -712,9 +713,8 @@ def test_load_missing_task_config_skipped(self, tmp_path): ) extractor = MockExtractor(config) - loaded = extractor._load_task_configs(["nonexistent_task"]) - - assert len(loaded) == 0 + with pytest.raises(FileNotFoundError, match="Task config not found"): + extractor._load_task_configs(["nonexistent_task"]) def test_load_multiple_task_configs(self, tmp_path): """Test loading multiple task configs at once.""" @@ -831,6 +831,39 @@ def test_observation_window_validation_skipped_when_none(self, tmp_path): assert loaded[0].observation_window_hours is None +class TestValidationFailures: + """Tests for fail-closed validation behavior.""" + + def test_validate_stays_raises_on_invalid_los(self, temp_parquet_structure): + config = ExtractorConfig(parquet_root=str(temp_parquet_structure)) + extractor = MockExtractor(config) + stays = extractor.extract_stays().with_columns( + pl.when(pl.col("stay_id") == 2) + .then(None) + .otherwise(pl.col("los_days")) + .alias("los_days") + ) + + with pytest.raises(ValueError, match="invalid LOS"): + extractor._validate_stays(stays) + + def test_validate_timeseries_raises_on_missing_stays(self, temp_parquet_structure): + config = ExtractorConfig(parquet_root=str(temp_parquet_structure)) + extractor = MockExtractor(config) + timeseries = extractor.extract_timeseries([1, 2]) + + with pytest.raises(ValueError, match="no timeseries data"): + extractor._validate_timeseries(timeseries, [1, 2, 3], ["heart_rate", "sbp"]) + + def test_validate_labels_raises_on_missing_stays(self, temp_parquet_structure): + config = ExtractorConfig(parquet_root=str(temp_parquet_structure)) + extractor = MockExtractor(config) + labels = pl.DataFrame({"stay_id": [1, 2], "mortality_24h": [0, 1]}) + + with pytest.raises(ValueError, match="no labels"): + extractor._validate_labels(labels, [1, 2, 3]) + + class TestPathResolution: """Tests for path resolution methods.""" @@ -968,6 +1001,7 @@ def test_resume_extraction_with_atomic_writes(self, tmp_path): config = ExtractorConfig( parquet_root=str(parquet_root), output_dir=str(output_dir), + tasks=["mortality_hospital"], ) extractor = MockExtractor(config) diff --git a/tests/test_config_schemas.py b/tests/test_config_schemas.py index a5c455e..5976559 100644 --- a/tests/test_config_schemas.py +++ b/tests/test_config_schemas.py @@ -9,8 +9,12 @@ import pytest import yaml from pydantic import ValidationError + +from slices.constants import THESIS_TASKS +from slices.data.config_schemas import DataConfig from slices.training.config_schemas import ( OptimizerConfig, + PretrainTrainingConfig, SchedulerConfig, TaskConfig, TrainingConfig, @@ -77,6 +81,11 @@ def test_defaults_applied(self): """Missing optional fields should use defaults.""" cfg = TaskConfig(task_name="mortality_24h") assert cfg.task_type == "binary" + assert cfg.prediction_window_hours is None + assert cfg.observation_window_hours is None + assert cfg.gap_hours == 0 + assert cfg.label_sources == [] + assert cfg.label_params == {} assert cfg.head_type == "mlp" assert cfg.hidden_dims == [64] assert cfg.dropout == 0.1 @@ -84,6 +93,66 @@ def test_defaults_applied(self): assert cfg.n_classes is None assert cfg.use_layer_norm is False + def test_label_definition_fields_are_accepted(self): + """Training task configs may include label semantics for freshness validation.""" + cfg = TaskConfig( + task_name="mortality_24h", + prediction_window_hours=24, + observation_window_hours=24, + gap_hours=0, + label_sources=["stays", "mortality_info"], + label_params={}, + ) + + assert cfg.prediction_window_hours == 24 + assert cfg.observation_window_hours == 24 + assert cfg.label_sources == ["stays", "mortality_info"] + + +class TestDataConfigValidation: + """Tests for data loading config bounds.""" + + def test_default_tasks_are_thesis_tasks(self): + cfg = DataConfig(processed_dir="data/processed/miiv") + + assert cfg.tasks == list(THESIS_TASKS) + + def test_split_ratios_must_be_bounded(self): + with pytest.raises(ValidationError, match="train_ratio"): + DataConfig( + processed_dir="data/processed/miiv", + train_ratio=1.2, + val_ratio=-0.1, + test_ratio=-0.1, + ) + + def test_split_ratios_must_sum_to_one(self): + with pytest.raises(ValidationError, match="Split ratios must sum to 1.0"): + DataConfig( + processed_dir="data/processed/miiv", + train_ratio=0.5, + val_ratio=0.3, + test_ratio=0.3, + ) + + def test_known_hydra_data_fields_pass(self): + cfg = DataConfig( + processed_dir="data/processed/miiv", + ricu_output_dir="data/ricu_output/miiv", + enable_sliding_windows=True, + window_size=24, + window_stride=12, + ) + + assert cfg.ricu_output_dir == "data/ricu_output/miiv" + assert cfg.enable_sliding_windows is True + assert cfg.window_size == 24 + assert cfg.window_stride == 12 + + def test_unknown_data_key_rejected(self): + with pytest.raises(ValidationError, match="procesed_dir"): + DataConfig(processed_dir="data/processed/miiv", procesed_dir="typo") + class TestTrainingConfigValidation: """Tests that TrainingConfig catches invalid/misspelled keys.""" @@ -126,6 +195,25 @@ def test_use_missing_token_field_exists(self): cfg = TrainingConfig(use_missing_token=False) assert cfg.use_missing_token is False + def test_allow_best_ckpt_fallback_field_exists(self): + """allow_best_ckpt_fallback should be accepted as a validated training flag.""" + cfg = TrainingConfig(allow_best_ckpt_fallback=True) + assert cfg.allow_best_ckpt_fallback is True + + +class TestPretrainTrainingConfigValidation: + """Tests that pretraining config typos are rejected.""" + + def test_valid_config_passes(self): + cfg = PretrainTrainingConfig(max_epochs=100, batch_size=256) + + assert cfg.max_epochs == 100 + assert cfg.batch_size == 256 + + def test_extra_key_rejected(self): + with pytest.raises(ValidationError, match="max_epoch"): + PretrainTrainingConfig(max_epoch=100) + class TestOptimizerConfigValidation: """Tests that OptimizerConfig catches invalid configs.""" @@ -176,3 +264,54 @@ def test_finetune_yaml_has_use_missing_token(self): "so users can discover and override it" ) assert raw["training"]["use_missing_token"] is True + + +class TestCheckpointFallbackYAML: + """Test that best-checkpoint fallback is discoverable in training configs.""" + + @pytest.mark.parametrize( + ("config_name", "expected"), + [ + ("finetune.yaml", False), + ("supervised.yaml", False), + ], + ) + def test_training_yaml_has_allow_best_ckpt_fallback(self, config_name, expected): + import pathlib + + yaml_path = pathlib.Path(__file__).parent.parent / "configs" / config_name + with open(yaml_path) as f: + raw = yaml.safe_load(f) + + assert "training" in raw + assert "allow_best_ckpt_fallback" in raw["training"] + assert raw["training"]["allow_best_ckpt_fallback"] is expected + + +class TestExperimentClassMetadataYAML: + """Test final rerun metadata fields in launch configs.""" + + @pytest.mark.parametrize( + "config_name", + [ + "pretrain.yaml", + "finetune.yaml", + "supervised.yaml", + "gru_d.yaml", + "xgboost.yaml", + ], + ) + def test_configs_use_experiment_class_metadata_not_sprint(self, config_name): + import pathlib + + yaml_path = pathlib.Path(__file__).parent.parent / "configs" / config_name + with open(yaml_path) as f: + raw = yaml.safe_load(f) + + assert "sprint" not in raw + assert raw["experiment_class"] is None + assert raw["experiment_subtype"] is None + assert raw["revision"] is None + assert raw["rerun_reason"] is None + assert "s${sprint}" not in raw["logging"]["run_name"] + assert "${experiment_class}" in raw["logging"]["run_name"] diff --git a/tests/test_contrastive_objective.py b/tests/test_contrastive_objective.py index 28cda32..8f8687e 100644 --- a/tests/test_contrastive_objective.py +++ b/tests/test_contrastive_objective.py @@ -1,10 +1,9 @@ -"""Tests for observation-level Contrastive (SimCLR-style) SSL objective.""" +"""Tests for timestep-level Contrastive SSL objective (temporal + instance modes).""" import pytest import torch + from slices.models.encoders import ( - ObservationTransformerConfig, - ObservationTransformerEncoder, TransformerConfig, TransformerEncoder, ) @@ -38,7 +37,6 @@ def test_l2_normalized(self): x = torch.randn(4, 32) z = head(x) - # Each row should have unit norm norms = z.norm(dim=-1) assert torch.allclose(norms, torch.ones_like(norms), atol=1e-5) @@ -53,14 +51,22 @@ class TestContrastiveInit: @pytest.fixture def encoder(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - return ObservationTransformerEncoder(config) + return TransformerEncoder(config) @pytest.fixture def contrastive_config(self): return ContrastiveConfig( + mode="instance", mask_ratio=0.75, proj_hidden_dim=64, proj_output_dim=16, @@ -73,22 +79,30 @@ def test_initialization(self, encoder, contrastive_config): assert obj.config == contrastive_config assert hasattr(obj, "projection_head") assert obj.missing_token is None - # No target encoder in SimCLR-style assert not hasattr(obj, "target_encoder") - def test_requires_observation_encoder(self): + def test_requires_obs_aware(self): config = TransformerConfig(d_input=10, d_model=32, n_layers=1, n_heads=4, pooling="none") encoder = TransformerEncoder(config) cont_config = ContrastiveConfig() - with pytest.raises(ValueError, match="tokenize.*encode"): + with pytest.raises(ValueError, match="obs_aware=True"): ContrastiveObjective(encoder, cont_config) + def test_invalid_mode_raises(self): + with pytest.raises(ValueError, match="must be 'temporal' or 'instance'"): + ContrastiveConfig(mode="bad") + def test_requires_no_pooling(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, pooling="mean" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + pooling="mean", + obs_aware=True, ) - encoder = ObservationTransformerEncoder(config) + encoder = TransformerEncoder(config) cont_config = ContrastiveConfig() with pytest.raises(ValueError, match="pooling='none'"): @@ -105,14 +119,22 @@ class TestContrastiveForward: @pytest.fixture def encoder(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - return ObservationTransformerEncoder(config) + return TransformerEncoder(config) @pytest.fixture def contrastive_config(self): return ContrastiveConfig( + mode="instance", mask_ratio=0.75, proj_hidden_dim=64, proj_output_dim=16, @@ -136,7 +158,7 @@ def test_forward_returns_loss_and_metrics(self, encoder, contrastive_config): assert "contrastive_accuracy" in metrics assert "contrastive_pos_similarity" in metrics assert "contrastive_temperature" in metrics - assert "contrastive_n_tokens_per_sample" in metrics + assert "contrastive_n_timesteps" in metrics assert "contrastive_n_visible_view1" in metrics assert "contrastive_n_visible_view2" in metrics @@ -163,10 +185,70 @@ def test_two_views_encoded(self, encoder, contrastive_config): _, metrics = obj(x, obs_mask) - # Both views should have visible tokens assert metrics["contrastive_n_visible_view1"] > 0 assert metrics["contrastive_n_visible_view2"] > 0 + def test_empty_timesteps_are_excluded_from_visible_counts(self, encoder, contrastive_config): + """Contrastive diagnostics should ignore hours with no observed variables.""" + obj = ContrastiveObjective(encoder, contrastive_config) + + B, T, D = 2, 8, 10 + x = torch.randn(B, T, D) + obs_mask = torch.zeros(B, T, D, dtype=torch.bool) + obs_mask[:, 2, :4] = True + obs_mask[:, 6, :4] = True + + _, metrics = obj(x, obs_mask) + + assert metrics["contrastive_n_visible_view1"] <= 2 + assert metrics["contrastive_n_visible_view2"] <= 2 + assert metrics["contrastive_n_visible_view1"] + metrics[ + "contrastive_n_masked_view1" + ] == pytest.approx(2.0) + assert metrics["contrastive_n_visible_view2"] + metrics[ + "contrastive_n_masked_view2" + ] == pytest.approx(2.0) + + def test_single_eligible_timestep_falls_back_to_shared_view(self, encoder): + """Complementary masks should not create an all-padding view on sparse samples.""" + config = ContrastiveConfig( + mode="instance", + mask_ratio=0.5, + complementary_masks=True, + proj_hidden_dim=64, + proj_output_dim=16, + temperature=0.1, + ) + obj = ContrastiveObjective(encoder, config) + + B, T, D = 2, 8, 10 + x = torch.randn(B, T, D) + obs_mask = torch.zeros(B, T, D, dtype=torch.bool) + obs_mask[:, 3, :4] = True + + loss, metrics = obj(x, obs_mask) + + assert torch.isfinite(loss) + assert torch.isfinite(metrics["contrastive_loss"]) + assert metrics["contrastive_n_visible_view1"] > 0 + assert metrics["contrastive_n_visible_view2"] > 0 + + def test_instance_mode_skips_all_empty_samples(self, encoder, contrastive_config): + """All-empty rows should be excluded from instance NT-Xent denominators.""" + obj = ContrastiveObjective(encoder, contrastive_config) + + B, T, D = 3, 8, 10 + x = torch.randn(B, T, D) + obs_mask = torch.zeros(B, T, D, dtype=torch.bool) + obs_mask[1:, 2:5, :4] = True + + loss, metrics = obj(x, obs_mask) + + assert torch.isfinite(loss) + assert torch.isfinite(metrics["contrastive_loss"]) + assert metrics["contrastive_n_samples_used"] == 2 + assert metrics["contrastive_n_samples_skipped"] == 1 + # ============================================================================= # NT-Xent loss tests @@ -178,13 +260,21 @@ class TestNTXentLoss: @pytest.fixture def encoder(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - return ObservationTransformerEncoder(config) + return TransformerEncoder(config) def test_accuracy_in_range(self, encoder): config = ContrastiveConfig( + mode="instance", mask_ratio=0.75, proj_hidden_dim=64, proj_output_dim=16, @@ -199,16 +289,13 @@ def test_accuracy_in_range(self, encoder): def test_perfect_alignment_low_loss(self, encoder): """When z1 == z2, NT-Xent loss should be near-minimal.""" - B, proj_dim = 8, 16 temperature = 0.1 - # Create identical, L2-normalized embeddings for both views z = torch.randn(B, proj_dim) z = torch.nn.functional.normalize(z, dim=-1) z1, z2 = z.clone(), z.clone() - # Compute NT-Xent manually z_cat = torch.cat([z1, z2], dim=0) sim_matrix = torch.mm(z_cat, z_cat.t()) / temperature labels = torch.cat([torch.arange(B, 2 * B), torch.arange(B)]) @@ -216,11 +303,8 @@ def test_perfect_alignment_low_loss(self, encoder): sim_matrix = sim_matrix.masked_fill(mask, float("-inf")) loss = torch.nn.functional.cross_entropy(sim_matrix, labels) - # With perfect alignment, positive sim = 1/tau and negatives < 1/tau - # Loss should be low (near zero for distinct samples) - assert loss.item() < 2.0 # Much less than log(2B-1) ≈ 2.7 + assert loss.item() < 2.0 - # Also: accuracy should be 1.0 with perfect alignment preds = sim_matrix.argmax(dim=1) accuracy = (preds == labels).float().mean() assert accuracy.item() == 1.0 @@ -228,12 +312,14 @@ def test_perfect_alignment_low_loss(self, encoder): def test_temperature_effect_on_loss(self, encoder): """Lower temperature should sharpen the distribution, increasing loss magnitude.""" config_low_temp = ContrastiveConfig( + mode="instance", mask_ratio=0.75, proj_hidden_dim=64, proj_output_dim=16, temperature=0.01, ) config_high_temp = ContrastiveConfig( + mode="instance", mask_ratio=0.75, proj_hidden_dim=64, proj_output_dim=16, @@ -242,12 +328,10 @@ def test_temperature_effect_on_loss(self, encoder): torch.manual_seed(123) obj_low = ContrastiveObjective(encoder, config_low_temp) - # Need a separate encoder copy for fair comparison from copy import deepcopy encoder2 = deepcopy(encoder) obj_high = ContrastiveObjective(encoder2, config_high_temp) - # Copy weights so both objectives are identical except temperature obj_high.load_state_dict(obj_low.state_dict()) x = torch.randn(8, 8, 10) @@ -258,12 +342,12 @@ def test_temperature_effect_on_loss(self, encoder): torch.manual_seed(99) loss_high, metrics_high = obj_high(x, obs_mask) - # Lower temperature produces higher loss (sharper distribution, harder task) assert loss_low.item() > loss_high.item() def test_temperature_in_metrics(self, encoder): temp = 0.07 config = ContrastiveConfig( + mode="instance", mask_ratio=0.75, proj_hidden_dim=64, proj_output_dim=16, @@ -288,13 +372,21 @@ class TestContrastiveEdgeCases: @pytest.fixture def encoder(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - return ObservationTransformerEncoder(config) + return TransformerEncoder(config) def test_sparse_data(self, encoder): config = ContrastiveConfig( + mode="instance", mask_ratio=0.5, proj_hidden_dim=64, proj_output_dim=16, @@ -303,31 +395,13 @@ def test_sparse_data(self, encoder): B, T, D = 4, 8, 10 x = torch.randn(B, T, D) - obs_mask = torch.rand(B, T, D) > 0.9 # ~10% observed - # Ensure at least some observations per sample + obs_mask = torch.rand(B, T, D) > 0.9 for b in range(B): obs_mask[b, 0, 0] = True loss, metrics = obj(x, obs_mask) assert torch.isfinite(loss) - def test_single_observation(self, encoder): - config = ContrastiveConfig( - mask_ratio=0.5, - proj_hidden_dim=64, - proj_output_dim=16, - ) - obj = ContrastiveObjective(encoder, config) - - B, T, D = 2, 4, 10 - x = torch.randn(B, T, D) - obs_mask = torch.zeros(B, T, D, dtype=torch.bool) - obs_mask[0, 0, 0] = True - obs_mask[1, 1, 3] = True - - loss, metrics = obj(x, obs_mask) - assert torch.isfinite(loss) - # ============================================================================= # Gradient flow tests @@ -339,13 +413,21 @@ class TestContrastiveGradientFlow: @pytest.fixture def encoder(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - return ObservationTransformerEncoder(config) + return TransformerEncoder(config) def test_gradients_to_encoder_and_projection(self, encoder): config = ContrastiveConfig( + mode="instance", mask_ratio=0.75, proj_hidden_dim=64, proj_output_dim=16, @@ -358,13 +440,11 @@ def test_gradients_to_encoder_and_projection(self, encoder): loss, _ = obj(x, obs_mask) loss.backward() - # Encoder should have gradients encoder_has_grad = any( p.grad is not None and p.grad.abs().sum() > 0 for p in obj.encoder.parameters() ) assert encoder_has_grad - # Projection head should have gradients proj_has_grad = any( p.grad is not None and p.grad.abs().sum() > 0 for p in obj.projection_head.parameters() ) @@ -380,12 +460,20 @@ class TestContrastiveConvergence: """Test that loss decreases during training.""" def test_loss_decreases(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=2, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=2, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - encoder = ObservationTransformerEncoder(config) + encoder = TransformerEncoder(config) cont_config = ContrastiveConfig( + mode="instance", mask_ratio=0.75, proj_hidden_dim=64, proj_output_dim=16, @@ -415,6 +503,318 @@ def test_loss_decreases(self): ), f"Loss should decrease: initial={initial_loss:.4f}, final={final_loss:.4f}" +# ============================================================================= +# Temporal mode tests +# ============================================================================= + + +class TestTemporalContrastive: + """Tests for temporal contrastive mode (per-timestep overlap pairs).""" + + @pytest.fixture + def encoder(self): + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, + ) + return TransformerEncoder(config) + + @pytest.fixture + def temporal_config(self): + return ContrastiveConfig( + mode="temporal", + mask_ratio=0.5, + proj_hidden_dim=64, + proj_output_dim=16, + temperature=0.1, + complementary_masks=False, + ) + + def test_scatter_to_full_shape_and_zeros(self, encoder, temporal_config): + """_scatter_to_full produces correct shape with zeros at masked positions.""" + from slices.models.pretraining.contrastive import ContrastiveObjective + + B, T, d_enc = 4, 8, 32 + # Create a mask with exactly 4 visible per sample + ssl_mask = torch.zeros(B, T, dtype=torch.bool) + ssl_mask[:, :4] = True # first 4 timesteps visible + + encoded = torch.randn(B, 4, d_enc) + full = ContrastiveObjective._scatter_to_full(encoded, ssl_mask, T) + + assert full.shape == (B, T, d_enc) + # Masked positions (indices 4-7) should be zero + assert (full[:, 4:, :] == 0).all() + # Visible positions should be non-zero (with high probability) + assert full[:, :4, :].abs().sum() > 0 + + def test_scatter_to_full_ignores_padded_visible_tokens(self, encoder, temporal_config): + """Uneven visible counts should not scatter padded tokens into masked slots.""" + from slices.models.pretraining.contrastive import ContrastiveObjective + + encoded = torch.tensor( + [ + [[10.0], [11.0]], + [[20.0], [99.0]], + ] + ) + ssl_mask = torch.tensor( + [ + [True, True, False], + [True, False, False], + ] + ) + + full = ContrastiveObjective._scatter_to_full(encoded, ssl_mask, n_timesteps=3) + + assert torch.allclose(full[1, 0], torch.tensor([20.0])) + assert torch.allclose(full[1, 1], torch.tensor([0.0])) + assert torch.allclose(full[1, 2], torch.tensor([0.0])) + + def test_scatter_to_full_gradient_flow(self): + """Gradients flow through scatter into torch.zeros back to the source tensor.""" + from slices.models.pretraining.contrastive import ContrastiveObjective + + B, T, d = 2, 6, 4 + ssl_mask = torch.zeros(B, T, dtype=torch.bool) + ssl_mask[:, :3] = True # 3 visible per sample + + encoded = torch.randn(B, 3, d, requires_grad=True) + full = ContrastiveObjective._scatter_to_full(encoded, ssl_mask, T) + + # Sum visible positions only — gradient should reach encoded + loss = full[ssl_mask].sum() + loss.backward() + + assert encoded.grad is not None + assert encoded.grad.abs().sum() > 0 + # Each visible encoded token contributes d elements to the sum, + # so each gradient element should be 1.0 + assert torch.allclose(encoded.grad, torch.ones_like(encoded.grad)) + + def test_scatter_to_full_roundtrip(self, encoder, temporal_config): + """Scatter should place tokens back at the correct original positions.""" + from slices.models.pretraining.contrastive import ContrastiveObjective + from slices.models.pretraining.masking import ( + create_timestep_mask, + extract_visible_timesteps, + ) + + B, T, d = 4, 8, 32 + tokens = torch.randn(B, T, d) + + torch.manual_seed(42) + ssl_mask = create_timestep_mask(B, T, 0.5, tokens.device) + + vis_tokens, vis_padding = extract_visible_timesteps(tokens, ssl_mask) + full = ContrastiveObjective._scatter_to_full(vis_tokens, ssl_mask, T) + + # At visible positions, scattered values should match original tokens + for b in range(B): + for t in range(T): + if ssl_mask[b, t]: + assert torch.allclose(full[b, t], tokens[b, t], atol=1e-6) + + def test_forward_returns_loss_and_overlap_metrics(self, encoder, temporal_config): + obj = ContrastiveObjective(encoder, temporal_config) + + B, T, D = 8, 8, 10 + x = torch.randn(B, T, D) + obs_mask = torch.ones(B, T, D, dtype=torch.bool) + + loss, metrics = obj(x, obs_mask) + + assert loss.shape == () + assert loss.item() >= 0 + assert not torch.isnan(loss) + + # Standard metrics + assert "contrastive_loss" in metrics + assert "ssl_loss" in metrics + assert "contrastive_accuracy" in metrics + assert "contrastive_pos_similarity" in metrics + + # Temporal-specific overlap metrics + assert "contrastive_n_overlap_tokens" in metrics + assert "contrastive_n_overlap_per_sample" in metrics + assert metrics["contrastive_n_overlap_tokens"] > 0 + + def test_gradient_flow_temporal(self, encoder, temporal_config): + """Gradients flow to encoder and projection head in temporal mode.""" + obj = ContrastiveObjective(encoder, temporal_config) + + x = torch.randn(8, 8, 10) + obs_mask = torch.ones(8, 8, 10, dtype=torch.bool) + + loss, _ = obj(x, obs_mask) + loss.backward() + + encoder_has_grad = any( + p.grad is not None and p.grad.abs().sum() > 0 for p in obj.encoder.parameters() + ) + assert encoder_has_grad + + proj_has_grad = any( + p.grad is not None and p.grad.abs().sum() > 0 for p in obj.projection_head.parameters() + ) + assert proj_has_grad + + def test_convergence_temporal(self): + """Loss decreases during training in temporal mode.""" + enc_config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=2, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, + ) + encoder = TransformerEncoder(enc_config) + + cont_config = ContrastiveConfig( + mode="temporal", + mask_ratio=0.5, + proj_hidden_dim=64, + proj_output_dim=16, + temperature=0.1, + complementary_masks=False, + ) + obj = ContrastiveObjective(encoder, cont_config) + optimizer = torch.optim.Adam(obj.parameters(), lr=1e-3) + + torch.manual_seed(42) + B, T, D = 8, 8, 10 + x = torch.randn(B, T, D) + obs_mask = torch.ones(B, T, D, dtype=torch.bool) + + initial_loss = None + for step in range(30): + optimizer.zero_grad() + loss, _ = obj(x, obs_mask) + loss.backward() + optimizer.step() + + if initial_loss is None: + initial_loss = loss.item() + + final_loss = loss.item() + assert ( + final_loss < initial_loss + ), f"Loss should decrease: initial={initial_loss:.4f}, final={final_loss:.4f}" + + def test_overlap_count_reasonable(self, encoder, temporal_config): + """With mask_ratio=0.5, ~25% of timesteps should overlap.""" + obj = ContrastiveObjective(encoder, temporal_config) + + B, T, D = 64, 48, 10 + x = torch.randn(B, T, D) + obs_mask = torch.ones(B, T, D, dtype=torch.bool) + + _, metrics = obj(x, obs_mask) + + # Expected overlap per sample ≈ 0.5 * 0.5 * 48 = 12 + avg_overlap = metrics["contrastive_n_overlap_per_sample"] + assert 4 < avg_overlap < 24, f"Unexpected overlap: {avg_overlap}" + + +# ============================================================================= +# Complementary masks tests +# ============================================================================= + + +class TestComplementaryMasks: + """Tests for complementary mask behavior.""" + + @pytest.fixture + def encoder(self): + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, + ) + return TransformerEncoder(config) + + def test_complementary_masks_zero_overlap(self, encoder): + """Complementary masks should have zero overlap.""" + from slices.models.pretraining.masking import create_timestep_mask + + B, T = 8, 48 + ssl_mask_1 = create_timestep_mask(B, T, 0.5, torch.device("cpu")) + ssl_mask_2 = ~ssl_mask_1 + + overlap = ssl_mask_1 & ssl_mask_2 + assert overlap.sum() == 0 + + def test_complementary_masks_full_coverage(self, encoder): + """Complementary masks should cover all timesteps.""" + from slices.models.pretraining.masking import create_timestep_mask + + B, T = 8, 48 + ssl_mask_1 = create_timestep_mask(B, T, 0.5, torch.device("cpu")) + ssl_mask_2 = ~ssl_mask_1 + + coverage = ssl_mask_1 | ssl_mask_2 + assert coverage.all() + + def test_complementary_temporal_raises(self): + """complementary_masks=True + mode='temporal' should raise ValueError.""" + with pytest.raises(ValueError, match="incompatible with mode='temporal'"): + ContrastiveConfig( + mode="temporal", + complementary_masks=True, + ) + + def test_forward_with_complementary_masks(self, encoder): + """Forward pass works with complementary masks (default).""" + config = ContrastiveConfig( + mode="instance", + mask_ratio=0.5, + proj_hidden_dim=64, + proj_output_dim=16, + complementary_masks=True, + ) + obj = ContrastiveObjective(encoder, config) + + B, T, D = 8, 16, 10 + x = torch.randn(B, T, D) + obs_mask = torch.ones(B, T, D, dtype=torch.bool) + + loss, metrics = obj(x, obs_mask) + assert torch.isfinite(loss) + assert "contrastive_loss" in metrics + + def test_forward_with_independent_masks(self, encoder): + """Forward pass works with independent masks (complementary_masks=False).""" + config = ContrastiveConfig( + mode="instance", + mask_ratio=0.5, + proj_hidden_dim=64, + proj_output_dim=16, + complementary_masks=False, + ) + obj = ContrastiveObjective(encoder, config) + + B, T, D = 8, 16, 10 + x = torch.randn(B, T, D) + obs_mask = torch.ones(B, T, D, dtype=torch.bool) + + loss, metrics = obj(x, obs_mask) + assert torch.isfinite(loss) + + # ============================================================================= # Factory integration tests # ============================================================================= @@ -431,10 +831,15 @@ def test_in_registry(self): assert get_ssl_config_class("contrastive") == ContrastiveConfig def test_build_works(self): - encoder_config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, pooling="none" + encoder_config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + pooling="none", + obs_aware=True, ) - encoder = ObservationTransformerEncoder(encoder_config) + encoder = TransformerEncoder(encoder_config) cont_config = ContrastiveConfig(mask_ratio=0.75) ssl_objective = build_ssl_objective(encoder, cont_config) @@ -443,10 +848,15 @@ def test_build_works(self): assert ssl_objective.encoder is encoder def test_get_encoder(self): - encoder_config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, pooling="none" + encoder_config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + pooling="none", + obs_aware=True, ) - encoder = ObservationTransformerEncoder(encoder_config) + encoder = TransformerEncoder(encoder_config) cont_config = ContrastiveConfig() obj = ContrastiveObjective(encoder, cont_config) diff --git a/tests/test_dataset_datamodule.py b/tests/test_dataset_datamodule.py index ea10b6f..53b69a6 100644 --- a/tests/test_dataset_datamodule.py +++ b/tests/test_dataset_datamodule.py @@ -8,8 +8,73 @@ import pytest import torch import yaml + from slices.data.datamodule import ICUDataModule, icu_collate_fn from slices.data.dataset import ICUDataset +from slices.data.splits import compute_patient_level_splits, load_cached_splits + + +def _patients_for_indices( + static_df: pl.DataFrame, + stay_ids: list[int], + indices: list[int], +) -> set[int]: + stay_to_patient = dict(zip(static_df["stay_id"].to_list(), static_df["patient_id"].to_list())) + return {stay_to_patient[stay_ids[i]] for i in indices} + + +def _write_split_matrix_fixture(data_dir, task_names: list[str]) -> None: + """Create a small processed dataset with task-specific label missingness.""" + data_dir.mkdir(parents=True) + + n_stays = 30 + stay_ids = list(range(1, n_stays + 1)) + patient_ids = [stay_id + 1000 for stay_id in stay_ids] + + with open(data_dir / "metadata.yaml", "w") as f: + yaml.dump( + { + "dataset": data_dir.name, + "feature_set": "core", + "feature_names": ["heart_rate"], + "n_features": 1, + "seq_length_hours": 2, + "min_stay_hours": 1, + "task_names": task_names, + "n_stays": n_stays, + }, + f, + ) + + pl.DataFrame( + { + "stay_id": stay_ids, + "patient_id": patient_ids, + "age": [60] * n_stays, + "gender": ["F" if i % 2 else "M" for i in stay_ids], + } + ).write_parquet(data_dir / "static.parquet") + + pl.DataFrame( + { + "stay_id": stay_ids, + "timeseries": [[[float(i)], [float(i + 1)]] for i in stay_ids], + "mask": [[[True], [True]] for _ in stay_ids], + } + ).write_parquet(data_dir / "timeseries.parquet") + + labels = {"stay_id": stay_ids} + for task_idx, task_name in enumerate(task_names): + values = [] + for stay_id in stay_ids: + if (stay_id + task_idx) % 5 == 0: + values.append(None) + elif task_name == "los_remaining": + values.append(float((stay_id % 7) + 1)) + else: + values.append(stay_id % 2) + labels[task_name] = values + pl.DataFrame(labels).write_parquet(data_dir / "labels.parquet") @pytest.fixture @@ -24,7 +89,7 @@ def mock_extracted_data(tmp_path): "feature_set": "core", "feature_names": ["heart_rate", "sbp", "resp_rate"], "n_features": 3, - "seq_length_hours": 48, + "seq_length_hours": 24, "min_stay_hours": 6, "task_names": ["mortality_24h", "mortality_hospital"], "n_stays": 10, @@ -57,8 +122,8 @@ def mock_extracted_data(tmp_path): static_df.write_parquet(data_dir / "static.parquet") # Create timeseries (dense format with nested lists) - # Each stay has 48 hours x 3 features - seq_length = 48 + # Each stay has 24 hours x 3 features + seq_length = 24 n_features = 3 timeseries_data = [] @@ -107,7 +172,7 @@ def test_initialization(self, mock_extracted_data): assert len(dataset) == 10 assert dataset.n_features == 3 - assert dataset.seq_length == 48 + assert dataset.seq_length == 24 assert dataset.feature_names == ["heart_rate", "sbp", "resp_rate"] assert dataset.task_names == ["mortality_24h", "mortality_hospital"] @@ -150,8 +215,8 @@ def test_getitem_returns_correct_tensors(self, mock_extracted_data): assert "static" in sample # Check shapes - assert sample["timeseries"].shape == (48, 3) - assert sample["mask"].shape == (48, 3) + assert sample["timeseries"].shape == (24, 3) + assert sample["mask"].shape == (24, 3) # Check types assert sample["timeseries"].dtype == torch.float32 @@ -159,8 +224,8 @@ def test_getitem_returns_correct_tensors(self, mock_extracted_data): assert isinstance(sample["stay_id"], int) assert sample["label"].dtype == torch.float32 - def test_zero_fill_after_preprocessing(self, mock_extracted_data): - """Test that zero-fill removes all NaN values.""" + def test_mean_imputation_after_preprocessing(self, mock_extracted_data): + """Missing values should be imputed with feature means when normalize=False.""" dataset = ICUDataset( mock_extracted_data, task_name="mortality_24h", @@ -169,14 +234,29 @@ def test_zero_fill_after_preprocessing(self, mock_extracted_data): sample = dataset[0] - # After zero-fill, there should be no NaN values + # After imputation, there should be no NaN values assert not torch.isnan(sample["timeseries"]).any() - # Missing positions (mask=False) should be zero - mask = sample["mask"] - missing_values = sample["timeseries"][~mask] - if len(missing_values) > 0: - assert (missing_values == 0.0).all() + # Missing positions should receive the per-feature imputation value. + missing_mask = ~sample["mask"] + if missing_mask.any(): + expected = dataset.feature_means.unsqueeze(0).expand(dataset.seq_length, -1) + assert torch.allclose(sample["timeseries"][missing_mask], expected[missing_mask]) + + def test_train_scoped_imputation_means_are_used_when_available(self, mock_extracted_data): + """normalize=False should still use train-split means when train_indices are provided.""" + dataset = ICUDataset( + mock_extracted_data, + task_name="mortality_24h", + normalize=False, + train_indices=[0, 1, 2, 3], + ) + + sample = dataset[5] + missing_mask = ~sample["mask"] + if missing_mask.any(): + expected = dataset.feature_means.unsqueeze(0).expand(dataset.seq_length, -1) + assert torch.allclose(sample["timeseries"][missing_mask], expected[missing_mask]) def test_normalization(self, mock_extracted_data): """Test that normalization is applied correctly.""" @@ -218,6 +298,71 @@ def test_get_label_statistics(self, mock_extracted_data): assert stats["mortality_24h"]["negative"] == 7 assert stats["mortality_24h"]["prevalence"] == pytest.approx(0.3) + def test_get_label_statistics_regression_task(self, tmp_path): + """Regression tasks should report numeric summaries, not binary prevalence.""" + data_dir = tmp_path / "processed_regression" + data_dir.mkdir(parents=True) + + metadata = { + "dataset": "mock", + "feature_set": "core", + "feature_names": ["heart_rate"], + "n_features": 1, + "seq_length_hours": 4, + "min_stay_hours": 4, + "task_names": ["los_remaining"], + "n_stays": 3, + } + + with open(data_dir / "metadata.yaml", "w") as f: + yaml.dump(metadata, f) + + static_df = pl.DataFrame( + { + "stay_id": [1, 2, 3], + "patient_id": [11, 12, 13], + "age": [60, 61, 62], + "gender": ["M", "F", "M"], + } + ) + static_df.write_parquet(data_dir / "static.parquet") + + timeseries_df = pl.DataFrame( + { + "stay_id": [1, 2, 3], + "timeseries": [ + [[1.0], [2.0], [3.0], [4.0]], + [[2.0], [3.0], [4.0], [5.0]], + [[3.0], [4.0], [5.0], [6.0]], + ], + "mask": [ + [[True], [True], [True], [True]], + [[True], [True], [True], [True]], + [[True], [True], [True], [True]], + ], + } + ) + timeseries_df.write_parquet(data_dir / "timeseries.parquet") + + labels_df = pl.DataFrame( + { + "stay_id": [1, 2, 3], + "los_remaining": [1.5, 2.5, 4.0], + } + ) + labels_df.write_parquet(data_dir / "labels.parquet") + + dataset = ICUDataset(data_dir, task_name="los_remaining", normalize=False) + stats = dataset.get_label_statistics() + + assert stats["los_remaining"]["task_type"] == "regression" + assert stats["los_remaining"]["total"] == 3 + assert stats["los_remaining"]["mean"] == pytest.approx((1.5 + 2.5 + 4.0) / 3) + assert stats["los_remaining"]["min"] == pytest.approx(1.5) + assert stats["los_remaining"]["max"] == pytest.approx(4.0) + assert "positive" not in stats["los_remaining"] + assert "prevalence" not in stats["los_remaining"] + def test_override_seq_length(self, mock_extracted_data): """Test that seq_length can be overridden.""" dataset = ICUDataset( @@ -258,7 +403,8 @@ def test_static_features_in_sample(self, mock_extracted_data): assert "static" in sample assert "age" in sample["static"] assert "gender" in sample["static"] - assert "los_days" in sample["static"] + assert "los_days" not in sample["static"] + assert "los_days" in dataset.static_df.columns def test_normalize_without_train_indices_raises_error(self, mock_extracted_data): """Test that normalize=True without train_indices raises ValueError (Issue #4 fix). @@ -624,7 +770,7 @@ def test_dataloaders_return_batches(self, mock_extracted_data): # Check batch size (may be smaller due to drop_last=True and small dataset) assert batch["timeseries"].shape[0] <= 2 - assert batch["timeseries"].shape[1] == 48 + assert batch["timeseries"].shape[1] == 24 assert batch["timeseries"].shape[2] == 3 def test_reproducible_splits(self, mock_extracted_data): @@ -704,6 +850,14 @@ def test_invalid_split_ratios(self): test_ratio=0.3, # Sum = 1.1 ) + with pytest.raises(ValueError, match="train_ratio must be in \\[0, 1\\]"): + ICUDataModule( + processed_dir=".", + train_ratio=1.2, + val_ratio=-0.1, + test_ratio=-0.1, + ) + def test_train_dataloader_before_setup_raises_error(self, mock_extracted_data): """Test that calling train_dataloader before setup raises RuntimeError.""" dm = ICUDataModule( @@ -752,7 +906,7 @@ def test_get_seq_length(self, mock_extracted_data): ) dm.setup() - assert dm.get_seq_length() == 48 + assert dm.get_seq_length() == 24 def test_get_label_statistics(self, mock_extracted_data): """Test get_label_statistics returns correct information.""" @@ -770,6 +924,28 @@ def test_get_label_statistics(self, mock_extracted_data): assert "positive" in stats["mortality_24h"] assert "prevalence" in stats["mortality_24h"] + def test_get_train_label_statistics_uses_optimization_subset(self, mock_extracted_data): + """Train label statistics should match the actual optimization subset.""" + dm = ICUDataModule( + processed_dir=mock_extracted_data, + task_name="mortality_24h", + seed=42, + label_fraction=0.5, + ) + dm.setup() + + subset_stats = dm.get_train_label_statistics() + full_train_stats = dm.get_train_label_statistics(use_full_train=True) + + subset_stay_ids = [dm.dataset.stay_ids[i] for i in dm.train_indices] + subset_labels = dm.dataset.labels_df.filter(pl.col("stay_id").is_in(subset_stay_ids)) + expected_positive = int((subset_labels["mortality_24h"] == 1).sum()) + + assert subset_stats["mortality_24h"]["total"] == len(dm.train_indices) + assert subset_stats["mortality_24h"]["positive"] == expected_positive + assert full_train_stats["mortality_24h"]["total"] == len(dm.full_train_indices) + assert full_train_stats["mortality_24h"]["total"] >= subset_stats["mortality_24h"]["total"] + def test_custom_split_ratios(self, mock_extracted_data): """Test custom split ratios are applied.""" dm = ICUDataModule( @@ -875,6 +1051,236 @@ def test_cached_splits_load_correctly(self, mock_extracted_data): assert dm1.val_indices == dm2.val_indices, "Val indices should match after reload" assert dm1.test_indices == dm2.test_indices, "Test indices should match after reload" + def test_cached_full_cohort_splits_are_filtered_within_patient_sets(self, mock_extracted_data): + """Task-filtered cohorts should reuse the full-cohort patient assignment.""" + dm = ICUDataModule( + processed_dir=mock_extracted_data, + task_name=None, + seed=42, + ) + dm.setup() + + metadata_path = mock_extracted_data / "metadata.yaml" + with open(metadata_path) as f: + metadata = yaml.safe_load(f) + metadata["task_names"] = [*metadata["task_names"], "aki_kdigo"] + with open(metadata_path, "w") as f: + yaml.safe_dump(metadata, f) + + labels_path = mock_extracted_data / "labels.parquet" + labels_df = pl.read_parquet(labels_path).with_columns( + pl.Series( + "aki_kdigo", + [None, None, 0, 1, 1, 0, 1, 1, 0, 0], + dtype=pl.Int64, + ) + ) + labels_df.write_parquet(labels_path) + + static_df = pl.read_parquet( + mock_extracted_data / "static.parquet", + columns=["stay_id", "patient_id"], + ) + all_stay_ids = pl.read_parquet( + mock_extracted_data / "timeseries.parquet", + columns=["stay_id"], + )["stay_id"].to_list() + filtered_stay_ids = labels_df.filter(pl.col("aki_kdigo").is_not_null())["stay_id"].to_list() + + cached = load_cached_splits( + processed_dir=mock_extracted_data, + static_df=static_df, + stay_ids=filtered_stay_ids, + seed=42, + train_ratio=0.7, + val_ratio=0.15, + test_ratio=0.15, + cohort_stay_ids=all_stay_ids, + ) + + assert cached is not None + train_indices, val_indices, test_indices = cached + + with open(mock_extracted_data / "splits.yaml") as f: + cached_yaml = yaml.safe_load(f) + + ssl_train_patients = set(cached_yaml["train_patients"]) + downstream_train_patients = _patients_for_indices( + static_df, filtered_stay_ids, train_indices + ) + downstream_val_patients = _patients_for_indices(static_df, filtered_stay_ids, val_indices) + downstream_test_patients = _patients_for_indices(static_df, filtered_stay_ids, test_indices) + + assert downstream_train_patients.issubset(ssl_train_patients) + assert downstream_val_patients.issubset(set(cached_yaml["val_patients"])) + assert downstream_test_patients.issubset(set(cached_yaml["test_patients"])) + assert ssl_train_patients.isdisjoint(downstream_val_patients) + assert ssl_train_patients.isdisjoint(downstream_test_patients) + + def test_aki_normalization_uses_full_ssl_train_split(self, tmp_path): + """AKI normalizer should use the full SSL train cohort, not AKI-labelable train.""" + data_dir = tmp_path / "aki_normalization" + _write_split_matrix_fixture(data_dir, ["aki_kdigo"]) + + ssl_dm = ICUDataModule( + processed_dir=data_dir, + task_name=None, + seed=42, + train_ratio=1.0, + val_ratio=0.0, + test_ratio=0.0, + normalize=True, + ) + ssl_dm.setup() + + aki_dm = ICUDataModule( + processed_dir=data_dir, + task_name="aki_kdigo", + seed=42, + train_ratio=1.0, + val_ratio=0.0, + test_ratio=0.0, + normalize=True, + ) + aki_dm.setup() + + assert len(ssl_dm.normalization_train_indices) == len(ssl_dm.train_indices) + assert aki_dm.normalization_train_indices == ssl_dm.normalization_train_indices + assert len(aki_dm.normalization_train_indices) == len(ssl_dm.train_indices) + assert len(aki_dm.full_train_indices) < len(aki_dm.normalization_train_indices) + assert len(aki_dm.full_train_indices) == len(aki_dm.dataset.train_indices) + assert aki_dm.dataset.normalization_train_indices == ssl_dm.normalization_train_indices + assert torch.allclose(aki_dm.dataset.feature_means, ssl_dm.dataset.feature_means) + assert torch.allclose(aki_dm.dataset.feature_stds, ssl_dm.dataset.feature_stds) + + stats_counts = set() + for stats_path in data_dir.glob("normalization_stats_*.yaml"): + with open(stats_path) as f: + stats_counts.add(yaml.safe_load(f)["train_indices_count"]) + + assert len(ssl_dm.train_indices) in stats_counts + assert len(aki_dm.full_train_indices) not in stats_counts + + def test_ssl_train_disjoint_from_task_filtered_eval_splits_for_runner_matrix(self, tmp_path): + """Every thesis dataset/task/seed should share the global patient split.""" + from scripts.internal.run_experiments import DATASETS, SEEDS_EXTENDED, TASKS + + for dataset_name in DATASETS: + data_dir = tmp_path / dataset_name + _write_split_matrix_fixture(data_dir, TASKS) + + for seed in SEEDS_EXTENDED: + ( + ssl_train_indices, + _, + _, + ssl_static_df, + _, + ssl_stay_ids, + _, + _, + _, + ) = compute_patient_level_splits( + data_dir, + task_name=None, + seed=seed, + train_ratio=0.7, + val_ratio=0.15, + test_ratio=0.15, + ) + ssl_train_patients = _patients_for_indices( + ssl_static_df, + ssl_stay_ids, + ssl_train_indices, + ) + + for task_name in TASKS: + ( + _, + downstream_val_indices, + downstream_test_indices, + downstream_static_df, + _, + _, + downstream_stay_ids, + _, + _, + ) = compute_patient_level_splits( + data_dir, + task_name=task_name, + seed=seed, + train_ratio=0.7, + val_ratio=0.15, + test_ratio=0.15, + ) + downstream_val_patients = _patients_for_indices( + downstream_static_df, + downstream_stay_ids, + downstream_val_indices, + ) + downstream_test_patients = _patients_for_indices( + downstream_static_df, + downstream_stay_ids, + downstream_test_indices, + ) + + assert ssl_train_patients.isdisjoint( + downstream_val_patients + ), f"{dataset_name}/{task_name}/seed{seed}: SSL train overlaps val" + assert ssl_train_patients.isdisjoint( + downstream_test_patients + ), f"{dataset_name}/{task_name}/seed{seed}: SSL train overlaps test" + + def test_label_efficiency_splits_preserve_full_split_provenance(self, mock_extracted_data): + """splits.yaml should keep the full split and record the train subset separately.""" + dm = ICUDataModule( + processed_dir=mock_extracted_data, + task_name="mortality_24h", + seed=42, + label_fraction=0.5, + ) + dm.setup() + + with open(mock_extracted_data / "splits.yaml") as f: + saved_splits = yaml.safe_load(f) + + full_train_patients = sorted( + { + dm.dataset.static_df.filter(pl.col("stay_id") == dm.dataset.stay_ids[i])[ + "patient_id" + ].item() + for i in dm.full_train_indices + } + ) + subset_stay_ids = sorted(dm.dataset.stay_ids[i] for i in dm.train_indices) + + assert saved_splits["train_stays"] == len(dm.full_train_indices) + assert saved_splits["train_patients"] == full_train_patients + assert saved_splits["train_subset_stays"] == len(dm.train_indices) + assert sorted(saved_splits["train_subset_stay_ids"]) == subset_stay_ids + + def test_existing_splits_yaml_is_not_rewritten_by_label_efficiency_runs( + self, mock_extracted_data + ): + """Prepared canonical splits.yaml should remain stable across downstream runs.""" + dm_full = ICUDataModule( + processed_dir=mock_extracted_data, + task_name="mortality_24h", + seed=42, + ) + dm_full.setup() + canonical_text = (mock_extracted_data / "splits.yaml").read_text() + + dm_subset = ICUDataModule( + processed_dir=mock_extracted_data, + task_name="mortality_24h", + seed=42, + label_fraction=0.5, + ) + dm_subset.setup() + + assert (mock_extracted_data / "splits.yaml").read_text() == canonical_text + class TestCollateFn: """Tests for the collate function.""" @@ -1019,8 +1425,9 @@ def test_normalization_stats_persistence(self, mock_extracted_data): train_indices=[0, 1, 2, 3], # Only first 4 samples are training ) - stats_file = data_dir / "normalization_stats.yaml" - assert stats_file.exists(), "normalization_stats.yaml should be created" + # Stats are now hash-keyed: normalization_stats_.yaml + stats_files = list(data_dir.glob("normalization_stats_*.yaml")) + assert len(stats_files) >= 1, "normalization_stats_.yaml should be created" # Get stats from first dataset means1 = dataset1.feature_means.clone() @@ -1250,14 +1657,17 @@ def test_no_inline_import_warnings(self): if child.module == "warnings": inline_imports.append(node.name) - assert len(inline_imports) == 0, ( - f"Found 'import warnings' inside functions: {inline_imports}. " "Move to module level." - ) + assert ( + len(inline_imports) == 0 + ), f"Found 'import warnings' inside functions: {inline_imports}. Move to module level." - def test_module_level_import_exists(self): - """dataset.py should have warnings imported at module level.""" - import slices.data.dataset as mod + def test_tensor_cache_has_warnings_import(self): + """tensor_cache.py should have warnings imported at module level. + + Warning logic for normalization stats lives in tensor_cache.py, not dataset.py. + """ + import slices.data.tensor_cache as mod assert hasattr(mod, "warnings") or "warnings" in dir( mod - ), "dataset.py should import warnings at module level" + ), "tensor_cache.py should import warnings at module level" diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py index 6c2b54b..6eca227 100644 --- a/tests/test_embeddings.py +++ b/tests/test_embeddings.py @@ -6,6 +6,7 @@ import numpy as np import pytest import torch + from slices.debug.embeddings import ( CollapseMetrics, DimensionalityMetrics, diff --git a/tests/test_encoder_wrapper.py b/tests/test_encoder_wrapper.py index f253f71..255bdfa 100644 --- a/tests/test_encoder_wrapper.py +++ b/tests/test_encoder_wrapper.py @@ -11,6 +11,7 @@ import pytest import torch import torch.nn as nn + from slices.models.encoders import ( EncoderWithMissingToken, TransformerConfig, diff --git a/tests/test_evaluate_fairness.py b/tests/test_evaluate_fairness.py new file mode 100644 index 0000000..d08e7ea --- /dev/null +++ b/tests/test_evaluate_fairness.py @@ -0,0 +1,538 @@ +"""Focused tests for the standalone fairness rerun script.""" + +import importlib +import sys +from types import SimpleNamespace + +import pandas as pd +import pytest + +from slices.eval.fairness_metadata import ( + EVAL_ARTIFACT_SHA256_KEY, + FAIRNESS_ARTIFACT_PATH_KEY, + FAIRNESS_ARTIFACT_SHA256_KEY, + FAIRNESS_ARTIFACT_SOURCE_KEY, + FAIRNESS_CHECKPOINT_SOURCE_KEY, + FAIRNESS_MIN_SUBGROUP_SIZE_KEY, + FAIRNESS_PROTECTED_ATTRIBUTES_KEY, + FAIRNESS_SCHEMA_VERSION_KEY, + FAIRNESS_SCRIPT_VERSION, + FAIRNESS_SCRIPT_VERSION_KEY, + FAIRNESS_SUMMARY_SCHEMA_VERSION, + encode_protected_attributes, +) + + +def _binary_fairness_summary(attr: str, base: float = 0.7) -> dict[str, float]: + prefix = f"fairness/{attr}/" + return { + f"{prefix}n_valid_groups": 2, + f"{prefix}n_metric_valid_groups": 2, + f"{prefix}worst_group_auroc": base, + f"{prefix}worst_group_auprc": base - 0.1, + f"{prefix}auroc_gap": 0.05, + f"{prefix}auprc_gap": 0.04, + f"{prefix}demographic_parity_diff": 0.03, + f"{prefix}equalized_odds_diff": 0.02, + f"{prefix}disparate_impact_ratio": 0.9, + } + + +def _regression_fairness_summary(attr: str) -> dict[str, float]: + prefix = f"fairness/{attr}/" + return { + f"{prefix}n_valid_groups": 2, + f"{prefix}worst_group_mse": 2.0, + f"{prefix}worst_group_mae": 1.1, + f"{prefix}mse_gap": 0.4, + f"{prefix}mae_gap": 0.2, + } + + +def _fresh_fairness_metadata( + protected_attributes: list[str], + min_subgroup_size: int = 50, + artifact_path: str = "outputs/run/checkpoints/last.ckpt", + artifact_source: str = "recorded_final", + checkpoint_source: str = "final", + artifact_sha256: str = "abc123", +) -> dict[str, object]: + metadata = { + "_eval_checkpoint_source": ( + checkpoint_source if checkpoint_source in {"best", "final"} else None + ), + EVAL_ARTIFACT_SHA256_KEY: artifact_sha256, + FAIRNESS_SCHEMA_VERSION_KEY: FAIRNESS_SUMMARY_SCHEMA_VERSION, + FAIRNESS_SCRIPT_VERSION_KEY: FAIRNESS_SCRIPT_VERSION, + FAIRNESS_ARTIFACT_PATH_KEY: artifact_path, + FAIRNESS_ARTIFACT_SHA256_KEY: artifact_sha256, + FAIRNESS_ARTIFACT_SOURCE_KEY: artifact_source, + FAIRNESS_CHECKPOINT_SOURCE_KEY: checkpoint_source, + FAIRNESS_PROTECTED_ATTRIBUTES_KEY: encode_protected_attributes(protected_attributes), + FAIRNESS_MIN_SUBGROUP_SIZE_KEY: min_subgroup_size, + } + return {key: value for key, value in metadata.items() if value is not None} + + +def test_default_phases_include_baseline(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + assert "baseline" in mod.DEFAULT_PHASES + + +def test_default_experiment_classes_include_classical_baselines(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + assert "classical_baselines" in mod.DEFAULT_EXPERIMENT_CLASSES + + +def test_default_experiment_classes_include_downstream_families(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + assert {"label_efficiency", "cross_dataset_transfer", "hp_robustness"}.issubset( + set(mod.DEFAULT_EXPERIMENT_CLASSES) + ) + + +def test_parse_args_exposes_allow_incomplete_escape_hatch(monkeypatch): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + monkeypatch.delenv("WANDB_PROJECT", raising=False) + monkeypatch.setattr( + sys, + "argv", + ["evaluate_fairness.py", "--revision", "thesis-v1", "--allow-incomplete"], + ) + + args = mod.parse_args() + + assert args.allow_incomplete is True + assert args.project == "slices-thesis" + + +def test_main_exits_nonzero_when_no_runs_match(monkeypatch): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + monkeypatch.setattr(sys, "argv", ["evaluate_fairness.py", "--revision", "thesis-v1"]) + monkeypatch.setattr(mod, "fetch_eval_runs", lambda **_: []) + + with pytest.raises(SystemExit) as excinfo: + mod.main() + + assert excinfo.value.code == 1 + + +def test_resolve_evaluation_artifact_supports_xgboost(tmp_path): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + run_dir = tmp_path / "run-xgb" + run_dir.mkdir(parents=True) + model_path = run_dir / "xgboost_model.json" + model_path.write_text("{}") + + run = SimpleNamespace( + id="run-xgb", + config={ + "paradigm": "xgboost", + "output_dir": "outputs/run-xgb", + }, + summary_metrics={}, + ) + + artifact_path, source = mod.resolve_evaluation_artifact( + run, + outputs_root=str(tmp_path), + task_type="binary", + ) + + assert artifact_path == model_path + assert source == "xgboost_model" + + +def test_build_fairness_summary_metadata_stamps_artifact_and_settings(tmp_path): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + ckpt_path = tmp_path / "outputs" / "run" / "checkpoints" / "best.ckpt" + ckpt_path.parent.mkdir(parents=True) + ckpt_path.write_text("best checkpoint") + run = SimpleNamespace( + config={"paradigm": "mae", "output_dir": "outputs/run"}, + summary_metrics={ + "_eval_checkpoint_source": "best", + "_best_ckpt_path": "outputs/run/checkpoints/best.ckpt", + }, + ) + + metadata = mod.build_fairness_summary_metadata( + run, + ckpt_path, + "recorded_best", + ["gender", "age_group", "race"], + 50, + ) + + assert metadata[FAIRNESS_ARTIFACT_PATH_KEY] == str(ckpt_path) + assert metadata[FAIRNESS_ARTIFACT_SHA256_KEY] + assert metadata[FAIRNESS_ARTIFACT_SOURCE_KEY] == "recorded_best" + assert metadata[FAIRNESS_CHECKPOINT_SOURCE_KEY] == "best" + assert metadata[FAIRNESS_SCRIPT_VERSION_KEY] == FAIRNESS_SCRIPT_VERSION + assert metadata[FAIRNESS_PROTECTED_ATTRIBUTES_KEY] == encode_protected_attributes( + ["gender", "age_group", "race"] + ) + assert metadata[FAIRNESS_MIN_SUBGROUP_SIZE_KEY] == 50 + + +def test_write_fairness_to_wandb_replaces_existing_fairness_keys_by_default(monkeypatch): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + class FakeSummary(dict): + def save(self): + self["saved"] = True + + fake_summary = FakeSummary( + { + "fairness/race/worst_group_auroc": 0.11, + "_fairness_old_setting": "stale", + "test/auprc": 0.42, + } + ) + + class FakeRun: + summary = fake_summary + + class FakeApi: + def __init__(self, timeout): + self.timeout = timeout + + def run(self, run_path): + assert run_path == "entity/project/run" + return FakeRun() + + fake_wandb = SimpleNamespace(Api=FakeApi) + monkeypatch.setitem(sys.modules, "wandb", fake_wandb) + + mod.write_fairness_to_wandb( + "entity/project/run", + { + "fairness/gender/n_valid_groups": 2, + "_fairness_summary_schema_version": "fresh", + }, + ) + + assert fake_summary["fairness/race/worst_group_auroc"] is None + assert fake_summary["_fairness_old_setting"] is None + assert fake_summary["fairness/gender/n_valid_groups"] == 2 + assert fake_summary["test/auprc"] == 0.42 + assert fake_summary["saved"] is True + + +def test_has_fairness_metrics_requires_requested_attribute_completeness(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + run = SimpleNamespace( + config={"dataset": "miiv", "task": {"task_type": "binary"}}, + summary_metrics={ + **_binary_fairness_summary("gender", 0.71), + **_binary_fairness_summary("age_group", 0.69), + **_fresh_fairness_metadata(["gender", "age_group"]), + }, + ) + + assert mod.has_fairness_metrics(run, ["gender", "age_group"]) is True + assert mod.has_fairness_metrics(run, ["gender", "age_group", "race"]) is False + + +def test_has_fairness_metrics_ignores_race_for_eicu(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + run = SimpleNamespace( + config={"dataset": "eicu", "task": {"task_type": "binary"}}, + summary_metrics={ + **_binary_fairness_summary("gender", 0.71), + **_binary_fairness_summary("age_group", 0.69), + **_fresh_fairness_metadata(["gender", "age_group", "race"]), + }, + ) + + assert mod.has_fairness_metrics(run, ["gender", "age_group", "race"]) is True + + +def test_has_fairness_metrics_rejects_partial_binary_summaries(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + run = SimpleNamespace( + config={"dataset": "miiv", "task": {"task_type": "binary"}}, + summary_metrics={ + "fairness/gender/n_valid_groups": 2, + "fairness/gender/worst_group_auroc": 0.71, + }, + ) + + assert mod.has_fairness_metrics(run, ["gender"]) is False + + +def test_has_fairness_metrics_accepts_written_nan_summary_values(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + summary = _binary_fairness_summary("gender") + summary["fairness/gender/worst_group_auroc"] = float("nan") + summary.update(_fresh_fairness_metadata(["gender"])) + run = SimpleNamespace( + config={"dataset": "miiv", "task": {"task_type": "binary"}}, + summary_metrics=summary, + ) + + assert mod.has_fairness_metrics(run, ["gender"]) is True + + +def test_has_fairness_metrics_rejects_complete_but_unstamped_summaries(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + run = SimpleNamespace( + config={"dataset": "miiv", "task": {"task_type": "binary"}}, + summary_metrics=_binary_fairness_summary("gender"), + ) + + assert mod.has_fairness_metrics(run, ["gender"]) is False + + +def test_has_fairness_metrics_rejects_stale_metadata_settings(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + run = SimpleNamespace( + config={"dataset": "miiv", "task": {"task_type": "binary"}}, + summary_metrics={ + **_binary_fairness_summary("gender"), + **_fresh_fairness_metadata(["gender"], min_subgroup_size=25), + }, + ) + + assert mod.has_fairness_metrics(run, ["gender"], min_subgroup_size=50) is False + assert any( + "min subgroup size mismatch" in issue + for issue in mod.fairness_summary_metadata_issues(run, ["gender"], 50) + ) + + +def test_has_fairness_metrics_rejects_artifact_digest_mismatch(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + run = SimpleNamespace( + config={"dataset": "miiv", "task": {"task_type": "binary"}}, + summary_metrics={ + **_binary_fairness_summary("gender"), + **_fresh_fairness_metadata(["gender"], artifact_sha256="old-digest"), + EVAL_ARTIFACT_SHA256_KEY: "new-digest", + }, + ) + + assert mod.has_fairness_metrics(run, ["gender"], min_subgroup_size=50) is False + assert any( + "artifact sha256 mismatch" in issue + for issue in mod.fairness_summary_metadata_issues(run, ["gender"], 50) + ) + + +def test_has_fairness_metrics_requires_regression_metric_family(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + complete = SimpleNamespace( + config={ + "dataset": "miiv", + "task": {"task_name": "los_remaining", "task_type": "regression"}, + }, + summary_metrics={ + **_regression_fairness_summary("gender"), + **_regression_fairness_summary("age_group"), + **_fresh_fairness_metadata(["gender", "age_group"]), + }, + ) + partial = SimpleNamespace( + config={ + "dataset": "miiv", + "task": {"task_name": "los_remaining", "task_type": "regression"}, + }, + summary_metrics={ + "fairness/gender/n_valid_groups": 2, + "fairness/gender/worst_group_mse": 2.0, + **_regression_fairness_summary("age_group"), + **_fresh_fairness_metadata(["gender", "age_group"]), + }, + ) + + assert mod.has_fairness_metrics(complete, ["gender", "age_group"]) is True + assert mod.has_fairness_metrics(partial, ["gender", "age_group"]) is False + + +def test_missing_fairness_report_requirements_flags_missing_requested_attribute(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + run = SimpleNamespace(config={"dataset": "miiv", "task": {"task_type": "binary"}}) + report = { + "gender": { + "n_valid_groups": 2, + "n_metric_valid_groups": 2, + "worst_group_auroc": 0.71, + "worst_group_auprc": 0.61, + "auroc_gap": 0.05, + "auprc_gap": 0.04, + "demographic_parity_diff": 0.03, + "equalized_odds_diff": 0.02, + "disparate_impact_ratio": 0.9, + } + } + + missing = mod.missing_fairness_report_requirements(run, report, ["gender", "race"]) + + assert "race: no valid fairness groups" in missing + + +def test_missing_fairness_report_requirements_ignores_eicu_race(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + run = SimpleNamespace(config={"dataset": "eicu", "task": {"task_type": "binary"}}) + report = { + "gender": { + "n_valid_groups": 2, + "n_metric_valid_groups": 2, + "worst_group_auroc": 0.71, + "worst_group_auprc": 0.61, + "auroc_gap": 0.05, + "auprc_gap": 0.04, + "demographic_parity_diff": 0.03, + "equalized_odds_diff": 0.02, + "disparate_impact_ratio": 0.9, + } + } + + missing = mod.missing_fairness_report_requirements(run, report, ["gender", "race"]) + + assert missing == [] + + +def test_missing_fairness_report_requirements_accepts_written_nan_required_metric(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + run = SimpleNamespace(config={"dataset": "miiv", "task": {"task_type": "binary"}}) + report = { + "gender": { + "n_valid_groups": 2, + "n_metric_valid_groups": 1, + "worst_group_auroc": float("nan"), + "worst_group_auprc": 0.61, + "auroc_gap": 0.05, + "auprc_gap": 0.04, + "demographic_parity_diff": 0.03, + "equalized_odds_diff": 0.02, + "disparate_impact_ratio": 0.9, + } + } + + missing = mod.missing_fairness_report_requirements(run, report, ["gender"]) + + assert missing == [] + + +def test_filter_thesis_task_runs_drops_non_thesis_tasks(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + main = SimpleNamespace(config={"task": {"task_name": "mortality_24h"}}) + extension = SimpleNamespace(config={"task": {"task_name": "sepsis"}}) + + assert mod.filter_thesis_task_runs([main, extension]) == [main] + + +def test_fairness_matrix_coverage_accepts_complete_scoped_corpus(monkeypatch): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + export_mod = importlib.import_module("scripts.export_results") + row = { + "experiment_class": "core_ssl_benchmark", + "experiment_type": "core_ssl_benchmark", + "experiment_subtype": None, + "paradigm": "mae", + "dataset": "miiv", + "task": "mortality_24h", + "protocol": "full_finetune", + "label_fraction": 1.0, + "model_size": "default", + "source_dataset": None, + "phase": "finetune", + "upstream_pretrain_lr": None, + "upstream_pretrain_mask_ratio": None, + "seed": 42, + } + monkeypatch.setattr(export_mod, "build_expected_matrix_df", lambda **_: pd.DataFrame([row])) + run = SimpleNamespace( + config={ + "experiment_class": "core_ssl_benchmark", + "paradigm": "mae", + "dataset": "miiv", + "phase": "finetune", + "protocol": "full_finetune", + "seed": 42, + "task": {"task_name": "mortality_24h"}, + }, + tags=[], + ) + + issues = mod.fairness_matrix_coverage_issues( + [run], + experiment_classes=["core_ssl_benchmark"], + paradigms=None, + datasets=None, + phases=["finetune"], + ) + + assert issues == [] + + +def test_fairness_matrix_coverage_flags_missing_scoped_rows(monkeypatch): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + export_mod = importlib.import_module("scripts.export_results") + expected = { + "experiment_class": "core_ssl_benchmark", + "experiment_type": "core_ssl_benchmark", + "experiment_subtype": None, + "paradigm": "mae", + "dataset": "miiv", + "task": "mortality_24h", + "protocol": "full_finetune", + "label_fraction": 1.0, + "model_size": "default", + "source_dataset": None, + "phase": "finetune", + "upstream_pretrain_lr": None, + "upstream_pretrain_mask_ratio": None, + "seed": 42, + } + monkeypatch.setattr( + export_mod, "build_expected_matrix_df", lambda **_: pd.DataFrame([expected]) + ) + + issues = mod.fairness_matrix_coverage_issues( + [], + experiment_classes=["core_ssl_benchmark"], + paradigms=None, + datasets=None, + phases=["finetune"], + ) + + assert any("missing 1/1 expected downstream matrix rows" in issue for issue in issues) + + +def test_validate_evaluation_artifact_digest_rejects_stale_local_file(tmp_path): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + artifact = tmp_path / "outputs" / "run" / "checkpoints" / "last.ckpt" + artifact.parent.mkdir(parents=True) + artifact.write_text("local stale artifact") + run = SimpleNamespace( + id="run", + config={"paradigm": "mae", "output_dir": "outputs/run"}, + summary_metrics={ + "_eval_checkpoint_source": "final", + EVAL_ARTIFACT_SHA256_KEY: "logged-digest", + }, + ) + + with pytest.raises(RuntimeError, match="sha256 mismatch"): + mod.validate_evaluation_artifact_digest(run, artifact) diff --git a/tests/test_export_results.py b/tests/test_export_results.py new file mode 100644 index 0000000..7998cdc --- /dev/null +++ b/tests/test_export_results.py @@ -0,0 +1,975 @@ +"""Tests for the class-based results export pipeline.""" + +import importlib +import json +import sys + +import pandas as pd +import pytest + + +class DummyRun: + def __init__(self, config, tags, summary=None, name="dummy"): + self.config = config + self.tags = tags + self.summary_metrics = summary or {} + self.id = "dummy-id" + self.url = "https://wandb.test/dummy-id" + self.name = name + self.group = None + self.created_at = "2026-04-21T00:00:00" + + +def _row(experiment_class, paradigm, seed, task="mortality_24h", offset=0.0): + return { + "experiment_class": experiment_class, + "experiment_type": experiment_class, + "experiment_subtype": None, + "paradigm": paradigm, + "dataset": "miiv", + "task": task, + "seed": seed, + "protocol": "full_finetune", + "label_fraction": 1.0, + "model_size": "default", + "source_dataset": None, + "phase": "baseline" if paradigm in {"gru_d", "xgboost"} else "finetune", + "upstream_pretrain_lr": None, + "upstream_pretrain_mask_ratio": None, + "test/auprc": 0.3 + offset + (seed / 100000.0), + } + + +def _fresh_fairness_metadata( + mod, + protected_attributes=None, + artifact_path="outputs/run/checkpoints/last.ckpt", + artifact_source="recorded_final", + checkpoint_source="final", + min_subgroup_size=50, + artifact_sha256="abc123", +): + protected_attributes = protected_attributes or mod.FAIRNESS_ATTRIBUTES + return { + mod.EVAL_ARTIFACT_SHA256_KEY: artifact_sha256, + mod.FAIRNESS_SCHEMA_VERSION_KEY: mod.FAIRNESS_SUMMARY_SCHEMA_VERSION, + mod.FAIRNESS_SCRIPT_VERSION_KEY: mod.FAIRNESS_SCRIPT_VERSION, + mod.FAIRNESS_ARTIFACT_PATH_KEY: artifact_path, + mod.FAIRNESS_ARTIFACT_SHA256_KEY: artifact_sha256, + mod.FAIRNESS_ARTIFACT_SOURCE_KEY: artifact_source, + mod.FAIRNESS_CHECKPOINT_SOURCE_KEY: checkpoint_source, + mod.FAIRNESS_PROTECTED_ATTRIBUTES_KEY: json.dumps(protected_attributes), + mod.FAIRNESS_MIN_SUBGROUP_SIZE_KEY: min_subgroup_size, + } + + +def _add_binary_fairness_summary(row, mod, attrs=None, value=0.0): + attrs = attrs or mod.FAIRNESS_ATTRIBUTES + required_keys = [ + f"fairness/{attr}/{metric_name}" + for attr in attrs + for metric_name in mod.BINARY_FAIRNESS_REQUIRED_METRICS + ] + row[mod.FAIRNESS_SUMMARY_KEY_COLUMN] = json.dumps(required_keys) + for key in required_keys: + row[key] = value + return row + + +def test_build_statistical_tests_df_produces_pairwise_significance_rows(): + mod = importlib.import_module("scripts.export_results") + + rows = [] + for paradigm, offset in [("mae", 0.10), ("supervised", 0.0)]: + for seed in [42, 123, 456]: + for task, base in [("mortality_24h", 0.30), ("aki_kdigo", 0.25)]: + row = _row("core_ssl_benchmark", paradigm, seed, task, offset) + row["test/auprc"] = base + offset + (seed / 100000.0) + rows.append(row) + + stats_df = mod.build_statistical_tests_df(pd.DataFrame(rows)) + + omnibus = stats_df[stats_df["comparison_scope"] == "omnibus_primary_metric"] + assert len(omnibus) == 1 + row = omnibus.iloc[0] + assert row["primary_metric_name"] == "test/auprc" + assert row["paradigm_a"] == "mae" + assert row["paradigm_b"] == "supervised" + assert row["n_pairs"] == 6 + assert row["better_paradigm"] == "mae" + assert row["n_shared_task_seed_pairs"] == 6 + + per_task = stats_df[stats_df["comparison_scope"] == "per_task"] + assert set(per_task["task"]) == {"mortality_24h", "aki_kdigo"} + assert set(per_task["n_pairs"]) == {3} + + +def test_extract_run_requires_experiment_class(): + mod = importlib.import_module("scripts.export_results") + + run = DummyRun( + config={ + "dataset": "miiv", + "paradigm": "mae", + "seed": 42, + "training": {"freeze_encoder": False}, + "task": {"task_name": "mortality_24h"}, + }, + tags=["phase:finetune"], + ) + + with pytest.raises(RuntimeError, match="no experiment_class"): + mod.extract_run(run, []) + + +def test_extract_run_uses_class_metadata_from_config_or_tags(): + mod = importlib.import_module("scripts.export_results") + + config_run = DummyRun( + config={ + "experiment_class": "core_ssl_benchmark", + "dataset": "miiv", + "paradigm": "mae", + "seed": 42, + "training": {"freeze_encoder": False}, + "task": {"task_name": "mortality_24h"}, + "encoder": {"d_model": 64, "n_layers": 2}, + }, + tags=["phase:finetune", "revision:thesis-v1"], + ) + tag_run = DummyRun( + config={ + "dataset": "miiv", + "paradigm": "xgboost", + "seed": 42, + "task": {"task_name": "mortality_24h"}, + }, + tags=["experiment_class:classical_baselines", "phase:baseline"], + ) + + assert mod.extract_run(config_run, [])["experiment_class"] == "core_ssl_benchmark" + tag_row = mod.extract_run(tag_run, []) + assert tag_row["experiment_class"] == "classical_baselines" + assert tag_row["protocol"] == "full_finetune" + + +def test_extract_run_exports_launch_commit_from_config_or_tags(): + mod = importlib.import_module("scripts.export_results") + + config_run = DummyRun( + config={ + "experiment_class": "core_ssl_benchmark", + "dataset": "miiv", + "paradigm": "mae", + "seed": 42, + "launch_commit": "abcdef1234567890", + "training": {"freeze_encoder": False}, + "task": {"task_name": "mortality_24h"}, + }, + tags=["phase:finetune", "commit:tag-commit"], + ) + tag_run = DummyRun( + config={ + "experiment_class": "classical_baselines", + "dataset": "miiv", + "paradigm": "xgboost", + "seed": 42, + "task": {"task_name": "mortality_24h"}, + }, + tags=["phase:baseline", "commit:tag-commit"], + ) + + assert mod.extract_run(config_run, [])["launch_commit"] == "abcdef1234567890" + assert mod.extract_run(tag_run, [])["launch_commit"] == "tag-commit" + + +def test_fetch_all_runs_filters_by_config_before_tags(monkeypatch): + mod = importlib.import_module("scripts.export_results") + captured = {} + keep_config = DummyRun( + config={ + "experiment_class": "core_ssl_benchmark", + "paradigm": "mae", + "dataset": "miiv", + "phase": "finetune", + "revision": "thesis-v1", + }, + tags=[], + ) + keep_config.id = "keep-config" + keep_tags = DummyRun( + config={}, + tags=[ + "experiment_class:core_ssl_benchmark", + "paradigm:mae", + "dataset:miiv", + "phase:finetune", + "revision:thesis-v1", + ], + ) + keep_tags.id = "keep-tags" + drop = DummyRun( + config={"experiment_class": "core_ssl_benchmark", "revision": "other"}, + tags=[], + ) + drop.id = "drop" + drop_stale_tag = DummyRun( + config={"revision": "other"}, + tags=["revision:thesis-v1"], + ) + drop_stale_tag.id = "drop-stale-tag" + + class DummyApi: + def __init__(self, timeout): + captured["timeout"] = timeout + + def runs(self, path, filters, order): + captured["path"] = path + captured["filters"] = filters + captured["order"] = order + return [keep_config, keep_tags, drop, drop_stale_tag] + + monkeypatch.setattr(mod.wandb, "Api", DummyApi) + + runs = mod.fetch_all_runs( + project="proj", + entity="entity", + experiment_class=["core_ssl_benchmark"], + paradigm=["mae"], + dataset=["miiv"], + phase=["finetune"], + revision=["thesis-v1"], + ) + + assert captured["path"] == "entity/proj" + assert captured["filters"] == {"state": "finished"} + assert [run.id for run in runs] == ["keep-config", "keep-tags"] + + +def test_build_per_seed_df_keeps_hp_robustness_rows_out_of_derived_views(): + mod = importlib.import_module("scripts.export_results") + + runs = [] + for run_id, lr_code in [("hp-lr-2e4", "00002"), ("hp-lr-5e4", "00005")]: + run = DummyRun( + config={ + "experiment_class": "hp_robustness", + "experiment_subtype": "lr_sensitivity", + "dataset": "miiv", + "paradigm": "mae", + "seed": 42, + "output_dir": ( + "outputs/hp_robustness/" f"finetune_mae_mortality_24h_miiv_seed42_lr{lr_code}" + ), + "encoder": {"d_model": 64, "n_layers": 2}, + "training": {"freeze_encoder": False}, + "task": {"task_name": "mortality_24h"}, + "label_fraction": 1.0, + }, + tags=["experiment_class:hp_robustness", "phase:finetune"], + name=f"finetune_mae_mortality_24h_miiv_seed42_lr{lr_code}", + ) + run.id = run_id + run.url = f"https://wandb.test/{run_id}" + runs.append(run) + + df = mod.build_per_seed_df(runs) + + assert len(df) == 2 + assert set(df["experiment_type"]) == {"hp_robustness"} + assert mod.build_label_efficiency_curve_df(df).empty + assert mod.build_capacity_study_comparison_df(df).empty + assert sorted(df["upstream_pretrain_lr"].tolist()) == [2e-4, 5e-4] + + +def test_build_per_seed_df_fails_closed_on_extraction_errors(): + mod = importlib.import_module("scripts.export_results") + + with pytest.raises(RuntimeError, match="failed extraction"): + mod.build_per_seed_df([object()]) + + assert mod.build_per_seed_df([object()], allow_extraction_failures=True).empty + + +def test_build_per_seed_df_fails_closed_on_exact_duplicate_fingerprints(): + mod = importlib.import_module("scripts.export_results") + + base_config = { + "experiment_class": "core_ssl_benchmark", + "dataset": "miiv", + "paradigm": "mae", + "seed": 42, + "encoder": {"d_model": 64, "n_layers": 2}, + "training": {"freeze_encoder": False}, + "optimizer": {"lr": 3e-4}, + "task": {"task_name": "mortality_24h"}, + } + run_a = DummyRun(base_config, ["phase:finetune"], name="run-a") + run_a.id = "run-a" + run_a.created_at = "2026-04-21T00:00:01" + run_b = DummyRun(base_config, ["phase:finetune"], name="run-b") + run_b.id = "run-b" + run_b.created_at = "2026-04-21T00:00:02" + + with pytest.raises(RuntimeError, match="Duplicate exact export fingerprints"): + mod.build_per_seed_df([run_a, run_b]) + + deduped = mod.build_per_seed_df([run_a, run_b], allow_duplicate_fingerprints=True) + assert deduped["wandb_run_id"].tolist() == ["run-b"] + + +def test_extract_run_uses_xgboost_learning_rate_for_lr_metadata(): + mod = importlib.import_module("scripts.export_results") + + run = DummyRun( + config={ + "experiment_class": "classical_baselines", + "dataset": "miiv", + "paradigm": "xgboost", + "seed": 42, + "protocol": "full_finetune", + "optimizer": {"lr": 3e-4}, + "xgboost": {"learning_rate": 0.1}, + "task": {"task_name": "mortality_24h"}, + }, + tags=["phase:baseline"], + ) + + assert mod.extract_run(run, [])["lr"] == pytest.approx(0.1) + + +def test_label_efficiency_view_adds_core_full_label_endpoint(): + mod = importlib.import_module("scripts.export_results") + + rows = [ + _row("core_ssl_benchmark", "mae", 42, offset=0.1), + {**_row("label_efficiency", "mae", 42, offset=0.0), "label_fraction": 0.1}, + ] + + view = mod.build_label_efficiency_curve_df(pd.DataFrame(rows)) + + assert set(view["source_experiment_class"]) == {"core_ssl_benchmark", "label_efficiency"} + assert set(view["experiment_class"]) == {"label_efficiency"} + assert set(view["comparison_view"]) == {"label_efficiency_with_core_endpoint"} + + +def test_build_statistical_tests_df_adds_classical_context_rows(): + mod = importlib.import_module("scripts.export_results") + + rows = [] + for experiment_class, paradigm, offset in [ + ("core_ssl_benchmark", "mae", 0.10), + ("core_ssl_benchmark", "supervised", 0.0), + ("classical_baselines", "xgboost", 0.03), + ("classical_baselines", "gru_d", 0.05), + ]: + for seed in [42, 123]: + for task in ["mortality_24h", "aki_kdigo"]: + rows.append(_row(experiment_class, paradigm, seed, task, offset)) + + stats_df = mod.build_statistical_tests_df(pd.DataFrame(rows)) + contextual = stats_df[stats_df["experiment_class"] == "classical_context_full"] + pairs = { + tuple(sorted((row["paradigm_a"], row["paradigm_b"]))) for _, row in contextual.iterrows() + } + + assert ("mae", "xgboost") in pairs + assert ("gru_d", "supervised") in pairs + assert ("gru_d", "xgboost") not in pairs + + +def test_build_ts2vec_vs_core_contrastive_df_compares_across_classes(): + mod = importlib.import_module("scripts.export_results") + + rows = [] + for experiment_class, paradigm, offset in [ + ("ts2vec_extension", "ts2vec", 0.04), + ("core_ssl_benchmark", "contrastive", 0.0), + ]: + for seed in [42, 123]: + for task in ["mortality_24h", "aki_kdigo"]: + rows.append(_row(experiment_class, paradigm, seed, task, offset)) + + comparison_df = mod.build_ts2vec_vs_core_contrastive_df(pd.DataFrame(rows)) + + assert len(comparison_df) == 1 + row = comparison_df.iloc[0] + assert row["comparison_type"] == "ts2vec_vs_core_contrastive" + assert row["paradigm_a"] == "ts2vec" + assert row["paradigm_b"] == "contrastive" + assert row["n_pairs"] == 4 + + +def test_validate_flags_wrong_fixed_seed_set_with_correct_count(): + mod = importlib.import_module("scripts.export_results") + + aggregated_df = pd.DataFrame( + [ + { + "experiment_class": "core_ssl_benchmark", + "paradigm": "mae", + "dataset": "miiv", + "task": "mortality_24h", + "protocol": "full_finetune", + "n_seeds": 5, + "seed_list": "[1, 2, 3, 4, 5]", + } + ] + ) + warnings = mod.validate(pd.DataFrame([{"test/auprc": 0.5}]), aggregated_df) + + assert any("expected seed set" in warning for warning in warnings) + assert any("unexpected=[1, 2, 3, 4, 5]" in warning for warning in warnings) + + +def test_validate_flags_absent_expected_matrix_rows(): + mod = importlib.import_module("scripts.export_results") + + expected = mod.build_expected_matrix_df( + experiment_class=["core_ssl_benchmark"], + paradigm=["mae"], + dataset=["miiv"], + phase=["finetune"], + ) + present = expected[ + (expected["task"] == "mortality_24h") & (expected["protocol"] == "full_finetune") + ].copy() + present["wandb_run_id"] = [f"run-{seed}" for seed in present["seed"]] + present["test/auprc"] = 0.5 + aggregated = mod.build_aggregated_df(present) + + warnings = mod.validate(present, aggregated, expected_matrix_df=expected) + joined = "\n".join(warnings) + + assert "expected matrix evaluation rows are absent" in joined + assert "missing experiment_class=core_ssl_benchmark" in joined + + +def test_build_aggregated_df_records_per_metric_non_null_counts(): + mod = importlib.import_module("scripts.export_results") + + rows = [ + {**_row("core_ssl_benchmark", "mae", 42), "test/auroc": 0.7, "wandb_run_id": "run-42"}, + { + **_row("core_ssl_benchmark", "mae", 123), + "test/auprc": None, + "test/auroc": 0.8, + "wandb_run_id": "run-123", + }, + ] + + aggregated = mod.build_aggregated_df(pd.DataFrame(rows)) + + assert aggregated.iloc[0]["n_seeds"] == 2 + assert aggregated.iloc[0]["test/auprc/n"] == 1 + assert aggregated.iloc[0]["test/auroc/n"] == 2 + assert aggregated.iloc[0]["test/auroc/ci95_lower"] < aggregated.iloc[0]["test/auroc/mean"] + assert aggregated.iloc[0]["test/auroc/ci95_upper"] > aggregated.iloc[0]["test/auroc/mean"] + assert pd.isna(aggregated.iloc[0]["test/auprc/ci95_lower"]) + + +def test_validate_warns_when_evaluation_row_missing_primary_metric(): + mod = importlib.import_module("scripts.export_results") + + warnings = mod.validate( + pd.DataFrame( + [ + { + "wandb_run_id": "run-1", + "paradigm": "mae", + "dataset": "miiv", + "task": "mortality_24h", + "seed": 42, + "phase": "finetune", + "test/auprc": None, + } + ] + ), + pd.DataFrame(), + ) + + assert any("missing their primary test metric" in warning for warning in warnings) + + +def test_extract_run_exports_los_regression_fairness_keys(): + mod = importlib.import_module("scripts.export_results") + + run = DummyRun( + config={ + "experiment_class": "core_ssl_benchmark", + "dataset": "miiv", + "paradigm": "mae", + "seed": 42, + "encoder": {"d_model": 64, "n_layers": 2}, + "training": {"freeze_encoder": False}, + "task": {"task_name": "los_remaining", "task_type": "regression"}, + }, + tags=["experiment_class:core_ssl_benchmark", "phase:finetune"], + summary={ + "test/mae": 1.4, + "fairness/gender/worst_group_mse": 2.5, + "fairness/gender/worst_group_mae": 1.2, + "fairness/gender/mse_gap": 0.7, + "fairness/gender/mae_gap": 0.3, + "fairness/gender/per_group_mse/F": 1.8, + "fairness/gender/per_group_mse/M": 2.5, + "fairness/gender/per_group_mae/F": 0.9, + "fairness/gender/per_group_mae/M": 1.2, + }, + name="los_fairness", + ) + + row = mod.extract_run(run, mod.ALL_METRICS) + assert row["fairness/gender/worst_group_mse"] == pytest.approx(2.5) + assert row["fairness/gender/mse_gap"] == pytest.approx(0.7) + assert row["fairness/gender/per_group_mse/M"] == pytest.approx(2.5) + + agg = mod.build_aggregated_df(pd.DataFrame([row])) + assert agg.iloc[0]["fairness/gender/worst_group_mse/mean"] == pytest.approx(2.5) + assert agg.iloc[0]["fairness/gender/per_group_mae/F/mean"] == pytest.approx(0.9) + + +def test_validate_warns_on_missing_or_failed_checkpoint_provenance(): + mod = importlib.import_module("scripts.export_results") + fairness_keys = [ + f"fairness/{attr}/{metric_name}" + for attr in mod.FAIRNESS_ATTRIBUTES + for metric_name in mod.BINARY_FAIRNESS_REQUIRED_METRICS + ] + + per_seed_df = pd.DataFrame( + [ + { + "wandb_run_id": "run-missing", + "paradigm": "mae", + "dataset": "miiv", + "task": "mortality_24h", + "seed": 42, + "phase": "finetune", + "test/auprc": 0.4, + }, + { + "wandb_run_id": "run-failed", + "paradigm": "jepa", + "dataset": "miiv", + "task": "mortality_24h", + "seed": 123, + "phase": "finetune", + "test/auprc": 0.4, + "_eval_checkpoint_source": "failed", + "_best_ckpt_error": "no best checkpoint", + }, + { + "wandb_run_id": "run-bad-load", + "paradigm": "contrastive", + "dataset": "miiv", + "task": "mortality_24h", + "seed": 456, + "phase": "finetune", + "test/auprc": 0.4, + "_eval_checkpoint_source": "best", + "_best_ckpt_path": "outputs/run/checkpoints/best.ckpt", + "_best_ckpt_load_ok": False, + }, + { + "wandb_run_id": "run-xgb", + "paradigm": "xgboost", + "dataset": "miiv", + "task": "mortality_24h", + "seed": 42, + "phase": "baseline", + "test/auprc": 0.4, + "_output_dir": "outputs/run-xgb", + mod.FAIRNESS_SUMMARY_KEY_COLUMN: json.dumps(fairness_keys), + **_fresh_fairness_metadata( + mod, + artifact_path="outputs/run-xgb/xgboost_model.json", + artifact_source="xgboost_model", + checkpoint_source="xgboost_model", + ), + }, + ] + ) + + warnings = mod.validate(per_seed_df, pd.DataFrame()) + joined = "\n".join(warnings) + + assert "checkpoint provenance" in joined + assert "run-missing" in joined + assert "missing _eval_checkpoint_source" in joined + assert "run-failed" in joined + assert "run-bad-load" in joined + assert "run-xgb" not in joined + + +def test_validate_warns_when_fairness_summary_metrics_are_missing(): + mod = importlib.import_module("scripts.export_results") + + row = { + **_row("core_ssl_benchmark", "mae", 42), + "wandb_run_id": "run-no-fairness", + "_eval_checkpoint_source": "final", + } + per_seed_df = pd.DataFrame([row]) + aggregated_df = mod.build_aggregated_df(per_seed_df) + + warnings = mod.validate(per_seed_df, aggregated_df, expected_seeds={42}) + joined = "\n".join(warnings) + + assert "missing required fairness summary metrics" in joined + assert "fairness/gender/n_metric_valid_groups" in joined + + +def test_validate_accepts_written_nan_fairness_summary_metrics(): + mod = importlib.import_module("scripts.export_results") + + required_keys = [ + f"fairness/{attr}/{metric_name}" + for attr in mod.FAIRNESS_ATTRIBUTES + for metric_name in mod.BINARY_FAIRNESS_REQUIRED_METRICS + ] + row = { + **_row("core_ssl_benchmark", "mae", 42), + "wandb_run_id": "run-nan-fairness", + "_eval_checkpoint_source": "final", + mod.FAIRNESS_SUMMARY_KEY_COLUMN: json.dumps(required_keys), + } + for key in required_keys: + row[key] = float("nan") + + per_seed_df = pd.DataFrame([row]) + aggregated_df = mod.build_aggregated_df(per_seed_df) + warnings = mod.validate(per_seed_df, aggregated_df, expected_seeds={42}) + + assert not any("missing required fairness summary metrics" in warning for warning in warnings) + + +def test_validate_warns_when_fairness_summary_metadata_is_missing_or_stale(): + mod = importlib.import_module("scripts.export_results") + + row = { + **_row("core_ssl_benchmark", "mae", 42), + "wandb_run_id": "run-stale-fairness", + "_eval_checkpoint_source": "final", + "_output_dir": "outputs/run-stale-fairness", + **_fresh_fairness_metadata( + mod, + artifact_path="outputs/old-run/checkpoints/last.ckpt", + ), + } + _add_binary_fairness_summary(row, mod, value=0.0) + + per_seed_df = pd.DataFrame([row]) + aggregated_df = mod.build_aggregated_df(per_seed_df) + warnings = mod.validate(per_seed_df, aggregated_df, expected_seeds={42}) + joined = "\n".join(warnings) + + assert "missing or stale fairness summary metadata" in joined + assert "artifact path mismatch" in joined + assert "run-stale-fairness" in joined + + +def test_validate_accepts_fresh_fairness_summary_metadata(): + mod = importlib.import_module("scripts.export_results") + + row = { + **_row("core_ssl_benchmark", "mae", 42), + "wandb_run_id": "run-fresh-fairness", + "_eval_checkpoint_source": "final", + "_output_dir": "outputs/run-fresh-fairness", + **_fresh_fairness_metadata( + mod, + artifact_path="outputs/run-fresh-fairness/checkpoints/last.ckpt", + ), + } + _add_binary_fairness_summary(row, mod, value=0.0) + + per_seed_df = pd.DataFrame([row]) + aggregated_df = mod.build_aggregated_df(per_seed_df) + warnings = mod.validate(per_seed_df, aggregated_df, expected_seeds={42}) + + assert not warnings + + +def test_validate_warns_when_fairness_artifact_digest_differs_from_test_artifact(): + mod = importlib.import_module("scripts.export_results") + + row = { + **_row("core_ssl_benchmark", "mae", 42), + "wandb_run_id": "run-digest-mismatch", + "_eval_checkpoint_source": "final", + "_output_dir": "outputs/run-digest-mismatch", + **_fresh_fairness_metadata( + mod, + artifact_path="outputs/run-digest-mismatch/checkpoints/last.ckpt", + artifact_sha256="train-digest", + ), + mod.FAIRNESS_ARTIFACT_SHA256_KEY: "fairness-digest", + } + _add_binary_fairness_summary(row, mod, value=0.0) + + per_seed_df = pd.DataFrame([row]) + aggregated_df = mod.build_aggregated_df(per_seed_df) + warnings = mod.validate(per_seed_df, aggregated_df, expected_seeds={42}) + joined = "\n".join(warnings) + + assert "artifact sha256 mismatch" in joined + assert "run-digest-mismatch" in joined + + +def test_validate_does_not_require_eicu_race_fairness(): + mod = importlib.import_module("scripts.export_results") + + attrs = ["gender", "age_group"] + required_keys = [ + f"fairness/{attr}/{metric_name}" + for attr in attrs + for metric_name in mod.BINARY_FAIRNESS_REQUIRED_METRICS + ] + row = { + **_row("core_ssl_benchmark", "mae", 42), + "dataset": "eicu", + "wandb_run_id": "run-eicu-fairness", + "_eval_checkpoint_source": "final", + mod.FAIRNESS_SUMMARY_KEY_COLUMN: json.dumps(required_keys), + } + for key in required_keys: + row[key] = 0.0 + + per_seed_df = pd.DataFrame([row]) + aggregated_df = mod.build_aggregated_df(per_seed_df) + warnings = mod.validate(per_seed_df, aggregated_df, expected_seeds={42}) + + assert not any("missing required fairness summary metrics" in warning for warning in warnings) + + +def test_filter_thesis_tasks_excludes_optional_mortality_by_default(): + mod = importlib.import_module("scripts.export_results") + + df = pd.DataFrame( + [ + {"task": "mortality_24h", "wandb_run_id": "main"}, + {"task": "sepsis", "wandb_run_id": "extension"}, + ] + ) + + filtered = mod.filter_thesis_tasks(df) + unfiltered = mod.filter_thesis_tasks(df, include_extension_tasks=True) + + assert filtered["wandb_run_id"].tolist() == ["main"] + assert unfiltered["wandb_run_id"].tolist() == ["main", "extension"] + + +def test_parse_args_exposes_extension_task_escape_hatch(monkeypatch): + mod = importlib.import_module("scripts.export_results") + + monkeypatch.setattr( + sys, + "argv", + ["export_results.py", "--revision", "thesis-v1", "--include-extension-tasks"], + ) + + args = mod.parse_args() + + assert args.include_extension_tasks is True + + +def test_parse_args_uses_revision_env_when_cli_omits_it(monkeypatch): + mod = importlib.import_module("scripts.export_results") + + monkeypatch.setenv("REVISION", "thesis-v1") + monkeypatch.delenv("WANDB_REVISION", raising=False) + monkeypatch.delenv("WANDB_PROJECT", raising=False) + monkeypatch.setattr(sys, "argv", ["export_results.py"]) + + args = mod.parse_args() + + assert args.revision == ["thesis-v1"] + assert args.project == "slices-thesis" + + +def test_parse_args_exposes_allow_incomplete_escape_hatch(monkeypatch): + mod = importlib.import_module("scripts.export_results") + + monkeypatch.setattr( + sys, + "argv", + ["export_results.py", "--revision", "thesis-v1", "--allow-incomplete"], + ) + + assert mod.parse_args().allow_incomplete is True + + +def test_parse_args_exposes_extraction_failure_escape_hatch(monkeypatch): + mod = importlib.import_module("scripts.export_results") + + monkeypatch.setattr( + sys, + "argv", + ["export_results.py", "--revision", "thesis-v1", "--allow-extraction-failures"], + ) + + assert mod.parse_args().allow_extraction_failures is True + + +def test_parse_args_exposes_duplicate_fingerprint_escape_hatch(monkeypatch): + mod = importlib.import_module("scripts.export_results") + + monkeypatch.setattr( + sys, + "argv", + ["export_results.py", "--revision", "thesis-v1", "--allow-duplicate-fingerprints"], + ) + + assert mod.parse_args().allow_duplicate_fingerprints is True + + +def test_parse_args_rejects_multiple_revisions_by_default(monkeypatch): + mod = importlib.import_module("scripts.export_results") + + monkeypatch.setattr(sys, "argv", ["export_results.py", "--revision", "v1", "v2"]) + + with pytest.raises(SystemExit) as excinfo: + mod.parse_args() + + assert excinfo.value.code == 2 + + +def test_parse_args_allows_multiple_revisions_for_exploratory_export(monkeypatch): + mod = importlib.import_module("scripts.export_results") + + monkeypatch.setattr( + sys, + "argv", + ["export_results.py", "--revision", "v1", "v2", "--allow-multiple-revisions"], + ) + + args = mod.parse_args() + + assert args.revision == ["v1", "v2"] + assert args.allow_multiple_revisions is True + + +def test_validate_warns_when_group_mixes_revisions(): + mod = importlib.import_module("scripts.export_results") + + rows = [] + for seed, revision in [(42, "v1"), (123, "v2")]: + rows.append( + { + **_row("core_ssl_benchmark", "mae", seed), + "wandb_run_id": f"run-{seed}", + "revision": revision, + "_eval_checkpoint_source": "final", + } + ) + + per_seed_df = pd.DataFrame(rows) + aggregated_df = mod.build_aggregated_df(per_seed_df) + warnings = mod.validate(per_seed_df, aggregated_df, expected_seeds={42, 123}) + + assert any("mixes multiple revisions" in warning for warning in warnings) + + +def test_validate_warns_when_revision_mixes_launch_commits(): + mod = importlib.import_module("scripts.export_results") + + rows = [] + for seed, launch_commit in [(42, "commit-a"), (123, "commit-b")]: + rows.append( + { + **_row("core_ssl_benchmark", "mae", seed), + "wandb_run_id": f"run-{seed}", + "revision": "thesis-v1", + "launch_commit": launch_commit, + "_eval_checkpoint_source": "final", + } + ) + + per_seed_df = pd.DataFrame(rows) + aggregated_df = mod.build_aggregated_df(per_seed_df) + warnings = mod.validate(per_seed_df, aggregated_df, expected_seeds={42, 123}) + + assert any( + "revision=thesis-v1 mixes multiple launch commits" in warning for warning in warnings + ) + + +def test_validate_warns_when_revision_lacks_launch_commit(): + mod = importlib.import_module("scripts.export_results") + + per_seed_df = pd.DataFrame( + [ + { + **_row("core_ssl_benchmark", "mae", 42), + "wandb_run_id": "run-42", + "revision": "thesis-v1", + "launch_commit": None, + "_eval_checkpoint_source": "final", + } + ] + ) + aggregated_df = mod.build_aggregated_df(per_seed_df) + warnings = mod.validate(per_seed_df, aggregated_df, expected_seeds={42}) + + assert any( + "revision=thesis-v1 has 1 runs missing launch_commit" in warning for warning in warnings + ) + + +def test_main_exits_nonzero_when_no_runs_match(monkeypatch, tmp_path): + mod = importlib.import_module("scripts.export_results") + + monkeypatch.setattr( + sys, + "argv", + ["export_results.py", "--revision", "thesis-v1", "--output-dir", str(tmp_path)], + ) + monkeypatch.setattr(mod, "fetch_all_runs", lambda **_: []) + + with pytest.raises(SystemExit) as excinfo: + mod.main() + + assert excinfo.value.code == 1 + + +def test_main_does_not_write_parquet_when_validation_fails(monkeypatch, tmp_path): + mod = importlib.import_module("scripts.export_results") + + monkeypatch.setattr( + sys, + "argv", + ["export_results.py", "--revision", "thesis-v1", "--output-dir", str(tmp_path)], + ) + per_seed = pd.DataFrame([_row("core_ssl_benchmark", "mae", 42)]) + + monkeypatch.setattr(mod, "fetch_all_runs", lambda **_: [object()]) + monkeypatch.setattr(mod, "build_per_seed_df", lambda *_, **__: per_seed) + monkeypatch.setattr(mod, "filter_thesis_tasks", lambda df, **_: df) + monkeypatch.setattr(mod, "build_aggregated_df", lambda _: pd.DataFrame([{"n_seeds": 1}])) + monkeypatch.setattr(mod, "build_label_efficiency_curve_df", lambda _: pd.DataFrame()) + monkeypatch.setattr(mod, "build_capacity_study_comparison_df", lambda _: pd.DataFrame()) + monkeypatch.setattr(mod, "build_classical_context_df", lambda _: pd.DataFrame()) + monkeypatch.setattr(mod, "build_statistical_tests_df", lambda _: pd.DataFrame()) + monkeypatch.setattr(mod, "build_ts2vec_vs_core_contrastive_df", lambda _: pd.DataFrame()) + monkeypatch.setattr(mod, "build_expected_matrix_df", lambda **_: pd.DataFrame()) + monkeypatch.setattr(mod, "validate", lambda *_, **__: ["WARNING: incomplete"]) + + with pytest.raises(SystemExit) as excinfo: + mod.main() + + assert excinfo.value.code == 1 + assert not list(tmp_path.glob("*.parquet")) + + +def test_parse_args_requires_revision_without_cli_or_env(monkeypatch): + mod = importlib.import_module("scripts.export_results") + + monkeypatch.delenv("REVISION", raising=False) + monkeypatch.delenv("WANDB_REVISION", raising=False) + monkeypatch.setattr(sys, "argv", ["export_results.py"]) + + with pytest.raises(SystemExit) as excinfo: + mod.parse_args() + + assert excinfo.value.code == 2 diff --git a/tests/test_extractor_config.py b/tests/test_extractor_config.py index 7a951cd..70ce8d8 100644 --- a/tests/test_extractor_config.py +++ b/tests/test_extractor_config.py @@ -2,6 +2,8 @@ import pytest from pydantic import ValidationError + +from slices.constants import THESIS_TASKS from slices.data.extractors.base import ExtractorConfig @@ -24,10 +26,10 @@ def test_default_values(self): config = ExtractorConfig(parquet_root="/path/to/parquet") assert config.output_dir == "data/processed" - assert config.seq_length_hours == 48 + assert config.seq_length_hours == 24 assert config.feature_set == "core" assert config.tasks_dir is None - assert config.min_stay_hours == 48 + assert config.min_stay_hours == 24 assert "mortality_24h" in config.tasks def test_custom_values(self): @@ -175,11 +177,10 @@ class TestExtractorConfigTasksParameter: """Test tasks parameter handling.""" def test_default_tasks_list(self): - """Test that default tasks list contains expected mortality tasks.""" + """Test that default tasks list contains the thesis task surface.""" config = ExtractorConfig(parquet_root="/path") - assert "mortality_24h" in config.tasks - assert "mortality_hospital" in config.tasks + assert config.tasks == list(THESIS_TASKS) def test_custom_tasks_list(self): """Test custom tasks list.""" diff --git a/tests/test_extractor_integration.py b/tests/test_extractor_integration.py index 84e774a..71be804 100644 --- a/tests/test_extractor_integration.py +++ b/tests/test_extractor_integration.py @@ -5,6 +5,7 @@ import polars as pl import pytest + from slices.data.extractors.base import BaseExtractor, ExtractorConfig from slices.data.labels import LabelConfig diff --git a/tests/test_factories.py b/tests/test_factories.py index 0793ece..3bb4679 100644 --- a/tests/test_factories.py +++ b/tests/test_factories.py @@ -6,6 +6,7 @@ import pytest import torch import torch.nn as nn + from slices.models.encoders.factory import ( ENCODER_CONFIG_REGISTRY, ENCODER_REGISTRY, @@ -351,19 +352,19 @@ def test_config_registry_contains_mae(self): def test_build_ssl_objective_mae(self): """build_ssl_objective should create MAEObjective correctly.""" - from slices.models.encoders import ( - ObservationTransformerConfig, - ObservationTransformerEncoder, - ) + from slices.models.encoders import TransformerConfig, TransformerEncoder - encoder_config = ObservationTransformerConfig( + encoder_config = TransformerConfig( d_input=35, d_model=64, n_layers=2, n_heads=4, + d_ff=256, pooling="none", + obs_aware=True, + max_seq_length=48, ) - encoder = ObservationTransformerEncoder(encoder_config) + encoder = TransformerEncoder(encoder_config) mae_config = MAEConfig( name="mae", @@ -398,19 +399,19 @@ def test_get_ssl_config_class_unknown_raises(self): def test_built_ssl_objective_forward_pass(self): """Built SSL objective should perform forward pass correctly.""" - from slices.models.encoders import ( - ObservationTransformerConfig, - ObservationTransformerEncoder, - ) + from slices.models.encoders import TransformerConfig, TransformerEncoder - encoder_config = ObservationTransformerConfig( + encoder_config = TransformerConfig( d_input=35, d_model=64, n_layers=2, n_heads=4, + d_ff=256, pooling="none", + obs_aware=True, + max_seq_length=48, ) - encoder = ObservationTransformerEncoder(encoder_config) + encoder = TransformerEncoder(encoder_config) mae_config = MAEConfig( name="mae", diff --git a/tests/test_fairness.py b/tests/test_fairness.py index 0cf7ee0..b06ff92 100644 --- a/tests/test_fairness.py +++ b/tests/test_fairness.py @@ -10,6 +10,7 @@ import pytest import torch + from slices.eval.fairness import ( FairnessConfig, compute_demographic_parity, diff --git a/tests/test_fairness_evaluator.py b/tests/test_fairness_evaluator.py index d1466fd..fba4770 100644 --- a/tests/test_fairness_evaluator.py +++ b/tests/test_fairness_evaluator.py @@ -13,7 +13,8 @@ import polars as pl import pytest import torch -from slices.eval.fairness_evaluator import FairnessEvaluator + +from slices.eval.fairness_evaluator import FairnessEvaluator, flatten_fairness_report def make_static_df( @@ -122,6 +123,68 @@ def test_selects_minimum(self): valid_aurocs = [v for v in per_group.values() if v == v] # Filter NaN assert worst == pytest.approx(min(valid_aurocs)) + def test_one_class_group_makes_worst_group_discrimination_non_comparable(self): + """A size-valid one-class group should not be omitted from worst-group accounting.""" + static_df = make_static_df( + n=6, + genders=["M", "M", "M", "F", "F", "F"], + ages=[50.0] * 6, + ) + stay_ids = list(range(6)) + labels = torch.tensor([1.0, 1.0, 1.0, 0.0, 1.0, 0.0]) + predictions = torch.tensor([0.9, 0.8, 0.7, 0.2, 0.8, 0.3]) + + evaluator = FairnessEvaluator( + static_df, + protected_attributes=["gender"], + min_subgroup_size=1, + ) + report = evaluator.evaluate(predictions, labels, stay_ids) + gender_report = report["gender"] + + assert gender_report["n_valid_groups"] == 2 + assert gender_report["n_metric_valid_groups"] == 1 + assert gender_report["n_single_class_groups"] == 1 + assert gender_report["per_group_auroc"]["M"] != gender_report["per_group_auroc"]["M"] + assert gender_report["worst_group_auroc"] != gender_report["worst_group_auroc"] + assert gender_report["worst_group_auprc"] != gender_report["worst_group_auprc"] + assert gender_report["auroc_gap"] != gender_report["auroc_gap"] + assert gender_report["auprc_gap"] != gender_report["auprc_gap"] + + +class TestRegressionFairness: + """Tests for regression-task fairness metrics.""" + + def test_singleton_prediction_dimension_is_squeezed(self): + """Regression fairness should treat (N, 1) predictions as samplewise outputs.""" + static_df = make_static_df( + n=4, + genders=["M", "M", "F", "F"], + ages=[50.0, 50.0, 50.0, 50.0], + ) + stay_ids = [0, 1, 2, 3] + labels = torch.tensor([0.0, 1.0, 0.0, 1.0]) + predictions = torch.tensor([[0.0], [1.5], [0.0], [0.5]]) + + evaluator = FairnessEvaluator( + static_df, + protected_attributes=["gender"], + min_subgroup_size=1, + task_type="regression", + ) + report = evaluator.evaluate(predictions, labels, stay_ids) + + per_group_mse = report["gender"]["per_group_mse"] + per_group_mae = report["gender"]["per_group_mae"] + assert per_group_mse["M"] == pytest.approx(0.125) + assert per_group_mse["F"] == pytest.approx(0.125) + assert per_group_mae["M"] == pytest.approx(0.25) + assert per_group_mae["F"] == pytest.approx(0.25) + assert report["gender"]["worst_group_mse"] == pytest.approx(0.125) + assert report["gender"]["worst_group_mae"] == pytest.approx(0.25) + assert report["gender"]["mse_gap"] == pytest.approx(0.0) + assert report["gender"]["mae_gap"] == pytest.approx(0.0) + class TestAgeBinning: """Tests for age binning.""" @@ -155,6 +218,24 @@ def test_boundary_ages(self): assert groups[5].item() == 2 # 79 -> 65-79 assert groups[6].item() == 3 # 80 -> 80+ + def test_null_and_nan_ages_are_unknown(self): + """Null and numeric NaN ages should not fall into the youngest bin.""" + static_df = pl.DataFrame( + { + "stay_id": [1, 2, 3], + "patient_id": [1, 2, 3], + "gender": ["M", "F", "M"], + "age": [None, float("nan"), 30.0], + "los_days": [5.0, 5.0, 5.0], + } + ) + evaluator = FairnessEvaluator(static_df, protected_attributes=["age_group"]) + + group_ids, group_names, _ = evaluator._encode_attribute([1, 2, 3], "age_group") + + assert group_ids.tolist() == [-1, -1, 0] + assert group_names[-1] == "unknown" + class TestAttributeAutoDetection: """Tests for attribute auto-detection.""" @@ -187,6 +268,26 @@ def test_missing_columns_filtered(self): evaluator = FairnessEvaluator(static_df, protected_attributes=["gender", "race"]) assert evaluator._available_attributes == [] + def test_unknown_gender_values_are_excluded(self): + """String unknown gender labels should not form a fairness subgroup.""" + static_df = pl.DataFrame( + { + "stay_id": [1, 2, 3, 4], + "patient_id": [1, 2, 3, 4], + "gender": ["M", "Unknown", "F", "not specified"], + "age": [50.0, 50.0, 50.0, 50.0], + "los_days": [5.0, 5.0, 5.0, 5.0], + } + ) + evaluator = FairnessEvaluator( + static_df, protected_attributes=["gender"], min_subgroup_size=1 + ) + + group_ids, group_names, _ = evaluator._encode_attribute([1, 2, 3, 4], "gender") + + assert group_ids.tolist().count(-1) == 2 + assert set(group_names.values()) == {"M", "F", "unknown"} + class TestMinSubgroupSize: """Tests for min_subgroup_size filtering.""" @@ -232,6 +333,49 @@ def test_all_groups_large_enough(self): assert "gender" in report assert len(report["gender"]["per_group_auroc"]) == 2 + def test_threshold_counts_unique_patients_not_stays(self): + """Min subgroup size should be enforced on patients, not repeated stays.""" + rows = [] + stay_id = 0 + for patient_id in range(40): + for _ in range(2): + rows.append( + { + "stay_id": stay_id, + "patient_id": f"m-{patient_id}", + "gender": "M", + "age": 50.0, + "los_days": 5.0, + } + ) + stay_id += 1 + for patient_id in range(55): + rows.append( + { + "stay_id": stay_id, + "patient_id": f"f-{patient_id}", + "gender": "F", + "age": 50.0, + "los_days": 5.0, + } + ) + stay_id += 1 + + static_df = pl.DataFrame(rows) + predictions = torch.tensor([0.8 if i % 2 == 0 else 0.2 for i in range(len(rows))]) + labels = torch.tensor([1.0 if i % 2 == 0 else 0.0 for i in range(len(rows))]) + stay_ids = static_df["stay_id"].to_list() + + evaluator = FairnessEvaluator( + static_df, + protected_attributes=["gender"], + min_subgroup_size=50, + ) + report = evaluator.evaluate(predictions, labels, stay_ids) + + # Male has 80 stays but only 40 patients, so the attribute should be skipped. + assert "gender" not in report + class TestEvaluateStructure: """Tests for evaluate() return structure.""" @@ -260,6 +404,7 @@ def test_returns_correct_keys(self): assert "equalized_odds_diff" in gender_report assert "disparate_impact_ratio" in gender_report assert "n_valid_groups" in gender_report + assert "n_metric_valid_groups" in gender_report assert "group_sizes" in gender_report def test_multiple_attributes(self): @@ -293,6 +438,83 @@ def test_multiple_attributes(self): assert "gender" in report assert "age_group" in report + def test_group_sizes_report_patients_and_samples_separately(self): + """The report should expose patient counts separately from stay counts.""" + rows = [] + stay_id = 0 + for patient_id in range(55): + for _ in range(2): + rows.append( + { + "stay_id": stay_id, + "patient_id": f"m-{patient_id}", + "gender": "M", + "age": 50.0, + "los_days": 5.0, + } + ) + stay_id += 1 + for patient_id in range(55): + for _ in range(2): + rows.append( + { + "stay_id": stay_id, + "patient_id": f"f-{patient_id}", + "gender": "F", + "age": 50.0, + "los_days": 5.0, + } + ) + stay_id += 1 + + static_df = pl.DataFrame(rows) + predictions = torch.tensor([0.8 if i % 2 == 0 else 0.2 for i in range(len(rows))]) + labels = torch.tensor([1.0 if i % 2 == 0 else 0.0 for i in range(len(rows))]) + stay_ids = static_df["stay_id"].to_list() + + evaluator = FairnessEvaluator( + static_df, + protected_attributes=["gender"], + min_subgroup_size=50, + ) + report = evaluator.evaluate(predictions, labels, stay_ids) + + assert report["gender"]["group_sizes"]["M"] == 55 + assert report["gender"]["group_sample_sizes"]["M"] == 110 + + def test_race_uses_canonical_miiv_only_schema_for_combined(self): + """Combined race fairness should only include canonical MIMIC race bins.""" + static_df = pl.DataFrame( + { + "stay_id": [1, 2, 3, 4], + "patient_id": ["a", "b", "c", "d"], + "gender": ["M", "F", "M", "F"], + "age": [50.0, 60.0, 55.0, 65.0], + "race": [ + "WHITE", + "BLACK/AFRICAN AMERICAN", + "HISPANIC/LATINO", + "ASIAN", + ], + "source_dataset": ["miiv", "miiv", "eicu", "eicu"], + "los_days": [5.0, 5.0, 5.0, 5.0], + } + ) + predictions = torch.tensor([0.9, 0.1, 0.8, 0.2]) + labels = torch.tensor([1.0, 0.0, 1.0, 0.0]) + stay_ids = [1, 2, 3, 4] + + evaluator = FairnessEvaluator( + static_df, + protected_attributes=["race"], + min_subgroup_size=1, + dataset_name="combined", + ) + report = evaluator.evaluate(predictions, labels, stay_ids) + + assert "race" in report + assert set(report["race"]["per_group_auroc"]) == {"White", "Black"} + class TestPrintReport: """Tests for print_report().""" @@ -315,3 +537,24 @@ def test_print_without_error(self, capsys): captured = capsys.readouterr() assert "Fairness Evaluation Report" in captured.out assert "gender" in captured.out + + +class TestFlattenFairnessReport: + """Tests for fairness report flattening.""" + + def test_flattens_nested_metrics_for_logging(self): + report = { + "gender": { + "worst_group_auroc": 0.71, + "per_group_auroc": {"M": 0.83, "F": 0.71}, + "group_sizes": {"M": 55, "F": 60}, + } + } + + flat = flatten_fairness_report(report) + + assert flat["fairness/gender/worst_group_auroc"] == pytest.approx(0.71) + assert flat["fairness/gender/per_group_auroc/M"] == pytest.approx(0.83) + assert flat["fairness/gender/per_group_auroc/F"] == pytest.approx(0.71) + assert flat["fairness/gender/group_sizes/M"] == 55 + assert flat["fairness/gender/group_sizes/F"] == 60 diff --git a/tests/test_finetune_module.py b/tests/test_finetune_module.py index abf3e80..99b44d2 100644 --- a/tests/test_finetune_module.py +++ b/tests/test_finetune_module.py @@ -8,6 +8,7 @@ import torch from omegaconf import OmegaConf from pydantic import ValidationError + from slices.models.encoders import TransformerConfig, TransformerEncoder from slices.models.heads import ( LinearTaskHead, @@ -412,12 +413,45 @@ def test_module_creation(self, sample_config): assert module.task_head is not None assert isinstance(module.task_head, MLPTaskHead) + def test_obs_aware_transformer_forward_masks_empty_timesteps(self, sample_config): + """Downstream pooling should exclude fully unobserved obs-aware timesteps.""" + config = OmegaConf.create(OmegaConf.to_container(sample_config, resolve=True)) + config.encoder.obs_aware = True + config.training.use_missing_token = False + module = FineTuneModule(config) + + captured = {} + + def capture_forward(timeseries, mask=None, padding_mask=None): + captured["padding_mask"] = padding_mask.detach().clone() + return torch.zeros(timeseries.size(0), module.encoder.get_output_dim()) + + module.encoder.forward = capture_forward + + timeseries = torch.randn(2, 4, 35) + obs_mask = torch.ones(2, 4, 35, dtype=torch.bool) + obs_mask[0, 1, :] = False + obs_mask[1, 2, :] = False + + module(timeseries, obs_mask) + + assert torch.equal(captured["padding_mask"], obs_mask.any(dim=-1)) + def test_load_encoder_checkpoint(self, sample_config, encoder_checkpoint): """Test loading encoder from checkpoint.""" module = FineTuneModule(sample_config, checkpoint_path=encoder_checkpoint) assert module.encoder is not None + def test_rejects_ambiguous_checkpoint_sources(self, sample_config): + """Finetuning must not silently prefer one checkpoint source over another.""" + with pytest.raises(ValueError, match="exactly one checkpoint source"): + FineTuneModule( + sample_config, + checkpoint_path="outputs/pretrain/encoder.pt", + pretrain_checkpoint_path="outputs/pretrain/last.ckpt", + ) + def test_load_pretrain_checkpoint_auto_detect_encoder(self, sample_config, tmp_path): """Test that pretrain checkpoint auto-detects encoder architecture. @@ -657,6 +691,7 @@ def test_observation_encoder_not_wrapped(self): # Encoder should be ObservationTransformerEncoder directly, not wrapped assert isinstance(module.encoder, ObservationTransformerEncoder) + assert module.encoder.handles_missingness_intrinsically() def test_smart_encoder_not_wrapped(self): """SMARTEncoder should not be wrapped (handles missingness via MLPEmbedder).""" @@ -700,6 +735,7 @@ def test_smart_encoder_not_wrapped(self): # SMART encoder should NOT be wrapped assert isinstance(module.encoder, SMARTEncoder) + assert module.encoder.handles_missingness_intrinsically() class TestFineTuneModuleGradualUnfreeze: @@ -1297,3 +1333,57 @@ def test_explicit_class_weights(self, base_config): assert module.criterion.weight is not None assert torch.allclose(module.criterion.weight, torch.tensor([0.3, 0.7])) + + +class TestMetricThresholdConfig: + """Test finetune metric construction from eval config.""" + + def test_binary_metric_threshold_is_forwarded(self): + config = OmegaConf.create( + { + "encoder": { + "name": "transformer", + "d_input": 10, + "d_model": 32, + "n_layers": 1, + "n_heads": 2, + "d_ff": 64, + "dropout": 0.0, + "max_seq_length": 24, + "pooling": "mean", + "use_positional_encoding": True, + "prenorm": True, + "activation": "gelu", + "layer_norm_eps": 1e-5, + }, + "task": { + "task_name": "mortality_24h", + "task_type": "binary", + "n_classes": None, + "head_type": "mlp", + "hidden_dims": [16], + "dropout": 0.0, + "activation": "relu", + }, + "training": { + "freeze_encoder": False, + "unfreeze_epoch": None, + }, + "optimizer": { + "name": "adam", + "lr": 1e-3, + }, + "eval": { + "metrics": { + "names": ["accuracy", "precision", "recall", "specificity"], + "threshold": 0.8, + } + }, + } + ) + + module = FineTuneModule(config) + + assert module.train_metrics["accuracy"].threshold == pytest.approx(0.8) + assert module.val_metrics["precision"].threshold == pytest.approx(0.8) + assert module.test_metrics["specificity"].threshold == pytest.approx(0.8) diff --git a/tests/test_fixes.py b/tests/test_fixes.py new file mode 100644 index 0000000..7da56f2 --- /dev/null +++ b/tests/test_fixes.py @@ -0,0 +1,4066 @@ +"""Regression tests for FIXES.md Issues 2-9. + +Tests label manifest validation, normalization stats integrity, +sequence-length override, checkpoint provenance, mortality precision, +and exporter dedup logic. +""" + +import importlib +import os +import re +import shutil +import subprocess +import sys +from datetime import datetime, timedelta, timezone +from pathlib import Path +from types import SimpleNamespace + +import pandas as pd +import polars as pl +import pytest +import torch +import yaml +from omegaconf import OmegaConf + +from slices.constants import THESIS_TASKS +from slices.data.labels import LabelBuilder, LabelBuilderFactory, LabelConfig +from slices.data.labels.aki import AKILabelBuilder +from slices.data.labels.los import LOSLabelBuilder +from slices.data.labels.mortality import MortalityLabelBuilder +from slices.data.tensor_cache import ( + _compute_split_hash, + get_data_fingerprint, + get_preprocessing_fingerprint, + load_cached_tensors, + load_normalization_stats, + save_cached_tensors, + save_normalization_stats, +) +from slices.data.tensor_preprocessing import extract_tensors_from_dataframe + + +class DummyWandbRun: + """Minimal W&B run stub for export pipeline tests.""" + + def __init__(self, run_id, config, tags, name, created_at="2026-04-07T00:00:00"): + self.config = config + self.summary_metrics = {} + self.id = run_id + self.url = f"https://example.com/{run_id}" + self.name = name + self.tags = tags + self.group = "g" + self.created_at = created_at + + +# ============================================================================ +# Issue 3: Label manifest freshness checking +# ============================================================================ + + +class TestLabelManifest: + """Tests for label builder versioning and config hashing.""" + + def test_all_builders_have_semantic_version(self): + """Every label builder subclass must define SEMANTIC_VERSION.""" + assert hasattr(MortalityLabelBuilder, "SEMANTIC_VERSION") + assert hasattr(AKILabelBuilder, "SEMANTIC_VERSION") + assert hasattr(LOSLabelBuilder, "SEMANTIC_VERSION") + + def test_config_hash_deterministic(self): + """Same config produces same hash.""" + config = LabelConfig( + task_name="mortality_24h", + task_type="binary", + prediction_window_hours=24, + observation_window_hours=48, + gap_hours=0, + ) + h1 = LabelBuilder.config_hash(config) + h2 = LabelBuilder.config_hash(config) + assert h1 == h2 + assert len(h1) == 16 + + def test_config_hash_changes_on_window_change(self): + """Different prediction_window_hours produces different hash.""" + config_24 = LabelConfig( + task_name="mortality_24h", + task_type="binary", + prediction_window_hours=24, + observation_window_hours=24, + ) + config_48 = LabelConfig( + task_name="mortality_24h", + task_type="binary", + prediction_window_hours=48, + observation_window_hours=24, + ) + assert LabelBuilder.config_hash(config_24) != LabelBuilder.config_hash(config_48) + + def test_config_hash_changes_on_gap_change(self): + """Different gap_hours produces different hash.""" + config_0 = LabelConfig( + task_name="mortality_24h", + task_type="binary", + prediction_window_hours=24, + gap_hours=0, + ) + config_6 = LabelConfig( + task_name="mortality_24h", + task_type="binary", + prediction_window_hours=24, + gap_hours=6, + ) + assert LabelBuilder.config_hash(config_0) != LabelBuilder.config_hash(config_6) + + def test_config_hash_changes_on_label_sources_change(self): + """Different label_sources must invalidate stale emitted labels.""" + config_stays = LabelConfig( + task_name="mortality_24h", + task_type="binary", + label_sources=["stays"], + ) + config_mortality = LabelConfig( + task_name="mortality_24h", + task_type="binary", + label_sources=["stays", "mortality_info"], + ) + + assert LabelBuilder.config_hash(config_stays) != LabelBuilder.config_hash(config_mortality) + + def test_config_hash_changes_on_supported_datasets_change(self): + """Dataset support affects whether a task emits labels.""" + config_miiv = LabelConfig( + task_name="mortality_24h", + task_type="binary", + supported_datasets=["miiv"], + ) + config_all = LabelConfig( + task_name="mortality_24h", + task_type="binary", + supported_datasets=None, + ) + + assert LabelBuilder.config_hash(config_miiv) != LabelBuilder.config_hash(config_all) + + def test_validate_data_prerequisites_version_mismatch(self, tmp_path): + """Builder version mismatch should raise RuntimeError.""" + from slices.training.utils import validate_data_prerequisites + + # Create metadata with old version + metadata = { + "label_manifest": { + "mortality_24h": { + "builder_version": "0.0.1", # old version + "config_hash": "will_not_match", + } + } + } + (tmp_path / "splits.yaml").write_text("seed: 42\n") + with open(tmp_path / "metadata.yaml", "w") as f: + yaml.dump(metadata, f) + + with pytest.raises(RuntimeError, match="Label builder version mismatch"): + validate_data_prerequisites(str(tmp_path), "miiv", task_names=["mortality_24h"]) + + def test_validate_data_prerequisites_missing_manifest_raises(self, tmp_path): + """Missing label_manifest should abort training.""" + from slices.training.utils import validate_data_prerequisites + + metadata = {"task_names": ["mortality_24h"]} # no label_manifest + (tmp_path / "splits.yaml").write_text("seed: 42\n") + with open(tmp_path / "metadata.yaml", "w") as f: + yaml.dump(metadata, f) + + with pytest.raises(RuntimeError, match="no label_manifest"): + validate_data_prerequisites(str(tmp_path), "miiv", task_names=["mortality_24h"]) + + def test_validate_data_prerequisites_missing_metadata_raises(self, tmp_path): + """Missing metadata.yaml should abort training.""" + from slices.training.utils import validate_data_prerequisites + + (tmp_path / "splits.yaml").write_text("seed: 42\n") + # No metadata.yaml at all + + with pytest.raises(FileNotFoundError, match="metadata.yaml not found"): + validate_data_prerequisites(str(tmp_path), "miiv", task_names=["mortality_24h"]) + + def test_validate_data_prerequisites_missing_task_in_manifest_raises(self, tmp_path): + """Task missing from manifest should abort training.""" + from slices.training.utils import validate_data_prerequisites + + metadata = { + "label_manifest": { + "mortality_hospital": { + "builder_version": "1.0.0", + "config_hash": "abc123", + } + # mortality_24h is NOT in the manifest + } + } + (tmp_path / "splits.yaml").write_text("seed: 42\n") + with open(tmp_path / "metadata.yaml", "w") as f: + yaml.dump(metadata, f) + + with pytest.raises(RuntimeError, match="not found in label manifest"): + validate_data_prerequisites(str(tmp_path), "miiv", task_names=["mortality_24h"]) + + def test_validate_data_prerequisites_matching_passes(self, tmp_path): + """Matching manifest should pass cleanly.""" + from slices.data.utils import get_package_data_dir + from slices.training.utils import validate_data_prerequisites + + with open(get_package_data_dir() / "tasks" / "mortality_24h.yaml") as f: + config = LabelConfig(**yaml.safe_load(f)) + builder = LabelBuilderFactory.create(config) + metadata = { + "label_manifest": { + "mortality_24h": { + "builder_version": builder.SEMANTIC_VERSION, + "config_hash": LabelBuilder.config_hash(config), + } + } + } + (tmp_path / "splits.yaml").write_text("seed: 42\n") + with open(tmp_path / "metadata.yaml", "w") as f: + yaml.dump(metadata, f) + + # Should not raise when the manifest matches the checked task config. + validate_data_prerequisites(str(tmp_path), "miiv", task_names=["mortality_24h"]) + + def test_validate_data_prerequisites_rejects_wrong_benchmark_window(self, tmp_path): + """Declared benchmark windows must match the fixed thesis contract.""" + from slices.training.utils import validate_data_prerequisites + + metadata = {"seq_length_hours": 48} + (tmp_path / "splits.yaml").write_text("seed: 42\n") + with open(tmp_path / "metadata.yaml", "w") as f: + yaml.dump(metadata, f) + + with pytest.raises(RuntimeError, match="seq_length_hours"): + validate_data_prerequisites(str(tmp_path), "miiv") + + def test_validate_data_prerequisites_rejects_blocked_features(self, tmp_path): + """Processed tensors must not contain leakage or future-derived features.""" + from slices.training.utils import validate_data_prerequisites + + metadata = {"feature_names": ["hr", "los_icu"]} + (tmp_path / "splits.yaml").write_text("seed: 42\n") + with open(tmp_path / "metadata.yaml", "w") as f: + yaml.dump(metadata, f) + + with pytest.raises(RuntimeError, match="blocked leakage"): + validate_data_prerequisites(str(tmp_path), "miiv") + + def test_validate_data_prerequisites_uses_provided_task_config(self, tmp_path): + """Active Hydra task configs should override the fallback task-definition tree.""" + from slices.training.utils import validate_data_prerequisites + + active_task = { + "task_name": "mortality_24h", + "task_type": "binary", + "prediction_window_hours": 24, + "observation_window_hours": 6, + "gap_hours": 6, + "label_sources": ["stays", "mortality_info"], + "label_params": {}, + "head_type": "mlp", + "hidden_dims": [64], + "dropout": 0.1, + "activation": "relu", + } + label_config = LabelConfig( + task_name=active_task["task_name"], + task_type=active_task["task_type"], + prediction_window_hours=active_task["prediction_window_hours"], + observation_window_hours=active_task["observation_window_hours"], + gap_hours=active_task["gap_hours"], + label_sources=active_task["label_sources"], + label_params=active_task["label_params"], + ) + builder = LabelBuilderFactory.create(label_config) + metadata = { + "label_manifest": { + "mortality_24h": { + "builder_version": builder.SEMANTIC_VERSION, + "config_hash": LabelBuilder.config_hash(label_config), + } + } + } + (tmp_path / "splits.yaml").write_text("seed: 42\n") + with open(tmp_path / "metadata.yaml", "w") as f: + yaml.dump(metadata, f) + + validate_data_prerequisites( + str(tmp_path), + "miiv", + task_names=["mortality_24h"], + task_configs=[active_task], + ) + + def test_validate_data_prerequisites_missing_task_config_raises(self, tmp_path): + """Missing task config should fail closed instead of skipping validation.""" + from slices.training.utils import validate_data_prerequisites + + metadata = { + "label_manifest": { + "definitely_missing_task": { + "builder_version": "1.0.0", + "config_hash": "abc123", + } + } + } + (tmp_path / "splits.yaml").write_text("seed: 42\n") + with open(tmp_path / "metadata.yaml", "w") as f: + yaml.dump(metadata, f) + + with pytest.raises(FileNotFoundError, match="Task config not found"): + validate_data_prerequisites( + str(tmp_path), + "miiv", + task_names=["definitely_missing_task"], + ) + + +# ============================================================================ +# Issue 9: Hash-keyed normalization stats cache +# ============================================================================ + + +class TestNormalizationStatsCache: + """Tests for hash-keyed normalization stats files.""" + + def test_split_hash_deterministic(self): + """Same indices produce same hash.""" + indices = [0, 5, 10, 15] + h1 = _compute_split_hash(indices, normalize=True) + h2 = _compute_split_hash(indices, normalize=True) + assert h1 == h2 + + def test_split_hash_differs_for_different_indices(self): + """Different indices produce different hash.""" + h1 = _compute_split_hash([0, 1, 2], normalize=True) + h2 = _compute_split_hash([0, 1, 3], normalize=True) + assert h1 != h2 + + def test_split_hash_differs_for_normalize_flag(self): + """Same indices but different normalize flag produce different hash.""" + h1 = _compute_split_hash([0, 1, 2], normalize=True) + h2 = _compute_split_hash([0, 1, 2], normalize=False) + assert h1 != h2 + + def test_split_hash_order_invariant(self): + """Hash is the same regardless of input order (sorted internally).""" + h1 = _compute_split_hash([5, 0, 10], normalize=True) + h2 = _compute_split_hash([0, 5, 10], normalize=True) + assert h1 == h2 + + def test_save_load_roundtrip(self, tmp_path): + """Save and load stats with hash-keyed filenames.""" + means = torch.tensor([1.0, 2.0, 3.0]) + stds = torch.tensor([0.5, 1.0, 1.5]) + names = ["a", "b", "c"] + indices = [0, 1, 2, 3, 4] + + save_normalization_stats(tmp_path, means, stds, names, indices, normalize=True) + + loaded = load_normalization_stats(tmp_path, indices, normalize=True) + assert loaded is not None + assert loaded["feature_means"] == means.tolist() + assert loaded["feature_stds"] == stds.tolist() + + def test_different_splits_different_files(self, tmp_path): + """Different splits write to different files.""" + means = torch.tensor([1.0, 2.0]) + stds = torch.tensor([0.5, 1.0]) + names = ["a", "b"] + + save_normalization_stats(tmp_path, means, stds, names, [0, 1, 2], normalize=True) + save_normalization_stats(tmp_path, means * 2, stds * 2, names, [3, 4, 5], normalize=True) + + stats_files = list(tmp_path.glob("normalization_stats_*.yaml")) + assert len(stats_files) == 2 + + loaded1 = load_normalization_stats(tmp_path, [0, 1, 2], normalize=True) + loaded2 = load_normalization_stats(tmp_path, [3, 4, 5], normalize=True) + assert loaded1["feature_means"] != loaded2["feature_means"] + + def test_legacy_fallback(self, tmp_path): + """Legacy normalization_stats.yaml is loaded when hash-keyed file doesn't exist.""" + indices = [0, 1, 2] + legacy_stats = { + "feature_means": [1.0, 2.0], + "feature_stds": [0.5, 1.0], + "feature_names": ["a", "b"], + "train_indices": indices, + "normalize": True, + "data_fingerprint": get_data_fingerprint(tmp_path), + "preprocessing_fingerprint": get_preprocessing_fingerprint(), + } + with open(tmp_path / "normalization_stats.yaml", "w") as f: + yaml.dump(legacy_stats, f) + + loaded = load_normalization_stats(tmp_path, indices, normalize=True) + assert loaded is not None + assert loaded["feature_means"] == [1.0, 2.0] + + def test_legacy_stats_without_fingerprints_are_ignored(self, tmp_path): + """Legacy stats without freshness fingerprints should not be trusted.""" + indices = [0, 1, 2] + legacy_stats = { + "feature_means": [1.0, 2.0], + "feature_stds": [0.5, 1.0], + "feature_names": ["a", "b"], + "train_indices": indices, + "normalize": True, + } + with open(tmp_path / "normalization_stats.yaml", "w") as f: + yaml.dump(legacy_stats, f) + + loaded = load_normalization_stats(tmp_path, indices, normalize=True) + assert loaded is None + + def test_hash_keyed_stats_invalidated_on_fingerprint_mismatch(self, tmp_path, monkeypatch): + """Hash-keyed normalization stats should be invalidated when fingerprints change.""" + import slices.data.tensor_cache as cache_mod + + monkeypatch.setattr(cache_mod, "get_data_fingerprint", lambda data_dir: "data-v1") + monkeypatch.setattr(cache_mod, "get_preprocessing_fingerprint", lambda: "prep-v1") + + means = torch.tensor([1.0, 2.0]) + stds = torch.tensor([0.5, 1.0]) + save_normalization_stats(tmp_path, means, stds, ["a", "b"], [0, 1, 2], normalize=True) + assert load_normalization_stats(tmp_path, [0, 1, 2], normalize=True) is not None + + monkeypatch.setattr(cache_mod, "get_preprocessing_fingerprint", lambda: "prep-v2") + assert load_normalization_stats(tmp_path, [0, 1, 2], normalize=True) is None + + +class TestRawTensorCacheFreshness: + """Tests for raw tensor cache invalidation.""" + + def test_tensor_cache_invalidated_on_data_fingerprint_mismatch(self, tmp_path, monkeypatch): + """Raw tensor cache should be ignored when processed data changes.""" + import slices.data.tensor_cache as cache_mod + + monkeypatch.setattr(cache_mod, "get_data_fingerprint", lambda data_dir: "data-v1") + monkeypatch.setattr(cache_mod, "get_preprocessing_fingerprint", lambda: "prep-v1") + + timeseries = torch.zeros((2, 4, 3), dtype=torch.float32) + masks = torch.ones((2, 4, 3), dtype=torch.bool) + save_cached_tensors(tmp_path, timeseries, masks, seq_length=4, n_features=3) + loaded = load_cached_tensors(tmp_path, seq_length=4, n_features=3) + assert loaded is not None + + monkeypatch.setattr(cache_mod, "get_data_fingerprint", lambda data_dir: "data-v2") + assert load_cached_tensors(tmp_path, seq_length=4, n_features=3) is None + + def test_failed_tensor_cache_save_removes_temp_file(self, tmp_path, monkeypatch): + """Failed atomic writes should not leave a target cache or stale temp file.""" + import slices.data.tensor_cache as cache_mod + + monkeypatch.setattr(cache_mod, "get_data_fingerprint", lambda data_dir: "data-v1") + monkeypatch.setattr(cache_mod, "get_preprocessing_fingerprint", lambda: "prep-v1") + + def fail_save(cache_data, path): + raise RuntimeError("simulated save failure") + + monkeypatch.setattr(cache_mod.torch, "save", fail_save) + + timeseries = torch.zeros((2, 4, 3), dtype=torch.float32) + masks = torch.ones((2, 4, 3), dtype=torch.bool) + save_cached_tensors(tmp_path, timeseries, masks, seq_length=4, n_features=3) + + cache_path = cache_mod.get_tensor_cache_path(tmp_path, seq_length=4, n_features=3) + assert not cache_path.exists() + assert list(cache_path.parent.glob("*.tmp")) == [] + + +class TestCombinedDatasetValidation: + """Tests for combined-dataset compatibility checks.""" + + def test_feature_order_mismatch_raises(self): + """Same feature set in different order should fail closed.""" + import importlib + + mod = importlib.import_module("scripts.preprocessing.create_combined_dataset") + + meta_a = {"feature_names": ["hr", "map"], "seq_length_hours": 48, "min_stay_hours": 48} + meta_b = {"feature_names": ["map", "hr"], "seq_length_hours": 48, "min_stay_hours": 48} + + with pytest.raises(ValueError, match="Feature order mismatch"): + mod.validate_feature_compatibility(meta_a, meta_b) + + def test_invariant_mismatch_raises(self): + """Different preprocessing invariants should fail before merge.""" + import importlib + + mod = importlib.import_module("scripts.preprocessing.create_combined_dataset") + + meta_a = {"feature_names": ["hr", "map"], "seq_length_hours": 48, "min_stay_hours": 48} + meta_b = {"feature_names": ["hr", "map"], "seq_length_hours": 72, "min_stay_hours": 48} + + with pytest.raises(ValueError, match="preprocessing invariants"): + mod.validate_feature_compatibility(meta_a, meta_b) + + def test_patient_ids_are_namespaced_even_without_stay_collision(self): + """Combined setup should namespace patient IDs independently of stay_id overlap.""" + import importlib + + mod = importlib.import_module("scripts.preprocessing.create_combined_dataset") + + static_a = pl.DataFrame({"stay_id": [1], "patient_id": [10]}) + static_b = pl.DataFrame({"stay_id": [2], "patient_id": [10]}) + + namespaced_a = mod.namespace_patient_ids(static_a, "miiv") + namespaced_b = mod.namespace_patient_ids(static_b, "eicu") + + mod.validate_no_id_collision(namespaced_a, namespaced_b) + assert namespaced_a["patient_id"].item() == "miiv:10" + assert namespaced_b["patient_id"].item() == "eicu:10" + + def test_required_thesis_task_missing_raises(self): + """Combined data should not silently narrow the thesis task surface.""" + import importlib + + mod = importlib.import_module("scripts.preprocessing.create_combined_dataset") + + task_names = list(THESIS_TASKS) + meta_a = { + "task_names": task_names, + "label_manifest": {task: {} for task in task_names}, + } + meta_b = { + "task_names": ["mortality_24h", "mortality_hospital"], + "label_manifest": {"mortality_24h": {}, "mortality_hospital": {}}, + } + labels_a = pl.DataFrame({"stay_id": [1], **{task: [0] for task in task_names}}) + labels_b = pl.DataFrame( + { + "stay_id": [2], + "mortality_24h": [0], + "mortality_hospital": [0], + } + ) + + with pytest.raises(ValueError, match="all thesis tasks"): + mod.validate_required_task_coverage( + meta_a, + meta_b, + labels_a, + labels_b, + ("miiv", "eicu"), + ) + + +class TestCombinedSetupPath: + """Regression tests for the combined-dataset setup flow.""" + + def _write_processed_dataset( + self, + processed_dir, + dataset_name: str, + stay_id: int, + patient_id: int, + ) -> None: + processed_dir.mkdir(parents=True, exist_ok=True) + + static_df = pl.DataFrame( + { + "stay_id": [stay_id], + "patient_id": [patient_id], + "age": [65], + "gender": ["M"], + } + ) + static_df.write_parquet(processed_dir / "static.parquet") + + timeseries_df = pl.DataFrame( + { + "stay_id": [stay_id], + "timeseries": [[[[1.0, 2.0], [3.0, 4.0]]]], + "mask": [[[[True, True], [True, True]]]], + } + ) + timeseries_df.write_parquet(processed_dir / "timeseries.parquet") + + labels_df = pl.DataFrame( + { + "stay_id": [stay_id], + "mortality_24h": [0], + "mortality_hospital": [0], + "aki_kdigo": [0], + "los_remaining": [1.0], + } + ) + labels_df.write_parquet(processed_dir / "labels.parquet") + + label_manifest = { + task: { + "builder_version": "1.0.0", + "config_hash": "abc123", + } + for task in THESIS_TASKS + } + metadata = { + "dataset": dataset_name, + "feature_set": "core", + "feature_names": ["hr", "map"], + "n_features": 2, + "seq_length_hours": 2, + "min_stay_hours": 2, + "label_horizon_hours": 24, + "task_names": list(THESIS_TASKS), + "label_manifest": label_manifest, + } + with open(processed_dir / "metadata.yaml", "w") as f: + yaml.dump(metadata, f) + + def test_create_combined_dataset_prepares_by_default(self, tmp_path, monkeypatch): + """Combined creation should trigger splits/stat prep unless explicitly skipped.""" + mod = importlib.import_module("scripts.preprocessing.create_combined_dataset") + source_a = tmp_path / "miiv" + source_b = tmp_path / "eicu" + output_dir = tmp_path / "combined" + + self._write_processed_dataset(source_a, "miiv", stay_id=1, patient_id=11) + self._write_processed_dataset(source_b, "eicu", stay_id=2, patient_id=22) + + captured = {} + + def fake_prepare_processed_dataset(processed_dir, seed, dataset_name=None): + captured["processed_dir"] = processed_dir + captured["seed"] = seed + captured["dataset_name"] = dataset_name + return {}, {} + + monkeypatch.setattr(mod, "prepare_processed_dataset", fake_prepare_processed_dataset) + monkeypatch.setattr( + sys, + "argv", + [ + "create_combined_dataset.py", + "--source", + str(source_a), + str(source_b), + "--output", + str(output_dir), + ], + ) + + mod.main() + + assert captured["processed_dir"] == output_dir + assert captured["seed"] == 42 + assert captured["dataset_name"] == "combined" + + def test_setup_and_extract_default_path_includes_combined(self, tmp_path): + """The default setup path should build the combined dataset as part of readiness prep.""" + repo_root = Path(__file__).resolve().parents[1] + temp_repo = tmp_path / "repo" + (temp_repo / "scripts" / "internal").mkdir(parents=True, exist_ok=True) + shutil.copy2(repo_root / "scripts" / "setup_and_extract.sh", temp_repo / "scripts") + + for raw_dir in ("data/raw/mimiciv", "data/raw/eicu-crd"): + (temp_repo / raw_dir).mkdir(parents=True, exist_ok=True) + + for ds in ("miiv", "eicu"): + ricu_output = temp_repo / "data" / "ricu_output" / ds + ricu_output.mkdir(parents=True, exist_ok=True) + (ricu_output / "done.txt").write_text("ok\n") + + processed_dir = temp_repo / "data" / "processed" / ds + processed_dir.mkdir(parents=True, exist_ok=True) + for name in ( + "timeseries.parquet", + "static.parquet", + "labels.parquet", + "metadata.yaml", + "splits.yaml", + "normalization_stats.yaml", + ): + (processed_dir / name).write_text("") + + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + uv_log = tmp_path / "uv.log" + uv_script = bin_dir / "uv" + uv_script.write_text( + "#!/usr/bin/env bash\n" 'printf \'%s\\n\' "$*" >> "$UV_LOG"\n' "exit 0\n" + ) + uv_script.chmod(0o755) + + env = os.environ.copy() + env["PATH"] = f"{bin_dir}:{env['PATH']}" + env["UV_LOG"] = str(uv_log) + + result = subprocess.run( + ["bash", "scripts/setup_and_extract.sh", "--skip-deps"], + cwd=temp_repo, + env=env, + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stdout + result.stderr + log_text = uv_log.read_text() + assert "scripts/preprocessing/create_combined_dataset.py" in log_text + assert "--source data/processed/miiv data/processed/eicu" in log_text + + def test_build_processed_data_commands_include_combined_flow(self): + """The lightweight wrapper should orchestrate the existing Python steps.""" + mod = importlib.import_module("scripts.preprocessing.build_processed_data") + + commands = mod.build_commands( + ["miiv", "eicu"], + build_combined=True, + seed=7, + python_executable="python", + ) + + assert commands == [ + ["python", "scripts/preprocessing/extract_ricu.py", "dataset=miiv"], + ["python", "scripts/preprocessing/prepare_dataset.py", "dataset=miiv", "seed=7"], + ["python", "scripts/preprocessing/extract_ricu.py", "dataset=eicu"], + ["python", "scripts/preprocessing/prepare_dataset.py", "dataset=eicu", "seed=7"], + [ + "python", + "scripts/preprocessing/create_combined_dataset.py", + "--source", + "data/processed/miiv", + "data/processed/eicu", + "--names", + "miiv", + "eicu", + "--output", + "data/processed/combined", + "--seed", + "7", + ], + ] + + def test_build_processed_data_combined_adds_missing_source(self): + """Requesting combined should ensure both source datasets are processed.""" + mod = importlib.import_module("scripts.preprocessing.build_processed_data") + + commands = mod.build_commands( + ["miiv"], + build_combined=True, + seed=42, + python_executable="python", + ) + + assert ["python", "scripts/preprocessing/extract_ricu.py", "dataset=eicu"] in commands + + +class TestFairnessRevisionScoping: + """Regression tests for revision-scoped standalone fairness evaluation.""" + + def test_parse_args_accepts_revision_filter(self, monkeypatch): + """The fairness CLI should expose a revision filter for rerun scoping.""" + mod = importlib.import_module("scripts.eval.evaluate_fairness") + monkeypatch.delenv("REVISION", raising=False) + monkeypatch.delenv("WANDB_REVISION", raising=False) + + monkeypatch.setattr( + sys, + "argv", + ["evaluate_fairness.py", "--revision", "thesis-v1"], + ) + args = mod.parse_args() + + assert args.revision == ["thesis-v1"] + + def test_parse_args_uses_revision_env_when_cli_omits_it(self, monkeypatch): + """The standalone script may inherit an explicit revision from the environment.""" + mod = importlib.import_module("scripts.eval.evaluate_fairness") + monkeypatch.setenv("REVISION", "thesis-v1") + monkeypatch.delenv("WANDB_REVISION", raising=False) + monkeypatch.setattr(sys, "argv", ["evaluate_fairness.py"]) + + args = mod.parse_args() + + assert args.revision == ["thesis-v1"] + + def test_parse_args_requires_revision_when_no_scope_is_available(self, monkeypatch): + """Fail closed instead of running an unscoped fairness sweep.""" + mod = importlib.import_module("scripts.eval.evaluate_fairness") + monkeypatch.delenv("REVISION", raising=False) + monkeypatch.delenv("WANDB_REVISION", raising=False) + monkeypatch.setattr(sys, "argv", ["evaluate_fairness.py"]) + + with pytest.raises(SystemExit) as excinfo: + mod.parse_args() + + assert excinfo.value.code == 2 + + def test_default_experiment_classes_include_all_downstream_training_classes(self): + """The standalone thesis fairness corpus should include every downstream class.""" + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + expected = { + "core_ssl_benchmark", + "label_efficiency", + "cross_dataset_transfer", + "hp_robustness", + "capacity_study", + "classical_baselines", + "ts2vec_extension", + "smart_external_reference", + } + assert expected.issubset(set(mod.DEFAULT_EXPERIMENT_CLASSES)) + + def test_fetch_eval_runs_single_revision_filters_by_config_before_tags(self, monkeypatch): + """Revision-scoped fairness should not depend on W&B tags alone.""" + mod = importlib.import_module("scripts.eval.evaluate_fairness") + captured = {} + keep_config = SimpleNamespace( + id="keep_config", + config={ + "experiment_class": "core_ssl_benchmark", + "paradigm": "mae", + "dataset": "miiv", + "phase": "finetune", + "revision": "thesis-v1", + }, + tags=[], + ) + keep_tags = SimpleNamespace( + id="keep_tags", + config={}, + tags=[ + "experiment_class:core_ssl_benchmark", + "paradigm:mae", + "dataset:miiv", + "phase:finetune", + "revision:thesis-v1", + ], + ) + drop = SimpleNamespace(id="drop", config={"revision": "other"}, tags=[]) + drop_stale_tag = SimpleNamespace( + id="drop_stale_tag", + config={"revision": "other"}, + tags=["revision:thesis-v1"], + ) + + class DummyApi: + def __init__(self, timeout): + captured["timeout"] = timeout + + def runs(self, path, filters, order): + captured["path"] = path + captured["filters"] = filters + captured["order"] = order + return [keep_config, keep_tags, drop, drop_stale_tag] + + monkeypatch.setitem(sys.modules, "wandb", SimpleNamespace(Api=DummyApi)) + + runs = mod.fetch_eval_runs( + project="proj", + entity="entity", + experiment_classes=["core_ssl_benchmark"], + paradigms=["mae"], + datasets=["miiv"], + phases=["finetune"], + revisions=["thesis-v1"], + ) + + assert captured["path"] == "entity/proj" + assert captured["filters"] == {"state": "finished"} + assert [run.id for run in runs] == ["keep_config", "keep_tags"] + + def test_fetch_eval_runs_multi_revision_filters_client_side(self, monkeypatch): + """Multiple revisions should be filtered client-side without mixing reruns.""" + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + runs = [ + SimpleNamespace(id="keep_v1", tags=["phase:finetune", "revision:v1"]), + SimpleNamespace(id="keep_v2", tags=["phase:supervised", "revision:v2"]), + SimpleNamespace(id="drop_phase", tags=["phase:baseline", "revision:v1"]), + SimpleNamespace(id="drop_revision", tags=["phase:finetune", "revision:v3"]), + ] + + class DummyApi: + def __init__(self, timeout): + pass + + def runs(self, path, filters, order): + return runs + + monkeypatch.setitem(sys.modules, "wandb", SimpleNamespace(Api=DummyApi)) + + filtered = mod.fetch_eval_runs( + project="proj", + entity=None, + experiment_classes=None, + paradigms=None, + datasets=None, + phases=["finetune", "supervised"], + revisions=["v1", "v2"], + ) + + assert [run.id for run in filtered] == ["keep_v1", "keep_v2"] + + def test_recorded_best_checkpoint_is_used_for_fairness(self, tmp_path): + """Standalone fairness should honor the run's recorded best checkpoint.""" + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + run_dir = tmp_path / "run123" / "checkpoints" + run_dir.mkdir(parents=True) + best_ckpt = run_dir / "best.ckpt" + last_ckpt = run_dir / "last.ckpt" + best_ckpt.write_text("best") + last_ckpt.write_text("last") + + run = SimpleNamespace( + id="run123", + config={"output_dir": "outputs/run123"}, + summary_metrics={ + "_eval_checkpoint_source": "best", + "_best_ckpt_path": "outputs/run123/checkpoints/best.ckpt", + }, + ) + + ckpt_path, source = mod.resolve_evaluation_checkpoint( + run, + outputs_root=str(tmp_path), + task_type="binary", + ) + + assert ckpt_path == best_ckpt + assert source == "recorded_best" + + def test_recorded_final_checkpoint_uses_last_ckpt(self, tmp_path): + """Standalone fairness should use last.ckpt when test metrics used the final model.""" + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + run_dir = tmp_path / "run456" / "checkpoints" + run_dir.mkdir(parents=True) + (run_dir / "last.ckpt").write_text("last") + + run = SimpleNamespace( + id="run456", + config={"output_dir": "outputs/run456"}, + summary_metrics={"_eval_checkpoint_source": "final"}, + ) + + ckpt_path, source = mod.resolve_evaluation_checkpoint( + run, + outputs_root=str(tmp_path), + task_type="binary", + ) + + assert ckpt_path == run_dir / "last.ckpt" + assert source == "recorded_final" + + def test_missing_checkpoint_provenance_raises(self, tmp_path): + """Runs without recorded evaluation provenance must not use heuristic checkpoint lookup.""" + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + run_dir = tmp_path / "run789" / "checkpoints" + run_dir.mkdir(parents=True) + (run_dir / "best.ckpt").write_text("best") + (run_dir / "last.ckpt").write_text("last") + + run = SimpleNamespace( + id="run789", + config={"output_dir": "outputs/run789"}, + summary_metrics={}, + ) + + with pytest.raises(RuntimeError, match="lacks recorded checkpoint provenance"): + mod.resolve_evaluation_checkpoint( + run, + outputs_root=str(tmp_path), + task_type="binary", + ) + + def test_checkpoint_loader_rejects_cls_pooling_override(self): + """CLS pooling changes learned parameters and must not be treated as aggregation-only.""" + from slices.training.checkpoint_loading import _apply_finetune_pooling_override + + cfg = OmegaConf.create({"encoder": {"pooling": "cls"}}) + encoder_config = { + "d_input": 4, + "d_model": 8, + "n_layers": 1, + "n_heads": 2, + "d_ff": 16, + "pooling": "none", + } + + with pytest.raises(RuntimeError, match="CLS pooling adds learned parameters"): + _apply_finetune_pooling_override(encoder_config, cfg) + + def test_tmux_launcher_passes_revision_to_fairness(self, tmp_path): + """The thesis launcher should thread revision tags through the fairness sweep.""" + repo_root = Path(__file__).resolve().parents[1] + temp_repo = tmp_path / "repo" + (temp_repo / "scripts" / "internal").mkdir(parents=True, exist_ok=True) + shutil.copy2( + repo_root / "scripts" / "internal" / "launch_thesis_tmux.sh", + temp_repo / "scripts" / "internal", + ) + + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + tmux_log = tmp_path / "tmux.log" + tmux_script = bin_dir / "tmux" + tmux_script.write_text( + "#!/usr/bin/env bash\n" + 'printf \'%s\\n\' "$*" >> "$TMUX_LOG"\n' + 'if [ "$1" = "has-session" ]; then\n' + " exit 1\n" + "fi\n" + "exit 0\n" + ) + tmux_script.chmod(0o755) + + env = os.environ.copy() + env["PATH"] = f"{bin_dir}:{env['PATH']}" + env["TMUX_LOG"] = str(tmux_log) + env["SESSION_NAME"] = "test-session" + env["WANDB_ENTITY"] = "test-entity" + env["REVISION"] = "thesis-v2" + env["RUN_EXPORT"] = "0" + env["SKIP_LAUNCH_GIT_CHECK"] = "1" + env["VALIDATE_PROCESSED_ARTIFACTS"] = "0" + env["PURGE_RUNTIME_CACHES"] = "0" + + result = subprocess.run( + ["bash", "scripts/internal/launch_thesis_tmux.sh"], + cwd=temp_repo, + env=env, + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stdout + result.stderr + runner_scripts = list((temp_repo / "logs" / "runner").glob("thesis-run-*.sh")) + assert len(runner_scripts) == 1 + runner_text = runner_scripts[0].read_text() + + assert "scripts/eval/evaluate_fairness.py" in runner_text + assert "--revision thesis-v2" in runner_text + assert "--project slices-thesis" in runner_text + assert "--entity test-entity" in runner_text + assert "--force" in runner_text + assert "uv sync --dev --locked" in runner_text + assert "export SLICES_SKIP_LAUNCH_GIT_CHECK=1" in runner_text + + def test_auto_shutdown_counts_fairness_and_export_processes(self): + """Fairness and export jobs should keep final-run VMs alive.""" + script = Path("scripts/internal/auto_shutdown.sh").read_text() + pattern = re.search(r'pgrep -f "([^"]+)"', script).group(1) + + assert re.search( + pattern, + "uv run python scripts/eval/evaluate_fairness.py --revision thesis-v1", + ) + assert re.search( + pattern, + "uv run python scripts/export_results.py --revision thesis-v1", + ) + + def test_tmux_launcher_rejects_offline_wandb_mode(self, tmp_path): + """Final thesis launcher should not silently create offline W&B runs.""" + repo_root = Path(__file__).resolve().parents[1] + temp_repo = tmp_path / "repo" + (temp_repo / "scripts" / "internal").mkdir(parents=True, exist_ok=True) + shutil.copy2( + repo_root / "scripts" / "internal" / "launch_thesis_tmux.sh", + temp_repo / "scripts" / "internal", + ) + + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + tmux_script = bin_dir / "tmux" + tmux_script.write_text( + "#!/usr/bin/env bash\n" + 'if [ "$1" = "has-session" ]; then\n' + " exit 1\n" + "fi\n" + "exit 0\n" + ) + tmux_script.chmod(0o755) + + env = os.environ.copy() + env["PATH"] = f"{bin_dir}:{env['PATH']}" + env["WANDB_ENTITY"] = "test-entity" + env["WANDB_MODE"] = "offline" + env["SKIP_LAUNCH_GIT_CHECK"] = "1" + env["VALIDATE_PROCESSED_ARTIFACTS"] = "0" + + result = subprocess.run( + ["bash", "scripts/internal/launch_thesis_tmux.sh"], + cwd=temp_repo, + env=env, + capture_output=True, + text=True, + ) + + assert result.returncode != 0 + assert "WANDB_MODE must be unset or 'online'" in result.stderr + + def test_tmux_launcher_includes_classical_baselines_in_fairness(self, tmp_path): + """The fairness sweep should include the canonical baseline class by default.""" + repo_root = Path(__file__).resolve().parents[1] + temp_repo = tmp_path / "repo" + (temp_repo / "scripts" / "internal").mkdir(parents=True, exist_ok=True) + shutil.copy2( + repo_root / "scripts" / "internal" / "launch_thesis_tmux.sh", + temp_repo / "scripts" / "internal", + ) + + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + tmux_log = tmp_path / "tmux.log" + tmux_script = bin_dir / "tmux" + tmux_script.write_text( + "#!/usr/bin/env bash\n" + 'printf \'%s\\n\' "$*" >> "$TMUX_LOG"\n' + 'if [ "$1" = "has-session" ]; then\n' + " exit 1\n" + "fi\n" + "exit 0\n" + ) + tmux_script.chmod(0o755) + + env = os.environ.copy() + env["PATH"] = f"{bin_dir}:{env['PATH']}" + env["TMUX_LOG"] = str(tmux_log) + env["SESSION_NAME"] = "test-session" + env["WANDB_ENTITY"] = "test-entity" + env["RUN_EXPORT"] = "0" + env["SKIP_LAUNCH_GIT_CHECK"] = "1" + env["VALIDATE_PROCESSED_ARTIFACTS"] = "0" + env["PURGE_RUNTIME_CACHES"] = "0" + + result = subprocess.run( + ["bash", "scripts/internal/launch_thesis_tmux.sh"], + cwd=temp_repo, + env=env, + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stdout + result.stderr + runner_scripts = list((temp_repo / "logs" / "runner").glob("thesis-run-*.sh")) + assert len(runner_scripts) == 1 + runner_text = runner_scripts[0].read_text() + + assert "--experiment-class" in runner_text + assert "classical_baselines" in runner_text + assert "core_ssl_benchmark" in runner_text + + def test_tmux_launcher_rejects_partial_final_corpus(self, tmp_path): + """Final thesis launch should fail closed when INCLUDE_* drops planned runs.""" + repo_root = Path(__file__).resolve().parents[1] + temp_repo = tmp_path / "repo" + (temp_repo / "scripts" / "internal").mkdir(parents=True, exist_ok=True) + shutil.copy2( + repo_root / "scripts" / "internal" / "launch_thesis_tmux.sh", + temp_repo / "scripts" / "internal", + ) + + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + tmux_script = bin_dir / "tmux" + tmux_script.write_text( + "#!/usr/bin/env bash\n" + 'if [ "$1" = "has-session" ]; then\n' + " exit 1\n" + "fi\n" + "exit 0\n" + ) + tmux_script.chmod(0o755) + + env = os.environ.copy() + env["PATH"] = f"{bin_dir}:{env['PATH']}" + env["WANDB_ENTITY"] = "test-entity" + env["RUN_EXPORT"] = "0" + env["SKIP_LAUNCH_GIT_CHECK"] = "1" + env["VALIDATE_PROCESSED_ARTIFACTS"] = "0" + env["PURGE_RUNTIME_CACHES"] = "0" + env["INCLUDE_CAPACITY_STUDY"] = "0" + + result = subprocess.run( + ["bash", "scripts/internal/launch_thesis_tmux.sh"], + cwd=temp_repo, + env=env, + capture_output=True, + text=True, + ) + + assert result.returncode != 0 + assert "Final corpus preflight failed: selected 2490 runs, expected 2590" in result.stderr + + def test_tmux_status_pane_uses_uv_run_python(self, tmp_path): + """The status pane should use uv-managed Python.""" + repo_root = Path(__file__).resolve().parents[1] + temp_repo = tmp_path / "repo" + (temp_repo / "scripts" / "internal").mkdir(parents=True, exist_ok=True) + shutil.copy2( + repo_root / "scripts" / "internal" / "launch_thesis_tmux.sh", + temp_repo / "scripts" / "internal", + ) + + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + tmux_log = tmp_path / "tmux.log" + tmux_script = bin_dir / "tmux" + tmux_script.write_text( + "#!/usr/bin/env bash\n" + 'printf \'%s\\n\' "$*" >> "$TMUX_LOG"\n' + 'if [ "$1" = "has-session" ]; then\n' + " exit 1\n" + "fi\n" + "exit 0\n" + ) + tmux_script.chmod(0o755) + + env = os.environ.copy() + env["PATH"] = f"{bin_dir}:{env['PATH']}" + env["TMUX_LOG"] = str(tmux_log) + env["SESSION_NAME"] = "test-session" + env["WANDB_ENTITY"] = "test-entity" + env["RUN_EXPORT"] = "0" + env["SKIP_LAUNCH_GIT_CHECK"] = "1" + env["VALIDATE_PROCESSED_ARTIFACTS"] = "0" + env["PURGE_RUNTIME_CACHES"] = "0" + + result = subprocess.run( + ["bash", "scripts/internal/launch_thesis_tmux.sh"], + cwd=temp_repo, + env=env, + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stdout + result.stderr + assert ( + "uv\\ run\\ python\\ scripts/internal/run_experiments.py\\ status" + in tmux_log.read_text() + ) + + +# ============================================================================ +# Issue 4: Sequence-length override +# ============================================================================ + + +class TestSeqLengthOverride: + """Tests for sequence-length mismatch handling in tensor extraction.""" + + def _make_df(self, n_samples: int, stored_len: int, n_features: int) -> pl.DataFrame: + """Create a test DataFrame with nested list columns.""" + rows = [] + for i in range(n_samples): + ts_row = [[float(j + i * 0.1) for _ in range(n_features)] for j in range(stored_len)] + mask_row = [[True for _ in range(n_features)] for _ in range(stored_len)] + rows.append({"timeseries": ts_row, "mask": mask_row}) + return pl.DataFrame(rows) + + def test_override_shorter_than_stored(self): + """Requesting shorter seq_length should truncate.""" + df = self._make_df(n_samples=3, stored_len=48, n_features=2) + ts, masks = extract_tensors_from_dataframe(df, seq_length=24, n_features=2) + assert ts.shape == (3, 24, 2) + assert masks.shape == (3, 24, 2) + + def test_override_longer_than_stored(self): + """Requesting longer seq_length should pad.""" + df = self._make_df(n_samples=3, stored_len=24, n_features=2) + ts, masks = extract_tensors_from_dataframe(df, seq_length=48, n_features=2) + assert ts.shape == (3, 48, 2) + assert masks.shape == (3, 48, 2) + # Padded positions should have NaN in timeseries and False in masks + assert torch.isnan(ts[0, 24, 0]) + assert not masks[0, 24, 0] + + def test_override_equal_to_stored(self): + """Requesting same seq_length should use fast path.""" + df = self._make_df(n_samples=3, stored_len=48, n_features=2) + ts, masks = extract_tensors_from_dataframe(df, seq_length=48, n_features=2) + assert ts.shape == (3, 48, 2) + assert masks.shape == (3, 48, 2) + + +# ============================================================================ +# Issue 2: Mortality timestamp precision +# ============================================================================ + + +class TestMortalityPrecision: + """Tests for precision-aware mortality label building.""" + + def _make_mortality_data(self, stays, deaths): + """Helper to create raw_data dict for mortality builder. + + stays: list of (stay_id, intime, outtime) + deaths: list of (stay_id, death_time, death_date, precision, source, + hospital_expire_flag) + """ + stays_df = pl.DataFrame( + { + "stay_id": [s[0] for s in stays], + "intime": [s[1] for s in stays], + "outtime": [s[2] for s in stays], + } + ) + mort_df = pl.DataFrame( + { + "stay_id": [d[0] for d in deaths], + "death_time": [d[1] for d in deaths], + "death_date": [d[2] for d in deaths], + "death_time_precision": [d[3] for d in deaths], + "death_source": [d[4] for d in deaths], + "hospital_expire_flag": [d[5] for d in deaths], + "dischtime": [None for _ in deaths], + "discharge_location": [None for _ in deaths], + "date_of_death": [d[1] if d[1] else None for d in deaths], + }, + schema={ + "stay_id": pl.Int64, + "death_time": pl.Datetime("us"), + "death_date": pl.Date, + "death_time_precision": pl.Utf8, + "death_source": pl.Utf8, + "hospital_expire_flag": pl.Int32, + "dischtime": pl.Datetime("us"), + "discharge_location": pl.Utf8, + "date_of_death": pl.Datetime("us"), + }, + ) + return {"stays": stays_df, "mortality_info": mort_df} + + def test_timestamp_precision_exact(self): + """Exact timestamp within prediction window → positive.""" + base = datetime(2020, 1, 1, 0, 0, 0) + config = LabelConfig( + task_name="mortality_24h", + task_type="binary", + prediction_window_hours=24, + observation_window_hours=48, + gap_hours=0, + label_sources=["stays", "mortality_info"], + ) + builder = MortalityLabelBuilder(config) + + # Death at hour 60 (within 48-72 prediction window) + raw_data = self._make_mortality_data( + stays=[(1, base, base + timedelta(hours=96))], + deaths=[(1, base + timedelta(hours=60), None, "timestamp", "admissions.deathtime", 1)], + ) + labels = builder.build_labels(raw_data) + assert labels["label"][0] == 1 + + def test_timestamp_death_during_observation_excluded_even_if_outtime_after_obs(self): + """Observation-window timestamp deaths do not require outtime < obs_end.""" + base = datetime(2020, 1, 1, 0, 0, 0) + config = LabelConfig( + task_name="mortality_24h", + task_type="binary", + prediction_window_hours=24, + observation_window_hours=48, + gap_hours=0, + label_sources=["stays", "mortality_info"], + ) + builder = MortalityLabelBuilder(config) + + raw_data = self._make_mortality_data( + stays=[(1, base, base + timedelta(hours=96))], + deaths=[(1, base + timedelta(hours=12), None, "timestamp", "admissions.deathtime", 1)], + ) + + labels = builder.build_labels(raw_data) + + assert labels["label"][0] is None + + def test_hospital_timestamp_death_during_observation_excluded_even_if_outtime_after_obs(self): + """Hospital mortality with obs window must exclude early observed deaths.""" + base = datetime(2020, 1, 1, 0, 0, 0) + config = LabelConfig( + task_name="mortality_hospital", + task_type="binary", + prediction_window_hours=None, + observation_window_hours=24, + gap_hours=0, + label_sources=["stays", "mortality_info"], + ) + builder = MortalityLabelBuilder(config) + + raw_data = self._make_mortality_data( + stays=[(1, base, base + timedelta(hours=96))], + deaths=[(1, base + timedelta(hours=12), None, "timestamp", "admissions.deathtime", 1)], + ) + + labels = builder.build_labels(raw_data) + + assert labels["label"][0] is None + + def test_timestamp_precision_outside_window(self): + """Exact timestamp after prediction window → negative.""" + base = datetime(2020, 1, 1, 0, 0, 0) + config = LabelConfig( + task_name="mortality_24h", + task_type="binary", + prediction_window_hours=24, + observation_window_hours=48, + gap_hours=0, + label_sources=["stays", "mortality_info"], + ) + builder = MortalityLabelBuilder(config) + + # Death at hour 80 (after 72h prediction end) + raw_data = self._make_mortality_data( + stays=[(1, base, base + timedelta(hours=96))], + deaths=[(1, base + timedelta(hours=80), None, "timestamp", "admissions.deathtime", 1)], + ) + labels = builder.build_labels(raw_data) + assert labels["label"][0] == 0 + + def test_date_only_fully_inside_window(self): + """Date-only death where entire day falls inside prediction window → positive.""" + base = datetime(2020, 1, 1, 0, 0, 0) + config = LabelConfig( + task_name="mortality_24h", + task_type="binary", + prediction_window_hours=24, + observation_window_hours=48, + gap_hours=0, + label_sources=["stays", "mortality_info"], + ) + builder = MortalityLabelBuilder(config) + + # Prediction window: hour 48-72 = Jan 3 00:00 to Jan 4 00:00 + # Death date: Jan 3 (00:00-23:59) — entirely within window + from datetime import date + + raw_data = self._make_mortality_data( + stays=[(1, base, base + timedelta(hours=96))], + deaths=[(1, None, date(2020, 1, 3), "date", "patients.dod", 1)], + ) + labels = builder.build_labels(raw_data) + assert labels["label"][0] == 1 + + def test_date_only_overlapping_boundary_is_null(self): + """Date-only death overlapping prediction window boundary → null.""" + base = datetime(2020, 1, 1, 0, 0, 0) + config = LabelConfig( + task_name="mortality_24h", + task_type="binary", + prediction_window_hours=24, + observation_window_hours=48, + gap_hours=0, + label_sources=["stays", "mortality_info"], + ) + builder = MortalityLabelBuilder(config) + + # Prediction window: hour 48-72 = Jan 3 00:00 to Jan 4 00:00 + # Death date: Jan 4 (00:00-23:59) — overlaps pred_end boundary + from datetime import date + + raw_data = self._make_mortality_data( + stays=[(1, base, base + timedelta(hours=120))], + deaths=[(1, None, date(2020, 1, 4), "date", "patients.dod", 1)], + ) + labels = builder.build_labels(raw_data) + assert labels["label"][0] is None + + def test_date_only_before_observation_end(self): + """Date-only death before observation window end → null (excluded).""" + base = datetime(2020, 1, 1, 0, 0, 0) + config = LabelConfig( + task_name="mortality_24h", + task_type="binary", + prediction_window_hours=24, + observation_window_hours=48, + gap_hours=0, + label_sources=["stays", "mortality_info"], + ) + builder = MortalityLabelBuilder(config) + + # Death date: Jan 1 — before obs_end (Jan 3 00:00). + # Patient also left ICU during obs (outtime < obs_end). + from datetime import date + + raw_data = self._make_mortality_data( + stays=[(1, base, base + timedelta(hours=20))], # short stay + deaths=[(1, None, date(2020, 1, 1), "date", "patients.dod", 1)], + ) + labels = builder.build_labels(raw_data) + assert labels["label"][0] is None + + def test_timezone_aware_stays_with_date_only_death(self): + """UTC-aware stay times should not break date-only mortality comparisons.""" + base = datetime(2020, 1, 1, 0, 0, 0, tzinfo=timezone.utc) + config = LabelConfig( + task_name="mortality_24h", + task_type="binary", + prediction_window_hours=24, + observation_window_hours=48, + gap_hours=0, + label_sources=["stays", "mortality_info"], + ) + builder = MortalityLabelBuilder(config) + + from datetime import date + + raw_data = self._make_mortality_data( + stays=[(1, base, base + timedelta(hours=96))], + deaths=[(1, None, date(2020, 1, 3), "date", "patients.dod", 1)], + ) + + assert raw_data["stays"]["intime"].dtype == pl.Datetime("us", "UTC") + + labels = builder.build_labels(raw_data) + assert labels["label"][0] == 1 + + def test_survived_patient(self): + """Patient who survived → label 0.""" + base = datetime(2020, 1, 1, 0, 0, 0) + config = LabelConfig( + task_name="mortality_24h", + task_type="binary", + prediction_window_hours=24, + observation_window_hours=48, + gap_hours=0, + label_sources=["stays", "mortality_info"], + ) + builder = MortalityLabelBuilder(config) + + raw_data = self._make_mortality_data( + stays=[(1, base, base + timedelta(hours=96))], + deaths=[(1, None, None, None, None, 0)], + ) + labels = builder.build_labels(raw_data) + assert labels["label"][0] == 0 + + def test_legacy_schema_migration(self): + """Legacy datetime date_of_death should be migrated conservatively.""" + merged = pl.DataFrame( + { + "stay_id": [1], + "intime": [datetime(2020, 1, 1)], + "outtime": [datetime(2020, 1, 5)], + "date_of_death": [datetime(2020, 1, 3, 12, 0, 0)], + "hospital_expire_flag": [1], + } + ) + result = MortalityLabelBuilder._ensure_precision_columns(merged) + assert "death_time_precision" in result.columns + assert "death_time" in result.columns + assert result["death_time_precision"][0] == "date" + assert result["death_time"][0] is None + assert result["death_date"][0].isoformat() == "2020-01-03" + + +# ============================================================================ +# Issue 6: Exporter LR/mask_ratio decode and historical recovery +# ============================================================================ + + +class TestExporterClassMetadata: + """Tests for class-based exporter metadata recovery.""" + + def test_lr_decode_matches_run_experiments_encoding(self): + """Decode table must match str(v).replace('.','') encoding from run_experiments.py.""" + # These are the exact values from LR_ABLATION = [2e-4, 5e-4, 2e-3] + assert str(2e-4).replace(".", "") == "00002" + assert str(5e-4).replace(".", "") == "00005" + assert str(2e-3).replace(".", "") == "0002" + + from scripts.export_results import _LR_DECODE + + assert _LR_DECODE["00002"] == 2e-4 + assert _LR_DECODE["00005"] == 5e-4 + assert _LR_DECODE["0002"] == 2e-3 + + def test_mr_decode_matches_run_experiments_encoding(self): + """Decode table must match str(v).replace('.','') encoding.""" + assert str(0.3).replace(".", "") == "03" + assert str(0.75).replace(".", "") == "075" + + from scripts.export_results import _MR_DECODE + + assert _MR_DECODE["03"] == 0.3 + assert _MR_DECODE["075"] == 0.75 + + def test_recover_lr_from_output_dir(self): + """LR sensitivity finetunes recover upstream LR from output_dir pattern.""" + from scripts.export_results import _recover_pretrain_metadata + + config = { + "output_dir": "outputs/hp_robustness/finetune_mae_mortality_24h_miiv_seed789_lr00002" + } + up_lr, up_mr, subtype = _recover_pretrain_metadata("some_run", config) + assert up_lr == 2e-4 + assert subtype == "lr_sensitivity" + + def test_recover_mask_ratio_from_output_dir(self): + """Mask-ratio sensitivity finetunes recover upstream mask ratio from output_dir.""" + from scripts.export_results import _recover_pretrain_metadata + + for suffix in ["mask_ratio03", "maskratio03"]: + config = { + "output_dir": ( + "outputs/hp_robustness/" f"finetune_jepa_mortality_24h_miiv_seed42_{suffix}" + ) + } + up_lr, up_mr, subtype = _recover_pretrain_metadata("some_run", config) + assert up_mr == 0.3 + assert subtype == "mask_ratio_sensitivity" + + def test_new_runs_use_config_directly(self): + """New runs with explicit upstream fields skip name parsing.""" + from scripts.export_results import _recover_pretrain_metadata + + config = {"upstream_pretrain_lr": 5e-4, "experiment_subtype": "lr_sensitivity"} + up_lr, up_mr, subtype = _recover_pretrain_metadata("irrelevant_name", config) + assert up_lr == 5e-4 + assert subtype == "lr_sensitivity" + + def test_view_mask_subtype_recovers_nested_mask_ratio(self): + """Contrastive view/mask sensitivity rows should preserve the upstream ratio.""" + from scripts.export_results import _recover_pretrain_metadata + + config = { + "experiment_subtype": "view_mask_sensitivity", + "ssl": {"mask_ratio": 0.75}, + } + up_lr, up_mr, subtype = _recover_pretrain_metadata("irrelevant_name", config) + + assert up_lr is None + assert up_mr == 0.75 + assert subtype == "view_mask_sensitivity" + + def test_core_run_returns_none(self): + """Core runs (no HP ablation) return None for all fields.""" + from scripts.export_results import _recover_pretrain_metadata + + config = {"output_dir": "outputs/core_ssl_benchmark/finetune_mae_mortality_24h_miiv_seed42"} + up_lr, up_mr, subtype = _recover_pretrain_metadata("some_run", config) + assert up_lr is None + assert up_mr is None + assert subtype is None + + def test_infer_model_size_for_capacity_variants(self): + """Exporter should preserve medium/large labels.""" + from scripts.export_results import _infer_model_size + + assert _infer_model_size({"encoder": {"d_model": 64, "n_layers": 2}}) == "default" + assert _infer_model_size({"encoder": {"d_model": 128, "n_layers": 4}}) == "medium" + assert _infer_model_size({"encoder": {"d_model": 256, "n_layers": 4}}) == "large" + assert ( + _infer_model_size( + {"paradigm": "smart", "encoder": {"name": "smart", "d_model": 32, "n_layers": 2}} + ) + == "default" + ) + + def test_extract_run_uses_experiment_class_tags(self): + """Runs should export with explicit final experiment classes.""" + from scripts.export_results import extract_run + + transfer = DummyWandbRun( + run_id="transfer", + config={ + "experiment_class": "cross_dataset_transfer", + "dataset": "eicu", + "paradigm": "mae", + "source_dataset": "miiv", + "seed": 789, + "encoder": {"d_model": 64, "n_layers": 2}, + "training": {"freeze_encoder": False}, + "task": {"task_name": "mortality_24h"}, + }, + tags=["phase:finetune"], + name="finetune_mae_mortality_24h_eicu_from_miiv_seed789", + ) + assert extract_run(transfer, [])["experiment_class"] == "cross_dataset_transfer" + + ts2vec = DummyWandbRun( + run_id="ts2vec", + config={ + "experiment_class": "ts2vec_extension", + "dataset": "miiv", + "paradigm": "ts2vec", + "seed": 42, + "encoder": {"d_model": 64, "n_layers": 2}, + "training": {"freeze_encoder": False}, + "task": {"task_name": "mortality_24h"}, + }, + tags=["phase:finetune"], + name="finetune_ts2vec_mortality_24h_miiv_seed42", + ) + assert extract_run(ts2vec, [])["experiment_class"] == "ts2vec_extension" + + def test_extract_run_recovers_transfer_from_output_dir(self): + """Transfer runs should recover source_dataset from output_dir.""" + from scripts.export_results import extract_run + + transfer = DummyWandbRun( + run_id="transfer_output_dir", + config={ + "experiment_class": "cross_dataset_transfer", + "dataset": "eicu", + "paradigm": "mae", + "seed": 789, + "output_dir": ( + "outputs/cross_dataset_transfer/" + "finetune_mae_mortality_24h_eicu_seed789_from_miiv" + ), + "encoder": {"d_model": 64, "n_layers": 2}, + "training": {"freeze_encoder": False}, + "task": {"task_name": "mortality_24h"}, + }, + tags=["phase:finetune"], + name="s10_finetune_eicu_mae_mortality_24h_seed789", + ) + + row = extract_run(transfer, []) + assert row["source_dataset"] == "miiv" + assert row["experiment_class"] == "cross_dataset_transfer" + + def test_extract_run_assigns_full_finetune_protocol_to_xgboost_baselines(self): + """XGBoost baselines must align with the full-finetune comparison family.""" + from scripts.export_results import extract_run + + run = DummyWandbRun( + run_id="xgb", + config={ + "experiment_class": "classical_baselines", + "dataset": "miiv", + "paradigm": "xgboost", + "seed": 42, + "task": {"task_name": "mortality_24h"}, + }, + tags=["phase:baseline"], + name="s11_xgboost_miiv_mortality_24h_seed42", + ) + + assert extract_run(run, [])["protocol"] == "full_finetune" + + def test_xgboost_wandb_tags_include_revision_scope(self): + """XGBoost runs must be visible to revision-scoped export and fairness jobs.""" + from scripts.training.xgboost_baseline import _build_wandb_tags + + cfg = OmegaConf.create( + { + "logging": {"wandb_tags": ["phase:baseline"]}, + "experiment_class": "classical_baselines", + "revision": "thesis-v1", + "rerun_reason": "rerun canonical thesis baseline sweep with fixed tags", + "label_fraction": 0.1, + } + ) + + tags = _build_wandb_tags(cfg) + + assert "phase:baseline" in tags + assert "experiment_class:classical_baselines" in tags + assert "revision:thesis-v1" in tags + assert "label_fraction:0.1" in tags + assert "ablation:label-efficiency" in tags + assert any(tag.startswith("rerun-reason:") and len(tag) <= 64 for tag in tags) + + def test_xgboost_initializes_wandb_before_data_validation(self): + """Failed XGBoost jobs should appear in final W&B run accounting.""" + script = Path("scripts/training/xgboost_baseline.py").read_text() + + assert script.index("wandb_module, wandb_run = _init_wandb_run") < script.index( + "validate_data_prerequisites(" + ) + + def test_build_per_seed_df_keeps_distinct_hp_ablation_configs(self): + """Distinct upstream HP-ablation configs must not deduplicate together.""" + from scripts.export_results import build_per_seed_df + + runs = [ + DummyWandbRun( + run_id="hp_lr_2e4", + config={ + "experiment_class": "hp_robustness", + "dataset": "miiv", + "paradigm": "mae", + "seed": 42, + "output_dir": ( + "outputs/hp_robustness/" "finetune_mae_mortality_24h_miiv_seed42_lr00002" + ), + "encoder": {"d_model": 64, "n_layers": 2}, + "training": {"freeze_encoder": False}, + "task": {"task_name": "mortality_24h"}, + }, + tags=["phase:finetune"], + name="finetune_mae_mortality_24h_miiv_seed42_lr00002", + created_at="2026-04-07T00:00:01", + ), + DummyWandbRun( + run_id="hp_lr_5e4", + config={ + "experiment_class": "hp_robustness", + "dataset": "miiv", + "paradigm": "mae", + "seed": 42, + "output_dir": ( + "outputs/hp_robustness/" "finetune_mae_mortality_24h_miiv_seed42_lr00005" + ), + "encoder": {"d_model": 64, "n_layers": 2}, + "training": {"freeze_encoder": False}, + "task": {"task_name": "mortality_24h"}, + }, + tags=["phase:finetune"], + name="finetune_mae_mortality_24h_miiv_seed42_lr00005", + created_at="2026-04-07T00:00:02", + ), + ] + + df = build_per_seed_df(runs) + + assert len(df) == 2 + assert set(df["experiment_class"]) == {"hp_robustness"} + assert sorted(df["upstream_pretrain_lr"].tolist()) == [2e-4, 5e-4] + + def test_build_aggregated_df_merges_hp_robustness_across_seeds(self): + """HP robustness rows should merge across seeds by scientific fingerprint.""" + from scripts.export_results import build_aggregated_df + + per_seed_df = pd.DataFrame( + [ + { + "wandb_run_id": f"run_{i}", + "experiment_class": "hp_robustness", + "experiment_type": "hp_robustness", + "experiment_subtype": "lr_sensitivity", + "paradigm": "mae", + "dataset": "miiv", + "task": "mortality_24h", + "seed": seed, + "protocol": "full_finetune", + "label_fraction": 1.0, + "model_size": "default", + "source_dataset": None, + "phase": "finetune", + "upstream_pretrain_lr": 2e-4, + "upstream_pretrain_mask_ratio": None, + } + for i, seed in enumerate([42, 123, 456, 789, 1011]) + ] + ) + + agg = build_aggregated_df(per_seed_df) + + assert len(agg) == 1 + assert agg.iloc[0]["experiment_class"] == "hp_robustness" + assert agg.iloc[0]["experiment_subtype"] == "lr_sensitivity" + assert agg.iloc[0]["n_seeds"] == 5 + assert set(yaml.safe_load(agg.iloc[0]["experiment_class_list"])) == {"hp_robustness"} + + def test_build_per_seed_df_raises_on_ambiguous_core_collision(self): + """Distinct configs that a fingerprint would collapse should fail closed.""" + from scripts.export_results import build_per_seed_df + + runs = [ + DummyWandbRun( + run_id="core_lr_1e4", + config={ + "experiment_class": "core_ssl_benchmark", + "dataset": "miiv", + "paradigm": "mae", + "seed": 42, + "encoder": {"d_model": 64, "n_layers": 2}, + "training": {"freeze_encoder": False}, + "optimizer": {"lr": 1e-4}, + "task": {"task_name": "mortality_24h"}, + }, + tags=["phase:finetune"], + name="finetune_mae_mortality_24h_miiv_seed42_lr1e4", + created_at="2026-04-07T00:00:01", + ), + DummyWandbRun( + run_id="core_lr_3e4", + config={ + "experiment_class": "core_ssl_benchmark", + "dataset": "miiv", + "paradigm": "mae", + "seed": 42, + "encoder": {"d_model": 64, "n_layers": 2}, + "training": {"freeze_encoder": False}, + "optimizer": {"lr": 3e-4}, + "task": {"task_name": "mortality_24h"}, + }, + tags=["phase:finetune"], + name="finetune_mae_mortality_24h_miiv_seed42_lr3e4", + created_at="2026-04-07T00:00:02", + ), + ] + + with pytest.raises(RuntimeError, match="Ambiguous export fingerprint"): + build_per_seed_df(runs) + + +class TestExperimentRunnerWandbOverrides: + """Tests for clean project/entity overrides in the experiment runner.""" + + def test_apply_wandb_target_injects_project_and_entity(self): + from scripts.internal.run_experiments import Run, apply_wandb_target + + run = Run( + id="core_ssl_benchmark_supervised_mortality_24h_miiv_seed42", + experiment_class="core_ssl_benchmark", + run_type="supervised", + paradigm="supervised", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/supervised_mortality_24h_miiv_seed42", + ) + + result = apply_wandb_target([run], project="slices-thesis", entity="hannes-ill") + assert result[0].extra_overrides["project_name"] == "slices-thesis" + assert result[0].extra_overrides["logging.wandb_project"] == "slices-thesis" + assert result[0].extra_overrides["logging.wandb_entity"] == "hannes-ill" + + def test_launch_commit_mismatch_resets_stale_completed_state(self): + from scripts.internal.run_experiments import Run, reset_state_for_launch_commit_mismatch + + run = Run( + id="core_ssl_benchmark_rev-thesis-v1_supervised_mortality_24h_miiv_seed42", + experiment_class="core_ssl_benchmark", + run_type="supervised", + paradigm="supervised", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/supervised_mortality_24h_miiv_seed42", + task="mortality_24h", + extra_overrides={"+launch_commit": "newcommit"}, + ) + state = { + "version": 1, + "runs": { + run.id: { + "status": "completed", + "launch_commit": "oldcommit", + "command": "uv run python scripts/training/supervised.py", + } + }, + } + + reset_state_for_launch_commit_mismatch([run], state) + + assert state["runs"][run.id]["status"] == "pending" + assert "launch_commit changed" in state["runs"][run.id]["reset_reason"] + + def test_launch_commit_mismatch_quarantines_stale_resume_artifacts(self, tmp_path): + from scripts.internal.run_experiments import Run, reset_state_for_launch_commit_mismatch + + output_dir = tmp_path / "supervised_mortality_24h_miiv_seed42" + checkpoint_dir = output_dir / "checkpoints" + checkpoint_dir.mkdir(parents=True) + (checkpoint_dir / "last.ckpt").write_text("old checkpoint") + + run = Run( + id="core_ssl_benchmark_rev-thesis-v1_supervised_mortality_24h_miiv_seed42", + experiment_class="core_ssl_benchmark", + run_type="supervised", + paradigm="supervised", + dataset="miiv", + seed=42, + output_dir=str(output_dir), + task="mortality_24h", + extra_overrides={"+launch_commit": "newcommit"}, + ) + state = { + "version": 1, + "runs": { + run.id: { + "status": "completed", + "launch_commit": "oldcommit", + "command": "uv run python scripts/training/supervised.py", + } + }, + } + + reset_state_for_launch_commit_mismatch([run], state) + + state_entry = state["runs"][run.id] + quarantined_output_dir = Path(state_entry["quarantined_output_dir"]) + assert state_entry["status"] == "pending" + assert not output_dir.exists() + assert (quarantined_output_dir / "checkpoints" / "last.ckpt").read_text() == ( + "old checkpoint" + ) + assert not any(part.startswith("ckpt_path=") for part in run.build_command({})) + + def test_wandb_target_mismatch_resets_completed_state_and_output(self, tmp_path): + from scripts.internal.run_experiments import ( + Run, + apply_wandb_target, + reset_state_for_launch_identity_mismatch, + ) + + output_dir = tmp_path / "supervised_mortality_24h_miiv_seed42" + checkpoint_dir = output_dir / "checkpoints" + checkpoint_dir.mkdir(parents=True) + (checkpoint_dir / "last.ckpt").write_text("old checkpoint") + + run = Run( + id="core_ssl_benchmark_rev-thesis-v1_supervised_mortality_24h_miiv_seed42", + experiment_class="core_ssl_benchmark", + run_type="supervised", + paradigm="supervised", + dataset="miiv", + seed=42, + output_dir=str(output_dir), + task="mortality_24h", + extra_overrides={"revision": "thesis-v1"}, + ) + apply_wandb_target([run], project="new-project", entity="team-a") + state = { + "version": 1, + "runs": { + run.id: { + "status": "completed", + "command": ( + "uv run python scripts/training/supervised.py " + "revision=thesis-v1 project_name=old-project " + "logging.wandb_project=old-project logging.wandb_entity=team-a" + ), + } + }, + } + + reset_state_for_launch_identity_mismatch([run], state) + + state_entry = state["runs"][run.id] + assert state_entry["status"] == "pending" + assert "launch identity changed" in state_entry["reset_reason"] + assert not output_dir.exists() + assert Path(state_entry["quarantined_output_dir"]).exists() + + def test_scoped_run_does_not_resume_unmarked_checkpoint(self, tmp_path): + from scripts.internal.run_experiments import Run + + output_dir = tmp_path / "pretrain_mae_miiv_seed42" + checkpoint_dir = output_dir / "checkpoints" + checkpoint_dir.mkdir(parents=True) + (checkpoint_dir / "last.ckpt").write_text("old checkpoint") + + run = Run( + id="core_ssl_benchmark_rev-thesis-v1_pretrain_mae_miiv_seed42", + experiment_class="core_ssl_benchmark", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir=str(output_dir), + extra_overrides={"revision": "thesis-v1"}, + ) + + assert not any(part.startswith("ckpt_path=") for part in run.build_command({})) + + def test_scoped_run_resumes_matching_marked_checkpoint(self, tmp_path): + from scripts.internal.run_experiments import Run, _write_output_launch_identity + + output_dir = tmp_path / "pretrain_mae_miiv_seed42" + checkpoint_dir = output_dir / "checkpoints" + checkpoint_dir.mkdir(parents=True) + last_ckpt = checkpoint_dir / "last.ckpt" + last_ckpt.write_text("checkpoint") + + run = Run( + id="core_ssl_benchmark_rev-thesis-v1_pretrain_mae_miiv_seed42", + experiment_class="core_ssl_benchmark", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir=str(output_dir), + extra_overrides={"revision": "thesis-v1"}, + ) + _write_output_launch_identity(run) + + assert f"ckpt_path={last_ckpt}" in run.build_command({}) + + def test_revised_status_handles_base_and_revision_groups(self, monkeypatch, capsys): + import scripts.internal.run_experiments as runner + + base = runner.Run( + id="core_ssl_benchmark_supervised_mortality_24h_miiv_seed42", + experiment_class="core_ssl_benchmark", + run_type="supervised", + paradigm="supervised", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/supervised_mortality_24h_miiv_seed42", + task="mortality_24h", + ) + revised_id = "core_ssl_benchmark_rev-thesis-v1_supervised_mortality_24h_miiv_seed42" + + monkeypatch.setattr(runner, "generate_all_runs", lambda: [base]) + monkeypatch.setattr( + runner, + "load_state", + lambda: { + "version": 1, + "runs": { + revised_id: {"status": "completed"}, + }, + }, + ) + + runner.print_status() + + output = capsys.readouterr().out + assert "core_ssl_benchmark" in output + assert "core_ssl_benchmark/thesis-v1" in output + + def test_revised_status_filter_uses_revisioned_generated_ids(self, monkeypatch, capsys): + import scripts.internal.run_experiments as runner + + base = runner.Run( + id="core_ssl_benchmark_supervised_mortality_24h_miiv_seed42", + run_key="supervised_mortality_24h_miiv_seed42", + experiment_class="core_ssl_benchmark", + run_type="supervised", + paradigm="supervised", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/supervised_mortality_24h_miiv_seed42", + task="mortality_24h", + ) + revised_id = "core_ssl_benchmark_rev-thesis-v1_supervised_mortality_24h_miiv_seed42" + + monkeypatch.setattr(runner, "generate_all_runs", lambda: [base]) + monkeypatch.setattr( + runner, + "load_state", + lambda: {"version": 1, "runs": {revised_id: {"status": "completed"}}}, + ) + + runner.print_status(revision_filter="thesis-v1") + + output = capsys.readouterr().out + assert "core_ssl_benchmark/thesis-v1" in output + assert " 1 | 1 |" in output + assert "core_ssl_benchmark |" not in output + + def test_main_run_requires_revision_and_project(self, monkeypatch): + import scripts.internal.run_experiments as runner + + monkeypatch.delenv("WANDB_PROJECT", raising=False) + monkeypatch.setattr( + sys, + "argv", + ["run_experiments.py", "run", "--experiment-class", "core_ssl_benchmark"], + ) + + with pytest.raises(SystemExit) as excinfo: + runner.main() + + assert excinfo.value.code == 2 + + def test_main_retry_requires_revision_scope_and_project(self, monkeypatch): + import scripts.internal.run_experiments as runner + + monkeypatch.delenv("WANDB_PROJECT", raising=False) + monkeypatch.setattr(sys, "argv", ["run_experiments.py", "retry", "--failed"]) + + with pytest.raises(SystemExit) as excinfo: + runner.main() + + assert excinfo.value.code == 2 + + def test_main_run_dry_run_warns_without_launch_commit(self, monkeypatch, capsys): + import scripts.internal.run_experiments as runner + + monkeypatch.delenv("SLICES_LAUNCH_COMMIT", raising=False) + monkeypatch.setattr( + sys, + "argv", + [ + "run_experiments.py", + "run", + "--experiment-class", + "core_ssl_benchmark", + "--revision", + "dryrun-audit", + "--project", + "slices-thesis", + "--entity", + "hannes-ill", + "--dry-run", + ], + ) + monkeypatch.setattr(runner, "cmd_run", lambda args: 0) + + with pytest.raises(SystemExit) as excinfo: + runner.main() + + assert excinfo.value.code == 0 + assert "dry run has no launch provenance" in capsys.readouterr().err + + def test_main_run_dry_run_rejects_invalid_launch_commit(self, monkeypatch): + import scripts.internal.run_experiments as runner + + def fail_git_stdout(args): + raise RuntimeError("Needed a single revision") + + monkeypatch.delenv("SLICES_LAUNCH_COMMIT", raising=False) + monkeypatch.setattr(runner, "_git_stdout", fail_git_stdout) + monkeypatch.setattr( + sys, + "argv", + [ + "run_experiments.py", + "run", + "--experiment-class", + "core_ssl_benchmark", + "--revision", + "dryrun-audit", + "--project", + "slices-thesis", + "--entity", + "hannes-ill", + "--launch-commit", + "not-a-commit", + "--dry-run", + ], + ) + monkeypatch.setattr( + runner, "cmd_run", lambda args: pytest.fail("cmd_run should not execute") + ) + + with pytest.raises(SystemExit) as excinfo: + runner.main() + + assert excinfo.value.code == 2 + + def test_main_run_dry_run_resolves_launch_commit_before_execution(self, monkeypatch): + import scripts.internal.run_experiments as runner + + resolved_commit = "b" * 40 + captured = {} + + def fake_git_stdout(args): + if args == ["rev-parse", "--verify", "HEAD^{commit}"]: + return resolved_commit + raise AssertionError(f"unexpected git command: {args}") + + def fake_cmd_run(args): + captured["args"] = args + return 0 + + monkeypatch.delenv("SLICES_LAUNCH_COMMIT", raising=False) + monkeypatch.setattr(runner, "_git_stdout", fake_git_stdout) + monkeypatch.setattr(runner, "cmd_run", fake_cmd_run) + monkeypatch.setattr( + sys, + "argv", + [ + "run_experiments.py", + "run", + "--experiment-class", + "core_ssl_benchmark", + "--revision", + "dryrun-audit", + "--project", + "slices-thesis", + "--entity", + "hannes-ill", + "--launch-commit", + "HEAD", + "--dry-run", + ], + ) + + with pytest.raises(SystemExit) as excinfo: + runner.main() + + assert excinfo.value.code == 0 + assert captured["args"].launch_commit == resolved_commit + + def test_main_run_requires_launch_commit_for_final_launch(self, monkeypatch): + import scripts.internal.run_experiments as runner + + monkeypatch.delenv("SLICES_LAUNCH_COMMIT", raising=False) + monkeypatch.setattr( + sys, + "argv", + [ + "run_experiments.py", + "run", + "--experiment-class", + "core_ssl_benchmark", + "--revision", + "thesis-v1", + "--project", + "slices-thesis", + "--entity", + "hannes-ill", + ], + ) + + with pytest.raises(SystemExit) as excinfo: + runner.main() + + assert excinfo.value.code == 2 + + def test_main_run_resolves_symbolic_launch_commit_before_execution(self, monkeypatch): + import scripts.internal.run_experiments as runner + + resolved_commit = "a" * 40 + captured = {} + + def fake_git_stdout(args): + if args == ["rev-parse", "--verify", "HEAD^{commit}"]: + return resolved_commit + if args == ["rev-parse", "--verify", "HEAD"]: + return resolved_commit + raise AssertionError(f"unexpected git command: {args}") + + def fake_cmd_run(args): + captured["launch_commit"] = args.launch_commit + return 0 + + monkeypatch.delenv("WANDB_MODE", raising=False) + monkeypatch.delenv("SLICES_SKIP_LAUNCH_GIT_CHECK", raising=False) + monkeypatch.setattr(runner, "_git_stdout", fake_git_stdout) + monkeypatch.setattr(runner, "_git_quiet", lambda args: True) + monkeypatch.setattr(runner, "cmd_run", fake_cmd_run) + monkeypatch.setattr( + sys, + "argv", + [ + "run_experiments.py", + "run", + "--experiment-class", + "core_ssl_benchmark", + "--revision", + "thesis-v1", + "--project", + "slices-thesis", + "--entity", + "hannes-ill", + "--launch-commit", + "HEAD", + ], + ) + + with pytest.raises(SystemExit) as excinfo: + runner.main() + + assert excinfo.value.code == 0 + assert captured["launch_commit"] == resolved_commit + + def test_main_run_requires_wandb_entity_for_final_launch(self, monkeypatch): + import scripts.internal.run_experiments as runner + + monkeypatch.delenv("WANDB_ENTITY", raising=False) + monkeypatch.setattr( + sys, + "argv", + [ + "run_experiments.py", + "run", + "--experiment-class", + "core_ssl_benchmark", + "--revision", + "thesis-v1", + "--project", + "slices-thesis", + "--launch-commit", + "dummy", + ], + ) + + with pytest.raises(SystemExit) as excinfo: + runner.main() + + assert excinfo.value.code == 2 + + def test_main_run_rejects_offline_wandb_mode_for_final_launch(self, monkeypatch): + import scripts.internal.run_experiments as runner + + monkeypatch.setenv("WANDB_MODE", "offline") + monkeypatch.setattr( + sys, + "argv", + [ + "run_experiments.py", + "run", + "--experiment-class", + "core_ssl_benchmark", + "--revision", + "thesis-v1", + "--project", + "slices-thesis", + "--entity", + "hannes-ill", + "--launch-commit", + "dummy", + ], + ) + + with pytest.raises(SystemExit) as excinfo: + runner.main() + + assert excinfo.value.code == 2 + + def test_main_run_skip_git_check_accepts_unchecked_launch_commit( + self, + monkeypatch, + capsys, + ): + import scripts.internal.run_experiments as runner + + monkeypatch.setenv("SLICES_SKIP_LAUNCH_GIT_CHECK", "1") + monkeypatch.delenv("WANDB_MODE", raising=False) + monkeypatch.setattr( + sys, + "argv", + [ + "run_experiments.py", + "run", + "--experiment-class", + "core_ssl_benchmark", + "--revision", + "thesis-v1", + "--project", + "slices-thesis", + "--entity", + "hannes-ill", + "--launch-commit", + "unchecked", + ], + ) + monkeypatch.setattr(runner, "cmd_run", lambda args: 0) + + with pytest.raises(SystemExit) as excinfo: + runner.main() + + assert excinfo.value.code == 0 + assert "skipping final launch git provenance validation" in capsys.readouterr().err + + def test_cmd_run_respects_requested_class_order(self, monkeypatch): + import scripts.internal.run_experiments as runner + + capacity = runner.Run( + id="capacity_study_supervised_mortality_24h_miiv_seed42", + experiment_class="capacity_study", + run_type="supervised", + paradigm="supervised", + dataset="miiv", + seed=42, + output_dir="outputs/capacity_study/supervised_mortality_24h_miiv_seed42", + task="mortality_24h", + ) + classical = runner.Run( + id="classical_baselines_supervised_mortality_24h_miiv_seed789", + experiment_class="classical_baselines", + run_type="supervised", + paradigm="supervised", + dataset="miiv", + seed=789, + output_dir="outputs/classical_baselines/supervised_mortality_24h_miiv_seed789", + task="mortality_24h", + ) + scheduled = {} + + monkeypatch.setattr(runner, "generate_all_runs", lambda: [capacity, classical]) + monkeypatch.setattr(runner, "load_state", lambda: {"version": 1, "runs": {}}) + monkeypatch.setattr(runner, "recover_stale_running", lambda state: None) + monkeypatch.setattr( + runner, + "run_scheduler", + lambda runs, state, parallel, dry_run: scheduled.setdefault("runs", runs), + ) + + args = SimpleNamespace( + experiment_class=["classical_baselines", "capacity_study"], + revision=None, + reason=None, + project=None, + entity=None, + parallel=1, + dry_run=True, + ) + + runner.cmd_run(args) + + assert [run.experiment_class for run in scheduled["runs"]] == [ + "classical_baselines", + "capacity_study", + ] + + def test_pretrain_and_finetune_commands_resume_from_last_checkpoint(self, tmp_path): + from scripts.internal.run_experiments import Run + + pretrain_dir = tmp_path / "pretrain" + finetune_dir = tmp_path / "finetune" + (pretrain_dir / "checkpoints").mkdir(parents=True) + (finetune_dir / "checkpoints").mkdir(parents=True) + (pretrain_dir / "checkpoints" / "last.ckpt").write_text("checkpoint") + (finetune_dir / "checkpoints" / "last.ckpt").write_text("checkpoint") + + pretrain = Run( + id="pretrain", + experiment_class="core_ssl_benchmark", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir=str(pretrain_dir), + ) + finetune = Run( + id="finetune", + experiment_class="core_ssl_benchmark", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir=str(finetune_dir), + depends_on=[pretrain.id], + task="mortality_24h", + freeze_encoder=False, + ) + + pretrain_cmd = pretrain.build_command({}) + assert f"ckpt_path={pretrain_dir / 'checkpoints' / 'last.ckpt'}" in pretrain_cmd + assert f"ckpt_path={finetune_dir / 'checkpoints' / 'last.ckpt'}" in finetune.build_command( + {pretrain.id: pretrain} + ) + + def test_recover_stale_running_resets_pid_reuse(self, monkeypatch): + import scripts.internal.run_experiments as runner + + state = { + "version": 1, + "runs": {"run-a": {"status": "running", "pid": 123, "command": "uv run expected"}}, + } + + monkeypatch.setattr(runner, "is_pid_alive", lambda pid: True) + monkeypatch.setattr(runner, "_pid_matches_command", lambda pid, command: False) + + recovered = runner.recover_stale_running(state) + + assert recovered == 1 + assert state["runs"]["run-a"]["status"] == "pending" + assert "pid" not in state["runs"]["run-a"] + assert "command" not in state["runs"]["run-a"] + + def test_status_recovers_and_persists_stale_running_entries(self, monkeypatch, capsys): + import scripts.internal.run_experiments as runner + + state = { + "version": 1, + "runs": { + "core_ssl_benchmark_rev-thesis-v1_pretrain_mae_miiv_seed42": { + "status": "running", + "pid": 123, + "command": "uv run stale", + } + }, + } + saved = {} + + monkeypatch.setattr(runner, "generate_all_runs", lambda: []) + monkeypatch.setattr(runner, "load_state", lambda: state) + monkeypatch.setattr( + runner, "save_state", lambda updated: saved.setdefault("state", updated) + ) + monkeypatch.setattr(runner, "is_pid_alive", lambda pid: False) + + runner.print_status(["core_ssl_benchmark"], "thesis-v1") + + assert ( + state["runs"]["core_ssl_benchmark_rev-thesis-v1_pretrain_mae_miiv_seed42"]["status"] + == "pending" + ) + assert saved["state"] is state + assert "Recovered stale run" in capsys.readouterr().out + + def test_scheduler_exit_code_flags_pending_and_running_runs(self): + import scripts.internal.run_experiments as runner + + runs = [ + runner.Run( + id="pending-run", + experiment_class="core_ssl_benchmark", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/pending", + ), + runner.Run( + id="running-run", + experiment_class="core_ssl_benchmark", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=123, + output_dir="outputs/running", + ), + ] + state = { + "version": 1, + "runs": { + "pending-run": {"status": "pending"}, + "running-run": {"status": "running"}, + }, + } + + assert runner._scheduler_exit_code(runs, state) == 1 + + def test_runner_no_longer_exposes_inherited_tag_command(self): + import scripts.internal.run_experiments as runner + + assert not hasattr(runner, "cmd_tag") + + def test_wandb_logger_includes_model_size_metadata(self, monkeypatch): + import slices.training.utils as training_utils + + captured = {} + + class DummyConfig: + def update(self, cfg): + captured["config_update"] = cfg + + class DummyExperiment: + config = DummyConfig() + + class DummyWandbLogger: + def __init__(self, **kwargs): + captured["kwargs"] = kwargs + self.experiment = DummyExperiment() + + monkeypatch.setattr(training_utils, "WandbLogger", DummyWandbLogger) + + cfg = OmegaConf.create( + { + "output_dir": "outputs/capacity_study/example", + "experiment_class": "capacity_study", + "model_size": "medium", + "source_dataset": "eicu", + "label_fraction": 0.1, + "training": {"freeze_encoder": False}, + "logging": { + "use_wandb": True, + "wandb_project": "slices", + "wandb_entity": None, + "run_name": "capacity_study_finetune_miiv_mae_mortality_24h_seed42", + "wandb_group": "finetune_miiv_mae_mortality_24h", + "wandb_tags": ["phase:finetune"], + }, + } + ) + + logger = training_utils.setup_wandb_logger(cfg) + + assert isinstance(logger, DummyWandbLogger) + assert captured["kwargs"]["name"] == ( + "capacity_study_finetune_miiv_mae_mortality_24h_seed42_from_eicu_medium_frac01" + ) + assert ( + captured["kwargs"]["group"] == "finetune_miiv_mae_mortality_24h_from_eicu_medium_frac01" + ) + assert "experiment_class:capacity_study" in captured["kwargs"]["tags"] + assert "model_size:medium" in captured["kwargs"]["tags"] + assert "source_dataset:eicu" in captured["kwargs"]["tags"] + assert "ablation:label-efficiency" in captured["kwargs"]["tags"] + assert "ablation:transfer" in captured["kwargs"]["tags"] + + def test_setup_wandb_logger_raises_clear_error_for_missing_entity(self, monkeypatch): + import wandb + + import slices.training.utils as training_utils + + class DummyWandbLogger: + def __init__(self, **kwargs): + self._entity = kwargs["entity"] + + @property + def experiment(self): + if self._entity == "missing-team": + raise wandb.errors.CommError( + "entity missing-team not found during upsertBucket" + ) + raise AssertionError("unexpected entity") + + monkeypatch.setattr(training_utils, "WandbLogger", DummyWandbLogger) + + cfg = OmegaConf.create( + { + "output_dir": "outputs/smoke/supervised_miiv_mortality_24h", + "training": {"freeze_encoder": False}, + "logging": { + "use_wandb": True, + "wandb_project": "slices-final", + "wandb_entity": "missing-team", + "run_name": "supervised_miiv_mortality_24h_seed42", + "wandb_group": "supervised_miiv_mortality_24h", + "wandb_tags": [], + }, + } + ) + + with pytest.raises( + training_utils.WandbEntityNotFoundError, + match="W&B entity 'missing-team' could not be found", + ): + training_utils.setup_wandb_logger(cfg) + + def test_training_entrypoints_exit_nonzero_on_wandb_entity_error(self): + import ast + + scripts = [ + Path("scripts/training/pretrain.py"), + Path("scripts/training/finetune.py"), + Path("scripts/training/supervised.py"), + Path("scripts/training/xgboost_baseline.py"), + ] + + for script in scripts: + tree = ast.parse(script.read_text()) + handlers = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.ExceptHandler) + and isinstance(node.type, ast.Name) + and node.type.id == "WandbEntityNotFoundError" + ] + assert handlers, f"{script} does not catch WandbEntityNotFoundError" + + for handler in handlers: + assert not any(isinstance(node, ast.Return) for node in ast.walk(handler)) + assert any( + isinstance(node, ast.Raise) + and isinstance(node.exc, ast.Call) + and isinstance(node.exc.func, ast.Name) + and node.exc.func.id == "SystemExit" + and len(node.exc.args) == 1 + and isinstance(node.exc.args[0], ast.Constant) + and node.exc.args[0].value == 1 + for node in ast.walk(handler) + ), f"{script} must exit nonzero on W&B entity errors" + + def test_transfer_finetune_command_propagates_source_dataset(self): + from scripts.internal.run_experiments import Run + + pretrain = Run( + id="pretrain_mae_miiv_seed42", + experiment_class="cross_dataset_transfer", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/cross_dataset_transfer/pretrain_mae_miiv_seed42", + ) + finetune = Run( + id="finetune_mae_eicu_seed42", + experiment_class="cross_dataset_transfer", + run_type="finetune", + paradigm="mae", + dataset="eicu", + seed=42, + output_dir="outputs/cross_dataset_transfer/finetune_mae_mortality_24h_eicu_seed42_from_miiv", + depends_on=[pretrain.id], + task="mortality_24h", + freeze_encoder=False, + source_dataset="miiv", + ) + + cmd = finetune.build_command({pretrain.id: pretrain, finetune.id: finetune}) + assert "+source_dataset=miiv" in cmd + + def test_linear_probe_finetune_command_uses_strict_linear_probe(self): + from scripts.internal.run_experiments import Run + + pretrain = Run( + id="pretrain_mae_miiv_seed42", + experiment_class="core_ssl_benchmark", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/pretrain_mae_miiv_seed42", + ) + probe = Run( + id="probe_mae_miiv_seed42", + experiment_class="core_ssl_benchmark", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/probe_mae_mortality_24h_miiv_seed42", + depends_on=[pretrain.id], + task="mortality_24h", + freeze_encoder=True, + ) + + cmd = probe.build_command({pretrain.id: pretrain, probe.id: probe}) + + assert "protocol=linear_probe" in cmd + assert "training.freeze_encoder=true" not in cmd + assert "task.head_type=linear" not in cmd + assert "task.hidden_dims=[]" not in cmd + assert "task.dropout=0.0" not in cmd + assert "protocol=a" not in cmd + assert "protocol=A" not in cmd + assert "+protocol=A" not in cmd + + def test_full_finetune_command_keeps_task_mlp_head(self): + from scripts.internal.run_experiments import Run + + pretrain = Run( + id="pretrain_mae_miiv_seed42", + experiment_class="core_ssl_benchmark", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/pretrain_mae_miiv_seed42", + ) + finetune = Run( + id="finetune_mae_miiv_seed42", + experiment_class="core_ssl_benchmark", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/finetune_mae_mortality_24h_miiv_seed42", + depends_on=[pretrain.id], + task="mortality_24h", + freeze_encoder=False, + ) + + cmd = finetune.build_command({pretrain.id: pretrain, finetune.id: finetune}) + + assert "protocol=full_finetune" in cmd + assert "training.freeze_encoder=false" not in cmd + assert "task.head_type=linear" not in cmd + assert "protocol=b" not in cmd + assert "protocol=B" not in cmd + assert "+protocol=B" not in cmd + + def test_xgboost_command_does_not_load_protocol_config_group(self): + from scripts.internal.run_experiments import Run + + run = Run( + id="xgboost_miiv_seed42", + experiment_class="classical_baselines", + run_type="xgboost", + paradigm="xgboost", + dataset="miiv", + seed=42, + output_dir="outputs/classical_baselines/xgboost_mortality_24h_miiv_seed42", + task="mortality_24h", + ) + + cmd = run.build_command({run.id: run}) + + assert "protocol=full_finetune" not in cmd + assert "+protocol=B" not in cmd + assert "protocol=B" not in cmd + + def test_generated_finetune_protocol_selectors_are_public_names(self): + import scripts.internal.run_experiments as runner + + runs = runner.generate_all_runs() + runs_by_id = {run.id: run for run in runs} + commands = [" ".join(run.build_command(runs_by_id)) for run in runs] + + assert not any("protocol=A" in command or "protocol=B" in command for command in commands) + assert not any("protocol=a" in command or "protocol=b" in command for command in commands) + assert sum("protocol=linear_probe" in command for command in commands) == 825 + assert sum("protocol=full_finetune" in command for command in commands) == 1020 + + def test_manual_finetune_protocol_configs_encode_safe_defaults(self): + linear_probe = OmegaConf.load("configs/protocol/linear_probe.yaml") + full_finetune = OmegaConf.load("configs/protocol/full_finetune.yaml") + + assert linear_probe.protocol == "linear_probe" + assert linear_probe.training.freeze_encoder is True + assert linear_probe.training.max_epochs == 50 + assert linear_probe.training.early_stopping_patience == 10 + assert linear_probe.optimizer.lr == pytest.approx(1.0e-4) + assert linear_probe.task.head_type == "linear" + assert list(linear_probe.task.hidden_dims) == [] + assert linear_probe.task.dropout == 0.0 + + assert full_finetune.protocol == "full_finetune" + assert full_finetune.training.freeze_encoder is False + assert full_finetune.training.max_epochs == 100 + assert full_finetune.training.early_stopping_patience == 10 + assert full_finetune.optimizer.lr == pytest.approx(3.0e-4) + assert full_finetune.task.head_type == "mlp" + assert list(full_finetune.task.hidden_dims) == [64] + assert full_finetune.task.dropout == 0.3 + + +class TestExperimentRunnerRetry: + """Tests for retry scoping and dependency preservation.""" + + def test_cmd_retry_respects_class_filter_but_keeps_dependencies(self, monkeypatch): + import scripts.internal.run_experiments as runner + + dependency = runner.Run( + id="dep_pretrain", + experiment_class="core_ssl_benchmark", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/pretrain_mae_miiv_seed42", + ) + target = runner.Run( + id="target_finetune", + experiment_class="label_efficiency", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/label_efficiency/finetune_mae_miiv_seed42", + depends_on=[dependency.id], + task="mortality_24h", + ) + other = runner.Run( + id="other_failed", + experiment_class="classical_baselines", + run_type="supervised", + paradigm="supervised", + dataset="miiv", + seed=42, + output_dir="outputs/classical_baselines/supervised_miiv_seed42", + task="mortality_24h", + ) + + state = { + "version": 1, + "runs": { + dependency.id: {"status": "completed"}, + target.id: {"status": "failed"}, + other.id: {"status": "failed"}, + }, + } + scheduled = {} + + monkeypatch.setattr(runner, "generate_all_runs", lambda: [dependency, target, other]) + monkeypatch.setattr(runner, "load_state", lambda: state) + monkeypatch.setattr(runner, "recover_stale_running", lambda state: None) + monkeypatch.setattr(runner, "save_state", lambda state: None) + monkeypatch.setattr( + runner, + "run_scheduler", + lambda runs, state, parallel, dry_run: scheduled.setdefault("runs", runs), + ) + + args = SimpleNamespace( + experiment_class=["label_efficiency"], + failed=True, + skipped=False, + revision=None, + reason=None, + project=None, + entity=None, + parallel=1, + ) + + runner.cmd_retry(args) + + scheduled_ids = [run.id for run in scheduled["runs"]] + assert target.id in scheduled_ids + assert dependency.id in scheduled_ids + assert other.id not in scheduled_ids + + def test_cmd_retry_skipped_reports_failed_dependencies(self, monkeypatch, capsys): + import scripts.internal.run_experiments as runner + + dependency = runner.Run( + id="dep_pretrain", + experiment_class="core_ssl_benchmark", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/pretrain_mae_miiv_seed42", + ) + target = runner.Run( + id="target_finetune", + experiment_class="label_efficiency", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/label_efficiency/finetune_mae_miiv_seed42", + depends_on=[dependency.id], + task="mortality_24h", + ) + + state = { + "version": 1, + "runs": { + dependency.id: {"status": "failed"}, + target.id: {"status": "skipped"}, + }, + } + + monkeypatch.setattr(runner, "generate_all_runs", lambda: [dependency, target]) + monkeypatch.setattr(runner, "load_state", lambda: state) + monkeypatch.setattr(runner, "recover_stale_running", lambda state: None) + monkeypatch.setattr( + runner, + "run_scheduler", + lambda runs, state, parallel, dry_run: pytest.fail("scheduler should not run"), + ) + + args = SimpleNamespace( + experiment_class=["label_efficiency"], + failed=False, + skipped=True, + revision=None, + reason=None, + project=None, + entity=None, + parallel=1, + ) + + with pytest.raises(SystemExit): + runner.cmd_retry(args) + + output = capsys.readouterr().out + assert "dep_pretrain [failed]" in output + assert ( + "uv run python scripts/internal/run_experiments.py retry --failed --skipped" in output + ) + assert state["runs"][target.id]["status"] == "skipped" + assert state["runs"][dependency.id]["status"] == "failed" + + def test_cmd_retry_skipped_resets_failed_dependencies_when_requested(self, monkeypatch): + import scripts.internal.run_experiments as runner + + dependency = runner.Run( + id="dep_pretrain", + experiment_class="core_ssl_benchmark", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/pretrain_mae_miiv_seed42", + ) + target = runner.Run( + id="target_finetune", + experiment_class="label_efficiency", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/label_efficiency/finetune_mae_miiv_seed42", + depends_on=[dependency.id], + task="mortality_24h", + ) + + state = { + "version": 1, + "runs": { + dependency.id: {"status": "failed"}, + target.id: {"status": "skipped"}, + }, + } + scheduled = {} + + monkeypatch.setattr(runner, "generate_all_runs", lambda: [dependency, target]) + monkeypatch.setattr(runner, "load_state", lambda: state) + monkeypatch.setattr(runner, "recover_stale_running", lambda state: None) + monkeypatch.setattr(runner, "save_state", lambda state: None) + monkeypatch.setattr( + runner, + "run_scheduler", + lambda runs, state, parallel, dry_run: scheduled.setdefault("runs", runs), + ) + + args = SimpleNamespace( + experiment_class=["label_efficiency"], + failed=True, + skipped=True, + revision=None, + reason=None, + project=None, + entity=None, + parallel=1, + ) + + runner.cmd_retry(args) + + scheduled_ids = [run.id for run in scheduled["runs"]] + assert dependency.id in scheduled_ids + assert target.id in scheduled_ids + assert state["runs"][dependency.id]["status"] == "pending" + assert state["runs"][target.id]["status"] == "pending" + + def test_cmd_retry_failed_surfaces_skipped_dependents_for_accounting(self, monkeypatch, capsys): + import scripts.internal.run_experiments as runner + + dependency = runner.Run( + id="dep_pretrain", + experiment_class="core_ssl_benchmark", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/pretrain_mae_miiv_seed42", + ) + target = runner.Run( + id="target_finetune", + experiment_class="label_efficiency", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/label_efficiency/finetune_mae_miiv_seed42", + depends_on=[dependency.id], + task="mortality_24h", + ) + + state = { + "version": 1, + "runs": { + dependency.id: {"status": "failed"}, + target.id: {"status": "skipped"}, + }, + } + scheduled = {} + + monkeypatch.setattr(runner, "generate_all_runs", lambda: [dependency, target]) + monkeypatch.setattr(runner, "load_state", lambda: state) + monkeypatch.setattr(runner, "recover_stale_running", lambda state: None) + monkeypatch.setattr(runner, "save_state", lambda state: None) + monkeypatch.setattr( + runner, + "run_scheduler", + lambda runs, state, parallel, dry_run: scheduled.setdefault("runs", runs), + ) + + args = SimpleNamespace( + experiment_class=["core_ssl_benchmark"], + failed=True, + skipped=False, + revision=None, + reason=None, + project=None, + entity=None, + parallel=1, + ) + + runner.cmd_retry(args) + + output = capsys.readouterr().out + scheduled_ids = [run.id for run in scheduled["runs"]] + assert dependency.id in scheduled_ids + assert target.id in scheduled_ids + assert state["runs"][dependency.id]["status"] == "pending" + assert state["runs"][target.id]["status"] == "skipped" + assert "target_finetune [skipped]" in output + assert "retry --failed --skipped" in output + + def test_cmd_retry_revises_dependencies_for_revision_scoped_retry(self, monkeypatch): + import scripts.internal.run_experiments as runner + + dependency = runner.Run( + id="core_ssl_benchmark_pretrain_mae_miiv_seed42", + run_key="pretrain_mae_miiv_seed42", + experiment_class="core_ssl_benchmark", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/pretrain_mae_miiv_seed42", + ) + target = runner.Run( + id="label_efficiency_finetune_mae_mortality_24h_miiv_seed42", + run_key="finetune_mae_mortality_24h_miiv_seed42", + experiment_class="label_efficiency", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/label_efficiency/finetune_mae_mortality_24h_miiv_seed42", + depends_on=[dependency.id], + task="mortality_24h", + ) + + revised_dep_id = "core_ssl_benchmark_rev-thesis-v1_pretrain_mae_miiv_seed42" + revised_target_id = "label_efficiency_rev-thesis-v1_finetune_mae_mortality_24h_miiv_seed42" + state = { + "version": 1, + "runs": { + revised_dep_id: {"status": "completed"}, + revised_target_id: {"status": "failed"}, + }, + } + scheduled = {} + + monkeypatch.setattr(runner, "generate_all_runs", lambda: [dependency, target]) + monkeypatch.setattr(runner, "load_state", lambda: state) + monkeypatch.setattr(runner, "recover_stale_running", lambda state: None) + monkeypatch.setattr(runner, "save_state", lambda state: None) + monkeypatch.setattr( + runner, + "run_scheduler", + lambda runs, state, parallel, dry_run: scheduled.setdefault("runs", runs), + ) + + args = SimpleNamespace( + experiment_class=["label_efficiency"], + failed=True, + skipped=False, + revision="thesis-v1", + reason="retry downstream failure", + project=None, + entity=None, + parallel=1, + ) + + runner.cmd_retry(args) + + scheduled_by_id = {run.id: run for run in scheduled["runs"]} + assert revised_dep_id in scheduled_by_id + assert revised_target_id in scheduled_by_id + assert scheduled_by_id[revised_target_id].depends_on == [revised_dep_id] + assert scheduled_by_id[revised_dep_id].extra_overrides["revision"] == "thesis-v1" + assert scheduled_by_id[revised_target_id].extra_overrides["revision"] == "thesis-v1" + + def test_cmd_retry_revision_scoped_failed_skipped_includes_cross_class_dependents( + self, + monkeypatch, + capsys, + ): + import scripts.internal.run_experiments as runner + + dependency = runner.Run( + id="core_ssl_benchmark_pretrain_mae_miiv_seed42", + run_key="pretrain_mae_miiv_seed42", + experiment_class="core_ssl_benchmark", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/pretrain_mae_miiv_seed42", + ) + target = runner.Run( + id="label_efficiency_finetune_mae_mortality_24h_miiv_seed42", + run_key="finetune_mae_mortality_24h_miiv_seed42", + experiment_class="label_efficiency", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/label_efficiency/finetune_mae_mortality_24h_miiv_seed42", + depends_on=[dependency.id], + task="mortality_24h", + ) + + revised_dep_id = "core_ssl_benchmark_rev-thesis-v1_pretrain_mae_miiv_seed42" + revised_target_id = "label_efficiency_rev-thesis-v1_finetune_mae_mortality_24h_miiv_seed42" + state = { + "version": 1, + "runs": { + revised_dep_id: {"status": "failed"}, + revised_target_id: {"status": "skipped"}, + }, + } + scheduled = {} + + monkeypatch.setattr(runner, "generate_all_runs", lambda: [dependency, target]) + monkeypatch.setattr(runner, "load_state", lambda: state) + monkeypatch.setattr(runner, "recover_stale_running", lambda state: None) + monkeypatch.setattr(runner, "save_state", lambda state: None) + monkeypatch.setattr( + runner, + "run_scheduler", + lambda runs, state, parallel, dry_run: scheduled.setdefault("runs", runs), + ) + + args = SimpleNamespace( + experiment_class=["core_ssl_benchmark"], + failed=True, + skipped=True, + revision="thesis-v1", + reason="retry upstream failure", + project=None, + entity=None, + parallel=1, + ) + + runner.cmd_retry(args) + + output = capsys.readouterr().out + scheduled_by_id = {run.id: run for run in scheduled["runs"]} + assert revised_dep_id in scheduled_by_id + assert revised_target_id in scheduled_by_id + assert scheduled_by_id[revised_target_id].depends_on == [revised_dep_id] + assert state["runs"][revised_dep_id]["status"] == "pending" + assert state["runs"][revised_target_id]["status"] == "pending" + assert "Including skipped downstream dependents" in output + + def test_select_ready_runs_prioritizes_pretrains_with_slot_budget(self): + import scripts.internal.run_experiments as runner + + pretrain = runner.Run( + id="pretrain", + experiment_class="core_ssl_benchmark", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/pretrain_mae_miiv_seed42", + ) + finetune = runner.Run( + id="finetune", + experiment_class="core_ssl_benchmark", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/finetune_mae_miiv_seed42", + task="mortality_24h", + ) + + selected = runner._select_ready_runs( + [finetune, pretrain], + active_run_ids=set(), + runs_by_id={pretrain.id: pretrain, finetune.id: finetune}, + slot_budget=4, + ) + + assert [run.id for run in selected] == [pretrain.id] + + def test_select_ready_runs_does_not_mix_pretrain_with_active_gpu_work(self): + import scripts.internal.run_experiments as runner + + pretrain = runner.Run( + id="pretrain", + experiment_class="core_ssl_benchmark", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/pretrain_mae_miiv_seed42", + ) + finetune = runner.Run( + id="finetune", + experiment_class="core_ssl_benchmark", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/finetune_mae_miiv_seed42", + task="mortality_24h", + ) + + selected = runner._select_ready_runs( + [pretrain], + active_run_ids={finetune.id}, + runs_by_id={pretrain.id: pretrain, finetune.id: finetune}, + slot_budget=4, + ) + + assert selected == [] + + def test_run_scheduler_returns_nonzero_for_failed_run(self, tmp_path, monkeypatch): + import scripts.internal.run_experiments as runner + + failed_run = runner.Run( + id="failed", + experiment_class="core_ssl_benchmark", + run_type="supervised", + paradigm="supervised", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/supervised_miiv_seed42", + task="mortality_24h", + ) + state = {"version": 1, "runs": {failed_run.id: {"status": "failed"}}} + + monkeypatch.setattr(runner, "LOG_DIR", tmp_path) + monkeypatch.setattr(runner, "save_state", lambda state: None) + + exit_code = runner.run_scheduler([failed_run], state, parallel=1, dry_run=False) + + assert exit_code == 1 + + def test_scheduler_rejects_completed_pretrain_missing_encoder(self, tmp_path, monkeypatch): + import scripts.internal.run_experiments as runner + + pretrain_dir = tmp_path / "pretrain" + finetune_dir = tmp_path / "finetune" + pretrain_dir.mkdir() + finetune_dir.mkdir() + pretrain = runner.Run( + id="pretrain", + experiment_class="core_ssl_benchmark", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir=str(pretrain_dir), + ) + finetune = runner.Run( + id="finetune", + experiment_class="core_ssl_benchmark", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir=str(finetune_dir), + depends_on=[pretrain.id], + task="mortality_24h", + freeze_encoder=False, + ) + state = { + "version": 1, + "runs": { + pretrain.id: {"status": "completed"}, + finetune.id: {"status": "pending"}, + }, + } + + monkeypatch.setattr(runner, "LOG_DIR", tmp_path / "logs") + monkeypatch.setattr(runner, "save_state", lambda state: None) + + exit_code = runner.run_scheduler([pretrain, finetune], state, parallel=1, dry_run=False) + + assert exit_code == 1 + assert state["runs"][pretrain.id]["status"] == "failed" + assert state["runs"][finetune.id]["status"] == "skipped" + assert "missing encoder.pt" in state["runs"][pretrain.id]["reason"] + + def test_scheduler_exit_code_nonzero_for_dependency_skips(self): + import scripts.internal.run_experiments as runner + + skipped_run = runner.Run( + id="skipped", + experiment_class="core_ssl_benchmark", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/core_ssl_benchmark/finetune_mae_miiv_seed42", + task="mortality_24h", + ) + state = {"version": 1, "runs": {skipped_run.id: {"status": "skipped"}}} + + assert runner._scheduler_exit_code([skipped_run], state) == 1 + + +class TestExportStatisticalScope: + """Regression tests for classical-baseline significance comparisons.""" + + def test_build_statistical_tests_df_includes_xgboost_comparisons(self): + from scripts.export_results import build_statistical_tests_df + + rows = [] + for paradigm, phase, protocol, offset in [ + ("xgboost", "baseline", "full_finetune", 0.08), + ("gru_d", "baseline", "full_finetune", 0.03), + ("supervised", "supervised", "full_finetune", 0.0), + ]: + for seed in [42, 123]: + for task, base in [("mortality_24h", 0.30), ("aki_kdigo", 0.25)]: + rows.append( + { + "experiment_class": "classical_baselines", + "experiment_type": "classical_baselines", + "paradigm": paradigm, + "dataset": "miiv", + "task": task, + "seed": seed, + "protocol": protocol, + "label_fraction": 1.0, + "model_size": "default", + "source_dataset": None, + "phase": phase, + "test/auprc": base + offset + (seed / 100000.0), + } + ) + + stats_df = build_statistical_tests_df(pd.DataFrame(rows)) + pairs = { + tuple(sorted((row["paradigm_a"], row["paradigm_b"]))) for _, row in stats_df.iterrows() + } + + assert ("gru_d", "supervised") in pairs + assert ("supervised", "xgboost") in pairs + assert ("gru_d", "xgboost") in pairs + per_task = stats_df[stats_df["comparison_scope"] == "per_task"] + assert set(per_task["task"]) == {"mortality_24h", "aki_kdigo"} + + +class TestXGBoostBaseline: + """Regression tests for XGBoost configuration choices.""" + + def test_xgboost_uses_primary_metric_for_early_stopping(self): + from scripts.training.xgboost_baseline import _xgboost_eval_metric + + assert _xgboost_eval_metric("binary") == "aucpr" + assert _xgboost_eval_metric("regression") == "mae" + + def test_xgboost_binary_ece_matches_uniform_bin_definition(self): + from scripts.training.xgboost_baseline import _binary_ece + + y_true = [0, 0, 1, 1] + y_prob = [0.1, 0.3, 0.7, 0.9] + + assert _binary_ece(y_true, y_prob, n_bins=2) == pytest.approx(0.2) + + def test_xgboost_config_caps_cpu_parallelism(self): + cfg = OmegaConf.load("configs/xgboost.yaml") + + assert cfg.xgboost.n_jobs == 4 + + def test_xgboost_balanced_scale_pos_weight_uses_native_ratio(self): + from scripts.training.xgboost_baseline import _resolve_scale_pos_weight + + labels = torch.tensor([0, 0, 0, 0, 1], dtype=torch.float32) + + assert _resolve_scale_pos_weight("balanced", labels.numpy()) == pytest.approx(4.0) + + +class TestTrainingScriptClassWeighting: + """Regression tests for entrypoint-level class weight resolution.""" + + @pytest.mark.parametrize( + "module_name", + ["scripts.training.finetune", "scripts.training.supervised"], + ) + def test_balanced_class_weight_uses_train_split_stats(self, monkeypatch, module_name): + """'balanced' must resolve from train-split stats before module construction.""" + module = importlib.import_module(module_name) + captured = {} + + class DummyDataModule: + def __init__(self, *args, **kwargs): + pass + + def setup(self): + return None + + def get_feature_dim(self): + return 3 + + def get_seq_length(self): + return 24 + + def get_split_info(self): + return { + "train_patients": 5, + "train_stays": 10, + "val_patients": 2, + "val_stays": 4, + "test_patients": 2, + "test_stays": 4, + } + + def get_label_statistics(self): + return { + "mortality_24h": { + "total": 100, + "positive": 20, + "negative": 80, + "prevalence": 0.2, + } + } + + def get_train_label_statistics(self, use_full_train: bool = False): + return { + "mortality_24h": { + "total": 10, + "positive": 1, + "negative": 9, + "prevalence": 0.1, + } + } + + class StopAfterCaptureError(Exception): + pass + + def fake_finetune_module(config, checkpoint_path=None, pretrain_checkpoint_path=None): + captured["class_weight"] = list(config.training.class_weight) + captured["d_input"] = config.encoder.d_input + captured["max_seq_length"] = config.encoder.max_seq_length + raise StopAfterCaptureError + + monkeypatch.setattr(module, "validate_data_prerequisites", lambda *args, **kwargs: None) + monkeypatch.setattr(module.pl, "seed_everything", lambda *args, **kwargs: None) + monkeypatch.setattr(module, "ICUDataModule", DummyDataModule) + monkeypatch.setattr(module, "FineTuneModule", fake_finetune_module) + if hasattr(module, "_detect_paradigm_from_checkpoint"): + monkeypatch.setattr(module, "_detect_paradigm_from_checkpoint", lambda *a, **k: None) + + cfg = OmegaConf.create( + { + "dataset": "miiv", + "seed": 42, + "paradigm": "mae", + "checkpoint": None, + "pretrain_checkpoint": "dummy.ckpt", + "data": { + "processed_dir": "/tmp/processed", + "num_workers": 0, + }, + "task": { + "task_name": "mortality_24h", + "task_type": "binary", + "head_type": "mlp", + "hidden_dims": [16], + "dropout": 0.0, + "activation": "relu", + }, + "encoder": { + "name": "transformer", + "d_input": 0, + "max_seq_length": 0, + }, + "training": { + "batch_size": 8, + "class_weight": "balanced", + "freeze_encoder": False, + }, + "optimizer": { + "name": "adam", + "lr": 1e-3, + }, + "logging": { + "use_wandb": False, + }, + } + ) + + with pytest.raises(StopAfterCaptureError): + module.main.__wrapped__(cfg) + + expected = [(10 / (2 * 9)) ** 0.5, (10 / (2 * 1)) ** 0.5] + assert captured["class_weight"] == pytest.approx(expected) + assert captured["d_input"] == 3 + assert captured["max_seq_length"] == 24 + + @pytest.mark.parametrize( + "module_name", + ["scripts.training.finetune", "scripts.training.supervised"], + ) + def test_regression_task_disables_balanced_class_weight(self, monkeypatch, module_name): + """Regression runs should not derive binary class weights.""" + module = importlib.import_module(module_name) + captured = {} + + class DummyDataModule: + def __init__(self, *args, **kwargs): + pass + + def setup(self): + return None + + def get_feature_dim(self): + return 3 + + def get_seq_length(self): + return 24 + + def get_split_info(self): + return { + "train_patients": 5, + "train_stays": 10, + "val_patients": 2, + "val_stays": 4, + "test_patients": 2, + "test_stays": 4, + } + + def get_label_statistics(self): + return { + "los_remaining": { + "task_type": "regression", + "total": 10, + "mean": 4.2, + "std": 1.1, + "min": 1.0, + "max": 8.0, + } + } + + def get_train_label_statistics(self, use_full_train: bool = False): + return { + "los_remaining": { + "task_type": "regression", + "total": 10, + "mean": 4.2, + "std": 1.1, + "min": 1.0, + "max": 8.0, + } + } + + class StopAfterCaptureError(Exception): + pass + + def fake_finetune_module(config, checkpoint_path=None, pretrain_checkpoint_path=None): + captured["class_weight"] = config.training.class_weight + raise StopAfterCaptureError + + monkeypatch.setattr(module, "validate_data_prerequisites", lambda *args, **kwargs: None) + monkeypatch.setattr(module, "report_and_validate_train_label_support", lambda *a, **k: None) + monkeypatch.setattr(module.pl, "seed_everything", lambda *args, **kwargs: None) + monkeypatch.setattr(module, "ICUDataModule", DummyDataModule) + monkeypatch.setattr(module, "FineTuneModule", fake_finetune_module) + if hasattr(module, "_detect_paradigm_from_checkpoint"): + monkeypatch.setattr(module, "_detect_paradigm_from_checkpoint", lambda *a, **k: None) + + cfg = OmegaConf.create( + { + "dataset": "miiv", + "seed": 42, + "paradigm": "mae", + "checkpoint": None, + "pretrain_checkpoint": "dummy.ckpt", + "data": { + "processed_dir": "/tmp/processed", + "num_workers": 0, + }, + "task": { + "task_name": "los_remaining", + "task_type": "regression", + "head_type": "mlp", + "hidden_dims": [16], + "dropout": 0.0, + "activation": "relu", + }, + "encoder": { + "name": "transformer", + "d_input": 0, + "max_seq_length": 0, + }, + "training": { + "batch_size": 8, + "class_weight": "balanced", + "freeze_encoder": False, + }, + "optimizer": { + "name": "adam", + "lr": 1e-3, + }, + "logging": { + "use_wandb": False, + }, + } + ) + + with pytest.raises(StopAfterCaptureError): + module.main.__wrapped__(cfg) + + assert captured["class_weight"] is None + + +class TestSupervisedBaselineMetrics: + """Regression tests for supervised trivial baselines.""" + + def test_classification_baseline_uses_train_majority_class(self): + module = importlib.import_module("scripts.training.supervised") + + class DummyDataset: + def __init__(self, labels): + self.labels = labels + + def __len__(self): + return len(self.labels) + + def __getitem__(self, idx): + return {"label": float(self.labels[idx])} + + class DummyLoader: + def __init__(self, dataset): + self.dataset = dataset + + class DummyDataModule: + def train_dataloader(self): + return DummyLoader(DummyDataset([1, 1, 1, 0])) + + def test_dataloader(self): + return DummyLoader(DummyDataset([0, 0, 0])) + + metrics = module.compute_baseline_metrics( + DummyDataModule(), + task_name="mortality_24h", + task_type="binary", + ) + + assert metrics["baseline/train_positive_ratio"] == pytest.approx(0.75) + assert metrics["baseline/majority_accuracy"] == pytest.approx(0.0) + + def test_regression_baseline_fits_on_train_labels_only(self): + module = importlib.import_module("scripts.training.supervised") + + class DummyDataset: + def __init__(self, labels): + self.labels = labels + + def __len__(self): + return len(self.labels) + + def __getitem__(self, idx): + return {"label": float(self.labels[idx])} + + class DummyLoader: + def __init__(self, dataset): + self.dataset = dataset + + class DummyDataModule: + def train_dataloader(self): + return DummyLoader(DummyDataset([10.0, 10.0])) + + def test_dataloader(self): + return DummyLoader(DummyDataset([0.0, 0.0])) + + metrics = module.compute_baseline_metrics( + DummyDataModule(), + task_name="los_remaining", + task_type="regression", + ) + + assert metrics["baseline/train_label_mean"] == pytest.approx(10.0) + assert metrics["baseline/mean_predictor_mse"] == pytest.approx(100.0) + assert metrics["baseline/median_predictor_mae"] == pytest.approx(10.0) + + +class TestFinetuneCheckpointParadigmDetection: + """Regression tests for checkpoint-driven finetune metadata.""" + + def test_entrypoint_rejects_ambiguous_checkpoint_sources(self): + module = importlib.import_module("scripts.training.finetune") + cfg = OmegaConf.create( + { + "checkpoint": "outputs/pretrain/encoder.pt", + "pretrain_checkpoint": "outputs/pretrain/last.ckpt", + } + ) + + with pytest.raises(ValueError, match="exactly one checkpoint source"): + module.main.__wrapped__(cfg) + + def test_pretrain_checkpoint_auto_detects_paradigm(self, monkeypatch, tmp_path): + module = importlib.import_module("scripts.training.finetune") + captured = {} + + class DummyDataModule: + def __init__(self, *args, **kwargs): + pass + + def setup(self): + return None + + def get_feature_dim(self): + return 3 + + def get_seq_length(self): + return 24 + + def get_split_info(self): + return { + "train_patients": 5, + "train_stays": 10, + "val_patients": 2, + "val_stays": 4, + "test_patients": 2, + "test_stays": 4, + } + + def get_label_statistics(self): + return { + "mortality_24h": { + "total": 100, + "positive": 20, + "negative": 80, + "prevalence": 0.2, + } + } + + def get_train_label_statistics(self, use_full_train: bool = False): + return { + "mortality_24h": { + "total": 10, + "positive": 2, + "negative": 8, + "prevalence": 0.2, + } + } + + class StopAfterCaptureError(Exception): + pass + + def fake_finetune_module(config, checkpoint_path=None, pretrain_checkpoint_path=None): + captured["paradigm"] = config.paradigm + captured["pretrain_checkpoint_path"] = pretrain_checkpoint_path + raise StopAfterCaptureError + + pretrain_checkpoint = tmp_path / "pretrain.ckpt" + torch.save( + { + "state_dict": {"encoder.weight": torch.zeros(1)}, + "hyper_parameters": { + "config": { + "ssl": {"name": "jepa"}, + "paradigm": "jepa", + } + }, + }, + pretrain_checkpoint, + ) + + monkeypatch.setattr(module, "validate_data_prerequisites", lambda *args, **kwargs: None) + monkeypatch.setattr(module, "report_and_validate_train_label_support", lambda *a, **k: None) + monkeypatch.setattr(module.pl, "seed_everything", lambda *args, **kwargs: None) + monkeypatch.setattr(module, "ICUDataModule", DummyDataModule) + monkeypatch.setattr(module, "FineTuneModule", fake_finetune_module) + + cfg = OmegaConf.create( + { + "dataset": "miiv", + "seed": 42, + "paradigm": "mae", + "checkpoint": None, + "pretrain_checkpoint": str(pretrain_checkpoint), + "data": { + "processed_dir": "/tmp/processed", + "num_workers": 0, + }, + "task": { + "task_name": "mortality_24h", + "task_type": "binary", + "prediction_window_hours": 24, + "observation_window_hours": 24, + "gap_hours": 0, + "label_sources": ["stays", "mortality_info"], + "label_params": {}, + "head_type": "mlp", + "hidden_dims": [16], + "dropout": 0.0, + "activation": "relu", + }, + "encoder": { + "name": "transformer", + "d_input": 0, + "max_seq_length": 0, + }, + "training": { + "batch_size": 8, + "class_weight": None, + "freeze_encoder": False, + }, + "optimizer": { + "name": "adam", + "lr": 1e-3, + }, + "logging": { + "use_wandb": False, + }, + } + ) + + with pytest.raises(StopAfterCaptureError): + module.main.__wrapped__(cfg) + + assert captured["paradigm"] == "jepa" + assert captured["pretrain_checkpoint_path"] == str(pretrain_checkpoint) + + +class TestExperimentRunnerMatrix: + """Tests for thesis-final experiment-class matrix coverage.""" + + def test_matrix_counts_match_final_experiment_classes(self): + from collections import Counter + + from scripts.internal.run_experiments import ( + EXPECTED_CLASS_COUNTS, + generate_all_runs, + scientific_fingerprint, + ) + + runs = generate_all_runs() + by_class = Counter(run.experiment_class for run in runs) + + assert len(runs) == 2590 + assert sum(run.experiment_class != "smart_external_reference" for run in runs) == 2455 + assert by_class == EXPECTED_CLASS_COUNTS + assert len({scientific_fingerprint(run) for run in runs}) == len(runs) + assert not any(run.task == "mortality" for run in runs) + assert not any( + run.task == "mortality_24h" + and run.label_fraction == 0.01 + and run.experiment_class + in {"label_efficiency", "capacity_study", "classical_baselines"} + for run in runs + ) + + label_eff_hospital_fractions = { + run.label_fraction + for run in runs + if run.experiment_class == "label_efficiency" + and run.task == "mortality_hospital" + and run.label_fraction < 1.0 + } + assert label_eff_hospital_fractions == {0.01, 0.05, 0.1, 0.25, 0.5} + + non_classical_baselines = [ + run + for run in runs + if run.run_type in ("gru_d", "xgboost") + and run.experiment_class != "classical_baselines" + ] + assert non_classical_baselines == [] + + contrastive_mask_sweep = [ + run + for run in runs + if run.experiment_class == "hp_robustness" + and run.paradigm == "contrastive" + and run.experiment_subtype == "view_mask_sensitivity" + and run.run_type == "pretrain" + ] + assert contrastive_mask_sweep + assert all( + run.extra_overrides["ssl.complementary_masks"] is False + for run in contrastive_mask_sweep + ) + + def test_capacity_study_expanded_to_five_seeds(self): + from scripts.internal.run_experiments import SEEDS_EXTENDED, MatrixBuilder + + builder = MatrixBuilder() + builder.build_capacity_study() + + assert len(builder.runs) == 100 + assert sorted({run.seed for run in builder.runs}) == SEEDS_EXTENDED + assert {run.model_size for run in builder.runs} == {"medium", "large"} + + def test_ts2vec_extension_includes_both_protocols(self): + from scripts.internal.run_experiments import MatrixBuilder + + builder = MatrixBuilder() + builder.build_ts2vec_extension() + + assert len(builder.runs) == 135 + pretrains = [run for run in builder.runs if run.run_type == "pretrain"] + finetunes = [run for run in builder.runs if run.run_type == "finetune"] + + assert len(pretrains) == 15 + assert sum(run.freeze_encoder is True for run in finetunes) == 60 + assert sum(run.freeze_encoder is False for run in finetunes) == 60 + + def test_smart_external_reference_uses_stable_default_model_size(self): + from scripts.internal.run_experiments import MatrixBuilder + + builder = MatrixBuilder() + builder.build_smart_external_reference() + + assert len(builder.runs) == 135 + assert {run.model_size for run in builder.runs} == {"default"} diff --git a/tests/test_gru_d_encoder.py b/tests/test_gru_d_encoder.py new file mode 100644 index 0000000..473ecd3 --- /dev/null +++ b/tests/test_gru_d_encoder.py @@ -0,0 +1,37 @@ +"""Tests for canonical GRU-D decay equations.""" + +import math + +import torch + +from slices.models.encoders.gru_d import GRUDConfig, GRUDEncoder + + +def test_gru_d_hidden_decay_uses_feature_wise_deltas(): + encoder = GRUDEncoder(GRUDConfig(d_input=2, d_model=2)) + with torch.no_grad(): + encoder.W_gamma_h.copy_(torch.eye(2)) + encoder.b_gamma_h.zero_() + + delta_t = torch.tensor([[1.0, 3.0]]) + + gamma_h = encoder._compute_hidden_decay(delta_t) + + torch.testing.assert_close(gamma_h, torch.tensor([[math.exp(-1.0), math.exp(-3.0)]])) + + +def test_gru_d_updates_last_observed_only_from_raw_observations(): + encoder = GRUDEncoder(GRUDConfig(d_input=2, d_model=2)) + with torch.no_grad(): + encoder.x_mean.zero_() + + x_last = torch.tensor([[2.0, 10.0]]) + x_t = torch.tensor([[99.0, 20.0]]) + mask_t = torch.tensor([[False, True]]) + gamma_x = torch.tensor([[0.5, 0.5]]) + + x_decayed = encoder._impute_inputs(x_t, mask_t, gamma_x, x_last) + updated = encoder._update_last_observed(x_t, mask_t, x_last) + + torch.testing.assert_close(x_decayed, torch.tensor([[1.0, 20.0]])) + torch.testing.assert_close(updated, torch.tensor([[2.0, 20.0]])) diff --git a/tests/test_imputation_eval.py b/tests/test_imputation_eval.py index 8ef3cb9..bb4b33b 100644 --- a/tests/test_imputation_eval.py +++ b/tests/test_imputation_eval.py @@ -9,12 +9,16 @@ - from_encoder_checkpoint() creates linear decoder """ +import importlib + import pytest import torch import torch.nn as nn -from slices.eval.imputation import ImputationEvaluator +from omegaconf import OmegaConf from torch.utils.data import DataLoader +from slices.eval.imputation import ImputationEvaluator + class SimpleEncoder(nn.Module): """Minimal encoder for testing.""" @@ -248,6 +252,43 @@ def test_evaluate_returns_feature_names(self): for key in results["nrmse_per_feature"]: assert key.startswith("feat_"), f"Expected feature name, got {key}" + def test_evaluate_uses_provided_feature_stds(self, monkeypatch): + """Explicit feature stds should be used instead of recomputing from the eval loader.""" + encoder = SimpleEncoder(d_input=5, d_model=8) + decoder = nn.Linear(8, 5) + nn.init.zeros_(decoder.weight) + nn.init.zeros_(decoder.bias) + + evaluator = ImputationEvaluator(encoder=encoder, decoder=decoder, d_input=5) + + timeseries = torch.ones(16, 12, 5) * 2.0 + mask = torch.ones(16, 12, 5, dtype=torch.bool) + + class ConstDataset: + def __len__(self): + return 16 + + def __getitem__(self, idx): + return {"timeseries": timeseries[idx], "mask": mask[idx]} + + loader = DataLoader(ConstDataset(), batch_size=16) + + def fail_if_called(_): + raise AssertionError( + "compute_feature_stds should not be called when feature_stds is set" + ) + + monkeypatch.setattr(evaluator, "compute_feature_stds", fail_if_called) + + results = evaluator.evaluate( + loader, + mask_strategy="random", + mask_ratio=0.5, + feature_stds={i: 2.0 for i in range(5)}, + ) + + assert results["nrmse_overall"] == pytest.approx(1.0, rel=1e-3) + class TestNRMSEComputation: """Tests for NRMSE computation correctness.""" @@ -353,6 +394,49 @@ def test_creates_linear_decoder(self, tmp_path): dec_out = evaluator.decoder(enc_out) assert dec_out.shape == (2, 12, 10) + def test_smart_encoder_checkpoint_adapts_to_timestep_probe(self, tmp_path): + """SMART pooling=none checkpoints should produce (B, T, D) reconstructions.""" + from slices.models.encoders import build_encoder + + encoder = build_encoder( + "smart", + { + "d_input": 5, + "d_model": 8, + "n_layers": 1, + "n_heads": 2, + "d_ff": 16, + "max_seq_length": 12, + "pooling": "none", + }, + ) + ckpt = { + "encoder_state_dict": encoder.state_dict(), + "encoder_config": { + "name": "smart", + "d_input": 5, + "d_model": 8, + "n_layers": 1, + "n_heads": 2, + "d_ff": 16, + "max_seq_length": 12, + "pooling": "none", + }, + "version": 3, + } + ckpt_path = tmp_path / "smart_encoder.pt" + torch.save(ckpt, ckpt_path) + + evaluator = ImputationEvaluator.from_encoder_checkpoint(str(ckpt_path), d_input=5) + dummy = torch.randn(2, 12, 5) + mask = torch.ones_like(dummy, dtype=torch.bool) + with torch.no_grad(): + enc_out = evaluator.encoder(dummy, mask=mask) + dec_out = evaluator.decoder(enc_out) + + assert enc_out.shape == (2, 12, 8) + assert dec_out.shape == (2, 12, 5) + class TestInit: """Tests for __init__ edge cases.""" @@ -369,3 +453,125 @@ def test_custom_decoder(self): decoder = nn.Linear(16, 10) evaluator = ImputationEvaluator(encoder, decoder=decoder) assert evaluator.decoder is decoder + + +class TestEvaluateImputationScript: + """Regression tests for the standalone imputation evaluation script.""" + + def test_pretrain_checkpoint_still_trains_probe_decoder(self, monkeypatch, tmp_path): + """MAE checkpoints must not skip probe-decoder training.""" + module = importlib.import_module("scripts.eval.evaluate_imputation") + captured = {} + + class DummyDataset: + def __len__(self): + return 4 + + def __getitem__(self, idx): + return { + "timeseries": torch.zeros(2, 3), + "mask": torch.ones(2, 3, dtype=torch.bool), + } + + def get_feature_names(self): + return ["a", "b", "c"] + + class DummyDataModule: + def __init__(self, *args, **kwargs): + self.dataset = DummyDataset() + self.train_indices = [0, 1] + self.pin_memory = False + + def setup(self): + return None + + def get_feature_dim(self): + return 3 + + def get_seq_length(self): + return 2 + + def train_dataloader(self): + return "train_loader" + + def test_dataloader(self): + return "test_loader" + + class DummyEvaluator: + def train_decoder(self, dataloader, max_epochs, lr): + captured["train_decoder"] = { + "dataloader": dataloader, + "max_epochs": max_epochs, + "lr": lr, + } + return {"train_losses": [1.0]} + + def compute_feature_stds(self, dataloader): + captured["feature_stds_loader"] = dataloader + return {0: 1.0, 1: 1.0, 2: 1.0} + + def evaluate(self, dataloader, mask_strategy, mask_ratio, feature_stds): + captured["evaluate"] = { + "dataloader": dataloader, + "mask_strategy": mask_strategy, + "mask_ratio": mask_ratio, + "feature_stds": feature_stds, + } + return { + "nrmse_per_feature": {"a": 0.1, "b": 0.2, "c": 0.3}, + "mae_overall": 0.2, + "nrmse_overall": 0.25, + } + + dummy_evaluator = DummyEvaluator() + + monkeypatch.setattr(module.pl, "seed_everything", lambda *args, **kwargs: None) + monkeypatch.setattr(module, "validate_data_prerequisites", lambda *args, **kwargs: None) + monkeypatch.setattr(module, "ICUDataModule", DummyDataModule) + monkeypatch.setattr(module, "DataLoader", lambda dataset, **kwargs: "train_stats_loader") + monkeypatch.setattr(module, "Subset", lambda dataset, indices: dataset) + monkeypatch.setattr( + module.ImputationEvaluator, + "from_mae_checkpoint", + staticmethod(lambda *args, **kwargs: dummy_evaluator), + ) + + cfg = OmegaConf.create( + { + "dataset": "miiv", + "seed": 42, + "data": { + "processed_dir": "data/processed/miiv", + "num_workers": 0, + }, + "batch_size": 8, + "checkpoint": None, + "pretrain_checkpoint": "outputs/pretrain/ssl-last.ckpt", + "masking": { + "strategies": ["random"], + "mask_ratio": 0.15, + }, + "reconstruction_head": { + "max_epochs": 3, + "lr": 1e-4, + }, + "logging": { + "use_wandb": False, + }, + "output_dir": str(tmp_path / "test-imputation"), + } + ) + + module.main.__wrapped__(cfg) + + assert captured["train_decoder"] == { + "dataloader": "train_loader", + "max_epochs": 3, + "lr": 1e-4, + } + assert captured["feature_stds_loader"] == "train_stats_loader" + assert captured["evaluate"]["dataloader"] == "test_loader" + assert captured["evaluate"]["mask_strategy"] == "random" + results_path = tmp_path / "test-imputation" / "imputation_results.json" + assert results_path.exists() + assert '"pretrain_checkpoint": "outputs/pretrain/ssl-last.ckpt"' in results_path.read_text() diff --git a/tests/test_inference.py b/tests/test_inference.py new file mode 100644 index 0000000..ca2e76d --- /dev/null +++ b/tests/test_inference.py @@ -0,0 +1,57 @@ +"""Tests for shared inference utilities.""" + +import torch + +from slices.eval.inference import run_inference + + +class _DummyRegressionModel(torch.nn.Module): + def forward(self, timeseries: torch.Tensor, mask: torch.Tensor): + del timeseries, mask + return {"probs": torch.tensor([[0.25], [0.75]], dtype=torch.float32)} + + +def test_run_inference_squeezes_singleton_regression_outputs(): + model = _DummyRegressionModel() + dataloader = [ + { + "timeseries": torch.zeros(2, 3, 4), + "mask": torch.ones(2, 3, 4, dtype=torch.bool), + "label": torch.tensor([1.0, 2.0]), + "stay_id": torch.tensor([10, 11]), + } + ] + + predictions, labels, stay_ids = run_inference(model, dataloader) + + assert predictions.shape == (2,) + assert labels.shape == (2,) + assert torch.allclose(predictions, torch.tensor([0.25, 0.75])) + assert torch.allclose(labels, torch.tensor([1.0, 2.0])) + assert stay_ids == [10, 11] + + +class _DummySingletonModel(torch.nn.Module): + def forward(self, timeseries: torch.Tensor, mask: torch.Tensor): + del timeseries, mask + return {"probs": torch.tensor([[0.5]], dtype=torch.float32)} + + +def test_run_inference_handles_singleton_batches_without_scalar_squeeze(): + model = _DummySingletonModel() + dataloader = [ + { + "timeseries": torch.zeros(1, 3, 4), + "mask": torch.ones(1, 3, 4, dtype=torch.bool), + "label": torch.tensor([1.0]), + "stay_id": torch.tensor([10]), + } + ] + + predictions, labels, stay_ids = run_inference(model, dataloader) + + assert predictions.shape == (1,) + assert labels.shape == (1,) + assert torch.allclose(predictions, torch.tensor([0.5])) + assert torch.allclose(labels, torch.tensor([1.0])) + assert stay_ids == [10] diff --git a/tests/test_jepa_objective.py b/tests/test_jepa_objective.py index ef0a2d1..37ef8d7 100644 --- a/tests/test_jepa_objective.py +++ b/tests/test_jepa_objective.py @@ -1,10 +1,9 @@ -"""Tests for observation-level JEPA (Joint-Embedding Predictive Architecture) SSL objective.""" +"""Tests for timestep-level JEPA (Joint-Embedding Predictive Architecture) SSL objective.""" import pytest import torch + from slices.models.encoders import ( - ObservationTransformerConfig, - ObservationTransformerEncoder, TransformerConfig, TransformerEncoder, ) @@ -15,6 +14,95 @@ get_ssl_config_class, ) +# ============================================================================= +# Block masking tests +# ============================================================================= + + +class TestJEPABlockMasking: + """Tests for JEPA block-mask budgeting.""" + + @pytest.mark.parametrize("mask_ratio,expected_masked", [(0.3, 7), (0.5, 12), (0.75, 18)]) + def test_block_mask_hits_requested_valid_budget(self, mask_ratio, expected_masked): + from slices.models.pretraining.masking import create_block_timestep_mask + + B, T = 128, 24 + valid_timestep_mask = torch.ones(B, T, dtype=torch.bool) + + ssl_mask = create_block_timestep_mask( + B, + T, + mask_ratio, + torch.device("cpu"), + n_blocks=3, + valid_timestep_mask=valid_timestep_mask, + ) + + actual_masked = ((~ssl_mask) & valid_timestep_mask).sum(dim=1) + assert actual_masked.unique().tolist() == [expected_masked] + + def test_block_mask_budgets_over_valid_timesteps_only(self): + from slices.models.pretraining.masking import create_block_timestep_mask + + B, T = 3, 12 + valid_timestep_mask = torch.zeros(B, T, dtype=torch.bool) + valid_timestep_mask[0, [0, 1, 2, 4, 5, 8, 9, 11]] = True + valid_timestep_mask[1, [1, 3, 5, 7, 9]] = True + valid_timestep_mask[2, [6]] = True + + ssl_mask = create_block_timestep_mask( + B, + T, + 0.75, + torch.device("cpu"), + n_blocks=3, + valid_timestep_mask=valid_timestep_mask, + ) + + actual_masked = ((~ssl_mask) & valid_timestep_mask).sum(dim=1) + actual_visible = (ssl_mask & valid_timestep_mask).sum(dim=1) + assert actual_masked.tolist() == [6, 4, 0] + assert actual_visible.tolist() == [2, 1, 1] + assert torch.all(ssl_mask[~valid_timestep_mask]) + + +class TestSharedTimestepMaskBudget: + """Core SSL objectives should share the same per-sample timestep budget.""" + + @pytest.mark.parametrize("mask_ratio", [0.3, 0.5, 0.75]) + def test_random_and_block_masks_share_integer_budget(self, mask_ratio): + from slices.models.pretraining.masking import ( + create_block_timestep_mask, + create_timestep_mask, + ) + + B, T = 24, 24 + valid_timestep_mask = torch.zeros(B, T, dtype=torch.bool) + for b, n_valid in enumerate(range(1, B + 1)): + valid_timestep_mask[b, :n_valid] = True + + random_mask = create_timestep_mask( + B, + T, + mask_ratio, + torch.device("cpu"), + valid_timestep_mask=valid_timestep_mask, + ) + block_mask = create_block_timestep_mask( + B, + T, + mask_ratio, + torch.device("cpu"), + n_blocks=3, + valid_timestep_mask=valid_timestep_mask, + ) + + random_masked = ((~random_mask) & valid_timestep_mask).sum(dim=1) + block_masked = ((~block_mask) & valid_timestep_mask).sum(dim=1) + assert random_masked.tolist() == block_masked.tolist() + assert torch.all(random_mask[~valid_timestep_mask]) + + # ============================================================================= # Predictor tests # ============================================================================= @@ -32,21 +120,19 @@ def test_predictor_output_shape(self): predictor_n_heads=2, predictor_d_ff=64, ) - predictor = JEPAPredictor(d_encoder=32, n_features=10, max_seq_length=48, config=config) + predictor = JEPAPredictor(d_encoder=32, max_seq_length=48, config=config) - B, max_obs, n_vis = 2, 20, 5 + B, T, n_vis = 2, 8, 3 encoded_visible = torch.randn(B, n_vis, 32) - ssl_mask = torch.ones(B, max_obs, dtype=torch.bool) - ssl_mask[:, :15] = False # First 15 masked + ssl_mask = torch.ones(B, T, dtype=torch.bool) + ssl_mask[:, :5] = False # First 5 masked token_info = { - "timestep_idx": torch.randint(0, 8, (B, max_obs)), - "feature_idx": torch.randint(0, 10, (B, max_obs)), + "timestep_idx": torch.arange(T).unsqueeze(0).expand(B, -1), } - padding_mask = torch.ones(B, max_obs, dtype=torch.bool) - pred = predictor(encoded_visible, ssl_mask, token_info, max_obs, padding_mask) - # Output should be (B, max_obs, d_encoder) - assert pred.shape == (B, max_obs, 32) + pred = predictor(encoded_visible, ssl_mask, token_info, T) + # Output should be (B, T, d_encoder) + assert pred.shape == (B, T, 32) def test_predictor_mask_token_is_learnable(self): from slices.models.pretraining.jepa import JEPAPredictor @@ -57,7 +143,7 @@ def test_predictor_mask_token_is_learnable(self): predictor_n_heads=2, predictor_d_ff=32, ) - predictor = JEPAPredictor(d_encoder=16, n_features=5, max_seq_length=10, config=config) + predictor = JEPAPredictor(d_encoder=16, max_seq_length=10, config=config) assert predictor.mask_token.requires_grad assert predictor.mask_token.shape == (1, 1, 16) @@ -72,23 +158,59 @@ def test_predictor_output_in_encoder_space(self): predictor_n_layers=1, predictor_n_heads=2, ) - predictor = JEPAPredictor( - d_encoder=d_encoder, n_features=10, max_seq_length=48, config=config - ) + predictor = JEPAPredictor(d_encoder=d_encoder, max_seq_length=48, config=config) - B, max_obs, n_vis = 2, 10, 3 + B, T, n_vis = 2, 8, 3 encoded_visible = torch.randn(B, n_vis, d_encoder) - ssl_mask = torch.ones(B, max_obs, dtype=torch.bool) - ssl_mask[:, :7] = False + ssl_mask = torch.ones(B, T, dtype=torch.bool) + ssl_mask[:, :5] = False token_info = { - "timestep_idx": torch.randint(0, 8, (B, max_obs)), - "feature_idx": torch.randint(0, 10, (B, max_obs)), + "timestep_idx": torch.arange(T).unsqueeze(0).expand(B, -1), } - padding_mask = torch.ones(B, max_obs, dtype=torch.bool) - pred = predictor(encoded_visible, ssl_mask, token_info, max_obs, padding_mask) + pred = predictor(encoded_visible, ssl_mask, token_info, T) assert pred.shape[-1] == d_encoder + def test_predictor_ignores_padded_visible_tokens(self): + """Padded visible-token rows must not overwrite predictor mask tokens.""" + from slices.models.pretraining.jepa import JEPAPredictor + + class KwargIdentity(torch.nn.Module): + def forward(self, x, **kwargs): + return x + + config = JEPAConfig(predictor_d_model=2, predictor_n_layers=1, predictor_n_heads=1) + predictor = JEPAPredictor(d_encoder=2, max_seq_length=3, config=config) + predictor.encoder_proj = torch.nn.Identity() + predictor.predictor = KwargIdentity() + predictor.embed_dropout = torch.nn.Identity() + predictor.output_proj = torch.nn.Identity() + predictor.time_pe.zero_() + with torch.no_grad(): + predictor.mask_token.fill_(-1.0) + + encoded_visible = torch.tensor( + [ + [[10.0, 10.0], [11.0, 11.0]], + [[20.0, 20.0], [99.0, 99.0]], + ] + ) + ssl_mask = torch.tensor( + [ + [True, True, False], + [True, False, False], + ] + ) + token_info = { + "timestep_idx": torch.arange(3).unsqueeze(0).expand(2, -1), + } + + pred = predictor(encoded_visible, ssl_mask, token_info, n_timesteps=3) + + assert torch.allclose(pred[1, 0], torch.tensor([20.0, 20.0])) + assert torch.allclose(pred[1, 1], torch.tensor([-1.0, -1.0])) + assert torch.allclose(pred[1, 2], torch.tensor([-1.0, -1.0])) + # ============================================================================= # Init validation tests @@ -100,10 +222,17 @@ class TestJEPAInit: @pytest.fixture def encoder(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - return ObservationTransformerEncoder(config) + return TransformerEncoder(config) @pytest.fixture def jepa_config(self): @@ -123,19 +252,24 @@ def test_initialization(self, encoder, jepa_config): assert hasattr(jepa, "target_encoder") assert jepa.missing_token is None - def test_requires_observation_encoder(self): + def test_requires_obs_aware(self): config = TransformerConfig(d_input=10, d_model=32, n_layers=1, n_heads=4, pooling="none") encoder = TransformerEncoder(config) jepa_config = JEPAConfig() - with pytest.raises(ValueError, match="tokenize.*encode"): + with pytest.raises(ValueError, match="obs_aware=True"): JEPAObjective(encoder, jepa_config) def test_requires_no_pooling(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, pooling="mean" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + pooling="mean", + obs_aware=True, ) - encoder = ObservationTransformerEncoder(config) + encoder = TransformerEncoder(config) jepa_config = JEPAConfig() with pytest.raises(ValueError, match="pooling='none'"): @@ -148,9 +282,7 @@ def test_target_encoder_is_frozen(self, encoder, jepa_config): def test_target_encoder_is_copy(self, encoder, jepa_config): jepa = JEPAObjective(encoder, jepa_config) - # Target should be a separate copy, not the same object assert jepa.target_encoder is not jepa.encoder - # But should have the same weights initially for p_online, p_target in zip(encoder.parameters(), jepa.target_encoder.parameters()): assert torch.allclose(p_online, p_target) @@ -165,10 +297,17 @@ class TestJEPAForward: @pytest.fixture def encoder(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - return ObservationTransformerEncoder(config) + return TransformerEncoder(config) @pytest.fixture def jepa_config(self): @@ -195,7 +334,7 @@ def test_forward_returns_loss_and_metrics(self, encoder, jepa_config): assert "jepa_loss" in metrics assert "ssl_loss" in metrics assert "jepa_mask_ratio_actual" in metrics - assert "jepa_n_tokens_per_sample" in metrics + assert "jepa_n_timesteps" in metrics assert "jepa_n_visible_per_sample" in metrics assert "jepa_n_masked_per_sample" in metrics assert "jepa_momentum" in metrics @@ -209,7 +348,6 @@ def test_backward(self, encoder, jepa_config): loss, _ = jepa(x, obs_mask) loss.backward() - # Check encoder and predictor have gradients has_encoder_grad = False for param in jepa.encoder.parameters(): if param.grad is not None and param.grad.abs().sum() > 0: @@ -224,7 +362,7 @@ def test_backward(self, encoder, jepa_config): break assert has_pred_grad - def test_encoder_sees_fewer_tokens(self, encoder, jepa_config): + def test_encoder_sees_fewer_timesteps(self, encoder, jepa_config): jepa = JEPAObjective(encoder, jepa_config) B, T, D = 4, 12, 10 @@ -233,11 +371,73 @@ def test_encoder_sees_fewer_tokens(self, encoder, jepa_config): _, metrics = jepa(x, obs_mask) - n_total = metrics["jepa_n_tokens_per_sample"] * B - n_visible = metrics["jepa_n_visible_per_sample"] * B - ratio = n_visible / n_total + n_visible = metrics["jepa_n_visible_per_sample"] + n_masked = metrics["jepa_n_masked_per_sample"] + ratio = n_visible / (n_visible + n_masked) assert 0.10 <= ratio <= 0.50 # ~25% visible + def test_empty_timesteps_are_excluded_from_ssl_counts(self, encoder, jepa_config): + """JEPA masking counts should only cover timesteps with observations.""" + jepa = JEPAObjective(encoder, jepa_config) + + B, T, D = 2, 8, 10 + x = torch.randn(B, T, D) + obs_mask = torch.zeros(B, T, D, dtype=torch.bool) + obs_mask[:, 1, :3] = True + obs_mask[:, 5, :3] = True + obs_mask[:, 6, :3] = True + + _, metrics = jepa(x, obs_mask) + + assert metrics["jepa_n_visible_per_sample"] + metrics[ + "jepa_n_masked_per_sample" + ] == pytest.approx(3.0) + + def test_empty_timesteps_are_masked_in_predictor_attention(self, encoder, jepa_config): + """Fully unobserved hours should not participate in JEPA predictor attention.""" + jepa = JEPAObjective(encoder, jepa_config) + captured: dict[str, torch.Tensor] = {} + + def capture_padding_mask(module, args, kwargs): + padding_mask = kwargs.get("src_key_padding_mask") + assert padding_mask is not None + captured["src_key_padding_mask"] = padding_mask.detach().clone() + + handle = jepa.predictor.predictor.register_forward_pre_hook( + capture_padding_mask, + with_kwargs=True, + ) + + try: + B, T, D = 2, 8, 10 + x = torch.randn(B, T, D) + obs_mask = torch.zeros(B, T, D, dtype=torch.bool) + obs_mask[:, 1, :3] = True + obs_mask[:, 5, :3] = True + obs_mask[:, 6, :3] = True + + jepa(x, obs_mask) + finally: + handle.remove() + + expected_padding_mask = ~obs_mask.any(dim=-1) + assert "src_key_padding_mask" in captured + assert torch.equal(captured["src_key_padding_mask"], expected_padding_mask) + + def test_all_empty_sample_keeps_loss_finite(self, encoder, jepa_config): + """A fully unobserved sample in the batch should not make JEPA loss NaN.""" + jepa = JEPAObjective(encoder, jepa_config) + + B, T, D = 2, 8, 10 + x = torch.randn(B, T, D) + obs_mask = torch.zeros(B, T, D, dtype=torch.bool) + obs_mask[1, 1:4, :3] = True + + loss, metrics = jepa(x, obs_mask) + + assert torch.isfinite(loss) + assert torch.isfinite(metrics["jepa_loss"]) + # ============================================================================= # Momentum tests @@ -249,10 +449,17 @@ class TestJEPAMomentum: @pytest.fixture def encoder(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - return ObservationTransformerEncoder(config) + return TransformerEncoder(config) @pytest.fixture def jepa_config(self): @@ -269,10 +476,8 @@ def jepa_config(self): def test_momentum_update_changes_target(self, encoder, jepa_config): jepa = JEPAObjective(encoder, jepa_config) - # Store original target params original_params = [p.clone() for p in jepa.target_encoder.parameters()] - # Do a training step to change online encoder x = torch.randn(2, 8, 10) obs_mask = torch.ones(2, 8, 10, dtype=torch.bool) loss, _ = jepa(x, obs_mask) @@ -280,10 +485,8 @@ def test_momentum_update_changes_target(self, encoder, jepa_config): optimizer = torch.optim.Adam([p for p in jepa.parameters() if p.requires_grad], lr=1e-2) optimizer.step() - # Now update momentum jepa.momentum_update(progress=0.5) - # Target should have changed changed = False for orig, new in zip(original_params, jepa.target_encoder.parameters()): if not torch.allclose(orig, new): @@ -307,7 +510,6 @@ def test_momentum_schedule(self, encoder, jepa_config): def test_ema_formula(self, encoder, jepa_config): jepa = JEPAObjective(encoder, jepa_config) - # Diverge online from target so EMA is non-trivial with torch.no_grad(): for p in jepa.encoder.parameters(): p.add_(torch.randn_like(p) * 0.5) @@ -315,10 +517,9 @@ def test_ema_formula(self, encoder, jepa_config): online_params = [p.clone() for p in jepa.encoder.parameters()] target_params = [p.clone() for p in jepa.target_encoder.parameters()] - m = 0.996 # momentum at progress=0 + m = 0.996 jepa.momentum_update(progress=0.0) - # Verify EMA: target = m * old_target + (1-m) * online for online, old_target, new_target in zip( online_params, target_params, @@ -326,7 +527,6 @@ def test_ema_formula(self, encoder, jepa_config): ): expected = m * old_target + (1 - m) * online assert torch.allclose(new_target, expected, atol=1e-6) - # Ensure the update actually changed the target (non-trivial test) assert not torch.allclose(new_target, old_target, atol=1e-6) @@ -340,10 +540,17 @@ class TestJEPAEdgeCases: @pytest.fixture def encoder(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - return ObservationTransformerEncoder(config) + return TransformerEncoder(config) def test_very_sparse_data(self, encoder): config = JEPAConfig( @@ -357,32 +564,13 @@ def test_very_sparse_data(self, encoder): B, T, D = 2, 8, 10 x = torch.randn(B, T, D) - obs_mask = torch.rand(B, T, D) > 0.95 # ~5% observed + obs_mask = torch.rand(B, T, D) > 0.95 obs_mask[0, 0, 0] = True obs_mask[1, 0, 0] = True loss, metrics = jepa(x, obs_mask) assert torch.isfinite(loss) - def test_single_observation(self, encoder): - config = JEPAConfig( - mask_ratio=0.75, - predictor_d_model=16, - predictor_n_layers=1, - predictor_n_heads=2, - predictor_d_ff=32, - ) - jepa = JEPAObjective(encoder, config) - - B, T, D = 2, 4, 10 - x = torch.randn(B, T, D) - obs_mask = torch.zeros(B, T, D, dtype=torch.bool) - obs_mask[0, 0, 0] = True - obs_mask[1, 1, 3] = True - - loss, metrics = jepa(x, obs_mask) - assert torch.isfinite(loss) - def test_batch_with_varying_sparsity(self, encoder): config = JEPAConfig( mask_ratio=0.75, @@ -396,9 +584,9 @@ def test_batch_with_varying_sparsity(self, encoder): B, T, D = 3, 8, 10 x = torch.randn(B, T, D) obs_mask = torch.zeros(B, T, D, dtype=torch.bool) - obs_mask[0] = True # All observed - obs_mask[1, :4, :5] = True # Half observed - obs_mask[2, 0, 0] = True # Almost nothing + obs_mask[0] = True + obs_mask[1, :4, :5] = True + obs_mask[2, 0, 0] = True loss, metrics = jepa(x, obs_mask) assert torch.isfinite(loss) @@ -414,10 +602,17 @@ class TestJEPAGradientFlow: @pytest.fixture def encoder(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - return ObservationTransformerEncoder(config) + return TransformerEncoder(config) def test_gradients_to_encoder_and_predictor(self, encoder): config = JEPAConfig( @@ -435,13 +630,11 @@ def test_gradients_to_encoder_and_predictor(self, encoder): loss, _ = jepa(x, obs_mask) loss.backward() - # Encoder should have gradients encoder_has_grad = any( p.grad is not None and p.grad.abs().sum() > 0 for p in jepa.encoder.parameters() ) assert encoder_has_grad - # Predictor should have gradients pred_has_grad = any( p.grad is not None and p.grad.abs().sum() > 0 for p in jepa.predictor.parameters() ) @@ -463,7 +656,6 @@ def test_no_gradients_to_target(self, encoder): loss, _ = jepa(x, obs_mask) loss.backward() - # Target encoder should NOT have gradients for param in jepa.target_encoder.parameters(): assert param.grad is None @@ -477,10 +669,17 @@ class TestJEPAConvergence: """Test that loss decreases during training.""" def test_loss_decreases(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=2, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=2, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - encoder = ObservationTransformerEncoder(config) + encoder = TransformerEncoder(config) jepa_config = JEPAConfig( mask_ratio=0.75, @@ -524,10 +723,17 @@ class TestJEPALossTypes: @pytest.fixture def encoder(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - return ObservationTransformerEncoder(config) + return TransformerEncoder(config) def test_mse_loss(self, encoder): config = JEPAConfig( @@ -580,10 +786,15 @@ def test_in_registry(self): assert get_ssl_config_class("jepa") == JEPAConfig def test_build_works(self): - encoder_config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, pooling="none" + encoder_config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + pooling="none", + obs_aware=True, ) - encoder = ObservationTransformerEncoder(encoder_config) + encoder = TransformerEncoder(encoder_config) jepa_config = JEPAConfig(mask_ratio=0.75) ssl_objective = build_ssl_objective(encoder, jepa_config) @@ -592,10 +803,15 @@ def test_build_works(self): assert ssl_objective.encoder is encoder def test_get_encoder(self): - encoder_config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, pooling="none" + encoder_config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + pooling="none", + obs_aware=True, ) - encoder = ObservationTransformerEncoder(encoder_config) + encoder = TransformerEncoder(encoder_config) jepa_config = JEPAConfig() jepa = JEPAObjective(encoder, jepa_config) diff --git a/tests/test_los_builder.py b/tests/test_los_builder.py index 01fc839..0d3cdd6 100644 --- a/tests/test_los_builder.py +++ b/tests/test_los_builder.py @@ -2,6 +2,7 @@ import polars as pl import pytest + from slices.data.labels import LabelBuilderFactory, LabelConfig from slices.data.labels.los import LOSLabelBuilder diff --git a/tests/test_mae_objective.py b/tests/test_mae_objective.py index ad72eb0..5234782 100644 --- a/tests/test_mae_objective.py +++ b/tests/test_mae_objective.py @@ -1,7 +1,8 @@ -"""Tests for observation-level MAE (Masked Autoencoder) SSL objective.""" +"""Tests for timestep-level MAE (Masked Autoencoder) SSL objective.""" import pytest import torch + from slices.data.transforms import ( MaskingStrategy, apply_mask, @@ -12,8 +13,6 @@ create_timestep_mask, ) from slices.models.encoders import ( - ObservationTransformerConfig, - ObservationTransformerEncoder, TransformerConfig, TransformerEncoder, ) @@ -92,88 +91,6 @@ def test_all_masking_strategies(self, strategy: MaskingStrategy): assert mask.dtype == torch.bool -# ============================================================================= -# Observation tokenization tests -# ============================================================================= - - -class TestObservationTokenization: - """Tests for observation-level tokenization.""" - - @pytest.fixture - def encoder(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="none" - ) - return ObservationTransformerEncoder(config) - - def test_token_count_equals_observed(self, encoder): - """Number of valid tokens should equal number of observed values.""" - B, T, D = 2, 8, 10 - x = torch.randn(B, T, D) - obs_mask = torch.rand(B, T, D) > 0.5 - - tokens, padding_mask, token_info = encoder.tokenize(x, obs_mask) - - for b in range(B): - expected_n = obs_mask[b].sum().item() - actual_n = padding_mask[b].sum().item() - assert actual_n == expected_n - - def test_token_values_match_input(self, encoder): - """Token values should correspond to actual observed input values.""" - B, T, D = 1, 4, 3 - x = torch.randn(B, T, D) - obs_mask = torch.ones(B, T, D, dtype=torch.bool) - obs_mask[0, 2, 1] = False # One missing value - - _, _, token_info = encoder.tokenize(x, obs_mask) - n_obs = token_info["n_obs"][0].item() - - # All observed values should be in token_info["values"] - observed_vals = x[obs_mask].sort()[0] - token_vals = token_info["values"][0, :n_obs].sort()[0] - assert torch.allclose(observed_vals, token_vals) - - def test_empty_obs_mask_gives_at_least_one_token(self, encoder): - """Even with no observations, we get at least 1 token (padded).""" - B, T, D = 1, 4, 3 - x = torch.randn(B, T, D) - obs_mask = torch.zeros(B, T, D, dtype=torch.bool) - - tokens, padding_mask, token_info = encoder.tokenize(x, obs_mask) - - assert tokens.shape[1] >= 1 - assert padding_mask.shape[1] >= 1 - # No valid tokens - assert padding_mask.sum().item() == 0 - - def test_all_observed_token_count(self, encoder): - """With all observed, tokens = T * D.""" - B, T, D = 1, 4, 10 - x = torch.randn(B, T, D) - obs_mask = torch.ones(B, T, D, dtype=torch.bool) - - _, padding_mask, token_info = encoder.tokenize(x, obs_mask) - - assert padding_mask.sum().item() == T * D - - def test_feature_and_timestep_indices_valid(self, encoder): - """Feature and timestep indices should be within valid ranges.""" - B, T, D = 2, 8, 10 - x = torch.randn(B, T, D) - obs_mask = torch.rand(B, T, D) > 0.3 - - _, padding_mask, token_info = encoder.tokenize(x, obs_mask) - - for b in range(B): - n = padding_mask[b].sum().item() - assert token_info["timestep_idx"][b, :n].max() < T - assert token_info["feature_idx"][b, :n].max() < D - assert token_info["timestep_idx"][b, :n].min() >= 0 - assert token_info["feature_idx"][b, :n].min() >= 0 - - # ============================================================================= # MAE Decoder tests # ============================================================================= @@ -190,18 +107,16 @@ def test_decoder_output_shape(self): ) decoder = MAEDecoder(d_encoder=32, n_features=10, max_seq_length=48, config=config) - B, max_obs, n_vis = 2, 20, 5 + B, T, n_vis = 2, 8, 3 encoded_visible = torch.randn(B, n_vis, 32) - ssl_mask = torch.ones(B, max_obs, dtype=torch.bool) - ssl_mask[:, :15] = False # First 15 are masked + ssl_mask = torch.ones(B, T, dtype=torch.bool) + ssl_mask[:, :5] = False # First 5 are masked token_info = { - "timestep_idx": torch.randint(0, 8, (B, max_obs)), - "feature_idx": torch.randint(0, 10, (B, max_obs)), + "timestep_idx": torch.arange(T).unsqueeze(0).expand(B, -1), } - padding_mask = torch.ones(B, max_obs, dtype=torch.bool) - pred = decoder(encoded_visible, ssl_mask, token_info, max_obs, padding_mask) - assert pred.shape == (B, max_obs) + pred = decoder(encoded_visible, ssl_mask, token_info, T) + assert pred.shape == (B, T, 10) # (B, T, D) def test_decoder_mask_token_is_learnable(self): from slices.models.pretraining.mae import MAEDecoder @@ -224,6 +139,76 @@ def test_decoder_mask_token_in_decoder_space(self): assert decoder.mask_token.shape[-1] == d_decoder + def test_decoder_ignores_padded_visible_tokens(self): + """Padded visible-token rows must not overwrite real masked positions.""" + from slices.models.pretraining.mae import MAEDecoder + + config = MAEConfig(decoder_d_model=2, decoder_n_layers=1, decoder_n_heads=1) + decoder = MAEDecoder(d_encoder=2, n_features=2, max_seq_length=3, config=config) + decoder.encoder_proj = torch.nn.Identity() + decoder.decoder = torch.nn.Identity() + decoder.embed_dropout = torch.nn.Identity() + decoder.output_proj = torch.nn.Identity() + decoder.time_pe.zero_() + with torch.no_grad(): + decoder.mask_token.fill_(-1.0) + + encoded_visible = torch.tensor( + [ + [[10.0, 10.0], [11.0, 11.0]], + [[20.0, 20.0], [99.0, 99.0]], + ] + ) + ssl_mask = torch.tensor( + [ + [True, True, False], + [True, False, False], + ] + ) + token_info = { + "timestep_idx": torch.arange(3).unsqueeze(0).expand(2, -1), + } + + pred = decoder(encoded_visible, ssl_mask, token_info, n_timesteps=3) + + assert torch.allclose(pred[1, 0], torch.tensor([20.0, 20.0])) + assert torch.allclose(pred[1, 1], torch.tensor([-1.0, -1.0])) + assert torch.allclose(pred[1, 2], torch.tensor([-1.0, -1.0])) + + def test_decoder_masks_empty_timesteps_from_attention_context(self): + """Fully unobserved timesteps should not contribute decoder context.""" + from slices.models.pretraining.mae import MAEDecoder + + class ContextSumDecoder(torch.nn.Module): + def forward(self, src, src_key_padding_mask=None): + if src_key_padding_mask is not None: + src = src.masked_fill(src_key_padding_mask.unsqueeze(-1), 0.0) + context = src.sum(dim=1, keepdim=True) + return context.expand_as(src) + + config = MAEConfig(decoder_d_model=2, decoder_n_layers=1, decoder_n_heads=1) + decoder = MAEDecoder(d_encoder=2, n_features=2, max_seq_length=3, config=config) + decoder.encoder_proj = torch.nn.Identity() + decoder.decoder = ContextSumDecoder() + decoder.embed_dropout = torch.nn.Identity() + decoder.output_proj = torch.nn.Identity() + decoder.time_pe.zero_() + with torch.no_grad(): + decoder.mask_token.fill_(100.0) + + encoded_visible = torch.tensor([[[1.0, 0.0], [3.0, 0.0]]]) + ssl_mask = torch.tensor([[True, False, True]]) + token_info = { + "timestep_idx": torch.arange(3).unsqueeze(0), + "valid_timestep_mask": torch.tensor([[True, False, True]]), + } + + pred = decoder(encoded_visible, ssl_mask, token_info, n_timesteps=3) + + expected_context = torch.tensor([4.0, 0.0]) + assert torch.allclose(pred[0, 0], expected_context) + assert torch.allclose(pred[0, 2], expected_context) + # ============================================================================= # MAE Objective tests @@ -231,14 +216,21 @@ def test_decoder_mask_token_in_decoder_space(self): class TestMAEObjective: - """Tests for observation-level MAE objective.""" + """Tests for timestep-level MAE objective.""" @pytest.fixture def encoder(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - return ObservationTransformerEncoder(config) + return TransformerEncoder(config) @pytest.fixture def mae_config(self): @@ -255,7 +247,6 @@ def test_mae_initialization(self, encoder, mae_config): assert mae.encoder is encoder assert mae.config == mae_config assert hasattr(mae, "decoder") - # No MISSING_TOKEN in observation-level MAE assert mae.missing_token is None def test_mae_forward(self, encoder, mae_config): @@ -274,9 +265,10 @@ def test_mae_forward(self, encoder, mae_config): assert "mae_recon_loss_masked" in metrics assert "mae_recon_loss_visible" in metrics assert "mae_mask_ratio_actual" in metrics - assert "mae_n_tokens_per_sample" in metrics + assert "mae_n_timesteps" in metrics assert "mae_n_visible_per_sample" in metrics assert "mae_n_masked_per_sample" in metrics + assert "mae_n_loss_positions" in metrics def test_mae_backward(self, encoder, mae_config): mae = MAEObjective(encoder, mae_config) @@ -291,8 +283,8 @@ def test_mae_backward(self, encoder, mae_config): if param.requires_grad: assert param.grad is not None, f"No gradient for {name}" - def test_encoder_sees_fewer_tokens(self, encoder, mae_config): - """Encoder should process fewer tokens than total observations (75% masked).""" + def test_encoder_sees_fewer_timesteps(self, encoder, mae_config): + """Encoder should process fewer timesteps than total (75% masked).""" mae = MAEObjective(encoder, mae_config) B, T, D = 4, 12, 10 @@ -301,16 +293,14 @@ def test_encoder_sees_fewer_tokens(self, encoder, mae_config): loss, metrics = mae(x, obs_mask) - B = 4 - n_total = metrics["mae_n_tokens_per_sample"] * B - n_visible = metrics["mae_n_visible_per_sample"] * B + n_visible = metrics["mae_n_visible_per_sample"] + n_masked = metrics["mae_n_masked_per_sample"] # Visible should be ~25% of total - assert n_visible < n_total - ratio = n_visible / n_total - assert 0.10 <= ratio <= 0.50 # Allow tolerance around 0.25 + ratio = n_visible / (n_visible + n_masked) + assert 0.10 <= ratio <= 0.50 - def test_loss_only_on_masked(self, encoder, mae_config): - """Loss should only be computed on masked positions.""" + def test_loss_only_on_masked_and_observed(self, encoder, mae_config): + """Loss should only be computed on observed features at masked timesteps.""" mae = MAEObjective(encoder, mae_config) B, T, D = 2, 8, 10 @@ -319,27 +309,30 @@ def test_loss_only_on_masked(self, encoder, mae_config): loss, metrics = mae(x, obs_mask) - B = 2 n_masked = metrics["mae_n_masked_per_sample"] * B - n_total = metrics["mae_n_tokens_per_sample"] * B assert n_masked > 0 - assert n_masked < n_total + assert metrics["mae_n_loss_positions"] > 0 - def test_mae_requires_observation_encoder(self): - """MAE should reject encoders without tokenize/encode methods.""" + def test_mae_requires_obs_aware(self): + """MAE should reject encoders without obs_aware=True.""" config = TransformerConfig(d_input=10, d_model=32, n_layers=1, n_heads=4, pooling="none") encoder = TransformerEncoder(config) mae_config = MAEConfig() - with pytest.raises(ValueError, match="tokenize.*encode"): + with pytest.raises(ValueError, match="obs_aware=True"): MAEObjective(encoder, mae_config) def test_mae_requires_no_pooling(self): """MAE should reject encoder with pooling != 'none'.""" - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, pooling="mean" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + pooling="mean", + obs_aware=True, ) - encoder = ObservationTransformerEncoder(config) + encoder = TransformerEncoder(config) mae_config = MAEConfig() with pytest.raises(ValueError, match="pooling='none'"): @@ -360,7 +353,7 @@ def test_mae_training_step(self, encoder, mae_config): assert torch.isfinite(loss) def test_mae_respects_obs_mask(self, encoder, mae_config): - """With sparse data, tokens should match observed count.""" + """With sparse data, loss positions should only be on observed features.""" mae = MAEObjective(encoder, mae_config) B, T, D = 2, 8, 10 @@ -370,9 +363,95 @@ def test_mae_respects_obs_mask(self, encoder, mae_config): loss, metrics = mae(x, obs_mask) assert torch.isfinite(loss) - # n_tokens should reflect observed count - expected_obs = obs_mask.sum().item() - assert metrics["mae_n_tokens_per_sample"] * B == expected_obs + # Loss positions should be <= observed features at masked timesteps + n_loss = metrics["mae_n_loss_positions"] + n_observed = obs_mask.sum().item() + assert n_loss <= n_observed + + def test_empty_timesteps_are_excluded_from_mae_counts(self, encoder, mae_config): + """MAE masking metrics should only count timesteps with observations.""" + mae = MAEObjective(encoder, mae_config) + + B, T, D = 2, 8, 10 + x = torch.randn(B, T, D) + obs_mask = torch.zeros(B, T, D, dtype=torch.bool) + obs_mask[:, 1, :3] = True + obs_mask[:, 5, :3] = True + obs_mask[:, 6, :3] = True + + _, metrics = mae(x, obs_mask) + + assert metrics["mae_n_visible_per_sample"] + metrics[ + "mae_n_masked_per_sample" + ] == pytest.approx(3.0) + + def test_mae_passes_valid_timestep_mask_to_masking(self, encoder, mae_config, monkeypatch): + """Fully unobserved timesteps should be excluded from MAE masking.""" + mae = MAEObjective(encoder, mae_config) + captured = {} + + def fake_create_timestep_mask( + batch_size, + n_timesteps, + mask_ratio, + device, + valid_timestep_mask=None, + ): + captured["batch_size"] = batch_size + captured["n_timesteps"] = n_timesteps + captured["mask_ratio"] = mask_ratio + captured["device"] = device + captured["masking_valid_timestep_mask"] = valid_timestep_mask.clone() + return torch.ones(batch_size, n_timesteps, dtype=torch.bool, device=device) + + def fake_extract_visible_timesteps( + tokens, + ssl_mask, + valid_timestep_mask=None, + ): + captured["extract_valid_timestep_mask"] = valid_timestep_mask.clone() + visible_tokens = tokens[:, :1, :] + vis_padding = torch.ones( + tokens.shape[0], + 1, + dtype=torch.bool, + device=tokens.device, + ) + return visible_tokens, vis_padding + + monkeypatch.setattr( + "slices.models.pretraining.mae.create_timestep_mask", + fake_create_timestep_mask, + ) + monkeypatch.setattr( + "slices.models.pretraining.mae.extract_visible_timesteps", + fake_extract_visible_timesteps, + ) + + x = torch.randn(2, 4, 10) + obs_mask = torch.tensor( + [ + [ + [True] * 10, + [False] * 10, + [True] * 10, + [False] * 10, + ], + [ + [False] * 10, + [True] * 10, + [True] * 10, + [False] * 10, + ], + ], + dtype=torch.bool, + ) + + mae(x, obs_mask) + + expected = obs_mask.any(dim=-1) + assert torch.equal(captured["masking_valid_timestep_mask"], expected) + assert torch.equal(captured["extract_valid_timestep_mask"], expected) def test_mae_get_encoder(self, encoder, mae_config): mae = MAEObjective(encoder, mae_config) @@ -389,10 +468,17 @@ class TestMaskRatios: @pytest.fixture def encoder(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - return ObservationTransformerEncoder(config) + return TransformerEncoder(config) @pytest.mark.parametrize("mask_ratio", [0.25, 0.5, 0.75, 0.9]) def test_mask_ratios(self, encoder, mask_ratio): @@ -405,18 +491,18 @@ def test_mask_ratios(self, encoder, mask_ratio): ) mae = MAEObjective(encoder, config) - x = torch.randn(4, 8, 10) - obs_mask = torch.ones(4, 8, 10, dtype=torch.bool) + x = torch.randn(4, 32, 10) + obs_mask = torch.ones(4, 32, 10, dtype=torch.bool) loss, metrics = mae(x, obs_mask) assert torch.isfinite(loss) - # Check actual ratio is close to target + # Check actual ratio is close to target (32 timesteps → low discretization noise) actual = metrics["mae_mask_ratio_actual"] assert abs(actual - mask_ratio) < 0.15 def test_high_mask_ratio_at_least_one_visible(self, encoder): - """Even with very high mask ratio, at least 1 visible token per sample.""" + """Even with very high mask ratio, at least 1 visible timestep per sample.""" config = MAEConfig( mask_ratio=0.99, decoder_d_model=16, @@ -431,7 +517,7 @@ def test_high_mask_ratio_at_least_one_visible(self, encoder): loss, metrics = mae(x, obs_mask) assert torch.isfinite(loss) - assert metrics["mae_n_visible_per_sample"] * 4 >= 4 # At least 1 per sample + assert metrics["mae_n_visible_per_sample"] >= 1 # At least 1 per sample # ============================================================================= @@ -444,10 +530,17 @@ class TestEdgeCases: @pytest.fixture def encoder(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - return ObservationTransformerEncoder(config) + return TransformerEncoder(config) def test_very_sparse_data(self, encoder): """Test with very few observations per sample.""" @@ -472,29 +565,6 @@ def test_very_sparse_data(self, encoder): loss, metrics = mae(x, obs_mask) assert torch.isfinite(loss) - def test_single_observation(self, encoder): - """Test with exactly 1 observed value per sample.""" - mae_config = MAEConfig( - mask_ratio=0.75, - decoder_d_model=16, - decoder_n_layers=1, - decoder_n_heads=2, - decoder_d_ff=32, - ) - mae = MAEObjective(encoder, mae_config) - - B, T, D = 2, 4, 10 - x = torch.randn(B, T, D) - obs_mask = torch.zeros(B, T, D, dtype=torch.bool) - obs_mask[0, 0, 0] = True - obs_mask[1, 1, 3] = True - - loss, metrics = mae(x, obs_mask) - assert torch.isfinite(loss) - # With 1 obs, at least 1 must be visible, so 0 masked → loss should be 0 - # (no masked tokens to compute loss on) - assert loss.item() < 1e-7 - def test_batch_with_varying_sparsity(self, encoder): """Test batch where samples have different observation counts.""" mae_config = MAEConfig( @@ -527,10 +597,17 @@ class TestGradientFlow: @pytest.fixture def encoder(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - return ObservationTransformerEncoder(config) + return TransformerEncoder(config) def test_gradients_flow_to_encoder(self, encoder): mae_config = MAEConfig( @@ -576,24 +653,6 @@ def test_gradients_flow_to_decoder(self, encoder): assert mae.decoder.mask_token.grad is not None assert mae.decoder.mask_token.grad.abs().sum() > 0 - def test_gradients_flow_to_decoder_feature_embed(self, encoder): - mae_config = MAEConfig( - mask_ratio=0.75, - decoder_d_model=16, - decoder_n_layers=1, - decoder_n_heads=2, - decoder_d_ff=32, - ) - mae = MAEObjective(encoder, mae_config) - - x = torch.randn(2, 8, 10) - obs_mask = torch.ones(2, 8, 10, dtype=torch.bool) - - loss, _ = mae(x, obs_mask) - loss.backward() - - assert mae.decoder.feature_embed.weight.grad is not None - # ============================================================================= # Training convergence test @@ -604,10 +663,17 @@ class TestTrainingConvergence: """Test that loss decreases during training.""" def test_loss_decreases(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=2, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=2, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - encoder = ObservationTransformerEncoder(config) + encoder = TransformerEncoder(config) mae_config = MAEConfig( mask_ratio=0.75, @@ -642,10 +708,17 @@ def test_loss_decreases(self): def test_training_masks_vary(self): """Training masks should vary between forward passes.""" - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="none" + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, ) - encoder = ObservationTransformerEncoder(config) + encoder = TransformerEncoder(config) mae_config = MAEConfig( mask_ratio=0.5, @@ -675,13 +748,18 @@ def test_training_masks_vary(self): class TestSSLFactory: - """Tests for SSL factory with observation-level MAE.""" + """Tests for SSL factory with timestep-level MAE.""" def test_build_ssl_objective_mae(self): - encoder_config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, pooling="none" + encoder_config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + pooling="none", + obs_aware=True, ) - encoder = ObservationTransformerEncoder(encoder_config) + encoder = TransformerEncoder(encoder_config) mae_config = MAEConfig(mask_ratio=0.75) ssl_objective = build_ssl_objective(encoder, mae_config) @@ -690,8 +768,13 @@ def test_build_ssl_objective_mae(self): assert ssl_objective.encoder is encoder def test_build_ssl_objective_unknown(self): - encoder_config = ObservationTransformerConfig(d_input=10, d_model=32, pooling="none") - encoder = ObservationTransformerEncoder(encoder_config) + encoder_config = TransformerConfig( + d_input=10, + d_model=32, + pooling="none", + obs_aware=True, + ) + encoder = TransformerEncoder(encoder_config) from slices.models.pretraining.base import SSLConfig @@ -709,56 +792,5 @@ def test_get_ssl_config_class_unknown(self): get_ssl_config_class("unknown_objective") -# ============================================================================= -# Observation encoder forward (finetuning mode) tests -# ============================================================================= - - -class TestObservationEncoderForward: - """Test observation encoder in forward/finetuning mode.""" - - def test_forward_with_mean_pooling(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="mean" - ) - encoder = ObservationTransformerEncoder(config) - - B, T, D = 4, 8, 10 - x = torch.randn(B, T, D) - obs_mask = torch.rand(B, T, D) > 0.3 - - out = encoder(x, mask=obs_mask) - assert out.shape == (B, 32) # (B, d_model) - - def test_forward_with_no_pooling(self): - config = ObservationTransformerConfig( - d_input=10, d_model=32, n_layers=1, n_heads=4, d_ff=64, pooling="none" - ) - encoder = ObservationTransformerEncoder(config) - - B, T, D = 2, 4, 10 - x = torch.randn(B, T, D) - obs_mask = torch.ones(B, T, D, dtype=torch.bool) - - out = encoder(x, mask=obs_mask) - assert out.shape == (B, T * D, 32) # (B, max_obs, d_model) - - def test_forward_without_mask(self): - """Without mask, all values treated as observed.""" - config = ObservationTransformerConfig( - d_input=5, d_model=16, n_layers=1, n_heads=2, d_ff=32, pooling="mean" - ) - encoder = ObservationTransformerEncoder(config) - - x = torch.randn(2, 4, 5) - out = encoder(x) - assert out.shape == (2, 16) - - def test_encoder_output_dim(self): - config = ObservationTransformerConfig(d_input=10, d_model=64) - encoder = ObservationTransformerEncoder(config) - assert encoder.get_output_dim() == 64 - - if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/test_metrics.py b/tests/test_metrics.py index 468d43a..335fdd2 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -8,8 +8,11 @@ - Error handling """ +import numpy as np import pytest import torch +from torchmetrics import MetricCollection + from slices.eval.metrics import ( AVAILABLE_METRICS, DEFAULT_METRICS, @@ -18,7 +21,6 @@ build_metrics, get_default_metrics, ) -from torchmetrics import MetricCollection class TestMetricConfig: @@ -341,11 +343,39 @@ def test_default_metrics_structure(self): for metric in defaults: assert metric in available - def test_binary_defaults_minimal(self): - """Test binary defaults are minimal.""" + def test_binary_defaults_match_public_export_surface(self): + """Test binary defaults include threshold and calibration metrics.""" defaults = DEFAULT_METRICS["binary"] - assert len(defaults) <= 4 # Core discrimination + calibration metrics - assert "auroc" in defaults # Core clinical metric + assert defaults == [ + "auroc", + "auprc", + "accuracy", + "f1", + "precision", + "recall", + "specificity", + "brier_score", + "ece", + ] + + def test_xgboost_binary_metrics_match_publication_binary_surface(self): + """XGBoost binary metrics should match documented publication metric keys.""" + from pathlib import Path + + from scripts.training.xgboost_baseline import _binary_classification_metrics + + y_true = np.array([0, 0, 1, 1]) + y_pred_proba = np.array([0.05, 0.35, 0.65, 0.95]) + metrics = _binary_classification_metrics(y_true, y_pred_proba) + expected_keys = {f"test/{metric_name}" for metric_name in DEFAULT_METRICS["binary"]} + + assert expected_keys.issubset(metrics) + + config_text = Path("configs/eval/default.yaml").read_text() + readme_text = Path("src/slices/eval/README.md").read_text() + for metric_name in DEFAULT_METRICS["binary"]: + assert metric_name in config_text + assert metric_name in readme_text def test_all_task_types_covered(self): """Test all task types have entries.""" diff --git a/tests/test_model_common.py b/tests/test_model_common.py index 1b575fb..0436dd7 100644 --- a/tests/test_model_common.py +++ b/tests/test_model_common.py @@ -11,10 +11,13 @@ import pytest import torch import torch.nn as nn + from slices.models.common import PositionalEncoding, apply_pooling, get_activation from slices.models.encoders import ( LinearConfig, LinearEncoder, + ObservationTransformerConfig, + ObservationTransformerEncoder, TransformerConfig, TransformerEncoder, ) @@ -56,6 +59,12 @@ def test_max_pooling_with_mask(self): out = apply_pooling(x, "max", padding_mask=mask) assert torch.allclose(out, torch.tensor([[3.0, 4.0]])) + def test_max_pooling_all_empty_mask_returns_zero(self): + x = torch.tensor([[[1.0, 2.0], [3.0, 4.0]]]) + mask = torch.tensor([[False, False]]) + out = apply_pooling(x, "max", padding_mask=mask) + assert torch.allclose(out, torch.zeros(1, 2)) + def test_last_pooling_without_mask(self): x = torch.tensor([[[1.0], [2.0], [3.0]]]) out = apply_pooling(x, "last") @@ -67,6 +76,12 @@ def test_last_pooling_with_mask(self): out = apply_pooling(x, "last", padding_mask=mask) assert torch.allclose(out, torch.tensor([[2.0]])) + def test_last_pooling_all_empty_mask_returns_zero(self): + x = torch.tensor([[[99.0], [100.0]]]) + mask = torch.tensor([[False, False]]) + out = apply_pooling(x, "last", padding_mask=mask) + assert torch.allclose(out, torch.zeros(1, 1)) + def test_cls_pooling(self): x = torch.tensor([[[10.0, 20.0], [1.0, 2.0]]]) out = apply_pooling(x, "cls") @@ -87,6 +102,27 @@ def test_transformer_uses_shared_pooling(self): out = encoder(x) assert out.shape == (2, 8) + def test_observation_transformer_all_empty_row_is_finite(self): + config = ObservationTransformerConfig( + d_input=4, + d_model=8, + n_layers=1, + n_heads=2, + d_ff=16, + pooling="mean", + ) + encoder = ObservationTransformerEncoder(config) + encoder.eval() + x = torch.zeros(2, 5, 4) + mask = torch.zeros(2, 5, 4, dtype=torch.bool) + mask[1, 0, 0] = True + + with torch.no_grad(): + out = encoder(x, mask=mask) + + assert out.shape == (2, 8) + assert torch.isfinite(out).all() + def test_linear_uses_shared_pooling(self): """Verify LinearEncoder._apply_pooling delegates to shared utility.""" config = LinearConfig(d_input=4, d_model=8, pooling="max") diff --git a/tests/test_package.py b/tests/test_package.py index 231c227..3721caa 100644 --- a/tests/test_package.py +++ b/tests/test_package.py @@ -76,15 +76,17 @@ def test_mortality_builder_inherits_base(self) -> None: def test_datamodule_inherits_lightning(self) -> None: """Test that ICUDataModule inherits from LightningDataModule.""" import lightning.pytorch as L + from slices.data.datamodule import ICUDataModule assert issubclass(ICUDataModule, L.LightningDataModule) def test_dataset_inherits_torch_dataset(self) -> None: """Test that ICUDataset inherits from torch Dataset.""" - from slices.data.dataset import ICUDataset from torch.utils.data import Dataset + from slices.data.dataset import ICUDataset + assert issubclass(ICUDataset, Dataset) diff --git a/tests/test_pretrain_module.py b/tests/test_pretrain_module.py index 4dd9482..3e28c5f 100644 --- a/tests/test_pretrain_module.py +++ b/tests/test_pretrain_module.py @@ -12,6 +12,7 @@ import pytest import torch from omegaconf import OmegaConf + from slices.training.pretrain_module import SSLPretrainModule @@ -21,7 +22,7 @@ def minimal_config(): return OmegaConf.create( { "encoder": { - "name": "observation_transformer", + "name": "transformer", "d_input": 9, "d_model": 32, "n_layers": 1, @@ -33,6 +34,8 @@ def minimal_config(): "prenorm": True, "activation": "gelu", "layer_norm_eps": 1e-5, + "obs_aware": True, + "use_positional_encoding": True, }, "ssl": { "name": "mae", @@ -65,10 +68,10 @@ def test_module_initialization(self, minimal_config): def test_encoder_type(self, minimal_config): """Test that encoder is correct type.""" - from slices.models.encoders import ObservationTransformerEncoder + from slices.models.encoders import TransformerEncoder module = SSLPretrainModule(minimal_config) - assert isinstance(module.encoder, ObservationTransformerEncoder) + assert isinstance(module.encoder, TransformerEncoder) def test_ssl_objective_type(self, minimal_config): """Test that SSL objective is correct type.""" @@ -115,7 +118,7 @@ def test_forward_with_missing_values(self, minimal_config): loss, metrics = module(timeseries, mask) assert torch.isfinite(loss) - assert "mae_obs_ratio" in metrics + assert "mae_n_loss_positions" in metrics class TestSSLPretrainModuleTraining: @@ -311,7 +314,7 @@ def test_warmup_starts_nonzero(self): config = OmegaConf.create( { "encoder": { - "name": "observation_transformer", + "name": "transformer", "d_input": 9, "d_model": 32, "n_layers": 1, @@ -323,6 +326,8 @@ def test_warmup_starts_nonzero(self): "prenorm": True, "activation": "gelu", "layer_norm_eps": 1e-5, + "obs_aware": True, + "use_positional_encoding": True, }, "ssl": { "name": "mae", @@ -370,7 +375,7 @@ def test_warmup_schedule_values(self): config = OmegaConf.create( { "encoder": { - "name": "observation_transformer", + "name": "transformer", "d_input": 9, "d_model": 32, "n_layers": 1, @@ -382,6 +387,8 @@ def test_warmup_schedule_values(self): "prenorm": True, "activation": "gelu", "layer_norm_eps": 1e-5, + "obs_aware": True, + "use_positional_encoding": True, }, "ssl": { "name": "mae", @@ -440,7 +447,7 @@ def test_warmup_reaches_full_lr(self): config = OmegaConf.create( { "encoder": { - "name": "observation_transformer", + "name": "transformer", "d_input": 9, "d_model": 32, "n_layers": 1, @@ -452,6 +459,8 @@ def test_warmup_reaches_full_lr(self): "prenorm": True, "activation": "gelu", "layer_norm_eps": 1e-5, + "obs_aware": True, + "use_positional_encoding": True, }, "ssl": { "name": "mae", @@ -511,6 +520,14 @@ def test_typo_in_scheduler_config_rejected(self, minimal_config): with pytest.raises(ValidationError, match="t_max"): SSLPretrainModule(OmegaConf.create(cfg)) + def test_typo_in_training_config_rejected(self, minimal_config): + from pydantic import ValidationError + + cfg = OmegaConf.to_container(minimal_config, resolve=True) + cfg["training"] = {"max_epoch": 100} + with pytest.raises(ValidationError, match="max_epoch"): + SSLPretrainModule(OmegaConf.create(cfg)) + def test_no_scheduler_passes(self, minimal_config): """Module without scheduler should still pass validation.""" module = SSLPretrainModule(minimal_config) diff --git a/tests/test_regression_metrics.py b/tests/test_regression_metrics.py index 744ff50..421424a 100644 --- a/tests/test_regression_metrics.py +++ b/tests/test_regression_metrics.py @@ -12,6 +12,9 @@ import pytest import torch +from torchmetrics import MetricCollection +from torchmetrics.regression import MeanAbsoluteError, MeanSquaredError, R2Score + from slices.eval.metrics import ( AVAILABLE_METRICS, DEFAULT_METRICS, @@ -19,8 +22,6 @@ _build_metric, build_metrics, ) -from torchmetrics import MetricCollection -from torchmetrics.regression import MeanAbsoluteError, MeanSquaredError, R2Score class TestRegressionMetricConfig: diff --git a/tests/test_ricu_extractor.py b/tests/test_ricu_extractor.py index bcda728..3e461b2 100644 --- a/tests/test_ricu_extractor.py +++ b/tests/test_ricu_extractor.py @@ -7,6 +7,7 @@ import polars as pl import pytest import yaml + from slices.data.extractors.base import ExtractorConfig from slices.data.extractors.ricu import RicuExtractor @@ -27,6 +28,7 @@ def ricu_output_dir(tmp_path: Path) -> Path: "feature_names": ["hr", "sbp", "crea"], "n_features": 3, "seq_length_hours": 6, + "raw_export_horizon_hours": 6, "n_stays": 3, "ricu_version": "0.7.0", } @@ -110,6 +112,88 @@ def ricu_output_dir(tmp_path: Path) -> Path: return ricu_dir +def _create_aki_benchmark_ricu_output(tmp_path: Path) -> Path: + """Create mock RICU output for the 24h-input / 48h-raw AKI benchmark.""" + ricu_dir = tmp_path / "ricu_output_aki" + ricu_dir.mkdir() + + metadata = { + "dataset": "miiv", + "feature_names": ["crea"], + "n_features": 1, + "seq_length_hours": 48, + "raw_export_horizon_hours": 48, + "n_stays": 3, + "ricu_version": "0.7.0", + } + with open(ricu_dir / "ricu_metadata.yaml", "w") as f: + yaml.dump(metadata, f) + + stays = pl.DataFrame( + { + "stay_id": [100, 200, 300], + "patient_id": [10, 20, 30], + "hadm_id": [1000, 2000, 3000], + "intime": [ + datetime(2020, 1, 1, 8, 0), + datetime(2020, 1, 2, 8, 0), + datetime(2020, 1, 3, 8, 0), + ], + "outtime": [ + datetime(2020, 1, 4, 12, 0), + datetime(2020, 1, 5, 12, 0), + datetime(2020, 1, 6, 12, 0), + ], + "los_days": [3.2, 3.2, 3.2], + "age": [65.0, 70.0, 58.0], + "gender": ["M", "F", "M"], + "race": ["WHITE", "BLACK", "ASIAN"], + "admission_type": ["EMERGENCY", "EMERGENCY", "EMERGENCY"], + "insurance": ["Medicare", "Private", "Medicaid"], + "first_careunit": ["MICU", "MICU", "MICU"], + "height": [175.0, 165.0, 180.0], + "weight": [80.0, 68.0, 92.0], + } + ) + stays.write_parquet(ricu_dir / "ricu_stays.parquet") + + timeseries = pl.DataFrame( + { + "stay_id": [100, 100, 100, 100, 200, 200, 200, 200, 300, 300, 300], + "hour": [0, 12, 23, 30, 0, 12, 23, 30, 0, 12, 23], + "crea": [1.0, 0.9, 1.0, 1.6, 1.0, 1.1, 1.0, 1.1, 0.8, 0.9, 0.85], + "crea_mask": [True] * 11, + } + ) + timeseries.write_parquet(ricu_dir / "ricu_timeseries.parquet") + + mortality = pl.DataFrame( + { + "stay_id": [100, 200, 300], + "date_of_death": [None, None, None], + "hospital_expire_flag": [0, 0, 0], + "dischtime": [ + datetime(2020, 1, 7, 8, 0), + datetime(2020, 1, 8, 8, 0), + datetime(2020, 1, 9, 8, 0), + ], + "discharge_location": ["HOME", "HOME", "HOME"], + } + ) + mortality.write_parquet(ricu_dir / "ricu_mortality.parquet") + + diagnoses = pl.DataFrame( + { + "stay_id": [100, 200, 300], + "icd_code": ["N179", "I10", "E119"], + "icd_version": [10, 10, 10], + } + ) + diagnoses.write_parquet(ricu_dir / "ricu_diagnoses.parquet") + + return ricu_dir + + @pytest.fixture def ricu_extractor(ricu_output_dir: Path, tmp_path: Path) -> RicuExtractor: """Create a RicuExtractor instance with mock data.""" @@ -135,6 +219,8 @@ def test_init_success(self, ricu_output_dir: Path, tmp_path: Path) -> None: config = ExtractorConfig( parquet_root=str(ricu_output_dir), output_dir=str(tmp_path / "processed"), + seq_length_hours=6, + tasks=[], ) extractor = RicuExtractor(config) @@ -164,6 +250,147 @@ def test_no_duckdb_connection(self, ricu_extractor: RicuExtractor) -> None: """Verify that RicuExtractor does not have a DuckDB connection.""" assert not hasattr(ricu_extractor, "conn") + def test_init_raises_when_python_horizon_exceeds_ricu_export( + self, ricu_output_dir: Path, tmp_path: Path + ) -> None: + config = ExtractorConfig( + parquet_root=str(ricu_output_dir), + output_dir=str(tmp_path / "processed"), + seq_length_hours=8, + ) + + with pytest.raises(ValueError, match="upstream RICU export only contains 6 hours"): + RicuExtractor(config) + + def test_init_raises_when_task_requires_longer_raw_export_horizon( + self, ricu_output_dir: Path, tmp_path: Path + ) -> None: + tasks_dir = tmp_path / "tasks" + tasks_dir.mkdir() + + task_config = { + "task_name": "aki_kdigo", + "task_type": "binary", + "observation_window_hours": 6, + "prediction_window_hours": 4, + "gap_hours": 0, + "label_sources": ["stays", "timeseries"], + "label_params": {"creatinine_col": "crea"}, + "primary_metric": "auprc", + } + with open(tasks_dir / "aki_kdigo.yaml", "w") as f: + yaml.dump(task_config, f) + + config = ExtractorConfig( + parquet_root=str(ricu_output_dir), + output_dir=str(tmp_path / "processed"), + seq_length_hours=6, + min_stay_hours=0, + tasks=["aki_kdigo"], + tasks_dir=str(tasks_dir), + ) + + with pytest.raises( + ValueError, + match="Active task labels require a longer upstream raw timeseries horizon", + ): + RicuExtractor(config) + + def test_init_accepts_explicit_raw_export_horizon_longer_than_input_window( + self, ricu_output_dir: Path, tmp_path: Path + ) -> None: + with open(ricu_output_dir / "ricu_metadata.yaml") as f: + metadata = yaml.safe_load(f) + metadata["raw_export_horizon_hours"] = 10 + with open(ricu_output_dir / "ricu_metadata.yaml", "w") as f: + yaml.dump(metadata, f) + + tasks_dir = tmp_path / "tasks" + tasks_dir.mkdir() + task_config = { + "task_name": "aki_kdigo", + "task_type": "binary", + "observation_window_hours": 6, + "prediction_window_hours": 4, + "gap_hours": 0, + "label_sources": ["stays", "timeseries"], + "label_params": {"creatinine_col": "crea"}, + "primary_metric": "auprc", + } + with open(tasks_dir / "aki_kdigo.yaml", "w") as f: + yaml.dump(task_config, f) + + config = ExtractorConfig( + parquet_root=str(ricu_output_dir), + output_dir=str(tmp_path / "processed"), + seq_length_hours=6, + min_stay_hours=0, + tasks=["aki_kdigo"], + tasks_dir=str(tasks_dir), + ) + + extractor = RicuExtractor(config) + assert extractor._get_raw_export_horizon_hours() == 10 + + def test_existing_extraction_is_invalidated_when_upstream_inputs_change( + self, ricu_output_dir: Path, tmp_path: Path + ) -> None: + output_dir = tmp_path / "processed" + config = ExtractorConfig( + parquet_root=str(ricu_output_dir), + output_dir=str(output_dir), + seq_length_hours=6, + min_stay_hours=0, + tasks=[], + ) + + RicuExtractor(config).run() + + resumed = RicuExtractor(config) + assert resumed._check_existing_extraction() is not None + + timeseries_path = ricu_output_dir / "ricu_timeseries.parquet" + timeseries = pl.read_parquet(timeseries_path) + updated = pl.concat( + [ + timeseries, + pl.DataFrame( + { + "stay_id": [999], + "hour": [0], + "hr": [88.0], + "sbp": [120.0], + "crea": [1.2], + "hr_mask": [True], + "sbp_mask": [True], + "crea_mask": [True], + } + ), + ], + how="vertical_relaxed", + ) + updated.write_parquet(timeseries_path) + + after_change = RicuExtractor(config) + assert after_change._check_existing_extraction() is None + + def test_upstream_signature_includes_feature_blocklist( + self, ricu_output_dir: Path, tmp_path: Path + ) -> None: + config = ExtractorConfig( + parquet_root=str(ricu_output_dir), + output_dir=str(tmp_path / "processed"), + seq_length_hours=6, + min_stay_hours=0, + tasks=[], + ) + + signature = RicuExtractor(config)._get_upstream_source_signature() + + assert "feature_blocklist" in signature + assert "norepi_dur" in signature["feature_blocklist"] + assert "qsofa" in signature["feature_blocklist"] + # --------------------------------------------------------------------------- # Core extraction methods @@ -228,6 +455,32 @@ def test_hours_per_stay(self, ricu_extractor: RicuExtractor) -> None: ts = ricu_extractor.extract_timeseries([100]) assert len(ts) == 6 # 6 hours of data + def test_reads_from_partitioned_directory(self, ricu_output_dir: Path, tmp_path: Path) -> None: + """Timeseries stored as a directory of part files (chunked output).""" + ts_file = ricu_output_dir / "ricu_timeseries.parquet" + original = pl.read_parquet(ts_file) + + # Split into two part files in a directory with the same name + ts_file.unlink() + ts_file.mkdir() + part1 = original.filter(pl.col("stay_id") == 100) + part2 = original.filter(pl.col("stay_id").is_in([200, 300])) + part1.write_parquet(ts_file / "part_0001.parquet") + part2.write_parquet(ts_file / "part_0002.parquet") + + config = ExtractorConfig( + parquet_root=str(ricu_output_dir), + output_dir=str(tmp_path / "processed_partitioned"), + seq_length_hours=6, + tasks=[], + ) + extractor = RicuExtractor(config) + ts = extractor.extract_timeseries([100, 200, 300]) + + assert set(ts["stay_id"].unique().to_list()) == {100, 200, 300} + assert len(ts) == len(original) + assert set(ts.columns) == set(original.columns) + class TestExtractDataSource: """Tests for extract_data_source().""" @@ -265,6 +518,16 @@ def test_filter_by_stay_ids(self, ricu_extractor: RicuExtractor) -> None: assert len(mort) == 1 assert mort["stay_id"][0] == 200 + def test_legacy_mortality_datetime_is_migrated_conservatively( + self, ricu_extractor: RicuExtractor + ) -> None: + mort = ricu_extractor.extract_data_source("mortality_info", [300]) + row = mort.filter(pl.col("stay_id") == 300) + + assert row["death_time_precision"][0] == "date" + assert row["death_time"][0] is None + assert row["death_date"][0].isoformat() == "2020-01-08" + class TestExtractRawEvents: """Tests for _extract_raw_events() — should raise NotImplementedError.""" @@ -396,9 +659,72 @@ def test_run_metadata_content(self, ricu_output_dir: Path, tmp_path: Path) -> No assert meta["dataset"] == "miiv" assert meta["n_features"] == 3 assert meta["seq_length_hours"] == 6 + assert meta["input_seq_length_hours"] == 6 + assert meta["raw_export_horizon_hours"] == 6 + assert meta["required_raw_export_horizon_hours"] == 6 assert meta["n_stays"] == 3 assert meta["feature_names"] == ["hr", "sbp", "crea"] + assert meta["label_quality_stats"] == {} assert "ricu_metadata" in meta + assert "upstream_source_signature" in meta + assert len(meta["upstream_source_signature"]["files"]) >= 4 + + def test_run_excludes_derived_blocklisted_features( + self, ricu_output_dir: Path, tmp_path: Path + ) -> None: + """Derived RICU concepts should not enter processed model inputs.""" + metadata_path = ricu_output_dir / "ricu_metadata.yaml" + with open(metadata_path) as f: + metadata = yaml.safe_load(f) + metadata["feature_names"] = [ + "hr", + "norepi_dur", + "norepi60", + "sofa", + "pafi", + "qsofa", + "sbp", + "crea", + ] + metadata["n_features"] = len(metadata["feature_names"]) + with open(metadata_path, "w") as f: + yaml.safe_dump(metadata, f) + + timeseries_path = ricu_output_dir / "ricu_timeseries.parquet" + timeseries = pl.read_parquet(timeseries_path).with_columns( + pl.lit(30.0).alias("norepi_dur"), + pl.lit(0.2).alias("norepi60"), + pl.lit(4.0).alias("sofa"), + pl.lit(250.0).alias("pafi"), + pl.lit(2.0).alias("qsofa"), + pl.lit(True).alias("norepi_dur_mask"), + pl.lit(True).alias("norepi60_mask"), + pl.lit(True).alias("sofa_mask"), + pl.lit(True).alias("pafi_mask"), + pl.lit(True).alias("qsofa_mask"), + ) + timeseries.write_parquet(timeseries_path) + + output_dir = tmp_path / "processed" + config = ExtractorConfig( + parquet_root=str(ricu_output_dir), + output_dir=str(output_dir), + seq_length_hours=6, + min_stay_hours=0, + tasks=[], + ) + RicuExtractor(config).run() + + with open(output_dir / "metadata.yaml") as f: + meta = yaml.safe_load(f) + assert meta["feature_names"] == ["hr", "sbp", "crea"] + assert meta["n_features"] == 3 + + ts = pl.read_parquet(output_dir / "timeseries.parquet") + first_ts = ts["timeseries"].to_list()[0] + first_mask = ts["mask"].to_list()[0] + assert len(first_ts[0]) == 3 + assert len(first_mask[0]) == 3 def test_run_filters_short_stays(self, ricu_output_dir: Path, tmp_path: Path) -> None: output_dir = tmp_path / "processed" @@ -477,6 +803,70 @@ def test_run_with_task_labels(self, ricu_output_dir: Path, tmp_path: Path) -> No assert "mortality_hospital" in labels.columns assert len(labels) == 3 + def test_run_aki_24h_input_with_48h_raw_export(self, tmp_path: Path) -> None: + """AKI extraction should use 48h raw creatinine while saving 24h model inputs.""" + ricu_output_dir = _create_aki_benchmark_ricu_output(tmp_path) + output_dir = tmp_path / "processed_aki" + tasks_dir = tmp_path / "tasks_aki" + tasks_dir.mkdir() + + task_config = { + "task_name": "aki_kdigo", + "task_type": "binary", + "observation_window_hours": 24, + "prediction_window_hours": 24, + "gap_hours": 0, + "label_sources": ["stays", "timeseries"], + "label_params": { + "creatinine_col": "crea", + "baseline_window_hours": 24, + "absolute_rise_threshold": 0.3, + "relative_rise_threshold": 1.5, + "relative_window_hours": 168, + }, + "quality_checks": {"max_missing_percentage": 28.0}, + "primary_metric": "auprc", + } + with open(tasks_dir / "aki_kdigo.yaml", "w") as f: + yaml.dump(task_config, f) + + config = ExtractorConfig( + parquet_root=str(ricu_output_dir), + output_dir=str(output_dir), + seq_length_hours=24, + min_stay_hours=24, + tasks=["aki_kdigo"], + tasks_dir=str(tasks_dir), + ) + RicuExtractor(config).run() + + labels = pl.read_parquet(output_dir / "labels.parquet") + labels_dict = dict(zip(labels["stay_id"], labels["aki_kdigo"])) + assert labels_dict[100] == 1 + assert labels_dict[200] == 0 + assert labels_dict[300] is None + + dense_timeseries = pl.read_parquet(output_dir / "timeseries.parquet") + stay100_ts = dense_timeseries.filter(pl.col("stay_id") == 100)["timeseries"].to_list()[0] + assert len(stay100_ts) == 24 # model input remains 24h + + with open(output_dir / "metadata.yaml") as f: + meta = yaml.safe_load(f) + + assert meta["input_seq_length_hours"] == 24 + assert meta["raw_export_horizon_hours"] == 48 + assert meta["required_raw_export_horizon_hours"] == 48 + + quality_stats = meta["label_quality_stats"]["aki_kdigo"] + assert quality_stats["total_stays"] == 3 + assert quality_stats["positive_labels"] == 1 + assert quality_stats["negative_labels"] == 1 + assert quality_stats["null_labels"] == 1 + assert quality_stats["null_reason_counts"] == { + "no_creatinine_or_baseline": 0, + "no_post_obs_creatinine": 1, + } + def test_run_resume_skips_existing(self, ricu_output_dir: Path, tmp_path: Path) -> None: """Test that a second run() resumes without duplicating stays.""" output_dir = tmp_path / "processed" @@ -499,3 +889,220 @@ def test_run_resume_skips_existing(self, ricu_output_dir: Path, tmp_path: Path) # Should have same number of stays (no duplicates) assert len(static_second) == n_first + + def test_run_rebuilds_when_task_config_changes( + self, ricu_output_dir: Path, tmp_path: Path + ) -> None: + """Task-config drift should invalidate resume instead of silently skipping work.""" + output_dir = tmp_path / "processed" + tasks_dir = tmp_path / "tasks" + tasks_dir.mkdir() + + task_config = { + "task_name": "mortality_24h", + "task_type": "binary", + "observation_window_hours": 6, + "prediction_window_hours": 24, + "gap_hours": 0, + "label_sources": ["stays", "mortality_info"], + "label_params": {}, + } + task_path = tasks_dir / "mortality_24h.yaml" + with open(task_path, "w") as f: + yaml.dump(task_config, f) + + config = ExtractorConfig( + parquet_root=str(ricu_output_dir), + output_dir=str(output_dir), + seq_length_hours=6, + min_stay_hours=0, + tasks=["mortality_24h"], + tasks_dir=str(tasks_dir), + ) + + RicuExtractor(config).run() + with open(output_dir / "metadata.yaml") as f: + initial_metadata = yaml.safe_load(f) + + task_config["gap_hours"] = 6 + with open(task_path, "w") as f: + yaml.dump(task_config, f) + + RicuExtractor(config).run() + with open(output_dir / "metadata.yaml") as f: + updated_metadata = yaml.safe_load(f) + + assert ( + updated_metadata["label_manifest"]["mortality_24h"]["config_hash"] + != initial_metadata["label_manifest"]["mortality_24h"]["config_hash"] + ) + + def test_run_rebuilds_when_builder_version_changes( + self, ricu_output_dir: Path, tmp_path: Path, monkeypatch + ) -> None: + """Builder-version drift should invalidate resume instead of merging stale labels.""" + from slices.data.labels.mortality import MortalityLabelBuilder + + output_dir = tmp_path / "processed" + tasks_dir = tmp_path / "tasks" + tasks_dir.mkdir() + + task_config = { + "task_name": "mortality_hospital", + "task_type": "binary", + "observation_window_hours": 6, + "prediction_window_hours": None, + "gap_hours": 0, + "label_sources": ["stays", "mortality_info"], + "label_params": {}, + } + with open(tasks_dir / "mortality_hospital.yaml", "w") as f: + yaml.dump(task_config, f) + + config = ExtractorConfig( + parquet_root=str(ricu_output_dir), + output_dir=str(output_dir), + seq_length_hours=6, + min_stay_hours=0, + tasks=["mortality_hospital"], + tasks_dir=str(tasks_dir), + ) + + RicuExtractor(config).run() + with open(output_dir / "metadata.yaml") as f: + initial_metadata = yaml.safe_load(f) + + initial_version = MortalityLabelBuilder.SEMANTIC_VERSION + monkeypatch.setattr(MortalityLabelBuilder, "SEMANTIC_VERSION", "9.9.9") + + RicuExtractor(config).run() + with open(output_dir / "metadata.yaml") as f: + updated_metadata = yaml.safe_load(f) + + assert ( + initial_metadata["label_manifest"]["mortality_hospital"]["builder_version"] + == initial_version + ) + assert ( + updated_metadata["label_manifest"]["mortality_hospital"]["builder_version"] == "9.9.9" + ) + + def test_run_excludes_invalid_los_before_validation( + self, ricu_output_dir: Path, tmp_path: Path + ) -> None: + """Exclude invalid LOS rows before validating the benchmark cohort.""" + output_dir = tmp_path / "processed" + + stays_path = ricu_output_dir / "ricu_stays.parquet" + stays = pl.read_parquet(stays_path).with_columns( + pl.when(pl.col("stay_id") == 200) + .then(None) + .otherwise(pl.col("los_days")) + .alias("los_days") + ) + stays.write_parquet(stays_path) + + config = ExtractorConfig( + parquet_root=str(ricu_output_dir), + output_dir=str(output_dir), + seq_length_hours=6, + min_stay_hours=0, + tasks=[], + ) + + RicuExtractor(config).run() + + static = pl.read_parquet(output_dir / "static.parquet") + assert static["stay_id"].to_list() == [100, 300] + + def test_run_fails_closed_on_missing_timeseries_coverage( + self, ricu_output_dir: Path, tmp_path: Path + ) -> None: + """Missing timeseries rows for an extracted stay should abort the run.""" + output_dir = tmp_path / "processed" + + timeseries_path = ricu_output_dir / "ricu_timeseries.parquet" + timeseries = pl.read_parquet(timeseries_path).filter(pl.col("stay_id") != 300) + timeseries.write_parquet(timeseries_path) + + config = ExtractorConfig( + parquet_root=str(ricu_output_dir), + output_dir=str(output_dir), + seq_length_hours=6, + min_stay_hours=0, + tasks=[], + ) + + with pytest.raises(ValueError, match="no timeseries data"): + RicuExtractor(config).run() + + def test_run_fails_closed_on_missing_labels( + self, ricu_output_dir: Path, tmp_path: Path, monkeypatch + ) -> None: + """Missing label rows should abort extraction instead of producing a degraded dataset.""" + output_dir = tmp_path / "processed" + tasks_dir = tmp_path / "tasks" + tasks_dir.mkdir() + + task_config = { + "task_name": "mortality_hospital", + "task_type": "binary", + "observation_window_hours": 6, + "prediction_window_hours": None, + "gap_hours": 0, + "label_sources": ["stays", "mortality_info"], + "label_params": {}, + } + with open(tasks_dir / "mortality_hospital.yaml", "w") as f: + yaml.dump(task_config, f) + + config = ExtractorConfig( + parquet_root=str(ricu_output_dir), + output_dir=str(output_dir), + seq_length_hours=6, + min_stay_hours=0, + tasks=["mortality_hospital"], + tasks_dir=str(tasks_dir), + ) + + extractor = RicuExtractor(config) + original_extract_labels = extractor.extract_labels + + def drop_one_label(stay_ids, task_configs): + labels = original_extract_labels(stay_ids, task_configs) + return labels.filter(pl.col("stay_id") != 300) + + monkeypatch.setattr(extractor, "extract_labels", drop_one_label) + + with pytest.raises(ValueError, match="no labels"): + extractor.run() + + def test_run_rebuilds_when_upstream_timeseries_changes( + self, ricu_output_dir: Path, tmp_path: Path + ) -> None: + """Changing upstream parquet inputs should invalidate resume and rebuild outputs.""" + output_dir = tmp_path / "processed" + config = ExtractorConfig( + parquet_root=str(ricu_output_dir), + output_dir=str(output_dir), + seq_length_hours=6, + min_stay_hours=0, + tasks=[], + ) + + RicuExtractor(config).run() + + ts_path = ricu_output_dir / "ricu_timeseries.parquet" + updated = pl.read_parquet(ts_path).with_columns( + pl.when((pl.col("stay_id") == 100) & (pl.col("hour") == 0)) + .then(pl.lit(999.0)) + .otherwise(pl.col("hr")) + .alias("hr") + ) + updated.write_parquet(ts_path) + + RicuExtractor(config).run() + + timeseries = pl.read_parquet(output_dir / "timeseries.parquet") + stay100 = timeseries.filter(pl.col("stay_id") == 100) + assert stay100["timeseries"].to_list()[0][0][0] == 999.0 diff --git a/tests/test_sampling.py b/tests/test_sampling.py index 6d54e23..b73af85 100644 --- a/tests/test_sampling.py +++ b/tests/test_sampling.py @@ -6,6 +6,7 @@ import numpy as np import polars as pl import pytest + from slices.debug.sampling import ( SelectionStrategy, SentinelConfig, diff --git a/tests/test_sliding_window.py b/tests/test_sliding_window.py index 52dcd2d..878a415 100644 --- a/tests/test_sliding_window.py +++ b/tests/test_sliding_window.py @@ -5,6 +5,7 @@ import pytest import torch import yaml + from slices.data.datamodule import ICUDataModule, icu_collate_fn from slices.data.dataset import ICUDataset from slices.data.sliding_window import SlidingWindowDataset diff --git a/tests/test_smart_encoder.py b/tests/test_smart_encoder.py index 6779f5b..d28e2ce 100644 --- a/tests/test_smart_encoder.py +++ b/tests/test_smart_encoder.py @@ -13,6 +13,7 @@ import pytest import torch import torch.nn as nn + from slices.models.encoders.smart import ( BasicBlock, MLPBlock, @@ -203,6 +204,34 @@ def test_seq_attention_additive_mask_bias(self): # Outputs should differ due to different attention biases assert not torch.allclose(out_all, out_half, atol=1e-5) + def test_seq_attention_missing_pairs_use_finite_graded_bias(self, monkeypatch): + """SMART appendix attention should not hard-block missing/missing pairs.""" + import slices.models.encoders.smart as smart_module + + captured = {} + + def fake_scaled_dot_product_attention(q, k, v, attn_mask=None): + captured["attn_mask"] = attn_mask.detach().clone() + return torch.zeros_like(q) + + monkeypatch.setattr( + smart_module.F, + "scaled_dot_product_attention", + fake_scaled_dot_product_attention, + ) + + block = SeqAttentionBlock(d_model=4, n_heads=1, dropout=0.0) + x = torch.randn(1, 1, 3, 4) + obs_mask = torch.zeros(1, 1, 2, dtype=torch.bool) + + block(x, obs_mask) + + attn_mask = captured["attn_mask"] + assert attn_mask.shape == (1, 1, 3, 3) + assert torch.isfinite(attn_mask).all() + assert attn_mask[0, 0, 1, 2].item() == 0.0 + assert attn_mask[0, 0, 0, 1].item() == 1.0 + class TestVarAttentionBlock: """Tests for VarAttentionBlock (cross-variable attention).""" @@ -337,7 +366,7 @@ def test_smart_encoder_init_default(self): assert hasattr(encoder, "query") assert hasattr(encoder, "pos_encoder") assert hasattr(encoder, "blocks") - assert hasattr(encoder, "final_norm") + # No final_norm — matches original SMART which has no post-block LayerNorm def test_smart_encoder_query_token_shape(self): """Test that query tokens have correct shape.""" diff --git a/tests/test_smart_objective.py b/tests/test_smart_objective.py index d32aa82..3a75f2e 100644 --- a/tests/test_smart_objective.py +++ b/tests/test_smart_objective.py @@ -12,6 +12,7 @@ import pytest import torch + from slices.models.encoders.smart import SMARTEncoder, SMARTEncoderConfig from slices.models.encoders.transformer import TransformerConfig, TransformerEncoder from slices.models.pretraining.factory import ( @@ -131,6 +132,12 @@ def test_target_encoder_is_frozen(self, smart_encoder, smart_config): for param in smart.target_encoder.parameters(): assert not param.requires_grad + def test_target_encoder_defaults_to_eval_mode(self, smart_encoder, smart_config): + """The EMA teacher should start in eval mode to disable dropout noise.""" + smart = SMARTObjective(smart_encoder, smart_config) + + assert not smart.target_encoder.training + def test_target_encoder_is_copy(self, smart_encoder, smart_config): """Test that target encoder is a deep copy of online encoder.""" smart = SMARTObjective(smart_encoder, smart_config) @@ -148,6 +155,16 @@ def test_online_encoder_is_trainable(self, smart_encoder, smart_config): for param in smart.encoder.parameters(): assert param.requires_grad + def test_target_encoder_stays_in_eval_mode_when_training(self, smart_encoder, smart_config): + """Calling train() on the objective must not re-enable teacher dropout.""" + smart = SMARTObjective(smart_encoder, smart_config) + + smart.train() + + assert smart.training + assert smart.encoder.training + assert not smart.target_encoder.training + class TestSMARTObjectiveForward: """Tests for SMART objective forward pass.""" diff --git a/tests/test_snapshots.py b/tests/test_snapshots.py new file mode 100644 index 0000000..a1f9c78 --- /dev/null +++ b/tests/test_snapshots.py @@ -0,0 +1,42 @@ +from pathlib import Path + +import polars as pl + +from slices.debug.snapshots import ( + LegacyPipelineStage, + capture_dense_snapshot, + capture_labels_snapshot, + capture_stays_snapshot, + export_snapshot, +) + + +def test_legacy_snapshot_helpers_use_valid_stages(tmp_path: Path) -> None: + stays = pl.DataFrame({"stay_id": [1, 2], "subject_id": [10, 20]}) + labels = pl.DataFrame({"stay_id": [1, 2], "mortality_24h": [0, 1]}) + dense = pl.DataFrame( + { + "stay_id": [1], + "timeseries": [[[1.0, 2.0], [3.0, 4.0]]], + "mask": [[[True, False], [True, True]]], + } + ) + + stays_snapshot = capture_stays_snapshot(stays, stay_ids=[1]) + labels_snapshot = capture_labels_snapshot(labels, stay_ids=[1]) + dense_snapshot = capture_dense_snapshot( + dense, + feature_names=["heart_rate", "spo2"], + stay_ids=[1], + flatten=True, + ) + + assert stays_snapshot.stage == LegacyPipelineStage.STAYS + assert labels_snapshot.stage == LegacyPipelineStage.LABELS + assert dense_snapshot.stage == LegacyPipelineStage.DENSE_TIMESERIES + assert dense_snapshot.data.height == 4 + + exported_path = export_snapshot(stays_snapshot, tmp_path) + assert exported_path == tmp_path / "stays.csv" + assert exported_path.exists() + assert (tmp_path / "stays_metadata.yaml").exists() diff --git a/tests/test_statistical.py b/tests/test_statistical.py index 10a8d92..f27983e 100644 --- a/tests/test_statistical.py +++ b/tests/test_statistical.py @@ -4,11 +4,21 @@ - bootstrap_ci: deterministic metric yields zero-width CI, known value coverage - paired_bootstrap_test: identical models produce p ~0.5, better model produces p ~0.0, edge cases (single sample, all-same predictions, NaN handling) +- paired_wilcoxon_signed_rank / Bonferroni / Cohen's d helpers """ +import pandas as pd import pytest import torch -from slices.eval.statistical import bootstrap_ci, paired_bootstrap_test +from torchmetrics import AUROC + +from slices.eval.statistical import ( + bonferroni_correction, + bootstrap_ci, + cohens_d, + paired_bootstrap_test, + paired_wilcoxon_signed_rank, +) # ── Helpers ────────────────────────────────────────────────────────────────── @@ -87,6 +97,18 @@ def test_reproducibility_with_seed(self): r2 = bootstrap_ci(accuracy_metric, preds, targets, seed=123) assert r1 == r2 + def test_auroc_ci_skips_single_class_resamples(self): + """Single-class AUROC bootstrap samples should not become artificial zeros.""" + preds = torch.tensor([0.9] + [0.1] * 9) + targets = torch.tensor([1] + [0] * 9) + + result = bootstrap_ci(AUROC(task="binary"), preds, targets, n_bootstraps=200, seed=0) + + assert result["point"] == pytest.approx(1.0) + assert result["ci_lower"] == pytest.approx(1.0) + assert result["ci_upper"] == pytest.approx(1.0) + assert result["std"] == pytest.approx(0.0) + # ── paired_bootstrap_test ──────────────────────────────────────────────────── @@ -206,3 +228,56 @@ def test_reproducibility_with_seed(self): r1 = paired_bootstrap_test(accuracy_metric, preds_a, preds_b, targets, seed=99) r2 = paired_bootstrap_test(accuracy_metric, preds_a, preds_b, targets, seed=99) assert r1 == r2 + + def test_auroc_paired_bootstrap_skips_single_class_resamples(self): + preds = torch.tensor([0.9] + [0.1] * 9) + targets = torch.tensor([1] + [0] * 9) + + result = paired_bootstrap_test( + AUROC(task="binary"), + preds, + preds, + targets, + n_bootstraps=200, + seed=0, + ) + + assert result["score_a"] == pytest.approx(1.0) + assert result["score_b"] == pytest.approx(1.0) + assert result["delta"] == pytest.approx(0.0) + assert result["p_value"] == pytest.approx(1.0) + + +class TestWilcoxonHelpers: + def test_paired_wilcoxon_detects_consistent_improvement(self): + """Consistent paired gains should yield a small two-sided p-value.""" + values_a = [0.70, 0.72, 0.68, 0.75, 0.73, 0.71] + values_b = [0.60, 0.62, 0.58, 0.65, 0.63, 0.61] + + result = paired_wilcoxon_signed_rank(values_a, values_b) + + assert result["n_pairs"] == pytest.approx(6.0) + assert result["n_nonzero_pairs"] == pytest.approx(6.0) + assert result["p_value"] < 0.05 + + def test_paired_wilcoxon_uses_exact_small_sample_with_tied_differences(self): + result = paired_wilcoxon_signed_rank([1, 1, 1, 1, 1], [0, 0, 0, 0, 0]) + + assert result["n_pairs"] == pytest.approx(5.0) + assert result["n_nonzero_pairs"] == pytest.approx(5.0) + assert result["p_value"] == pytest.approx(0.0625) + + def test_paired_wilcoxon_handles_all_zero_differences(self): + result = paired_wilcoxon_signed_rank([1.0, 2.0, 3.0], [1.0, 2.0, 3.0]) + assert result["p_value"] == pytest.approx(1.0) + assert result["n_nonzero_pairs"] == pytest.approx(0.0) + + def test_bonferroni_preserves_nans(self): + corrected = bonferroni_correction([0.01, 0.02, float("nan")]) + assert corrected[0] == pytest.approx(0.02) + assert corrected[1] == pytest.approx(0.04) + assert pd.isna(corrected[2]) + + def test_cohens_d_paired_uses_difference_scale(self): + effect = cohens_d([0.8, 0.9, 1.0, 1.1], [0.6, 0.7, 0.8, 0.9], paired=True) + assert effect > 0 diff --git a/tests/test_task_builders.py b/tests/test_task_builders.py index 4f3b1cd..814511d 100644 --- a/tests/test_task_builders.py +++ b/tests/test_task_builders.py @@ -3,14 +3,37 @@ Comprehensive tests for mortality tasks and task builder infrastructure. """ -from datetime import datetime +from datetime import date, datetime import polars as pl import pytest + from slices.data.labels import LabelBuilderFactory, LabelConfig from slices.data.labels.mortality import MortalityLabelBuilder +def _as_timestamp_precision_mortality(df: pl.DataFrame) -> pl.DataFrame: + """Convert legacy-style test fixtures to the precision-aware mortality schema.""" + if "death_time_precision" in df.columns or "date_of_death" not in df.columns: + return df + + return df.with_columns( + pl.col("date_of_death").cast(pl.Datetime("us")).alias("death_time"), + pl.when(pl.col("date_of_death").is_not_null()) + .then(pl.col("date_of_death").cast(pl.Date)) + .otherwise(pl.lit(None).cast(pl.Date)) + .alias("death_date"), + pl.when(pl.col("date_of_death").is_not_null()) + .then(pl.lit("timestamp")) + .otherwise(pl.lit(None)) + .alias("death_time_precision"), + pl.when(pl.col("date_of_death").is_not_null()) + .then(pl.lit("test_fixture")) + .otherwise(pl.lit(None)) + .alias("death_source"), + ) + + class TestLabelConfig: """Tests for LabelConfig dataclass.""" @@ -92,24 +115,26 @@ def sample_stays(self): @pytest.fixture def sample_mortality_info(self): """Sample mortality data.""" - return pl.DataFrame( - { - "stay_id": [1, 2, 3, 4], - "date_of_death": [ - datetime(2020, 1, 1, 20, 0), # Died 10h after admission - None, # Survived - datetime(2020, 1, 5, 10, 0), # Died 50h after admission - datetime(2020, 1, 6, 14, 0), # Died 48h after admission - ], - "hospital_expire_flag": [1, 0, 1, 1], - "dischtime": [ - datetime(2020, 1, 3, 10, 0), - datetime(2020, 1, 5, 12, 0), - datetime(2020, 1, 6, 8, 0), - datetime(2020, 1, 10, 14, 0), - ], - "discharge_location": ["DIED", "HOME", "DIED", "DIED"], - } + return _as_timestamp_precision_mortality( + pl.DataFrame( + { + "stay_id": [1, 2, 3, 4], + "date_of_death": [ + datetime(2020, 1, 1, 20, 0), # Died 10h after admission + None, # Survived + datetime(2020, 1, 5, 10, 0), # Died 50h after admission + datetime(2020, 1, 6, 14, 0), # Died 48h after admission + ], + "hospital_expire_flag": [1, 0, 1, 1], + "dischtime": [ + datetime(2020, 1, 3, 10, 0), + datetime(2020, 1, 5, 12, 0), + datetime(2020, 1, 6, 8, 0), + datetime(2020, 1, 10, 14, 0), + ], + "discharge_location": ["DIED", "HOME", "DIED", "DIED"], + } + ) ) def test_hospital_mortality(self, sample_stays, sample_mortality_info): @@ -183,9 +208,8 @@ def test_hospital_mortality_with_obs_window(self, sample_stays, sample_mortality labels = builder.build_labels(raw_data) labels_dict = dict(zip(labels["stay_id"], labels["label"])) - # Stay 1: died at 10h (during 48h obs window), outtime=48h (still in ICU at obs end) - # outtime (Jan 3 10:00) >= obs_end (Jan 3 10:00), so NOT left_icu_during_obs -> label=1 - assert labels_dict[1] == 1 # hospital_expire_flag=1 + # Stay 1: died at 10h, during the observation window, so excluded. + assert labels_dict[1] is None assert labels_dict[2] == 0 # Survived assert labels_dict[3] == 1 # Died in hospital (hospital_expire_flag=1) assert labels_dict[4] == 1 # Died in hospital (hospital_expire_flag=1) @@ -193,24 +217,26 @@ def test_hospital_mortality_with_obs_window(self, sample_stays, sample_mortality def test_icu_mortality(self, sample_stays): """Test ICU mortality prediction (death during ICU stay, window_hours=-1).""" # Create mortality info where some patients died during ICU, some after - mortality_info = pl.DataFrame( - { - "stay_id": [1, 2, 3, 4], - "date_of_death": [ - datetime(2020, 1, 2, 10, 0), # Died during ICU stay - None, # Survived - datetime(2020, 1, 10, 10, 0), # Died after ICU discharge - datetime(2020, 1, 6, 8, 0), # Died at exact ICU discharge time - ], - "hospital_expire_flag": [1, 0, 1, 1], - "dischtime": [ - datetime(2020, 1, 3, 10, 0), - datetime(2020, 1, 5, 12, 0), - datetime(2020, 1, 10, 10, 0), - datetime(2020, 1, 10, 14, 0), - ], - "discharge_location": ["DIED", "HOME", "DIED", "DIED"], - } + mortality_info = _as_timestamp_precision_mortality( + pl.DataFrame( + { + "stay_id": [1, 2, 3, 4], + "date_of_death": [ + datetime(2020, 1, 2, 10, 0), # Died during ICU stay + None, # Survived + datetime(2020, 1, 10, 10, 0), # Died after ICU discharge + datetime(2020, 1, 6, 8, 0), # Died at exact ICU discharge time + ], + "hospital_expire_flag": [1, 0, 1, 1], + "dischtime": [ + datetime(2020, 1, 3, 10, 0), + datetime(2020, 1, 5, 12, 0), + datetime(2020, 1, 10, 10, 0), + datetime(2020, 1, 10, 14, 0), + ], + "discharge_location": ["DIED", "HOME", "DIED", "DIED"], + } + ) ) config = LabelConfig( @@ -293,20 +319,22 @@ def boundary_stays(self): def test_death_exactly_at_24h_boundary(self, boundary_stays): """Test death exactly at 24-hour boundary (should be included).""" - mortality_info = pl.DataFrame( - { - "stay_id": [1, 2, 3, 4, 5], - "date_of_death": [ - datetime(2020, 1, 2, 0, 0), # Exactly 24h (should be included) - datetime(2020, 1, 1, 23, 59, 59), # Just before 24h (included) - datetime(2020, 1, 2, 0, 0, 1), # Just after 24h (excluded) - None, # Survived - datetime(2020, 1, 1, 12, 0), # 12h (included) - ], - "hospital_expire_flag": [1, 1, 1, 0, 1], - "dischtime": [datetime(2020, 1, 3, 0, 0)] * 5, - "discharge_location": ["DIED", "DIED", "DIED", "HOME", "DIED"], - } + mortality_info = _as_timestamp_precision_mortality( + pl.DataFrame( + { + "stay_id": [1, 2, 3, 4, 5], + "date_of_death": [ + datetime(2020, 1, 2, 0, 0), # Exactly 24h (should be included) + datetime(2020, 1, 1, 23, 59, 59), # Just before 24h (included) + datetime(2020, 1, 2, 0, 0, 1), # Just after 24h (excluded) + None, # Survived + datetime(2020, 1, 1, 12, 0), # 12h (included) + ], + "hospital_expire_flag": [1, 1, 1, 0, 1], + "dischtime": [datetime(2020, 1, 3, 0, 0)] * 5, + "discharge_location": ["DIED", "DIED", "DIED", "HOME", "DIED"], + } + ) ) config = LabelConfig( @@ -333,18 +361,20 @@ def test_death_exactly_at_24h_boundary(self, boundary_stays): def test_legacy_time_bounded_boundary(self, boundary_stays): """Test death exactly at time boundary for legacy time-bounded mortality.""" - mortality_info = pl.DataFrame( - { - "stay_id": [1, 2, 3], - "date_of_death": [ - datetime(2020, 1, 3, 0, 0), # Exactly 48h - datetime(2020, 1, 3, 0, 0, 1), # Just after 48h - datetime(2020, 1, 2, 23, 59, 59), # Just before 48h - ], - "hospital_expire_flag": [1, 1, 1], - "dischtime": [datetime(2020, 1, 3, 0, 0)] * 3, - "discharge_location": ["DIED", "DIED", "DIED"], - } + mortality_info = _as_timestamp_precision_mortality( + pl.DataFrame( + { + "stay_id": [1, 2, 3], + "date_of_death": [ + datetime(2020, 1, 3, 0, 0), # Exactly 48h + datetime(2020, 1, 3, 0, 0, 1), # Just after 48h + datetime(2020, 1, 2, 23, 59, 59), # Just before 48h + ], + "hospital_expire_flag": [1, 1, 1], + "dischtime": [datetime(2020, 1, 3, 0, 0)] * 3, + "discharge_location": ["DIED", "DIED", "DIED"], + } + ) ) config = LabelConfig( @@ -369,14 +399,16 @@ def test_legacy_time_bounded_boundary(self, boundary_stays): def test_all_survivors(self, boundary_stays): """Test dataset where all patients survive.""" - mortality_info = pl.DataFrame( - { - "stay_id": [1, 2, 3, 4, 5], - "date_of_death": [None] * 5, - "hospital_expire_flag": [0] * 5, - "dischtime": [datetime(2020, 1, 3, 0, 0)] * 5, - "discharge_location": ["HOME"] * 5, - } + mortality_info = _as_timestamp_precision_mortality( + pl.DataFrame( + { + "stay_id": [1, 2, 3, 4, 5], + "date_of_death": [None] * 5, + "hospital_expire_flag": [0] * 5, + "dischtime": [datetime(2020, 1, 3, 0, 0)] * 5, + "discharge_location": ["HOME"] * 5, + } + ) ) config = LabelConfig( @@ -400,14 +432,16 @@ def test_all_survivors(self, boundary_stays): def test_all_deaths(self, boundary_stays): """Test dataset where all patients die.""" - mortality_info = pl.DataFrame( - { - "stay_id": [1, 2, 3, 4, 5], - "date_of_death": [datetime(2020, 1, 1, 12, 0)] * 5, - "hospital_expire_flag": [1] * 5, - "dischtime": [datetime(2020, 1, 1, 12, 0)] * 5, - "discharge_location": ["DIED"] * 5, - } + mortality_info = _as_timestamp_precision_mortality( + pl.DataFrame( + { + "stay_id": [1, 2, 3, 4, 5], + "date_of_death": [datetime(2020, 1, 1, 12, 0)] * 5, + "hospital_expire_flag": [1] * 5, + "dischtime": [datetime(2020, 1, 1, 12, 0)] * 5, + "discharge_location": ["DIED"] * 5, + } + ) ) config = LabelConfig( @@ -439,14 +473,16 @@ def test_single_stay(self): } ) - mortality_info = pl.DataFrame( - { - "stay_id": [1], - "date_of_death": [datetime(2020, 1, 1, 12, 0)], - "hospital_expire_flag": [1], - "dischtime": [datetime(2020, 1, 1, 12, 0)], - "discharge_location": ["DIED"], - } + mortality_info = _as_timestamp_precision_mortality( + pl.DataFrame( + { + "stay_id": [1], + "date_of_death": [datetime(2020, 1, 1, 12, 0)], + "hospital_expire_flag": [1], + "dischtime": [datetime(2020, 1, 1, 12, 0)], + "discharge_location": ["DIED"], + } + ) ) config = LabelConfig( @@ -478,14 +514,16 @@ def test_missing_mortality_info_for_stay(self): ) # Only mortality info for stay 1 - mortality_info = pl.DataFrame( - { - "stay_id": [1], - "date_of_death": [datetime(2020, 1, 1, 12, 0)], - "hospital_expire_flag": [1], - "dischtime": [datetime(2020, 1, 1, 12, 0)], - "discharge_location": ["DIED"], - } + mortality_info = _as_timestamp_precision_mortality( + pl.DataFrame( + { + "stay_id": [1], + "date_of_death": [datetime(2020, 1, 1, 12, 0)], + "hospital_expire_flag": [1], + "dischtime": [datetime(2020, 1, 1, 12, 0)], + "discharge_location": ["DIED"], + } + ) ) config = LabelConfig( @@ -509,6 +547,174 @@ def test_missing_mortality_info_for_stay(self): assert labels_dict[3] is None # Missing info -> null (unknown outcome) +class TestMortalityUnknownOutcomeRegressions: + """Regressions for nullable eICU hospital mortality evidence.""" + + @staticmethod + def _single_stay() -> pl.DataFrame: + base = datetime(2020, 1, 1, 0, 0) + return pl.DataFrame( + { + "stay_id": [1], + "intime": [base], + "outtime": [base.replace(day=5)], + } + ) + + @staticmethod + def _mortality_info( + *, + hospital_expire_flag: int | None, + discharge_location: str | None, + dischtime: datetime | None = None, + death_date: date | None = None, + death_source: str | None = None, + ) -> pl.DataFrame: + return pl.DataFrame( + { + "stay_id": [1], + "death_time": [None], + "death_date": [death_date], + "death_time_precision": ["date" if death_date is not None else None], + "death_source": [death_source], + "hospital_expire_flag": [hospital_expire_flag], + "dischtime": [dischtime], + "discharge_location": [discharge_location], + "date_of_death": [None], + }, + schema={ + "stay_id": pl.Int64, + "death_time": pl.Datetime("us"), + "death_date": pl.Date, + "death_time_precision": pl.Utf8, + "death_source": pl.Utf8, + "hospital_expire_flag": pl.Int32, + "dischtime": pl.Datetime("us"), + "discharge_location": pl.Utf8, + "date_of_death": pl.Datetime("us"), + }, + ) + + def test_null_hospital_flag_without_death_evidence_is_null_for_windowed_mortality(self): + """Unknown windowed mortality outcome must not become a negative label.""" + config = LabelConfig( + task_name="mortality_24h", + task_type="binary_classification", + prediction_window_hours=24, + observation_window_hours=24, + label_sources=["stays", "mortality_info"], + ) + builder = MortalityLabelBuilder(config) + + labels = builder.build_labels( + { + "stays": self._single_stay(), + "mortality_info": self._mortality_info( + hospital_expire_flag=None, + discharge_location=None, + ), + } + ) + + assert labels["label"][0] is None + + def test_mimic_post_discharge_dod_is_not_hospital_mortality(self): + """MIMIC patients.dod after discharge should not count as hospital death.""" + base = datetime(2020, 1, 1, 0, 0) + config = LabelConfig( + task_name="mortality_hospital", + task_type="binary_classification", + prediction_window_hours=None, + observation_window_hours=24, + label_sources=["stays", "mortality_info"], + ) + builder = MortalityLabelBuilder(config) + + labels = builder.build_labels( + { + "stays": self._single_stay(), + "mortality_info": self._mortality_info( + hospital_expire_flag=0, + discharge_location="HOME", + dischtime=base.replace(day=3), + death_date=date(2020, 2, 1), + death_source="patients.dod", + ), + } + ) + + assert labels["label"][0] == 0 + + def test_eicu_death_discharge_location_is_death_evidence(self): + """eICU discharge_location == Death should override a missing death flag.""" + config = LabelConfig( + task_name="mortality_hospital", + task_type="binary_classification", + prediction_window_hours=None, + label_sources=["stays", "mortality_info"], + ) + builder = MortalityLabelBuilder(config) + + labels = builder.build_labels( + { + "stays": self._single_stay(), + "mortality_info": self._mortality_info( + hospital_expire_flag=None, + discharge_location="Death", + ), + } + ) + + assert labels["label"][0] == 1 + + def test_windowed_mortality_uses_death_discharge_time_inside_prediction_window(self): + """Death-coded discharge time inside the target window should be positive.""" + base = datetime(2020, 1, 1, 0, 0) + config = LabelConfig( + task_name="mortality_24h", + task_type="binary_classification", + prediction_window_hours=24, + observation_window_hours=24, + label_sources=["stays", "mortality_info"], + ) + builder = MortalityLabelBuilder(config) + + labels = builder.build_labels( + { + "stays": self._single_stay(), + "mortality_info": self._mortality_info( + hospital_expire_flag=None, + discharge_location="Death", + dischtime=base.replace(day=2, hour=6), + ), + } + ) + + assert labels["label"][0] == 1 + + def test_hospital_mortality_unknown_outcome_is_null(self): + """Unknown hospital mortality outcome should not silently become 0.""" + config = LabelConfig( + task_name="mortality_hospital", + task_type="binary_classification", + prediction_window_hours=None, + label_sources=["stays", "mortality_info"], + ) + builder = MortalityLabelBuilder(config) + + labels = builder.build_labels( + { + "stays": self._single_stay(), + "mortality_info": self._mortality_info( + hospital_expire_flag=None, + discharge_location=None, + ), + } + ) + + assert labels["label"][0] is None + + class TestWindowedMortalityLabels: """Tests for windowed mortality prediction (observation + prediction windows). @@ -554,22 +760,32 @@ def test_windowed_mortality_basic(self): } ) - mortality_info = pl.DataFrame( - { - "stay_id": [1, 2, 3, 4, 5, 6, 7], - "date_of_death": [ - datetime(2020, 1, 2, 0, 0), # 24h - during obs - datetime(2020, 1, 3, 2, 0), # 50h - during pred - datetime(2020, 1, 4, 0, 0), # 72h - at pred boundary - datetime(2020, 1, 4, 8, 0), # 80h - after pred - None, # Survived - datetime(2020, 1, 2, 23, 0), # 47h - during obs - datetime(2020, 1, 3, 0, 0, 1), # 48h + 1s - just after obs - ], - "hospital_expire_flag": [1, 1, 1, 1, 0, 1, 1], - "dischtime": [datetime(2020, 1, 10, 0, 0)] * 7, - "discharge_location": ["DIED", "DIED", "DIED", "DIED", "HOME", "DIED", "DIED"], - } + mortality_info = _as_timestamp_precision_mortality( + pl.DataFrame( + { + "stay_id": [1, 2, 3, 4, 5, 6, 7], + "date_of_death": [ + datetime(2020, 1, 2, 0, 0), # 24h - during obs + datetime(2020, 1, 3, 2, 0), # 50h - during pred + datetime(2020, 1, 4, 0, 0), # 72h - at pred boundary + datetime(2020, 1, 4, 8, 0), # 80h - after pred + None, # Survived + datetime(2020, 1, 2, 23, 0), # 47h - during obs + datetime(2020, 1, 3, 0, 0, 1), # 48h + 1s - just after obs + ], + "hospital_expire_flag": [1, 1, 1, 1, 0, 1, 1], + "dischtime": [datetime(2020, 1, 10, 0, 0)] * 7, + "discharge_location": [ + "DIED", + "DIED", + "DIED", + "DIED", + "HOME", + "DIED", + "DIED", + ], + } + ) ) config = LabelConfig( @@ -623,20 +839,22 @@ def test_windowed_mortality_with_gap(self): } ) - mortality_info = pl.DataFrame( - { - "stay_id": [1, 2, 3, 4, 5], - "date_of_death": [ - datetime(2020, 1, 1, 20, 0), # 20h - during obs - datetime(2020, 1, 2, 3, 0), # 27h - during gap - datetime(2020, 1, 2, 8, 0), # 32h - during pred - datetime(2020, 1, 3, 10, 0), # 58h - after pred - None, # Survived - ], - "hospital_expire_flag": [1, 1, 1, 1, 0], - "dischtime": [datetime(2020, 1, 10, 0, 0)] * 5, - "discharge_location": ["DIED", "DIED", "DIED", "DIED", "HOME"], - } + mortality_info = _as_timestamp_precision_mortality( + pl.DataFrame( + { + "stay_id": [1, 2, 3, 4, 5], + "date_of_death": [ + datetime(2020, 1, 1, 20, 0), # 20h - during obs + datetime(2020, 1, 2, 3, 0), # 27h - during gap + datetime(2020, 1, 2, 8, 0), # 32h - during pred + datetime(2020, 1, 3, 10, 0), # 58h - after pred + None, # Survived + ], + "hospital_expire_flag": [1, 1, 1, 1, 0], + "dischtime": [datetime(2020, 1, 10, 0, 0)] * 5, + "discharge_location": ["DIED", "DIED", "DIED", "DIED", "HOME"], + } + ) ) config = LabelConfig( @@ -683,18 +901,20 @@ def test_windowed_mortality_boundary_at_obs_end(self): } ) - mortality_info = pl.DataFrame( - { - "stay_id": [1, 2, 3], - "date_of_death": [ - datetime(2020, 1, 3, 0, 0), # Exactly 48h - datetime(2020, 1, 2, 23, 59, 59), # 1 second before 48h - datetime(2020, 1, 3, 0, 0, 1), # 1 second after 48h - ], - "hospital_expire_flag": [1, 1, 1], - "dischtime": [datetime(2020, 1, 10, 0, 0)] * 3, - "discharge_location": ["DIED", "DIED", "DIED"], - } + mortality_info = _as_timestamp_precision_mortality( + pl.DataFrame( + { + "stay_id": [1, 2, 3], + "date_of_death": [ + datetime(2020, 1, 3, 0, 0), # Exactly 48h + datetime(2020, 1, 2, 23, 59, 59), # 1 second before 48h + datetime(2020, 1, 3, 0, 0, 1), # 1 second after 48h + ], + "hospital_expire_flag": [1, 1, 1], + "dischtime": [datetime(2020, 1, 10, 0, 0)] * 3, + "discharge_location": ["DIED", "DIED", "DIED"], + } + ) ) config = LabelConfig( @@ -734,18 +954,20 @@ def test_windowed_mortality_boundary_at_pred_end(self): } ) - mortality_info = pl.DataFrame( - { - "stay_id": [1, 2, 3], - "date_of_death": [ - datetime(2020, 1, 4, 0, 0), # Exactly 72h (48+24) - datetime(2020, 1, 3, 23, 59, 59), # 1 second before 72h - datetime(2020, 1, 4, 0, 0, 1), # 1 second after 72h - ], - "hospital_expire_flag": [1, 1, 1], - "dischtime": [datetime(2020, 1, 10, 0, 0)] * 3, - "discharge_location": ["DIED", "DIED", "DIED"], - } + mortality_info = _as_timestamp_precision_mortality( + pl.DataFrame( + { + "stay_id": [1, 2, 3], + "date_of_death": [ + datetime(2020, 1, 4, 0, 0), # Exactly 72h (48+24) + datetime(2020, 1, 3, 23, 59, 59), # 1 second before 72h + datetime(2020, 1, 4, 0, 0, 1), # 1 second after 72h + ], + "hospital_expire_flag": [1, 1, 1], + "dischtime": [datetime(2020, 1, 10, 0, 0)] * 3, + "discharge_location": ["DIED", "DIED", "DIED"], + } + ) ) config = LabelConfig( @@ -795,19 +1017,21 @@ def test_death_exactly_at_prediction_start_is_positive(self): } ) - mortality_info = pl.DataFrame( - { - "stay_id": [1, 2, 3, 4], - "date_of_death": [ - datetime(2020, 1, 3, 0, 0, 1), # 48h + 1s - just after obs ends - datetime(2020, 1, 3, 0, 0), # Exactly 48h - at prediction start - datetime(2020, 1, 3, 12, 0), # 60h - middle of pred window - datetime(2020, 1, 4, 0, 0), # 72h - at pred end - ], - "hospital_expire_flag": [1, 1, 1, 1], - "dischtime": [datetime(2020, 1, 10, 0, 0)] * 4, - "discharge_location": ["DIED", "DIED", "DIED", "DIED"], - } + mortality_info = _as_timestamp_precision_mortality( + pl.DataFrame( + { + "stay_id": [1, 2, 3, 4], + "date_of_death": [ + datetime(2020, 1, 3, 0, 0, 1), # 48h + 1s - just after obs ends + datetime(2020, 1, 3, 0, 0), # Exactly 48h - at prediction start + datetime(2020, 1, 3, 12, 0), # 60h - middle of pred window + datetime(2020, 1, 4, 0, 0), # 72h - at pred end + ], + "hospital_expire_flag": [1, 1, 1, 1], + "dischtime": [datetime(2020, 1, 10, 0, 0)] * 4, + "discharge_location": ["DIED", "DIED", "DIED", "DIED"], + } + ) ) config = LabelConfig( @@ -853,18 +1077,20 @@ def test_windowed_mortality_all_excluded(self): } ) - mortality_info = pl.DataFrame( - { - "stay_id": [1, 2, 3], - "date_of_death": [ - datetime(2020, 1, 1, 12, 0), # 12h - datetime(2020, 1, 2, 0, 0), # 24h - datetime(2020, 1, 2, 23, 59, 59), # 47h 59m 59s - ], - "hospital_expire_flag": [1, 1, 1], - "dischtime": [datetime(2020, 1, 3, 0, 0)] * 3, - "discharge_location": ["DIED", "DIED", "DIED"], - } + mortality_info = _as_timestamp_precision_mortality( + pl.DataFrame( + { + "stay_id": [1, 2, 3], + "date_of_death": [ + datetime(2020, 1, 1, 12, 0), # 12h + datetime(2020, 1, 2, 0, 0), # 24h + datetime(2020, 1, 2, 23, 59, 59), # 47h 59m 59s + ], + "hospital_expire_flag": [1, 1, 1], + "dischtime": [datetime(2020, 1, 3, 0, 0)] * 3, + "discharge_location": ["DIED", "DIED", "DIED"], + } + ) ) config = LabelConfig( @@ -910,21 +1136,23 @@ def test_windowed_mortality_until_icu_discharge(self): } ) - mortality_info = pl.DataFrame( - { - "stay_id": [1, 2, 3, 4, 5, 6], - "date_of_death": [ - datetime(2020, 1, 2, 0, 0), # 24h - during obs - datetime(2020, 1, 4, 0, 0), # 72h - during pred, at outtime - datetime(2020, 1, 5, 0, 0), # 96h - at outtime - datetime(2020, 1, 4, 0, 0), # 72h - at outtime (short stay) - None, # Survived (long stay) - datetime(2020, 1, 6, 0, 0), # 120h - after outtime (died post-discharge) - ], - "hospital_expire_flag": [1, 1, 1, 1, 0, 1], - "dischtime": [datetime(2020, 1, 10, 0, 0)] * 6, - "discharge_location": ["DIED", "DIED", "DIED", "DIED", "HOME", "DIED"], - } + mortality_info = _as_timestamp_precision_mortality( + pl.DataFrame( + { + "stay_id": [1, 2, 3, 4, 5, 6], + "date_of_death": [ + datetime(2020, 1, 2, 0, 0), # 24h - during obs + datetime(2020, 1, 4, 0, 0), # 72h - during pred, at outtime + datetime(2020, 1, 5, 0, 0), # 96h - at outtime + datetime(2020, 1, 4, 0, 0), # 72h - at outtime (short stay) + None, # Survived (long stay) + datetime(2020, 1, 6, 0, 0), # 120h - after outtime (died post-discharge) + ], + "hospital_expire_flag": [1, 1, 1, 1, 0, 1], + "dischtime": [datetime(2020, 1, 10, 0, 0)] * 6, + "discharge_location": ["DIED", "DIED", "DIED", "DIED", "HOME", "DIED"], + } + ) ) config = LabelConfig( @@ -997,3 +1225,33 @@ def test_factory_list_available(self): assert "mortality" in available assert isinstance(available["mortality"], type) + + +class TestRequiredRawTimeseriesHorizon: + """Tests for task-specific raw timeseries horizon requirements.""" + + def test_aki_requires_post_observation_raw_horizon(self): + config = LabelConfig( + task_name="aki_kdigo", + task_type="binary_classification", + observation_window_hours=24, + prediction_window_hours=24, + gap_hours=0, + label_sources=["stays", "timeseries"], + ) + + builder = LabelBuilderFactory.create(config) + assert builder.required_raw_timeseries_horizon_hours() == 48 + + def test_mortality_does_not_request_raw_timeseries_horizon(self): + config = LabelConfig( + task_name="mortality_24h", + task_type="binary_classification", + observation_window_hours=24, + prediction_window_hours=24, + gap_hours=0, + label_sources=["stays", "mortality_info"], + ) + + builder = LabelBuilderFactory.create(config) + assert builder.required_raw_timeseries_horizon_hours() == 0 diff --git a/tests/test_training_utils.py b/tests/test_training_utils.py index b94232e..c60814c 100644 --- a/tests/test_training_utils.py +++ b/tests/test_training_utils.py @@ -7,15 +7,24 @@ - Supervised script save_encoder_weights: v3 format, loadable by FineTuneModule """ +from dataclasses import fields +from pathlib import Path + import pytest import torch import torch.nn as nn +import yaml from omegaconf import OmegaConf + +from slices.data.labels import LabelBuilder, LabelConfig from slices.training import FineTuneModule from slices.training.utils import ( build_optimizer, build_scheduler, + report_and_validate_train_label_support, + run_fairness_evaluation, save_encoder_checkpoint, + setup_finetune_callbacks, ) @@ -177,6 +186,247 @@ def test_state_dict_is_loadable(self, tmp_path): assert torch.allclose(new_encoder.bias, encoder.bias) +class TestTrainLabelSupportGuard: + """Tests for label-efficiency train support reporting.""" + + class _DummyDataModule: + def __init__(self, subset_stats, full_stats=None): + self._subset_stats = subset_stats + self._full_stats = full_stats or subset_stats + + def get_train_label_statistics(self, use_full_train: bool = False): + return self._full_stats if use_full_train else self._subset_stats + + def test_reports_binary_train_support(self, capsys): + dm = self._DummyDataModule( + subset_stats={ + "mortality_24h": { + "total": 320, + "positive": 5, + "negative": 315, + "prevalence": 5 / 320, + } + }, + full_stats={ + "mortality_24h": { + "total": 32000, + "positive": 544, + "negative": 31456, + "prevalence": 544 / 32000, + } + }, + ) + + stats = report_and_validate_train_label_support( + datamodule=dm, + task_name="mortality_24h", + task_type="binary", + dataset="miiv", + seed=42, + label_fraction=0.01, + min_train_positives=3, + ) + + captured = capsys.readouterr() + assert "miiv / mortality_24h / 42 / 0.01" in captured.out + assert stats["train_subset_positive"] == 5 + assert stats["full_train_positive"] == 544 + + def test_raises_when_label_efficiency_subset_has_too_few_positives(self): + dm = self._DummyDataModule( + subset_stats={ + "mortality_24h": { + "total": 320, + "positive": 2, + "negative": 318, + "prevalence": 2 / 320, + } + } + ) + + with pytest.raises(ValueError, match="Too few positive training examples"): + report_and_validate_train_label_support( + datamodule=dm, + task_name="mortality_24h", + task_type="binary", + dataset="eicu", + seed=7, + label_fraction=0.01, + min_train_positives=3, + ) + + def test_raises_when_binary_train_subset_loses_a_class(self): + dm = self._DummyDataModule( + subset_stats={ + "mortality_24h": { + "total": 320, + "positive": 0, + "negative": 320, + "prevalence": 0.0, + } + } + ) + + with pytest.raises(ValueError, match="lost a class"): + report_and_validate_train_label_support( + datamodule=dm, + task_name="mortality_24h", + task_type="binary", + dataset="miiv", + seed=1, + label_fraction=0.01, + min_train_positives=3, + ) + + +class TestSetupFinetuneCallbacks: + """Tests task-aware checkpoint monitor selection.""" + + @staticmethod + def _make_cfg(task_cfg: dict, training_cfg: dict | None = None): + return OmegaConf.create( + { + "task": task_cfg, + "training": {"early_stopping_patience": 10, **(training_cfg or {})}, + "checkpoint_dir": "checkpoints", + } + ) + + def test_classification_defaults_to_auprc_when_no_primary_metric(self): + cfg = self._make_cfg({"task_name": "mortality_24h", "task_type": "binary"}) + + callbacks = setup_finetune_callbacks(cfg) + + assert callbacks[0].monitor == "val/auprc" + assert callbacks[0].mode == "max" + assert callbacks[1].monitor == "val/auprc" + assert callbacks[1].mode == "max" + + def test_regression_uses_task_primary_metric(self): + cfg = self._make_cfg( + {"task_name": "los_remaining", "task_type": "regression", "primary_metric": "mae"} + ) + + callbacks = setup_finetune_callbacks(cfg) + + assert callbacks[0].monitor == "val/mae" + assert callbacks[0].mode == "min" + assert callbacks[1].monitor == "val/mae" + assert callbacks[1].mode == "min" + + def test_explicit_monitor_override_still_wins(self): + cfg = self._make_cfg( + {"task_name": "los_remaining", "task_type": "regression", "primary_metric": "mae"}, + {"early_stopping_monitor": "val/r2"}, + ) + + callbacks = setup_finetune_callbacks(cfg) + + assert callbacks[0].monitor == "val/r2" + assert callbacks[0].mode == "max" + assert callbacks[1].monitor == "val/r2" + assert callbacks[1].mode == "max" + + def test_training_task_configs_cover_package_tasks(self): + package_tasks_dir = Path("src/slices/data/tasks") + hydra_tasks_dir = Path("configs/tasks") + label_config_fields = {field.name for field in fields(LabelConfig)} + + package_tasks = { + yaml.safe_load(path.read_text())["task_name"] + for path in package_tasks_dir.glob("*.yaml") + } + hydra_tasks = { + yaml.safe_load(path.read_text())["task_name"] for path in hydra_tasks_dir.glob("*.yaml") + } + + assert hydra_tasks == package_tasks + + for task_name in sorted(package_tasks): + package_config = LabelConfig( + **yaml.safe_load((package_tasks_dir / f"{task_name}.yaml").read_text()) + ) + hydra_raw = yaml.safe_load((hydra_tasks_dir / f"{task_name}.yaml").read_text()) + hydra_config = LabelConfig( + **{key: value for key, value in hydra_raw.items() if key in label_config_fields} + ) + assert LabelBuilder.config_hash(hydra_config) == LabelBuilder.config_hash( + package_config + ) + + +class TestRunFairnessEvaluation: + """Tests for inline fairness evaluator setup.""" + + def test_passes_task_type_to_fairness_evaluator(self, monkeypatch): + import slices.eval.fairness_evaluator as fairness_module + import slices.eval.inference as inference_module + + captured = {} + + class DummyEvaluator: + def __init__( + self, + static_df, + protected_attributes, + min_subgroup_size, + task_type, + dataset_name, + ): + captured["task_type"] = task_type + captured["dataset_name"] = dataset_name + + def evaluate(self, predictions, labels_tensor, all_stay_ids): + return {"gender": {"n_valid_groups": 1}} + + def print_report(self, fairness_report): + captured["printed"] = fairness_report + + class DummyModel: + device = torch.device("cpu") + + class DummyDataset: + static_df = object() + + class DummyDataModule: + dataset = DummyDataset() + processed_dir = Path("data/processed/miiv") + + def test_dataloader(self): + return [] + + monkeypatch.setattr( + inference_module, + "run_inference", + lambda *args, **kwargs: (torch.tensor([1.2]), torch.tensor([1.0]), [10]), + ) + monkeypatch.setattr(fairness_module, "FairnessEvaluator", DummyEvaluator) + monkeypatch.setattr( + fairness_module, + "flatten_fairness_report", + lambda report: {"fairness/gender/n_valid_groups": 1}, + ) + + cfg = OmegaConf.create( + { + "task": {"task_type": "regression"}, + "eval": { + "fairness": { + "enabled": True, + "protected_attributes": ["gender"], + "min_subgroup_size": 5, + } + }, + } + ) + + report = run_fairness_evaluation(DummyModel(), DummyDataModule(), cfg) + + assert report == {"gender": {"n_valid_groups": 1}} + assert captured["task_type"] == "regression" + assert captured["dataset_name"] == "miiv" + + class TestSupervisedCheckpointFormat: """Test that save_encoder_weights from supervised script produces v3 format.""" diff --git a/tests/test_transformer_encoder.py b/tests/test_transformer_encoder.py index 9c4ae44..5db3d46 100644 --- a/tests/test_transformer_encoder.py +++ b/tests/test_transformer_encoder.py @@ -13,6 +13,7 @@ import pytest import torch import torch.nn as nn + from slices.models.common import PositionalEncoding from slices.models.encoders import TransformerConfig, TransformerEncoder from slices.models.encoders.transformer import TransformerEncoderLayer @@ -637,7 +638,7 @@ def test_large_transformer(self): def test_encoder_with_realistic_icu_data_shape(self): """Test encoder with realistic ICU data dimensions.""" - # Typical ICU dataset: 35 features, 48-hour window, batch of 32 + # Benchmark-shaped ICU tensor: 35 features, 24-hour window, batch of 32 config = TransformerConfig( d_input=35, d_model=128, @@ -651,7 +652,7 @@ def test_encoder_with_realistic_icu_data_shape(self): encoder = TransformerEncoder(config) batch_size = 32 - seq_length = 48 + seq_length = 24 n_features = 35 x = torch.randn(batch_size, seq_length, n_features) @@ -661,7 +662,7 @@ def test_encoder_with_realistic_icu_data_shape(self): # Some sequences have variable lengths for i in range(batch_size): if i % 4 == 0: # 25% of sequences - length = torch.randint(24, seq_length, (1,)).item() + length = torch.randint(max(1, seq_length // 2), seq_length, (1,)).item() padding_mask[i, length:] = False out = encoder(x, mask=obs_mask, padding_mask=padding_mask) @@ -860,5 +861,128 @@ def test_realistic_icu_scenario(self): assert out.shape == (batch_size, 128) +class TestTransformerEncoderObsAware: + """Tests for obs_aware TransformerEncoder with tokenize/encode methods.""" + + @pytest.fixture + def config(self): + return TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + max_seq_length=48, + pooling="none", + obs_aware=True, + dropout=0.0, + ) + + @pytest.fixture + def encoder(self, config): + return TransformerEncoder(config) + + def test_tokenize_shapes(self, encoder): + """Verify tokenize returns (B, T, d_model) output.""" + B, T, D = 2, 8, 10 + x = torch.randn(B, T, D) + obs_mask = torch.rand(B, T, D) > 0.3 + + tokens, padding_mask, token_info = encoder.tokenize(x, obs_mask) + + assert tokens.shape == (B, T, 32) + assert padding_mask.shape == (B, T) + assert torch.equal(padding_mask, obs_mask.any(dim=-1)) + assert "timestep_idx" in token_info + assert "values" in token_info + assert "obs_mask" in token_info + assert "valid_timestep_mask" in token_info + assert token_info["timestep_idx"].shape == (B, T) + + def test_encode_shapes(self, encoder): + """Verify encoding a subset of tokens.""" + B, T, D = 2, 8, 10 + x = torch.randn(B, T, D) + obs_mask = torch.rand(B, T, D) > 0.3 + + tokens, padding_mask, _ = encoder.tokenize(x, obs_mask) + + # Take a subset (simulating visible tokens) + n_vis = 4 + vis_tokens = tokens[:, :n_vis, :] + vis_padding = torch.ones(B, n_vis, dtype=torch.bool) + + encoded = encoder.encode(vis_tokens, vis_padding) + assert encoded.shape == (B, n_vis, 32) + + def test_forward_obs_aware(self, encoder): + """Verify obs_proj used when mask provided, different masks -> different outputs.""" + B, T, D = 2, 8, 10 + x = torch.randn(B, T, D) + encoder.eval() + + mask_full = torch.ones(B, T, D, dtype=torch.bool) + mask_sparse = torch.zeros(B, T, D, dtype=torch.bool) + mask_sparse[:, :, :3] = True + + with torch.no_grad(): + out_full = encoder(x, mask=mask_full) + out_sparse = encoder(x, mask=mask_sparse) + + # Different masks should produce different outputs (obs_proj sees mask) + assert not torch.allclose(out_full, out_sparse, atol=1e-5) + + def test_all_empty_obs_aware_row_is_finite(self, encoder): + """Fully unobserved rows should not create all-masked attention NaNs.""" + B, T, D = 2, 8, 10 + x = torch.zeros(B, T, D) + obs_mask = torch.zeros(B, T, D, dtype=torch.bool) + obs_mask[1, 0, 0] = True + encoder.eval() + + with torch.no_grad(): + out = encoder(x, mask=obs_mask, padding_mask=obs_mask.any(dim=-1)) + + assert torch.isfinite(out).all() + + def test_backward_compatible(self, encoder): + """Verify forward(x) without mask still uses input_proj.""" + B, T, D = 2, 8, 10 + x = torch.randn(B, T, D) + encoder.eval() + + with torch.no_grad(): + out = encoder(x) + + assert out.shape == (B, T, 32) # pooling='none' + + def test_gradient_flow_obs_proj(self, encoder): + """Verify gradients flow through obs_proj path.""" + B, T, D = 2, 8, 10 + x = torch.randn(B, T, D) + obs_mask = torch.ones(B, T, D, dtype=torch.bool) + + out = encoder(x, mask=obs_mask) + loss = out.sum() + loss.backward() + + # obs_proj should have gradients + for name, param in encoder.obs_proj.named_parameters(): + assert param.grad is not None, f"No gradient for obs_proj.{name}" + assert param.grad.abs().sum() > 0 + + def test_has_tokenize_encode(self): + """Obs-aware encoder should have tokenize() and encode() methods.""" + config = TransformerConfig( + d_input=10, d_model=32, n_layers=1, n_heads=4, pooling="none", obs_aware=True + ) + encoder = TransformerEncoder(config) + + assert hasattr(encoder, "tokenize") + assert hasattr(encoder, "encode") + assert callable(encoder.tokenize) + assert callable(encoder.encode) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/test_transforms.py b/tests/test_transforms.py index 1af82bf..ce6098f 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -6,6 +6,7 @@ import pytest import torch + from slices.data.transforms import ( apply_gaussian_noise, apply_mask, diff --git a/tests/test_ts2vec_objective.py b/tests/test_ts2vec_objective.py new file mode 100644 index 0000000..32e9324 --- /dev/null +++ b/tests/test_ts2vec_objective.py @@ -0,0 +1,70 @@ +"""Tests for TS2Vec-style temporal contrastive objective.""" + +import pytest +import torch + +from slices.models.encoders import TransformerConfig, TransformerEncoder +from slices.models.pretraining import TS2VecConfig, TS2VecObjective + + +@pytest.fixture +def encoder(): + config = TransformerConfig( + d_input=10, + d_model=32, + n_layers=1, + n_heads=4, + d_ff=64, + pooling="none", + obs_aware=True, + max_seq_length=48, + ) + return TransformerEncoder(config) + + +def test_empty_timesteps_are_excluded_from_ts2vec_masks(encoder): + """TS2Vec should only count observed timesteps as eligible SSL tokens.""" + config = TS2VecConfig( + mask_ratio=0.5, + noise_scale=0.0, + crop_ratio=1.0, + proj_hidden_dim=32, + proj_output_dim=16, + n_hierarchical_scales=2, + ) + objective = TS2VecObjective(encoder, config) + + B, T, D = 2, 8, 10 + x = torch.randn(B, T, D) + obs_mask = torch.zeros(B, T, D, dtype=torch.bool) + obs_mask[:, 1, :3] = True + obs_mask[:, 4, :3] = True + obs_mask[:, 7, :3] = True + + _, metrics = objective(x, obs_mask) + + assert metrics["ts2vec_n_visible_view1"].item() <= 3 + assert metrics["ts2vec_n_visible_view2"].item() <= 3 + assert metrics["ts2vec_n_overlap_per_sample"].item() <= 3 + + +def test_ts2vec_scatter_to_full_ignores_padded_visible_tokens(): + """Uneven visible counts should not scatter padded tokens into masked slots.""" + encoded = torch.tensor( + [ + [[10.0], [11.0]], + [[20.0], [99.0]], + ] + ) + ssl_mask = torch.tensor( + [ + [True, True, False], + [True, False, False], + ] + ) + + full = TS2VecObjective._scatter_to_full(encoded, ssl_mask, n_timesteps=3) + + assert torch.allclose(full[1, 0], torch.tensor([20.0])) + assert torch.allclose(full[1, 1], torch.tensor([0.0])) + assert torch.allclose(full[1, 2], torch.tensor([0.0])) diff --git a/uv.lock b/uv.lock index 91b78d3..ca1df05 100644 --- a/uv.lock +++ b/uv.lock @@ -132,28 +132,6 @@ version = "4.9.3" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/3e/38/7859ff46355f76f8d19459005ca000b6e7012f2f1ca597746cbcd1fbfe5e/antlr4-python3-runtime-4.9.3.tar.gz", hash = "sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b", size = 117034, upload-time = "2021-11-06T17:52:23.524Z" } -[[package]] -name = "anyio" -version = "4.12.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "idna" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/96/f0/5eb65b2bb0d09ac6776f2eb54adee6abe8228ea05b20a5ad0e4945de8aac/anyio-4.12.1.tar.gz", hash = "sha256:41cfcc3a4c85d3f05c932da7c26d0201ac36f72abd4435ba90d0464a3ffed703", size = 228685, upload-time = "2026-01-06T11:45:21.246Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/38/0e/27be9fdef66e72d64c0cdc3cc2823101b80585f8119b5c112c2e8f5f7dab/anyio-4.12.1-py3-none-any.whl", hash = "sha256:d405828884fc140aa80a3c667b8beed277f1dfedec42ba031bd6ac3db606ab6c", size = 113592, upload-time = "2026-01-06T11:45:19.497Z" }, -] - -[[package]] -name = "appdirs" -version = "1.4.4" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/d7/d8/05696357e0311f5b5c316d7b95f46c669dd9c15aaeecbb48c7d0aeb88c40/appdirs-1.4.4.tar.gz", hash = "sha256:7d5d0167b2b1ba821647616af46a749d1c653740dd0d2415100fe26e27afdf41", size = 13470, upload-time = "2020-05-11T07:59:51.037Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3b/00/2344469e2084fb287c2e0b57b72910309874c3245463acd6cf5e3db69324/appdirs-1.4.4-py2.py3-none-any.whl", hash = "sha256:a841dacd6b99318a741b166adb07e19ee71a274450e68237b4650ca1055ab128", size = 9566, upload-time = "2020-05-11T07:59:49.499Z" }, -] - [[package]] name = "attrs" version = "25.4.0" @@ -163,62 +141,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3a/2a/7cc015f5b9f5db42b7d48157e23356022889fc354a2813c15934b7cb5c0e/attrs-25.4.0-py3-none-any.whl", hash = "sha256:adcf7e2a1fb3b36ac48d97835bb6d8ade15b8dcce26aba8bf1d14847b57a3373", size = 67615, upload-time = "2025-10-06T13:54:43.17Z" }, ] -[[package]] -name = "authlib" -version = "1.6.6" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cryptography" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/bb/9b/b1661026ff24bc641b76b78c5222d614776b0c085bcfdac9bd15a1cb4b35/authlib-1.6.6.tar.gz", hash = "sha256:45770e8e056d0f283451d9996fbb59b70d45722b45d854d58f32878d0a40c38e", size = 164894, upload-time = "2025-12-12T08:01:41.464Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/54/51/321e821856452f7386c4e9df866f196720b1ad0c5ea1623ea7399969ae3b/authlib-1.6.6-py2.py3-none-any.whl", hash = "sha256:7d9e9bc535c13974313a87f53e8430eb6ea3d1cf6ae4f6efcd793f2e949143fd", size = 244005, upload-time = "2025-12-12T08:01:40.209Z" }, -] - -[[package]] -name = "autograd" -version = "1.8.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/67/1c/3c24ec03c8ba4decc742b1df5a10c52f98c84ca8797757f313e7bdcdf276/autograd-1.8.0.tar.gz", hash = "sha256:107374ded5b09fc8643ac925348c0369e7b0e73bbed9565ffd61b8fd04425683", size = 2562146, upload-time = "2025-05-05T12:49:02.502Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/84/ea/e16f0c423f7d83cf8b79cae9452040fb7b2e020c7439a167ee7c317de448/autograd-1.8.0-py3-none-any.whl", hash = "sha256:4ab9084294f814cf56c280adbe19612546a35574d67c574b04933c7d2ecb7d78", size = 51478, upload-time = "2025-05-05T12:49:00.585Z" }, -] - -[[package]] -name = "autograd-gamma" -version = "0.5.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "autograd" }, - { name = "scipy" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/85/ae/7f2031ea76140444b2453fa139041e5afd4a09fc5300cfefeb1103291f80/autograd-gamma-0.5.0.tar.gz", hash = "sha256:f27abb7b8bb9cffc8badcbf59f3fe44a9db39e124ecacf1992b6d952934ac9c4", size = 3952, upload-time = "2020-10-15T16:51:06.785Z" } - -[[package]] -name = "beartype" -version = "0.22.9" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/c7/94/1009e248bbfbab11397abca7193bea6626806be9a327d399810d523a07cb/beartype-0.22.9.tar.gz", hash = "sha256:8f82b54aa723a2848a56008d18875f91c1db02c32ef6a62319a002e3e25a975f", size = 1608866, upload-time = "2025-12-13T06:50:30.72Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl", hash = "sha256:d16c9bbc61ea14637596c5f6fbff2ee99cbe3573e46a716401734ef50c3060c2", size = 1333658, upload-time = "2025-12-13T06:50:28.266Z" }, -] - -[[package]] -name = "beautifulsoup4" -version = "4.14.3" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "soupsieve" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/c3/b0/1c6a16426d389813b48d95e26898aff79abbde42ad353958ad95cc8c9b21/beautifulsoup4-4.14.3.tar.gz", hash = "sha256:6292b1c5186d356bba669ef9f7f051757099565ad9ada5dd630bd9de5fa7fb86", size = 627737, upload-time = "2025-11-30T15:08:26.084Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1a/39/47f9197bdd44df24d67ac8893641e16f386c984a0619ef2ee4c51fbbc019/beautifulsoup4-4.14.3-py3-none-any.whl", hash = "sha256:0918bfe44902e6ad8d57732ba310582e98da931428d231a5ecb9e7c703a735bb", size = 107721, upload-time = "2025-11-30T15:08:24.087Z" }, -] - [[package]] name = "black" version = "25.11.0" @@ -248,15 +170,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/00/5d/aed32636ed30a6e7f9efd6ad14e2a0b0d687ae7c8c7ec4e4a557174b895c/black-25.11.0-py3-none-any.whl", hash = "sha256:e3f562da087791e96cefcd9dda058380a442ab322a02e222add53736451f604b", size = 204918, upload-time = "2025-11-10T01:53:48.917Z" }, ] -[[package]] -name = "cachetools" -version = "6.2.4" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/bc/1d/ede8680603f6016887c062a2cf4fc8fdba905866a3ab8831aa8aa651320c/cachetools-6.2.4.tar.gz", hash = "sha256:82c5c05585e70b6ba2d3ae09ea60b79548872185d2f24ae1f2709d37299fd607", size = 31731, upload-time = "2025-12-15T18:24:53.744Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/2c/fc/1d7b80d0eb7b714984ce40efc78859c022cd930e402f599d8ca9e39c78a4/cachetools-6.2.4-py3-none-any.whl", hash = "sha256:69a7a52634fed8b8bf6e24a050fb60bff1c9bd8f6d24572b99c32d4e71e62a51", size = 11551, upload-time = "2025-12-15T18:24:52.332Z" }, -] - [[package]] name = "certifi" version = "2025.11.12" @@ -266,63 +179,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/70/7d/9bc192684cea499815ff478dfcdc13835ddf401365057044fb721ec6bddb/certifi-2025.11.12-py3-none-any.whl", hash = "sha256:97de8790030bbd5c2d96b7ec782fc2f7820ef8dba6db909ccf95449f2d062d4b", size = 159438, upload-time = "2025-11-12T02:54:49.735Z" }, ] -[[package]] -name = "cffi" -version = "2.0.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pycparser", marker = "implementation_name != 'PyPy'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8453301356628e8147c79dbb825bcbc73dc7401f9846/cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529", size = 523588, upload-time = "2025-09-08T23:24:04.541Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ea/47/4f61023ea636104d4f16ab488e268b93008c3d0bb76893b1b31db1f96802/cffi-2.0.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:6d02d6655b0e54f54c4ef0b94eb6be0607b70853c45ce98bd278dc7de718be5d", size = 185271, upload-time = "2025-09-08T23:22:44.795Z" }, - { url = "https://files.pythonhosted.org/packages/df/a2/781b623f57358e360d62cdd7a8c681f074a71d445418a776eef0aadb4ab4/cffi-2.0.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8eca2a813c1cb7ad4fb74d368c2ffbbb4789d377ee5bb8df98373c2cc0dee76c", size = 181048, upload-time = "2025-09-08T23:22:45.938Z" }, - { url = "https://files.pythonhosted.org/packages/ff/df/a4f0fbd47331ceeba3d37c2e51e9dfc9722498becbeec2bd8bc856c9538a/cffi-2.0.0-cp312-cp312-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:21d1152871b019407d8ac3985f6775c079416c282e431a4da6afe7aefd2bccbe", size = 212529, upload-time = "2025-09-08T23:22:47.349Z" }, - { url = "https://files.pythonhosted.org/packages/d5/72/12b5f8d3865bf0f87cf1404d8c374e7487dcf097a1c91c436e72e6badd83/cffi-2.0.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:b21e08af67b8a103c71a250401c78d5e0893beff75e28c53c98f4de42f774062", size = 220097, upload-time = "2025-09-08T23:22:48.677Z" }, - { url = "https://files.pythonhosted.org/packages/c2/95/7a135d52a50dfa7c882ab0ac17e8dc11cec9d55d2c18dda414c051c5e69e/cffi-2.0.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:1e3a615586f05fc4065a8b22b8152f0c1b00cdbc60596d187c2a74f9e3036e4e", size = 207983, upload-time = "2025-09-08T23:22:50.06Z" }, - { url = "https://files.pythonhosted.org/packages/3a/c8/15cb9ada8895957ea171c62dc78ff3e99159ee7adb13c0123c001a2546c1/cffi-2.0.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:81afed14892743bbe14dacb9e36d9e0e504cd204e0b165062c488942b9718037", size = 206519, upload-time = "2025-09-08T23:22:51.364Z" }, - { url = "https://files.pythonhosted.org/packages/78/2d/7fa73dfa841b5ac06c7b8855cfc18622132e365f5b81d02230333ff26e9e/cffi-2.0.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3e17ed538242334bf70832644a32a7aae3d83b57567f9fd60a26257e992b79ba", size = 219572, upload-time = "2025-09-08T23:22:52.902Z" }, - { url = "https://files.pythonhosted.org/packages/07/e0/267e57e387b4ca276b90f0434ff88b2c2241ad72b16d31836adddfd6031b/cffi-2.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3925dd22fa2b7699ed2617149842d2e6adde22b262fcbfada50e3d195e4b3a94", size = 222963, upload-time = "2025-09-08T23:22:54.518Z" }, - { url = "https://files.pythonhosted.org/packages/b6/75/1f2747525e06f53efbd878f4d03bac5b859cbc11c633d0fb81432d98a795/cffi-2.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:2c8f814d84194c9ea681642fd164267891702542f028a15fc97d4674b6206187", size = 221361, upload-time = "2025-09-08T23:22:55.867Z" }, - { url = "https://files.pythonhosted.org/packages/7b/2b/2b6435f76bfeb6bbf055596976da087377ede68df465419d192acf00c437/cffi-2.0.0-cp312-cp312-win32.whl", hash = "sha256:da902562c3e9c550df360bfa53c035b2f241fed6d9aef119048073680ace4a18", size = 172932, upload-time = "2025-09-08T23:22:57.188Z" }, - { url = "https://files.pythonhosted.org/packages/f8/ed/13bd4418627013bec4ed6e54283b1959cf6db888048c7cf4b4c3b5b36002/cffi-2.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:da68248800ad6320861f129cd9c1bf96ca849a2771a59e0344e88681905916f5", size = 183557, upload-time = "2025-09-08T23:22:58.351Z" }, - { url = "https://files.pythonhosted.org/packages/95/31/9f7f93ad2f8eff1dbc1c3656d7ca5bfd8fb52c9d786b4dcf19b2d02217fa/cffi-2.0.0-cp312-cp312-win_arm64.whl", hash = "sha256:4671d9dd5ec934cb9a73e7ee9676f9362aba54f7f34910956b84d727b0d73fb6", size = 177762, upload-time = "2025-09-08T23:22:59.668Z" }, - { url = "https://files.pythonhosted.org/packages/4b/8d/a0a47a0c9e413a658623d014e91e74a50cdd2c423f7ccfd44086ef767f90/cffi-2.0.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:00bdf7acc5f795150faa6957054fbbca2439db2f775ce831222b66f192f03beb", size = 185230, upload-time = "2025-09-08T23:23:00.879Z" }, - { url = "https://files.pythonhosted.org/packages/4a/d2/a6c0296814556c68ee32009d9c2ad4f85f2707cdecfd7727951ec228005d/cffi-2.0.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:45d5e886156860dc35862657e1494b9bae8dfa63bf56796f2fb56e1679fc0bca", size = 181043, upload-time = "2025-09-08T23:23:02.231Z" }, - { url = "https://files.pythonhosted.org/packages/b0/1e/d22cc63332bd59b06481ceaac49d6c507598642e2230f201649058a7e704/cffi-2.0.0-cp313-cp313-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:07b271772c100085dd28b74fa0cd81c8fb1a3ba18b21e03d7c27f3436a10606b", size = 212446, upload-time = "2025-09-08T23:23:03.472Z" }, - { url = "https://files.pythonhosted.org/packages/a9/f5/a2c23eb03b61a0b8747f211eb716446c826ad66818ddc7810cc2cc19b3f2/cffi-2.0.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d48a880098c96020b02d5a1f7d9251308510ce8858940e6fa99ece33f610838b", size = 220101, upload-time = "2025-09-08T23:23:04.792Z" }, - { url = "https://files.pythonhosted.org/packages/f2/7f/e6647792fc5850d634695bc0e6ab4111ae88e89981d35ac269956605feba/cffi-2.0.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:f93fd8e5c8c0a4aa1f424d6173f14a892044054871c771f8566e4008eaa359d2", size = 207948, upload-time = "2025-09-08T23:23:06.127Z" }, - { url = "https://files.pythonhosted.org/packages/cb/1e/a5a1bd6f1fb30f22573f76533de12a00bf274abcdc55c8edab639078abb6/cffi-2.0.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:dd4f05f54a52fb558f1ba9f528228066954fee3ebe629fc1660d874d040ae5a3", size = 206422, upload-time = "2025-09-08T23:23:07.753Z" }, - { url = "https://files.pythonhosted.org/packages/98/df/0a1755e750013a2081e863e7cd37e0cdd02664372c754e5560099eb7aa44/cffi-2.0.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c8d3b5532fc71b7a77c09192b4a5a200ea992702734a2e9279a37f2478236f26", size = 219499, upload-time = "2025-09-08T23:23:09.648Z" }, - { url = "https://files.pythonhosted.org/packages/50/e1/a969e687fcf9ea58e6e2a928ad5e2dd88cc12f6f0ab477e9971f2309b57c/cffi-2.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:d9b29c1f0ae438d5ee9acb31cadee00a58c46cc9c0b2f9038c6b0b3470877a8c", size = 222928, upload-time = "2025-09-08T23:23:10.928Z" }, - { url = "https://files.pythonhosted.org/packages/36/54/0362578dd2c9e557a28ac77698ed67323ed5b9775ca9d3fe73fe191bb5d8/cffi-2.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6d50360be4546678fc1b79ffe7a66265e28667840010348dd69a314145807a1b", size = 221302, upload-time = "2025-09-08T23:23:12.42Z" }, - { url = "https://files.pythonhosted.org/packages/eb/6d/bf9bda840d5f1dfdbf0feca87fbdb64a918a69bca42cfa0ba7b137c48cb8/cffi-2.0.0-cp313-cp313-win32.whl", hash = "sha256:74a03b9698e198d47562765773b4a8309919089150a0bb17d829ad7b44b60d27", size = 172909, upload-time = "2025-09-08T23:23:14.32Z" }, - { url = "https://files.pythonhosted.org/packages/37/18/6519e1ee6f5a1e579e04b9ddb6f1676c17368a7aba48299c3759bbc3c8b3/cffi-2.0.0-cp313-cp313-win_amd64.whl", hash = "sha256:19f705ada2530c1167abacb171925dd886168931e0a7b78f5bffcae5c6b5be75", size = 183402, upload-time = "2025-09-08T23:23:15.535Z" }, - { url = "https://files.pythonhosted.org/packages/cb/0e/02ceeec9a7d6ee63bb596121c2c8e9b3a9e150936f4fbef6ca1943e6137c/cffi-2.0.0-cp313-cp313-win_arm64.whl", hash = "sha256:256f80b80ca3853f90c21b23ee78cd008713787b1b1e93eae9f3d6a7134abd91", size = 177780, upload-time = "2025-09-08T23:23:16.761Z" }, - { url = "https://files.pythonhosted.org/packages/92/c4/3ce07396253a83250ee98564f8d7e9789fab8e58858f35d07a9a2c78de9f/cffi-2.0.0-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:fc33c5141b55ed366cfaad382df24fe7dcbc686de5be719b207bb248e3053dc5", size = 185320, upload-time = "2025-09-08T23:23:18.087Z" }, - { url = "https://files.pythonhosted.org/packages/59/dd/27e9fa567a23931c838c6b02d0764611c62290062a6d4e8ff7863daf9730/cffi-2.0.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c654de545946e0db659b3400168c9ad31b5d29593291482c43e3564effbcee13", size = 181487, upload-time = "2025-09-08T23:23:19.622Z" }, - { url = "https://files.pythonhosted.org/packages/d6/43/0e822876f87ea8a4ef95442c3d766a06a51fc5298823f884ef87aaad168c/cffi-2.0.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:24b6f81f1983e6df8db3adc38562c83f7d4a0c36162885ec7f7b77c7dcbec97b", size = 220049, upload-time = "2025-09-08T23:23:20.853Z" }, - { url = "https://files.pythonhosted.org/packages/b4/89/76799151d9c2d2d1ead63c2429da9ea9d7aac304603de0c6e8764e6e8e70/cffi-2.0.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:12873ca6cb9b0f0d3a0da705d6086fe911591737a59f28b7936bdfed27c0d47c", size = 207793, upload-time = "2025-09-08T23:23:22.08Z" }, - { url = "https://files.pythonhosted.org/packages/bb/dd/3465b14bb9e24ee24cb88c9e3730f6de63111fffe513492bf8c808a3547e/cffi-2.0.0-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:d9b97165e8aed9272a6bb17c01e3cc5871a594a446ebedc996e2397a1c1ea8ef", size = 206300, upload-time = "2025-09-08T23:23:23.314Z" }, - { url = "https://files.pythonhosted.org/packages/47/d9/d83e293854571c877a92da46fdec39158f8d7e68da75bf73581225d28e90/cffi-2.0.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:afb8db5439b81cf9c9d0c80404b60c3cc9c3add93e114dcae767f1477cb53775", size = 219244, upload-time = "2025-09-08T23:23:24.541Z" }, - { url = "https://files.pythonhosted.org/packages/2b/0f/1f177e3683aead2bb00f7679a16451d302c436b5cbf2505f0ea8146ef59e/cffi-2.0.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:737fe7d37e1a1bffe70bd5754ea763a62a066dc5913ca57e957824b72a85e205", size = 222828, upload-time = "2025-09-08T23:23:26.143Z" }, - { url = "https://files.pythonhosted.org/packages/c6/0f/cafacebd4b040e3119dcb32fed8bdef8dfe94da653155f9d0b9dc660166e/cffi-2.0.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:38100abb9d1b1435bc4cc340bb4489635dc2f0da7456590877030c9b3d40b0c1", size = 220926, upload-time = "2025-09-08T23:23:27.873Z" }, - { url = "https://files.pythonhosted.org/packages/3e/aa/df335faa45b395396fcbc03de2dfcab242cd61a9900e914fe682a59170b1/cffi-2.0.0-cp314-cp314-win32.whl", hash = "sha256:087067fa8953339c723661eda6b54bc98c5625757ea62e95eb4898ad5e776e9f", size = 175328, upload-time = "2025-09-08T23:23:44.61Z" }, - { url = "https://files.pythonhosted.org/packages/bb/92/882c2d30831744296ce713f0feb4c1cd30f346ef747b530b5318715cc367/cffi-2.0.0-cp314-cp314-win_amd64.whl", hash = "sha256:203a48d1fb583fc7d78a4c6655692963b860a417c0528492a6bc21f1aaefab25", size = 185650, upload-time = "2025-09-08T23:23:45.848Z" }, - { url = "https://files.pythonhosted.org/packages/9f/2c/98ece204b9d35a7366b5b2c6539c350313ca13932143e79dc133ba757104/cffi-2.0.0-cp314-cp314-win_arm64.whl", hash = "sha256:dbd5c7a25a7cb98f5ca55d258b103a2054f859a46ae11aaf23134f9cc0d356ad", size = 180687, upload-time = "2025-09-08T23:23:47.105Z" }, - { url = "https://files.pythonhosted.org/packages/3e/61/c768e4d548bfa607abcda77423448df8c471f25dbe64fb2ef6d555eae006/cffi-2.0.0-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:9a67fc9e8eb39039280526379fb3a70023d77caec1852002b4da7e8b270c4dd9", size = 188773, upload-time = "2025-09-08T23:23:29.347Z" }, - { url = "https://files.pythonhosted.org/packages/2c/ea/5f76bce7cf6fcd0ab1a1058b5af899bfbef198bea4d5686da88471ea0336/cffi-2.0.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7a66c7204d8869299919db4d5069a82f1561581af12b11b3c9f48c584eb8743d", size = 185013, upload-time = "2025-09-08T23:23:30.63Z" }, - { url = "https://files.pythonhosted.org/packages/be/b4/c56878d0d1755cf9caa54ba71e5d049479c52f9e4afc230f06822162ab2f/cffi-2.0.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7cc09976e8b56f8cebd752f7113ad07752461f48a58cbba644139015ac24954c", size = 221593, upload-time = "2025-09-08T23:23:31.91Z" }, - { url = "https://files.pythonhosted.org/packages/e0/0d/eb704606dfe8033e7128df5e90fee946bbcb64a04fcdaa97321309004000/cffi-2.0.0-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:92b68146a71df78564e4ef48af17551a5ddd142e5190cdf2c5624d0c3ff5b2e8", size = 209354, upload-time = "2025-09-08T23:23:33.214Z" }, - { url = "https://files.pythonhosted.org/packages/d8/19/3c435d727b368ca475fb8742ab97c9cb13a0de600ce86f62eab7fa3eea60/cffi-2.0.0-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:b1e74d11748e7e98e2f426ab176d4ed720a64412b6a15054378afdb71e0f37dc", size = 208480, upload-time = "2025-09-08T23:23:34.495Z" }, - { url = "https://files.pythonhosted.org/packages/d0/44/681604464ed9541673e486521497406fadcc15b5217c3e326b061696899a/cffi-2.0.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:28a3a209b96630bca57cce802da70c266eb08c6e97e5afd61a75611ee6c64592", size = 221584, upload-time = "2025-09-08T23:23:36.096Z" }, - { url = "https://files.pythonhosted.org/packages/25/8e/342a504ff018a2825d395d44d63a767dd8ebc927ebda557fecdaca3ac33a/cffi-2.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:7553fb2090d71822f02c629afe6042c299edf91ba1bf94951165613553984512", size = 224443, upload-time = "2025-09-08T23:23:37.328Z" }, - { url = "https://files.pythonhosted.org/packages/e1/5e/b666bacbbc60fbf415ba9988324a132c9a7a0448a9a8f125074671c0f2c3/cffi-2.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:6c6c373cfc5c83a975506110d17457138c8c63016b563cc9ed6e056a82f13ce4", size = 223437, upload-time = "2025-09-08T23:23:38.945Z" }, - { url = "https://files.pythonhosted.org/packages/a0/1d/ec1a60bd1a10daa292d3cd6bb0b359a81607154fb8165f3ec95fe003b85c/cffi-2.0.0-cp314-cp314t-win32.whl", hash = "sha256:1fc9ea04857caf665289b7a75923f2c6ed559b8298a1b8c49e59f7dd95c8481e", size = 180487, upload-time = "2025-09-08T23:23:40.423Z" }, - { url = "https://files.pythonhosted.org/packages/bf/41/4c1168c74fac325c0c8156f04b6749c8b6a8f405bbf91413ba088359f60d/cffi-2.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:d68b6cef7827e8641e8ef16f4494edda8b36104d79773a334beaa1e3521430f6", size = 191726, upload-time = "2025-09-08T23:23:41.742Z" }, - { url = "https://files.pythonhosted.org/packages/ae/3a/dbeec9d1ee0844c679f6bb5d6ad4e9f198b1224f4e7a32825f47f6192b0c/cffi-2.0.0-cp314-cp314t-win_arm64.whl", hash = "sha256:0a1527a803f0a659de1af2e1fd700213caba79377e27e4693648c2923da066f9", size = 184195, upload-time = "2025-09-08T23:23:43.004Z" }, -] - [[package]] name = "cfgv" version = "3.5.0" @@ -401,15 +257,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/98/78/01c019cdb5d6498122777c1a43056ebb3ebfeef2076d9d026bfe15583b2b/click-8.3.1-py3-none-any.whl", hash = "sha256:981153a64e25f12d547d3426c367a4857371575ee7ad18df2a6183ab0545b2a6", size = 108274, upload-time = "2025-11-15T20:45:41.139Z" }, ] -[[package]] -name = "cloudpickle" -version = "3.1.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, -] - [[package]] name = "colorama" version = "0.4.6" @@ -559,62 +406,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ce/a3/43b749004e3c09452e39bb56347a008f0a0668aad37324a99b5c8ca91d9e/coverage-7.12.0-py3-none-any.whl", hash = "sha256:159d50c0b12e060b15ed3d39f87ed43d4f7f7ad40b8a534f4dd331adbb51104a", size = 209503, upload-time = "2025-11-18T13:34:18.892Z" }, ] -[[package]] -name = "cryptography" -version = "46.0.3" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cffi", marker = "platform_python_implementation != 'PyPy'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/9f/33/c00162f49c0e2fe8064a62cb92b93e50c74a72bc370ab92f86112b33ff62/cryptography-46.0.3.tar.gz", hash = "sha256:a8b17438104fed022ce745b362294d9ce35b4c2e45c1d958ad4a4b019285f4a1", size = 749258, upload-time = "2025-10-15T23:18:31.74Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1d/42/9c391dd801d6cf0d561b5890549d4b27bafcc53b39c31a817e69d87c625b/cryptography-46.0.3-cp311-abi3-macosx_10_9_universal2.whl", hash = "sha256:109d4ddfadf17e8e7779c39f9b18111a09efb969a301a31e987416a0191ed93a", size = 7225004, upload-time = "2025-10-15T23:16:52.239Z" }, - { url = "https://files.pythonhosted.org/packages/1c/67/38769ca6b65f07461eb200e85fc1639b438bdc667be02cf7f2cd6a64601c/cryptography-46.0.3-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:09859af8466b69bc3c27bdf4f5d84a665e0f7ab5088412e9e2ec49758eca5cbc", size = 4296667, upload-time = "2025-10-15T23:16:54.369Z" }, - { url = "https://files.pythonhosted.org/packages/5c/49/498c86566a1d80e978b42f0d702795f69887005548c041636df6ae1ca64c/cryptography-46.0.3-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:01ca9ff2885f3acc98c29f1860552e37f6d7c7d013d7334ff2a9de43a449315d", size = 4450807, upload-time = "2025-10-15T23:16:56.414Z" }, - { url = "https://files.pythonhosted.org/packages/4b/0a/863a3604112174c8624a2ac3c038662d9e59970c7f926acdcfaed8d61142/cryptography-46.0.3-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:6eae65d4c3d33da080cff9c4ab1f711b15c1d9760809dad6ea763f3812d254cb", size = 4299615, upload-time = "2025-10-15T23:16:58.442Z" }, - { url = "https://files.pythonhosted.org/packages/64/02/b73a533f6b64a69f3cd3872acb6ebc12aef924d8d103133bb3ea750dc703/cryptography-46.0.3-cp311-abi3-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:e5bf0ed4490068a2e72ac03d786693adeb909981cc596425d09032d372bcc849", size = 4016800, upload-time = "2025-10-15T23:17:00.378Z" }, - { url = "https://files.pythonhosted.org/packages/25/d5/16e41afbfa450cde85a3b7ec599bebefaef16b5c6ba4ec49a3532336ed72/cryptography-46.0.3-cp311-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:5ecfccd2329e37e9b7112a888e76d9feca2347f12f37918facbb893d7bb88ee8", size = 4984707, upload-time = "2025-10-15T23:17:01.98Z" }, - { url = "https://files.pythonhosted.org/packages/c9/56/e7e69b427c3878352c2fb9b450bd0e19ed552753491d39d7d0a2f5226d41/cryptography-46.0.3-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:a2c0cd47381a3229c403062f764160d57d4d175e022c1df84e168c6251a22eec", size = 4482541, upload-time = "2025-10-15T23:17:04.078Z" }, - { url = "https://files.pythonhosted.org/packages/78/f6/50736d40d97e8483172f1bb6e698895b92a223dba513b0ca6f06b2365339/cryptography-46.0.3-cp311-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:549e234ff32571b1f4076ac269fcce7a808d3bf98b76c8dd560e42dbc66d7d91", size = 4299464, upload-time = "2025-10-15T23:17:05.483Z" }, - { url = "https://files.pythonhosted.org/packages/00/de/d8e26b1a855f19d9994a19c702fa2e93b0456beccbcfe437eda00e0701f2/cryptography-46.0.3-cp311-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:c0a7bb1a68a5d3471880e264621346c48665b3bf1c3759d682fc0864c540bd9e", size = 4950838, upload-time = "2025-10-15T23:17:07.425Z" }, - { url = "https://files.pythonhosted.org/packages/8f/29/798fc4ec461a1c9e9f735f2fc58741b0daae30688f41b2497dcbc9ed1355/cryptography-46.0.3-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:10b01676fc208c3e6feeb25a8b83d81767e8059e1fe86e1dc62d10a3018fa926", size = 4481596, upload-time = "2025-10-15T23:17:09.343Z" }, - { url = "https://files.pythonhosted.org/packages/15/8d/03cd48b20a573adfff7652b76271078e3045b9f49387920e7f1f631d125e/cryptography-46.0.3-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:0abf1ffd6e57c67e92af68330d05760b7b7efb243aab8377e583284dbab72c71", size = 4426782, upload-time = "2025-10-15T23:17:11.22Z" }, - { url = "https://files.pythonhosted.org/packages/fa/b1/ebacbfe53317d55cf33165bda24c86523497a6881f339f9aae5c2e13e57b/cryptography-46.0.3-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:a04bee9ab6a4da801eb9b51f1b708a1b5b5c9eb48c03f74198464c66f0d344ac", size = 4698381, upload-time = "2025-10-15T23:17:12.829Z" }, - { url = "https://files.pythonhosted.org/packages/96/92/8a6a9525893325fc057a01f654d7efc2c64b9de90413adcf605a85744ff4/cryptography-46.0.3-cp311-abi3-win32.whl", hash = "sha256:f260d0d41e9b4da1ed1e0f1ce571f97fe370b152ab18778e9e8f67d6af432018", size = 3055988, upload-time = "2025-10-15T23:17:14.65Z" }, - { url = "https://files.pythonhosted.org/packages/7e/bf/80fbf45253ea585a1e492a6a17efcb93467701fa79e71550a430c5e60df0/cryptography-46.0.3-cp311-abi3-win_amd64.whl", hash = "sha256:a9a3008438615669153eb86b26b61e09993921ebdd75385ddd748702c5adfddb", size = 3514451, upload-time = "2025-10-15T23:17:16.142Z" }, - { url = "https://files.pythonhosted.org/packages/2e/af/9b302da4c87b0beb9db4e756386a7c6c5b8003cd0e742277888d352ae91d/cryptography-46.0.3-cp311-abi3-win_arm64.whl", hash = "sha256:5d7f93296ee28f68447397bf5198428c9aeeab45705a55d53a6343455dcb2c3c", size = 2928007, upload-time = "2025-10-15T23:17:18.04Z" }, - { url = "https://files.pythonhosted.org/packages/f5/e2/a510aa736755bffa9d2f75029c229111a1d02f8ecd5de03078f4c18d91a3/cryptography-46.0.3-cp314-cp314t-macosx_10_9_universal2.whl", hash = "sha256:00a5e7e87938e5ff9ff5447ab086a5706a957137e6e433841e9d24f38a065217", size = 7158012, upload-time = "2025-10-15T23:17:19.982Z" }, - { url = "https://files.pythonhosted.org/packages/73/dc/9aa866fbdbb95b02e7f9d086f1fccfeebf8953509b87e3f28fff927ff8a0/cryptography-46.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c8daeb2d2174beb4575b77482320303f3d39b8e81153da4f0fb08eb5fe86a6c5", size = 4288728, upload-time = "2025-10-15T23:17:21.527Z" }, - { url = "https://files.pythonhosted.org/packages/c5/fd/bc1daf8230eaa075184cbbf5f8cd00ba9db4fd32d63fb83da4671b72ed8a/cryptography-46.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:39b6755623145ad5eff1dab323f4eae2a32a77a7abef2c5089a04a3d04366715", size = 4435078, upload-time = "2025-10-15T23:17:23.042Z" }, - { url = "https://files.pythonhosted.org/packages/82/98/d3bd5407ce4c60017f8ff9e63ffee4200ab3e23fe05b765cab805a7db008/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:db391fa7c66df6762ee3f00c95a89e6d428f4d60e7abc8328f4fe155b5ac6e54", size = 4293460, upload-time = "2025-10-15T23:17:24.885Z" }, - { url = "https://files.pythonhosted.org/packages/26/e9/e23e7900983c2b8af7a08098db406cf989d7f09caea7897e347598d4cd5b/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:78a97cf6a8839a48c49271cdcbd5cf37ca2c1d6b7fdd86cc864f302b5e9bf459", size = 3995237, upload-time = "2025-10-15T23:17:26.449Z" }, - { url = "https://files.pythonhosted.org/packages/91/15/af68c509d4a138cfe299d0d7ddb14afba15233223ebd933b4bbdbc7155d3/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_ppc64le.whl", hash = "sha256:dfb781ff7eaa91a6f7fd41776ec37c5853c795d3b358d4896fdbb5df168af422", size = 4967344, upload-time = "2025-10-15T23:17:28.06Z" }, - { url = "https://files.pythonhosted.org/packages/ca/e3/8643d077c53868b681af077edf6b3cb58288b5423610f21c62aadcbe99f4/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:6f61efb26e76c45c4a227835ddeae96d83624fb0d29eb5df5b96e14ed1a0afb7", size = 4466564, upload-time = "2025-10-15T23:17:29.665Z" }, - { url = "https://files.pythonhosted.org/packages/0e/43/c1e8726fa59c236ff477ff2b5dc071e54b21e5a1e51aa2cee1676f1c986f/cryptography-46.0.3-cp314-cp314t-manylinux_2_34_aarch64.whl", hash = "sha256:23b1a8f26e43f47ceb6d6a43115f33a5a37d57df4ea0ca295b780ae8546e8044", size = 4292415, upload-time = "2025-10-15T23:17:31.686Z" }, - { url = "https://files.pythonhosted.org/packages/42/f9/2f8fefdb1aee8a8e3256a0568cffc4e6d517b256a2fe97a029b3f1b9fe7e/cryptography-46.0.3-cp314-cp314t-manylinux_2_34_ppc64le.whl", hash = "sha256:b419ae593c86b87014b9be7396b385491ad7f320bde96826d0dd174459e54665", size = 4931457, upload-time = "2025-10-15T23:17:33.478Z" }, - { url = "https://files.pythonhosted.org/packages/79/30/9b54127a9a778ccd6d27c3da7563e9f2d341826075ceab89ae3b41bf5be2/cryptography-46.0.3-cp314-cp314t-manylinux_2_34_x86_64.whl", hash = "sha256:50fc3343ac490c6b08c0cf0d704e881d0d660be923fd3076db3e932007e726e3", size = 4466074, upload-time = "2025-10-15T23:17:35.158Z" }, - { url = "https://files.pythonhosted.org/packages/ac/68/b4f4a10928e26c941b1b6a179143af9f4d27d88fe84a6a3c53592d2e76bf/cryptography-46.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:22d7e97932f511d6b0b04f2bfd818d73dcd5928db509460aaf48384778eb6d20", size = 4420569, upload-time = "2025-10-15T23:17:37.188Z" }, - { url = "https://files.pythonhosted.org/packages/a3/49/3746dab4c0d1979888f125226357d3262a6dd40e114ac29e3d2abdf1ec55/cryptography-46.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:d55f3dffadd674514ad19451161118fd010988540cee43d8bc20675e775925de", size = 4681941, upload-time = "2025-10-15T23:17:39.236Z" }, - { url = "https://files.pythonhosted.org/packages/fd/30/27654c1dbaf7e4a3531fa1fc77986d04aefa4d6d78259a62c9dc13d7ad36/cryptography-46.0.3-cp314-cp314t-win32.whl", hash = "sha256:8a6e050cb6164d3f830453754094c086ff2d0b2f3a897a1d9820f6139a1f0914", size = 3022339, upload-time = "2025-10-15T23:17:40.888Z" }, - { url = "https://files.pythonhosted.org/packages/f6/30/640f34ccd4d2a1bc88367b54b926b781b5a018d65f404d409aba76a84b1c/cryptography-46.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:760f83faa07f8b64e9c33fc963d790a2edb24efb479e3520c14a45741cd9b2db", size = 3494315, upload-time = "2025-10-15T23:17:42.769Z" }, - { url = "https://files.pythonhosted.org/packages/ba/8b/88cc7e3bd0a8e7b861f26981f7b820e1f46aa9d26cc482d0feba0ecb4919/cryptography-46.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:516ea134e703e9fe26bcd1277a4b59ad30586ea90c365a87781d7887a646fe21", size = 2919331, upload-time = "2025-10-15T23:17:44.468Z" }, - { url = "https://files.pythonhosted.org/packages/fd/23/45fe7f376a7df8daf6da3556603b36f53475a99ce4faacb6ba2cf3d82021/cryptography-46.0.3-cp38-abi3-macosx_10_9_universal2.whl", hash = "sha256:cb3d760a6117f621261d662bccc8ef5bc32ca673e037c83fbe565324f5c46936", size = 7218248, upload-time = "2025-10-15T23:17:46.294Z" }, - { url = "https://files.pythonhosted.org/packages/27/32/b68d27471372737054cbd34c84981f9edbc24fe67ca225d389799614e27f/cryptography-46.0.3-cp38-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:4b7387121ac7d15e550f5cb4a43aef2559ed759c35df7336c402bb8275ac9683", size = 4294089, upload-time = "2025-10-15T23:17:48.269Z" }, - { url = "https://files.pythonhosted.org/packages/26/42/fa8389d4478368743e24e61eea78846a0006caffaf72ea24a15159215a14/cryptography-46.0.3-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:15ab9b093e8f09daab0f2159bb7e47532596075139dd74365da52ecc9cb46c5d", size = 4440029, upload-time = "2025-10-15T23:17:49.837Z" }, - { url = "https://files.pythonhosted.org/packages/5f/eb/f483db0ec5ac040824f269e93dd2bd8a21ecd1027e77ad7bdf6914f2fd80/cryptography-46.0.3-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:46acf53b40ea38f9c6c229599a4a13f0d46a6c3fa9ef19fc1a124d62e338dfa0", size = 4297222, upload-time = "2025-10-15T23:17:51.357Z" }, - { url = "https://files.pythonhosted.org/packages/fd/cf/da9502c4e1912cb1da3807ea3618a6829bee8207456fbbeebc361ec38ba3/cryptography-46.0.3-cp38-abi3-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:10ca84c4668d066a9878890047f03546f3ae0a6b8b39b697457b7757aaf18dbc", size = 4012280, upload-time = "2025-10-15T23:17:52.964Z" }, - { url = "https://files.pythonhosted.org/packages/6b/8f/9adb86b93330e0df8b3dcf03eae67c33ba89958fc2e03862ef1ac2b42465/cryptography-46.0.3-cp38-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:36e627112085bb3b81b19fed209c05ce2a52ee8b15d161b7c643a7d5a88491f3", size = 4978958, upload-time = "2025-10-15T23:17:54.965Z" }, - { url = "https://files.pythonhosted.org/packages/d1/a0/5fa77988289c34bdb9f913f5606ecc9ada1adb5ae870bd0d1054a7021cc4/cryptography-46.0.3-cp38-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:1000713389b75c449a6e979ffc7dcc8ac90b437048766cef052d4d30b8220971", size = 4473714, upload-time = "2025-10-15T23:17:56.754Z" }, - { url = "https://files.pythonhosted.org/packages/14/e5/fc82d72a58d41c393697aa18c9abe5ae1214ff6f2a5c18ac470f92777895/cryptography-46.0.3-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:b02cf04496f6576afffef5ddd04a0cb7d49cf6be16a9059d793a30b035f6b6ac", size = 4296970, upload-time = "2025-10-15T23:17:58.588Z" }, - { url = "https://files.pythonhosted.org/packages/78/06/5663ed35438d0b09056973994f1aec467492b33bd31da36e468b01ec1097/cryptography-46.0.3-cp38-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:71e842ec9bc7abf543b47cf86b9a743baa95f4677d22baa4c7d5c69e49e9bc04", size = 4940236, upload-time = "2025-10-15T23:18:00.897Z" }, - { url = "https://files.pythonhosted.org/packages/fc/59/873633f3f2dcd8a053b8dd1d38f783043b5fce589c0f6988bf55ef57e43e/cryptography-46.0.3-cp38-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:402b58fc32614f00980b66d6e56a5b4118e6cb362ae8f3fda141ba4689bd4506", size = 4472642, upload-time = "2025-10-15T23:18:02.749Z" }, - { url = "https://files.pythonhosted.org/packages/3d/39/8e71f3930e40f6877737d6f69248cf74d4e34b886a3967d32f919cc50d3b/cryptography-46.0.3-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:ef639cb3372f69ec44915fafcd6698b6cc78fbe0c2ea41be867f6ed612811963", size = 4423126, upload-time = "2025-10-15T23:18:04.85Z" }, - { url = "https://files.pythonhosted.org/packages/cd/c7/f65027c2810e14c3e7268353b1681932b87e5a48e65505d8cc17c99e36ae/cryptography-46.0.3-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:3b51b8ca4f1c6453d8829e1eb7299499ca7f313900dd4d89a24b8b87c0a780d4", size = 4686573, upload-time = "2025-10-15T23:18:06.908Z" }, - { url = "https://files.pythonhosted.org/packages/0a/6e/1c8331ddf91ca4730ab3086a0f1be19c65510a33b5a441cb334e7a2d2560/cryptography-46.0.3-cp38-abi3-win32.whl", hash = "sha256:6276eb85ef938dc035d59b87c8a7dc559a232f954962520137529d77b18ff1df", size = 3036695, upload-time = "2025-10-15T23:18:08.672Z" }, - { url = "https://files.pythonhosted.org/packages/90/45/b0d691df20633eff80955a0fc7695ff9051ffce8b69741444bd9ed7bd0db/cryptography-46.0.3-cp38-abi3-win_amd64.whl", hash = "sha256:416260257577718c05135c55958b674000baef9a1c7d9e8f306ec60d71db850f", size = 3501720, upload-time = "2025-10-15T23:18:10.632Z" }, - { url = "https://files.pythonhosted.org/packages/e8/cb/2da4cc83f5edb9c3257d09e1e7ab7b23f049c7962cae8d842bbef0a9cec9/cryptography-46.0.3-cp38-abi3-win_arm64.whl", hash = "sha256:d89c3468de4cdc4f08a57e214384d0471911a3830fcdaf7a8cc587e42a866372", size = 2918740, upload-time = "2025-10-15T23:18:12.277Z" }, -] - [[package]] name = "cycler" version = "0.12.1" @@ -624,45 +415,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321, upload-time = "2023-10-07T05:32:16.783Z" }, ] -[[package]] -name = "cyclopts" -version = "4.4.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "attrs" }, - { name = "docstring-parser" }, - { name = "rich" }, - { name = "rich-rst" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/43/c4/60b6068e703c78656d07b249919754f8f60e9e7da3325560574ee27b4e39/cyclopts-4.4.4.tar.gz", hash = "sha256:f30c591c971d974ab4f223e099f881668daed72de713713c984ca41479d393dd", size = 160046, upload-time = "2026-01-05T03:40:18.438Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/20/5b/0eceb9a5990de9025733a0d212ca43649ba9facd58b8552b6bf93c11439d/cyclopts-4.4.4-py3-none-any.whl", hash = "sha256:316f798fe2f2a30cb70e7140cfde2a46617bfbb575d31bbfdc0b2410a447bd83", size = 197398, upload-time = "2026-01-05T03:40:17.141Z" }, -] - -[[package]] -name = "db-dtypes" -version = "1.5.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy" }, - { name = "packaging" }, - { name = "pandas" }, - { name = "pyarrow" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/27/3d/8c021df7796bde3bc4a9f8129865676a45d95ba8c5fd14b1de82997ff0b4/db_dtypes-1.5.0.tar.gz", hash = "sha256:ad9e94243f53e104bc77dbf9ae44b580d83a770d3694483aba59c9767966daa5", size = 34800, upload-time = "2025-12-15T21:47:21.874Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/01/80/171c7c5b78e60ab25d6f11e3d38675fe7ef843ddc79a7fd26916d3a6ca05/db_dtypes-1.5.0-py3-none-any.whl", hash = "sha256:abdbb2e4eb965800ed6f98af0c5c1cafff9063ace09114be2d26a7f046be2c8a", size = 18267, upload-time = "2025-12-15T21:47:21.026Z" }, -] - -[[package]] -name = "diskcache" -version = "5.6.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/3f/21/1c1ffc1a039ddcc459db43cc108658f32c57d271d7289a2794e401d0fdb6/diskcache-5.6.3.tar.gz", hash = "sha256:2c3a3fa2743d8535d832ec61c2054a1641f41775aa7c556758a109941e33e4fc", size = 67916, upload-time = "2023-08-31T06:12:00.316Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3f/27/4570e78fc0bf5ea0ca45eb1de3818a23787af9b390c0b0a0033a1b8236f9/diskcache-5.6.3-py3-none-any.whl", hash = "sha256:5e31b2d5fbad117cc363ebaf6b689474db18a1f6438bc82358b024abd4c2ca19", size = 45550, upload-time = "2023-08-31T06:11:58.822Z" }, -] - [[package]] name = "distlib" version = "0.4.0" @@ -672,141 +424,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/33/6b/e0547afaf41bf2c42e52430072fa5658766e3d65bd4b03a563d1b6336f57/distlib-0.4.0-py2.py3-none-any.whl", hash = "sha256:9659f7d87e46584a30b5780e43ac7a2143098441670ff0a49d5f9034c54a6c16", size = 469047, upload-time = "2025-07-17T16:51:58.613Z" }, ] -[[package]] -name = "dnspython" -version = "2.8.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/8c/8b/57666417c0f90f08bcafa776861060426765fdb422eb10212086fb811d26/dnspython-2.8.0.tar.gz", hash = "sha256:181d3c6996452cb1189c4046c61599b84a5a86e099562ffde77d26984ff26d0f", size = 368251, upload-time = "2025-09-07T18:58:00.022Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ba/5a/18ad964b0086c6e62e2e7500f7edc89e3faa45033c71c1893d34eed2b2de/dnspython-2.8.0-py3-none-any.whl", hash = "sha256:01d9bbc4a2d76bf0db7c1f729812ded6d912bd318d3b1cf81d30c0f845dbf3af", size = 331094, upload-time = "2025-09-07T18:57:58.071Z" }, -] - -[[package]] -name = "docstring-parser" -version = "0.17.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/b2/9d/c3b43da9515bd270df0f80548d9944e389870713cc1fe2b8fb35fe2bcefd/docstring_parser-0.17.0.tar.gz", hash = "sha256:583de4a309722b3315439bb31d64ba3eebada841f2e2cee23b99df001434c912", size = 27442, upload-time = "2025-07-21T07:35:01.868Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/55/e2/2537ebcff11c1ee1ff17d8d0b6f4db75873e3b0fb32c2d4a2ee31ecb310a/docstring_parser-0.17.0-py3-none-any.whl", hash = "sha256:cf2569abd23dce8099b300f9b4fa8191e9582dda731fd533daf54c4551658708", size = 36896, upload-time = "2025-07-21T07:35:00.684Z" }, -] - -[[package]] -name = "docutils" -version = "0.22.4" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ae/b6/03bb70946330e88ffec97aefd3ea75ba575cb2e762061e0e62a213befee8/docutils-0.22.4.tar.gz", hash = "sha256:4db53b1fde9abecbb74d91230d32ab626d94f6badfc575d6db9194a49df29968", size = 2291750, upload-time = "2025-12-18T19:00:26.443Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/02/10/5da547df7a391dcde17f59520a231527b8571e6f46fc8efb02ccb370ab12/docutils-0.22.4-py3-none-any.whl", hash = "sha256:d0013f540772d1420576855455d050a2180186c91c15779301ac2ccb3eeb68de", size = 633196, upload-time = "2025-12-18T19:00:18.077Z" }, -] - -[[package]] -name = "duckdb" -version = "1.4.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/81/99/ac6c105118751cc3ccae980b12e44847273f3402e647ec3197aff2251e23/duckdb-1.4.2.tar.gz", hash = "sha256:df81acee3b15ecb2c72eb8f8579fb5922f6f56c71f5c8892ea3bc6fab39aa2c4", size = 18469786, upload-time = "2025-11-12T13:18:04.203Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c9/29/2f68c57e7c4242fedbf4b3fdc24fce2ffcf60640c936621d8a645593a161/duckdb-1.4.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:9356fe17af2711e0a5ace4b20a0373e03163545fd7516e0c3c40428f44597052", size = 29015814, upload-time = "2025-11-12T13:16:59.329Z" }, - { url = "https://files.pythonhosted.org/packages/34/b7/030cc278a4ae788800a833b2901b9a7da7a6993121053c4155c359328531/duckdb-1.4.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:946a8374c0252db3fa41165ab9952b48adc8de06561a6b5fd62025ac700e492f", size = 15403892, upload-time = "2025-11-12T13:17:02.141Z" }, - { url = "https://files.pythonhosted.org/packages/f7/a2/67f4798a7a29bd0813f8a1e94a83e857e57f5d1ba14cf3edc5551aad0095/duckdb-1.4.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:389fa9abe4ca37d091332a2f8c3ebd713f18e87dc4cb5e8efd3e5aa8ddf8885f", size = 13733622, upload-time = "2025-11-12T13:17:04.502Z" }, - { url = "https://files.pythonhosted.org/packages/6e/ac/d0d0e3feae9663334b2336f15785d280b54a56c3ffa10334e20a51a87ecd/duckdb-1.4.2-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7be8c0c40f2264b91500b89c688f743e1c7764966e988f680b1f19416b00052e", size = 18470220, upload-time = "2025-11-12T13:17:07.049Z" }, - { url = "https://files.pythonhosted.org/packages/a5/52/7570a50430cbffc8bd702443ac28a446b0fa4f77747a3821d4b37a852b15/duckdb-1.4.2-cp312-cp312-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c6a21732dd52a76f1e61484c06d65800b18f57fe29e8102a7466c201a2221604", size = 20481138, upload-time = "2025-11-12T13:17:09.459Z" }, - { url = "https://files.pythonhosted.org/packages/95/5e/be05f46a290ea27630c112ff9e01fd01f585e599967fc52fe2edc7bc2039/duckdb-1.4.2-cp312-cp312-win_amd64.whl", hash = "sha256:769440f4507c20542ae2e5b87f6c6c6d3f148c0aa8f912528f6c97e9aedf6a21", size = 12330737, upload-time = "2025-11-12T13:17:12.02Z" }, - { url = "https://files.pythonhosted.org/packages/70/c4/5054dbe79cf570b0c97db0c2eba7eb541cc561037360479059a3b57e4a32/duckdb-1.4.2-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:de646227fc2c53101ac84e86e444e7561aa077387aca8b37052f3803ee690a17", size = 29015784, upload-time = "2025-11-12T13:17:14.409Z" }, - { url = "https://files.pythonhosted.org/packages/2c/b8/97f4f07d9459f5d262751cccfb2f4256debb8fe5ca92370cebe21aab1ee2/duckdb-1.4.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:f1fac31babda2045d4cdefe6d0fd2ebdd8d4c2a333fbcc11607cfeaec202d18d", size = 15403788, upload-time = "2025-11-12T13:17:16.864Z" }, - { url = "https://files.pythonhosted.org/packages/a4/ea/112f33ace03682bafd4aaf0a3336da689b9834663e7032b3d678fd2902c9/duckdb-1.4.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:43ac632f40ab1aede9b4ce3c09ea043f26f3db97b83c07c632c84ebd7f7c0f4a", size = 13733603, upload-time = "2025-11-12T13:17:20.884Z" }, - { url = "https://files.pythonhosted.org/packages/34/83/8d6f845a9a946e8b47b6253b9edb084c45670763e815feed6cfefc957e89/duckdb-1.4.2-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:77db030b48321bf785767b7b1800bf657dd2584f6df0a77e05201ecd22017da2", size = 18473725, upload-time = "2025-11-12T13:17:23.074Z" }, - { url = "https://files.pythonhosted.org/packages/82/29/153d1b4fc14c68e6766d7712d35a7ab6272a801c52160126ac7df681f758/duckdb-1.4.2-cp313-cp313-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a456adbc3459c9dcd99052fad20bd5f8ef642be5b04d09590376b2eb3eb84f5c", size = 20481971, upload-time = "2025-11-12T13:17:26.703Z" }, - { url = "https://files.pythonhosted.org/packages/58/b7/8d3a58b5ebfb9e79ed4030a0f2fbd7e404c52602e977b1e7ab51651816c7/duckdb-1.4.2-cp313-cp313-win_amd64.whl", hash = "sha256:2f7c61617d2b1da3da5d7e215be616ad45aa3221c4b9e2c4d1c28ed09bc3c1c4", size = 12330535, upload-time = "2025-11-12T13:17:29.175Z" }, - { url = "https://files.pythonhosted.org/packages/25/46/0f316e4d0d6bada350b9da06691a2537c329c8948c78e8b5e0c4874bc5e2/duckdb-1.4.2-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:422be8c6bdc98366c97f464b204b81b892bf962abeae6b0184104b8233da4f19", size = 29028616, upload-time = "2025-11-12T13:17:31.599Z" }, - { url = "https://files.pythonhosted.org/packages/82/ab/e04a8f97865251b544aee9501088d4f0cb8e8b37339bd465c0d33857d411/duckdb-1.4.2-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:459b1855bd06a226a2838da4f14c8863fd87a62e63d414a7f7f682a7c616511a", size = 15410382, upload-time = "2025-11-12T13:17:34.14Z" }, - { url = "https://files.pythonhosted.org/packages/47/ec/b8229517c2f9fe88a38bb1a172a2da4d0ff34996d319d74554fda80b6358/duckdb-1.4.2-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:20c45b4ead1ea4d23a1be1cd4f1dfc635e58b55f0dd11e38781369be6c549903", size = 13737588, upload-time = "2025-11-12T13:17:36.515Z" }, - { url = "https://files.pythonhosted.org/packages/f2/9a/63d26da9011890a5b893e0c21845c0c0b43c634bf263af3bbca64be0db76/duckdb-1.4.2-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2e552451054534970dc999e69ca5ae5c606458548c43fb66d772117760485096", size = 18477886, upload-time = "2025-11-12T13:17:39.136Z" }, - { url = "https://files.pythonhosted.org/packages/23/35/b1fae4c5245697837f6f63e407fa81e7ccc7948f6ef2b124cd38736f4d1d/duckdb-1.4.2-cp314-cp314-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:128c97dab574a438d7c8d020670b21c68792267d88e65a7773667b556541fa9b", size = 20483292, upload-time = "2025-11-12T13:17:41.501Z" }, - { url = "https://files.pythonhosted.org/packages/25/5e/6f5ebaabc12c6db62f471f86b5c9c8debd57f11aa1b2acbbcc4c68683238/duckdb-1.4.2-cp314-cp314-win_amd64.whl", hash = "sha256:dfcc56a83420c0dec0b83e97a6b33addac1b7554b8828894f9d203955591218c", size = 12830520, upload-time = "2025-11-12T13:17:43.93Z" }, -] - -[[package]] -name = "ecdsa" -version = "0.19.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "six" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/c0/1f/924e3caae75f471eae4b26bd13b698f6af2c44279f67af317439c2f4c46a/ecdsa-0.19.1.tar.gz", hash = "sha256:478cba7b62555866fcb3bb3fe985e06decbdb68ef55713c4e5ab98c57d508e61", size = 201793, upload-time = "2025-03-13T11:52:43.25Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/cb/a3/460c57f094a4a165c84a1341c373b0a4f5ec6ac244b998d5021aade89b77/ecdsa-0.19.1-py2.py3-none-any.whl", hash = "sha256:30638e27cf77b7e15c4c4cc1973720149e1033827cfd00661ca5c8cc0cdb24c3", size = 150607, upload-time = "2025-03-13T11:52:41.757Z" }, -] - -[[package]] -name = "email-validator" -version = "2.3.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "dnspython" }, - { name = "idna" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/f5/22/900cb125c76b7aaa450ce02fd727f452243f2e91a61af068b40adba60ea9/email_validator-2.3.0.tar.gz", hash = "sha256:9fc05c37f2f6cf439ff414f8fc46d917929974a82244c20eb10231ba60c54426", size = 51238, upload-time = "2025-08-26T13:09:06.831Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/de/15/545e2b6cf2e3be84bc1ed85613edd75b8aea69807a71c26f4ca6a9258e82/email_validator-2.3.0-py3-none-any.whl", hash = "sha256:80f13f623413e6b197ae73bb10bf4eb0908faf509ad8362c5edeb0be7fd450b4", size = 35604, upload-time = "2025-08-26T13:09:05.858Z" }, -] - -[[package]] -name = "exceptiongroup" -version = "1.3.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/8a/0e/97c33bf5009bdbac74fd2beace167cab3f978feb69cc36f1ef79360d6c4e/exceptiongroup-1.3.1-py3-none-any.whl", hash = "sha256:a7a39a3bd276781e98394987d3a5701d0c4edffb633bb7a5144577f82c773598", size = 16740, upload-time = "2025-11-21T23:01:53.443Z" }, -] - -[[package]] -name = "fakeredis" -version = "2.33.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "redis" }, - { name = "sortedcontainers" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/5f/f9/57464119936414d60697fcbd32f38909bb5688b616ae13de6e98384433e0/fakeredis-2.33.0.tar.gz", hash = "sha256:d7bc9a69d21df108a6451bbffee23b3eba432c21a654afc7ff2d295428ec5770", size = 175187, upload-time = "2025-12-16T19:45:52.269Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/6e/78/a850fed8aeef96d4a99043c90b818b2ed5419cd5b24a4049fd7cfb9f1471/fakeredis-2.33.0-py3-none-any.whl", hash = "sha256:de535f3f9ccde1c56672ab2fdd6a8efbc4f2619fc2f1acc87b8737177d71c965", size = 119605, upload-time = "2025-12-16T19:45:51.08Z" }, -] - -[package.optional-dependencies] -lua = [ - { name = "lupa" }, -] - -[[package]] -name = "fastmcp" -version = "2.14.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "authlib" }, - { name = "cyclopts" }, - { name = "exceptiongroup" }, - { name = "httpx" }, - { name = "jsonschema-path" }, - { name = "mcp" }, - { name = "openapi-pydantic" }, - { name = "platformdirs" }, - { name = "py-key-value-aio", extra = ["disk", "keyring", "memory"] }, - { name = "pydantic", extra = ["email"] }, - { name = "pydocket" }, - { name = "pyperclip" }, - { name = "python-dotenv" }, - { name = "rich" }, - { name = "uvicorn" }, - { name = "websockets" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/d1/1e/e3528227688c248283f6d86869b1e900563ffc223eff00f4f923d2750365/fastmcp-2.14.2.tar.gz", hash = "sha256:bd23d1b808b6f446444f10114dac468b11bfb9153ed78628f5619763d0cf573e", size = 8272966, upload-time = "2025-12-31T15:26:13.433Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/0d/67/8456d39484fcb7afd0defed21918e773ed59a98b39e5b633328527c88367/fastmcp-2.14.2-py3-none-any.whl", hash = "sha256:e33cd622e1ebd5110af6a981804525b6cd41072e3c7d68268ed69ef3be651aca", size = 413279, upload-time = "2025-12-31T15:26:11.178Z" }, -] - [[package]] name = "filelock" version = "3.20.0" @@ -857,24 +474,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0c/14/634f7daea5ffe6a5f7a0322ba8e1a0e23c9257b80aa91458107896d1dfc7/fonttools-4.61.0-py3-none-any.whl", hash = "sha256:276f14c560e6f98d24ef7f5f44438e55ff5a67f78fa85236b218462c9f5d0635", size = 1144485, upload-time = "2025-11-28T17:05:47.573Z" }, ] -[[package]] -name = "formulaic" -version = "1.2.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "interface-meta" }, - { name = "narwhals" }, - { name = "numpy" }, - { name = "pandas" }, - { name = "scipy" }, - { name = "typing-extensions" }, - { name = "wrapt" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/6a/8b/8038d2af289a5cc194fa0a255fe964a1a04e0e6ca4426aed8841a4b571e6/formulaic-1.2.1.tar.gz", hash = "sha256:dc79476baa2d811b35798893eb2f2c1e51edee8d7a9c1429b400e56f4e0beccc", size = 655266, upload-time = "2025-09-21T05:27:31.488Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1a/9d/c2c8b51b32f829a16fe042db30ad1dcef7947bf1dcf77c2cfd7b6f37b83a/formulaic-1.2.1-py3-none-any.whl", hash = "sha256:661d6d2467aa961b9afb3a1e2a187494239793c63eb729e422d1307afa98b43b", size = 117290, upload-time = "2025-09-21T05:27:30.025Z" }, -] - [[package]] name = "frozenlist" version = "1.8.0" @@ -1002,259 +601,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/01/61/d4b89fec821f72385526e1b9d9a3a0385dda4a72b206d28049e2c7cd39b8/gitpython-3.1.45-py3-none-any.whl", hash = "sha256:8908cb2e02fb3b93b7eb0f2827125cb699869470432cc885f019b8fd0fccff77", size = 208168, upload-time = "2025-07-24T03:45:52.517Z" }, ] -[[package]] -name = "google-api-core" -version = "2.29.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "google-auth" }, - { name = "googleapis-common-protos" }, - { name = "proto-plus" }, - { name = "protobuf" }, - { name = "requests" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/0d/10/05572d33273292bac49c2d1785925f7bc3ff2fe50e3044cf1062c1dde32e/google_api_core-2.29.0.tar.gz", hash = "sha256:84181be0f8e6b04006df75ddfe728f24489f0af57c96a529ff7cf45bc28797f7", size = 177828, upload-time = "2026-01-08T22:21:39.269Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/77/b6/85c4d21067220b9a78cfb81f516f9725ea6befc1544ec9bd2c1acd97c324/google_api_core-2.29.0-py3-none-any.whl", hash = "sha256:d30bc60980daa36e314b5d5a3e5958b0200cb44ca8fa1be2b614e932b75a3ea9", size = 173906, upload-time = "2026-01-08T22:21:36.093Z" }, -] - -[package.optional-dependencies] -grpc = [ - { name = "grpcio" }, - { name = "grpcio-status" }, -] - -[[package]] -name = "google-auth" -version = "2.47.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pyasn1-modules" }, - { name = "rsa" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/60/3c/ec64b9a275ca22fa1cd3b6e77fefcf837b0732c890aa32d2bd21313d9b33/google_auth-2.47.0.tar.gz", hash = "sha256:833229070a9dfee1a353ae9877dcd2dec069a8281a4e72e72f77d4a70ff945da", size = 323719, upload-time = "2026-01-06T21:55:31.045Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/db/18/79e9008530b79527e0d5f79e7eef08d3b179b7f851cfd3a2f27822fbdfa9/google_auth-2.47.0-py3-none-any.whl", hash = "sha256:c516d68336bfde7cf0da26aab674a36fedcf04b37ac4edd59c597178760c3498", size = 234867, upload-time = "2026-01-06T21:55:28.6Z" }, -] - -[[package]] -name = "google-cloud-bigquery" -version = "3.40.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "google-api-core", extra = ["grpc"] }, - { name = "google-auth" }, - { name = "google-cloud-core" }, - { name = "google-resumable-media" }, - { name = "packaging" }, - { name = "python-dateutil" }, - { name = "requests" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/94/0a/62438ca138a095945468968696d9cca75a4cfd059e810402e70b0236d8ba/google_cloud_bigquery-3.40.0.tar.gz", hash = "sha256:b3ccb11caf0029f15b29569518f667553fe08f6f1459b959020c83fbbd8f2e68", size = 509287, upload-time = "2026-01-08T01:07:26.065Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/90/6a/90a04270dd60cc70259b73744f6e610ae9a158b21ab50fb695cca0056a3d/google_cloud_bigquery-3.40.0-py3-none-any.whl", hash = "sha256:0469bcf9e3dad3cab65b67cce98180c8c0aacf3253d47f0f8e976f299b49b5ab", size = 261335, upload-time = "2026-01-08T01:07:23.761Z" }, -] - -[[package]] -name = "google-cloud-core" -version = "2.5.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "google-api-core" }, - { name = "google-auth" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/a6/03/ef0bc99d0e0faf4fdbe67ac445e18cdaa74824fd93cd069e7bb6548cb52d/google_cloud_core-2.5.0.tar.gz", hash = "sha256:7c1b7ef5c92311717bd05301aa1a91ffbc565673d3b0b4163a52d8413a186963", size = 36027, upload-time = "2025-10-29T23:17:39.513Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/89/20/bfa472e327c8edee00f04beecc80baeddd2ab33ee0e86fd7654da49d45e9/google_cloud_core-2.5.0-py3-none-any.whl", hash = "sha256:67d977b41ae6c7211ee830c7912e41003ea8194bff15ae7d72fd6f51e57acabc", size = 29469, upload-time = "2025-10-29T23:17:38.548Z" }, -] - -[[package]] -name = "google-crc32c" -version = "1.8.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/03/41/4b9c02f99e4c5fb477122cd5437403b552873f014616ac1d19ac8221a58d/google_crc32c-1.8.0.tar.gz", hash = "sha256:a428e25fb7691024de47fecfbff7ff957214da51eddded0da0ae0e0f03a2cf79", size = 14192, upload-time = "2025-12-16T00:35:25.142Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e9/5f/7307325b1198b59324c0fa9807cafb551afb65e831699f2ce211ad5c8240/google_crc32c-1.8.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:4b8286b659c1335172e39563ab0a768b8015e88e08329fa5321f774275fc3113", size = 31300, upload-time = "2025-12-16T00:21:56.723Z" }, - { url = "https://files.pythonhosted.org/packages/21/8e/58c0d5d86e2220e6a37befe7e6a94dd2f6006044b1a33edf1ff6d9f7e319/google_crc32c-1.8.0-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:2a3dc3318507de089c5384cc74d54318401410f82aa65b2d9cdde9d297aca7cb", size = 30867, upload-time = "2025-12-16T00:38:31.302Z" }, - { url = "https://files.pythonhosted.org/packages/ce/a9/a780cc66f86335a6019f557a8aaca8fbb970728f0efd2430d15ff1beae0e/google_crc32c-1.8.0-cp312-cp312-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:14f87e04d613dfa218d6135e81b78272c3b904e2a7053b841481b38a7d901411", size = 33364, upload-time = "2025-12-16T00:40:22.96Z" }, - { url = "https://files.pythonhosted.org/packages/21/3f/3457ea803db0198c9aaca2dd373750972ce28a26f00544b6b85088811939/google_crc32c-1.8.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb5c869c2923d56cb0c8e6bcdd73c009c36ae39b652dbe46a05eb4ef0ad01454", size = 33740, upload-time = "2025-12-16T00:40:23.96Z" }, - { url = "https://files.pythonhosted.org/packages/df/c0/87c2073e0c72515bb8733d4eef7b21548e8d189f094b5dad20b0ecaf64f6/google_crc32c-1.8.0-cp312-cp312-win_amd64.whl", hash = "sha256:3cc0c8912038065eafa603b238abf252e204accab2a704c63b9e14837a854962", size = 34437, upload-time = "2025-12-16T00:35:21.395Z" }, - { url = "https://files.pythonhosted.org/packages/d1/db/000f15b41724589b0e7bc24bc7a8967898d8d3bc8caf64c513d91ef1f6c0/google_crc32c-1.8.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:3ebb04528e83b2634857f43f9bb8ef5b2bbe7f10f140daeb01b58f972d04736b", size = 31297, upload-time = "2025-12-16T00:23:20.709Z" }, - { url = "https://files.pythonhosted.org/packages/d7/0d/8ebed0c39c53a7e838e2a486da8abb0e52de135f1b376ae2f0b160eb4c1a/google_crc32c-1.8.0-cp313-cp313-macosx_12_0_x86_64.whl", hash = "sha256:450dc98429d3e33ed2926fc99ee81001928d63460f8538f21a5d6060912a8e27", size = 30867, upload-time = "2025-12-16T00:43:14.628Z" }, - { url = "https://files.pythonhosted.org/packages/ce/42/b468aec74a0354b34c8cbf748db20d6e350a68a2b0912e128cabee49806c/google_crc32c-1.8.0-cp313-cp313-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:3b9776774b24ba76831609ffbabce8cdf6fa2bd5e9df37b594221c7e333a81fa", size = 33344, upload-time = "2025-12-16T00:40:24.742Z" }, - { url = "https://files.pythonhosted.org/packages/1c/e8/b33784d6fc77fb5062a8a7854e43e1e618b87d5ddf610a88025e4de6226e/google_crc32c-1.8.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:89c17d53d75562edfff86679244830599ee0a48efc216200691de8b02ab6b2b8", size = 33694, upload-time = "2025-12-16T00:40:25.505Z" }, - { url = "https://files.pythonhosted.org/packages/92/b1/d3cbd4d988afb3d8e4db94ca953df429ed6db7282ed0e700d25e6c7bfc8d/google_crc32c-1.8.0-cp313-cp313-win_amd64.whl", hash = "sha256:57a50a9035b75643996fbf224d6661e386c7162d1dfdab9bc4ca790947d1007f", size = 34435, upload-time = "2025-12-16T00:35:22.107Z" }, - { url = "https://files.pythonhosted.org/packages/21/88/8ecf3c2b864a490b9e7010c84fd203ec8cf3b280651106a3a74dd1b0ca72/google_crc32c-1.8.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:e6584b12cb06796d285d09e33f63309a09368b9d806a551d8036a4207ea43697", size = 31301, upload-time = "2025-12-16T00:24:48.527Z" }, - { url = "https://files.pythonhosted.org/packages/36/c6/f7ff6c11f5ca215d9f43d3629163727a272eabc356e5c9b2853df2bfe965/google_crc32c-1.8.0-cp314-cp314-macosx_12_0_x86_64.whl", hash = "sha256:f4b51844ef67d6cf2e9425983274da75f18b1597bb2c998e1c0a0e8d46f8f651", size = 30868, upload-time = "2025-12-16T00:48:12.163Z" }, - { url = "https://files.pythonhosted.org/packages/56/15/c25671c7aad70f8179d858c55a6ae8404902abe0cdcf32a29d581792b491/google_crc32c-1.8.0-cp314-cp314-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:b0d1a7afc6e8e4635564ba8aa5c0548e3173e41b6384d7711a9123165f582de2", size = 33381, upload-time = "2025-12-16T00:40:26.268Z" }, - { url = "https://files.pythonhosted.org/packages/42/fa/f50f51260d7b0ef5d4898af122d8a7ec5a84e2984f676f746445f783705f/google_crc32c-1.8.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8b3f68782f3cbd1bce027e48768293072813469af6a61a86f6bb4977a4380f21", size = 33734, upload-time = "2025-12-16T00:40:27.028Z" }, - { url = "https://files.pythonhosted.org/packages/08/a5/7b059810934a09fb3ccb657e0843813c1fee1183d3bc2c8041800374aa2c/google_crc32c-1.8.0-cp314-cp314-win_amd64.whl", hash = "sha256:d511b3153e7011a27ab6ee6bb3a5404a55b994dc1a7322c0b87b29606d9790e2", size = 34878, upload-time = "2025-12-16T00:35:23.142Z" }, -] - -[[package]] -name = "google-resumable-media" -version = "2.8.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "google-crc32c" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/64/d7/520b62a35b23038ff005e334dba3ffc75fcf583bee26723f1fd8fd4b6919/google_resumable_media-2.8.0.tar.gz", hash = "sha256:f1157ed8b46994d60a1bc432544db62352043113684d4e030ee02e77ebe9a1ae", size = 2163265, upload-time = "2025-11-17T15:38:06.659Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1f/0b/93afde9cfe012260e9fe1522f35c9b72d6ee222f316586b1f23ecf44d518/google_resumable_media-2.8.0-py3-none-any.whl", hash = "sha256:dd14a116af303845a8d932ddae161a26e86cc229645bc98b39f026f9b1717582", size = 81340, upload-time = "2025-11-17T15:38:05.594Z" }, -] - -[[package]] -name = "googleapis-common-protos" -version = "1.72.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "protobuf" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/e5/7b/adfd75544c415c487b33061fe7ae526165241c1ea133f9a9125a56b39fd8/googleapis_common_protos-1.72.0.tar.gz", hash = "sha256:e55a601c1b32b52d7a3e65f43563e2aa61bcd737998ee672ac9b951cd49319f5", size = 147433, upload-time = "2025-11-06T18:29:24.087Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c4/ab/09169d5a4612a5f92490806649ac8d41e3ec9129c636754575b3553f4ea4/googleapis_common_protos-1.72.0-py3-none-any.whl", hash = "sha256:4299c5a82d5ae1a9702ada957347726b167f9f8d1fc352477702a1e851ff4038", size = 297515, upload-time = "2025-11-06T18:29:13.14Z" }, -] - -[[package]] -name = "greenlet" -version = "3.3.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/c7/e5/40dbda2736893e3e53d25838e0f19a2b417dfc122b9989c91918db30b5d3/greenlet-3.3.0.tar.gz", hash = "sha256:a82bb225a4e9e4d653dd2fb7b8b2d36e4fb25bc0165422a11e48b88e9e6f78fb", size = 190651, upload-time = "2025-12-04T14:49:44.05Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f8/0a/a3871375c7b9727edaeeea994bfff7c63ff7804c9829c19309ba2e058807/greenlet-3.3.0-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:b01548f6e0b9e9784a2c99c5651e5dc89ffcbe870bc5fb2e5ef864e9cc6b5dcb", size = 276379, upload-time = "2025-12-04T14:23:30.498Z" }, - { url = "https://files.pythonhosted.org/packages/43/ab/7ebfe34dce8b87be0d11dae91acbf76f7b8246bf9d6b319c741f99fa59c6/greenlet-3.3.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:349345b770dc88f81506c6861d22a6ccd422207829d2c854ae2af8025af303e3", size = 597294, upload-time = "2025-12-04T14:50:06.847Z" }, - { url = "https://files.pythonhosted.org/packages/a4/39/f1c8da50024feecd0793dbd5e08f526809b8ab5609224a2da40aad3a7641/greenlet-3.3.0-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:e8e18ed6995e9e2c0b4ed264d2cf89260ab3ac7e13555b8032b25a74c6d18655", size = 607742, upload-time = "2025-12-04T14:57:42.349Z" }, - { url = "https://files.pythonhosted.org/packages/77/cb/43692bcd5f7a0da6ec0ec6d58ee7cddb606d055ce94a62ac9b1aa481e969/greenlet-3.3.0-cp312-cp312-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:c024b1e5696626890038e34f76140ed1daf858e37496d33f2af57f06189e70d7", size = 622297, upload-time = "2025-12-04T15:07:13.552Z" }, - { url = "https://files.pythonhosted.org/packages/75/b0/6bde0b1011a60782108c01de5913c588cf51a839174538d266de15e4bf4d/greenlet-3.3.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:047ab3df20ede6a57c35c14bf5200fcf04039d50f908270d3f9a7a82064f543b", size = 609885, upload-time = "2025-12-04T14:26:02.368Z" }, - { url = "https://files.pythonhosted.org/packages/49/0e/49b46ac39f931f59f987b7cd9f34bfec8ef81d2a1e6e00682f55be5de9f4/greenlet-3.3.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2d9ad37fc657b1102ec880e637cccf20191581f75c64087a549e66c57e1ceb53", size = 1567424, upload-time = "2025-12-04T15:04:23.757Z" }, - { url = "https://files.pythonhosted.org/packages/05/f5/49a9ac2dff7f10091935def9165c90236d8f175afb27cbed38fb1d61ab6b/greenlet-3.3.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:83cd0e36932e0e7f36a64b732a6f60c2fc2df28c351bae79fbaf4f8092fe7614", size = 1636017, upload-time = "2025-12-04T14:27:29.688Z" }, - { url = "https://files.pythonhosted.org/packages/6c/79/3912a94cf27ec503e51ba493692d6db1e3cd8ac7ac52b0b47c8e33d7f4f9/greenlet-3.3.0-cp312-cp312-win_amd64.whl", hash = "sha256:a7a34b13d43a6b78abf828a6d0e87d3385680eaf830cd60d20d52f249faabf39", size = 301964, upload-time = "2025-12-04T14:36:58.316Z" }, - { url = "https://files.pythonhosted.org/packages/02/2f/28592176381b9ab2cafa12829ba7b472d177f3acc35d8fbcf3673d966fff/greenlet-3.3.0-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:a1e41a81c7e2825822f4e068c48cb2196002362619e2d70b148f20a831c00739", size = 275140, upload-time = "2025-12-04T14:23:01.282Z" }, - { url = "https://files.pythonhosted.org/packages/2c/80/fbe937bf81e9fca98c981fe499e59a3f45df2a04da0baa5c2be0dca0d329/greenlet-3.3.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9f515a47d02da4d30caaa85b69474cec77b7929b2e936ff7fb853d42f4bf8808", size = 599219, upload-time = "2025-12-04T14:50:08.309Z" }, - { url = "https://files.pythonhosted.org/packages/c2/ff/7c985128f0514271b8268476af89aee6866df5eec04ac17dcfbc676213df/greenlet-3.3.0-cp313-cp313-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:7d2d9fd66bfadf230b385fdc90426fcd6eb64db54b40c495b72ac0feb5766c54", size = 610211, upload-time = "2025-12-04T14:57:43.968Z" }, - { url = "https://files.pythonhosted.org/packages/79/07/c47a82d881319ec18a4510bb30463ed6891f2ad2c1901ed5ec23d3de351f/greenlet-3.3.0-cp313-cp313-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:30a6e28487a790417d036088b3bcb3f3ac7d8babaa7d0139edbaddebf3af9492", size = 624311, upload-time = "2025-12-04T15:07:14.697Z" }, - { url = "https://files.pythonhosted.org/packages/fd/8e/424b8c6e78bd9837d14ff7df01a9829fc883ba2ab4ea787d4f848435f23f/greenlet-3.3.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:087ea5e004437321508a8d6f20efc4cfec5e3c30118e1417ea96ed1d93950527", size = 612833, upload-time = "2025-12-04T14:26:03.669Z" }, - { url = "https://files.pythonhosted.org/packages/b5/ba/56699ff9b7c76ca12f1cdc27a886d0f81f2189c3455ff9f65246780f713d/greenlet-3.3.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:ab97cf74045343f6c60a39913fa59710e4bd26a536ce7ab2397adf8b27e67c39", size = 1567256, upload-time = "2025-12-04T15:04:25.276Z" }, - { url = "https://files.pythonhosted.org/packages/1e/37/f31136132967982d698c71a281a8901daf1a8fbab935dce7c0cf15f942cc/greenlet-3.3.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5375d2e23184629112ca1ea89a53389dddbffcf417dad40125713d88eb5f96e8", size = 1636483, upload-time = "2025-12-04T14:27:30.804Z" }, - { url = "https://files.pythonhosted.org/packages/7e/71/ba21c3fb8c5dce83b8c01f458a42e99ffdb1963aeec08fff5a18588d8fd7/greenlet-3.3.0-cp313-cp313-win_amd64.whl", hash = "sha256:9ee1942ea19550094033c35d25d20726e4f1c40d59545815e1128ac58d416d38", size = 301833, upload-time = "2025-12-04T14:32:23.929Z" }, - { url = "https://files.pythonhosted.org/packages/d7/7c/f0a6d0ede2c7bf092d00bc83ad5bafb7e6ec9b4aab2fbdfa6f134dc73327/greenlet-3.3.0-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:60c2ef0f578afb3c8d92ea07ad327f9a062547137afe91f38408f08aacab667f", size = 275671, upload-time = "2025-12-04T14:23:05.267Z" }, - { url = "https://files.pythonhosted.org/packages/44/06/dac639ae1a50f5969d82d2e3dd9767d30d6dbdbab0e1a54010c8fe90263c/greenlet-3.3.0-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0a5d554d0712ba1de0a6c94c640f7aeba3f85b3a6e1f2899c11c2c0428da9365", size = 646360, upload-time = "2025-12-04T14:50:10.026Z" }, - { url = "https://files.pythonhosted.org/packages/e0/94/0fb76fe6c5369fba9bf98529ada6f4c3a1adf19e406a47332245ef0eb357/greenlet-3.3.0-cp314-cp314-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:3a898b1e9c5f7307ebbde4102908e6cbfcb9ea16284a3abe15cab996bee8b9b3", size = 658160, upload-time = "2025-12-04T14:57:45.41Z" }, - { url = "https://files.pythonhosted.org/packages/93/79/d2c70cae6e823fac36c3bbc9077962105052b7ef81db2f01ec3b9bf17e2b/greenlet-3.3.0-cp314-cp314-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:dcd2bdbd444ff340e8d6bdf54d2f206ccddbb3ccfdcd3c25bf4afaa7b8f0cf45", size = 671388, upload-time = "2025-12-04T15:07:15.789Z" }, - { url = "https://files.pythonhosted.org/packages/b8/14/bab308fc2c1b5228c3224ec2bf928ce2e4d21d8046c161e44a2012b5203e/greenlet-3.3.0-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5773edda4dc00e173820722711d043799d3adb4f01731f40619e07ea2750b955", size = 660166, upload-time = "2025-12-04T14:26:05.099Z" }, - { url = "https://files.pythonhosted.org/packages/4b/d2/91465d39164eaa0085177f61983d80ffe746c5a1860f009811d498e7259c/greenlet-3.3.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ac0549373982b36d5fd5d30beb8a7a33ee541ff98d2b502714a09f1169f31b55", size = 1615193, upload-time = "2025-12-04T15:04:27.041Z" }, - { url = "https://files.pythonhosted.org/packages/42/1b/83d110a37044b92423084d52d5d5a3b3a73cafb51b547e6d7366ff62eff1/greenlet-3.3.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:d198d2d977460358c3b3a4dc844f875d1adb33817f0613f663a656f463764ccc", size = 1683653, upload-time = "2025-12-04T14:27:32.366Z" }, - { url = "https://files.pythonhosted.org/packages/7c/9a/9030e6f9aa8fd7808e9c31ba4c38f87c4f8ec324ee67431d181fe396d705/greenlet-3.3.0-cp314-cp314-win_amd64.whl", hash = "sha256:73f51dd0e0bdb596fb0417e475fa3c5e32d4c83638296e560086b8d7da7c4170", size = 305387, upload-time = "2025-12-04T14:26:51.063Z" }, - { url = "https://files.pythonhosted.org/packages/a0/66/bd6317bc5932accf351fc19f177ffba53712a202f9df10587da8df257c7e/greenlet-3.3.0-cp314-cp314t-macosx_11_0_universal2.whl", hash = "sha256:d6ed6f85fae6cdfdb9ce04c9bf7a08d666cfcfb914e7d006f44f840b46741931", size = 282638, upload-time = "2025-12-04T14:25:20.941Z" }, - { url = "https://files.pythonhosted.org/packages/30/cf/cc81cb030b40e738d6e69502ccbd0dd1bced0588e958f9e757945de24404/greenlet-3.3.0-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d9125050fcf24554e69c4cacb086b87b3b55dc395a8b3ebe6487b045b2614388", size = 651145, upload-time = "2025-12-04T14:50:11.039Z" }, - { url = "https://files.pythonhosted.org/packages/9c/ea/1020037b5ecfe95ca7df8d8549959baceb8186031da83d5ecceff8b08cd2/greenlet-3.3.0-cp314-cp314t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:87e63ccfa13c0a0f6234ed0add552af24cc67dd886731f2261e46e241608bee3", size = 654236, upload-time = "2025-12-04T14:57:47.007Z" }, - { url = "https://files.pythonhosted.org/packages/69/cc/1e4bae2e45ca2fa55299f4e85854606a78ecc37fead20d69322f96000504/greenlet-3.3.0-cp314-cp314t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:2662433acbca297c9153a4023fe2161c8dcfdcc91f10433171cf7e7d94ba2221", size = 662506, upload-time = "2025-12-04T15:07:16.906Z" }, - { url = "https://files.pythonhosted.org/packages/57/b9/f8025d71a6085c441a7eaff0fd928bbb275a6633773667023d19179fe815/greenlet-3.3.0-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3c6e9b9c1527a78520357de498b0e709fb9e2f49c3a513afd5a249007261911b", size = 653783, upload-time = "2025-12-04T14:26:06.225Z" }, - { url = "https://files.pythonhosted.org/packages/f6/c7/876a8c7a7485d5d6b5c6821201d542ef28be645aa024cfe1145b35c120c1/greenlet-3.3.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:286d093f95ec98fdd92fcb955003b8a3d054b4e2cab3e2707a5039e7b50520fd", size = 1614857, upload-time = "2025-12-04T15:04:28.484Z" }, - { url = "https://files.pythonhosted.org/packages/4f/dc/041be1dff9f23dac5f48a43323cd0789cb798342011c19a248d9c9335536/greenlet-3.3.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:6c10513330af5b8ae16f023e8ddbfb486ab355d04467c4679c5cfe4659975dd9", size = 1676034, upload-time = "2025-12-04T14:27:33.531Z" }, -] - -[[package]] -name = "grpcio" -version = "1.76.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/b6/e0/318c1ce3ae5a17894d5791e87aea147587c9e702f24122cc7a5c8bbaeeb1/grpcio-1.76.0.tar.gz", hash = "sha256:7be78388d6da1a25c0d5ec506523db58b18be22d9c37d8d3a32c08be4987bd73", size = 12785182, upload-time = "2025-10-21T16:23:12.106Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/bf/05/8e29121994b8d959ffa0afd28996d452f291b48cfc0875619de0bde2c50c/grpcio-1.76.0-cp312-cp312-linux_armv7l.whl", hash = "sha256:81fd9652b37b36f16138611c7e884eb82e0cec137c40d3ef7c3f9b3ed00f6ed8", size = 5799718, upload-time = "2025-10-21T16:21:17.939Z" }, - { url = "https://files.pythonhosted.org/packages/d9/75/11d0e66b3cdf998c996489581bdad8900db79ebd83513e45c19548f1cba4/grpcio-1.76.0-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:04bbe1bfe3a68bbfd4e52402ab7d4eb59d72d02647ae2042204326cf4bbad280", size = 11825627, upload-time = "2025-10-21T16:21:20.466Z" }, - { url = "https://files.pythonhosted.org/packages/28/50/2f0aa0498bc188048f5d9504dcc5c2c24f2eb1a9337cd0fa09a61a2e75f0/grpcio-1.76.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d388087771c837cdb6515539f43b9d4bf0b0f23593a24054ac16f7a960be16f4", size = 6359167, upload-time = "2025-10-21T16:21:23.122Z" }, - { url = "https://files.pythonhosted.org/packages/66/e5/bbf0bb97d29ede1d59d6588af40018cfc345b17ce979b7b45424628dc8bb/grpcio-1.76.0-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:9f8f757bebaaea112c00dba718fc0d3260052ce714e25804a03f93f5d1c6cc11", size = 7044267, upload-time = "2025-10-21T16:21:25.995Z" }, - { url = "https://files.pythonhosted.org/packages/f5/86/f6ec2164f743d9609691115ae8ece098c76b894ebe4f7c94a655c6b03e98/grpcio-1.76.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:980a846182ce88c4f2f7e2c22c56aefd515daeb36149d1c897f83cf57999e0b6", size = 6573963, upload-time = "2025-10-21T16:21:28.631Z" }, - { url = "https://files.pythonhosted.org/packages/60/bc/8d9d0d8505feccfdf38a766d262c71e73639c165b311c9457208b56d92ae/grpcio-1.76.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:f92f88e6c033db65a5ae3d97905c8fea9c725b63e28d5a75cb73b49bda5024d8", size = 7164484, upload-time = "2025-10-21T16:21:30.837Z" }, - { url = "https://files.pythonhosted.org/packages/67/e6/5d6c2fc10b95edf6df9b8f19cf10a34263b7fd48493936fffd5085521292/grpcio-1.76.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:4baf3cbe2f0be3289eb68ac8ae771156971848bb8aaff60bad42005539431980", size = 8127777, upload-time = "2025-10-21T16:21:33.577Z" }, - { url = "https://files.pythonhosted.org/packages/3f/c8/dce8ff21c86abe025efe304d9e31fdb0deaaa3b502b6a78141080f206da0/grpcio-1.76.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:615ba64c208aaceb5ec83bfdce7728b80bfeb8be97562944836a7a0a9647d882", size = 7594014, upload-time = "2025-10-21T16:21:41.882Z" }, - { url = "https://files.pythonhosted.org/packages/e0/42/ad28191ebf983a5d0ecef90bab66baa5a6b18f2bfdef9d0a63b1973d9f75/grpcio-1.76.0-cp312-cp312-win32.whl", hash = "sha256:45d59a649a82df5718fd9527ce775fd66d1af35e6d31abdcdc906a49c6822958", size = 3984750, upload-time = "2025-10-21T16:21:44.006Z" }, - { url = "https://files.pythonhosted.org/packages/9e/00/7bd478cbb851c04a48baccaa49b75abaa8e4122f7d86da797500cccdd771/grpcio-1.76.0-cp312-cp312-win_amd64.whl", hash = "sha256:c088e7a90b6017307f423efbb9d1ba97a22aa2170876223f9709e9d1de0b5347", size = 4704003, upload-time = "2025-10-21T16:21:46.244Z" }, - { url = "https://files.pythonhosted.org/packages/fc/ed/71467ab770effc9e8cef5f2e7388beb2be26ed642d567697bb103a790c72/grpcio-1.76.0-cp313-cp313-linux_armv7l.whl", hash = "sha256:26ef06c73eb53267c2b319f43e6634c7556ea37672029241a056629af27c10e2", size = 5807716, upload-time = "2025-10-21T16:21:48.475Z" }, - { url = "https://files.pythonhosted.org/packages/2c/85/c6ed56f9817fab03fa8a111ca91469941fb514e3e3ce6d793cb8f1e1347b/grpcio-1.76.0-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:45e0111e73f43f735d70786557dc38141185072d7ff8dc1829d6a77ac1471468", size = 11821522, upload-time = "2025-10-21T16:21:51.142Z" }, - { url = "https://files.pythonhosted.org/packages/ac/31/2b8a235ab40c39cbc141ef647f8a6eb7b0028f023015a4842933bc0d6831/grpcio-1.76.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:83d57312a58dcfe2a3a0f9d1389b299438909a02db60e2f2ea2ae2d8034909d3", size = 6362558, upload-time = "2025-10-21T16:21:54.213Z" }, - { url = "https://files.pythonhosted.org/packages/bd/64/9784eab483358e08847498ee56faf8ff6ea8e0a4592568d9f68edc97e9e9/grpcio-1.76.0-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:3e2a27c89eb9ac3d81ec8835e12414d73536c6e620355d65102503064a4ed6eb", size = 7049990, upload-time = "2025-10-21T16:21:56.476Z" }, - { url = "https://files.pythonhosted.org/packages/2b/94/8c12319a6369434e7a184b987e8e9f3b49a114c489b8315f029e24de4837/grpcio-1.76.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:61f69297cba3950a524f61c7c8ee12e55c486cb5f7db47ff9dcee33da6f0d3ae", size = 6575387, upload-time = "2025-10-21T16:21:59.051Z" }, - { url = "https://files.pythonhosted.org/packages/15/0f/f12c32b03f731f4a6242f771f63039df182c8b8e2cf8075b245b409259d4/grpcio-1.76.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:6a15c17af8839b6801d554263c546c69c4d7718ad4321e3166175b37eaacca77", size = 7166668, upload-time = "2025-10-21T16:22:02.049Z" }, - { url = "https://files.pythonhosted.org/packages/ff/2d/3ec9ce0c2b1d92dd59d1c3264aaec9f0f7c817d6e8ac683b97198a36ed5a/grpcio-1.76.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:25a18e9810fbc7e7f03ec2516addc116a957f8cbb8cbc95ccc80faa072743d03", size = 8124928, upload-time = "2025-10-21T16:22:04.984Z" }, - { url = "https://files.pythonhosted.org/packages/1a/74/fd3317be5672f4856bcdd1a9e7b5e17554692d3db9a3b273879dc02d657d/grpcio-1.76.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:931091142fd8cc14edccc0845a79248bc155425eee9a98b2db2ea4f00a235a42", size = 7589983, upload-time = "2025-10-21T16:22:07.881Z" }, - { url = "https://files.pythonhosted.org/packages/45/bb/ca038cf420f405971f19821c8c15bcbc875505f6ffadafe9ffd77871dc4c/grpcio-1.76.0-cp313-cp313-win32.whl", hash = "sha256:5e8571632780e08526f118f74170ad8d50fb0a48c23a746bef2a6ebade3abd6f", size = 3984727, upload-time = "2025-10-21T16:22:10.032Z" }, - { url = "https://files.pythonhosted.org/packages/41/80/84087dc56437ced7cdd4b13d7875e7439a52a261e3ab4e06488ba6173b0a/grpcio-1.76.0-cp313-cp313-win_amd64.whl", hash = "sha256:f9f7bd5faab55f47231ad8dba7787866b69f5e93bc306e3915606779bbfb4ba8", size = 4702799, upload-time = "2025-10-21T16:22:12.709Z" }, - { url = "https://files.pythonhosted.org/packages/b4/46/39adac80de49d678e6e073b70204091e76631e03e94928b9ea4ecf0f6e0e/grpcio-1.76.0-cp314-cp314-linux_armv7l.whl", hash = "sha256:ff8a59ea85a1f2191a0ffcc61298c571bc566332f82e5f5be1b83c9d8e668a62", size = 5808417, upload-time = "2025-10-21T16:22:15.02Z" }, - { url = "https://files.pythonhosted.org/packages/9c/f5/a4531f7fb8b4e2a60b94e39d5d924469b7a6988176b3422487be61fe2998/grpcio-1.76.0-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:06c3d6b076e7b593905d04fdba6a0525711b3466f43b3400266f04ff735de0cd", size = 11828219, upload-time = "2025-10-21T16:22:17.954Z" }, - { url = "https://files.pythonhosted.org/packages/4b/1c/de55d868ed7a8bd6acc6b1d6ddc4aa36d07a9f31d33c912c804adb1b971b/grpcio-1.76.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fd5ef5932f6475c436c4a55e4336ebbe47bd3272be04964a03d316bbf4afbcbc", size = 6367826, upload-time = "2025-10-21T16:22:20.721Z" }, - { url = "https://files.pythonhosted.org/packages/59/64/99e44c02b5adb0ad13ab3adc89cb33cb54bfa90c74770f2607eea629b86f/grpcio-1.76.0-cp314-cp314-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:b331680e46239e090f5b3cead313cc772f6caa7d0fc8de349337563125361a4a", size = 7049550, upload-time = "2025-10-21T16:22:23.637Z" }, - { url = "https://files.pythonhosted.org/packages/43/28/40a5be3f9a86949b83e7d6a2ad6011d993cbe9b6bd27bea881f61c7788b6/grpcio-1.76.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2229ae655ec4e8999599469559e97630185fdd53ae1e8997d147b7c9b2b72cba", size = 6575564, upload-time = "2025-10-21T16:22:26.016Z" }, - { url = "https://files.pythonhosted.org/packages/4b/a9/1be18e6055b64467440208a8559afac243c66a8b904213af6f392dc2212f/grpcio-1.76.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:490fa6d203992c47c7b9e4a9d39003a0c2bcc1c9aa3c058730884bbbb0ee9f09", size = 7176236, upload-time = "2025-10-21T16:22:28.362Z" }, - { url = "https://files.pythonhosted.org/packages/0f/55/dba05d3fcc151ce6e81327541d2cc8394f442f6b350fead67401661bf041/grpcio-1.76.0-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:479496325ce554792dba6548fae3df31a72cef7bad71ca2e12b0e58f9b336bfc", size = 8125795, upload-time = "2025-10-21T16:22:31.075Z" }, - { url = "https://files.pythonhosted.org/packages/4a/45/122df922d05655f63930cf42c9e3f72ba20aadb26c100ee105cad4ce4257/grpcio-1.76.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:1c9b93f79f48b03ada57ea24725d83a30284a012ec27eab2cf7e50a550cbbbcc", size = 7592214, upload-time = "2025-10-21T16:22:33.831Z" }, - { url = "https://files.pythonhosted.org/packages/4a/6e/0b899b7f6b66e5af39e377055fb4a6675c9ee28431df5708139df2e93233/grpcio-1.76.0-cp314-cp314-win32.whl", hash = "sha256:747fa73efa9b8b1488a95d0ba1039c8e2dca0f741612d80415b1e1c560febf4e", size = 4062961, upload-time = "2025-10-21T16:22:36.468Z" }, - { url = "https://files.pythonhosted.org/packages/19/41/0b430b01a2eb38ee887f88c1f07644a1df8e289353b78e82b37ef988fb64/grpcio-1.76.0-cp314-cp314-win_amd64.whl", hash = "sha256:922fa70ba549fce362d2e2871ab542082d66e2aaf0c19480ea453905b01f384e", size = 4834462, upload-time = "2025-10-21T16:22:39.772Z" }, -] - -[[package]] -name = "grpcio-status" -version = "1.76.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "googleapis-common-protos" }, - { name = "grpcio" }, - { name = "protobuf" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/3f/46/e9f19d5be65e8423f886813a2a9d0056ba94757b0c5007aa59aed1a961fa/grpcio_status-1.76.0.tar.gz", hash = "sha256:25fcbfec74c15d1a1cb5da3fab8ee9672852dc16a5a9eeb5baf7d7a9952943cd", size = 13679, upload-time = "2025-10-21T16:28:52.545Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/8c/cc/27ba60ad5a5f2067963e6a858743500df408eb5855e98be778eaef8c9b02/grpcio_status-1.76.0-py3-none-any.whl", hash = "sha256:380568794055a8efbbd8871162df92012e0228a5f6dffaf57f2a00c534103b18", size = 14425, upload-time = "2025-10-21T16:28:40.853Z" }, -] - -[[package]] -name = "h11" -version = "0.16.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/01/ee/02a2c011bdab74c6fb3c75474d40b3052059d95df7e73351460c8588d963/h11-0.16.0.tar.gz", hash = "sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1", size = 101250, upload-time = "2025-04-24T03:35:25.427Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86", size = 37515, upload-time = "2025-04-24T03:35:24.344Z" }, -] - -[[package]] -name = "httpcore" -version = "1.0.9" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "certifi" }, - { name = "h11" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/06/94/82699a10bca87a5556c9c59b5963f2d039dbd239f25bc2a63907a05a14cb/httpcore-1.0.9.tar.gz", hash = "sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8", size = 85484, upload-time = "2025-04-24T22:06:22.219Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" }, -] - -[[package]] -name = "httpx" -version = "0.28.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "anyio" }, - { name = "certifi" }, - { name = "httpcore" }, - { name = "idna" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/b1/df/48c586a5fe32a0f01324ee087459e112ebb7224f646c0b5023f5e79e9956/httpx-0.28.1.tar.gz", hash = "sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc", size = 141406, upload-time = "2024-12-06T15:37:23.222Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" }, -] - -[[package]] -name = "httpx-sse" -version = "0.4.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/0f/4c/751061ffa58615a32c31b2d82e8482be8dd4a89154f003147acee90f2be9/httpx_sse-0.4.3.tar.gz", hash = "sha256:9b1ed0127459a66014aec3c56bebd93da3c1bc8bb6618c8082039a44889a755d", size = 15943, upload-time = "2025-10-10T21:48:22.271Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d2/fd/6668e5aec43ab844de6fc74927e155a3b37bf40d7c3790e49fc0406b6578/httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc", size = 8960, upload-time = "2025-10-10T21:48:21.158Z" }, -] - [[package]] name = "hydra-core" version = "1.3.2" @@ -1287,18 +633,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0e/61/66938bbb5fc52dbdf84594873d5b51fb1f7c7794e9c0f5bd885f30bc507b/idna-3.11-py3-none-any.whl", hash = "sha256:771a87f49d9defaf64091e6e6fe9c18d4833f140bd19464795bc32d966ca37ea", size = 71008, upload-time = "2025-10-12T14:55:18.883Z" }, ] -[[package]] -name = "importlib-metadata" -version = "8.7.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "zipp" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/f3/49/3b30cad09e7771a4982d9975a8cbf64f00d4a1ececb53297f1d9a7be1b10/importlib_metadata-8.7.1.tar.gz", hash = "sha256:49fef1ae6440c182052f407c8d34a68f72efc36db9ca90dc0113398f2fdde8bb", size = 57107, upload-time = "2025-12-21T10:00:19.278Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fa/5e/f8e9a1d23b9c20a551a8a02ea3637b4642e22c2626e3a13a9a29cdea99eb/importlib_metadata-8.7.1-py3-none-any.whl", hash = "sha256:5a1f80bf1daa489495071efbb095d75a634cf28a8bc299581244063b53176151", size = 27865, upload-time = "2025-12-21T10:00:18.329Z" }, -] - [[package]] name = "iniconfig" version = "2.3.0" @@ -1308,57 +642,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, ] -[[package]] -name = "interface-meta" -version = "1.3.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/4d/75/10526292b332f3479c246750a96f6ec11a28e297839a9c25583b2aadc119/interface_meta-1.3.0.tar.gz", hash = "sha256:8a4493f8bdb73fb9655dcd5115bc897e207319e36c8835f39c516a2d7e9d79a1", size = 15007, upload-time = "2022-04-04T01:12:46.517Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/02/3f/a6ec28c88e2d8e54d32598a1e0b5208a4baa72a8e7f6e241beab5731eb9d/interface_meta-1.3.0-py3-none-any.whl", hash = "sha256:de35dc5241431886e709e20a14d6597ed07c9f1e8b4bfcffde2190ca5b700ee8", size = 14854, upload-time = "2022-04-04T01:12:45.183Z" }, -] - -[[package]] -name = "jaraco-classes" -version = "3.4.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "more-itertools" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/06/c0/ed4a27bc5571b99e3cff68f8a9fa5b56ff7df1c2251cc715a652ddd26402/jaraco.classes-3.4.0.tar.gz", hash = "sha256:47a024b51d0239c0dd8c8540c6c7f484be3b8fcf0b2d85c13825780d3b3f3acd", size = 11780, upload-time = "2024-03-31T07:27:36.643Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7f/66/b15ce62552d84bbfcec9a4873ab79d993a1dd4edb922cbfccae192bd5b5f/jaraco.classes-3.4.0-py3-none-any.whl", hash = "sha256:f662826b6bed8cace05e7ff873ce0f9283b5c924470fe664fff1c2f00f581790", size = 6777, upload-time = "2024-03-31T07:27:34.792Z" }, -] - -[[package]] -name = "jaraco-context" -version = "6.0.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/8d/7d/41acf8e22d791bde812cb6c2c36128bb932ed8ae066bcb5e39cb198e8253/jaraco_context-6.0.2.tar.gz", hash = "sha256:953ae8dddb57b1d791bf72ea1009b32088840a7dd19b9ba16443f62be919ee57", size = 14994, upload-time = "2025-12-24T19:21:35.784Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c7/0c/1e0096ced9c55f9c6c6655446798df74165780375d3f5ab5f33751e087ae/jaraco_context-6.0.2-py3-none-any.whl", hash = "sha256:55fc21af4b4f9ca94aa643b6ee7fe13b1e4c01abf3aeb98ca4ad9c80b741c786", size = 6988, upload-time = "2025-12-24T19:21:34.557Z" }, -] - -[[package]] -name = "jaraco-functools" -version = "4.4.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "more-itertools" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/0f/27/056e0638a86749374d6f57d0b0db39f29509cce9313cf91bdc0ac4d91084/jaraco_functools-4.4.0.tar.gz", hash = "sha256:da21933b0417b89515562656547a77b4931f98176eb173644c0d35032a33d6bb", size = 19943, upload-time = "2025-12-21T09:29:43.6Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fd/c4/813bb09f0985cb21e959f21f2464169eca882656849adf727ac7bb7e1767/jaraco_functools-4.4.0-py3-none-any.whl", hash = "sha256:9eec1e36f45c818d9bf307c8948eb03b2b56cd44087b3cdc989abca1f20b9176", size = 10481, upload-time = "2025-12-21T09:29:42.27Z" }, -] - -[[package]] -name = "jeepney" -version = "0.9.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/7b/6f/357efd7602486741aa73ffc0617fb310a29b588ed0fd69c2399acbb85b0c/jeepney-0.9.0.tar.gz", hash = "sha256:cf0e9e845622b81e4a28df94c40345400256ec608d0e55bb8a3feaa9163f5732", size = 106758, upload-time = "2025-02-27T18:51:01.684Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b2/a3/e137168c9c44d18eff0376253da9f1e9234d0239e0ee230d2fee6cea8e55/jeepney-0.9.0-py3-none-any.whl", hash = "sha256:97e5714520c16fc0a45695e5365a2e11b81ea79bba796e26f9f1d178cb182683", size = 49010, upload-time = "2025-02-27T18:51:00.104Z" }, -] - [[package]] name = "jinja2" version = "3.1.6" @@ -1380,65 +663,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1e/e8/685f47e0d754320684db4425a0967f7d3fa70126bffd76110b7009a0090f/joblib-1.5.2-py3-none-any.whl", hash = "sha256:4e1f0bdbb987e6d843c70cf43714cb276623def372df3c22fe5266b2670bc241", size = 308396, upload-time = "2025-08-27T12:15:45.188Z" }, ] -[[package]] -name = "jsonschema" -version = "4.26.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "attrs" }, - { name = "jsonschema-specifications" }, - { name = "referencing" }, - { name = "rpds-py" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/b3/fc/e067678238fa451312d4c62bf6e6cf5ec56375422aee02f9cb5f909b3047/jsonschema-4.26.0.tar.gz", hash = "sha256:0c26707e2efad8aa1bfc5b7ce170f3fccc2e4918ff85989ba9ffa9facb2be326", size = 366583, upload-time = "2026-01-07T13:41:07.246Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/69/90/f63fb5873511e014207a475e2bb4e8b2e570d655b00ac19a9a0ca0a385ee/jsonschema-4.26.0-py3-none-any.whl", hash = "sha256:d489f15263b8d200f8387e64b4c3a75f06629559fb73deb8fdfb525f2dab50ce", size = 90630, upload-time = "2026-01-07T13:41:05.306Z" }, -] - -[[package]] -name = "jsonschema-path" -version = "0.3.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pathable" }, - { name = "pyyaml" }, - { name = "referencing" }, - { name = "requests" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/6e/45/41ebc679c2a4fced6a722f624c18d658dee42612b83ea24c1caf7c0eb3a8/jsonschema_path-0.3.4.tar.gz", hash = "sha256:8365356039f16cc65fddffafda5f58766e34bebab7d6d105616ab52bc4297001", size = 11159, upload-time = "2025-01-24T14:33:16.547Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/cb/58/3485da8cb93d2f393bce453adeef16896751f14ba3e2024bc21dc9597646/jsonschema_path-0.3.4-py3-none-any.whl", hash = "sha256:f502191fdc2b22050f9a81c9237be9d27145b9001c55842bece5e94e382e52f8", size = 14810, upload-time = "2025-01-24T14:33:14.652Z" }, -] - -[[package]] -name = "jsonschema-specifications" -version = "2025.9.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "referencing" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/19/74/a633ee74eb36c44aa6d1095e7cc5569bebf04342ee146178e2d36600708b/jsonschema_specifications-2025.9.1.tar.gz", hash = "sha256:b540987f239e745613c7a9176f3edb72b832a4ac465cf02712288397832b5e8d", size = 32855, upload-time = "2025-09-08T01:34:59.186Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/41/45/1a4ed80516f02155c51f51e8cedb3c1902296743db0bbc66608a0db2814f/jsonschema_specifications-2025.9.1-py3-none-any.whl", hash = "sha256:98802fee3a11ee76ecaca44429fda8a41bff98b00a0f2838151b113f210cc6fe", size = 18437, upload-time = "2025-09-08T01:34:57.871Z" }, -] - -[[package]] -name = "keyring" -version = "25.7.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "jaraco-classes" }, - { name = "jaraco-context" }, - { name = "jaraco-functools" }, - { name = "jeepney", marker = "sys_platform == 'linux'" }, - { name = "pywin32-ctypes", marker = "sys_platform == 'win32'" }, - { name = "secretstorage", marker = "sys_platform == 'linux'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/43/4b/674af6ef2f97d56f0ab5153bf0bfa28ccb6c3ed4d1babf4305449668807b/keyring-25.7.0.tar.gz", hash = "sha256:fe01bd85eb3f8fb3dd0405defdeac9a5b4f6f0439edbb3149577f244a2e8245b", size = 63516, upload-time = "2025-11-16T16:26:09.482Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/81/db/e655086b7f3a705df045bf0933bdd9c2f79bb3c97bfef1384598bb79a217/keyring-25.7.0-py3-none-any.whl", hash = "sha256:be4a0b195f149690c166e850609a477c532ddbfbaed96a404d4e43f8d5e2689f", size = 39160, upload-time = "2025-11-16T16:26:08.402Z" }, -] - [[package]] name = "kiwisolver" version = "1.4.9" @@ -1563,24 +787,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0e/3d/72cc9ec90bb80b5b1a65f0bb74a0f540195837baaf3b98c7fa4a7aa9718e/librt-0.6.3-cp314-cp314t-win_arm64.whl", hash = "sha256:afb39550205cc5e5c935762c6bf6a2bb34f7d21a68eadb25e2db7bf3593fecc0", size = 20246, upload-time = "2025-11-29T14:01:44.13Z" }, ] -[[package]] -name = "lifelines" -version = "0.30.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "autograd" }, - { name = "autograd-gamma" }, - { name = "formulaic" }, - { name = "matplotlib" }, - { name = "numpy" }, - { name = "pandas" }, - { name = "scipy" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/ea/4f/f0b363278d40baf7d7a03217bee839cb880946c62109f243391c8754bb09/lifelines-0.30.0.tar.gz", hash = "sha256:f7f6f6275fcb167fe0f5b1ef98f868993f9c074cb74b1dd6e92736efa854be18", size = 383221, upload-time = "2024-10-29T12:00:43.513Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/14/f7/379e185a75ac8166ac70756d0ba68d9a2b02b555c7fde4983246752396bd/lifelines-0.30.0-py3-none-any.whl", hash = "sha256:ac7c602c8aceced9770d3977817c9d99c250ed8cd86f2567fa0d23e4e8014bf9", size = 349319, upload-time = "2024-10-29T12:00:41.749Z" }, -] - [[package]] name = "lightning" version = "2.6.0" @@ -1616,55 +822,27 @@ wheels = [ ] [[package]] -name = "lupa" -version = "2.6" +name = "llvmlite" +version = "0.47.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/b8/1c/191c3e6ec6502e3dbe25a53e27f69a5daeac3e56de1f73c0138224171ead/lupa-2.6.tar.gz", hash = "sha256:9a770a6e89576be3447668d7ced312cd6fd41d3c13c2462c9dc2c2ab570e45d9", size = 7240282, upload-time = "2025-10-24T07:20:29.738Z" } +sdist = { url = "https://files.pythonhosted.org/packages/01/88/a8952b6d5c21e74cbf158515b779666f692846502623e9e3c39d8e8ba25f/llvmlite-0.47.0.tar.gz", hash = "sha256:62031ce968ec74e95092184d4b0e857e444f8fdff0b8f9213707699570c33ccc", size = 193614, upload-time = "2026-03-31T18:29:53.497Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/94/86/ce243390535c39d53ea17ccf0240815e6e457e413e40428a658ea4ee4b8d/lupa-2.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:47ce718817ef1cc0c40d87c3d5ae56a800d61af00fbc0fad1ca9be12df2f3b56", size = 951707, upload-time = "2025-10-24T07:18:03.884Z" }, - { url = "https://files.pythonhosted.org/packages/86/85/cedea5e6cbeb54396fdcc55f6b741696f3f036d23cfaf986d50d680446da/lupa-2.6-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:7aba985b15b101495aa4b07112cdc08baa0c545390d560ad5cfde2e9e34f4d58", size = 1916703, upload-time = "2025-10-24T07:18:05.6Z" }, - { url = "https://files.pythonhosted.org/packages/24/be/3d6b5f9a8588c01a4d88129284c726017b2089f3a3fd3ba8bd977292fea0/lupa-2.6-cp312-cp312-macosx_11_0_x86_64.whl", hash = "sha256:b766f62f95b2739f2248977d29b0722e589dcf4f0ccfa827ccbd29f0148bd2e5", size = 985152, upload-time = "2025-10-24T07:18:08.561Z" }, - { url = "https://files.pythonhosted.org/packages/eb/23/9f9a05beee5d5dce9deca4cb07c91c40a90541fc0a8e09db4ee670da550f/lupa-2.6-cp312-cp312-manylinux2010_i686.manylinux_2_12_i686.manylinux_2_28_i686.whl", hash = "sha256:00a934c23331f94cb51760097ebfab14b005d55a6b30a2b480e3c53dd2fa290d", size = 1159599, upload-time = "2025-10-24T07:18:10.346Z" }, - { url = "https://files.pythonhosted.org/packages/40/4e/e7c0583083db9d7f1fd023800a9767d8e4391e8330d56c2373d890ac971b/lupa-2.6-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:21de9f38bd475303e34a042b7081aabdf50bd9bafd36ce4faea2f90fd9f15c31", size = 1038686, upload-time = "2025-10-24T07:18:12.112Z" }, - { url = "https://files.pythonhosted.org/packages/1c/9f/5a4f7d959d4feba5e203ff0c31889e74d1ca3153122be4a46dca7d92bf7c/lupa-2.6-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:cf3bda96d3fc41237e964a69c23647d50d4e28421111360274d4799832c560e9", size = 2071956, upload-time = "2025-10-24T07:18:14.572Z" }, - { url = "https://files.pythonhosted.org/packages/92/34/2f4f13ca65d01169b1720176aedc4af17bc19ee834598c7292db232cb6dc/lupa-2.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:5a76ead245da54801a81053794aa3975f213221f6542d14ec4b859ee2e7e0323", size = 1057199, upload-time = "2025-10-24T07:18:16.379Z" }, - { url = "https://files.pythonhosted.org/packages/35/2a/5f7d2eebec6993b0dcd428e0184ad71afb06a45ba13e717f6501bfed1da3/lupa-2.6-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:8dd0861741caa20886ddbda0a121d8e52fb9b5bb153d82fa9bba796962bf30e8", size = 1173693, upload-time = "2025-10-24T07:18:18.153Z" }, - { url = "https://files.pythonhosted.org/packages/e4/29/089b4d2f8e34417349af3904bb40bec40b65c8731f45e3fd8d497ca573e5/lupa-2.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:239e63948b0b23023f81d9a19a395e768ed3da6a299f84e7963b8f813f6e3f9c", size = 2164394, upload-time = "2025-10-24T07:18:20.403Z" }, - { url = "https://files.pythonhosted.org/packages/f3/1b/79c17b23c921f81468a111cad843b076a17ef4b684c4a8dff32a7969c3f0/lupa-2.6-cp312-cp312-win32.whl", hash = "sha256:325894e1099499e7a6f9c351147661a2011887603c71086d36fe0f964d52d1ce", size = 1420647, upload-time = "2025-10-24T07:18:23.368Z" }, - { url = "https://files.pythonhosted.org/packages/b8/15/5121e68aad3584e26e1425a5c9a79cd898f8a152292059e128c206ee817c/lupa-2.6-cp312-cp312-win_amd64.whl", hash = "sha256:c735a1ce8ee60edb0fe71d665f1e6b7c55c6021f1d340eb8c865952c602cd36f", size = 1688529, upload-time = "2025-10-24T07:18:25.523Z" }, - { url = "https://files.pythonhosted.org/packages/28/1d/21176b682ca5469001199d8b95fa1737e29957a3d185186e7a8b55345f2e/lupa-2.6-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:663a6e58a0f60e7d212017d6678639ac8df0119bc13c2145029dcba084391310", size = 947232, upload-time = "2025-10-24T07:18:27.878Z" }, - { url = "https://files.pythonhosted.org/packages/ce/4c/d327befb684660ca13cf79cd1f1d604331808f9f1b6fb6bf57832f8edf80/lupa-2.6-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:d1f5afda5c20b1f3217a80e9bc1b77037f8a6eb11612fd3ada19065303c8f380", size = 1908625, upload-time = "2025-10-24T07:18:29.944Z" }, - { url = "https://files.pythonhosted.org/packages/66/8e/ad22b0a19454dfd08662237a84c792d6d420d36b061f239e084f29d1a4f3/lupa-2.6-cp313-cp313-macosx_11_0_x86_64.whl", hash = "sha256:26f2b3c085fe76e9119e48c1013c1cccdc1f51585d456858290475aa38e7089e", size = 981057, upload-time = "2025-10-24T07:18:31.553Z" }, - { url = "https://files.pythonhosted.org/packages/5c/48/74859073ab276bd0566c719f9ca0108b0cfc1956ca0d68678d117d47d155/lupa-2.6-cp313-cp313-manylinux2010_i686.manylinux_2_12_i686.manylinux_2_28_i686.whl", hash = "sha256:60d2f902c7b96fb8ab98493dcff315e7bb4d0b44dc9dd76eb37de575025d5685", size = 1156227, upload-time = "2025-10-24T07:18:33.981Z" }, - { url = "https://files.pythonhosted.org/packages/09/6c/0e9ded061916877253c2266074060eb71ed99fb21d73c8c114a76725bce2/lupa-2.6-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a02d25dee3a3250967c36590128d9220ae02f2eda166a24279da0b481519cbff", size = 1035752, upload-time = "2025-10-24T07:18:36.32Z" }, - { url = "https://files.pythonhosted.org/packages/dd/ef/f8c32e454ef9f3fe909f6c7d57a39f950996c37a3deb7b391fec7903dab7/lupa-2.6-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6eae1ee16b886b8914ff292dbefbf2f48abfbdee94b33a88d1d5475e02423203", size = 2069009, upload-time = "2025-10-24T07:18:38.072Z" }, - { url = "https://files.pythonhosted.org/packages/53/dc/15b80c226a5225815a890ee1c11f07968e0aba7a852df41e8ae6fe285063/lupa-2.6-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:b0edd5073a4ee74ab36f74fe61450148e6044f3952b8d21248581f3c5d1a58be", size = 1056301, upload-time = "2025-10-24T07:18:40.165Z" }, - { url = "https://files.pythonhosted.org/packages/31/14/2086c1425c985acfb30997a67e90c39457122df41324d3c179d6ee2292c6/lupa-2.6-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:0c53ee9f22a8a17e7d4266ad48e86f43771951797042dd51d1494aaa4f5f3f0a", size = 1170673, upload-time = "2025-10-24T07:18:42.426Z" }, - { url = "https://files.pythonhosted.org/packages/10/e5/b216c054cf86576c0191bf9a9f05de6f7e8e07164897d95eea0078dca9b2/lupa-2.6-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:de7c0f157a9064a400d828789191a96da7f4ce889969a588b87ec80de9b14772", size = 2162227, upload-time = "2025-10-24T07:18:46.112Z" }, - { url = "https://files.pythonhosted.org/packages/59/2f/33ecb5bedf4f3bc297ceacb7f016ff951331d352f58e7e791589609ea306/lupa-2.6-cp313-cp313-win32.whl", hash = "sha256:ee9523941ae0a87b5b703417720c5d78f72d2f5bc23883a2ea80a949a3ed9e75", size = 1419558, upload-time = "2025-10-24T07:18:48.371Z" }, - { url = "https://files.pythonhosted.org/packages/f9/b4/55e885834c847ea610e111d87b9ed4768f0afdaeebc00cd46810f25029f6/lupa-2.6-cp313-cp313-win_amd64.whl", hash = "sha256:b1335a5835b0a25ebdbc75cf0bda195e54d133e4d994877ef025e218c2e59db9", size = 1683424, upload-time = "2025-10-24T07:18:50.976Z" }, - { url = "https://files.pythonhosted.org/packages/66/9d/d9427394e54d22a35d1139ef12e845fd700d4872a67a34db32516170b746/lupa-2.6-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:dcb6d0a3264873e1653bc188499f48c1fb4b41a779e315eba45256cfe7bc33c1", size = 953818, upload-time = "2025-10-24T07:18:53.378Z" }, - { url = "https://files.pythonhosted.org/packages/10/41/27bbe81953fb2f9ecfced5d9c99f85b37964cfaf6aa8453bb11283983721/lupa-2.6-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:a37e01f2128f8c36106726cb9d360bac087d58c54b4522b033cc5691c584db18", size = 1915850, upload-time = "2025-10-24T07:18:55.259Z" }, - { url = "https://files.pythonhosted.org/packages/a3/98/f9ff60db84a75ba8725506bbf448fb085bc77868a021998ed2a66d920568/lupa-2.6-cp314-cp314-macosx_11_0_x86_64.whl", hash = "sha256:458bd7e9ff3c150b245b0fcfbb9bd2593d1152ea7f0a7b91c1d185846da033fe", size = 982344, upload-time = "2025-10-24T07:18:57.05Z" }, - { url = "https://files.pythonhosted.org/packages/41/f7/f39e0f1c055c3b887d86b404aaf0ca197b5edfd235a8b81b45b25bac7fc3/lupa-2.6-cp314-cp314-manylinux2010_i686.manylinux_2_12_i686.manylinux_2_28_i686.whl", hash = "sha256:052ee82cac5206a02df77119c325339acbc09f5ce66967f66a2e12a0f3211cad", size = 1156543, upload-time = "2025-10-24T07:18:59.251Z" }, - { url = "https://files.pythonhosted.org/packages/9e/9c/59e6cffa0d672d662ae17bd7ac8ecd2c89c9449dee499e3eb13ca9cd10d9/lupa-2.6-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:96594eca3c87dd07938009e95e591e43d554c1dbd0385be03c100367141db5a8", size = 1047974, upload-time = "2025-10-24T07:19:01.449Z" }, - { url = "https://files.pythonhosted.org/packages/23/c6/a04e9cef7c052717fcb28fb63b3824802488f688391895b618e39be0f684/lupa-2.6-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e8faddd9d198688c8884091173a088a8e920ecc96cda2ffed576a23574c4b3f6", size = 2073458, upload-time = "2025-10-24T07:19:03.369Z" }, - { url = "https://files.pythonhosted.org/packages/e6/10/824173d10f38b51fc77785228f01411b6ca28826ce27404c7c912e0e442c/lupa-2.6-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:daebb3a6b58095c917e76ba727ab37b27477fb926957c825205fbda431552134", size = 1067683, upload-time = "2025-10-24T07:19:06.2Z" }, - { url = "https://files.pythonhosted.org/packages/b6/dc/9692fbcf3c924d9c4ece2d8d2f724451ac2e09af0bd2a782db1cef34e799/lupa-2.6-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:f3154e68972befe0f81564e37d8142b5d5d79931a18309226a04ec92487d4ea3", size = 1171892, upload-time = "2025-10-24T07:19:08.544Z" }, - { url = "https://files.pythonhosted.org/packages/84/ff/e318b628d4643c278c96ab3ddea07fc36b075a57383c837f5b11e537ba9d/lupa-2.6-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:e4dadf77b9fedc0bfa53417cc28dc2278a26d4cbd95c29f8927ad4d8fe0a7ef9", size = 2166641, upload-time = "2025-10-24T07:19:10.485Z" }, - { url = "https://files.pythonhosted.org/packages/12/f7/a6f9ec2806cf2d50826980cdb4b3cffc7691dc6f95e13cc728846d5cb793/lupa-2.6-cp314-cp314-win32.whl", hash = "sha256:cb34169c6fa3bab3e8ac58ca21b8a7102f6a94b6a5d08d3636312f3f02fafd8f", size = 1456857, upload-time = "2025-10-24T07:19:37.989Z" }, - { url = "https://files.pythonhosted.org/packages/c5/de/df71896f25bdc18360fdfa3b802cd7d57d7fede41a0e9724a4625b412c85/lupa-2.6-cp314-cp314-win_amd64.whl", hash = "sha256:b74f944fe46c421e25d0f8692aef1e842192f6f7f68034201382ac440ef9ea67", size = 1731191, upload-time = "2025-10-24T07:19:40.281Z" }, - { url = "https://files.pythonhosted.org/packages/47/3c/a1f23b01c54669465f5f4c4083107d496fbe6fb45998771420e9aadcf145/lupa-2.6-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:0e21b716408a21ab65723f8841cf7f2f37a844b7a965eeabb785e27fca4099cf", size = 999343, upload-time = "2025-10-24T07:19:12.519Z" }, - { url = "https://files.pythonhosted.org/packages/c5/6d/501994291cb640bfa2ccf7f554be4e6914afa21c4026bd01bff9ca8aac57/lupa-2.6-cp314-cp314t-macosx_11_0_universal2.whl", hash = "sha256:589db872a141bfff828340079bbdf3e9a31f2689f4ca0d88f97d9e8c2eae6142", size = 2000730, upload-time = "2025-10-24T07:19:14.869Z" }, - { url = "https://files.pythonhosted.org/packages/53/a5/457ffb4f3f20469956c2d4c4842a7675e884efc895b2f23d126d23e126cc/lupa-2.6-cp314-cp314t-macosx_11_0_x86_64.whl", hash = "sha256:cd852a91a4a9d4dcbb9a58100f820a75a425703ec3e3f049055f60b8533b7953", size = 1021553, upload-time = "2025-10-24T07:19:17.123Z" }, - { url = "https://files.pythonhosted.org/packages/51/6b/36bb5a5d0960f2a5c7c700e0819abb76fd9bf9c1d8a66e5106416d6e9b14/lupa-2.6-cp314-cp314t-manylinux2010_i686.manylinux_2_12_i686.manylinux_2_28_i686.whl", hash = "sha256:0334753be028358922415ca97a64a3048e4ed155413fc4eaf87dd0a7e2752983", size = 1133275, upload-time = "2025-10-24T07:19:20.51Z" }, - { url = "https://files.pythonhosted.org/packages/19/86/202ff4429f663013f37d2229f6176ca9f83678a50257d70f61a0a97281bf/lupa-2.6-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:661d895cd38c87658a34780fac54a690ec036ead743e41b74c3fb81a9e65a6aa", size = 1038441, upload-time = "2025-10-24T07:19:22.509Z" }, - { url = "https://files.pythonhosted.org/packages/a7/42/d8125f8e420714e5b52e9c08d88b5329dfb02dcca731b4f21faaee6cc5b5/lupa-2.6-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6aa58454ccc13878cc177c62529a2056be734da16369e451987ff92784994ca7", size = 2058324, upload-time = "2025-10-24T07:19:24.979Z" }, - { url = "https://files.pythonhosted.org/packages/2b/2c/47bf8b84059876e877a339717ddb595a4a7b0e8740bacae78ba527562e1c/lupa-2.6-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:1425017264e470c98022bba8cff5bd46d054a827f5df6b80274f9cc71dafd24f", size = 1060250, upload-time = "2025-10-24T07:19:27.262Z" }, - { url = "https://files.pythonhosted.org/packages/c2/06/d88add2b6406ca1bdec99d11a429222837ca6d03bea42ca75afa169a78cb/lupa-2.6-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:224af0532d216e3105f0a127410f12320f7c5f1aa0300bdf9646b8d9afb0048c", size = 1151126, upload-time = "2025-10-24T07:19:29.522Z" }, - { url = "https://files.pythonhosted.org/packages/b4/a0/89e6a024c3b4485b89ef86881c9d55e097e7cb0bdb74efb746f2fa6a9a76/lupa-2.6-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:9abb98d5a8fd27c8285302e82199f0e56e463066f88f619d6594a450bf269d80", size = 2153693, upload-time = "2025-10-24T07:19:31.379Z" }, - { url = "https://files.pythonhosted.org/packages/b6/36/a0f007dc58fc1bbf51fb85dcc82fcb1f21b8c4261361de7dab0e3d8521ef/lupa-2.6-cp314-cp314t-win32.whl", hash = "sha256:1849efeba7a8f6fb8aa2c13790bee988fd242ae404bd459509640eeea3d1e291", size = 1590104, upload-time = "2025-10-24T07:19:33.514Z" }, - { url = "https://files.pythonhosted.org/packages/7d/5e/db903ce9cf82c48d6b91bf6d63ae4c8d0d17958939a4e04ba6b9f38b8643/lupa-2.6-cp314-cp314t-win_amd64.whl", hash = "sha256:fc1498d1a4fc028bc521c26d0fad4ca00ed63b952e32fb95949bda76a04bad52", size = 1913818, upload-time = "2025-10-24T07:19:36.039Z" }, + { url = "https://files.pythonhosted.org/packages/fa/48/4b7fe0e34c169fa2f12532916133e0b219d2823b540733651b34fdac509a/llvmlite-0.47.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:306a265f408c259067257a732c8e159284334018b4083a9e35f67d19792b164f", size = 37232769, upload-time = "2026-03-31T18:28:43.735Z" }, + { url = "https://files.pythonhosted.org/packages/e6/4b/e3f2cd17822cf772a4a51a0a8080b0032e6d37b2dbe8cfb724eac4e31c52/llvmlite-0.47.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5853bf26160857c0c2573415ff4efe01c4c651e59e2c55c2a088740acfee51cd", size = 56275178, upload-time = "2026-03-31T18:28:48.342Z" }, + { url = "https://files.pythonhosted.org/packages/b6/55/a3b4a543185305a9bdf3d9759d53646ed96e55e7dfd43f53e7a421b8fbae/llvmlite-0.47.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:003bcf7fa579e14db59c1a1e113f93ab8a06b56a4be31c7f08264d1d4072d077", size = 55128632, upload-time = "2026-03-31T18:28:52.901Z" }, + { url = "https://files.pythonhosted.org/packages/2f/f5/d281ae0f79378a5a91f308ea9fdb9f9cc068fddd09629edc0725a5a8fde1/llvmlite-0.47.0-cp312-cp312-win_amd64.whl", hash = "sha256:f3079f25bdc24cd9d27c4b2b5e68f5f60c4fdb7e8ad5ee2b9b006007558f9df7", size = 38138692, upload-time = "2026-03-31T18:28:57.147Z" }, + { url = "https://files.pythonhosted.org/packages/77/6f/4615353e016799f80fa52ccb270a843c413b22361fadda2589b2922fb9b0/llvmlite-0.47.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:a3c6a735d4e1041808434f9d440faa3d78d9b4af2ee64d05a66f351883b6ceec", size = 37232771, upload-time = "2026-03-31T18:29:01.324Z" }, + { url = "https://files.pythonhosted.org/packages/31/b8/69f5565f1a280d032525878a86511eebed0645818492feeb169dfb20ae8e/llvmlite-0.47.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2699a74321189e812d476a43d6d7f652f51811e7b5aad9d9bba842a1c7927acb", size = 56275178, upload-time = "2026-03-31T18:29:05.748Z" }, + { url = "https://files.pythonhosted.org/packages/d6/da/b32cafcb926fb0ce2aa25553bf32cb8764af31438f40e2481df08884c947/llvmlite-0.47.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6c6951e2b29930227963e53ee152441f0e14be92e9d4231852102d986c761e40", size = 55128632, upload-time = "2026-03-31T18:29:11.235Z" }, + { url = "https://files.pythonhosted.org/packages/46/9f/4898b44e4042c60fafcb1162dfb7014f6f15b1ec19bf29cfea6bf26df90d/llvmlite-0.47.0-cp313-cp313-win_amd64.whl", hash = "sha256:c2e9adf8698d813a9a5efb2d4370caf344dbc1e145019851fee6a6f319ba760e", size = 38138695, upload-time = "2026-03-31T18:29:15.43Z" }, + { url = "https://files.pythonhosted.org/packages/1c/d4/33c8af00f0bf6f552d74f3a054f648af2c5bc6bece97972f3bfadce4f5ec/llvmlite-0.47.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:de966c626c35c9dff5ae7bf12db25637738d0df83fc370cf793bc94d43d92d14", size = 37232773, upload-time = "2026-03-31T18:29:19.453Z" }, + { url = "https://files.pythonhosted.org/packages/64/1d/a760e993e0c0ba6db38d46b9f48f6c7dceb8ac838824997fb9e25f97bc04/llvmlite-0.47.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ddbccff2aeaff8670368340a158abefc032fe9b3ccf7d9c496639263d00151aa", size = 56275176, upload-time = "2026-03-31T18:29:24.149Z" }, + { url = "https://files.pythonhosted.org/packages/84/3b/e679bc3b29127182a7f4aa2d2e9e5bea42adb93fb840484147d59c236299/llvmlite-0.47.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d4a7b778a2e144fc64468fb9bf509ac1226c9813a00b4d7afea5d988c4e22fca", size = 55128631, upload-time = "2026-03-31T18:29:29.536Z" }, + { url = "https://files.pythonhosted.org/packages/be/f7/19e2a09c62809c9e63bbd14ce71fb92c6ff7b7b3045741bb00c781efc3c9/llvmlite-0.47.0-cp314-cp314-win_amd64.whl", hash = "sha256:694e3c2cdc472ed2bd8bd4555ca002eec4310961dd58ef791d508f57b5cc4c94", size = 39153826, upload-time = "2026-03-31T18:29:33.681Z" }, + { url = "https://files.pythonhosted.org/packages/40/a1/581a8c707b5e80efdbbe1dd94527404d33fe50bceb71f39d5a7e11bd57b7/llvmlite-0.47.0-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:92ec8a169a20b473c1c54d4695e371bde36489fc1efa3688e11e99beba0abf9c", size = 37232772, upload-time = "2026-03-31T18:29:37.952Z" }, + { url = "https://files.pythonhosted.org/packages/11/03/16090dd6f74ba2b8b922276047f15962fbeea0a75d5601607edb301ba945/llvmlite-0.47.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fa1cbd800edd3b20bc141521f7fd45a6185a5b84109aa6855134e81397ffe72b", size = 56275178, upload-time = "2026-03-31T18:29:42.58Z" }, + { url = "https://files.pythonhosted.org/packages/f5/cb/0abf1dd4c5286a95ffe0c1d8c67aec06b515894a0dd2ac97f5e27b82ab0b/llvmlite-0.47.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f6725179b89f03b17dabe236ff3422cb8291b4c1bf40af152826dfd34e350ae8", size = 55128632, upload-time = "2026-03-31T18:29:46.939Z" }, + { url = "https://files.pythonhosted.org/packages/4f/79/d3bbab197e86e0ff4f9c07122895b66a3e0d024247fcff7f12c473cb36d9/llvmlite-0.47.0-cp314-cp314t-win_amd64.whl", hash = "sha256:6842cf6f707ec4be3d985a385ad03f72b2d724439e118fcbe99b2929964f0453", size = 39153839, upload-time = "2026-03-31T18:29:51.004Z" }, ] [[package]] @@ -1747,36 +925,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/92/aa/df863bcc39c5e0946263454aba394de8a9084dbaff8ad143846b0d844739/lxml-6.0.2-cp314-cp314t-win_arm64.whl", hash = "sha256:bb4c1847b303835d89d785a18801a883436cdfd5dc3d62947f9c49e24f0f5a2c", size = 3822205, upload-time = "2025-09-22T04:03:36.249Z" }, ] -[[package]] -name = "m4-infra" -version = "0.0.0.dev0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "appdirs" }, - { name = "beautifulsoup4" }, - { name = "cryptography" }, - { name = "db-dtypes" }, - { name = "duckdb" }, - { name = "fastmcp" }, - { name = "google-cloud-bigquery" }, - { name = "httpx" }, - { name = "lifelines" }, - { name = "matplotlib" }, - { name = "pandas" }, - { name = "polars", extra = ["pyarrow"] }, - { name = "pyjwt", extra = ["crypto"] }, - { name = "python-jose", extra = ["cryptography"] }, - { name = "requests" }, - { name = "rich" }, - { name = "sqlalchemy" }, - { name = "sqlparse" }, - { name = "typer" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/da/77/77727fa69a5ec211e5c29ef58f1cb2389cb1b2e732435f2713d88fc98311/m4_infra-0.0.0.dev0.tar.gz", hash = "sha256:a9d13530253cf93eb7891eca73c159c511aae5086609002869d2d8db2b729d3d", size = 153780, upload-time = "2026-01-12T16:25:26.138Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f0/11/c2bb4ac0a1bd777df4550c791aef2b246c19496ecdb0c8172d2804863e1e/m4_infra-0.0.0.dev0-py3-none-any.whl", hash = "sha256:2de20a73c7c3c11b9d9af9a9ee577501e70a051dc8da753d54f24efd0d5adcc8", size = 152844, upload-time = "2026-01-12T16:25:24.671Z" }, -] - [[package]] name = "markdown-it-py" version = "4.0.0" @@ -1906,31 +1054,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5d/49/d651878698a0b67f23aa28e17f45a6d6dd3d3f933fa29087fa4ce5947b5a/matplotlib-3.10.8-cp314-cp314t-win_arm64.whl", hash = "sha256:113bb52413ea508ce954a02c10ffd0d565f9c3bc7f2eddc27dfe1731e71c7b5f", size = 8192560, upload-time = "2025-12-10T22:56:38.008Z" }, ] -[[package]] -name = "mcp" -version = "1.25.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "anyio" }, - { name = "httpx" }, - { name = "httpx-sse" }, - { name = "jsonschema" }, - { name = "pydantic" }, - { name = "pydantic-settings" }, - { name = "pyjwt", extra = ["crypto"] }, - { name = "python-multipart" }, - { name = "pywin32", marker = "sys_platform == 'win32'" }, - { name = "sse-starlette" }, - { name = "starlette" }, - { name = "typing-extensions" }, - { name = "typing-inspection" }, - { name = "uvicorn", marker = "sys_platform != 'emscripten'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/d5/2d/649d80a0ecf6a1f82632ca44bec21c0461a9d9fc8934d38cb5b319f2db5e/mcp-1.25.0.tar.gz", hash = "sha256:56310361ebf0364e2d438e5b45f7668cbb124e158bb358333cd06e49e83a6802", size = 605387, upload-time = "2025-12-19T10:19:56.985Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e2/fc/6dc7659c2ae5ddf280477011f4213a74f806862856b796ef08f028e664bf/mcp-1.25.0-py3-none-any.whl", hash = "sha256:b37c38144a666add0862614cc79ec276e97d72aa8ca26d622818d4e278b9721a", size = 233076, upload-time = "2025-12-19T10:19:55.416Z" }, -] - [[package]] name = "mdurl" version = "0.1.2" @@ -1940,15 +1063,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, ] -[[package]] -name = "more-itertools" -version = "10.8.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ea/5d/38b681d3fce7a266dd9ab73c66959406d565b3e85f21d5e66e1181d93721/more_itertools-10.8.0.tar.gz", hash = "sha256:f638ddf8a1a0d134181275fb5d58b086ead7c6a72429ad725c67503f13ba30bd", size = 137431, upload-time = "2025-09-02T15:23:11.018Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl", hash = "sha256:52d4362373dcf7c52546bc4af9a86ee7c4579df9a8dc268be0a2f949d376cc9b", size = 69667, upload-time = "2025-09-02T15:23:09.635Z" }, -] - [[package]] name = "mpmath" version = "1.3.0" @@ -2099,15 +1213,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505", size = 4963, upload-time = "2025-04-22T14:54:22.983Z" }, ] -[[package]] -name = "narwhals" -version = "2.15.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/47/6d/b57c64e5038a8cf071bce391bb11551657a74558877ac961e7fa905ece27/narwhals-2.15.0.tar.gz", hash = "sha256:a9585975b99d95084268445a1fdd881311fa26ef1caa18020d959d5b2ff9a965", size = 603479, upload-time = "2026-01-06T08:10:13.27Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3d/2e/cf2ffeb386ac3763526151163ad7da9f1b586aac96d2b4f7de1eaebf0c61/narwhals-2.15.0-py3-none-any.whl", hash = "sha256:cbfe21ca19d260d9fd67f995ec75c44592d1f106933b03ddd375df7ac841f9d6", size = 432856, upload-time = "2026-01-06T08:10:11.511Z" }, -] - [[package]] name = "networkx" version = "3.6" @@ -2126,6 +1231,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d2/1d/1b658dbd2b9fa9c4c9f32accbfc0205d532c8c6194dc0f2a4c0428e7128a/nodeenv-1.9.1-py2.py3-none-any.whl", hash = "sha256:ba11c9782d29c27c70ffbdda2d7415098754709be8a7056d79a737cd901155c9", size = 22314, upload-time = "2024-06-04T18:44:08.352Z" }, ] +[[package]] +name = "numba" +version = "0.65.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "llvmlite" }, + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/49/61/7299643b9c18d669e04be7c5bcb64d985070d07553274817b45b049e7bfe/numba-0.65.0.tar.gz", hash = "sha256:edad0d9f6682e93624c00125a471ae4df186175d71fd604c983c377cdc03e68b", size = 2764131, upload-time = "2026-04-01T03:52:01.946Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6c/2f/8bd31a1ea43c01ac215283d83aa5f8d5acbe7a36c85b82f1757bfe9ccb31/numba-0.65.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:b27ee4847e1bfb17e9604d100417ee7c1d10f15a6711c6213404b3da13a0b2aa", size = 2680705, upload-time = "2026-04-01T03:51:32.597Z" }, + { url = "https://files.pythonhosted.org/packages/73/36/88406bd58600cc696417b8e5dd6a056478da808f3eaf48d18e2421e0c2d9/numba-0.65.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a52d92ffd297c10364bce60cd1fcb88f99284ab5df085f2c6bcd1cb33b529a6f", size = 3801411, upload-time = "2026-04-01T03:51:34.321Z" }, + { url = "https://files.pythonhosted.org/packages/0c/61/ce753a1d7646dd477e16d15e89473703faebb8995d2f71d7ad69a540b565/numba-0.65.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:da8e371e328c06d0010c3d8b44b21858652831b85bcfba78cb22c042e22dbd8e", size = 3501622, upload-time = "2026-04-01T03:51:36.348Z" }, + { url = "https://files.pythonhosted.org/packages/7d/86/db87a5393f1b1fabef53ac3ba4e6b938bb27e40a04ad7cc512098fcae032/numba-0.65.0-cp312-cp312-win_amd64.whl", hash = "sha256:59bb9f2bb9f1238dfd8e927ba50645c18ae769fef4f3d58ea0ea22a2683b91f5", size = 2749979, upload-time = "2026-04-01T03:51:37.88Z" }, + { url = "https://files.pythonhosted.org/packages/8b/f8/eee0f1ff456218db036bfc9023995ec1f85a9dc8f2422f1594f6a87829e0/numba-0.65.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:c6334094563a456a695c812e6846288376ca02327cf246cdcc83e1bb27862367", size = 2680679, upload-time = "2026-04-01T03:51:39.491Z" }, + { url = "https://files.pythonhosted.org/packages/1b/8f/3d116e4b8e92f6abace431afa4b2b944f4d65bdee83af886f5c4b263df95/numba-0.65.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b8a9008411615c69d083d1dcf477f75a5aa727b30beb16e139799e2be945cdfd", size = 3809537, upload-time = "2026-04-01T03:51:41.42Z" }, + { url = "https://files.pythonhosted.org/packages/b5/2c/6a3ca4128e253cb67affe06deb47688f51ce968f5111e2a06d010e6f1fa6/numba-0.65.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:af96c0cba53664efcb361528b8c75e011a6556c859c7e08424c2715201c6cf7a", size = 3508615, upload-time = "2026-04-01T03:51:43.444Z" }, + { url = "https://files.pythonhosted.org/packages/96/0e/267f9a36fb282c104a971d7eecb685b411c47dce2a740fe69cf5fc2945d9/numba-0.65.0-cp313-cp313-win_amd64.whl", hash = "sha256:6254e73b9c929dc736a1fbd3d6f5680789709a5067cae1fa7198707385129c04", size = 2749938, upload-time = "2026-04-01T03:51:45.218Z" }, + { url = "https://files.pythonhosted.org/packages/56/a4/90edb01e9176053578e343d7a7276bc28356741ee67059aed8ed2c1a4e59/numba-0.65.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:ee336b398a6fca51b1f626034de99f50cb1bd87d537a166275158a3cee744b82", size = 2680878, upload-time = "2026-04-01T03:51:46.91Z" }, + { url = "https://files.pythonhosted.org/packages/24/8d/e12d6ff4b9119db3cbf7b2db1ce257576441bd3c76388c786dea74f20b02/numba-0.65.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:05c0a9fdf75d85f57dee47b719e8d6415707b80aae45d75f63f9dc1b935c29f7", size = 3778456, upload-time = "2026-04-01T03:51:48.552Z" }, + { url = "https://files.pythonhosted.org/packages/17/89/abcd83e76f6a773276fe76244140671bcc5bf820f6e2ae1a15362ae4c8c9/numba-0.65.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:583680e0e8faf124d362df23b4b593f3221a8996341a63d1b664c122401bec2f", size = 3478464, upload-time = "2026-04-01T03:51:50.527Z" }, + { url = "https://files.pythonhosted.org/packages/73/5b/fbce55ce3d933afbc7ade04df826853e4a846aaa47d58d2fbb669b8f2d08/numba-0.65.0-cp314-cp314-win_amd64.whl", hash = "sha256:add297d3e1c08dd884f44100152612fa41e66a51d15fdf91307f9dde31d06830", size = 2752012, upload-time = "2026-04-01T03:51:52.691Z" }, + { url = "https://files.pythonhosted.org/packages/1e/ab/af705f4257d9388fb2fd6d7416573e98b6ca9c786e8b58f02720978557bd/numba-0.65.0-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:194a243ba53a9157c8538cbb3166ec015d785a8c5d584d06cdd88bee902233c7", size = 2683961, upload-time = "2026-04-01T03:51:54.281Z" }, + { url = "https://files.pythonhosted.org/packages/ff/e5/8267b0adb0c01b52b553df5062fbbb42c30ed5362d08b85cc913a36f838f/numba-0.65.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c7fa502960f7a2f3f5cb025bc7bff888a3551277b92431bfdc5ba2f11a375749", size = 3816373, upload-time = "2026-04-01T03:51:56.18Z" }, + { url = "https://files.pythonhosted.org/packages/b0/f5/b8397ca360971669a93706b9274592b6864e4367a37d498fbbcb62aa2d48/numba-0.65.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5046c63f783ca3eb6195f826a50797465e7c4ce811daa17c9bea47e310c9b964", size = 3532782, upload-time = "2026-04-01T03:51:58.387Z" }, + { url = "https://files.pythonhosted.org/packages/f5/21/1e73fa16bf0393ebb74c5bb208d712152ffdfc84600a8e93a3180317856e/numba-0.65.0-cp314-cp314t-win_amd64.whl", hash = "sha256:46fd679ae4f68c7a5d5721efbd29ecee0b0f3013211591891d79b51bfdf73113", size = 2757611, upload-time = "2026-04-01T03:52:00.083Z" }, +] + [[package]] name = "numpy" version = "2.3.5" @@ -2296,6 +1429,7 @@ name = "nvidia-nccl-cu12" version = "2.27.5" source = { registry = "https://pypi.org/simple" } wheels = [ + { url = "https://files.pythonhosted.org/packages/bb/1c/857979db0ef194ca5e21478a0612bcdbbe59458d7694361882279947b349/nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:31432ad4d1fb1004eb0c56203dc9bc2178a1ba69d1d9e02d64a6938ab5e40e7a", size = 322400625, upload-time = "2025-06-26T04:11:04.496Z" }, { url = "https://files.pythonhosted.org/packages/6e/89/f7a07dc961b60645dbbf42e80f2bc85ade7feb9a491b11a1e973aa00071f/nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ad730cf15cb5d25fe849c6e6ca9eb5b76db16a80f13f425ac68d8e2e55624457", size = 322348229, upload-time = "2025-06-26T04:11:28.385Z" }, ] @@ -2336,87 +1470,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e3/94/1843518e420fa3ed6919835845df698c7e27e183cb997394e4a670973a65/omegaconf-2.3.0-py3-none-any.whl", hash = "sha256:7b4df175cdb08ba400f45cae3bdcae7ba8365db4d165fc65fd04b050ab63b46b", size = 79500, upload-time = "2022-12-08T20:59:19.686Z" }, ] -[[package]] -name = "openapi-pydantic" -version = "0.5.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pydantic" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/02/2e/58d83848dd1a79cb92ed8e63f6ba901ca282c5f09d04af9423ec26c56fd7/openapi_pydantic-0.5.1.tar.gz", hash = "sha256:ff6835af6bde7a459fb93eb93bb92b8749b754fc6e51b2f1590a19dc3005ee0d", size = 60892, upload-time = "2025-01-08T19:29:27.083Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/12/cf/03675d8bd8ecbf4445504d8071adab19f5f993676795708e36402ab38263/openapi_pydantic-0.5.1-py3-none-any.whl", hash = "sha256:a3a09ef4586f5bd760a8df7f43028b60cafb6d9f61de2acba9574766255ab146", size = 96381, upload-time = "2025-01-08T19:29:25.275Z" }, -] - -[[package]] -name = "opentelemetry-api" -version = "1.39.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "importlib-metadata" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/97/b9/3161be15bb8e3ad01be8be5a968a9237c3027c5be504362ff800fca3e442/opentelemetry_api-1.39.1.tar.gz", hash = "sha256:fbde8c80e1b937a2c61f20347e91c0c18a1940cecf012d62e65a7caf08967c9c", size = 65767, upload-time = "2025-12-11T13:32:39.182Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/cf/df/d3f1ddf4bb4cb50ed9b1139cc7b1c54c34a1e7ce8fd1b9a37c0d1551a6bd/opentelemetry_api-1.39.1-py3-none-any.whl", hash = "sha256:2edd8463432a7f8443edce90972169b195e7d6a05500cd29e6d13898187c9950", size = 66356, upload-time = "2025-12-11T13:32:17.304Z" }, -] - -[[package]] -name = "opentelemetry-exporter-prometheus" -version = "0.60b1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "opentelemetry-api" }, - { name = "opentelemetry-sdk" }, - { name = "prometheus-client" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/14/39/7dafa6fff210737267bed35a8855b6ac7399b9e582b8cf1f25f842517012/opentelemetry_exporter_prometheus-0.60b1.tar.gz", hash = "sha256:a4011b46906323f71724649d301b4dc188aaa068852e814f4df38cc76eac616b", size = 14976, upload-time = "2025-12-11T13:32:42.944Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9b/0d/4be6bf5477a3eb3d917d2f17d3c0b6720cd6cb97898444a61d43cc983f5c/opentelemetry_exporter_prometheus-0.60b1-py3-none-any.whl", hash = "sha256:49f59178de4f4590e3cef0b8b95cf6e071aae70e1f060566df5546fad773b8fd", size = 13019, upload-time = "2025-12-11T13:32:23.974Z" }, -] - -[[package]] -name = "opentelemetry-instrumentation" -version = "0.60b1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "opentelemetry-api" }, - { name = "opentelemetry-semantic-conventions" }, - { name = "packaging" }, - { name = "wrapt" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/41/0f/7e6b713ac117c1f5e4e3300748af699b9902a2e5e34c9cf443dde25a01fa/opentelemetry_instrumentation-0.60b1.tar.gz", hash = "sha256:57ddc7974c6eb35865af0426d1a17132b88b2ed8586897fee187fd5b8944bd6a", size = 31706, upload-time = "2025-12-11T13:36:42.515Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/77/d2/6788e83c5c86a2690101681aeef27eeb2a6bf22df52d3f263a22cee20915/opentelemetry_instrumentation-0.60b1-py3-none-any.whl", hash = "sha256:04480db952b48fb1ed0073f822f0ee26012b7be7c3eac1a3793122737c78632d", size = 33096, upload-time = "2025-12-11T13:35:33.067Z" }, -] - -[[package]] -name = "opentelemetry-sdk" -version = "1.39.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "opentelemetry-api" }, - { name = "opentelemetry-semantic-conventions" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/eb/fb/c76080c9ba07e1e8235d24cdcc4d125ef7aa3edf23eb4e497c2e50889adc/opentelemetry_sdk-1.39.1.tar.gz", hash = "sha256:cf4d4563caf7bff906c9f7967e2be22d0d6b349b908be0d90fb21c8e9c995cc6", size = 171460, upload-time = "2025-12-11T13:32:49.369Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7c/98/e91cf858f203d86f4eccdf763dcf01cf03f1dae80c3750f7e635bfa206b6/opentelemetry_sdk-1.39.1-py3-none-any.whl", hash = "sha256:4d5482c478513ecb0a5d938dcc61394e647066e0cc2676bee9f3af3f3f45f01c", size = 132565, upload-time = "2025-12-11T13:32:35.069Z" }, -] - -[[package]] -name = "opentelemetry-semantic-conventions" -version = "0.60b1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "opentelemetry-api" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/91/df/553f93ed38bf22f4b999d9be9c185adb558982214f33eae539d3b5cd0858/opentelemetry_semantic_conventions-0.60b1.tar.gz", hash = "sha256:87c228b5a0669b748c76d76df6c364c369c28f1c465e50f661e39737e84bc953", size = 137935, upload-time = "2025-12-11T13:32:50.487Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7a/5e/5958555e09635d09b75de3c4f8b9cae7335ca545d77392ffe7331534c402/opentelemetry_semantic_conventions-0.60b1-py3-none-any.whl", hash = "sha256:9fa8c8b0c110da289809292b0591220d3a7b53c1526a23021e977d68597893fb", size = 219982, upload-time = "2025-12-11T13:32:36.955Z" }, -] - [[package]] name = "packaging" version = "25.0" @@ -2473,15 +1526,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/70/44/5191d2e4026f86a2a109053e194d3ba7a31a2d10a9c2348368c63ed4e85a/pandas-2.3.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:3869faf4bd07b3b66a9f462417d0ca3a9df29a9f6abd5d0d0dbab15dac7abe87", size = 13202175, upload-time = "2025-09-29T23:31:59.173Z" }, ] -[[package]] -name = "pathable" -version = "0.4.4" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/67/93/8f2c2075b180c12c1e9f6a09d1a985bc2036906b13dff1d8917e395f2048/pathable-0.4.4.tar.gz", hash = "sha256:6905a3cd17804edfac7875b5f6c9142a218c7caef78693c2dbbbfbac186d88b2", size = 8124, upload-time = "2025-01-10T18:43:13.247Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7d/eb/b6260b31b1a96386c0a880edebe26f89669098acea8e0318bff6adb378fd/pathable-0.4.4-py3-none-any.whl", hash = "sha256:5ae9e94793b6ef5a4cbe0a7ce9dbbefc1eec38df253763fd0aeeacf2762dbbc2", size = 9592, upload-time = "2025-01-10T18:43:11.88Z" }, -] - [[package]] name = "pathspec" version = "0.12.1" @@ -2491,15 +1535,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cc/20/ff623b09d963f88bfde16306a54e12ee5ea43e9b597108672ff3a408aad6/pathspec-0.12.1-py3-none-any.whl", hash = "sha256:a0d503e138a4c123b27490a4f7beda6a01c6f288df0e4a8b79c7eb0dc7b4cc08", size = 31191, upload-time = "2023-12-10T22:30:43.14Z" }, ] -[[package]] -name = "pathvalidate" -version = "3.3.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/fa/2a/52a8da6fe965dea6192eb716b357558e103aea0a1e9a8352ad575a8406ca/pathvalidate-3.3.1.tar.gz", hash = "sha256:b18c07212bfead624345bb8e1d6141cdcf15a39736994ea0b94035ad2b1ba177", size = 63262, upload-time = "2025-06-15T09:07:20.736Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9a/70/875f4a23bfc4731703a5835487d0d2fb999031bd415e7d17c0ae615c18b7/pathvalidate-3.3.1-py3-none-any.whl", hash = "sha256:5263baab691f8e1af96092fa5137ee17df5bdfbd6cff1fcac4d6ef4bc2e1735f", size = 24305, upload-time = "2025-06-15T09:07:19.117Z" }, -] - [[package]] name = "pillow" version = "12.0.0" @@ -2599,11 +1634,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b4/9a/24e4b890c7ee4358964aa92c4d1865df0e8831f7df6abaa3a39914521724/polars-1.35.2-py3-none-any.whl", hash = "sha256:5e8057c8289ac148c793478323b726faea933d9776bd6b8a554b0ab7c03db87e", size = 783597, upload-time = "2025-11-09T13:18:51.361Z" }, ] -[package.optional-dependencies] -pyarrow = [ - { name = "pyarrow" }, -] - [[package]] name = "polars-runtime-32" version = "1.35.2" @@ -2646,15 +1676,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5d/c4/b2d28e9d2edf4f1713eb3c29307f1a63f3d67cf09bdda29715a36a68921a/pre_commit-4.5.0-py2.py3-none-any.whl", hash = "sha256:25e2ce09595174d9c97860a95609f9f852c0614ba602de3561e267547f2335e1", size = 226429, upload-time = "2025-11-22T21:02:40.836Z" }, ] -[[package]] -name = "prometheus-client" -version = "0.23.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/23/53/3edb5d68ecf6b38fcbcc1ad28391117d2a322d9a1a3eff04bfdb184d8c3b/prometheus_client-0.23.1.tar.gz", hash = "sha256:6ae8f9081eaaaf153a2e959d2e6c4f4fb57b12ef76c8c7980202f1e57b48b2ce", size = 80481, upload-time = "2025-09-18T20:47:25.043Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b8/db/14bafcb4af2139e046d03fd00dea7873e48eafe18b7d2797e73d6681f210/prometheus_client-0.23.1-py3-none-any.whl", hash = "sha256:dd1913e6e76b59cfe44e7a4b83e01afc9873c1bdfd2ed8739f1e76aeca115f99", size = 61145, upload-time = "2025-09-18T20:47:23.875Z" }, -] - [[package]] name = "propcache" version = "0.4.1" @@ -2732,23 +1753,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cf/12/96e4664c82ca2f31e1c8dff86afb867348979eb78d3cb8546a680287a1e9/propcache-0.4.1-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:521a463429ef54143092c11a77e04056dd00636f72e8c45b70aaa3140d639726", size = 247641, upload-time = "2025-10-08T19:48:27.207Z" }, { url = "https://files.pythonhosted.org/packages/18/ed/e7a9cfca28133386ba52278136d42209d3125db08d0a6395f0cba0c0285c/propcache-0.4.1-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:120c964da3fdc75e3731aa392527136d4ad35868cc556fd09bb6d09172d9a367", size = 262510, upload-time = "2025-10-08T19:48:28.65Z" }, { url = "https://files.pythonhosted.org/packages/f5/76/16d8bf65e8845dd62b4e2b57444ab81f07f40caa5652b8969b87ddcf2ef6/propcache-0.4.1-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:d8f353eb14ee3441ee844ade4277d560cdd68288838673273b978e3d6d2c8f36", size = 263161, upload-time = "2025-10-08T19:48:30.133Z" }, - { url = "https://files.pythonhosted.org/packages/e7/70/c99e9edb5d91d5ad8a49fa3c1e8285ba64f1476782fed10ab251ff413ba1/propcache-0.4.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:ab2943be7c652f09638800905ee1bab2c544e537edb57d527997a24c13dc1455", size = 257393, upload-time = "2025-10-08T19:48:31.567Z" }, - { url = "https://files.pythonhosted.org/packages/08/02/87b25304249a35c0915d236575bc3574a323f60b47939a2262b77632a3ee/propcache-0.4.1-cp314-cp314t-win32.whl", hash = "sha256:05674a162469f31358c30bcaa8883cb7829fa3110bf9c0991fe27d7896c42d85", size = 42546, upload-time = "2025-10-08T19:48:32.872Z" }, - { url = "https://files.pythonhosted.org/packages/cb/ef/3c6ecf8b317aa982f309835e8f96987466123c6e596646d4e6a1dfcd080f/propcache-0.4.1-cp314-cp314t-win_amd64.whl", hash = "sha256:990f6b3e2a27d683cb7602ed6c86f15ee6b43b1194736f9baaeb93d0016633b1", size = 46259, upload-time = "2025-10-08T19:48:34.226Z" }, - { url = "https://files.pythonhosted.org/packages/c4/2d/346e946d4951f37eca1e4f55be0f0174c52cd70720f84029b02f296f4a38/propcache-0.4.1-cp314-cp314t-win_arm64.whl", hash = "sha256:ecef2343af4cc68e05131e45024ba34f6095821988a9d0a02aa7c73fcc448aa9", size = 40428, upload-time = "2025-10-08T19:48:35.441Z" }, - { url = "https://files.pythonhosted.org/packages/5b/5a/bc7b4a4ef808fa59a816c17b20c4bef6884daebbdf627ff2a161da67da19/propcache-0.4.1-py3-none-any.whl", hash = "sha256:af2a6052aeb6cf17d3e46ee169099044fd8224cbaf75c76a2ef596e8163e2237", size = 13305, upload-time = "2025-10-08T19:49:00.792Z" }, -] - -[[package]] -name = "proto-plus" -version = "1.27.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "protobuf" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/01/89/9cbe2f4bba860e149108b683bc2efec21f14d5f7ed6e25562ad86acbc373/proto_plus-1.27.0.tar.gz", hash = "sha256:873af56dd0d7e91836aee871e5799e1c6f1bda86ac9a983e0bb9f0c266a568c4", size = 56158, upload-time = "2025-12-16T13:46:25.729Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/cd/24/3b7a0818484df9c28172857af32c2397b6d8fcd99d9468bd4684f98ebf0a/proto_plus-1.27.0-py3-none-any.whl", hash = "sha256:1baa7f81cf0f8acb8bc1f6d085008ba4171eaf669629d1b6d1673b21ed1c0a82", size = 50205, upload-time = "2025-12-16T13:46:24.76Z" }, + { url = "https://files.pythonhosted.org/packages/e7/70/c99e9edb5d91d5ad8a49fa3c1e8285ba64f1476782fed10ab251ff413ba1/propcache-0.4.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:ab2943be7c652f09638800905ee1bab2c544e537edb57d527997a24c13dc1455", size = 257393, upload-time = "2025-10-08T19:48:31.567Z" }, + { url = "https://files.pythonhosted.org/packages/08/02/87b25304249a35c0915d236575bc3574a323f60b47939a2262b77632a3ee/propcache-0.4.1-cp314-cp314t-win32.whl", hash = "sha256:05674a162469f31358c30bcaa8883cb7829fa3110bf9c0991fe27d7896c42d85", size = 42546, upload-time = "2025-10-08T19:48:32.872Z" }, + { url = "https://files.pythonhosted.org/packages/cb/ef/3c6ecf8b317aa982f309835e8f96987466123c6e596646d4e6a1dfcd080f/propcache-0.4.1-cp314-cp314t-win_amd64.whl", hash = "sha256:990f6b3e2a27d683cb7602ed6c86f15ee6b43b1194736f9baaeb93d0016633b1", size = 46259, upload-time = "2025-10-08T19:48:34.226Z" }, + { url = "https://files.pythonhosted.org/packages/c4/2d/346e946d4951f37eca1e4f55be0f0174c52cd70720f84029b02f296f4a38/propcache-0.4.1-cp314-cp314t-win_arm64.whl", hash = "sha256:ecef2343af4cc68e05131e45024ba34f6095821988a9d0a02aa7c73fcc448aa9", size = 40428, upload-time = "2025-10-08T19:48:35.441Z" }, + { url = "https://files.pythonhosted.org/packages/5b/5a/bc7b4a4ef808fa59a816c17b20c4bef6884daebbdf627ff2a161da67da19/propcache-0.4.1-py3-none-any.whl", hash = "sha256:af2a6052aeb6cf17d3e46ee169099044fd8224cbaf75c76a2ef596e8163e2237", size = 13305, upload-time = "2025-10-08T19:49:00.792Z" }, ] [[package]] @@ -2766,47 +1775,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/08/b4/46310463b4f6ceef310f8348786f3cff181cea671578e3d9743ba61a459e/protobuf-6.33.1-py3-none-any.whl", hash = "sha256:d595a9fd694fdeb061a62fbe10eb039cc1e444df81ec9bb70c7fc59ebcb1eafa", size = 170477, upload-time = "2025-11-13T16:44:17.633Z" }, ] -[[package]] -name = "py-key-value-aio" -version = "0.3.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "beartype" }, - { name = "py-key-value-shared" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/93/ce/3136b771dddf5ac905cc193b461eb67967cf3979688c6696e1f2cdcde7ea/py_key_value_aio-0.3.0.tar.gz", hash = "sha256:858e852fcf6d696d231266da66042d3355a7f9871650415feef9fca7a6cd4155", size = 50801, upload-time = "2025-11-17T16:50:04.711Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/99/10/72f6f213b8f0bce36eff21fda0a13271834e9eeff7f9609b01afdc253c79/py_key_value_aio-0.3.0-py3-none-any.whl", hash = "sha256:1c781915766078bfd608daa769fefb97e65d1d73746a3dfb640460e322071b64", size = 96342, upload-time = "2025-11-17T16:50:03.801Z" }, -] - -[package.optional-dependencies] -disk = [ - { name = "diskcache" }, - { name = "pathvalidate" }, -] -keyring = [ - { name = "keyring" }, -] -memory = [ - { name = "cachetools" }, -] -redis = [ - { name = "redis" }, -] - -[[package]] -name = "py-key-value-shared" -version = "0.3.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "beartype" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/7b/e4/1971dfc4620a3a15b4579fe99e024f5edd6e0967a71154771a059daff4db/py_key_value_shared-0.3.0.tar.gz", hash = "sha256:8fdd786cf96c3e900102945f92aa1473138ebe960ef49da1c833790160c28a4b", size = 11666, upload-time = "2025-11-17T16:50:06.849Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/51/e4/b8b0a03ece72f47dce2307d36e1c34725b7223d209fc679315ffe6a4e2c3/py_key_value_shared-0.3.0-py3-none-any.whl", hash = "sha256:5b0efba7ebca08bb158b1e93afc2f07d30b8f40c2fc12ce24a4c0d84f42f9298", size = 19560, upload-time = "2025-11-17T16:50:05.954Z" }, -] - [[package]] name = "pyarrow" version = "22.0.0" @@ -2850,36 +1818,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7b/03/f335d6c52b4a4761bcc83499789a1e2e16d9d201a58c327a9b5cc9a41bd9/pyarrow-22.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:0c34fe18094686194f204a3b1787a27456897d8a2d62caf84b61e8dfbc0252ae", size = 29185594, upload-time = "2025-10-24T10:09:53.111Z" }, ] -[[package]] -name = "pyasn1" -version = "0.6.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ba/e9/01f1a64245b89f039897cb0130016d79f77d52669aae6ee7b159a6c4c018/pyasn1-0.6.1.tar.gz", hash = "sha256:6f580d2bdd84365380830acf45550f2511469f673cb4a5ae3857a3170128b034", size = 145322, upload-time = "2024-09-10T22:41:42.55Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c8/f1/d6a797abb14f6283c0ddff96bbdd46937f64122b8c925cab503dd37f8214/pyasn1-0.6.1-py3-none-any.whl", hash = "sha256:0d632f46f2ba09143da3a8afe9e33fb6f92fa2320ab7e886e2d0f7672af84629", size = 83135, upload-time = "2024-09-11T16:00:36.122Z" }, -] - -[[package]] -name = "pyasn1-modules" -version = "0.4.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pyasn1" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/e9/e6/78ebbb10a8c8e4b61a59249394a4a594c1a7af95593dc933a349c8d00964/pyasn1_modules-0.4.2.tar.gz", hash = "sha256:677091de870a80aae844b1ca6134f54652fa2c8c5a52aa396440ac3106e941e6", size = 307892, upload-time = "2025-03-28T02:41:22.17Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/47/8d/d529b5d697919ba8c11ad626e835d4039be708a35b0d22de83a269a6682c/pyasn1_modules-0.4.2-py3-none-any.whl", hash = "sha256:29253a9207ce32b64c3ac6600edc75368f98473906e8fd1043bd6b5b1de2c14a", size = 181259, upload-time = "2025-03-28T02:41:19.028Z" }, -] - -[[package]] -name = "pycparser" -version = "2.23" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/fe/cf/d2d3b9f5699fb1e4615c8e32ff220203e43b248e1dfcc6736ad9057731ca/pycparser-2.23.tar.gz", hash = "sha256:78816d4f24add8f10a06d6f05b4d424ad9e96cfebf68a4ddc99c65c0720d00c2", size = 173734, upload-time = "2025-09-09T13:23:47.91Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a0/e3/59cd50310fc9b59512193629e1984c1f95e5c8ae6e5d8c69532ccc65a7fe/pycparser-2.23-py3-none-any.whl", hash = "sha256:e5c6e8d3fbad53479cab09ac03729e0a9faf2bee3db8208a550daf5af81a5934", size = 118140, upload-time = "2025-09-09T13:23:46.651Z" }, -] - [[package]] name = "pydantic" version = "2.12.5" @@ -2895,11 +1833,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5a/87/b70ad306ebb6f9b585f114d0ac2137d792b48be34d732d60e597c2f8465a/pydantic-2.12.5-py3-none-any.whl", hash = "sha256:e561593fccf61e8a20fc46dfc2dfe075b8be7d0188df33f221ad1f0139180f9d", size = 463580, upload-time = "2025-11-26T15:11:44.605Z" }, ] -[package.optional-dependencies] -email = [ - { name = "email-validator" }, -] - [[package]] name = "pydantic-core" version = "2.41.5" @@ -2971,43 +1904,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f7/07/34573da085946b6a313d7c42f82f16e8920bfd730665de2d11c0c37a74b5/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:76d0819de158cd855d1cbb8fcafdf6f5cf1eb8e470abe056d5d161106e38062b", size = 2139017, upload-time = "2025-11-04T13:42:59.471Z" }, ] -[[package]] -name = "pydantic-settings" -version = "2.12.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pydantic" }, - { name = "python-dotenv" }, - { name = "typing-inspection" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/43/4b/ac7e0aae12027748076d72a8764ff1c9d82ca75a7a52622e67ed3f765c54/pydantic_settings-2.12.0.tar.gz", hash = "sha256:005538ef951e3c2a68e1c08b292b5f2e71490def8589d4221b95dab00dafcfd0", size = 194184, upload-time = "2025-11-10T14:25:47.013Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c1/60/5d4751ba3f4a40a6891f24eec885f51afd78d208498268c734e256fb13c4/pydantic_settings-2.12.0-py3-none-any.whl", hash = "sha256:fddb9fd99a5b18da837b29710391e945b1e30c135477f484084ee513adb93809", size = 51880, upload-time = "2025-11-10T14:25:45.546Z" }, -] - -[[package]] -name = "pydocket" -version = "0.16.6" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cloudpickle" }, - { name = "fakeredis", extra = ["lua"] }, - { name = "opentelemetry-api" }, - { name = "opentelemetry-exporter-prometheus" }, - { name = "opentelemetry-instrumentation" }, - { name = "prometheus-client" }, - { name = "py-key-value-aio", extra = ["memory", "redis"] }, - { name = "python-json-logger" }, - { name = "redis" }, - { name = "rich" }, - { name = "typer" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/72/00/26befe5f58df7cd1aeda4a8d10bc7d1908ffd86b80fd995e57a2a7b3f7bd/pydocket-0.16.6.tar.gz", hash = "sha256:b96c96ad7692827214ed4ff25fcf941ec38371314db5dcc1ae792b3e9d3a0294", size = 299054, upload-time = "2026-01-09T22:09:15.405Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/0a/3f/7483e5a6dc6326b6e0c640619b5c5bd1d6e3c20e54d58f5fb86267cef00e/pydocket-0.16.6-py3-none-any.whl", hash = "sha256:683d21e2e846aa5106274e7d59210331b242d7fb0dce5b08d3b82065663ed183", size = 67697, upload-time = "2026-01-09T22:09:13.436Z" }, -] - [[package]] name = "pygments" version = "2.19.2" @@ -3018,17 +1914,19 @@ wheels = [ ] [[package]] -name = "pyjwt" -version = "2.10.1" +name = "pynndescent" +version = "0.6.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e7/46/bd74733ff231675599650d3e47f361794b22ef3e3770998dda30d3b63726/pyjwt-2.10.1.tar.gz", hash = "sha256:3cc5772eb20009233caf06e9d8a0577824723b44e6648ee0a2aedb6cf9381953", size = 87785, upload-time = "2024-11-28T03:43:29.933Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/61/ad/689f02752eeec26aed679477e80e632ef1b682313be70793d798c1d5fc8f/PyJWT-2.10.1-py3-none-any.whl", hash = "sha256:dcdd193e30abefd5debf142f9adfcdd2b58004e644f25406ffaebd50bd98dacb", size = 22997, upload-time = "2024-11-28T03:43:27.893Z" }, +dependencies = [ + { name = "joblib" }, + { name = "llvmlite" }, + { name = "numba" }, + { name = "scikit-learn" }, + { name = "scipy" }, ] - -[package.optional-dependencies] -crypto = [ - { name = "cryptography" }, +sdist = { url = "https://files.pythonhosted.org/packages/4a/fb/7f58c397fb31666756457ee2ac4c0289ef2daad57f4ae4be8dec12f80b03/pynndescent-0.6.0.tar.gz", hash = "sha256:7ffde0fb5b400741e055a9f7d377e3702e02250616834231f6c209e39aac24f5", size = 2992987, upload-time = "2026-01-08T21:29:58.943Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b2/e6/94145d714402fd5ade00b5661f2d0ab981219e07f7db9bfa16786cdb9c04/pynndescent-0.6.0-py3-none-any.whl", hash = "sha256:dc8c74844e4c7f5cbd1e0cd6909da86fdc789e6ff4997336e344779c3d5538ef", size = 73511, upload-time = "2026-01-08T21:29:57.306Z" }, ] [[package]] @@ -3040,15 +1938,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/10/5e/1aa9a93198c6b64513c9d7752de7422c06402de6600a8767da1524f9570b/pyparsing-3.2.5-py3-none-any.whl", hash = "sha256:e38a4f02064cf41fe6593d328d0512495ad1f3d8a91c4f73fc401b3079a59a5e", size = 113890, upload-time = "2025-09-21T04:11:04.117Z" }, ] -[[package]] -name = "pyperclip" -version = "1.11.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e8/52/d87eba7cb129b81563019d1679026e7a112ef76855d6159d24754dbd2a51/pyperclip-1.11.0.tar.gz", hash = "sha256:244035963e4428530d9e3a6101a1ef97209c6825edab1567beac148ccc1db1b6", size = 12185, upload-time = "2025-09-26T14:40:37.245Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/df/80/fc9d01d5ed37ba4c42ca2b55b4339ae6e200b456be3a1aaddf4a9fa99b8c/pyperclip-1.11.0-py3-none-any.whl", hash = "sha256:299403e9ff44581cb9ba2ffeed69c7aa96a008622ad0c46cb575ca75b5b84273", size = 11063, upload-time = "2025-09-26T14:40:36.069Z" }, -] - [[package]] name = "pytest" version = "9.0.1" @@ -3091,52 +1980,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, ] -[[package]] -name = "python-dotenv" -version = "1.2.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/f0/26/19cadc79a718c5edbec86fd4919a6b6d3f681039a2f6d66d14be94e75fb9/python_dotenv-1.2.1.tar.gz", hash = "sha256:42667e897e16ab0d66954af0e60a9caa94f0fd4ecf3aaf6d2d260eec1aa36ad6", size = 44221, upload-time = "2025-10-26T15:12:10.434Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/14/1b/a298b06749107c305e1fe0f814c6c74aea7b2f1e10989cb30f544a1b3253/python_dotenv-1.2.1-py3-none-any.whl", hash = "sha256:b81ee9561e9ca4004139c6cbba3a238c32b03e4894671e181b671e8cb8425d61", size = 21230, upload-time = "2025-10-26T15:12:09.109Z" }, -] - -[[package]] -name = "python-jose" -version = "3.5.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "ecdsa" }, - { name = "pyasn1" }, - { name = "rsa" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/c6/77/3a1c9039db7124eb039772b935f2244fbb73fc8ee65b9acf2375da1c07bf/python_jose-3.5.0.tar.gz", hash = "sha256:fb4eaa44dbeb1c26dcc69e4bd7ec54a1cb8dd64d3b4d81ef08d90ff453f2b01b", size = 92726, upload-time = "2025-05-28T17:31:54.288Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d9/c3/0bd11992072e6a1c513b16500a5d07f91a24017c5909b02c72c62d7ad024/python_jose-3.5.0-py2.py3-none-any.whl", hash = "sha256:abd1202f23d34dfad2c3d28cb8617b90acf34132c7afd60abd0b0b7d3cb55771", size = 34624, upload-time = "2025-05-28T17:31:52.802Z" }, -] - -[package.optional-dependencies] -cryptography = [ - { name = "cryptography" }, -] - -[[package]] -name = "python-json-logger" -version = "4.0.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/29/bf/eca6a3d43db1dae7070f70e160ab20b807627ba953663ba07928cdd3dc58/python_json_logger-4.0.0.tar.gz", hash = "sha256:f58e68eb46e1faed27e0f574a55a0455eecd7b8a5b88b85a784519ba3cff047f", size = 17683, upload-time = "2025-10-06T04:15:18.984Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/51/e5/fecf13f06e5e5f67e8837d777d1bc43fac0ed2b77a676804df5c34744727/python_json_logger-4.0.0-py3-none-any.whl", hash = "sha256:af09c9daf6a813aa4cc7180395f50f2a9e5fa056034c9953aec92e381c5ba1e2", size = 15548, upload-time = "2025-10-06T04:15:17.553Z" }, -] - -[[package]] -name = "python-multipart" -version = "0.0.21" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/78/96/804520d0850c7db98e5ccb70282e29208723f0964e88ffd9d0da2f52ea09/python_multipart-0.0.21.tar.gz", hash = "sha256:7137ebd4d3bbf70ea1622998f902b97a29434a9e8dc40eb203bbcf7c2a2cba92", size = 37196, upload-time = "2025-12-17T09:24:22.446Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/aa/76/03af049af4dcee5d27442f71b6924f01f3efb5d2bd34f23fcd563f2cc5f5/python_multipart-0.0.21-py3-none-any.whl", hash = "sha256:cf7a6713e01c87aa35387f4774e812c4361150938d20d232800f75ffcf266090", size = 24541, upload-time = "2025-12-17T09:24:21.153Z" }, -] - [[package]] name = "python-pptx" version = "1.0.2" @@ -3205,15 +2048,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c0/d2/21af5c535501a7233e734b8af901574572da66fcc254cb35d0609c9080dd/pywin32-311-cp314-cp314-win_arm64.whl", hash = "sha256:a508e2d9025764a8270f93111a970e1d0fbfc33f4153b388bb649b7eec4f9b42", size = 8932540, upload-time = "2025-07-14T20:13:36.379Z" }, ] -[[package]] -name = "pywin32-ctypes" -version = "0.2.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/85/9f/01a1a99704853cb63f253eea009390c88e7131c67e66a0a02099a8c917cb/pywin32-ctypes-0.2.3.tar.gz", hash = "sha256:d162dc04946d704503b2edc4d55f3dba5c1d539ead017afa00142c38b9885755", size = 29471, upload-time = "2024-08-14T10:15:34.626Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/de/3d/8161f7711c017e01ac9f008dfddd9410dff3674334c233bde66e7ba65bbf/pywin32_ctypes-0.2.3-py3-none-any.whl", hash = "sha256:8a1513379d709975552d202d942d9837758905c8d01eb82b8bcc30918929e7b8", size = 30756, upload-time = "2024-08-14T10:15:33.187Z" }, -] - [[package]] name = "pyyaml" version = "6.0.3" @@ -3260,29 +2094,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" }, ] -[[package]] -name = "redis" -version = "7.1.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/43/c8/983d5c6579a411d8a99bc5823cc5712768859b5ce2c8afe1a65b37832c81/redis-7.1.0.tar.gz", hash = "sha256:b1cc3cfa5a2cb9c2ab3ba700864fb0ad75617b41f01352ce5779dabf6d5f9c3c", size = 4796669, upload-time = "2025-11-19T15:54:39.961Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/89/f0/8956f8a86b20d7bb9d6ac0187cf4cd54d8065bc9a1a09eb8011d4d326596/redis-7.1.0-py3-none-any.whl", hash = "sha256:23c52b208f92b56103e17c5d06bdc1a6c2c0b3106583985a76a18f83b265de2b", size = 354159, upload-time = "2025-11-19T15:54:38.064Z" }, -] - -[[package]] -name = "referencing" -version = "0.36.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "attrs" }, - { name = "rpds-py" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/2f/db/98b5c277be99dd18bfd91dd04e1b759cad18d1a338188c936e92f921c7e2/referencing-0.36.2.tar.gz", hash = "sha256:df2e89862cd09deabbdba16944cc3f10feb6b3e6f18e902f7cc25609a34775aa", size = 74744, upload-time = "2025-01-25T08:48:16.138Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c1/b1/3baf80dc6d2b7bc27a95a67752d0208e410351e3feb4eb78de5f77454d8d/referencing-0.36.2-py3-none-any.whl", hash = "sha256:e8699adbbf8b5c7de96d8ffa0eb5c158b3beafce084968e2ea8bb08c6794dcd0", size = 26775, upload-time = "2025-01-25T08:48:14.241Z" }, -] - [[package]] name = "requests" version = "2.32.5" @@ -3311,112 +2122,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/25/7a/b0178788f8dc6cafce37a212c99565fa1fe7872c70c6c9c1e1a372d9d88f/rich-14.2.0-py3-none-any.whl", hash = "sha256:76bc51fe2e57d2b1be1f96c524b890b816e334ab4c1e45888799bfaab0021edd", size = 243393, upload-time = "2025-10-09T14:16:51.245Z" }, ] -[[package]] -name = "rich-rst" -version = "1.3.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "docutils" }, - { name = "rich" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/bc/6d/a506aaa4a9eaa945ed8ab2b7347859f53593864289853c5d6d62b77246e0/rich_rst-1.3.2.tar.gz", hash = "sha256:a1196fdddf1e364b02ec68a05e8ff8f6914fee10fbca2e6b6735f166bb0da8d4", size = 14936, upload-time = "2025-10-14T16:49:45.332Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/13/2f/b4530fbf948867702d0a3f27de4a6aab1d156f406d72852ab902c4d04de9/rich_rst-1.3.2-py3-none-any.whl", hash = "sha256:a99b4907cbe118cf9d18b0b44de272efa61f15117c61e39ebdc431baf5df722a", size = 12567, upload-time = "2025-10-14T16:49:42.953Z" }, -] - -[[package]] -name = "rpds-py" -version = "0.30.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/20/af/3f2f423103f1113b36230496629986e0ef7e199d2aa8392452b484b38ced/rpds_py-0.30.0.tar.gz", hash = "sha256:dd8ff7cf90014af0c0f787eea34794ebf6415242ee1d6fa91eaba725cc441e84", size = 69469, upload-time = "2025-11-30T20:24:38.837Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/03/e7/98a2f4ac921d82f33e03f3835f5bf3a4a40aa1bfdc57975e74a97b2b4bdd/rpds_py-0.30.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:a161f20d9a43006833cd7068375a94d035714d73a172b681d8881820600abfad", size = 375086, upload-time = "2025-11-30T20:22:17.93Z" }, - { url = "https://files.pythonhosted.org/packages/4d/a1/bca7fd3d452b272e13335db8d6b0b3ecde0f90ad6f16f3328c6fb150c889/rpds_py-0.30.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6abc8880d9d036ecaafe709079969f56e876fcf107f7a8e9920ba6d5a3878d05", size = 359053, upload-time = "2025-11-30T20:22:19.297Z" }, - { url = "https://files.pythonhosted.org/packages/65/1c/ae157e83a6357eceff62ba7e52113e3ec4834a84cfe07fa4b0757a7d105f/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ca28829ae5f5d569bb62a79512c842a03a12576375d5ece7d2cadf8abe96ec28", size = 390763, upload-time = "2025-11-30T20:22:21.661Z" }, - { url = "https://files.pythonhosted.org/packages/d4/36/eb2eb8515e2ad24c0bd43c3ee9cd74c33f7ca6430755ccdb240fd3144c44/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:a1010ed9524c73b94d15919ca4d41d8780980e1765babf85f9a2f90d247153dd", size = 408951, upload-time = "2025-11-30T20:22:23.408Z" }, - { url = "https://files.pythonhosted.org/packages/d6/65/ad8dc1784a331fabbd740ef6f71ce2198c7ed0890dab595adb9ea2d775a1/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f8d1736cfb49381ba528cd5baa46f82fdc65c06e843dab24dd70b63d09121b3f", size = 514622, upload-time = "2025-11-30T20:22:25.16Z" }, - { url = "https://files.pythonhosted.org/packages/63/8e/0cfa7ae158e15e143fe03993b5bcd743a59f541f5952e1546b1ac1b5fd45/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:d948b135c4693daff7bc2dcfc4ec57237a29bd37e60c2fabf5aff2bbacf3e2f1", size = 414492, upload-time = "2025-11-30T20:22:26.505Z" }, - { url = "https://files.pythonhosted.org/packages/60/1b/6f8f29f3f995c7ffdde46a626ddccd7c63aefc0efae881dc13b6e5d5bb16/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:47f236970bccb2233267d89173d3ad2703cd36a0e2a6e92d0560d333871a3d23", size = 394080, upload-time = "2025-11-30T20:22:27.934Z" }, - { url = "https://files.pythonhosted.org/packages/6d/d5/a266341051a7a3ca2f4b750a3aa4abc986378431fc2da508c5034d081b70/rpds_py-0.30.0-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:2e6ecb5a5bcacf59c3f912155044479af1d0b6681280048b338b28e364aca1f6", size = 408680, upload-time = "2025-11-30T20:22:29.341Z" }, - { url = "https://files.pythonhosted.org/packages/10/3b/71b725851df9ab7a7a4e33cf36d241933da66040d195a84781f49c50490c/rpds_py-0.30.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:a8fa71a2e078c527c3e9dc9fc5a98c9db40bcc8a92b4e8858e36d329f8684b51", size = 423589, upload-time = "2025-11-30T20:22:31.469Z" }, - { url = "https://files.pythonhosted.org/packages/00/2b/e59e58c544dc9bd8bd8384ecdb8ea91f6727f0e37a7131baeff8d6f51661/rpds_py-0.30.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:73c67f2db7bc334e518d097c6d1e6fed021bbc9b7d678d6cc433478365d1d5f5", size = 573289, upload-time = "2025-11-30T20:22:32.997Z" }, - { url = "https://files.pythonhosted.org/packages/da/3e/a18e6f5b460893172a7d6a680e86d3b6bc87a54c1f0b03446a3c8c7b588f/rpds_py-0.30.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:5ba103fb455be00f3b1c2076c9d4264bfcb037c976167a6047ed82f23153f02e", size = 599737, upload-time = "2025-11-30T20:22:34.419Z" }, - { url = "https://files.pythonhosted.org/packages/5c/e2/714694e4b87b85a18e2c243614974413c60aa107fd815b8cbc42b873d1d7/rpds_py-0.30.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:7cee9c752c0364588353e627da8a7e808a66873672bcb5f52890c33fd965b394", size = 563120, upload-time = "2025-11-30T20:22:35.903Z" }, - { url = "https://files.pythonhosted.org/packages/6f/ab/d5d5e3bcedb0a77f4f613706b750e50a5a3ba1c15ccd3665ecc636c968fd/rpds_py-0.30.0-cp312-cp312-win32.whl", hash = "sha256:1ab5b83dbcf55acc8b08fc62b796ef672c457b17dbd7820a11d6c52c06839bdf", size = 223782, upload-time = "2025-11-30T20:22:37.271Z" }, - { url = "https://files.pythonhosted.org/packages/39/3b/f786af9957306fdc38a74cef405b7b93180f481fb48453a114bb6465744a/rpds_py-0.30.0-cp312-cp312-win_amd64.whl", hash = "sha256:a090322ca841abd453d43456ac34db46e8b05fd9b3b4ac0c78bcde8b089f959b", size = 240463, upload-time = "2025-11-30T20:22:39.021Z" }, - { url = "https://files.pythonhosted.org/packages/f3/d2/b91dc748126c1559042cfe41990deb92c4ee3e2b415f6b5234969ffaf0cc/rpds_py-0.30.0-cp312-cp312-win_arm64.whl", hash = "sha256:669b1805bd639dd2989b281be2cfd951c6121b65e729d9b843e9639ef1fd555e", size = 230868, upload-time = "2025-11-30T20:22:40.493Z" }, - { url = "https://files.pythonhosted.org/packages/ed/dc/d61221eb88ff410de3c49143407f6f3147acf2538c86f2ab7ce65ae7d5f9/rpds_py-0.30.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:f83424d738204d9770830d35290ff3273fbb02b41f919870479fab14b9d303b2", size = 374887, upload-time = "2025-11-30T20:22:41.812Z" }, - { url = "https://files.pythonhosted.org/packages/fd/32/55fb50ae104061dbc564ef15cc43c013dc4a9f4527a1f4d99baddf56fe5f/rpds_py-0.30.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:e7536cd91353c5273434b4e003cbda89034d67e7710eab8761fd918ec6c69cf8", size = 358904, upload-time = "2025-11-30T20:22:43.479Z" }, - { url = "https://files.pythonhosted.org/packages/58/70/faed8186300e3b9bdd138d0273109784eea2396c68458ed580f885dfe7ad/rpds_py-0.30.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2771c6c15973347f50fece41fc447c054b7ac2ae0502388ce3b6738cd366e3d4", size = 389945, upload-time = "2025-11-30T20:22:44.819Z" }, - { url = "https://files.pythonhosted.org/packages/bd/a8/073cac3ed2c6387df38f71296d002ab43496a96b92c823e76f46b8af0543/rpds_py-0.30.0-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:0a59119fc6e3f460315fe9d08149f8102aa322299deaa5cab5b40092345c2136", size = 407783, upload-time = "2025-11-30T20:22:46.103Z" }, - { url = "https://files.pythonhosted.org/packages/77/57/5999eb8c58671f1c11eba084115e77a8899d6e694d2a18f69f0ba471ec8b/rpds_py-0.30.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:76fec018282b4ead0364022e3c54b60bf368b9d926877957a8624b58419169b7", size = 515021, upload-time = "2025-11-30T20:22:47.458Z" }, - { url = "https://files.pythonhosted.org/packages/e0/af/5ab4833eadc36c0a8ed2bc5c0de0493c04f6c06de223170bd0798ff98ced/rpds_py-0.30.0-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:692bef75a5525db97318e8cd061542b5a79812d711ea03dbc1f6f8dbb0c5f0d2", size = 414589, upload-time = "2025-11-30T20:22:48.872Z" }, - { url = "https://files.pythonhosted.org/packages/b7/de/f7192e12b21b9e9a68a6d0f249b4af3fdcdff8418be0767a627564afa1f1/rpds_py-0.30.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9027da1ce107104c50c81383cae773ef5c24d296dd11c99e2629dbd7967a20c6", size = 394025, upload-time = "2025-11-30T20:22:50.196Z" }, - { url = "https://files.pythonhosted.org/packages/91/c4/fc70cd0249496493500e7cc2de87504f5aa6509de1e88623431fec76d4b6/rpds_py-0.30.0-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:9cf69cdda1f5968a30a359aba2f7f9aa648a9ce4b580d6826437f2b291cfc86e", size = 408895, upload-time = "2025-11-30T20:22:51.87Z" }, - { url = "https://files.pythonhosted.org/packages/58/95/d9275b05ab96556fefff73a385813eb66032e4c99f411d0795372d9abcea/rpds_py-0.30.0-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:a4796a717bf12b9da9d3ad002519a86063dcac8988b030e405704ef7d74d2d9d", size = 422799, upload-time = "2025-11-30T20:22:53.341Z" }, - { url = "https://files.pythonhosted.org/packages/06/c1/3088fc04b6624eb12a57eb814f0d4997a44b0d208d6cace713033ff1a6ba/rpds_py-0.30.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:5d4c2aa7c50ad4728a094ebd5eb46c452e9cb7edbfdb18f9e1221f597a73e1e7", size = 572731, upload-time = "2025-11-30T20:22:54.778Z" }, - { url = "https://files.pythonhosted.org/packages/d8/42/c612a833183b39774e8ac8fecae81263a68b9583ee343db33ab571a7ce55/rpds_py-0.30.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:ba81a9203d07805435eb06f536d95a266c21e5b2dfbf6517748ca40c98d19e31", size = 599027, upload-time = "2025-11-30T20:22:56.212Z" }, - { url = "https://files.pythonhosted.org/packages/5f/60/525a50f45b01d70005403ae0e25f43c0384369ad24ffe46e8d9068b50086/rpds_py-0.30.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:945dccface01af02675628334f7cf49c2af4c1c904748efc5cf7bbdf0b579f95", size = 563020, upload-time = "2025-11-30T20:22:58.2Z" }, - { url = "https://files.pythonhosted.org/packages/0b/5d/47c4655e9bcd5ca907148535c10e7d489044243cc9941c16ed7cd53be91d/rpds_py-0.30.0-cp313-cp313-win32.whl", hash = "sha256:b40fb160a2db369a194cb27943582b38f79fc4887291417685f3ad693c5a1d5d", size = 223139, upload-time = "2025-11-30T20:23:00.209Z" }, - { url = "https://files.pythonhosted.org/packages/f2/e1/485132437d20aa4d3e1d8b3fb5a5e65aa8139f1e097080c2a8443201742c/rpds_py-0.30.0-cp313-cp313-win_amd64.whl", hash = "sha256:806f36b1b605e2d6a72716f321f20036b9489d29c51c91f4dd29a3e3afb73b15", size = 240224, upload-time = "2025-11-30T20:23:02.008Z" }, - { url = "https://files.pythonhosted.org/packages/24/95/ffd128ed1146a153d928617b0ef673960130be0009c77d8fbf0abe306713/rpds_py-0.30.0-cp313-cp313-win_arm64.whl", hash = "sha256:d96c2086587c7c30d44f31f42eae4eac89b60dabbac18c7669be3700f13c3ce1", size = 230645, upload-time = "2025-11-30T20:23:03.43Z" }, - { url = "https://files.pythonhosted.org/packages/ff/1b/b10de890a0def2a319a2626334a7f0ae388215eb60914dbac8a3bae54435/rpds_py-0.30.0-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:eb0b93f2e5c2189ee831ee43f156ed34e2a89a78a66b98cadad955972548be5a", size = 364443, upload-time = "2025-11-30T20:23:04.878Z" }, - { url = "https://files.pythonhosted.org/packages/0d/bf/27e39f5971dc4f305a4fb9c672ca06f290f7c4e261c568f3dea16a410d47/rpds_py-0.30.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:922e10f31f303c7c920da8981051ff6d8c1a56207dbdf330d9047f6d30b70e5e", size = 353375, upload-time = "2025-11-30T20:23:06.342Z" }, - { url = "https://files.pythonhosted.org/packages/40/58/442ada3bba6e8e6615fc00483135c14a7538d2ffac30e2d933ccf6852232/rpds_py-0.30.0-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cdc62c8286ba9bf7f47befdcea13ea0e26bf294bda99758fd90535cbaf408000", size = 383850, upload-time = "2025-11-30T20:23:07.825Z" }, - { url = "https://files.pythonhosted.org/packages/14/14/f59b0127409a33c6ef6f5c1ebd5ad8e32d7861c9c7adfa9a624fc3889f6c/rpds_py-0.30.0-cp313-cp313t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:47f9a91efc418b54fb8190a6b4aa7813a23fb79c51f4bb84e418f5476c38b8db", size = 392812, upload-time = "2025-11-30T20:23:09.228Z" }, - { url = "https://files.pythonhosted.org/packages/b3/66/e0be3e162ac299b3a22527e8913767d869e6cc75c46bd844aa43fb81ab62/rpds_py-0.30.0-cp313-cp313t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1f3587eb9b17f3789ad50824084fa6f81921bbf9a795826570bda82cb3ed91f2", size = 517841, upload-time = "2025-11-30T20:23:11.186Z" }, - { url = "https://files.pythonhosted.org/packages/3d/55/fa3b9cf31d0c963ecf1ba777f7cf4b2a2c976795ac430d24a1f43d25a6ba/rpds_py-0.30.0-cp313-cp313t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:39c02563fc592411c2c61d26b6c5fe1e51eaa44a75aa2c8735ca88b0d9599daa", size = 408149, upload-time = "2025-11-30T20:23:12.864Z" }, - { url = "https://files.pythonhosted.org/packages/60/ca/780cf3b1a32b18c0f05c441958d3758f02544f1d613abf9488cd78876378/rpds_py-0.30.0-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:51a1234d8febafdfd33a42d97da7a43f5dcb120c1060e352a3fbc0c6d36e2083", size = 383843, upload-time = "2025-11-30T20:23:14.638Z" }, - { url = "https://files.pythonhosted.org/packages/82/86/d5f2e04f2aa6247c613da0c1dd87fcd08fa17107e858193566048a1e2f0a/rpds_py-0.30.0-cp313-cp313t-manylinux_2_31_riscv64.whl", hash = "sha256:eb2c4071ab598733724c08221091e8d80e89064cd472819285a9ab0f24bcedb9", size = 396507, upload-time = "2025-11-30T20:23:16.105Z" }, - { url = "https://files.pythonhosted.org/packages/4b/9a/453255d2f769fe44e07ea9785c8347edaf867f7026872e76c1ad9f7bed92/rpds_py-0.30.0-cp313-cp313t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:6bdfdb946967d816e6adf9a3d8201bfad269c67efe6cefd7093ef959683c8de0", size = 414949, upload-time = "2025-11-30T20:23:17.539Z" }, - { url = "https://files.pythonhosted.org/packages/a3/31/622a86cdc0c45d6df0e9ccb6becdba5074735e7033c20e401a6d9d0e2ca0/rpds_py-0.30.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:c77afbd5f5250bf27bf516c7c4a016813eb2d3e116139aed0096940c5982da94", size = 565790, upload-time = "2025-11-30T20:23:19.029Z" }, - { url = "https://files.pythonhosted.org/packages/1c/5d/15bbf0fb4a3f58a3b1c67855ec1efcc4ceaef4e86644665fff03e1b66d8d/rpds_py-0.30.0-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:61046904275472a76c8c90c9ccee9013d70a6d0f73eecefd38c1ae7c39045a08", size = 590217, upload-time = "2025-11-30T20:23:20.885Z" }, - { url = "https://files.pythonhosted.org/packages/6d/61/21b8c41f68e60c8cc3b2e25644f0e3681926020f11d06ab0b78e3c6bbff1/rpds_py-0.30.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:4c5f36a861bc4b7da6516dbdf302c55313afa09b81931e8280361a4f6c9a2d27", size = 555806, upload-time = "2025-11-30T20:23:22.488Z" }, - { url = "https://files.pythonhosted.org/packages/f9/39/7e067bb06c31de48de3eb200f9fc7c58982a4d3db44b07e73963e10d3be9/rpds_py-0.30.0-cp313-cp313t-win32.whl", hash = "sha256:3d4a69de7a3e50ffc214ae16d79d8fbb0922972da0356dcf4d0fdca2878559c6", size = 211341, upload-time = "2025-11-30T20:23:24.449Z" }, - { url = "https://files.pythonhosted.org/packages/0a/4d/222ef0b46443cf4cf46764d9c630f3fe4abaa7245be9417e56e9f52b8f65/rpds_py-0.30.0-cp313-cp313t-win_amd64.whl", hash = "sha256:f14fc5df50a716f7ece6a80b6c78bb35ea2ca47c499e422aa4463455dd96d56d", size = 225768, upload-time = "2025-11-30T20:23:25.908Z" }, - { url = "https://files.pythonhosted.org/packages/86/81/dad16382ebbd3d0e0328776d8fd7ca94220e4fa0798d1dc5e7da48cb3201/rpds_py-0.30.0-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:68f19c879420aa08f61203801423f6cd5ac5f0ac4ac82a2368a9fcd6a9a075e0", size = 362099, upload-time = "2025-11-30T20:23:27.316Z" }, - { url = "https://files.pythonhosted.org/packages/2b/60/19f7884db5d5603edf3c6bce35408f45ad3e97e10007df0e17dd57af18f8/rpds_py-0.30.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:ec7c4490c672c1a0389d319b3a9cfcd098dcdc4783991553c332a15acf7249be", size = 353192, upload-time = "2025-11-30T20:23:29.151Z" }, - { url = "https://files.pythonhosted.org/packages/bf/c4/76eb0e1e72d1a9c4703c69607cec123c29028bff28ce41588792417098ac/rpds_py-0.30.0-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f251c812357a3fed308d684a5079ddfb9d933860fc6de89f2b7ab00da481e65f", size = 384080, upload-time = "2025-11-30T20:23:30.785Z" }, - { url = "https://files.pythonhosted.org/packages/72/87/87ea665e92f3298d1b26d78814721dc39ed8d2c74b86e83348d6b48a6f31/rpds_py-0.30.0-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:ac98b175585ecf4c0348fd7b29c3864bda53b805c773cbf7bfdaffc8070c976f", size = 394841, upload-time = "2025-11-30T20:23:32.209Z" }, - { url = "https://files.pythonhosted.org/packages/77/ad/7783a89ca0587c15dcbf139b4a8364a872a25f861bdb88ed99f9b0dec985/rpds_py-0.30.0-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3e62880792319dbeb7eb866547f2e35973289e7d5696c6e295476448f5b63c87", size = 516670, upload-time = "2025-11-30T20:23:33.742Z" }, - { url = "https://files.pythonhosted.org/packages/5b/3c/2882bdac942bd2172f3da574eab16f309ae10a3925644e969536553cb4ee/rpds_py-0.30.0-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:4e7fc54e0900ab35d041b0601431b0a0eb495f0851a0639b6ef90f7741b39a18", size = 408005, upload-time = "2025-11-30T20:23:35.253Z" }, - { url = "https://files.pythonhosted.org/packages/ce/81/9a91c0111ce1758c92516a3e44776920b579d9a7c09b2b06b642d4de3f0f/rpds_py-0.30.0-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:47e77dc9822d3ad616c3d5759ea5631a75e5809d5a28707744ef79d7a1bcfcad", size = 382112, upload-time = "2025-11-30T20:23:36.842Z" }, - { url = "https://files.pythonhosted.org/packages/cf/8e/1da49d4a107027e5fbc64daeab96a0706361a2918da10cb41769244b805d/rpds_py-0.30.0-cp314-cp314-manylinux_2_31_riscv64.whl", hash = "sha256:b4dc1a6ff022ff85ecafef7979a2c6eb423430e05f1165d6688234e62ba99a07", size = 399049, upload-time = "2025-11-30T20:23:38.343Z" }, - { url = "https://files.pythonhosted.org/packages/df/5a/7ee239b1aa48a127570ec03becbb29c9d5a9eb092febbd1699d567cae859/rpds_py-0.30.0-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:4559c972db3a360808309e06a74628b95eaccbf961c335c8fe0d590cf587456f", size = 415661, upload-time = "2025-11-30T20:23:40.263Z" }, - { url = "https://files.pythonhosted.org/packages/70/ea/caa143cf6b772f823bc7929a45da1fa83569ee49b11d18d0ada7f5ee6fd6/rpds_py-0.30.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:0ed177ed9bded28f8deb6ab40c183cd1192aa0de40c12f38be4d59cd33cb5c65", size = 565606, upload-time = "2025-11-30T20:23:42.186Z" }, - { url = "https://files.pythonhosted.org/packages/64/91/ac20ba2d69303f961ad8cf55bf7dbdb4763f627291ba3d0d7d67333cced9/rpds_py-0.30.0-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:ad1fa8db769b76ea911cb4e10f049d80bf518c104f15b3edb2371cc65375c46f", size = 591126, upload-time = "2025-11-30T20:23:44.086Z" }, - { url = "https://files.pythonhosted.org/packages/21/20/7ff5f3c8b00c8a95f75985128c26ba44503fb35b8e0259d812766ea966c7/rpds_py-0.30.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:46e83c697b1f1c72b50e5ee5adb4353eef7406fb3f2043d64c33f20ad1c2fc53", size = 553371, upload-time = "2025-11-30T20:23:46.004Z" }, - { url = "https://files.pythonhosted.org/packages/72/c7/81dadd7b27c8ee391c132a6b192111ca58d866577ce2d9b0ca157552cce0/rpds_py-0.30.0-cp314-cp314-win32.whl", hash = "sha256:ee454b2a007d57363c2dfd5b6ca4a5d7e2c518938f8ed3b706e37e5d470801ed", size = 215298, upload-time = "2025-11-30T20:23:47.696Z" }, - { url = "https://files.pythonhosted.org/packages/3e/d2/1aaac33287e8cfb07aab2e6b8ac1deca62f6f65411344f1433c55e6f3eb8/rpds_py-0.30.0-cp314-cp314-win_amd64.whl", hash = "sha256:95f0802447ac2d10bcc69f6dc28fe95fdf17940367b21d34e34c737870758950", size = 228604, upload-time = "2025-11-30T20:23:49.501Z" }, - { url = "https://files.pythonhosted.org/packages/e8/95/ab005315818cc519ad074cb7784dae60d939163108bd2b394e60dc7b5461/rpds_py-0.30.0-cp314-cp314-win_arm64.whl", hash = "sha256:613aa4771c99f03346e54c3f038e4cc574ac09a3ddfb0e8878487335e96dead6", size = 222391, upload-time = "2025-11-30T20:23:50.96Z" }, - { url = "https://files.pythonhosted.org/packages/9e/68/154fe0194d83b973cdedcdcc88947a2752411165930182ae41d983dcefa6/rpds_py-0.30.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:7e6ecfcb62edfd632e56983964e6884851786443739dbfe3582947e87274f7cb", size = 364868, upload-time = "2025-11-30T20:23:52.494Z" }, - { url = "https://files.pythonhosted.org/packages/83/69/8bbc8b07ec854d92a8b75668c24d2abcb1719ebf890f5604c61c9369a16f/rpds_py-0.30.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:a1d0bc22a7cdc173fedebb73ef81e07faef93692b8c1ad3733b67e31e1b6e1b8", size = 353747, upload-time = "2025-11-30T20:23:54.036Z" }, - { url = "https://files.pythonhosted.org/packages/ab/00/ba2e50183dbd9abcce9497fa5149c62b4ff3e22d338a30d690f9af970561/rpds_py-0.30.0-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0d08f00679177226c4cb8c5265012eea897c8ca3b93f429e546600c971bcbae7", size = 383795, upload-time = "2025-11-30T20:23:55.556Z" }, - { url = "https://files.pythonhosted.org/packages/05/6f/86f0272b84926bcb0e4c972262f54223e8ecc556b3224d281e6598fc9268/rpds_py-0.30.0-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:5965af57d5848192c13534f90f9dd16464f3c37aaf166cc1da1cae1fd5a34898", size = 393330, upload-time = "2025-11-30T20:23:57.033Z" }, - { url = "https://files.pythonhosted.org/packages/cb/e9/0e02bb2e6dc63d212641da45df2b0bf29699d01715913e0d0f017ee29438/rpds_py-0.30.0-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:9a4e86e34e9ab6b667c27f3211ca48f73dba7cd3d90f8d5b11be56e5dbc3fb4e", size = 518194, upload-time = "2025-11-30T20:23:58.637Z" }, - { url = "https://files.pythonhosted.org/packages/ee/ca/be7bca14cf21513bdf9c0606aba17d1f389ea2b6987035eb4f62bd923f25/rpds_py-0.30.0-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e5d3e6b26f2c785d65cc25ef1e5267ccbe1b069c5c21b8cc724efee290554419", size = 408340, upload-time = "2025-11-30T20:24:00.2Z" }, - { url = "https://files.pythonhosted.org/packages/c2/c7/736e00ebf39ed81d75544c0da6ef7b0998f8201b369acf842f9a90dc8fce/rpds_py-0.30.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:626a7433c34566535b6e56a1b39a7b17ba961e97ce3b80ec62e6f1312c025551", size = 383765, upload-time = "2025-11-30T20:24:01.759Z" }, - { url = "https://files.pythonhosted.org/packages/4a/3f/da50dfde9956aaf365c4adc9533b100008ed31aea635f2b8d7b627e25b49/rpds_py-0.30.0-cp314-cp314t-manylinux_2_31_riscv64.whl", hash = "sha256:acd7eb3f4471577b9b5a41baf02a978e8bdeb08b4b355273994f8b87032000a8", size = 396834, upload-time = "2025-11-30T20:24:03.687Z" }, - { url = "https://files.pythonhosted.org/packages/4e/00/34bcc2565b6020eab2623349efbdec810676ad571995911f1abdae62a3a0/rpds_py-0.30.0-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:fe5fa731a1fa8a0a56b0977413f8cacac1768dad38d16b3a296712709476fbd5", size = 415470, upload-time = "2025-11-30T20:24:05.232Z" }, - { url = "https://files.pythonhosted.org/packages/8c/28/882e72b5b3e6f718d5453bd4d0d9cf8df36fddeb4ddbbab17869d5868616/rpds_py-0.30.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:74a3243a411126362712ee1524dfc90c650a503502f135d54d1b352bd01f2404", size = 565630, upload-time = "2025-11-30T20:24:06.878Z" }, - { url = "https://files.pythonhosted.org/packages/3b/97/04a65539c17692de5b85c6e293520fd01317fd878ea1995f0367d4532fb1/rpds_py-0.30.0-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:3e8eeb0544f2eb0d2581774be4c3410356eba189529a6b3e36bbbf9696175856", size = 591148, upload-time = "2025-11-30T20:24:08.445Z" }, - { url = "https://files.pythonhosted.org/packages/85/70/92482ccffb96f5441aab93e26c4d66489eb599efdcf96fad90c14bbfb976/rpds_py-0.30.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:dbd936cde57abfee19ab3213cf9c26be06d60750e60a8e4dd85d1ab12c8b1f40", size = 556030, upload-time = "2025-11-30T20:24:10.956Z" }, - { url = "https://files.pythonhosted.org/packages/20/53/7c7e784abfa500a2b6b583b147ee4bb5a2b3747a9166bab52fec4b5b5e7d/rpds_py-0.30.0-cp314-cp314t-win32.whl", hash = "sha256:dc824125c72246d924f7f796b4f63c1e9dc810c7d9e2355864b3c3a73d59ade0", size = 211570, upload-time = "2025-11-30T20:24:12.735Z" }, - { url = "https://files.pythonhosted.org/packages/d0/02/fa464cdfbe6b26e0600b62c528b72d8608f5cc49f96b8d6e38c95d60c676/rpds_py-0.30.0-cp314-cp314t-win_amd64.whl", hash = "sha256:27f4b0e92de5bfbc6f86e43959e6edd1425c33b5e69aab0984a72047f2bcf1e3", size = 226532, upload-time = "2025-11-30T20:24:14.634Z" }, -] - -[[package]] -name = "rsa" -version = "4.9.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pyasn1" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/da/8a/22b7beea3ee0d44b1916c0c1cb0ee3af23b700b6da9f04991899d0c555d4/rsa-4.9.1.tar.gz", hash = "sha256:e7bdbfdb5497da4c07dfd35530e1a902659db6ff241e39d9953cad06ebd0ae75", size = 29034, upload-time = "2025-04-16T09:51:18.218Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/64/8d/0133e4eb4beed9e425d9a98ed6e081a55d195481b7632472be1af08d2f6b/rsa-4.9.1-py3-none-any.whl", hash = "sha256:68635866661c6836b8d39430f97a996acbd61bfa49406748ea243539fe239762", size = 34696, upload-time = "2025-04-16T09:51:17.142Z" }, -] - [[package]] name = "ruff" version = "0.14.8" @@ -3538,33 +2243,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/64/47/a494741db7280eae6dc033510c319e34d42dd41b7ac0c7ead39354d1a2b5/scipy-1.16.3-cp314-cp314t-win_arm64.whl", hash = "sha256:21d9d6b197227a12dcbf9633320a4e34c6b0e51c57268df255a0942983bac562", size = 26464127, upload-time = "2025-10-28T17:38:11.34Z" }, ] -[[package]] -name = "seaborn" -version = "0.13.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "matplotlib" }, - { name = "numpy" }, - { name = "pandas" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/86/59/a451d7420a77ab0b98f7affa3a1d78a313d2f7281a57afb1a34bae8ab412/seaborn-0.13.2.tar.gz", hash = "sha256:93e60a40988f4d65e9f4885df477e2fdaff6b73a9ded434c1ab356dd57eefff7", size = 1457696, upload-time = "2024-01-25T13:21:52.551Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/83/11/00d3c3dfc25ad54e731d91449895a79e4bf2384dc3ac01809010ba88f6d5/seaborn-0.13.2-py3-none-any.whl", hash = "sha256:636f8336facf092165e27924f223d3c62ca560b1f2bb5dff7ab7fad265361987", size = 294914, upload-time = "2024-01-25T13:21:49.598Z" }, -] - -[[package]] -name = "secretstorage" -version = "3.5.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cryptography", marker = "sys_platform == 'linux'" }, - { name = "jeepney", marker = "sys_platform == 'linux'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/1c/03/e834bcd866f2f8a49a85eaff47340affa3bfa391ee9912a952a1faa68c7b/secretstorage-3.5.0.tar.gz", hash = "sha256:f04b8e4689cbce351744d5537bf6b1329c6fc68f91fa666f60a380edddcd11be", size = 19884, upload-time = "2025-11-23T19:02:53.191Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b7/46/f5af3402b579fd5e11573ce652019a67074317e18c1935cc0b4ba9b35552/secretstorage-3.5.0-py3-none-any.whl", hash = "sha256:0ce65888c0725fcb2c5bc0fdb8e5438eece02c523557ea40ce0703c266248137", size = 15554, upload-time = "2025-11-23T19:02:51.545Z" }, -] - [[package]] name = "sentry-sdk" version = "2.47.0" @@ -3587,15 +2265,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a3/dc/17031897dae0efacfea57dfd3a82fdd2a2aeb58e0ff71b77b87e44edc772/setuptools-80.9.0-py3-none-any.whl", hash = "sha256:062d34222ad13e0cc312a4c02d73f059e86a4acbfbdea8f8f76b28c99f306922", size = 1201486, upload-time = "2025-05-27T00:56:49.664Z" }, ] -[[package]] -name = "shellingham" -version = "1.5.4" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/58/15/8b3609fd3830ef7b27b655beb4b4e9c62313a4e8da8c676e142cc210d58e/shellingham-1.5.4.tar.gz", hash = "sha256:8dbca0739d487e5bd35ab3ca4b36e11c4078f3a234bfce294b0a0291363404de", size = 10310, upload-time = "2023-10-24T04:13:40.426Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" }, -] - [[package]] name = "six" version = "1.17.0" @@ -3610,22 +2279,24 @@ name = "slices" version = "0.1.0" source = { editable = "." } dependencies = [ - { name = "duckdb" }, { name = "hydra-core" }, { name = "lightning" }, - { name = "m4-infra" }, - { name = "matplotlib" }, { name = "numpy" }, { name = "omegaconf" }, + { name = "pandas" }, { name = "polars" }, { name = "portalocker" }, - { name = "pre-commit" }, { name = "pyarrow" }, + { name = "pydantic" }, + { name = "pyyaml" }, { name = "rich" }, { name = "scikit-learn" }, - { name = "seaborn" }, + { name = "scipy" }, { name = "torch" }, + { name = "torchmetrics" }, + { name = "tqdm" }, { name = "wandb" }, + { name = "xgboost" }, ] [package.optional-dependencies] @@ -3637,46 +2308,62 @@ dev = [ { name = "pytest-cov" }, { name = "ruff" }, ] +viz = [ + { name = "matplotlib" }, + { name = "umap-learn" }, +] [package.dev-dependencies] dev = [ + { name = "black" }, + { name = "mypy" }, { name = "pytest" }, + { name = "pytest-cov" }, { name = "python-pptx" }, { name = "ruff" }, + { name = "types-pyyaml" }, ] [package.metadata] requires-dist = [ { name = "black", marker = "extra == 'dev'", specifier = ">=23.0" }, - { name = "duckdb", specifier = ">=0.9" }, { name = "hydra-core", specifier = ">=1.3" }, { name = "lightning", specifier = ">=2.0" }, - { name = "m4-infra", specifier = ">=0.0.0.dev0" }, - { name = "matplotlib", specifier = ">=3.10.7" }, + { name = "matplotlib", marker = "extra == 'viz'", specifier = ">=3.10.7" }, { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.0" }, { name = "numpy", specifier = ">=1.24" }, { name = "omegaconf", specifier = ">=2.3" }, + { name = "pandas", specifier = ">=2.0" }, { name = "polars", specifier = ">=0.19" }, { name = "portalocker", specifier = ">=2.8.0" }, - { name = "pre-commit", specifier = ">=4.5.0" }, { name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.0" }, { name = "pyarrow", specifier = ">=14.0" }, + { name = "pydantic", specifier = ">=2.0" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0" }, { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=4.0" }, + { name = "pyyaml", specifier = ">=6.0" }, { name = "rich", specifier = ">=13.0" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.1" }, { name = "scikit-learn", specifier = ">=1.3" }, - { name = "seaborn", specifier = ">=0.13.2" }, + { name = "scipy", specifier = ">=1.10" }, { name = "torch", specifier = ">=2.0" }, + { name = "torchmetrics", specifier = ">=1.0" }, + { name = "tqdm", specifier = ">=4.0" }, + { name = "umap-learn", marker = "extra == 'viz'", specifier = ">=0.5" }, { name = "wandb", specifier = ">=0.15" }, + { name = "xgboost", specifier = ">=2.0" }, ] -provides-extras = ["dev"] +provides-extras = ["viz", "dev"] [package.metadata.requires-dev] dev = [ + { name = "black", specifier = ">=23.0" }, + { name = "mypy", specifier = ">=1.0" }, { name = "pytest", specifier = ">=9.0.1" }, + { name = "pytest-cov", specifier = ">=4.0" }, { name = "python-pptx", specifier = ">=1.0.2" }, { name = "ruff", specifier = ">=0.14.8" }, + { name = "types-pyyaml", specifier = ">=6.0" }, ] [[package]] @@ -3688,94 +2375,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/04/be/d09147ad1ec7934636ad912901c5fd7667e1c858e19d355237db0d0cd5e4/smmap-5.0.2-py3-none-any.whl", hash = "sha256:b30115f0def7d7531d22a0fb6502488d879e75b260a9db4d0819cfb25403af5e", size = 24303, upload-time = "2025-01-02T07:14:38.724Z" }, ] -[[package]] -name = "sortedcontainers" -version = "2.4.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e8/c4/ba2f8066cceb6f23394729afe52f3bf7adec04bf9ed2c820b39e19299111/sortedcontainers-2.4.0.tar.gz", hash = "sha256:25caa5a06cc30b6b83d11423433f65d1f9d76c4c6a0c90e3379eaa43b9bfdb88", size = 30594, upload-time = "2021-05-16T22:03:42.897Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/32/46/9cb0e58b2deb7f82b84065f37f3bffeb12413f947f9388e4cac22c4621ce/sortedcontainers-2.4.0-py2.py3-none-any.whl", hash = "sha256:a163dcaede0f1c021485e957a39245190e74249897e2ae4b2aa38595db237ee0", size = 29575, upload-time = "2021-05-16T22:03:41.177Z" }, -] - -[[package]] -name = "soupsieve" -version = "2.8.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/89/23/adf3796d740536d63a6fbda113d07e60c734b6ed5d3058d1e47fc0495e47/soupsieve-2.8.1.tar.gz", hash = "sha256:4cf733bc50fa805f5df4b8ef4740fc0e0fa6218cf3006269afd3f9d6d80fd350", size = 117856, upload-time = "2025-12-18T13:50:34.655Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/48/f3/b67d6ea49ca9154453b6d70b34ea22f3996b9fa55da105a79d8732227adc/soupsieve-2.8.1-py3-none-any.whl", hash = "sha256:a11fe2a6f3d76ab3cf2de04eb339c1be5b506a8a47f2ceb6d139803177f85434", size = 36710, upload-time = "2025-12-18T13:50:33.267Z" }, -] - -[[package]] -name = "sqlalchemy" -version = "2.0.45" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "greenlet", marker = "platform_machine == 'AMD64' or platform_machine == 'WIN32' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'ppc64le' or platform_machine == 'win32' or platform_machine == 'x86_64'" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/be/f9/5e4491e5ccf42f5d9cfc663741d261b3e6e1683ae7812114e7636409fcc6/sqlalchemy-2.0.45.tar.gz", hash = "sha256:1632a4bda8d2d25703fdad6363058d882541bdaaee0e5e3ddfa0cd3229efce88", size = 9869912, upload-time = "2025-12-09T21:05:16.737Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/2d/c7/1900b56ce19bff1c26f39a4ce427faec7716c81ac792bfac8b6a9f3dca93/sqlalchemy-2.0.45-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b3ee2aac15169fb0d45822983631466d60b762085bc4535cd39e66bea362df5f", size = 3333760, upload-time = "2025-12-09T22:11:02.66Z" }, - { url = "https://files.pythonhosted.org/packages/0a/93/3be94d96bb442d0d9a60e55a6bb6e0958dd3457751c6f8502e56ef95fed0/sqlalchemy-2.0.45-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ba547ac0b361ab4f1608afbc8432db669bd0819b3e12e29fb5fa9529a8bba81d", size = 3348268, upload-time = "2025-12-09T22:13:49.054Z" }, - { url = "https://files.pythonhosted.org/packages/48/4b/f88ded696e61513595e4a9778f9d3f2bf7332cce4eb0c7cedaabddd6687b/sqlalchemy-2.0.45-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:215f0528b914e5c75ef2559f69dca86878a3beeb0c1be7279d77f18e8d180ed4", size = 3278144, upload-time = "2025-12-09T22:11:04.14Z" }, - { url = "https://files.pythonhosted.org/packages/ed/6a/310ecb5657221f3e1bd5288ed83aa554923fb5da48d760a9f7622afeb065/sqlalchemy-2.0.45-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:107029bf4f43d076d4011f1afb74f7c3e2ea029ec82eb23d8527d5e909e97aa6", size = 3313907, upload-time = "2025-12-09T22:13:50.598Z" }, - { url = "https://files.pythonhosted.org/packages/5c/39/69c0b4051079addd57c84a5bfb34920d87456dd4c90cf7ee0df6efafc8ff/sqlalchemy-2.0.45-cp312-cp312-win32.whl", hash = "sha256:0c9f6ada57b58420a2c0277ff853abe40b9e9449f8d7d231763c6bc30f5c4953", size = 2112182, upload-time = "2025-12-09T21:39:30.824Z" }, - { url = "https://files.pythonhosted.org/packages/f7/4e/510db49dd89fc3a6e994bee51848c94c48c4a00dc905e8d0133c251f41a7/sqlalchemy-2.0.45-cp312-cp312-win_amd64.whl", hash = "sha256:8defe5737c6d2179c7997242d6473587c3beb52e557f5ef0187277009f73e5e1", size = 2139200, upload-time = "2025-12-09T21:39:32.321Z" }, - { url = "https://files.pythonhosted.org/packages/6a/c8/7cc5221b47a54edc72a0140a1efa56e0a2730eefa4058d7ed0b4c4357ff8/sqlalchemy-2.0.45-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fe187fc31a54d7fd90352f34e8c008cf3ad5d064d08fedd3de2e8df83eb4a1cf", size = 3277082, upload-time = "2025-12-09T22:11:06.167Z" }, - { url = "https://files.pythonhosted.org/packages/0e/50/80a8d080ac7d3d321e5e5d420c9a522b0aa770ec7013ea91f9a8b7d36e4a/sqlalchemy-2.0.45-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:672c45cae53ba88e0dad74b9027dddd09ef6f441e927786b05bec75d949fbb2e", size = 3293131, upload-time = "2025-12-09T22:13:52.626Z" }, - { url = "https://files.pythonhosted.org/packages/da/4c/13dab31266fc9904f7609a5dc308a2432a066141d65b857760c3bef97e69/sqlalchemy-2.0.45-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:470daea2c1ce73910f08caf10575676a37159a6d16c4da33d0033546bddebc9b", size = 3225389, upload-time = "2025-12-09T22:11:08.093Z" }, - { url = "https://files.pythonhosted.org/packages/74/04/891b5c2e9f83589de202e7abaf24cd4e4fa59e1837d64d528829ad6cc107/sqlalchemy-2.0.45-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:9c6378449e0940476577047150fd09e242529b761dc887c9808a9a937fe990c8", size = 3266054, upload-time = "2025-12-09T22:13:54.262Z" }, - { url = "https://files.pythonhosted.org/packages/f1/24/fc59e7f71b0948cdd4cff7a286210e86b0443ef1d18a23b0d83b87e4b1f7/sqlalchemy-2.0.45-cp313-cp313-win32.whl", hash = "sha256:4b6bec67ca45bc166c8729910bd2a87f1c0407ee955df110d78948f5b5827e8a", size = 2110299, upload-time = "2025-12-09T21:39:33.486Z" }, - { url = "https://files.pythonhosted.org/packages/c0/c5/d17113020b2d43073412aeca09b60d2009442420372123b8d49cc253f8b8/sqlalchemy-2.0.45-cp313-cp313-win_amd64.whl", hash = "sha256:afbf47dc4de31fa38fd491f3705cac5307d21d4bb828a4f020ee59af412744ee", size = 2136264, upload-time = "2025-12-09T21:39:36.801Z" }, - { url = "https://files.pythonhosted.org/packages/3d/8d/bb40a5d10e7a5f2195f235c0b2f2c79b0bf6e8f00c0c223130a4fbd2db09/sqlalchemy-2.0.45-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:83d7009f40ce619d483d26ac1b757dfe3167b39921379a8bd1b596cf02dab4a6", size = 3521998, upload-time = "2025-12-09T22:13:28.622Z" }, - { url = "https://files.pythonhosted.org/packages/75/a5/346128b0464886f036c039ea287b7332a410aa2d3fb0bb5d404cb8861635/sqlalchemy-2.0.45-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:d8a2ca754e5415cde2b656c27900b19d50ba076aa05ce66e2207623d3fe41f5a", size = 3473434, upload-time = "2025-12-09T22:13:30.188Z" }, - { url = "https://files.pythonhosted.org/packages/cc/64/4e1913772646b060b025d3fc52ce91a58967fe58957df32b455de5a12b4f/sqlalchemy-2.0.45-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7f46ec744e7f51275582e6a24326e10c49fbdd3fc99103e01376841213028774", size = 3272404, upload-time = "2025-12-09T22:11:09.662Z" }, - { url = "https://files.pythonhosted.org/packages/b3/27/caf606ee924282fe4747ee4fd454b335a72a6e018f97eab5ff7f28199e16/sqlalchemy-2.0.45-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:883c600c345123c033c2f6caca18def08f1f7f4c3ebeb591a63b6fceffc95cce", size = 3277057, upload-time = "2025-12-09T22:13:56.213Z" }, - { url = "https://files.pythonhosted.org/packages/85/d0/3d64218c9724e91f3d1574d12eb7ff8f19f937643815d8daf792046d88ab/sqlalchemy-2.0.45-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:2c0b74aa79e2deade948fe8593654c8ef4228c44ba862bb7c9585c8e0db90f33", size = 3222279, upload-time = "2025-12-09T22:11:11.1Z" }, - { url = "https://files.pythonhosted.org/packages/24/10/dd7688a81c5bc7690c2a3764d55a238c524cd1a5a19487928844cb247695/sqlalchemy-2.0.45-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:8a420169cef179d4c9064365f42d779f1e5895ad26ca0c8b4c0233920973db74", size = 3244508, upload-time = "2025-12-09T22:13:57.932Z" }, - { url = "https://files.pythonhosted.org/packages/aa/41/db75756ca49f777e029968d9c9fee338c7907c563267740c6d310a8e3f60/sqlalchemy-2.0.45-cp314-cp314-win32.whl", hash = "sha256:e50dcb81a5dfe4b7b4a4aa8f338116d127cb209559124f3694c70d6cd072b68f", size = 2113204, upload-time = "2025-12-09T21:39:38.365Z" }, - { url = "https://files.pythonhosted.org/packages/89/a2/0e1590e9adb292b1d576dbcf67ff7df8cf55e56e78d2c927686d01080f4b/sqlalchemy-2.0.45-cp314-cp314-win_amd64.whl", hash = "sha256:4748601c8ea959e37e03d13dcda4a44837afcd1b21338e637f7c935b8da06177", size = 2138785, upload-time = "2025-12-09T21:39:39.503Z" }, - { url = "https://files.pythonhosted.org/packages/42/39/f05f0ed54d451156bbed0e23eb0516bcad7cbb9f18b3bf219c786371b3f0/sqlalchemy-2.0.45-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:cd337d3526ec5298f67d6a30bbbe4ed7e5e68862f0bf6dd21d289f8d37b7d60b", size = 3522029, upload-time = "2025-12-09T22:13:32.09Z" }, - { url = "https://files.pythonhosted.org/packages/54/0f/d15398b98b65c2bce288d5ee3f7d0a81f77ab89d9456994d5c7cc8b2a9db/sqlalchemy-2.0.45-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:9a62b446b7d86a3909abbcd1cd3cc550a832f99c2bc37c5b22e1925438b9367b", size = 3475142, upload-time = "2025-12-09T22:13:33.739Z" }, - { url = "https://files.pythonhosted.org/packages/bf/e1/3ccb13c643399d22289c6a9786c1a91e3dcbb68bce4beb44926ac2c557bf/sqlalchemy-2.0.45-py3-none-any.whl", hash = "sha256:5225a288e4c8cc2308dbdd874edad6e7d0fd38eac1e9e5f23503425c8eee20d0", size = 1936672, upload-time = "2025-12-09T21:54:52.608Z" }, -] - -[[package]] -name = "sqlparse" -version = "0.5.5" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/90/76/437d71068094df0726366574cf3432a4ed754217b436eb7429415cf2d480/sqlparse-0.5.5.tar.gz", hash = "sha256:e20d4a9b0b8585fdf63b10d30066c7c94c5d7a7ec47c889a2d83a3caa93ff28e", size = 120815, upload-time = "2025-12-19T07:17:45.073Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/49/4b/359f28a903c13438ef59ebeee215fb25da53066db67b305c125f1c6d2a25/sqlparse-0.5.5-py3-none-any.whl", hash = "sha256:12a08b3bf3eec877c519589833aed092e2444e68240a3577e8e26148acc7b1ba", size = 46138, upload-time = "2025-12-19T07:17:46.573Z" }, -] - -[[package]] -name = "sse-starlette" -version = "3.1.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "anyio" }, - { name = "starlette" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/da/34/f5df66cb383efdbf4f2db23cabb27f51b1dcb737efaf8a558f6f1d195134/sse_starlette-3.1.2.tar.gz", hash = "sha256:55eff034207a83a0eb86de9a68099bd0157838f0b8b999a1b742005c71e33618", size = 26303, upload-time = "2025-12-31T08:02:20.023Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b7/95/8c4b76eec9ae574474e5d2997557cebf764bcd3586458956c30631ae08f4/sse_starlette-3.1.2-py3-none-any.whl", hash = "sha256:cd800dd349f4521b317b9391d3796fa97b71748a4da9b9e00aafab32dda375c8", size = 12484, upload-time = "2025-12-31T08:02:18.894Z" }, -] - -[[package]] -name = "starlette" -version = "0.51.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "anyio" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/e7/65/5a1fadcc40c5fdc7df421a7506b79633af8f5d5e3a95c3e72acacec644b9/starlette-0.51.0.tar.gz", hash = "sha256:4c4fda9b1bc67f84037d3d14a5112e523509c369d9d47b111b2f984b0cc5ba6c", size = 2647658, upload-time = "2026-01-10T20:23:15.043Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/18/c4/09985a03dba389d4fe16a9014147a7b02fa76ef3519bf5846462a485876d/starlette-0.51.0-py3-none-any.whl", hash = "sha256:fb460a3d6fd3c958d729fdd96aee297f89a51b0181f16401fe8fd4cb6129165d", size = 74133, upload-time = "2026-01-10T20:23:13.445Z" }, -] - [[package]] name = "sympy" version = "1.14.0" @@ -3889,18 +2488,12 @@ wheels = [ ] [[package]] -name = "typer" -version = "0.21.1" +name = "types-pyyaml" +version = "6.0.12.20250915" source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "click" }, - { name = "rich" }, - { name = "shellingham" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/36/bf/8825b5929afd84d0dabd606c67cd57b8388cb3ec385f7ef19c5cc2202069/typer-0.21.1.tar.gz", hash = "sha256:ea835607cd752343b6b2b7ce676893e5a0324082268b48f27aa058bdb7d2145d", size = 110371, upload-time = "2026-01-06T11:21:10.989Z" } +sdist = { url = "https://files.pythonhosted.org/packages/7e/69/3c51b36d04da19b92f9e815be12753125bd8bc247ba0470a982e6979e71c/types_pyyaml-6.0.12.20250915.tar.gz", hash = "sha256:0f8b54a528c303f0e6f7165687dd33fafa81c807fcac23f632b63aa624ced1d3", size = 17522, upload-time = "2025-09-15T03:01:00.728Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a0/1d/d9257dd49ff2ca23ea5f132edf1281a0c4f9de8a762b9ae399b670a59235/typer-0.21.1-py3-none-any.whl", hash = "sha256:7985e89081c636b88d172c2ee0cfe33c253160994d47bdfdc302defd7d1f1d01", size = 47381, upload-time = "2026-01-06T11:21:09.824Z" }, + { url = "https://files.pythonhosted.org/packages/bd/e0/1eed384f02555dde685fff1a1ac805c1c7dcb6dd019c916fe659b1c1f9ec/types_pyyaml-6.0.12.20250915-py3-none-any.whl", hash = "sha256:e7d4d9e064e89a3b3cae120b4990cd370874d2bf12fa5f46c97018dd5d3c9ab6", size = 20338, upload-time = "2025-09-15T03:00:59.218Z" }, ] [[package]] @@ -3934,25 +2527,29 @@ wheels = [ ] [[package]] -name = "urllib3" -version = "2.5.0" +name = "umap-learn" +version = "0.5.12" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/15/22/9ee70a2574a4f4599c47dd506532914ce044817c7752a79b6a51286319bc/urllib3-2.5.0.tar.gz", hash = "sha256:3fc47733c7e419d4bc3f6b3dc2b4f890bb743906a30d56ba4a5bfa4bbff92760", size = 393185, upload-time = "2025-06-18T14:07:41.644Z" } +dependencies = [ + { name = "numba" }, + { name = "numpy" }, + { name = "pynndescent" }, + { name = "scikit-learn" }, + { name = "scipy" }, + { name = "tqdm" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/02/ee/af4171241117f85c74b5ca6448ea1033cc28d599c13651d67289bacd4083/umap_learn-0.5.12.tar.gz", hash = "sha256:6aff02ecac5f2aad9f3c65ee518d7ae93e1a985ae38721fdcffceee4232c33c7", size = 96672, upload-time = "2026-04-08T20:03:54.012Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a7/c2/fe1e52489ae3122415c51f387e221dd0773709bad6c6cdaa599e8a2c5185/urllib3-2.5.0-py3-none-any.whl", hash = "sha256:e6b01673c0fa6a13e374b50871808eb3bf7046c4b125b216f6bf1cc604cff0dc", size = 129795, upload-time = "2025-06-18T14:07:40.39Z" }, + { url = "https://files.pythonhosted.org/packages/1b/98/f63318ccbe75c810011fe9233884c5d348d94d90005de1b79e5f93bef9c0/umap_learn-0.5.12-py3-none-any.whl", hash = "sha256:f2a85d2a2adcb52b541bed9b27a23ca169b56bb1b23283abeebfb8dfb8a42fe5", size = 91849, upload-time = "2026-04-08T20:03:52.561Z" }, ] [[package]] -name = "uvicorn" -version = "0.40.0" +name = "urllib3" +version = "2.5.0" source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "click" }, - { name = "h11" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/c3/d1/8f3c683c9561a4e6689dd3b1d345c815f10f86acd044ee1fb9a4dcd0b8c5/uvicorn-0.40.0.tar.gz", hash = "sha256:839676675e87e73694518b5574fd0f24c9d97b46bea16df7b8c05ea1a51071ea", size = 81761, upload-time = "2025-12-21T14:16:22.45Z" } +sdist = { url = "https://files.pythonhosted.org/packages/15/22/9ee70a2574a4f4599c47dd506532914ce044817c7752a79b6a51286319bc/urllib3-2.5.0.tar.gz", hash = "sha256:3fc47733c7e419d4bc3f6b3dc2b4f890bb743906a30d56ba4a5bfa4bbff92760", size = 393185, upload-time = "2025-06-18T14:07:41.644Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3d/d8/2083a1daa7439a66f3a48589a57d576aa117726762618f6bb09fe3798796/uvicorn-0.40.0-py3-none-any.whl", hash = "sha256:c6c8f55bc8bf13eb6fa9ff87ad62308bbbc33d0b67f84293151efe87e0d5f2ee", size = 68502, upload-time = "2025-12-21T14:16:21.041Z" }, + { url = "https://files.pythonhosted.org/packages/a7/c2/fe1e52489ae3122415c51f387e221dd0773709bad6c6cdaa599e8a2c5185/urllib3-2.5.0-py3-none-any.whl", hash = "sha256:e6b01673c0fa6a13e374b50871808eb3bf7046c4b125b216f6bf1cc604cff0dc", size = 129795, upload-time = "2025-06-18T14:07:40.39Z" }, ] [[package]] @@ -3999,97 +2596,21 @@ wheels = [ ] [[package]] -name = "websockets" -version = "16.0" +name = "xgboost" +version = "3.2.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/04/24/4b2031d72e840ce4c1ccb255f693b15c334757fc50023e4db9537080b8c4/websockets-16.0.tar.gz", hash = "sha256:5f6261a5e56e8d5c42a4497b364ea24d94d9563e8fbd44e78ac40879c60179b5", size = 179346, upload-time = "2026-01-10T09:23:47.181Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/84/7b/bac442e6b96c9d25092695578dda82403c77936104b5682307bd4deb1ad4/websockets-16.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:71c989cbf3254fbd5e84d3bff31e4da39c43f884e64f2551d14bb3c186230f00", size = 177365, upload-time = "2026-01-10T09:22:46.787Z" }, - { url = "https://files.pythonhosted.org/packages/b0/fe/136ccece61bd690d9c1f715baaeefd953bb2360134de73519d5df19d29ca/websockets-16.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:8b6e209ffee39ff1b6d0fa7bfef6de950c60dfb91b8fcead17da4ee539121a79", size = 175038, upload-time = "2026-01-10T09:22:47.999Z" }, - { url = "https://files.pythonhosted.org/packages/40/1e/9771421ac2286eaab95b8575b0cb701ae3663abf8b5e1f64f1fd90d0a673/websockets-16.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:86890e837d61574c92a97496d590968b23c2ef0aeb8a9bc9421d174cd378ae39", size = 175328, upload-time = "2026-01-10T09:22:49.809Z" }, - { url = "https://files.pythonhosted.org/packages/18/29/71729b4671f21e1eaa5d6573031ab810ad2936c8175f03f97f3ff164c802/websockets-16.0-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:9b5aca38b67492ef518a8ab76851862488a478602229112c4b0d58d63a7a4d5c", size = 184915, upload-time = "2026-01-10T09:22:51.071Z" }, - { url = "https://files.pythonhosted.org/packages/97/bb/21c36b7dbbafc85d2d480cd65df02a1dc93bf76d97147605a8e27ff9409d/websockets-16.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e0334872c0a37b606418ac52f6ab9cfd17317ac26365f7f65e203e2d0d0d359f", size = 186152, upload-time = "2026-01-10T09:22:52.224Z" }, - { url = "https://files.pythonhosted.org/packages/4a/34/9bf8df0c0cf88fa7bfe36678dc7b02970c9a7d5e065a3099292db87b1be2/websockets-16.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a0b31e0b424cc6b5a04b8838bbaec1688834b2383256688cf47eb97412531da1", size = 185583, upload-time = "2026-01-10T09:22:53.443Z" }, - { url = "https://files.pythonhosted.org/packages/47/88/4dd516068e1a3d6ab3c7c183288404cd424a9a02d585efbac226cb61ff2d/websockets-16.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:485c49116d0af10ac698623c513c1cc01c9446c058a4e61e3bf6c19dff7335a2", size = 184880, upload-time = "2026-01-10T09:22:55.033Z" }, - { url = "https://files.pythonhosted.org/packages/91/d6/7d4553ad4bf1c0421e1ebd4b18de5d9098383b5caa1d937b63df8d04b565/websockets-16.0-cp312-cp312-win32.whl", hash = "sha256:eaded469f5e5b7294e2bdca0ab06becb6756ea86894a47806456089298813c89", size = 178261, upload-time = "2026-01-10T09:22:56.251Z" }, - { url = "https://files.pythonhosted.org/packages/c3/f0/f3a17365441ed1c27f850a80b2bc680a0fa9505d733fe152fdf5e98c1c0b/websockets-16.0-cp312-cp312-win_amd64.whl", hash = "sha256:5569417dc80977fc8c2d43a86f78e0a5a22fee17565d78621b6bb264a115d4ea", size = 178693, upload-time = "2026-01-10T09:22:57.478Z" }, - { url = "https://files.pythonhosted.org/packages/cc/9c/baa8456050d1c1b08dd0ec7346026668cbc6f145ab4e314d707bb845bf0d/websockets-16.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:878b336ac47938b474c8f982ac2f7266a540adc3fa4ad74ae96fea9823a02cc9", size = 177364, upload-time = "2026-01-10T09:22:59.333Z" }, - { url = "https://files.pythonhosted.org/packages/7e/0c/8811fc53e9bcff68fe7de2bcbe75116a8d959ac699a3200f4847a8925210/websockets-16.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:52a0fec0e6c8d9a784c2c78276a48a2bdf099e4ccc2a4cad53b27718dbfd0230", size = 175039, upload-time = "2026-01-10T09:23:01.171Z" }, - { url = "https://files.pythonhosted.org/packages/aa/82/39a5f910cb99ec0b59e482971238c845af9220d3ab9fa76dd9162cda9d62/websockets-16.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:e6578ed5b6981005df1860a56e3617f14a6c307e6a71b4fff8c48fdc50f3ed2c", size = 175323, upload-time = "2026-01-10T09:23:02.341Z" }, - { url = "https://files.pythonhosted.org/packages/bd/28/0a25ee5342eb5d5f297d992a77e56892ecb65e7854c7898fb7d35e9b33bd/websockets-16.0-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:95724e638f0f9c350bb1c2b0a7ad0e83d9cc0c9259f3ea94e40d7b02a2179ae5", size = 184975, upload-time = "2026-01-10T09:23:03.756Z" }, - { url = "https://files.pythonhosted.org/packages/f9/66/27ea52741752f5107c2e41fda05e8395a682a1e11c4e592a809a90c6a506/websockets-16.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c0204dc62a89dc9d50d682412c10b3542d748260d743500a85c13cd1ee4bde82", size = 186203, upload-time = "2026-01-10T09:23:05.01Z" }, - { url = "https://files.pythonhosted.org/packages/37/e5/8e32857371406a757816a2b471939d51c463509be73fa538216ea52b792a/websockets-16.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:52ac480f44d32970d66763115edea932f1c5b1312de36df06d6b219f6741eed8", size = 185653, upload-time = "2026-01-10T09:23:06.301Z" }, - { url = "https://files.pythonhosted.org/packages/9b/67/f926bac29882894669368dc73f4da900fcdf47955d0a0185d60103df5737/websockets-16.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6e5a82b677f8f6f59e8dfc34ec06ca6b5b48bc4fcda346acd093694cc2c24d8f", size = 184920, upload-time = "2026-01-10T09:23:07.492Z" }, - { url = "https://files.pythonhosted.org/packages/3c/a1/3d6ccdcd125b0a42a311bcd15a7f705d688f73b2a22d8cf1c0875d35d34a/websockets-16.0-cp313-cp313-win32.whl", hash = "sha256:abf050a199613f64c886ea10f38b47770a65154dc37181bfaff70c160f45315a", size = 178255, upload-time = "2026-01-10T09:23:09.245Z" }, - { url = "https://files.pythonhosted.org/packages/6b/ae/90366304d7c2ce80f9b826096a9e9048b4bb760e44d3b873bb272cba696b/websockets-16.0-cp313-cp313-win_amd64.whl", hash = "sha256:3425ac5cf448801335d6fdc7ae1eb22072055417a96cc6b31b3861f455fbc156", size = 178689, upload-time = "2026-01-10T09:23:10.483Z" }, - { url = "https://files.pythonhosted.org/packages/f3/1d/e88022630271f5bd349ed82417136281931e558d628dd52c4d8621b4a0b2/websockets-16.0-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:8cc451a50f2aee53042ac52d2d053d08bf89bcb31ae799cb4487587661c038a0", size = 177406, upload-time = "2026-01-10T09:23:12.178Z" }, - { url = "https://files.pythonhosted.org/packages/f2/78/e63be1bf0724eeb4616efb1ae1c9044f7c3953b7957799abb5915bffd38e/websockets-16.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:daa3b6ff70a9241cf6c7fc9e949d41232d9d7d26fd3522b1ad2b4d62487e9904", size = 175085, upload-time = "2026-01-10T09:23:13.511Z" }, - { url = "https://files.pythonhosted.org/packages/bb/f4/d3c9220d818ee955ae390cf319a7c7a467beceb24f05ee7aaaa2414345ba/websockets-16.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:fd3cb4adb94a2a6e2b7c0d8d05cb94e6f1c81a0cf9dc2694fb65c7e8d94c42e4", size = 175328, upload-time = "2026-01-10T09:23:14.727Z" }, - { url = "https://files.pythonhosted.org/packages/63/bc/d3e208028de777087e6fb2b122051a6ff7bbcca0d6df9d9c2bf1dd869ae9/websockets-16.0-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:781caf5e8eee67f663126490c2f96f40906594cb86b408a703630f95550a8c3e", size = 185044, upload-time = "2026-01-10T09:23:15.939Z" }, - { url = "https://files.pythonhosted.org/packages/ad/6e/9a0927ac24bd33a0a9af834d89e0abc7cfd8e13bed17a86407a66773cc0e/websockets-16.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:caab51a72c51973ca21fa8a18bd8165e1a0183f1ac7066a182ff27107b71e1a4", size = 186279, upload-time = "2026-01-10T09:23:17.148Z" }, - { url = "https://files.pythonhosted.org/packages/b9/ca/bf1c68440d7a868180e11be653c85959502efd3a709323230314fda6e0b3/websockets-16.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:19c4dc84098e523fd63711e563077d39e90ec6702aff4b5d9e344a60cb3c0cb1", size = 185711, upload-time = "2026-01-10T09:23:18.372Z" }, - { url = "https://files.pythonhosted.org/packages/c4/f8/fdc34643a989561f217bb477cbc47a3a07212cbda91c0e4389c43c296ebf/websockets-16.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:a5e18a238a2b2249c9a9235466b90e96ae4795672598a58772dd806edc7ac6d3", size = 184982, upload-time = "2026-01-10T09:23:19.652Z" }, - { url = "https://files.pythonhosted.org/packages/dd/d1/574fa27e233764dbac9c52730d63fcf2823b16f0856b3329fc6268d6ae4f/websockets-16.0-cp314-cp314-win32.whl", hash = "sha256:a069d734c4a043182729edd3e9f247c3b2a4035415a9172fd0f1b71658a320a8", size = 177915, upload-time = "2026-01-10T09:23:21.458Z" }, - { url = "https://files.pythonhosted.org/packages/8a/f1/ae6b937bf3126b5134ce1f482365fde31a357c784ac51852978768b5eff4/websockets-16.0-cp314-cp314-win_amd64.whl", hash = "sha256:c0ee0e63f23914732c6d7e0cce24915c48f3f1512ec1d079ed01fc629dab269d", size = 178381, upload-time = "2026-01-10T09:23:22.715Z" }, - { url = "https://files.pythonhosted.org/packages/06/9b/f791d1db48403e1f0a27577a6beb37afae94254a8c6f08be4a23e4930bc0/websockets-16.0-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:a35539cacc3febb22b8f4d4a99cc79b104226a756aa7400adc722e83b0d03244", size = 177737, upload-time = "2026-01-10T09:23:24.523Z" }, - { url = "https://files.pythonhosted.org/packages/bd/40/53ad02341fa33b3ce489023f635367a4ac98b73570102ad2cdd770dacc9a/websockets-16.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:b784ca5de850f4ce93ec85d3269d24d4c82f22b7212023c974c401d4980ebc5e", size = 175268, upload-time = "2026-01-10T09:23:25.781Z" }, - { url = "https://files.pythonhosted.org/packages/74/9b/6158d4e459b984f949dcbbb0c5d270154c7618e11c01029b9bbd1bb4c4f9/websockets-16.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:569d01a4e7fba956c5ae4fc988f0d4e187900f5497ce46339c996dbf24f17641", size = 175486, upload-time = "2026-01-10T09:23:27.033Z" }, - { url = "https://files.pythonhosted.org/packages/e5/2d/7583b30208b639c8090206f95073646c2c9ffd66f44df967981a64f849ad/websockets-16.0-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:50f23cdd8343b984957e4077839841146f67a3d31ab0d00e6b824e74c5b2f6e8", size = 185331, upload-time = "2026-01-10T09:23:28.259Z" }, - { url = "https://files.pythonhosted.org/packages/45/b0/cce3784eb519b7b5ad680d14b9673a31ab8dcb7aad8b64d81709d2430aa8/websockets-16.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:152284a83a00c59b759697b7f9e9cddf4e3c7861dd0d964b472b70f78f89e80e", size = 186501, upload-time = "2026-01-10T09:23:29.449Z" }, - { url = "https://files.pythonhosted.org/packages/19/60/b8ebe4c7e89fb5f6cdf080623c9d92789a53636950f7abacfc33fe2b3135/websockets-16.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:bc59589ab64b0022385f429b94697348a6a234e8ce22544e3681b2e9331b5944", size = 186062, upload-time = "2026-01-10T09:23:31.368Z" }, - { url = "https://files.pythonhosted.org/packages/88/a8/a080593f89b0138b6cba1b28f8df5673b5506f72879322288b031337c0b8/websockets-16.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:32da954ffa2814258030e5a57bc73a3635463238e797c7375dc8091327434206", size = 185356, upload-time = "2026-01-10T09:23:32.627Z" }, - { url = "https://files.pythonhosted.org/packages/c2/b6/b9afed2afadddaf5ebb2afa801abf4b0868f42f8539bfe4b071b5266c9fe/websockets-16.0-cp314-cp314t-win32.whl", hash = "sha256:5a4b4cc550cb665dd8a47f868c8d04c8230f857363ad3c9caf7a0c3bf8c61ca6", size = 178085, upload-time = "2026-01-10T09:23:33.816Z" }, - { url = "https://files.pythonhosted.org/packages/9f/3e/28135a24e384493fa804216b79a6a6759a38cc4ff59118787b9fb693df93/websockets-16.0-cp314-cp314t-win_amd64.whl", hash = "sha256:b14dc141ed6d2dde437cddb216004bcac6a1df0935d79656387bd41632ba0bbd", size = 178531, upload-time = "2026-01-10T09:23:35.016Z" }, - { url = "https://files.pythonhosted.org/packages/6f/28/258ebab549c2bf3e64d2b0217b973467394a9cea8c42f70418ca2c5d0d2e/websockets-16.0-py3-none-any.whl", hash = "sha256:1637db62fad1dc833276dded54215f2c7fa46912301a24bd94d45d46a011ceec", size = 171598, upload-time = "2026-01-10T09:23:45.395Z" }, +dependencies = [ + { name = "numpy" }, + { name = "nvidia-nccl-cu12", marker = "sys_platform == 'linux'" }, + { name = "scipy" }, ] - -[[package]] -name = "wrapt" -version = "1.17.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/95/8f/aeb76c5b46e273670962298c23e7ddde79916cb74db802131d49a85e4b7d/wrapt-1.17.3.tar.gz", hash = "sha256:f66eb08feaa410fe4eebd17f2a2c8e2e46d3476e9f8c783daa8e09e0faa666d0", size = 55547, upload-time = "2025-08-12T05:53:21.714Z" } +sdist = { url = "https://files.pythonhosted.org/packages/91/bb/1eb0242409d22db725d7a88088e6cfd6556829fb0736f9ff69aa9f1e9455/xgboost-3.2.0.tar.gz", hash = "sha256:99b0e9a2a64896cdaf509c5e46372d336c692406646d20f2af505003c0c5d70d", size = 1263936, upload-time = "2026-02-10T11:03:05.542Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9f/41/cad1aba93e752f1f9268c77270da3c469883d56e2798e7df6240dcb2287b/wrapt-1.17.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:ab232e7fdb44cdfbf55fc3afa31bcdb0d8980b9b95c38b6405df2acb672af0e0", size = 53998, upload-time = "2025-08-12T05:51:47.138Z" }, - { url = "https://files.pythonhosted.org/packages/60/f8/096a7cc13097a1869fe44efe68dace40d2a16ecb853141394047f0780b96/wrapt-1.17.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:9baa544e6acc91130e926e8c802a17f3b16fbea0fd441b5a60f5cf2cc5c3deba", size = 39020, upload-time = "2025-08-12T05:51:35.906Z" }, - { url = "https://files.pythonhosted.org/packages/33/df/bdf864b8997aab4febb96a9ae5c124f700a5abd9b5e13d2a3214ec4be705/wrapt-1.17.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6b538e31eca1a7ea4605e44f81a48aa24c4632a277431a6ed3f328835901f4fd", size = 39098, upload-time = "2025-08-12T05:51:57.474Z" }, - { url = "https://files.pythonhosted.org/packages/9f/81/5d931d78d0eb732b95dc3ddaeeb71c8bb572fb01356e9133916cd729ecdd/wrapt-1.17.3-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:042ec3bb8f319c147b1301f2393bc19dba6e176b7da446853406d041c36c7828", size = 88036, upload-time = "2025-08-12T05:52:34.784Z" }, - { url = "https://files.pythonhosted.org/packages/ca/38/2e1785df03b3d72d34fc6252d91d9d12dc27a5c89caef3335a1bbb8908ca/wrapt-1.17.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3af60380ba0b7b5aeb329bc4e402acd25bd877e98b3727b0135cb5c2efdaefe9", size = 88156, upload-time = "2025-08-12T05:52:13.599Z" }, - { url = "https://files.pythonhosted.org/packages/b3/8b/48cdb60fe0603e34e05cffda0b2a4adab81fd43718e11111a4b0100fd7c1/wrapt-1.17.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:0b02e424deef65c9f7326d8c19220a2c9040c51dc165cddb732f16198c168396", size = 87102, upload-time = "2025-08-12T05:52:14.56Z" }, - { url = "https://files.pythonhosted.org/packages/3c/51/d81abca783b58f40a154f1b2c56db1d2d9e0d04fa2d4224e357529f57a57/wrapt-1.17.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:74afa28374a3c3a11b3b5e5fca0ae03bef8450d6aa3ab3a1e2c30e3a75d023dc", size = 87732, upload-time = "2025-08-12T05:52:36.165Z" }, - { url = "https://files.pythonhosted.org/packages/9e/b1/43b286ca1392a006d5336412d41663eeef1ad57485f3e52c767376ba7e5a/wrapt-1.17.3-cp312-cp312-win32.whl", hash = "sha256:4da9f45279fff3543c371d5ababc57a0384f70be244de7759c85a7f989cb4ebe", size = 36705, upload-time = "2025-08-12T05:53:07.123Z" }, - { url = "https://files.pythonhosted.org/packages/28/de/49493f962bd3c586ab4b88066e967aa2e0703d6ef2c43aa28cb83bf7b507/wrapt-1.17.3-cp312-cp312-win_amd64.whl", hash = "sha256:e71d5c6ebac14875668a1e90baf2ea0ef5b7ac7918355850c0908ae82bcb297c", size = 38877, upload-time = "2025-08-12T05:53:05.436Z" }, - { url = "https://files.pythonhosted.org/packages/f1/48/0f7102fe9cb1e8a5a77f80d4f0956d62d97034bbe88d33e94699f99d181d/wrapt-1.17.3-cp312-cp312-win_arm64.whl", hash = "sha256:604d076c55e2fdd4c1c03d06dc1a31b95130010517b5019db15365ec4a405fc6", size = 36885, upload-time = "2025-08-12T05:52:54.367Z" }, - { url = "https://files.pythonhosted.org/packages/fc/f6/759ece88472157acb55fc195e5b116e06730f1b651b5b314c66291729193/wrapt-1.17.3-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:a47681378a0439215912ef542c45a783484d4dd82bac412b71e59cf9c0e1cea0", size = 54003, upload-time = "2025-08-12T05:51:48.627Z" }, - { url = "https://files.pythonhosted.org/packages/4f/a9/49940b9dc6d47027dc850c116d79b4155f15c08547d04db0f07121499347/wrapt-1.17.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:54a30837587c6ee3cd1a4d1c2ec5d24e77984d44e2f34547e2323ddb4e22eb77", size = 39025, upload-time = "2025-08-12T05:51:37.156Z" }, - { url = "https://files.pythonhosted.org/packages/45/35/6a08de0f2c96dcdd7fe464d7420ddb9a7655a6561150e5fc4da9356aeaab/wrapt-1.17.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:16ecf15d6af39246fe33e507105d67e4b81d8f8d2c6598ff7e3ca1b8a37213f7", size = 39108, upload-time = "2025-08-12T05:51:58.425Z" }, - { url = "https://files.pythonhosted.org/packages/0c/37/6faf15cfa41bf1f3dba80cd3f5ccc6622dfccb660ab26ed79f0178c7497f/wrapt-1.17.3-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:6fd1ad24dc235e4ab88cda009e19bf347aabb975e44fd5c2fb22a3f6e4141277", size = 88072, upload-time = "2025-08-12T05:52:37.53Z" }, - { url = "https://files.pythonhosted.org/packages/78/f2/efe19ada4a38e4e15b6dff39c3e3f3f73f5decf901f66e6f72fe79623a06/wrapt-1.17.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0ed61b7c2d49cee3c027372df5809a59d60cf1b6c2f81ee980a091f3afed6a2d", size = 88214, upload-time = "2025-08-12T05:52:15.886Z" }, - { url = "https://files.pythonhosted.org/packages/40/90/ca86701e9de1622b16e09689fc24b76f69b06bb0150990f6f4e8b0eeb576/wrapt-1.17.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:423ed5420ad5f5529db9ce89eac09c8a2f97da18eb1c870237e84c5a5c2d60aa", size = 87105, upload-time = "2025-08-12T05:52:17.914Z" }, - { url = "https://files.pythonhosted.org/packages/fd/e0/d10bd257c9a3e15cbf5523025252cc14d77468e8ed644aafb2d6f54cb95d/wrapt-1.17.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:e01375f275f010fcbf7f643b4279896d04e571889b8a5b3f848423d91bf07050", size = 87766, upload-time = "2025-08-12T05:52:39.243Z" }, - { url = "https://files.pythonhosted.org/packages/e8/cf/7d848740203c7b4b27eb55dbfede11aca974a51c3d894f6cc4b865f42f58/wrapt-1.17.3-cp313-cp313-win32.whl", hash = "sha256:53e5e39ff71b3fc484df8a522c933ea2b7cdd0d5d15ae82e5b23fde87d44cbd8", size = 36711, upload-time = "2025-08-12T05:53:10.074Z" }, - { url = "https://files.pythonhosted.org/packages/57/54/35a84d0a4d23ea675994104e667ceff49227ce473ba6a59ba2c84f250b74/wrapt-1.17.3-cp313-cp313-win_amd64.whl", hash = "sha256:1f0b2f40cf341ee8cc1a97d51ff50dddb9fcc73241b9143ec74b30fc4f44f6cb", size = 38885, upload-time = "2025-08-12T05:53:08.695Z" }, - { url = "https://files.pythonhosted.org/packages/01/77/66e54407c59d7b02a3c4e0af3783168fff8e5d61def52cda8728439d86bc/wrapt-1.17.3-cp313-cp313-win_arm64.whl", hash = "sha256:7425ac3c54430f5fc5e7b6f41d41e704db073309acfc09305816bc6a0b26bb16", size = 36896, upload-time = "2025-08-12T05:52:55.34Z" }, - { url = "https://files.pythonhosted.org/packages/02/a2/cd864b2a14f20d14f4c496fab97802001560f9f41554eef6df201cd7f76c/wrapt-1.17.3-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:cf30f6e3c077c8e6a9a7809c94551203c8843e74ba0c960f4a98cd80d4665d39", size = 54132, upload-time = "2025-08-12T05:51:49.864Z" }, - { url = "https://files.pythonhosted.org/packages/d5/46/d011725b0c89e853dc44cceb738a307cde5d240d023d6d40a82d1b4e1182/wrapt-1.17.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:e228514a06843cae89621384cfe3a80418f3c04aadf8a3b14e46a7be704e4235", size = 39091, upload-time = "2025-08-12T05:51:38.935Z" }, - { url = "https://files.pythonhosted.org/packages/2e/9e/3ad852d77c35aae7ddebdbc3b6d35ec8013af7d7dddad0ad911f3d891dae/wrapt-1.17.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:5ea5eb3c0c071862997d6f3e02af1d055f381b1d25b286b9d6644b79db77657c", size = 39172, upload-time = "2025-08-12T05:51:59.365Z" }, - { url = "https://files.pythonhosted.org/packages/c3/f7/c983d2762bcce2326c317c26a6a1e7016f7eb039c27cdf5c4e30f4160f31/wrapt-1.17.3-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:281262213373b6d5e4bb4353bc36d1ba4084e6d6b5d242863721ef2bf2c2930b", size = 87163, upload-time = "2025-08-12T05:52:40.965Z" }, - { url = "https://files.pythonhosted.org/packages/e4/0f/f673f75d489c7f22d17fe0193e84b41540d962f75fce579cf6873167c29b/wrapt-1.17.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:dc4a8d2b25efb6681ecacad42fca8859f88092d8732b170de6a5dddd80a1c8fa", size = 87963, upload-time = "2025-08-12T05:52:20.326Z" }, - { url = "https://files.pythonhosted.org/packages/df/61/515ad6caca68995da2fac7a6af97faab8f78ebe3bf4f761e1b77efbc47b5/wrapt-1.17.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:373342dd05b1d07d752cecbec0c41817231f29f3a89aa8b8843f7b95992ed0c7", size = 86945, upload-time = "2025-08-12T05:52:21.581Z" }, - { url = "https://files.pythonhosted.org/packages/d3/bd/4e70162ce398462a467bc09e768bee112f1412e563620adc353de9055d33/wrapt-1.17.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:d40770d7c0fd5cbed9d84b2c3f2e156431a12c9a37dc6284060fb4bec0b7ffd4", size = 86857, upload-time = "2025-08-12T05:52:43.043Z" }, - { url = "https://files.pythonhosted.org/packages/2b/b8/da8560695e9284810b8d3df8a19396a6e40e7518059584a1a394a2b35e0a/wrapt-1.17.3-cp314-cp314-win32.whl", hash = "sha256:fbd3c8319de8e1dc79d346929cd71d523622da527cca14e0c1d257e31c2b8b10", size = 37178, upload-time = "2025-08-12T05:53:12.605Z" }, - { url = "https://files.pythonhosted.org/packages/db/c8/b71eeb192c440d67a5a0449aaee2310a1a1e8eca41676046f99ed2487e9f/wrapt-1.17.3-cp314-cp314-win_amd64.whl", hash = "sha256:e1a4120ae5705f673727d3253de3ed0e016f7cd78dc463db1b31e2463e1f3cf6", size = 39310, upload-time = "2025-08-12T05:53:11.106Z" }, - { url = "https://files.pythonhosted.org/packages/45/20/2cda20fd4865fa40f86f6c46ed37a2a8356a7a2fde0773269311f2af56c7/wrapt-1.17.3-cp314-cp314-win_arm64.whl", hash = "sha256:507553480670cab08a800b9463bdb881b2edeed77dc677b0a5915e6106e91a58", size = 37266, upload-time = "2025-08-12T05:52:56.531Z" }, - { url = "https://files.pythonhosted.org/packages/77/ed/dd5cf21aec36c80443c6f900449260b80e2a65cf963668eaef3b9accce36/wrapt-1.17.3-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:ed7c635ae45cfbc1a7371f708727bf74690daedc49b4dba310590ca0bd28aa8a", size = 56544, upload-time = "2025-08-12T05:51:51.109Z" }, - { url = "https://files.pythonhosted.org/packages/8d/96/450c651cc753877ad100c7949ab4d2e2ecc4d97157e00fa8f45df682456a/wrapt-1.17.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:249f88ed15503f6492a71f01442abddd73856a0032ae860de6d75ca62eed8067", size = 40283, upload-time = "2025-08-12T05:51:39.912Z" }, - { url = "https://files.pythonhosted.org/packages/d1/86/2fcad95994d9b572db57632acb6f900695a648c3e063f2cd344b3f5c5a37/wrapt-1.17.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:5a03a38adec8066d5a37bea22f2ba6bbf39fcdefbe2d91419ab864c3fb515454", size = 40366, upload-time = "2025-08-12T05:52:00.693Z" }, - { url = "https://files.pythonhosted.org/packages/64/0e/f4472f2fdde2d4617975144311f8800ef73677a159be7fe61fa50997d6c0/wrapt-1.17.3-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:5d4478d72eb61c36e5b446e375bbc49ed002430d17cdec3cecb36993398e1a9e", size = 108571, upload-time = "2025-08-12T05:52:44.521Z" }, - { url = "https://files.pythonhosted.org/packages/cc/01/9b85a99996b0a97c8a17484684f206cbb6ba73c1ce6890ac668bcf3838fb/wrapt-1.17.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:223db574bb38637e8230eb14b185565023ab624474df94d2af18f1cdb625216f", size = 113094, upload-time = "2025-08-12T05:52:22.618Z" }, - { url = "https://files.pythonhosted.org/packages/25/02/78926c1efddcc7b3aa0bc3d6b33a822f7d898059f7cd9ace8c8318e559ef/wrapt-1.17.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e405adefb53a435f01efa7ccdec012c016b5a1d3f35459990afc39b6be4d5056", size = 110659, upload-time = "2025-08-12T05:52:24.057Z" }, - { url = "https://files.pythonhosted.org/packages/dc/ee/c414501ad518ac3e6fe184753632fe5e5ecacdcf0effc23f31c1e4f7bfcf/wrapt-1.17.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:88547535b787a6c9ce4086917b6e1d291aa8ed914fdd3a838b3539dc95c12804", size = 106946, upload-time = "2025-08-12T05:52:45.976Z" }, - { url = "https://files.pythonhosted.org/packages/be/44/a1bd64b723d13bb151d6cc91b986146a1952385e0392a78567e12149c7b4/wrapt-1.17.3-cp314-cp314t-win32.whl", hash = "sha256:41b1d2bc74c2cac6f9074df52b2efbef2b30bdfe5f40cb78f8ca22963bc62977", size = 38717, upload-time = "2025-08-12T05:53:15.214Z" }, - { url = "https://files.pythonhosted.org/packages/79/d9/7cfd5a312760ac4dd8bf0184a6ee9e43c33e47f3dadc303032ce012b8fa3/wrapt-1.17.3-cp314-cp314t-win_amd64.whl", hash = "sha256:73d496de46cd2cdbdbcce4ae4bcdb4afb6a11234a1df9c085249d55166b95116", size = 41334, upload-time = "2025-08-12T05:53:14.178Z" }, - { url = "https://files.pythonhosted.org/packages/46/78/10ad9781128ed2f99dbc474f43283b13fea8ba58723e98844367531c18e9/wrapt-1.17.3-cp314-cp314t-win_arm64.whl", hash = "sha256:f38e60678850c42461d4202739f9bf1e3a737c7ad283638251e79cc49effb6b6", size = 38471, upload-time = "2025-08-12T05:52:57.784Z" }, - { url = "https://files.pythonhosted.org/packages/1f/f6/a933bd70f98e9cf3e08167fc5cd7aaaca49147e48411c0bd5ae701bb2194/wrapt-1.17.3-py3-none-any.whl", hash = "sha256:7171ae35d2c33d326ac19dd8facb1e82e5fd04ef8c6c0e394d7af55a55051c22", size = 23591, upload-time = "2025-08-12T05:53:20.674Z" }, + { url = "https://files.pythonhosted.org/packages/2d/49/6e4cdd877c24adf56cb3586bc96d93d4dcd780b5ea1efb32e1ee0de08bae/xgboost-3.2.0-py3-none-macosx_10_15_x86_64.whl", hash = "sha256:2f661966d3e322536d9c448090a870fcba1e32ee5760c10b7c46bac7a342079a", size = 2507014, upload-time = "2026-02-10T10:50:57.44Z" }, + { url = "https://files.pythonhosted.org/packages/93/f1/c09ef1add609453aa3ba5bafcd0d1c1a805c1263c0b60138ec968f8ec296/xgboost-3.2.0-py3-none-macosx_12_0_arm64.whl", hash = "sha256:eabbd40d474b8dbf6cb3536325f9150b9e6f0db32d18de9914fb3227d0bef5b7", size = 2328527, upload-time = "2026-02-10T10:51:17.502Z" }, + { url = "https://files.pythonhosted.org/packages/96/9f/d9914a7b8df842832850b1a18e5f47aaa071c217cdd1da2ae9deb291018b/xgboost-3.2.0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:852eabc6d3b3702a59bf78dbfdcd1cb9c4d3a3b6e5ed1f8781d8b9512354fdd2", size = 131100954, upload-time = "2026-02-10T11:02:42.704Z" }, + { url = "https://files.pythonhosted.org/packages/79/98/679de17c2caa4fd3b0b4386ecf7377301702cb0afb22930a07c142fcb1d8/xgboost-3.2.0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:99b4a6bbcb47212fec5cf5fbe12347215f073c08967431b0122cfbd1ee70312c", size = 131748579, upload-time = "2026-02-10T10:54:40.424Z" }, + { url = "https://files.pythonhosted.org/packages/1f/3d/1661dd114a914a67e3f7ab66fa1382e7599c2a8c340f314ad30a3e2b4d08/xgboost-3.2.0-py3-none-win_amd64.whl", hash = "sha256:0d169736fd836fc13646c7ab787167b3a8110351c2c6bc770c755ee1618f0442", size = 101681668, upload-time = "2026-02-10T10:59:31.202Z" }, ] [[package]] @@ -4194,12 +2715,3 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/48/b7/503c98092fb3b344a179579f55814b613c1fbb1c23b3ec14a7b008a66a6e/yarl-1.22.0-cp314-cp314t-win_arm64.whl", hash = "sha256:9f6d73c1436b934e3f01df1e1b21ff765cd1d28c77dfb9ace207f746d4610ee1", size = 85171, upload-time = "2025-10-06T14:12:16.935Z" }, { url = "https://files.pythonhosted.org/packages/73/ae/b48f95715333080afb75a4504487cbe142cae1268afc482d06692d605ae6/yarl-1.22.0-py3-none-any.whl", hash = "sha256:1380560bdba02b6b6c90de54133c81c9f2a453dee9912fe58c1dcced1edb7cff", size = 46814, upload-time = "2025-10-06T14:12:53.872Z" }, ] - -[[package]] -name = "zipp" -version = "3.23.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e3/02/0f2892c661036d50ede074e376733dca2ae7c6eb617489437771209d4180/zipp-3.23.0.tar.gz", hash = "sha256:a07157588a12518c9d4034df3fbbee09c814741a33ff63c05fa29d26a2404166", size = 25547, upload-time = "2025-06-08T17:06:39.4Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/2e/54/647ade08bf0db230bfea292f893923872fd20be6ac6f53b2b936ba839d75/zipp-3.23.0-py3-none-any.whl", hash = "sha256:071652d6115ed432f5ce1d34c336c0adfd6a884660d1e9712a256d3d3bd4b14e", size = 10276, upload-time = "2025-06-08T17:06:38.034Z" }, -]