From 0fbf6062ae8a2248c2b2349a384a5aa62a70a519 Mon Sep 17 00:00:00 2001 From: hill Date: Mon, 2 Mar 2026 16:47:24 -0500 Subject: [PATCH 001/121] fix: use bf16-mixed precision and gradient accumulation to avoid OOM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Observation-level MAE tokenization produces up to 5376 tokens per sample (48 timesteps × 112 features), causing OOM on L4 (22GB) with batch_size=256 in fp32 due to quadratic attention memory. - batch_size: 256 → 64 (4× less peak memory) - precision: 32 → bf16-mixed (2× less memory, native L4 tensor core support) - accumulate_grad_batches: 1 → 4 (effective batch size stays 256) --- configs/pretrain.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/configs/pretrain.yaml b/configs/pretrain.yaml index 5e1d0c7..61e9c01 100644 --- a/configs/pretrain.yaml +++ b/configs/pretrain.yaml @@ -51,16 +51,16 @@ encoder: # Training configuration (same budget across thesis paradigms for fair comparison) training: max_epochs: 500 - batch_size: 256 + batch_size: 64 # Compute settings accelerator: auto # auto | cpu | gpu | mps devices: auto # auto | 1 | [0,1] | etc. - precision: 32 # 32 | 16 | bf16 + precision: bf16-mixed # 32 | 16-mixed | bf16-mixed # Regularization gradient_clip_val: 1.0 - accumulate_grad_batches: 1 + accumulate_grad_batches: 4 # effective batch_size = 64 * 4 = 256 # Early stopping early_stopping_patience: 10 From f73e4b6ebf57c6d9485ee9845f12b593a0e6e8fd Mon Sep 17 00:00:00 2001 From: hill Date: Mon, 2 Mar 2026 16:51:23 -0500 Subject: [PATCH 002/121] fix: cast vis_proj dtype in MAE decoder scatter to avoid bf16-mixed mismatch --- src/slices/models/pretraining/mae.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/slices/models/pretraining/mae.py b/src/slices/models/pretraining/mae.py index b1b04c6..9f639a7 100644 --- a/src/slices/models/pretraining/mae.py +++ b/src/slices/models/pretraining/mae.py @@ -156,7 +156,7 @@ def forward( 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) + full_tokens.scatter_(1, scatter_idx_expanded, vis_proj.to(full_tokens.dtype)) # Add positional information to ALL token positions timestep_idx = token_info["timestep_idx"] # (B, max_tokens) From 4b44760a6f5e90a31951cc7ee20407c9574d2cfb Mon Sep 17 00:00:00 2001 From: hill Date: Mon, 2 Mar 2026 18:43:10 -0500 Subject: [PATCH 003/121] refactor: switch SSL objectives from observation-level to timestep-level tokenization Unify all three SSL paradigms (MAE, JEPA, Contrastive) to use TransformerEncoder with obs_aware mode instead of ObservationTransformerEncoder. Each token now represents a full timestep (all D features) rather than individual observations, simplifying masking and decoding while preserving observation-aware input projection. Key changes: - Add obs_aware flag to TransformerEncoder with tokenize()/encode() API - Replace observation masking with timestep masking throughout - MAE decoder now predicts D features per timestep (not scalar per obs) - Extract build_sinusoidal_pe() to common utils - Update configs, finetune module, and all SSL tests --- configs/finetune.yaml | 2 +- configs/model/transformer.yaml | 4 + configs/pretrain.yaml | 9 +- scripts/training/pretrain.py | 6 +- src/slices/eval/imputation.py | 10 +- src/slices/models/common.py | 28 +- src/slices/models/encoders/transformer.py | 116 +++++- src/slices/models/pretraining/contrastive.py | 73 ++-- src/slices/models/pretraining/jepa.py | 140 +++----- src/slices/models/pretraining/mae.py | 208 +++++------ src/slices/models/pretraining/masking.py | 63 ++++ src/slices/training/finetune_module.py | 12 + tests/test_contrastive_objective.py | 157 ++++---- tests/test_factories.py | 24 +- tests/test_jepa_objective.py | 220 +++++++----- tests/test_mae_objective.py | 356 ++++++------------- tests/test_pretrain_module.py | 22 +- tests/test_transformer_encoder.py | 109 ++++++ 18 files changed, 864 insertions(+), 695 deletions(-) diff --git a/configs/finetune.yaml b/configs/finetune.yaml index b317ba2..00ec561 100644 --- a/configs/finetune.yaml +++ b/configs/finetune.yaml @@ -21,7 +21,7 @@ defaults: - data: ricu - - model: observation_transformer + - model: transformer - tasks: mortality_24h - _self_ diff --git a/configs/model/transformer.yaml b/configs/model/transformer.yaml index 1d4d000..c937295 100644 --- a/configs/model/transformer.yaml +++ b/configs/model/transformer.yaml @@ -45,3 +45,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 +# Enable for SSL pretraining; disabled by default for supervised baselines +obs_aware: false diff --git a/configs/pretrain.yaml b/configs/pretrain.yaml index 61e9c01..e144d00 100644 --- a/configs/pretrain.yaml +++ b/configs/pretrain.yaml @@ -12,7 +12,7 @@ defaults: - data: ricu - - model: observation_transformer + - model: transformer - ssl: mae - _self_ @@ -44,9 +44,14 @@ data: window_size: 48 # Window size in hours (48h = typical ICU prediction horizon) window_stride: 24 # Stride for overlapping windows (50% overlap) -# Encoder overrides (base from configs/model/observation_transformer.yaml: d=128, L=4, H=8) +# Encoder overrides for SSL pretraining encoder: pooling: none # Required for SSL - need per-token outputs (vs 'mean' for tasks) + obs_aware: true # Enable obs-aware MLP projection for SSL + d_model: 128 + n_layers: 4 + n_heads: 8 + d_ff: 512 # Training configuration (same budget across thesis paradigms for fair comparison) training: diff --git a/scripts/training/pretrain.py b/scripts/training/pretrain.py index 03c91ca..7c6cecf 100644 --- a/scripts/training/pretrain.py +++ b/scripts/training/pretrain.py @@ -29,9 +29,9 @@ # SSL -> compatible model encoder mappings _SSL_MODEL_COMPAT = { - "mae": {"observation_transformer"}, - "jepa": {"observation_transformer"}, - "contrastive": {"observation_transformer"}, + "mae": {"transformer"}, + "jepa": {"transformer"}, + "contrastive": {"transformer"}, "smart": {"smart"}, } diff --git a/src/slices/eval/imputation.py b/src/slices/eval/imputation.py index 1cdf756..10b5404 100644 --- a/src/slices/eval/imputation.py +++ b/src/slices/eval/imputation.py @@ -23,6 +23,8 @@ import torch.nn as nn from torch.utils.data import DataLoader +from slices.models.encoders import ObservationTransformerEncoder + logger = logging.getLogger(__name__) @@ -189,7 +191,9 @@ 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) encoder = encoder.to(device).eval() @@ -239,8 +243,10 @@ 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) else: raise ValueError( diff --git a/src/slices/models/common.py b/src/slices/models/common.py index 0c97640..dd86806 100644 --- a/src/slices/models/common.py +++ b/src/slices/models/common.py @@ -92,6 +92,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 +127,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/transformer.py b/src/slices/models/encoders/transformer.py index 7f06224..7cb645a 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,20 @@ 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.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": @@ -229,6 +247,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) + + # Fixed-length: all timesteps valid + padding_mask = torch.ones(B, T, dtype=torch.bool, device=device) + + token_info = { + "timestep_idx": torch.arange(T, device=device).unsqueeze(0).expand(B, -1), + "values": x, + "obs_mask": obs_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. + """ + # Convert to PyTorch convention: True = ignore + key_padding_mask = ~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 +348,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,8 +368,9 @@ 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 diff --git a/src/slices/models/pretraining/contrastive.py b/src/slices/models/pretraining/contrastive.py index b67e010..a0645c4 100644 --- a/src/slices/models/pretraining/contrastive.py +++ b/src/slices/models/pretraining/contrastive.py @@ -1,15 +1,15 @@ """Contrastive (SimCLR-style) SSL objective for ICU time-series. -Observation-level tokenization variant: uses two different random masks as +Timestep-level tokenization variant: uses two different random timestep masks as two augmented "views" of the same sample, then applies NT-Xent contrastive loss on the pooled representations. 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 +1. TransformerEncoder.tokenize() -> one token per timestep (B, T, d_model) +2. Two independent random masks -> two different subsets of timestep 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). @@ -24,12 +24,12 @@ import torch.nn.functional as F from .base import BaseSSLObjective, SSLConfig -from .masking import create_observation_mask, extract_visible +from .masking import create_timestep_mask, extract_visible_timesteps @dataclass class ContrastiveConfig(SSLConfig): - """Configuration for observation-level contrastive objective.""" + """Configuration for timestep-level contrastive objective.""" name: str = "contrastive" @@ -78,27 +78,27 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class ContrastiveObjective(BaseSSLObjective): - """Observation-level contrastive (SimCLR-style) SSL for ICU time-series. + """Timestep-level contrastive (SimCLR-style) 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 + 1. encoder.tokenize(x, obs_mask) -> timestep tokens + 2. Two independent random timestep 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 - 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"): + # Validate encoder has obs-aware tokenization + if not getattr(getattr(encoder, "config", None), "obs_aware", False): raise ValueError( - "Contrastive requires an encoder with tokenize() and encode() " - "methods (e.g., ObservationTransformerEncoder). Got: " + "Contrastive requires an encoder with obs_aware=True " + "(e.g., TransformerEncoder with obs_aware=True). Got: " f"{type(encoder).__name__}" ) @@ -140,17 +140,18 @@ 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) - # 3. View 1: extract → encode → mean pool - vis_tokens_1, vis_padding_1 = extract_visible(tokens, ssl_mask_1, padding_mask) + # 2. Two independent random timestep masks + ssl_mask_1 = create_timestep_mask(B, T, self.config.mask_ratio, device) + ssl_mask_2 = create_timestep_mask(B, T, self.config.mask_ratio, device) + + # 3. View 1: extract -> encode -> mean pool + vis_tokens_1, vis_padding_1 = extract_visible_timesteps(tokens, ssl_mask_1) 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) + # 4. View 2: extract -> encode -> mean pool + vis_tokens_2, vis_padding_2 = extract_visible_timesteps(tokens, ssl_mask_2) encoded_2 = self.encoder.encode(vis_tokens_2, vis_padding_2) pooled_2 = self._mean_pool(encoded_2, vis_padding_2) # (B, d_model) @@ -159,7 +160,7 @@ def forward( 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) + loss, metrics = self._nt_xent_loss(z1, z2, ssl_mask_1, ssl_mask_2) return loss, metrics @@ -188,16 +189,14 @@ def _nt_xent_loss( z2: torch.Tensor, ssl_mask_1: torch.Tensor, ssl_mask_2: torch.Tensor, - padding_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). Returns: (loss, metrics_dict) @@ -231,12 +230,12 @@ 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() + # Timestep statistics + T = ssl_mask_1.shape[1] + n_vis_1 = ssl_mask_1.sum().item() + n_vis_2 = ssl_mask_2.sum().item() + n_masked_1 = (~ssl_mask_1).sum().item() + n_masked_2 = (~ssl_mask_2).sum().item() metrics = { "contrastive_loss": loss.detach(), @@ -244,7 +243,7 @@ def _nt_xent_loss( "contrastive_accuracy": accuracy, "contrastive_pos_similarity": pos_sim, "contrastive_temperature": temperature, - "contrastive_n_tokens_per_sample": n_total / B, + "contrastive_n_timesteps": T, "contrastive_n_visible_view1": n_vis_1 / B, "contrastive_n_visible_view2": n_vis_2 / B, "contrastive_n_masked_view1": n_masked_1 / B, diff --git a/src/slices/models/pretraining/jepa.py b/src/slices/models/pretraining/jepa.py index 53b676a..4ed70b7 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,13 +23,15 @@ import torch.nn as nn import torch.nn.functional as F +from slices.models.common import build_sinusoidal_pe + from .base import BaseSSLObjective, SSLConfig -from .masking import create_observation_mask, extract_visible +from .masking import create_timestep_mask, extract_visible_timesteps @dataclass class JEPAConfig(SSLConfig): - """Configuration for observation-level JEPA objective.""" + """Configuration for timestep-level JEPA objective.""" name: str = "jepa" @@ -56,13 +57,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 +75,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,93 +93,78 @@ 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). + 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() + full_tokens = self.mask_token.expand(B, n_timesteps, d_pred).clone() - # Scatter visible tokens to original positions (same logic as MAEDecoder) + # Scatter visible tokens to original positions 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) - # 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 - - decoded = self.predictor(full_tokens, src_key_padding_mask=key_padding_mask) + # Run predictor transformer (no padding mask, all T valid) + decoded = self.predictor(full_tokens) # 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"): + # Validate encoder has obs-aware tokenization + if not getattr(getattr(encoder, "config", None), "obs_aware", False): raise ValueError( - "JEPA requires an encoder with tokenize() and encode() methods " - "(e.g., ObservationTransformerEncoder). Got: " + "JEPA requires an encoder with obs_aware=True " + "(e.g., TransformerEncoder with obs_aware=True). Got: " f"{type(encoder).__name__}" ) @@ -195,7 +177,6 @@ def __init__(self, encoder: nn.Module, config: JEPAConfig) -> None: ) d_encoder = encoder.get_output_dim() - n_features = encoder.config.d_input max_seq_length = encoder.config.max_seq_length self.missing_token = None @@ -209,7 +190,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 +213,36 @@ 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) + # 2. Create SSL mask on timesteps + ssl_mask = create_timestep_mask(B, T, self.config.mask_ratio, device) # 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) # 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, 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) return loss, metrics @@ -276,29 +251,25 @@ def _compute_loss( predicted: torch.Tensor, target: torch.Tensor, ssl_mask: torch.Tensor, - padding_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 + loss_mask = ~ssl_mask # (B, T) 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,16 +277,15 @@ 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() + B, T = ssl_mask.shape n_masked = loss_mask.sum().item() - n_visible = (ssl_mask & padding_mask).sum().item() + n_visible = ssl_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(B * T, 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, diff --git a/src/slices/models/pretraining/mae.py b/src/slices/models/pretraining/mae.py index 9f639a7..9b9012f 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 slices.models.common import build_sinusoidal_pe + from .base import BaseSSLObjective, SSLConfig -from .masking import create_observation_mask, extract_visible +from .masking import create_timestep_mask, extract_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,37 +92,26 @@ 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 @@ -136,70 +119,52 @@ def forward( # 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. + # Build full decoder input: mask tokens everywhere, then scatter visible + full_tokens = self.mask_token.expand(B, n_timesteps, d_dec).clone() + + # Scatter visible tokens to their original positions 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) + scatter_idx_expanded = scatter_idx.unsqueeze(-1).expand(-1, -1, d_dec) full_tokens.scatter_(1, scatter_idx_expanded, vis_proj.to(full_tokens.dtype)) - # 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) + # 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) - - # Run decoder transformer - decoded = self.decoder(full_tokens, src_key_padding_mask=key_padding_mask) + # Run decoder transformer (no padding mask needed, all T positions valid) + decoded = self.decoder(full_tokens) - # 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"): + # Validate encoder has obs-aware tokenization + if not getattr(getattr(encoder, "config", None), "obs_aware", False): raise ValueError( - "MAE requires an encoder with tokenize() and encode() methods " - "(e.g., ObservationTransformerEncoder). Got: " + "MAE requires an encoder with obs_aware=True " + "(e.g., TransformerEncoder with obs_aware=True). Got: " f"{type(encoder).__name__}" ) @@ -215,7 +180,6 @@ def __init__(self, encoder: nn.Module, config: MAEConfig) -> None: 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 +195,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 +207,15 @@ 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) + # tokens: (B, T, d_model) - # 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) + # 2. Create SSL mask on timesteps + ssl_mask = create_timestep_mask(B, T, self.config.mask_ratio, device) # 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) # 4. Encode visible tokens only encoded_visible = self.encoder.encode(visible_tokens, vis_padding) @@ -264,12 +225,11 @@ 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 masked token values - loss, metrics = self._compute_loss(predictions, true_values, ssl_mask, padding_mask) + # 6. Compute loss on observed features at masked timesteps + loss, metrics = self._compute_loss(predictions, x, ssl_mask, obs_mask) return loss, metrics @@ -278,37 +238,38 @@ def _compute_loss( predictions: torch.Tensor, true_values: torch.Tensor, ssl_mask: torch.Tensor, - padding_mask: torch.Tensor, + obs_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. 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_masked_timesteps = (~ssl_mask).sum().item() + n_visible_timesteps = ssl_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 +278,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(B * n_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..8cd9764 100644 --- a/src/slices/models/pretraining/masking.py +++ b/src/slices/models/pretraining/masking.py @@ -79,3 +79,66 @@ 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, +) -> 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. + + Returns: + ssl_mask: (B, T) bool mask, True = visible, False = masked. + """ + rand_vals = torch.rand(batch_size, n_timesteps, device=device) + ssl_mask = rand_vals >= mask_ratio # True = visible + + # Ensure at least 1 visible timestep per sample + n_visible = ssl_mask.sum(dim=1) # (B,) + needs_fix = n_visible == 0 + if needs_fix.any(): + for b in range(batch_size): + if needs_fix[b]: + ssl_mask[b, 0] = True + + return ssl_mask + + +def extract_visible_timesteps( + tokens: torch.Tensor, + ssl_mask: torch.Tensor, +) -> 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. + + Returns: + visible_tokens: (B, max_vis, d_model) + vis_padding: (B, max_vis) True = valid visible token + """ + B, T, d_model = tokens.shape + + n_visible = ssl_mask.sum(dim=1) # (B,) + max_vis = max(int(n_visible.max().item()), 1) + + # Argsort: visible (True=1) first + sort_idx = ssl_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 diff --git a/src/slices/training/finetune_module.py b/src/slices/training/finetune_module.py index f1e7aae..9244c21 100644 --- a/src/slices/training/finetune_module.py +++ b/src/slices/training/finetune_module.py @@ -24,6 +24,7 @@ EncoderWithMissingToken, ObservationTransformerEncoder, SMARTEncoder, + TransformerEncoder, build_encoder, ) from slices.models.heads import TaskHeadConfig, build_task_head @@ -415,6 +416,8 @@ def _wrap_encoder_with_missing_token(self, missing_token: Optional[torch.Tensor] # - ObservationTransformerEncoder: only tokenizes observed values # - SMARTEncoder: MLPEmbedder jointly embeds (value, mask) pairs, # so it needs to see original values with the mask bit + # - TransformerEncoder with obs_aware=True: obs_proj handles + # missingness via concatenated (values, obs_mask) input if isinstance(self.encoder, (ObservationTransformerEncoder, SMARTEncoder)): logger.info( "Skipping EncoderWithMissingToken wrapper for %s " @@ -423,6 +426,15 @@ def _wrap_encoder_with_missing_token(self, missing_token: Optional[torch.Tensor] ) return + if isinstance(self.encoder, TransformerEncoder) and getattr( + self.encoder.config, "obs_aware", False + ): + logger.info( + "Skipping EncoderWithMissingToken wrapper for obs_aware " + "TransformerEncoder (forward() handles mask via obs_proj)" + ) + return + d_input = self.encoder.config.d_input if missing_token is not None: diff --git a/tests/test_contrastive_objective.py b/tests/test_contrastive_objective.py index 28cda32..644873d 100644 --- a/tests/test_contrastive_objective.py +++ b/tests/test_contrastive_objective.py @@ -1,10 +1,8 @@ -"""Tests for observation-level Contrastive (SimCLR-style) SSL objective.""" +"""Tests for timestep-level Contrastive (SimCLR-style) SSL objective.""" import pytest import torch from slices.models.encoders import ( - ObservationTransformerConfig, - ObservationTransformerEncoder, TransformerConfig, TransformerEncoder, ) @@ -38,7 +36,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,10 +50,17 @@ 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): @@ -73,22 +77,26 @@ 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_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,10 +113,17 @@ 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): @@ -136,7 +151,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,7 +178,6 @@ 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 @@ -178,10 +192,17 @@ 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( @@ -199,16 +220,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 +234,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 @@ -242,12 +257,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,7 +271,6 @@ 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): @@ -288,10 +300,17 @@ 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( @@ -303,31 +322,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,10 +340,17 @@ 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( @@ -358,13 +366,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,10 +386,17 @@ 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( mask_ratio=0.75, @@ -431,10 +444,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 +461,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_factories.py b/tests/test_factories.py index 0793ece..41f53a8 100644 --- a/tests/test_factories.py +++ b/tests/test_factories.py @@ -351,19 +351,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 +398,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_jepa_objective.py b/tests/test_jepa_objective.py index ef0a2d1..bd7328b 100644 --- a/tests/test_jepa_objective.py +++ b/tests/test_jepa_objective.py @@ -1,10 +1,8 @@ -"""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, ) @@ -32,21 +30,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 +53,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,21 +68,17 @@ 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 @@ -100,10 +92,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 +122,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 +152,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 +167,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 +204,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 +218,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 +232,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,9 +241,9 @@ 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 @@ -249,10 +257,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 +284,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 +293,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 +318,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 +325,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 +335,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 +348,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 +372,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 +392,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 +410,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 +438,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 +464,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 +477,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 +531,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 +594,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 +611,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_mae_objective.py b/tests/test_mae_objective.py index ad72eb0..c7a93cd 100644 --- a/tests/test_mae_objective.py +++ b/tests/test_mae_objective.py @@ -1,4 +1,4 @@ -"""Tests for observation-level MAE (Masked Autoencoder) SSL objective.""" +"""Tests for timestep-level MAE (Masked Autoencoder) SSL objective.""" import pytest import torch @@ -12,8 +12,6 @@ create_timestep_mask, ) from slices.models.encoders import ( - ObservationTransformerConfig, - ObservationTransformerEncoder, TransformerConfig, TransformerEncoder, ) @@ -92,88 +90,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 +106,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 @@ -231,14 +145,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 +176,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 +194,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 +212,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 +222,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 +238,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 +282,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 +292,10 @@ 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_mae_get_encoder(self, encoder, mae_config): mae = MAEObjective(encoder, mae_config) @@ -389,10 +312,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 +335,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 +361,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 +374,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 +409,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 +441,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 +497,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 +507,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 +552,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 +592,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 +612,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 +636,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_pretrain_module.py b/tests/test_pretrain_module.py index 4dd9482..a08faf7 100644 --- a/tests/test_pretrain_module.py +++ b/tests/test_pretrain_module.py @@ -21,7 +21,7 @@ def minimal_config(): return OmegaConf.create( { "encoder": { - "name": "observation_transformer", + "name": "transformer", "d_input": 9, "d_model": 32, "n_layers": 1, @@ -33,6 +33,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 +67,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 +117,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 +313,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 +325,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 +374,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 +386,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 +446,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 +458,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", diff --git a/tests/test_transformer_encoder.py b/tests/test_transformer_encoder.py index 9c4ae44..35a0f04 100644 --- a/tests/test_transformer_encoder.py +++ b/tests/test_transformer_encoder.py @@ -860,5 +860,114 @@ 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 padding_mask.all() # All True for fixed T + assert "timestep_idx" in token_info + assert "values" in token_info + assert "obs_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_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"]) From a2d47e679ed8dba7df714abd406a24f904a4253a Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 3 Mar 2026 08:09:35 -0500 Subject: [PATCH 004/121] refactor: extract data and training utilities into focused modules MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Split large files into smaller, single-responsibility modules: - dataset.py → tensor_preprocessing.py, tensor_cache.py - datamodule.py → splits.py - finetune_module.py → checkpoint_loading.py Also move EXPERIMENT_PLAN.md to docs/. --- EXPERIMENT_PLAN.md | 387 ----------------- src/slices/data/datamodule.py | 421 ++----------------- src/slices/data/dataset.py | 491 ++++------------------ src/slices/data/splits.py | 459 ++++++++++++++++++++ src/slices/data/tensor_cache.py | 350 +++++++++++++++ src/slices/data/tensor_preprocessing.py | 278 ++++++++++++ src/slices/training/checkpoint_loading.py | 326 ++++++++++++++ src/slices/training/finetune_module.py | 289 +------------ 8 files changed, 1538 insertions(+), 1463 deletions(-) delete mode 100644 EXPERIMENT_PLAN.md create mode 100644 src/slices/data/splits.py create mode 100644 src/slices/data/tensor_cache.py create mode 100644 src/slices/data/tensor_preprocessing.py create mode 100644 src/slices/training/checkpoint_loading.py 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/src/slices/data/datamodule.py b/src/slices/data/datamodule.py index d976b24..29084fe 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_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. @@ -173,379 +173,6 @@ def __init__( 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,15 +192,33 @@ 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, + ) = compute_patient_level_splits( + self.processed_dir, + self.task_name, + self.seed, + self.train_ratio, + self.val_ratio, + self.test_ratio, + ) # 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, + # missing labels in compute_patient_level_splits. If any are still missing, # it indicates a bug. logger.debug("[Step 2/3] Creating ICUDataset") self.dataset = ICUDataset( @@ -601,7 +246,17 @@ def setup(self, stage: Optional[str] = None) -> None: # Save split information for reproducibility logger.debug("[Step 3/3] Saving split information") - self._save_split_info() + save_split_info( + self.processed_dir, + self.dataset, + self.train_indices, + self.val_indices, + self.test_indices, + self.seed, + self.train_ratio, + self.val_ratio, + self.test_ratio, + ) logger.info( f"DataModule setup complete: " diff --git a/src/slices/data/dataset.py b/src/slices/data/dataset.py index 12f7899..05a0bea 100644 --- a/src/slices/data/dataset.py +++ b/src/slices/data/dataset.py @@ -4,24 +4,34 @@ returns (timeseries, mask, labels, static_features) tuples for training. """ -import hashlib import logging -import warnings 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, +) + # 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 @@ -171,7 +181,9 @@ 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 + ) def _load_data(self) -> None: """Load data from Parquet files into memory.""" @@ -221,7 +233,14 @@ def _load_data(self) -> None: 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) + cached_tensors = load_cached_tensors( + self.data_dir, + self.normalize, + self.seq_length, + self.n_features, + self.train_indices, + self._excluded_stay_ids, + ) if cached_tensors is not None: # Use cached tensors (stacked format) self._timeseries_tensor = cached_tensors["timeseries_tensor"] @@ -235,13 +254,26 @@ def _load_data(self) -> None: 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) + cached_stats = load_normalization_stats( + self.data_dir, self.train_indices, self.normalize + ) # 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) + save_cached_tensors( + self.data_dir, + self._timeseries_tensor, + self._mask_tensor, + self.feature_means, + self.feature_stds, + self.normalize, + self.seq_length, + self.n_features, + self.train_indices, + self._excluded_stay_ids, + ) # Pre-compute labels and static features self._precompute_labels_and_static() @@ -258,82 +290,21 @@ def _precompute_tensors( 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). 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). + train_indices: Optional list of indices for training set. + cached_stats: Optional cached normalization statistics. """ 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 - - # ===================================================================== - # Step 2: Compute normalization statistics (vectorized) - # ===================================================================== - logger.debug("[2/3] Computing normalization statistics...") + # Step 1: Convert nested lists to tensors + timeseries_tensor, masks_tensor = convert_raw_to_tensors( + raw_timeseries, raw_masks, self.seq_length, self.n_features + ) + # Step 2: Compute normalization statistics self.feature_means = torch.zeros(self.n_features) self.feature_stds = torch.ones(self.n_features) @@ -343,324 +314,39 @@ 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] + # 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. @@ -919,41 +605,12 @@ def get_preprocessing_stages(self, idx: int) -> Dict[str, Dict[str, torch.Tensor 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/splits.py b/src/slices/data/splits.py new file mode 100644 index 0000000..bd14b08 --- /dev/null +++ b/src/slices/data/splits.py @@ -0,0 +1,459 @@ +"""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, +) -> 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: + 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 in order from timeseries parquet. + 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) 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 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( + 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 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. + 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. " 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(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 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] +]: + """Compute 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. + + 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). + """ + 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) + + # CRITICAL: Filter stays with missing labels BEFORE computing splits + # This ensures indices computed here match the filtered Dataset + stay_ids, excluded_stay_ids = filter_stays_with_missing_labels( + all_stay_ids, labels_df, task_name + ) + + # 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 + ) + if cached_splits is not None: + train_indices, val_indices, test_indices = cached_splits + return ( + train_indices, + val_indices, + test_indices, + static_df, + labels_df, + all_stay_ids, + stay_ids, + excluded_stay_ids, + ) + + 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(static_df["stay_id"].to_list(), 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={seed})") + 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 :]) + 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, + static_df, + labels_df, + all_stay_ids, + stay_ids, + excluded_stay_ids, + ) + + +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, +) -> 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. + """ + 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, + "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), + } + + 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/tensor_cache.py b/src/slices/data/tensor_cache.py new file mode 100644 index 0000000..45a73b7 --- /dev/null +++ b/src/slices/data/tensor_cache.py @@ -0,0 +1,350 @@ +"""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 logging +import warnings +from pathlib import Path +from typing import Any, Dict, List, Optional, Set + +import torch +import yaml + +logger = logging.getLogger(__name__) + + +def load_normalization_stats( + data_dir: Path, + current_train_indices: Optional[List[int]], + normalize: bool, +) -> 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: + data_dir: Path to data directory containing normalization_stats.yaml. + 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. + """ + if not normalize: + return None + + stats_path = 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( + 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 file for reproducibility. + + 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. + """ + stats_path = data_dir / "normalization_stats.yaml" + + stats = { + "feature_means": feature_means.tolist(), + "feature_stds": feature_stds.tolist(), + "feature_names": feature_names, + "train_indices": train_indices, + "normalize": 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( + normalize: bool, + seq_length: int, + n_features: int, + train_indices: Optional[List[int]], + excluded_stay_ids: Optional[Set[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: + normalize: Whether normalization is enabled. + seq_length: Sequence length. + n_features: Number of features. + train_indices: Optional list of training indices. + excluded_stay_ids: Optional set of excluded stay IDs. + + Returns: + Hash string to use as cache identifier. + """ + # Create a deterministic string representation of parameters + params = { + "normalize": normalize, + "seq_length": seq_length, + "n_features": 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 excluded_stay_ids: + excluded_str = ",".join(map(str, sorted(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( + data_dir: Path, + normalize: bool, + seq_length: int, + n_features: int, + train_indices: Optional[List[int]], + excluded_stay_ids: Optional[Set[int]], +) -> Path: + """Get the path to the tensor cache file. + + Args: + data_dir: Path to data directory. + normalize: Whether normalization is enabled. + seq_length: Sequence length. + n_features: Number of features. + train_indices: Optional list of training indices. + excluded_stay_ids: Optional set of excluded stay IDs. + + Returns: + Path to the tensor cache file. + """ + cache_key = get_tensor_cache_key( + normalize, seq_length, n_features, train_indices, excluded_stay_ids + ) + cache_dir = data_dir / ".tensor_cache" + return cache_dir / f"tensors_{cache_key}.pt" + + +def load_cached_tensors( + data_dir: Path, + normalize: bool, + seq_length: int, + n_features: int, + train_indices: Optional[List[int]], + excluded_stay_ids: Optional[Set[int]], +) -> Optional[Dict[str, Any]]: + """Load cached preprocessed tensors if they exist and are valid. + + Args: + data_dir: Path to data directory. + normalize: Whether normalization is enabled. + seq_length: Sequence length. + n_features: Number of features. + train_indices: Optional list of training indices. + excluded_stay_ids: Optional set of excluded stay IDs. + + Returns: + Dictionary with cached tensors and metadata if valid, None otherwise. + """ + cache_path = get_tensor_cache_path( + data_dir, normalize, seq_length, n_features, train_indices, excluded_stay_ids + ) + 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 cached.get("normalize") != 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( + data_dir: Path, + timeseries_tensor: torch.Tensor, + mask_tensor: torch.Tensor, + feature_means: torch.Tensor, + feature_stds: torch.Tensor, + normalize: bool, + seq_length: int, + n_features: int, + train_indices: Optional[List[int]], + excluded_stay_ids: Optional[Set[int]], +) -> None: + """Save preprocessed tensors to cache file. + + Args: + data_dir: Path to data directory. + timeseries_tensor: Preprocessed timeseries tensor. + mask_tensor: Observation mask tensor. + feature_means: Per-feature means. + feature_stds: Per-feature stds. + normalize: Whether normalization is enabled. + seq_length: Sequence length. + n_features: Number of features. + train_indices: Optional list of training indices. + excluded_stay_ids: Optional set of excluded stay IDs. + """ + cache_path = get_tensor_cache_path( + data_dir, normalize, seq_length, n_features, train_indices, excluded_stay_ids + ) + cache_dir = cache_path.parent + cache_dir.mkdir(exist_ok=True) + + cache_data = { + "timeseries_tensor": timeseries_tensor, + "mask_tensor": mask_tensor, + "feature_means": feature_means, + "feature_stds": feature_stds, + "n_features": n_features, + "seq_length": seq_length, + "normalize": 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( + 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..8151541 --- /dev/null +++ b/src/slices/data/tensor_preprocessing.py @@ -0,0 +1,278 @@ +"""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 Dict, List + +import numpy as np +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 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. + + 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...") + + 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 >= seq_length: + # Truncate + ts_arr = np.array(ts_data[:seq_length], dtype=np.float32) + mask_arr = np.array(mask_data[:seq_length], dtype=bool) + else: + # Pad with NaN / False + ts_arr = np.full((seq_length, n_features), np.nan, dtype=np.float32) + mask_arr = np.zeros((seq_length, 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 + + 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 (avoids physiologically + impossible values like 0 heart rate). + + 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, 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(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/training/checkpoint_loading.py b/src/slices/training/checkpoint_loading.py new file mode 100644 index 0000000..761b7fa --- /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 import ( + EncoderWithMissingToken, + ObservationTransformerEncoder, + SMARTEncoder, + TransformerEncoder, + build_encoder, +) + +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()) + + # 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. + """ + # Skip for encoders that handle missingness intrinsically + if isinstance(encoder, (ObservationTransformerEncoder, SMARTEncoder)): + logger.info( + "Skipping EncoderWithMissingToken wrapper for %s " + "(handles missingness intrinsically)", + type(encoder).__name__, + ) + return encoder + + if isinstance(encoder, TransformerEncoder) and getattr(encoder.config, "obs_aware", False): + logger.info( + "Skipping EncoderWithMissingToken wrapper for obs_aware " + "TransformerEncoder (forward() handles mask via obs_proj)" + ) + 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 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") + # 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 = 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, + ) + 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"} + # 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 = 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, + ) + 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/finetune_module.py b/src/slices/training/finetune_module.py index 9244c21..6aaea83 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,14 +19,13 @@ from omegaconf import DictConfig, OmegaConf from slices.eval import MetricConfig, build_metrics -from slices.models.encoders import ( - EncoderWithMissingToken, - ObservationTransformerEncoder, - SMARTEncoder, - TransformerEncoder, - build_encoder, -) +from slices.models.encoders import 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 @@ -101,12 +99,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), @@ -189,271 +191,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 - # - TransformerEncoder with obs_aware=True: obs_proj handles - # missingness via concatenated (values, obs_mask) input - if isinstance(self.encoder, (ObservationTransformerEncoder, SMARTEncoder)): - logger.info( - "Skipping EncoderWithMissingToken wrapper for %s " - "(handles missingness intrinsically)", - type(self.encoder).__name__, - ) - return - - if isinstance(self.encoder, TransformerEncoder) and getattr( - self.encoder.config, "obs_aware", False - ): - logger.info( - "Skipping EncoderWithMissingToken wrapper for obs_aware " - "TransformerEncoder (forward() handles mask via obs_proj)" - ) - 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: From 22a93b37a905fb989796cdd88e8a5e9064834fba Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 3 Mar 2026 08:19:10 -0500 Subject: [PATCH 005/121] fix: cast vis_proj dtype in JEPA predictor scatter to avoid bf16-mixed mismatch --- src/slices/models/pretraining/jepa.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/slices/models/pretraining/jepa.py b/src/slices/models/pretraining/jepa.py index 4ed70b7..37cf357 100644 --- a/src/slices/models/pretraining/jepa.py +++ b/src/slices/models/pretraining/jepa.py @@ -126,7 +126,7 @@ def forward( 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_(1, scatter_idx_expanded, vis_proj.to(full_tokens.dtype)) # Add time PE timestep_idx = token_info["timestep_idx"] From 64dce231a4808e462355b19266d31e94506cd6a8 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 3 Mar 2026 11:06:49 -0500 Subject: [PATCH 006/121] perf: reduce memory footprint and revert to fp32 precision Revert pretraining precision to fp32 (bf16-mixed was a temporary OOM workaround, no longer needed with accumulate_grad_batches=4). Free timeseries_df (~3-5 GB) after tensor loading, use mmap for tensor cache deserialization, and delete temporary DataModule attributes after setup. --- configs/pretrain.yaml | 2 +- src/slices/data/datamodule.py | 4 ++++ src/slices/data/dataset.py | 19 ++++++++++++++----- src/slices/data/tensor_cache.py | 2 +- 4 files changed, 20 insertions(+), 7 deletions(-) diff --git a/configs/pretrain.yaml b/configs/pretrain.yaml index e144d00..9970351 100644 --- a/configs/pretrain.yaml +++ b/configs/pretrain.yaml @@ -61,7 +61,7 @@ training: # Compute settings accelerator: auto # auto | cpu | gpu | mps devices: auto # auto | 1 | [0,1] | etc. - precision: bf16-mixed # 32 | 16-mixed | bf16-mixed + precision: 32 # 32 | 16-mixed | bf16-mixed # Regularization gradient_clip_val: 1.0 diff --git a/src/slices/data/datamodule.py b/src/slices/data/datamodule.py index 29084fe..0054b6d 100644 --- a/src/slices/data/datamodule.py +++ b/src/slices/data/datamodule.py @@ -258,6 +258,10 @@ def setup(self, stage: Optional[str] = None) -> None: self.test_ratio, ) + # 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 + logger.info( f"DataModule setup complete: " f"Train={len(self.train_indices):,}, Val={len(self.val_indices):,}, " diff --git a/src/slices/data/dataset.py b/src/slices/data/dataset.py index 05a0bea..3b805bd 100644 --- a/src/slices/data/dataset.py +++ b/src/slices/data/dataset.py @@ -275,6 +275,10 @@ def _load_data(self) -> None: self._excluded_stay_ids, ) + # Free raw DataFrame — only _timeseries_tensor and _mask_tensor needed from here. + # get_preprocessing_stages() reloads from parquet on demand. + del self.timeseries_df + # Pre-compute labels and static features self._precompute_labels_and_static() @@ -572,9 +576,8 @@ def get_label_statistics(self) -> Dict[str, Dict[str, Any]]: 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 @@ -600,8 +603,14 @@ 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"] diff --git a/src/slices/data/tensor_cache.py b/src/slices/data/tensor_cache.py index 45a73b7..123881f 100644 --- a/src/slices/data/tensor_cache.py +++ b/src/slices/data/tensor_cache.py @@ -221,7 +221,7 @@ def load_cached_tensors( try: logger.debug(f"Loading cached tensors from {cache_path.name}") - cached = torch.load(cache_path, weights_only=True) + cached = torch.load(cache_path, weights_only=True, mmap=True) # Validate cache metadata if cached.get("n_features") != n_features: From 5f071d0c271d6f49e646b4d850f70102e7cb0da0 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 3 Mar 2026 11:46:39 -0500 Subject: [PATCH 007/121] feat: add temporal NT-Xent mode to contrastive objective Replace instance-level (mean-pool) contrastive loss with temporal NT-Xent as the default mode. With two random masks at mask_ratio=0.5, ~25% of timesteps are visible in both views. These overlapping timesteps have two independently contextualized representations (different self-attention contexts), forming natural positive pairs. All other tokens across the batch are negatives. This preserves temporal granularity matching MAE/JEPA, ensuring fair paradigm comparison. The original instance mode is retained for ablation. --- configs/ssl/contrastive.yaml | 3 + src/slices/models/pretraining/contrastive.py | 193 +++++++++++++++-- tests/test_contrastive_objective.py | 212 ++++++++++++++++++- 3 files changed, 384 insertions(+), 24 deletions(-) diff --git a/configs/ssl/contrastive.yaml b/configs/ssl/contrastive.yaml index 06714ec..659c532 100644 --- a/configs/ssl/contrastive.yaml +++ b/configs/ssl/contrastive.yaml @@ -10,6 +10,9 @@ name: contrastive +# Mode: "temporal" (per-timestep overlap pairs) or "instance" (mean-pool ablation) +mode: temporal + # Masking parameters (same as MAE/JEPA for fair comparison) mask_ratio: 0.5 # Experiment plan default (ablate over 0.3/0.5/0.75) diff --git a/src/slices/models/pretraining/contrastive.py b/src/slices/models/pretraining/contrastive.py index a0645c4..f95aec1 100644 --- a/src/slices/models/pretraining/contrastive.py +++ b/src/slices/models/pretraining/contrastive.py @@ -1,19 +1,27 @@ """Contrastive (SimCLR-style) SSL objective for ICU time-series. Timestep-level tokenization variant: uses two different random timestep masks as -two augmented "views" of the same sample, then applies NT-Xent contrastive -loss on the pooled representations. +two augmented "views" of the same sample, then applies NT-Xent contrastive loss. + +Supports two modes: +- **temporal** (default): 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. +- **instance**: Original SimCLR-style — mean-pool each view to a single + sequence-level embedding, then NT-Xent on a (2B x 2B) matrix. Architecture: 1. TransformerEncoder.tokenize() -> one token per timestep (B, T, d_model) 2. Two independent random masks -> two different subsets of timestep 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 +4. temporal: scatter to (B,T,d) -> gather overlap tokens -> project -> NT-Xent + instance: mean-pool -> project -> NT-Xent -Key difference from MAE: discriminative (not reconstructive), global (not local). -Key difference from JEPA: global invariance (not local positional prediction). +Key difference from MAE: discriminative (not reconstructive). +Key difference from JEPA: invariance (not positional prediction). """ from dataclasses import dataclass @@ -33,6 +41,9 @@ class ContrastiveConfig(SSLConfig): name: str = "contrastive" + # Mode: "temporal" (per-timestep overlap pairs) or "instance" (mean-pool) + mode: str = "temporal" + # Masking mask_ratio: float = 0.5 @@ -43,6 +54,12 @@ class ContrastiveConfig(SSLConfig): # Temperature for NT-Xent temperature: float = 0.1 + 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}'" + ) + class ProjectionHead(nn.Module): """MLP projection head with L2 normalization. @@ -78,14 +95,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class ContrastiveObjective(BaseSSLObjective): - """Timestep-level contrastive (SimCLR-style) SSL for ICU time-series. + """Timestep-level contrastive SSL for ICU time-series. - Flow: - 1. encoder.tokenize(x, obs_mask) -> timestep tokens - 2. Two independent random timestep 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 encoder with tokenize()/encode() and pooling='none'. """ @@ -145,22 +158,29 @@ def forward( ssl_mask_1 = create_timestep_mask(B, T, self.config.mask_ratio, device) ssl_mask_2 = create_timestep_mask(B, T, self.config.mask_ratio, device) - # 3. View 1: extract -> encode -> mean pool + # 3. Encode both views vis_tokens_1, vis_padding_1 = extract_visible_timesteps(tokens, ssl_mask_1) 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_timesteps(tokens, ssl_mask_2) 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) + if self.config.mode == "temporal": + # 4a. Scatter encoded tokens back to full (B, T, d) grid + full_1 = self._scatter_to_full(encoded_1, ssl_mask_1, T) + full_2 = self._scatter_to_full(encoded_2, ssl_mask_2, T) - # 6. NT-Xent loss - loss, metrics = self._nt_xent_loss(z1, z2, ssl_mask_1, ssl_mask_2) + # 5a. Temporal NT-Xent on overlapping timesteps + loss, metrics = self._temporal_nt_xent_loss(full_1, full_2, ssl_mask_1, ssl_mask_2) + 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, ssl_mask_1, ssl_mask_2) return loss, metrics @@ -183,6 +203,133 @@ def _mean_pool( counts = padding_mask.sum(dim=1, keepdim=True).clamp(min=1).float() # (B, 1) return summed / counts + @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. + + Uses the same argsort-scatter pattern as the MAE decoder and JEPA + predictor to place 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. + """ + B, n_vis, d_enc = encoded.shape + device = encoded.device + + full = torch.zeros(B, n_timesteps, d_enc, device=device, dtype=encoded.dtype) + + vis_indices = ssl_mask.float().argsort(dim=1, descending=True, stable=True) + scatter_idx = vis_indices[:, :n_vis].unsqueeze(-1).expand(-1, -1, d_enc) + full.scatter_(1, scatter_idx, encoded) + + return full + + def _temporal_nt_xent_loss( + self, + full_1: torch.Tensor, + full_2: torch.Tensor, + ssl_mask_1: torch.Tensor, + ssl_mask_2: 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. + + 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(): + 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_visible_view1": ssl_mask_1.sum().item() / B, + "contrastive_n_visible_view2": ssl_mask_2.sum().item() / B, + "contrastive_n_masked_view1": (~ssl_mask_1).sum().item() / B, + "contrastive_n_masked_view2": (~ssl_mask_2).sum().item() / B, + "contrastive_n_overlap_tokens": 0, + "contrastive_n_overlap_per_sample": 0.0, + } + 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() + + 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_visible_view1": ssl_mask_1.sum().item() / B, + "contrastive_n_visible_view2": ssl_mask_2.sum().item() / B, + "contrastive_n_masked_view1": (~ssl_mask_1).sum().item() / B, + "contrastive_n_masked_view2": (~ssl_mask_2).sum().item() / B, + "contrastive_n_overlap_tokens": N, + "contrastive_n_overlap_per_sample": N / B, + } + + return loss, metrics + def _nt_xent_loss( self, z1: torch.Tensor, diff --git a/tests/test_contrastive_objective.py b/tests/test_contrastive_objective.py index 644873d..1303bd2 100644 --- a/tests/test_contrastive_objective.py +++ b/tests/test_contrastive_objective.py @@ -1,4 +1,4 @@ -"""Tests for timestep-level Contrastive (SimCLR-style) SSL objective.""" +"""Tests for timestep-level Contrastive SSL objective (temporal + instance modes).""" import pytest import torch @@ -65,6 +65,7 @@ def encoder(self): @pytest.fixture def contrastive_config(self): return ContrastiveConfig( + mode="instance", mask_ratio=0.75, proj_hidden_dim=64, proj_output_dim=16, @@ -87,6 +88,10 @@ def test_requires_obs_aware(self): 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 = TransformerConfig( d_input=10, @@ -128,6 +133,7 @@ def encoder(self): @pytest.fixture def contrastive_config(self): return ContrastiveConfig( + mode="instance", mask_ratio=0.75, proj_hidden_dim=64, proj_output_dim=16, @@ -206,6 +212,7 @@ def encoder(self): def test_accuracy_in_range(self, encoder): config = ContrastiveConfig( + mode="instance", mask_ratio=0.75, proj_hidden_dim=64, proj_output_dim=16, @@ -243,12 +250,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, @@ -276,6 +285,7 @@ def test_temperature_effect_on_loss(self, encoder): 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, @@ -314,6 +324,7 @@ def encoder(self): def test_sparse_data(self, encoder): config = ContrastiveConfig( + mode="instance", mask_ratio=0.5, proj_hidden_dim=64, proj_output_dim=16, @@ -354,6 +365,7 @@ def encoder(self): 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, @@ -399,6 +411,7 @@ def test_loss_decreases(self): encoder = TransformerEncoder(config) cont_config = ContrastiveConfig( + mode="instance", mask_ratio=0.75, proj_hidden_dim=64, proj_output_dim=16, @@ -428,6 +441,203 @@ 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, + ) + + 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_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, + ) + 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}" + + # ============================================================================= # Factory integration tests # ============================================================================= From dee1e700fa3760dfb81be09984cdfd6e0d3812e6 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 3 Mar 2026 17:36:47 -0500 Subject: [PATCH 008/121] fix: tensor cache never loads due to `or` on multi-element tensors Python's `or` calls bool() on the first operand, which raises RuntimeError for multi-element tensors. The except clause silently swallowed this, making every cache lookup return None. --- src/slices/data/tensor_cache.py | 63 ++++++++++++++++++++++----------- 1 file changed, 43 insertions(+), 20 deletions(-) diff --git a/src/slices/data/tensor_cache.py b/src/slices/data/tensor_cache.py index 123881f..82b3679 100644 --- a/src/slices/data/tensor_cache.py +++ b/src/slices/data/tensor_cache.py @@ -192,30 +192,20 @@ def get_tensor_cache_path( return cache_dir / f"tensors_{cache_key}.pt" -def load_cached_tensors( - data_dir: Path, - normalize: bool, - seq_length: int, - n_features: int, - train_indices: Optional[List[int]], - excluded_stay_ids: Optional[Set[int]], +def _try_load_cache( + cache_path: Path, n_features: int, seq_length: int, normalize: bool ) -> Optional[Dict[str, Any]]: - """Load cached preprocessed tensors if they exist and are valid. + """Try to load and validate a tensor cache file. Args: - data_dir: Path to data directory. - normalize: Whether normalization is enabled. - seq_length: Sequence length. - n_features: Number of features. - train_indices: Optional list of training indices. - excluded_stay_ids: Optional set of excluded stay IDs. + cache_path: Path to the cache file. + n_features: Expected number of features. + seq_length: Expected sequence length. + normalize: Expected normalize flag. Returns: - Dictionary with cached tensors and metadata if valid, None otherwise. + Dictionary with cached tensors if valid, None otherwise. """ - cache_path = get_tensor_cache_path( - data_dir, normalize, seq_length, n_features, train_indices, excluded_stay_ids - ) if not cache_path.exists(): return None @@ -235,8 +225,14 @@ def load_cached_tensors( 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") + # 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 @@ -258,6 +254,33 @@ def load_cached_tensors( return None +def load_cached_tensors( + data_dir: Path, + normalize: bool, + seq_length: int, + n_features: int, + train_indices: Optional[List[int]], + excluded_stay_ids: Optional[Set[int]], +) -> Optional[Dict[str, Any]]: + """Load cached preprocessed tensors if they exist and are valid. + + Args: + data_dir: Path to data directory. + normalize: Whether normalization is enabled. + seq_length: Sequence length. + n_features: Number of features. + train_indices: Optional list of training indices. + excluded_stay_ids: Optional set of excluded stay IDs. + + Returns: + Dictionary with cached tensors and metadata if valid, None otherwise. + """ + cache_path = get_tensor_cache_path( + data_dir, normalize, seq_length, n_features, train_indices, excluded_stay_ids + ) + return _try_load_cache(cache_path, n_features, seq_length, normalize) + + def save_cached_tensors( data_dir: Path, timeseries_tensor: torch.Tensor, From 6c60e29bbd2e48296d0b0924f023e783360a71d6 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 3 Mar 2026 17:37:03 -0500 Subject: [PATCH 009/121] fix: handle PyTorch 2.6 weights_only default in checkpoint loading PyTorch 2.6 changed torch.load to use weights_only=True by default, which rejects OmegaConf objects stored in Lightning checkpoints. --- scripts/debug/inspect_embeddings.py | 2 +- scripts/training/pretrain.py | 4 +++- src/slices/debug/embeddings.py | 2 +- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/scripts/debug/inspect_embeddings.py b/scripts/debug/inspect_embeddings.py index ffac1aa..fd97461 100644 --- a/scripts/debug/inspect_embeddings.py +++ b/scripts/debug/inspect_embeddings.py @@ -111,7 +111,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/training/pretrain.py b/scripts/training/pretrain.py index 7c6cecf..d5fb2e1 100644 --- a/scripts/training/pretrain.py +++ b/scripts/training/pretrain.py @@ -17,6 +17,7 @@ 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 @@ -198,7 +199,8 @@ def main(cfg: DictConfig) -> None: if best_ckpt: print(f"\n Loading best checkpoint: {best_ckpt}") print(f" Best val/loss: {checkpoint_callback.best_model_score:.4f}") - model = SSLPretrainModule.load_from_checkpoint(best_ckpt) + best_state = torch.load(best_ckpt, map_location="cpu", weights_only=False) + model.load_state_dict(best_state["state_dict"]) encoder_path = Path(cfg.output_dir) / "encoder.pt" encoder_path.parent.mkdir(parents=True, exist_ok=True) 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 From 387d9d9dd28055564d4d316b00b2fe8bf928bfc0 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 3 Mar 2026 17:45:12 -0500 Subject: [PATCH 010/121] fix: prevent JEPA representational collapse MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three targeted fixes for JEPA collapse observed in run ho5cwh3n: - Remove gradient accumulation (batch_size 64×4 → 256×1) to fix EMA updating 4× per optimizer step instead of 1× - Add target layer normalization in JEPA loss to prevent scale/shift collapse (standard in I-JEPA/BYOL) - Bottleneck predictor capacity (d_model 128→64, d_ff 512→256) to force information routing through the online encoder --- configs/pretrain.yaml | 4 ++-- configs/ssl/jepa.yaml | 9 ++++++--- src/slices/models/pretraining/jepa.py | 4 ++++ 3 files changed, 12 insertions(+), 5 deletions(-) diff --git a/configs/pretrain.yaml b/configs/pretrain.yaml index 9970351..075de9c 100644 --- a/configs/pretrain.yaml +++ b/configs/pretrain.yaml @@ -56,7 +56,7 @@ encoder: # Training configuration (same budget across thesis paradigms for fair comparison) training: max_epochs: 500 - batch_size: 64 + batch_size: 256 # Compute settings accelerator: auto # auto | cpu | gpu | mps @@ -65,7 +65,7 @@ training: # Regularization gradient_clip_val: 1.0 - accumulate_grad_batches: 4 # effective batch_size = 64 * 4 = 256 + accumulate_grad_batches: 1 # single batch, no accumulation # Early stopping early_stopping_patience: 10 diff --git a/configs/ssl/jepa.yaml b/configs/ssl/jepa.yaml index 464a883..c07566c 100644 --- a/configs/ssl/jepa.yaml +++ b/configs/ssl/jepa.yaml @@ -13,11 +13,14 @@ name: jepa # Masking parameters (same as MAE for fair comparison) mask_ratio: 0.5 # Experiment plan default (ablate over 0.3/0.5/0.75) -# 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: 64 predictor_n_layers: 2 predictor_n_heads: 4 -predictor_d_ff: 512 +predictor_d_ff: 256 predictor_dropout: 0.1 # Momentum encoder parameters diff --git a/src/slices/models/pretraining/jepa.py b/src/slices/models/pretraining/jepa.py index 37cf357..bde9886 100644 --- a/src/slices/models/pretraining/jepa.py +++ b/src/slices/models/pretraining/jepa.py @@ -177,6 +177,7 @@ def __init__(self, encoder: nn.Module, config: JEPAConfig) -> None: ) d_encoder = encoder.get_output_dim() + self.d_encoder = d_encoder max_seq_length = encoder.config.max_seq_length self.missing_token = None @@ -265,6 +266,9 @@ def _compute_loss( # Loss only on masked timesteps loss_mask = ~ssl_mask # (B, T) + # 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") element_loss = element_loss.mean(dim=-1) # (B, T) From 5fee7f16fecdb45c00b21f66d2ad7a08c304961e Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 3 Mar 2026 17:51:58 -0500 Subject: [PATCH 011/121] fix: match checkpoint filename placeholders to logged metric names The ModelCheckpoint filename templates used underscores (val_loss, val_auprc) but Lightning logs metrics with slashes (val/loss, val/auprc). Unmatched placeholders render as 0.0000 in filenames. --- src/slices/training/utils.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index b5576aa..e7175a7 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -160,7 +160,7 @@ 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}-{val/loss:.4f}", monitor="val/loss", mode="min", save_top_k=3, @@ -200,11 +200,9 @@ def setup_finetune_callbacks(cfg: DictConfig, checkpoint_prefix: str = "finetune monitor = cfg.training.get("early_stopping_monitor", default_monitor) mode = cfg.training.get("early_stopping_mode", default_mode) - metric_filename = monitor.replace("/", "_") - 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}:.4f}}", monitor=monitor, mode=mode, save_top_k=3, From 7978678fcb4e2374f68a63a6b8ed27d05b106133 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 3 Mar 2026 18:35:31 -0500 Subject: [PATCH 012/121] fix: address JEPA collapse with higher momentum, block masking, and monitoring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Increase EMA momentum_base from 0.996 to 0.999 to slow target encoder updates (~12%/epoch vs ~50%/epoch at 125 steps/epoch) - Add vectorized block masking (create_block_timestep_mask) that masks contiguous segments instead of random timesteps, forcing long-range prediction over neighbor interpolation - Add collapse monitoring metrics: jepa_target_repr_std (feature std across positions, collapse → 0) and jepa_target_repr_mean_cos (mean pairwise cosine similarity, collapse → 1.0) --- configs/ssl/jepa.yaml | 6 +- src/slices/models/pretraining/jepa.py | 46 +++++++++++--- src/slices/models/pretraining/masking.py | 79 ++++++++++++++++++++++++ 3 files changed, 122 insertions(+), 9 deletions(-) diff --git a/configs/ssl/jepa.yaml b/configs/ssl/jepa.yaml index c07566c..63e5e34 100644 --- a/configs/ssl/jepa.yaml +++ b/configs/ssl/jepa.yaml @@ -10,8 +10,10 @@ 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) +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 — bottlenecked relative to encoder (50% d_model). # Unlike MAE decoder (which reconstructs raw features, d_model→D), JEPA predictor @@ -26,7 +28,7 @@ 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/src/slices/models/pretraining/jepa.py b/src/slices/models/pretraining/jepa.py index bde9886..c8cd383 100644 --- a/src/slices/models/pretraining/jepa.py +++ b/src/slices/models/pretraining/jepa.py @@ -26,7 +26,7 @@ from slices.models.common import build_sinusoidal_pe from .base import BaseSSLObjective, SSLConfig -from .masking import create_timestep_mask, extract_visible_timesteps +from .masking import create_block_timestep_mask, create_timestep_mask, extract_visible_timesteps @dataclass @@ -37,6 +37,8 @@ class JEPAConfig(SSLConfig): # 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 @@ -46,7 +48,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 @@ -218,7 +220,12 @@ def forward( tokens, padding_mask, token_info = self.encoder.tokenize(x, obs_mask) # 2. Create SSL mask on timesteps - ssl_mask = create_timestep_mask(B, T, self.config.mask_ratio, device) + 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 + ) + else: + ssl_mask = create_timestep_mask(B, T, self.config.mask_ratio, device) # 3. Extract visible tokens visible_tokens, vis_padding = extract_visible_timesteps(tokens, ssl_mask) @@ -266,6 +273,33 @@ def _compute_loss( # Loss only on masked timesteps loss_mask = ~ssl_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.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]) @@ -281,10 +315,6 @@ def _compute_loss( loss = (element_loss * loss_mask.float()).sum() / loss_mask.float().sum().clamp(min=1) with torch.no_grad(): - B, T = ssl_mask.shape - n_masked = loss_mask.sum().item() - n_visible = ssl_mask.sum().item() - metrics = { "jepa_loss": loss.detach(), "ssl_loss": loss.detach(), @@ -293,6 +323,8 @@ def _compute_loss( "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/masking.py b/src/slices/models/pretraining/masking.py index 8cd9764..00e9b4a 100644 --- a/src/slices/models/pretraining/masking.py +++ b/src/slices/models/pretraining/masking.py @@ -112,6 +112,85 @@ def create_timestep_mask( return ssl_mask +def create_block_timestep_mask( + batch_size: int, + n_timesteps: int, + mask_ratio: float, + device: torch.device, + n_blocks: int = 3, +) -> torch.Tensor: + """Create contiguous block mask at timestep level. + + Masks n_blocks contiguous segments that together cover approximately + mask_ratio of the sequence. Forces the model to predict from distant + context rather than interpolate from adjacent visible timesteps. + + Strategy: divides the sequence into n_blocks equal zones, then places one + randomly-sized masked block within each zone. Block sizes are drawn from a + Dirichlet-like split of the total masked budget. Fully vectorized — no + Python loops over batch elements. + + 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). + + Returns: + ssl_mask: (B, T) bool mask, True = visible, False = masked. + """ + n_masked_total = max(int(n_timesteps * mask_ratio), 1) + n_masked_total = min(n_masked_total, n_timesteps - 1) + + # Split total masked budget into n_blocks random lengths per sample. + # Use Dirichlet-like splitting: draw n_blocks uniform values, normalize, + # then scale to sum to n_masked_total. Add 1 to each to ensure min length 1. + raw = torch.rand(batch_size, n_blocks) # CPU for speed + raw = raw / raw.sum(dim=1, keepdim=True) # normalize to sum=1 + # Reserve 1 per block, distribute the rest proportionally + extra = n_masked_total - n_blocks + if extra > 0: + block_lengths = 1 + (raw * extra).int() + # Fix rounding: adjust last block to hit exact total + block_lengths[:, -1] = n_masked_total - block_lengths[:, :-1].sum(dim=1) + else: + # Edge case: fewer masked timesteps than blocks — give 1 to first blocks + block_lengths = torch.zeros(batch_size, n_blocks, dtype=torch.int) + block_lengths[:, :n_masked_total] = 1 + + # Clamp to valid range + block_lengths = block_lengths.clamp(min=0, max=n_timesteps) + + # Divide sequence into n_blocks equal zones. Each block is placed randomly + # within its zone, avoiding cross-zone overlap by construction. + zone_size = n_timesteps // n_blocks + zone_starts = torch.arange(n_blocks) * zone_size # (n_blocks,) + + # Random offset within each zone (vectorized over batch) + # Max offset = zone_size - block_length (so block fits within zone) + max_offsets = zone_size - block_lengths # (B, n_blocks) + max_offsets = max_offsets.clamp(min=0) + offsets = (torch.rand(batch_size, n_blocks) * (max_offsets.float() + 1)).int() + offsets = offsets.clamp(max=max_offsets) + + # Compute absolute start positions: zone_start + offset + starts = zone_starts.unsqueeze(0) + offsets # (B, n_blocks) + + # Build mask: create time indices and compare with block ranges + # (B, n_blocks, T) — True where timestep t falls in block k + t = torch.arange(n_timesteps).unsqueeze(0).unsqueeze(0) # (1, 1, T) + starts_3d = starts.unsqueeze(2) # (B, n_blocks, 1) + ends_3d = (starts + block_lengths).unsqueeze(2) # (B, n_blocks, 1) + in_block = (t >= starts_3d) & (t < ends_3d) # (B, n_blocks, T) + + # Union across blocks -> masked positions + masked = in_block.any(dim=1) # (B, T) + ssl_mask = (~masked).to(device=device) + + return ssl_mask + + def extract_visible_timesteps( tokens: torch.Tensor, ssl_mask: torch.Tensor, From e06a523981cf8e0401f0813df3aed079c5bf3641 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 3 Mar 2026 18:54:54 -0500 Subject: [PATCH 013/121] fix: shorten training schedule and prevent early stopping during warmup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reduced max_epochs 500→100 and warmup 50→10 (same 10% ratio). Added min_epochs=10 so early stopping cannot trigger during LR warmup, which was killing JEPA runs before they reached peak learning rate. --- configs/pretrain.yaml | 5 +++-- scripts/training/pretrain.py | 1 + 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/configs/pretrain.yaml b/configs/pretrain.yaml index 075de9c..e1dcbc5 100644 --- a/configs/pretrain.yaml +++ b/configs/pretrain.yaml @@ -55,7 +55,8 @@ encoder: # Training configuration (same budget across thesis paradigms for fair comparison) training: - max_epochs: 500 + max_epochs: 100 + min_epochs: 10 # >= warmup_epochs; prevents early stopping during LR warmup batch_size: 256 # Compute settings @@ -82,7 +83,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 diff --git a/scripts/training/pretrain.py b/scripts/training/pretrain.py index d5fb2e1..5e92dc4 100644 --- a/scripts/training/pretrain.py +++ b/scripts/training/pretrain.py @@ -154,6 +154,7 @@ def main(cfg: DictConfig) -> None: trainer = pl.Trainer( max_epochs=cfg.training.max_epochs, + min_epochs=cfg.training.get("min_epochs", None), accelerator=cfg.training.get("accelerator", "auto"), devices=cfg.training.get("devices", "auto"), precision=cfg.training.get("precision", 32), From da1520f4257937e8ab2e0c1456a045f029ffae25 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 3 Mar 2026 22:04:56 -0500 Subject: [PATCH 014/121] feat: add sprint 1 launch script and VM auto-shutdown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - launch_sprint1.sh: runs all 4 paradigms (MAE, JEPA, Contrastive, Supervised) in parallel on MIMIC-IV with automatic pretrain→finetune chaining - auto_shutdown.sh: powers off GCP VM after 59min of no training activity (cron-based) - Update experiment plan with LR sensitivity ablation (Sprint 1b) --- scripts/auto_shutdown.sh | 36 +++++++++ scripts/launch_sprint1.sh | 152 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 188 insertions(+) create mode 100755 scripts/auto_shutdown.sh create mode 100755 scripts/launch_sprint1.sh diff --git a/scripts/auto_shutdown.sh b/scripts/auto_shutdown.sh new file mode 100755 index 0000000..dffcb6a --- /dev/null +++ b/scripts/auto_shutdown.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +# ============================================================================= +# Auto-shutdown: powers off the VM if no training has run for 30 minutes. +# 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 training process is running +if pgrep -f "scripts/training/(pretrain|finetune|supervised)\.py" > /dev/null 2>&1; then + # Training active — 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/launch_sprint1.sh b/scripts/launch_sprint1.sh new file mode 100755 index 0000000..15490dc --- /dev/null +++ b/scripts/launch_sprint1.sh @@ -0,0 +1,152 @@ +#!/usr/bin/env bash +# ============================================================================= +# Sprint 1: Sanity Check +# MIMIC-IV | All 4 paradigms | mortality_24h | seed=42 +# +# Runs 3 SSL pretrain jobs in parallel + 1 supervised baseline. +# After each pretrain finishes, automatically launches finetuning. +# ============================================================================= +set -euo pipefail + +SPRINT=1 +DATASET=miiv +SEED=42 +TASK=mortality_24h +LOG_DIR="logs/sprint1" + +mkdir -p "$LOG_DIR" + +echo "============================================================" +echo "SLICES Sprint 1: Sanity Check" +echo " Dataset: $DATASET" +echo " Task: $TASK" +echo " Seed: $SEED" +echo " Paradigms: MAE, JEPA, Contrastive, Supervised" +echo "============================================================" + +# Use predictable output dirs (override Hydra's timestamped dirs) +MAE_DIR="outputs/sprint1/pretrain_mae" +JEPA_DIR="outputs/sprint1/pretrain_jepa" +CONTRASTIVE_DIR="outputs/sprint1/pretrain_contrastive" +SUPERVISED_DIR="outputs/sprint1/supervised" +FT_MAE_DIR="outputs/sprint1/finetune_mae_${TASK}" +FT_JEPA_DIR="outputs/sprint1/finetune_jepa_${TASK}" +FT_CONTRASTIVE_DIR="outputs/sprint1/finetune_contrastive_${TASK}" + +# ----------------------------------------------------------------------------- +# Helper: run pretrain then finetune +# ----------------------------------------------------------------------------- +pretrain_and_finetune() { + local ssl_name="$1" + local pretrain_dir="$2" + local finetune_dir="$3" + local log_prefix="$LOG_DIR/${ssl_name}" + + echo "[$(date '+%H:%M:%S')] Starting $ssl_name pretraining..." + + uv run python scripts/training/pretrain.py \ + dataset=$DATASET \ + ssl=$ssl_name \ + seed=$SEED \ + sprint=$SPRINT \ + hydra.run.dir=$pretrain_dir \ + 2>&1 | tee "${log_prefix}_pretrain.log" + + echo "[$(date '+%H:%M:%S')] $ssl_name pretraining complete." + + # Check encoder was saved + if [ ! -f "${pretrain_dir}/encoder.pt" ]; then + echo "ERROR: ${pretrain_dir}/encoder.pt not found!" + return 1 + fi + + echo "[$(date '+%H:%M:%S')] Starting $ssl_name finetuning on $TASK..." + + uv run python scripts/training/finetune.py \ + dataset=$DATASET \ + checkpoint=${pretrain_dir}/encoder.pt \ + tasks=$TASK \ + seed=$SEED \ + sprint=$SPRINT \ + hydra.run.dir=$finetune_dir \ + 2>&1 | tee "${log_prefix}_finetune.log" + + echo "[$(date '+%H:%M:%S')] $ssl_name finetuning complete." +} + +# ----------------------------------------------------------------------------- +# Launch all 4 runs in parallel +# ----------------------------------------------------------------------------- + +# MAE: pretrain -> finetune +pretrain_and_finetune mae "$MAE_DIR" "$FT_MAE_DIR" & +PID_MAE=$! + +# JEPA: pretrain -> finetune +pretrain_and_finetune jepa "$JEPA_DIR" "$FT_JEPA_DIR" & +PID_JEPA=$! + +# Contrastive: pretrain -> finetune +pretrain_and_finetune contrastive "$CONTRASTIVE_DIR" "$FT_CONTRASTIVE_DIR" & +PID_CONTRASTIVE=$! + +# Supervised: single run (no pretrain step) +( + echo "[$(date '+%H:%M:%S')] Starting supervised baseline..." + + uv run python scripts/training/supervised.py \ + dataset=$DATASET \ + tasks=$TASK \ + seed=$SEED \ + sprint=$SPRINT \ + hydra.run.dir=$SUPERVISED_DIR \ + 2>&1 | tee "${LOG_DIR}/supervised.log" + + echo "[$(date '+%H:%M:%S')] Supervised baseline complete." +) & +PID_SUPERVISED=$! + +echo "" +echo "All 4 jobs launched in parallel:" +echo " MAE: PID=$PID_MAE" +echo " JEPA: PID=$PID_JEPA" +echo " Contrastive: PID=$PID_CONTRASTIVE" +echo " Supervised: PID=$PID_SUPERVISED" +echo "" +echo "Logs: $LOG_DIR/" +echo " mae_pretrain.log / mae_finetune.log" +echo " jepa_pretrain.log / jepa_finetune.log" +echo " contrastive_pretrain.log / contrastive_finetune.log" +echo " supervised.log" +echo "" +echo "Waiting for all jobs to finish..." + +# Wait and report results +FAILED=0 + +wait $PID_MAE || { echo "FAILED: MAE pipeline"; FAILED=$((FAILED + 1)); } +wait $PID_JEPA || { echo "FAILED: JEPA pipeline"; FAILED=$((FAILED + 1)); } +wait $PID_CONTRASTIVE || { echo "FAILED: Contrastive pipeline"; FAILED=$((FAILED + 1)); } +wait $PID_SUPERVISED || { echo "FAILED: Supervised baseline"; FAILED=$((FAILED + 1)); } + +echo "" +echo "============================================================" +if [ $FAILED -eq 0 ]; then + echo "Sprint 1 COMPLETE — all 4 paradigms finished successfully." +else + echo "Sprint 1 FINISHED WITH $FAILED FAILURE(S). Check logs." +fi +echo "============================================================" +echo "" +echo "Output directories:" +echo " MAE pretrain: $MAE_DIR" +echo " MAE finetune: $FT_MAE_DIR" +echo " JEPA pretrain: $JEPA_DIR" +echo " JEPA finetune: $FT_JEPA_DIR" +echo " Contrastive pretrain: $CONTRASTIVE_DIR" +echo " Contrastive finetune: $FT_CONTRASTIVE_DIR" +echo " Supervised: $SUPERVISED_DIR" +echo "" +echo "Check W&B for training curves: https://wandb.ai/slices" + +exit $FAILED From f630e2868d957a78a1c1ec6feb1f32829b31fa85 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 3 Mar 2026 23:09:56 -0500 Subject: [PATCH 015/121] fix: use fixed-schedule training and save both last/best encoders MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Disable early stopping for pretraining — JEPA's val loss rises as representations improve, making it unreliable as a stopping criterion. Follow I-JEPA/DINO/BYOL convention of fixed-schedule training. Save both last-epoch encoder (primary) and best-val-loss encoder (robustness check). Update experiment plan to reflect shared obs-aware timestep-level tokenization across all SSL paradigms. --- configs/pretrain.yaml | 7 ++++--- scripts/training/pretrain.py | 34 +++++++++++++++++++--------------- 2 files changed, 23 insertions(+), 18 deletions(-) diff --git a/configs/pretrain.yaml b/configs/pretrain.yaml index e1dcbc5..a2c3993 100644 --- a/configs/pretrain.yaml +++ b/configs/pretrain.yaml @@ -56,7 +56,6 @@ encoder: # Training configuration (same budget across thesis paradigms for fair comparison) training: max_epochs: 100 - min_epochs: 10 # >= warmup_epochs; prevents early stopping during LR warmup batch_size: 256 # Compute settings @@ -68,8 +67,10 @@ training: gradient_clip_val: 1.0 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 diff --git a/scripts/training/pretrain.py b/scripts/training/pretrain.py index 5e92dc4..eeecf2c 100644 --- a/scripts/training/pretrain.py +++ b/scripts/training/pretrain.py @@ -154,7 +154,6 @@ def main(cfg: DictConfig) -> None: trainer = pl.Trainer( max_epochs=cfg.training.max_epochs, - min_epochs=cfg.training.get("min_epochs", None), accelerator=cfg.training.get("accelerator", "auto"), devices=cfg.training.get("devices", "auto"), precision=cfg.training.get("precision", 32), @@ -189,26 +188,33 @@ def main(cfg: DictConfig) -> None: trainer.fit(model, datamodule=datamodule) # ========================================================================= - # 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}") best_state = torch.load(best_ckpt, map_location="cpu", weights_only=False) model.load_state_dict(best_state["state_dict"]) - - 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") + 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 @@ -217,13 +223,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( From e91a6bbcd900799dac925c5751f9ad04d1d32da9 Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 4 Mar 2026 08:33:10 -0500 Subject: [PATCH 016/121] refactor: unify encoder architecture across all paradigms (174K params) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All paradigms (MAE, JEPA, Contrastive, Supervised) now share the same encoder: d=64, L=2, H=4, obs_aware=True (~174K params). This eliminates model capacity as a confounding variable in the SSL comparison. - Enable obs_aware=True as default in base transformer config - Remove encoder size overrides from pretrain.yaml and supervised.yaml - Scale down SSL heads to match smaller encoder: MAE decoder: d=64, JEPA predictor: d=32, Contrastive proj: 256→64 - Update experiment plan to reflect shared architecture --- configs/model/transformer.yaml | 15 ++++----------- configs/pretrain.yaml | 6 +----- configs/ssl/contrastive.yaml | 4 ++-- configs/ssl/jepa.yaml | 6 +++--- configs/ssl/mae.yaml | 6 +++--- configs/supervised.yaml | 1 + 6 files changed, 14 insertions(+), 24 deletions(-) diff --git a/configs/model/transformer.yaml b/configs/model/transformer.yaml index c937295..b20d29b 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 @@ -47,5 +40,5 @@ use_positional_encoding: true # Add sinusoidal positional encoding pooling: mean # Observation-aware mode: concat (values, obs_mask) -> MLP -> d_model -# Enable for SSL pretraining; disabled by default for supervised baselines -obs_aware: false +# Enabled for all paradigms (SSL and supervised) for fair comparison +obs_aware: true diff --git a/configs/pretrain.yaml b/configs/pretrain.yaml index a2c3993..515d4b2 100644 --- a/configs/pretrain.yaml +++ b/configs/pretrain.yaml @@ -45,13 +45,9 @@ data: window_stride: 24 # Stride for overlapping windows (50% overlap) # 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) - obs_aware: true # Enable obs-aware MLP projection for SSL - d_model: 128 - n_layers: 4 - n_heads: 8 - d_ff: 512 # Training configuration (same budget across thesis paradigms for fair comparison) training: diff --git a/configs/ssl/contrastive.yaml b/configs/ssl/contrastive.yaml index 659c532..6e360ae 100644 --- a/configs/ssl/contrastive.yaml +++ b/configs/ssl/contrastive.yaml @@ -17,8 +17,8 @@ mode: temporal 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 temperature: 0.1 diff --git a/configs/ssl/jepa.yaml b/configs/ssl/jepa.yaml index 63e5e34..d39dbb2 100644 --- a/configs/ssl/jepa.yaml +++ b/configs/ssl/jepa.yaml @@ -19,10 +19,10 @@ mask_n_blocks: 3 # Number of contiguous blocks (only for block strategy) # 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: 64 +predictor_d_model: 32 # 50% of encoder d_model=64 predictor_n_layers: 2 -predictor_n_heads: 4 -predictor_d_ff: 256 +predictor_n_heads: 4 # 32/4 = 8 per head +predictor_d_ff: 128 # 4 * predictor_d_model predictor_dropout: 0.1 # Momentum encoder parameters diff --git a/configs/ssl/mae.yaml b/configs/ssl/mae.yaml index a30afe9..45ac05d 100644 --- a/configs/ssl/mae.yaml +++ b/configs/ssl/mae.yaml @@ -11,9 +11,9 @@ 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/supervised.yaml b/configs/supervised.yaml index aa59d20..9884aec 100644 --- a/configs/supervised.yaml +++ b/configs/supervised.yaml @@ -47,6 +47,7 @@ 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. From 497916b41d2b39d3878538b3120988f383348aad Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 4 Mar 2026 11:44:51 -0500 Subject: [PATCH 017/121] feat: switch contrastive to instance mode with collapse monitoring Switch contrastive objective from temporal to instance-level (SimCLR-style) with temperature 0.07. Add alignment, uniformity, and effective rank metrics for collapse detection. Update experiment plan with dual evaluation protocols (linear probe + full finetune) and expanded Sprint 1 validation scope. --- configs/ssl/contrastive.yaml | 8 +++---- src/slices/models/pretraining/contrastive.py | 25 ++++++++++++++++++-- 2 files changed, 27 insertions(+), 6 deletions(-) diff --git a/configs/ssl/contrastive.yaml b/configs/ssl/contrastive.yaml index 6e360ae..5ea2cf7 100644 --- a/configs/ssl/contrastive.yaml +++ b/configs/ssl/contrastive.yaml @@ -10,8 +10,8 @@ name: contrastive -# Mode: "temporal" (per-timestep overlap pairs) or "instance" (mean-pool ablation) -mode: temporal +# Mode: "temporal" (per-timestep overlap pairs) or "instance" (mean-pool, SimCLR-style) +mode: instance # Masking parameters (same as MAE/JEPA for fair comparison) mask_ratio: 0.5 # Experiment plan default (ablate over 0.3/0.5/0.75) @@ -20,5 +20,5 @@ mask_ratio: 0.5 # Experiment plan default (ablate over 0.3/0.5/0.75) 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 -temperature: 0.1 +# NT-Xent temperature (0.07 suits instance-level with ~511 negatives at B=256) +temperature: 0.07 diff --git a/src/slices/models/pretraining/contrastive.py b/src/slices/models/pretraining/contrastive.py index f95aec1..d212e06 100644 --- a/src/slices/models/pretraining/contrastive.py +++ b/src/slices/models/pretraining/contrastive.py @@ -42,7 +42,7 @@ class ContrastiveConfig(SSLConfig): name: str = "contrastive" # Mode: "temporal" (per-timestep overlap pairs) or "instance" (mean-pool) - mode: str = "temporal" + mode: str = "instance" # Masking mask_ratio: float = 0.5 @@ -52,7 +52,7 @@ class ContrastiveConfig(SSLConfig): proj_output_dim: int = 128 # Temperature for NT-Xent - temperature: float = 0.1 + temperature: float = 0.07 def __post_init__(self) -> None: if self.mode not in ("temporal", "instance"): @@ -377,6 +377,24 @@ def _nt_xent_loss( # Positive pair similarities (before temperature scaling) pos_sim = F.cosine_similarity(z1, z2, dim=-1).mean() + # --- 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 T = ssl_mask_1.shape[1] n_vis_1 = ssl_mask_1.sum().item() @@ -389,6 +407,9 @@ def _nt_xent_loss( "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_timesteps": T, "contrastive_n_visible_view1": n_vis_1 / B, From df8b612ceb25e7c7f93d7950600718192b8538df Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 4 Mar 2026 12:47:10 -0500 Subject: [PATCH 018/121] fix: add dropout to obs_proj and increase supervised regularization The obs_proj MLP was the only component in the encoder without dropout, causing severe overfitting when obs_aware=True (train auroc 0.993, val auroc 0.807). Added nn.Dropout to obs_proj for all paradigms. Also increased supervised baseline regularization: encoder dropout 0.2, label smoothing 0.1, weight decay 0.05. Added label_smoothing support to CrossEntropyLoss in FineTuneModule. --- configs/finetune.yaml | 3 +++ configs/supervised.yaml | 10 +++++++--- src/slices/models/encoders/transformer.py | 1 + src/slices/training/config_schemas.py | 1 + src/slices/training/finetune_module.py | 3 ++- 5 files changed, 14 insertions(+), 4 deletions(-) diff --git a/configs/finetune.yaml b/configs/finetune.yaml index 00ec561..dfb507f 100644 --- a/configs/finetune.yaml +++ b/configs/finetune.yaml @@ -101,6 +101,9 @@ training: # Override here only if you need a non-default metric. early_stopping_patience: 10 + # Label smoothing (default off for finetuning — frozen encoder doesn't overfit) + label_smoothing: 0.0 + # Debug: overfit on N batches (0 = disabled) overfit_batches: 0 diff --git a/configs/supervised.yaml b/configs/supervised.yaml index 9884aec..e73cebd 100644 --- a/configs/supervised.yaml +++ b/configs/supervised.yaml @@ -46,8 +46,9 @@ ckpt_path: null 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 +# Encoder overrides — increase dropout for training from scratch +encoder: + dropout: 0.2 # Task configuration inherited from configs/tasks/*.yaml # Override with: task=mortality_hospital, task=los_remaining, etc. @@ -79,6 +80,9 @@ training: # [w0, w1, ...]: explicit per-class weights class_weight: null + # Label smoothing — helps calibration with highly imbalanced tasks + label_smoothing: 0.1 + # Regularization gradient_clip_val: 1.0 accumulate_grad_batches: 1 @@ -97,7 +101,7 @@ training: optimizer: name: adamw lr: 1.0e-3 # Higher LR for training from scratch - weight_decay: 0.01 + weight_decay: 0.05 # Learning rate scheduler scheduler: diff --git a/src/slices/models/encoders/transformer.py b/src/slices/models/encoders/transformer.py index 7cb645a..0a26252 100644 --- a/src/slices/models/encoders/transformer.py +++ b/src/slices/models/encoders/transformer.py @@ -189,6 +189,7 @@ def __init__(self, config: TransformerConfig) -> None: 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( diff --git a/src/slices/training/config_schemas.py b/src/slices/training/config_schemas.py index 0a88f59..6a216a2 100644 --- a/src/slices/training/config_schemas.py +++ b/src/slices/training/config_schemas.py @@ -86,6 +86,7 @@ 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 diff --git a/src/slices/training/finetune_module.py b/src/slices/training/finetune_module.py index 6aaea83..0e61a03 100644 --- a/src/slices/training/finetune_module.py +++ b/src/slices/training/finetune_module.py @@ -219,7 +219,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": From 8099dd7caf5b5bab2de40e9159b3f9b1144dcadb Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 4 Mar 2026 18:45:53 -0500 Subject: [PATCH 019/121] feat: add parallel experiment runner for all 9 sprints Dependency-aware scheduler that generates 915 runs across sprints 1-8, deduplicates shared pretrains, and executes jobs in parallel with crash recovery and atomic state persistence. --- scripts/run_experiments.py | 852 +++++++++++++++++++++++++++++++++++++ 1 file changed, 852 insertions(+) create mode 100644 scripts/run_experiments.py diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py new file mode 100644 index 0000000..6b5a622 --- /dev/null +++ b/scripts/run_experiments.py @@ -0,0 +1,852 @@ +#!/usr/bin/env python3 +""" +Parallel experiment runner for SLICES. + +Generates all experiment configurations across sprints, resolves dependencies, +and executes them in parallel with crash recovery and state persistence. + +Usage: + uv run python scripts/run_experiments.py run --sprint 1 --parallel 4 + uv run python scripts/run_experiments.py run --sprint 1 2 3 --parallel 6 --dry-run + uv run python scripts/run_experiments.py status + uv run python scripts/run_experiments.py status --sprint 1 + uv run python scripts/run_experiments.py retry --failed --parallel 4 +""" +from __future__ import annotations + +import argparse +import json +import os +import signal +import subprocess +import sys +import tempfile +import time +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- +SSL_PARADIGMS = ["mae", "jepa", "contrastive"] +DATASETS = ["miiv", "eicu", "combined"] +TASKS = ["mortality_24h", "mortality_hospital", "aki_kdigo", "los_remaining"] +SEEDS = [42, 123, 456] +LABEL_FRACTIONS_FULL = [0.01, 0.05, 0.1, 0.25, 0.5] +LABEL_FRACTIONS_TREND = [0.1] +LR_ABLATION = [5e-4, 2e-3] # 1e-3 reused from Phase 1 +MASK_RATIO_ABLATION = [0.3, 0.75] # 0.5 reused from Phase 1 +TRANSFER_PAIRS = [("miiv", "eicu"), ("eicu", "miiv")] + +STATE_FILE = Path("outputs/experiment_state.json") +LOG_DIR = Path("logs/runner") + +# Protocol defaults +PROTO_A = {"freeze_encoder": True, "max_epochs": 50, "patience": 10, "lr": 1e-4} +PROTO_B = {"freeze_encoder": False, "max_epochs": 100, "patience": 20, "lr": 1e-3} + + +# --------------------------------------------------------------------------- +# Data Model +# --------------------------------------------------------------------------- +@dataclass +class Run: + id: str + sprint: str + run_type: str # "pretrain" | "finetune" | "supervised" + paradigm: str # "mae" | "jepa" | "contrastive" | "supervised" + dataset: str + seed: int + output_dir: str + 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) + # For transfer learning: dataset the encoder was pretrained on + source_dataset: str | None = None + + 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() + elif self.run_type == "finetune": + return self._finetune_cmd(runs_by_id) + elif self.run_type == "supervised": + return self._supervised_cmd() + else: + raise ValueError(f"Unknown run_type: {self.run_type}") + + 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"sprint={self.sprint}", + f"hydra.run.dir={self.output_dir}", + ] + for k, v in self.extra_overrides.items(): + cmd.append(f"{k}={v}") + return cmd + + def _finetune_cmd(self, runs_by_id: dict[str, Run]) -> list[str]: + # Find the pretrain dependency to get encoder path + pretrain_id = [d for d in self.depends_on if d in runs_by_id] + if not pretrain_id: + raise ValueError(f"Finetune run {self.id} has no pretrain dependency") + pretrain_dir = runs_by_id[pretrain_id[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"sprint={self.sprint}", + f"hydra.run.dir={self.output_dir}", + ] + if self.freeze_encoder is True: + cmd += [ + "training.freeze_encoder=true", + f"training.max_epochs={PROTO_A['max_epochs']}", + f"training.early_stopping_patience={PROTO_A['patience']}", + f"optimizer.lr={PROTO_A['lr']}", + ] + elif self.freeze_encoder is False: + cmd += [ + "training.freeze_encoder=false", + f"training.max_epochs={PROTO_B['max_epochs']}", + f"training.early_stopping_patience={PROTO_B['patience']}", + f"optimizer.lr={PROTO_B['lr']}", + ] + if self.label_fraction < 1.0: + cmd.append(f"label_fraction={self.label_fraction}") + for k, v in self.extra_overrides.items(): + cmd.append(f"{k}={v}") + 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"sprint={self.sprint}", + f"hydra.run.dir={self.output_dir}", + ] + if self.label_fraction < 1.0: + cmd.append(f"label_fraction={self.label_fraction}") + for k, v in self.extra_overrides.items(): + cmd.append(f"{k}={v}") + return cmd + + +# --------------------------------------------------------------------------- +# Matrix Generation +# --------------------------------------------------------------------------- +def _pretrain_key(paradigm: str, dataset: str, seed: int, extra: dict | None = None) -> str: + """Canonical key for deduplicating pretrain runs.""" + key = f"pretrain_{paradigm}_{dataset}_seed{seed}" + if extra: + for k, v in sorted(extra.items()): + short_k = k.split(".")[-1] + short_v = str(v).replace(".", "") + key += f"_{short_k}{short_v}" + return key + + +def _output_dir(sprint: str, name: str) -> str: + return f"outputs/sprint{sprint}/{name}" + + +class MatrixBuilder: + """Generates all Run objects across sprints with pretrain deduplication.""" + + def __init__(self): + self.runs: list[Run] = [] + self.pretrain_index: dict[str, Run] = {} # canonical_key -> Run + + def _add_pretrain( + self, sprint: str, paradigm: str, dataset: str, seed: int, extra: dict | None = None + ) -> Run: + """Add a pretrain run, deduplicating by canonical key.""" + extra = extra or {} + key = _pretrain_key(paradigm, dataset, seed, extra) + if key in self.pretrain_index: + return self.pretrain_index[key] + + dir_name = _pretrain_key(paradigm, dataset, seed, extra) + run = Run( + id=f"s{sprint}_{dir_name}", + sprint=sprint, + run_type="pretrain", + paradigm=paradigm, + dataset=dataset, + seed=seed, + output_dir=_output_dir(sprint, dir_name), + extra_overrides=extra, + ) + self.runs.append(run) + self.pretrain_index[key] = run + return run + + def _get_pretrain( + self, paradigm: str, dataset: str, seed: int, extra: dict | None = None + ) -> Run: + """Retrieve an existing pretrain run.""" + key = _pretrain_key(paradigm, dataset, seed, extra) + if key not in self.pretrain_index: + raise KeyError(f"Pretrain not found: {key}") + return self.pretrain_index[key] + + def _add_finetune( + self, + sprint: 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, + ) -> Run: + """Add a finetune run. + + Args: + extra: Hydra overrides added to the finetune command AND name. + name_extra: Dict used only for directory/ID disambiguation + (e.g. pretrain ablation params). Not passed to command. + """ + prefix = "probe" if freeze else "finetune" + name = f"{prefix}_{paradigm}_{task}_{dataset}_seed{seed}" + if source_dataset: + name += f"_from_{source_dataset}" + if label_fraction < 1.0: + frac_str = str(label_fraction).replace(".", "") + name += f"_frac{frac_str}" + # Append both extra and name_extra to the directory name + for d in [name_extra, extra]: + if d: + for k, v in sorted(d.items()): + short_k = k.split(".")[-1] + short_v = str(v).replace(".", "") + name += f"_{short_k}{short_v}" + + run = Run( + id=f"s{sprint}_{name}", + sprint=sprint, + run_type="finetune", + paradigm=paradigm, + dataset=dataset, + seed=seed, + output_dir=_output_dir(sprint, name), + depends_on=[pretrain_run.id], + task=task, + label_fraction=label_fraction, + freeze_encoder=freeze, + extra_overrides=extra or {}, + source_dataset=source_dataset, + ) + self.runs.append(run) + return run + + def _add_supervised( + self, sprint: str, dataset: str, seed: int, task: str, label_fraction: float = 1.0 + ) -> Run: + name = f"supervised_{task}_{dataset}_seed{seed}" + if label_fraction < 1.0: + frac_str = str(label_fraction).replace(".", "") + name += f"_frac{frac_str}" + run = Run( + id=f"s{sprint}_{name}", + sprint=sprint, + run_type="supervised", + paradigm="supervised", + dataset=dataset, + seed=seed, + output_dir=_output_dir(sprint, name), + task=task, + label_fraction=label_fraction, + ) + self.runs.append(run) + return run + + # --- Sprint builders --- + + def build_sprint1(self): + """MIMIC, all tasks, Protocol B + supervised, seed=42.""" + ds, seed, sprint = "miiv", 42, "1" + for p in SSL_PARADIGMS: + pt = self._add_pretrain(sprint, p, ds, seed) + for task in TASKS: + self._add_finetune(sprint, p, ds, seed, task, False, pt) + for task in TASKS: + self._add_supervised(sprint, ds, seed, task) + + def build_sprint1b(self): + """LR sensitivity, MIMIC, mortality_24h, seed=42.""" + ds, seed, task, sprint = "miiv", 42, "mortality_24h", "1b" + for p in SSL_PARADIGMS: + for lr in LR_ABLATION: + extra = {"optimizer.lr": lr} + pt = self._add_pretrain(sprint, p, ds, seed, extra) + self._add_finetune(sprint, p, ds, seed, task, False, pt, name_extra=extra) + + def build_sprint2(self): + """MIMIC Protocol A, seed=42 — reuses Sprint 1 pretrains.""" + ds, seed, sprint = "miiv", 42, "2" + for p in SSL_PARADIGMS: + pt = self._get_pretrain(p, ds, seed) + for task in TASKS: + self._add_finetune(sprint, p, ds, seed, task, True, pt) + + def build_sprint3(self): + """eICU, both protocols + supervised, seed=42.""" + ds, seed, sprint = "eicu", 42, "3" + for p in SSL_PARADIGMS: + pt = self._add_pretrain(sprint, p, ds, seed) + for task in TASKS: + self._add_finetune(sprint, p, ds, seed, task, False, pt) + self._add_finetune(sprint, p, ds, seed, task, True, pt) + for task in TASKS: + self._add_supervised(sprint, ds, seed, task) + + def build_sprint4(self): + """Combined, both protocols + supervised, seed=42.""" + ds, seed, sprint = "combined", 42, "4" + for p in SSL_PARADIGMS: + pt = self._add_pretrain(sprint, p, ds, seed) + for task in TASKS: + self._add_finetune(sprint, p, ds, seed, task, False, pt) + self._add_finetune(sprint, p, ds, seed, task, True, pt) + for task in TASKS: + self._add_supervised(sprint, ds, seed, task) + + def build_sprint5(self): + """Seeds 123, 456 for datasets miiv, eicu, combined.""" + sprint = "5" + for seed in [123, 456]: + for ds in DATASETS: + for p in SSL_PARADIGMS: + pt = self._add_pretrain(sprint, p, ds, seed) + for task in TASKS: + self._add_finetune(sprint, p, ds, seed, task, False, pt) + self._add_finetune(sprint, p, ds, seed, task, True, pt) + for task in TASKS: + self._add_supervised(sprint, ds, seed, task) + + def build_sprint6(self): + """Label efficiency ablation — reuses Phase 1 encoders.""" + sprint = "6" + for seed in SEEDS: + for ds in DATASETS: + for p in SSL_PARADIGMS: + pt = self._get_pretrain(p, ds, seed) + # mortality_24h gets full sweep + for frac in LABEL_FRACTIONS_FULL: + self._add_finetune(sprint, p, ds, seed, "mortality_24h", False, pt, frac) + self._add_finetune(sprint, p, ds, seed, "mortality_24h", True, pt, frac) + # Other tasks get trend check + for task in TASKS[1:]: + for frac in LABEL_FRACTIONS_TREND: + self._add_finetune(sprint, p, ds, seed, task, False, pt, frac) + self._add_finetune(sprint, p, ds, seed, task, True, pt, frac) + # Supervised label efficiency + for frac in LABEL_FRACTIONS_FULL: + self._add_supervised(sprint, ds, seed, "mortality_24h", frac) + for task in TASKS[1:]: + for frac in LABEL_FRACTIONS_TREND: + self._add_supervised(sprint, ds, seed, task, frac) + + def build_sprint7(self): + """Cross-dataset transfer — Protocol B only (full finetune).""" + sprint = "7" + for seed in SEEDS: + for source_ds, target_ds in TRANSFER_PAIRS: + for p in SSL_PARADIGMS: + pt = self._get_pretrain(p, source_ds, seed) + for task in TASKS: + self._add_finetune( + sprint, p, target_ds, seed, task, False, pt, source_dataset=source_ds + ) + + def build_sprint8(self): + """LR + mask ablations, seeds 123, 456.""" + sprint = "8" + for seed in [123, 456]: + for p in SSL_PARADIGMS: + # LR ablation + for lr in LR_ABLATION: + extra = {"optimizer.lr": lr} + pt = self._add_pretrain(sprint, p, "miiv", seed, extra) + self._add_finetune( + sprint, p, "miiv", seed, "mortality_24h", False, pt, name_extra=extra + ) + # Mask ratio ablation + for mr in MASK_RATIO_ABLATION: + extra = {"ssl.mask_ratio": mr} + pt = self._add_pretrain(sprint, p, "miiv", seed, extra) + self._add_finetune( + sprint, p, "miiv", seed, "mortality_24h", False, pt, name_extra=extra + ) + + def build_all(self) -> list[Run]: + """Build full experiment matrix. Order matters for dedup.""" + self.build_sprint1() + self.build_sprint1b() + self.build_sprint2() + self.build_sprint3() + self.build_sprint4() + self.build_sprint5() + self.build_sprint6() + self.build_sprint7() + self.build_sprint8() + return self.runs + + +def generate_all_runs() -> list[Run]: + builder = MatrixBuilder() + return builder.build_all() + + +# --------------------------------------------------------------------------- +# 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): + STATE_FILE.parent.mkdir(parents=True, exist_ok=True) + # Atomic write + 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 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): + if run_id not in state["runs"]: + state["runs"][run_id] = {} + entry = state["runs"][run_id] + entry["status"] = status + entry.update(kwargs) + + +def is_pid_alive(pid: int) -> bool: + try: + os.kill(pid, 0) + return True + except (OSError, ProcessLookupError): + return False + + +def recover_stale_running(state: dict): + """Reset running entries whose PIDs are dead back to pending.""" + 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(pid): + info["status"] = "pending" + info.pop("pid", None) + print(f" Recovered stale run: {run_id}") + + +# --------------------------------------------------------------------------- +# Scheduler +# --------------------------------------------------------------------------- +def run_scheduler(runs: list[Run], state: dict, parallel: int, dry_run: bool): + runs_by_id = {r.id: r for r in runs} + active: dict[str, subprocess.Popen] = {} # run_id -> 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 rid, proc in active.items(): + try: + proc.terminate() + except OSError: + pass + # Wait up to 30s for children + deadline = time.time() + 30 + for rid, proc in list(active.items()): + remaining = max(0, deadline - time.time()) + try: + proc.wait(timeout=remaining) + except subprocess.TimeoutExpired: + proc.kill() + # Reset interrupted runs to pending + for rid in active: + set_run_status(state, rid, "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 + + print(f"Scheduler started (parallel={parallel})") + print(f"State file: {STATE_FILE}") + print() + + while not shutting_down: + # 1. Poll active processes + finished = [] + for rid, proc in active.items(): + ret = proc.poll() + if ret is not None: + finished.append((rid, ret)) + + # 2. Handle completions + for rid, exit_code in finished: + proc = active.pop(rid) + fh = getattr(proc, "_log_fh", None) + if fh: + fh.close() + run = runs_by_id[rid] + now = datetime.now(timezone.utc).isoformat() + + elapsed = _format_elapsed(state, rid, now) + + if exit_code == 0: + # For pretrain, verify encoder.pt exists + if run.run_type == "pretrain": + encoder_path = Path(run.output_dir) / "encoder.pt" + if not encoder_path.exists(): + print(f" FAILED {rid}: exit 0 but encoder.pt " f"missing {elapsed}") + set_run_status(state, rid, "failed", finished_at=now, exit_code=exit_code) + _propagate_failure(rid, runs, state) + continue + print(f" DONE {rid} {elapsed}") + set_run_status(state, rid, "completed", finished_at=now, exit_code=0) + else: + print(f" FAILED {rid} (exit {exit_code}) {elapsed}") + set_run_status(state, rid, "failed", finished_at=now, exit_code=exit_code) + _propagate_failure(rid, runs, state) + + # 3. Find ready runs + ready = [] + for r in runs: + status = get_run_status(state, r.id) + if status != "pending": + continue + if r.id in active: + continue + # Check all dependencies completed + deps_ok = all(get_run_status(state, d) == "completed" for d in r.depends_on) + if deps_ok: + ready.append(r) + + # 4. Launch runs up to parallel limit + slots = parallel - len(active) + for r in ready[:slots]: + cmd = r.build_command(runs_by_id) + log_file = LOG_DIR / f"{r.id}.log" + now = datetime.now(timezone.utc).isoformat() + + print(f" START {r.id}") + log_fh = open(log_file, "w") + proc = subprocess.Popen( + cmd, + stdout=log_fh, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + # Store log handle for cleanup + proc._log_fh = log_fh # type: ignore[attr-defined] + active[r.id] = proc + set_run_status( + state, r.id, "running", started_at=now, pid=proc.pid, log_file=str(log_file) + ) + + # 5. Save state + save_state(state) + + # 6. Exit condition + if not active and not ready: + break + + time.sleep(2) + + # Close any remaining log handles + for proc in active.values(): + fh = getattr(proc, "_log_fh", None) + if fh: + fh.close() + + # Print summary + print() + _print_summary(runs, state) + + +def _format_elapsed(state: dict, run_id: str, now_iso: str) -> str: + """Format elapsed time as (Xh Ym Zs).""" + 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)" + elif secs < 3600: + return f"({secs // 60}m{secs % 60:02d}s)" + else: + h, rem = divmod(secs, 3600) + m, s = divmod(rem, 60) + return f"({h}h{m:02d}m{s:02d}s)" + except (ValueError, TypeError): + return "" + + +def _propagate_failure(failed_id: str, runs: list[Run], state: dict): + """Skip runs that depend on a failed run (transitive).""" + queue = [failed_id] + while queue: + fid = queue.pop(0) + for r in runs: + if fid in r.depends_on and get_run_status(state, r.id) == "pending": + set_run_status(state, r.id, "skipped", reason=f"dependency {fid} failed") + print(f" SKIP {r.id} (dep {fid} failed)") + queue.append(r.id) + + +def _print_dry_run(runs: list[Run], runs_by_id: dict[str, Run]): + """Print all runs and their commands without executing.""" + print(f"DRY RUN: {len(runs)} runs\n") + for r in runs: + deps = ", ".join(r.depends_on) if r.depends_on else "(none)" + cmd = " ".join(r.build_command(runs_by_id)) + print(f"[{r.id}]") + print(f" sprint={r.sprint} type={r.run_type} deps={deps}") + print(f" dir={r.output_dir}") + print(f" cmd: {cmd}") + print() + + +# --------------------------------------------------------------------------- +# Status / Retry +# --------------------------------------------------------------------------- +def print_status(sprint_filter: list[str] | None = None): + all_runs = generate_all_runs() + state = load_state() + + # Group by sprint + sprints: dict[str, list[Run]] = {} + for r in all_runs: + sprints.setdefault(r.sprint, []).append(r) + + sprint_keys = sorted(sprints.keys(), key=_sprint_sort_key) + if sprint_filter: + sprint_keys = [s for s in sprint_keys if s in sprint_filter] + + print( + f"{'Sprint':>6} | {'Total':>5} | {'Done':>4} | {'Run':>3} | " + f"{'Fail':>4} | {'Pend':>4} | {'Skip':>4}" + ) + print("-" * 52) + + totals = {"total": 0, "completed": 0, "running": 0, "failed": 0, "pending": 0, "skipped": 0} + + for s in sprint_keys: + runs = sprints[s] + counts = { + "total": len(runs), + "completed": 0, + "running": 0, + "failed": 0, + "pending": 0, + "skipped": 0, + } + for r in runs: + status = get_run_status(state, r.id) + if status in counts: + counts[status] += 1 + else: + counts["pending"] += 1 # unknown → pending + print( + f"{s:>6} | {counts['total']:>5} | {counts['completed']:>4} | " + f"{counts['running']:>3} | {counts['failed']:>4} | " + f"{counts['pending']:>4} | {counts['skipped']:>4}" + ) + for k in totals: + totals[k] += counts[k] + + print("-" * 52) + print( + f"{'TOTAL':>6} | {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): + """Print completion summary for a scheduler run.""" + counts = {"completed": 0, "failed": 0, "skipped": 0, "pending": 0} + for r in runs: + s = get_run_status(state, r.id) + counts[s] = counts.get(s, 0) + 1 + total = len(runs) + print( + f"Summary: {counts['completed']}/{total} completed, " + f"{counts['failed']} failed, {counts['skipped']} skipped, " + f"{counts['pending']} pending" + ) + + +def _sprint_sort_key(s: str) -> tuple[int, str]: + """Sort sprints: 1, 1b, 2, 3, ...""" + num = "" + suffix = "" + for c in s: + if c.isdigit(): + num += c + else: + suffix += c + return (int(num) if num else 0, suffix) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- +def cmd_run(args): + all_runs = generate_all_runs() + state = load_state() + recover_stale_running(state) + + # Filter to requested sprints + sprints = [str(s) for s in args.sprint] + runs = [r for r in all_runs if r.sprint in sprints] + + # Also include dependency runs from earlier sprints (pretrain reuse) + run_ids = {r.id for r in runs} + deps_needed = set() + for r in runs: + for d in r.depends_on: + if d not in run_ids: + deps_needed.add(d) + if deps_needed: + all_by_id = {r.id: r for r in all_runs} + extra = [all_by_id[d] for d in deps_needed if d in all_by_id] + runs = extra + runs + if extra and not args.dry_run: + print(f"Including {len(extra)} dependency run(s) from earlier " f"sprints") + + if not runs: + print(f"No runs found for sprint(s): {', '.join(sprints)}") + return + + print(f"Sprint(s) {', '.join(sprints)}: {len(runs)} runs") + run_scheduler(runs, state, args.parallel, args.dry_run) + + +def cmd_status(args): + sprint_filter = [str(s) for s in args.sprint] if args.sprint else None + print_status(sprint_filter) + + +def cmd_retry(args): + all_runs = generate_all_runs() + state = load_state() + recover_stale_running(state) + + runs_to_retry = [] + for r in all_runs: + status = get_run_status(state, r.id) + if args.failed and status == "failed": + set_run_status(state, r.id, "pending") + runs_to_retry.append(r) + elif args.skipped and status == "skipped": + set_run_status(state, r.id, "pending") + runs_to_retry.append(r) + + if not runs_to_retry: + print("No runs to retry.") + return + + # Include their dependencies that are completed (already fine) or pending + all_by_id = {r.id: r for r in all_runs} + retry_ids = {r.id for r in runs_to_retry} + deps_needed = set() + for r in runs_to_retry: + for d in r.depends_on: + if d not in retry_ids: + deps_needed.add(d) + extra = [all_by_id[d] for d in deps_needed if d in all_by_id] + runs_to_retry = extra + runs_to_retry + + save_state(state) + print(f"Retrying {len(runs_to_retry)} runs") + run_scheduler(runs_to_retry, state, args.parallel, dry_run=False) + + +def main(): + parser = argparse.ArgumentParser(description="SLICES parallel experiment runner") + sub = parser.add_subparsers(dest="command", required=True) + + # run + p_run = sub.add_parser("run", help="Run experiments for given sprints") + p_run.add_argument("--sprint", nargs="+", required=True, help="Sprint(s) to run (e.g. 1 1b 2)") + p_run.add_argument("--parallel", type=int, default=4, help="Max parallel jobs (default: 4)") + p_run.add_argument("--dry-run", action="store_true", help="Print runs without executing") + + # status + p_status = sub.add_parser("status", help="Show experiment status") + p_status.add_argument("--sprint", nargs="*", default=None, help="Filter by sprint(s)") + + # retry + 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 skipped runs") + p_retry.add_argument("--parallel", type=int, default=4, help="Max parallel jobs (default: 4)") + + args = parser.parse_args() + + if args.command == "run": + cmd_run(args) + elif args.command == "status": + cmd_status(args) + elif args.command == "retry": + cmd_retry(args) + + +if __name__ == "__main__": + main() From 38ac46cde10ae0af1a0dca24886e885a007aad59 Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 4 Mar 2026 18:53:05 -0500 Subject: [PATCH 020/121] fix: revert supervised-specific regularization, enable balanced class weights MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Revert encoder.dropout (0.2→0.1), weight_decay (0.05→0.01), and label_smoothing (0.1→removed) overrides in supervised.yaml to maintain controlled variable parity with SSL paradigms per experiment plan. Enable class_weight: balanced in both supervised.yaml and finetune.yaml to match experiment plan specification of inverse frequency weighting for imbalanced classification tasks. --- configs/finetune.yaml | 2 +- configs/supervised.yaml | 12 +-- scripts/launch_sprint1.sh | 152 -------------------------------------- 3 files changed, 5 insertions(+), 161 deletions(-) delete mode 100755 scripts/launch_sprint1.sh diff --git a/configs/finetune.yaml b/configs/finetune.yaml index dfb507f..9ace414 100644 --- a/configs/finetune.yaml +++ b/configs/finetune.yaml @@ -84,7 +84,7 @@ 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 # Compute settings accelerator: auto # auto | cpu | gpu | mps diff --git a/configs/supervised.yaml b/configs/supervised.yaml index e73cebd..97e9f3b 100644 --- a/configs/supervised.yaml +++ b/configs/supervised.yaml @@ -46,9 +46,8 @@ ckpt_path: null output_dir: ${hydra:runtime.output_dir} checkpoint_dir: ${output_dir}/checkpoints -# Encoder overrides — increase dropout for training from scratch -encoder: - dropout: 0.2 +# 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. @@ -78,10 +77,7 @@ training: # null: no weighting (default) # "balanced": auto-compute inverse frequency weights from training labels # [w0, w1, ...]: explicit per-class weights - class_weight: null - - # Label smoothing — helps calibration with highly imbalanced tasks - label_smoothing: 0.1 + class_weight: balanced # Regularization gradient_clip_val: 1.0 @@ -101,7 +97,7 @@ training: optimizer: name: adamw lr: 1.0e-3 # Higher LR for training from scratch - weight_decay: 0.05 + weight_decay: 0.01 # Learning rate scheduler scheduler: diff --git a/scripts/launch_sprint1.sh b/scripts/launch_sprint1.sh deleted file mode 100755 index 15490dc..0000000 --- a/scripts/launch_sprint1.sh +++ /dev/null @@ -1,152 +0,0 @@ -#!/usr/bin/env bash -# ============================================================================= -# Sprint 1: Sanity Check -# MIMIC-IV | All 4 paradigms | mortality_24h | seed=42 -# -# Runs 3 SSL pretrain jobs in parallel + 1 supervised baseline. -# After each pretrain finishes, automatically launches finetuning. -# ============================================================================= -set -euo pipefail - -SPRINT=1 -DATASET=miiv -SEED=42 -TASK=mortality_24h -LOG_DIR="logs/sprint1" - -mkdir -p "$LOG_DIR" - -echo "============================================================" -echo "SLICES Sprint 1: Sanity Check" -echo " Dataset: $DATASET" -echo " Task: $TASK" -echo " Seed: $SEED" -echo " Paradigms: MAE, JEPA, Contrastive, Supervised" -echo "============================================================" - -# Use predictable output dirs (override Hydra's timestamped dirs) -MAE_DIR="outputs/sprint1/pretrain_mae" -JEPA_DIR="outputs/sprint1/pretrain_jepa" -CONTRASTIVE_DIR="outputs/sprint1/pretrain_contrastive" -SUPERVISED_DIR="outputs/sprint1/supervised" -FT_MAE_DIR="outputs/sprint1/finetune_mae_${TASK}" -FT_JEPA_DIR="outputs/sprint1/finetune_jepa_${TASK}" -FT_CONTRASTIVE_DIR="outputs/sprint1/finetune_contrastive_${TASK}" - -# ----------------------------------------------------------------------------- -# Helper: run pretrain then finetune -# ----------------------------------------------------------------------------- -pretrain_and_finetune() { - local ssl_name="$1" - local pretrain_dir="$2" - local finetune_dir="$3" - local log_prefix="$LOG_DIR/${ssl_name}" - - echo "[$(date '+%H:%M:%S')] Starting $ssl_name pretraining..." - - uv run python scripts/training/pretrain.py \ - dataset=$DATASET \ - ssl=$ssl_name \ - seed=$SEED \ - sprint=$SPRINT \ - hydra.run.dir=$pretrain_dir \ - 2>&1 | tee "${log_prefix}_pretrain.log" - - echo "[$(date '+%H:%M:%S')] $ssl_name pretraining complete." - - # Check encoder was saved - if [ ! -f "${pretrain_dir}/encoder.pt" ]; then - echo "ERROR: ${pretrain_dir}/encoder.pt not found!" - return 1 - fi - - echo "[$(date '+%H:%M:%S')] Starting $ssl_name finetuning on $TASK..." - - uv run python scripts/training/finetune.py \ - dataset=$DATASET \ - checkpoint=${pretrain_dir}/encoder.pt \ - tasks=$TASK \ - seed=$SEED \ - sprint=$SPRINT \ - hydra.run.dir=$finetune_dir \ - 2>&1 | tee "${log_prefix}_finetune.log" - - echo "[$(date '+%H:%M:%S')] $ssl_name finetuning complete." -} - -# ----------------------------------------------------------------------------- -# Launch all 4 runs in parallel -# ----------------------------------------------------------------------------- - -# MAE: pretrain -> finetune -pretrain_and_finetune mae "$MAE_DIR" "$FT_MAE_DIR" & -PID_MAE=$! - -# JEPA: pretrain -> finetune -pretrain_and_finetune jepa "$JEPA_DIR" "$FT_JEPA_DIR" & -PID_JEPA=$! - -# Contrastive: pretrain -> finetune -pretrain_and_finetune contrastive "$CONTRASTIVE_DIR" "$FT_CONTRASTIVE_DIR" & -PID_CONTRASTIVE=$! - -# Supervised: single run (no pretrain step) -( - echo "[$(date '+%H:%M:%S')] Starting supervised baseline..." - - uv run python scripts/training/supervised.py \ - dataset=$DATASET \ - tasks=$TASK \ - seed=$SEED \ - sprint=$SPRINT \ - hydra.run.dir=$SUPERVISED_DIR \ - 2>&1 | tee "${LOG_DIR}/supervised.log" - - echo "[$(date '+%H:%M:%S')] Supervised baseline complete." -) & -PID_SUPERVISED=$! - -echo "" -echo "All 4 jobs launched in parallel:" -echo " MAE: PID=$PID_MAE" -echo " JEPA: PID=$PID_JEPA" -echo " Contrastive: PID=$PID_CONTRASTIVE" -echo " Supervised: PID=$PID_SUPERVISED" -echo "" -echo "Logs: $LOG_DIR/" -echo " mae_pretrain.log / mae_finetune.log" -echo " jepa_pretrain.log / jepa_finetune.log" -echo " contrastive_pretrain.log / contrastive_finetune.log" -echo " supervised.log" -echo "" -echo "Waiting for all jobs to finish..." - -# Wait and report results -FAILED=0 - -wait $PID_MAE || { echo "FAILED: MAE pipeline"; FAILED=$((FAILED + 1)); } -wait $PID_JEPA || { echo "FAILED: JEPA pipeline"; FAILED=$((FAILED + 1)); } -wait $PID_CONTRASTIVE || { echo "FAILED: Contrastive pipeline"; FAILED=$((FAILED + 1)); } -wait $PID_SUPERVISED || { echo "FAILED: Supervised baseline"; FAILED=$((FAILED + 1)); } - -echo "" -echo "============================================================" -if [ $FAILED -eq 0 ]; then - echo "Sprint 1 COMPLETE — all 4 paradigms finished successfully." -else - echo "Sprint 1 FINISHED WITH $FAILED FAILURE(S). Check logs." -fi -echo "============================================================" -echo "" -echo "Output directories:" -echo " MAE pretrain: $MAE_DIR" -echo " MAE finetune: $FT_MAE_DIR" -echo " JEPA pretrain: $JEPA_DIR" -echo " JEPA finetune: $FT_JEPA_DIR" -echo " Contrastive pretrain: $CONTRASTIVE_DIR" -echo " Contrastive finetune: $FT_CONTRASTIVE_DIR" -echo " Supervised: $SUPERVISED_DIR" -echo "" -echo "Check W&B for training curves: https://wandb.ai/slices" - -exit $FAILED From 0fc68211475b6a05958b5700e052adad428a24a4 Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 4 Mar 2026 23:14:27 -0500 Subject: [PATCH 021/121] fix: AKI KDIGO label builder detection window and class weight guard Detection window was filtering for hour >= obs_hours (48), but timeseries data only covers hours 0-47, producing zero positive labels. Changed to detect AKI after baseline period (hour >= baseline_hours, default 24). Also fixed absolute criterion to use all creatinine data with detection window filter, catching rises that span the baseline/detection boundary. This increased AKI positive rate from 6.4% to 12.1%. Added ValueError guard in supervised.py and finetune.py for zero-count classes when computing balanced weights. --- scripts/training/finetune.py | 9 ++++++-- scripts/training/supervised.py | 9 ++++++-- src/slices/data/labels/aki.py | 42 +++++++++++++++++++--------------- 3 files changed, 37 insertions(+), 23 deletions(-) diff --git a/scripts/training/finetune.py b/scripts/training/finetune.py index 17742b4..a3566c1 100644 --- a/scripts/training/finetune.py +++ b/scripts/training/finetune.py @@ -142,9 +142,14 @@ def main(cfg: DictConfig) -> None: 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_pos = stats.get("positive", 0) + n_neg = 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." + ) cfg.training.class_weight = [n_total / (2 * n_neg), n_total / (2 * n_pos)] print(f"\n Balanced class weights: {cfg.training.class_weight}") else: diff --git a/scripts/training/supervised.py b/scripts/training/supervised.py index 061d337..4368f7e 100644 --- a/scripts/training/supervised.py +++ b/scripts/training/supervised.py @@ -208,9 +208,14 @@ def main(cfg: DictConfig) -> None: 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_pos = stats.get("positive", 0) + n_neg = 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." + ) cfg.training.class_weight = [n_total / (2 * n_neg), n_total / (2 * n_pos)] print(f"\n Balanced class weights: {cfg.training.class_weight}") else: diff --git a/src/slices/data/labels/aki.py b/src/slices/data/labels/aki.py index d497be9..45e2098 100644 --- a/src/slices/data/labels/aki.py +++ b/src/slices/data/labels/aki.py @@ -17,8 +17,8 @@ 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. + Baseline: minimum creatinine in first baseline_window_hours (default 24h). + Evaluates criteria AFTER the baseline window within the observation period. Label = null if no creatinine measurements available. """ @@ -43,7 +43,6 @@ 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) abs_threshold = self.config.label_params.get("absolute_rise_threshold", 0.3) @@ -76,16 +75,13 @@ 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) + # Post-baseline creatinine values (detection window) + post_obs = crea_ts.filter(pl.col("hour") >= baseline_hours) - # 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. rel_aki = ( post_with_baseline.filter(pl.col("hour") < rel_window_hours) .filter(pl.col(crea_col) >= pl.col("baseline") * rel_threshold) @@ -95,8 +91,12 @@ 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 (>= baseline_hours), so we catch rises + # that span the baseline/detection boundary. + abs_aki = self._check_absolute_criterion_vectorized( + crea_ts, crea_col, abs_threshold, baseline_hours + ) # Combine criteria: AKI = 1 if either criterion is met stay_df = pl.DataFrame({"stay_id": pl.Series(stay_ids, dtype=pl.Int64)}) @@ -107,7 +107,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 = ( @@ -132,7 +132,7 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: # Build label: # - null if no creatinine data or no baseline - # - 0 if no post-obs data (survived observation without AKI evidence after) + # - 0 if no post-baseline data (no creatinine to evaluate after baseline period) # - 1 if either AKI criterion met # - 0 otherwise result = result.with_columns( @@ -163,19 +163,21 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: def _check_absolute_criterion_vectorized( self, - post_obs: pl.DataFrame, + crea_ts: pl.DataFrame, crea_col: str, abs_threshold: float, + detection_start_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 (hour >= detection_start_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 +186,23 @@ 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 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("crea_j") - pl.col("crea_i") >= abs_threshold) ) From 3eaa2836f292bf4fb7070e179fce1ee464962c3c Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 4 Mar 2026 23:40:40 -0500 Subject: [PATCH 022/121] fix: add downstream regularization to combat overfitting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Head dropout 0.1 → 0.3 (all task configs) - Class weights: balanced → sqrt(balanced) to reduce gradient amplification - Downstream LR 1e-3 → 3e-4 (supervised config + PROTO_B) - Weight decay 0.01 → 0.05 (supervised + finetune configs) - Early stopping patience 20 → 10 (supervised config + PROTO_B) - Label smoothing 0.0 → 0.1 (supervised + finetune configs) --- configs/finetune.yaml | 6 +++--- configs/supervised.yaml | 9 ++++++--- configs/tasks/aki_kdigo.yaml | 2 +- configs/tasks/los_remaining.yaml | 2 +- configs/tasks/mortality_24h.yaml | 2 +- configs/tasks/mortality_hospital.yaml | 2 +- scripts/run_experiments.py | 2 +- scripts/training/finetune.py | 5 +++-- scripts/training/supervised.py | 5 +++-- 9 files changed, 20 insertions(+), 15 deletions(-) diff --git a/configs/finetune.yaml b/configs/finetune.yaml index 9ace414..499e680 100644 --- a/configs/finetune.yaml +++ b/configs/finetune.yaml @@ -101,8 +101,8 @@ training: # Override here only if you need a non-default metric. early_stopping_patience: 10 - # Label smoothing (default off for finetuning — frozen encoder doesn't overfit) - label_smoothing: 0.0 + # Label smoothing for classification tasks + label_smoothing: 0.1 # Debug: overfit on N batches (0 = disabled) overfit_batches: 0 @@ -112,7 +112,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: diff --git a/configs/supervised.yaml b/configs/supervised.yaml index 97e9f3b..5af387d 100644 --- a/configs/supervised.yaml +++ b/configs/supervised.yaml @@ -87,7 +87,10 @@ training: # Monitor/mode are set automatically by task type in setup_callbacks(): # 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 @@ -96,8 +99,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: diff --git a/configs/tasks/aki_kdigo.yaml b/configs/tasks/aki_kdigo.yaml index c2b5e11..0aff0a8 100644 --- a/configs/tasks/aki_kdigo.yaml +++ b/configs/tasks/aki_kdigo.yaml @@ -7,7 +7,7 @@ task_type: binary 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..ce2562c 100644 --- a/configs/tasks/los_remaining.yaml +++ b/configs/tasks/los_remaining.yaml @@ -7,7 +7,7 @@ task_type: regression 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_24h.yaml b/configs/tasks/mortality_24h.yaml index cd21f21..f0c0cfa 100644 --- a/configs/tasks/mortality_24h.yaml +++ b/configs/tasks/mortality_24h.yaml @@ -7,7 +7,7 @@ task_type: binary 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..6a1a5c9 100644 --- a/configs/tasks/mortality_hospital.yaml +++ b/configs/tasks/mortality_hospital.yaml @@ -7,7 +7,7 @@ task_type: binary 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/scripts/run_experiments.py b/scripts/run_experiments.py index 6b5a622..c6befe3 100644 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -44,7 +44,7 @@ # Protocol defaults PROTO_A = {"freeze_encoder": True, "max_epochs": 50, "patience": 10, "lr": 1e-4} -PROTO_B = {"freeze_encoder": False, "max_epochs": 100, "patience": 20, "lr": 1e-3} +PROTO_B = {"freeze_encoder": False, "max_epochs": 100, "patience": 10, "lr": 3e-4} # --------------------------------------------------------------------------- diff --git a/scripts/training/finetune.py b/scripts/training/finetune.py index a3566c1..041e9ba 100644 --- a/scripts/training/finetune.py +++ b/scripts/training/finetune.py @@ -150,8 +150,9 @@ def main(cfg: DictConfig) -> None: f"Cannot compute balanced class weights for '{task_name}': " f"{n_pos} positive, {n_neg} negative. Check label extraction." ) - cfg.training.class_weight = [n_total / (2 * n_neg), n_total / (2 * n_pos)] - print(f"\n Balanced class weights: {cfg.training.class_weight}") + raw = [n_total / (2 * n_neg), n_total / (2 * n_pos)] + cfg.training.class_weight = [w**0.5 for w in raw] + 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") cfg.training.class_weight = None diff --git a/scripts/training/supervised.py b/scripts/training/supervised.py index 4368f7e..73b607d 100644 --- a/scripts/training/supervised.py +++ b/scripts/training/supervised.py @@ -216,8 +216,9 @@ def main(cfg: DictConfig) -> None: f"Cannot compute balanced class weights for '{task_name}': " f"{n_pos} positive, {n_neg} negative. Check label extraction." ) - cfg.training.class_weight = [n_total / (2 * n_neg), n_total / (2 * n_pos)] - print(f"\n Balanced class weights: {cfg.training.class_weight}") + raw = [n_total / (2 * n_neg), n_total / (2 * n_pos)] + cfg.training.class_weight = [w**0.5 for w in raw] + 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") cfg.training.class_weight = None From 174a5f8a732badfa36dde763d87b356e7b8db356 Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 5 Mar 2026 00:11:16 -0500 Subject: [PATCH 023/121] feat: add warmup subcommand to pre-build tensor caches before parallel runs Adds a `warmup` command that sequentially builds tensor caches for all unique (dataset, task, seed, label_fraction) combinations needed by the requested sprints. This prevents OOM kills when running parallel experiments on the first run, where each process would otherwise independently load and preprocess raw parquet files. --- scripts/run_experiments.py | 78 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index c6befe3..e7dedef 100644 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -6,6 +6,7 @@ and executes them in parallel with crash recovery and state persistence. Usage: + uv run python scripts/run_experiments.py warmup --sprint 1 uv run python scripts/run_experiments.py run --sprint 1 --parallel 4 uv run python scripts/run_experiments.py run --sprint 1 2 3 --parallel 6 --dry-run uv run python scripts/run_experiments.py status @@ -818,6 +819,73 @@ def cmd_retry(args): run_scheduler(runs_to_retry, state, args.parallel, dry_run=False) +def cmd_warmup(args): + """Pre-build tensor caches for requested sprints. + + Instantiates ICUDataModule for each unique (dataset, task, seed, label_fraction) + combination sequentially. This populates the tensor cache so that parallel + experiment runs can load preprocessed data from disk instead of each process + independently loading and converting raw parquet files (which causes OOM). + """ + all_runs = generate_all_runs() + + # Filter to requested sprints + sprints = [str(s) for s in args.sprint] + runs = [r for r in all_runs if r.sprint in sprints] + + # Also include dependency runs + run_ids = {r.id for r in runs} + deps_needed = set() + for r in runs: + for d in r.depends_on: + if d not in run_ids: + deps_needed.add(d) + if deps_needed: + all_by_id = {r.id: r for r in all_runs} + runs = [all_by_id[d] for d in deps_needed if d in all_by_id] + runs + + # Collect unique (dataset, task, seed, label_fraction) combos + # task=None for pretrain, task= for finetune/supervised + combos: dict[tuple, None] = {} # ordered set via dict + for r in runs: + if r.run_type == "pretrain": + key = (r.dataset, None, r.seed, 1.0) + else: + key = (r.dataset, r.task, r.seed, r.label_fraction) + combos[key] = None + + print(f"Warming up tensor caches for sprint(s) {', '.join(sprints)}") + print(f" {len(combos)} unique (dataset, task, seed, label_fraction) combinations\n") + + # Import here to avoid loading heavy deps when not needed + from slices.data.datamodule import ICUDataModule + + for i, (dataset, task, seed, label_fraction) in enumerate(combos, 1): + processed_dir = f"data/processed/{dataset}" + task_str = task or "(pretrain/no task)" + frac_str = f", frac={label_fraction}" if label_fraction < 1.0 else "" + print(f"[{i}/{len(combos)}] {dataset} / {task_str} / seed={seed}{frac_str}") + + try: + dm = ICUDataModule( + processed_dir=processed_dir, + task_name=task, + batch_size=1, # doesn't matter, we just need setup() + num_workers=0, + seed=seed, + label_fraction=label_fraction, + ) + dm.setup() + print(f" -> Cached ({len(dm.dataset):,} samples)") + # Free memory before next combo + del dm + except Exception as e: + print(f" -> ERROR: {e}") + + print("\nWarmup complete. Tensor caches saved to data/processed//.tensor_cache/") + print("You can now run experiments in parallel without OOM.") + + def main(): parser = argparse.ArgumentParser(description="SLICES parallel experiment runner") sub = parser.add_subparsers(dest="command", required=True) @@ -838,6 +906,14 @@ def main(): p_retry.add_argument("--skipped", action="store_true", help="Retry skipped runs") p_retry.add_argument("--parallel", type=int, default=4, help="Max parallel jobs (default: 4)") + # warmup + p_warmup = sub.add_parser( + "warmup", help="Pre-build tensor caches to avoid OOM during parallel runs" + ) + p_warmup.add_argument( + "--sprint", nargs="+", required=True, help="Sprint(s) to warmup (e.g. 1 1b 2)" + ) + args = parser.parse_args() if args.command == "run": @@ -846,6 +922,8 @@ def main(): cmd_status(args) elif args.command == "retry": cmd_retry(args) + elif args.command == "warmup": + cmd_warmup(args) if __name__ == "__main__": From 9efcafd00ea0b3a293352c232ee4956e83a9382c Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 12 Mar 2026 12:22:23 -0400 Subject: [PATCH 024/121] feat: add lr=2e-4 to LR ablation sweep for sprint 1b --- scripts/run_experiments.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index e7dedef..87da808 100644 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -36,7 +36,7 @@ SEEDS = [42, 123, 456] LABEL_FRACTIONS_FULL = [0.01, 0.05, 0.1, 0.25, 0.5] LABEL_FRACTIONS_TREND = [0.1] -LR_ABLATION = [5e-4, 2e-3] # 1e-3 reused from Phase 1 +LR_ABLATION = [2e-4, 5e-4, 2e-3] # 1e-3 reused from Phase 1 MASK_RATIO_ABLATION = [0.3, 0.75] # 0.5 reused from Phase 1 TRANSFER_PAIRS = [("miiv", "eicu"), ("eicu", "miiv")] From 506b95793b14d33bfe241015039ccfaf38e9101c Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 12 Mar 2026 12:57:05 -0400 Subject: [PATCH 025/121] feat: add complementary masks for contrastive SSL objective Switch default contrastive masking from two independent random masks (~25% overlap) to complementary masks (view 2 = ~view 1, zero overlap). This forces the encoder to learn abstract temporal semantics from non-overlapping windows instead of exploiting shared timesteps. - Add complementary_masks config flag (default True) - Validate incompatibility with temporal mode - Update docstrings and YAML config - Add tests for zero overlap, full coverage, and validation --- configs/ssl/contrastive.yaml | 4 + src/slices/models/pretraining/contrastive.py | 39 ++++++--- tests/test_contrastive_objective.py | 92 ++++++++++++++++++++ 3 files changed, 124 insertions(+), 11 deletions(-) diff --git a/configs/ssl/contrastive.yaml b/configs/ssl/contrastive.yaml index 5ea2cf7..4a94bc8 100644 --- a/configs/ssl/contrastive.yaml +++ b/configs/ssl/contrastive.yaml @@ -22,3 +22,7 @@ proj_output_dim: 64 # Output dimension (contrastive embedding space, = encode # NT-Xent temperature (0.07 suits instance-level with ~511 negatives at B=256) temperature: 0.07 + +# 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/src/slices/models/pretraining/contrastive.py b/src/slices/models/pretraining/contrastive.py index d212e06..6bb2c64 100644 --- a/src/slices/models/pretraining/contrastive.py +++ b/src/slices/models/pretraining/contrastive.py @@ -1,24 +1,30 @@ """Contrastive (SimCLR-style) SSL objective for ICU time-series. -Timestep-level tokenization variant: uses two different random timestep masks as -two augmented "views" of the same sample, then applies NT-Xent contrastive loss. +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: -- **temporal** (default): Positive pairs are formed from timesteps visible in +- **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. -- **instance**: Original SimCLR-style — mean-pool each view to a single - sequence-level embedding, then NT-Xent on a (2B x 2B) matrix. + N = number of overlap tokens. Requires complementary_masks=False. Architecture: 1. TransformerEncoder.tokenize() -> one token per timestep (B, T, d_model) -2. Two independent random masks -> two different subsets of timestep tokens (views) +2. Complementary masks (default) or two independent random masks -> two views 3. Encoder processes each view separately -> per-token representations -4. temporal: scatter to (B,T,d) -> gather overlap tokens -> project -> NT-Xent - instance: mean-pool -> project -> NT-Xent +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). @@ -54,11 +60,19 @@ class ContrastiveConfig(SSLConfig): # Temperature for NT-Xent temperature: float = 0.07 + # Use complementary (non-overlapping) masks for views + 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): @@ -154,9 +168,12 @@ def forward( # 1. Tokenize (shared between both views) tokens, padding_mask, token_info = self.encoder.tokenize(x, obs_mask) - # 2. Two independent random timestep masks + # 2. Create two views via timestep masks ssl_mask_1 = create_timestep_mask(B, T, self.config.mask_ratio, device) - ssl_mask_2 = create_timestep_mask(B, T, self.config.mask_ratio, device) + if self.config.complementary_masks: + ssl_mask_2 = ~ssl_mask_1 + else: + ssl_mask_2 = create_timestep_mask(B, T, self.config.mask_ratio, device) # 3. Encode both views vis_tokens_1, vis_padding_1 = extract_visible_timesteps(tokens, ssl_mask_1) diff --git a/tests/test_contrastive_objective.py b/tests/test_contrastive_objective.py index 1303bd2..65b32f1 100644 --- a/tests/test_contrastive_objective.py +++ b/tests/test_contrastive_objective.py @@ -471,6 +471,7 @@ def temporal_config(self): 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): @@ -599,6 +600,7 @@ def test_convergence_temporal(self): 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) @@ -638,6 +640,96 @@ def test_overlap_count_reasonable(self, encoder, temporal_config): 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 # ============================================================================= From 011d1f92d366a309c82781ea96f72171f954f33c Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 12 Mar 2026 12:58:52 -0400 Subject: [PATCH 026/121] feat: add Sprint 1c mask ratio sensitivity study and fix run counts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add Sprint 1c: mask ratio sensitivity (0.3, 0.5, 0.75) on MIMIC-IV, mortality_24h, seed=42 — validates shared mask_ratio=0.5 before committing to full experiment matrix - Add build_sprint1c() to runner script - Update Section 5.3 LR ablation to include 2e-4 (matching runner) - Fix Sprint 1b run count: 18 runs (9 pretrain + 9 finetune), not 12 - Update Sprint 8 to 60 runs (extra seeds for both LR and mask ratio) - Add sprint execution summary table to Section 6 - Update risk mitigation to reference both sensitivity studies --- scripts/run_experiments.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index 87da808..1701afc 100644 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -307,6 +307,15 @@ def build_sprint1b(self): pt = self._add_pretrain(sprint, p, ds, seed, extra) self._add_finetune(sprint, p, ds, seed, task, False, pt, name_extra=extra) + def build_sprint1c(self): + """Mask ratio sensitivity, MIMIC, mortality_24h, seed=42.""" + ds, seed, task, sprint = "miiv", 42, "mortality_24h", "1c" + for p in SSL_PARADIGMS: + for mr in MASK_RATIO_ABLATION: + extra = {"ssl.mask_ratio": mr} + pt = self._add_pretrain(sprint, p, ds, seed, extra) + self._add_finetune(sprint, p, ds, seed, task, False, pt, name_extra=extra) + def build_sprint2(self): """MIMIC Protocol A, seed=42 — reuses Sprint 1 pretrains.""" ds, seed, sprint = "miiv", 42, "2" @@ -409,6 +418,7 @@ def build_all(self) -> list[Run]: """Build full experiment matrix. Order matters for dedup.""" self.build_sprint1() self.build_sprint1b() + self.build_sprint1c() self.build_sprint2() self.build_sprint3() self.build_sprint4() From ffdee89037f7e48909849995c12583304559aa60 Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 12 Mar 2026 13:00:20 -0400 Subject: [PATCH 027/121] feat: add forward-looking prediction window for AKI labeling Refactor AKI label builder to use a proper forward-looking detection window (hours 48-72) instead of evaluating within the observation period. This prevents data leakage by ensuring the model never sees creatinine values used for label construction. Key changes: - AKI detection now uses [obs_window, obs_window + prediction_window) - Baseline window extended to full 48h observation period - RICU extraction horizon increased to 72h (LABEL_HORIZON_HOURS) - Stays with no creatinine in prediction window get null labels - prediction_window_hours is now required for AKI task - Tests updated to reflect forward-looking label semantics --- scripts/preprocessing/extract_with_ricu.R | 2 +- src/slices/constants.py | 5 +- src/slices/data/extractors/ricu.py | 3 +- src/slices/data/labels/aki.py | 46 ++++--- src/slices/data/tasks/aki_kdigo.yaml | 4 +- tests/test_aki_builder.py | 142 ++++++++++++++-------- 6 files changed, 129 insertions(+), 73 deletions(-) diff --git a/scripts/preprocessing/extract_with_ricu.R b/scripts/preprocessing/extract_with_ricu.R index 0e5432b..04130e2 100644 --- a/scripts/preprocessing/extract_with_ricu.R +++ b/scripts/preprocessing/extract_with_ricu.R @@ -43,7 +43,7 @@ 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, + make_option("--seq_length_hours", type = "integer", default = 72L, help = "Max hours per stay [default: %default]"), make_option("--raw_data_dir", type = "character", default = NULL, help = "Path to raw CSV files for ricu import (auto-detected if not set)") diff --git a/src/slices/constants.py b/src/slices/constants.py index 20e26b0..fbe193a 100644 --- a/src/slices/constants.py +++ b/src/slices/constants.py @@ -11,8 +11,11 @@ # ============================================================================= # Observation Window # ============================================================================= -SEQ_LENGTH_HOURS: int = 48 +SEQ_LENGTH_HOURS: int = 48 # Model input window MIN_STAY_HOURS: int = 48 +# RICU extraction horizon — longer than SEQ_LENGTH_HOURS to enable +# forward-looking labels (e.g. AKI prediction from hours 48-72). +LABEL_HORIZON_HOURS: int = 72 # ============================================================================= # Extraction diff --git a/src/slices/data/extractors/ricu.py b/src/slices/data/extractors/ricu.py index 9f00ff6..1c3f020 100644 --- a/src/slices/data/extractors/ricu.py +++ b/src/slices/data/extractors/ricu.py @@ -25,7 +25,7 @@ TextColumn, ) -from slices.constants import FEATURE_BLOCKLIST +from slices.constants import FEATURE_BLOCKLIST, LABEL_HORIZON_HOURS from slices.data.config_schemas import TimeSeriesConceptConfig from .base import BaseExtractor, ExtractorConfig @@ -308,6 +308,7 @@ def run(self) -> None: "feature_names": feature_names, "n_features": len(feature_names), "seq_length_hours": self.config.seq_length_hours, + "label_horizon_hours": LABEL_HORIZON_HOURS, "min_stay_hours": self.config.min_stay_hours, "task_names": task_names, "n_stays": len(stays_filtered), diff --git a/src/slices/data/labels/aki.py b/src/slices/data/labels/aki.py index 45e2098..d7b756f 100644 --- a/src/slices/data/labels/aki.py +++ b/src/slices/data/labels/aki.py @@ -17,9 +17,10 @@ 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 baseline_window_hours (default 24h). - Evaluates criteria AFTER the baseline window within the observation period. - Label = null if no creatinine measurements available. + Baseline: minimum creatinine in first baseline_window_hours (default 48h). + 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. """ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: @@ -44,11 +45,22 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: ) 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", 48) 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 @@ -75,13 +87,16 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: .agg(pl.col(crea_col).min().alias("baseline")) ) - # Post-baseline creatinine values (detection window) - post_obs = crea_ts.filter(pl.col("hour") >= baseline_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-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) --- + # 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) @@ -92,10 +107,9 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: # --- KDIGO Criterion 2: Absolute rise (>= 0.3 mg/dL within 48h window) --- # Use all creatinine data but only flag rises where the later measurement - # falls in the detection window (>= baseline_hours), so we catch rises - # that span the baseline/detection boundary. + # falls in the detection window [detection_start, detection_end). abs_aki = self._check_absolute_criterion_vectorized( - crea_ts, crea_col, abs_threshold, baseline_hours + crea_ts, crea_col, abs_threshold, detection_start, detection_end ) # Combine criteria: AKI = 1 if either criterion is met @@ -131,15 +145,14 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: ) # Build label: - # - null if no creatinine data or no baseline - # - 0 if no post-baseline data (no creatinine to evaluate after baseline period) + # - 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,7 +169,8 @@ 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 (no creatinine/baseline or no data in prediction window)" ) return result_df @@ -167,12 +181,13 @@ def _check_absolute_criterion_vectorized( 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 and the later measurement - falls in the detection window (hour >= detection_start_hour). + falls in the detection window [detection_start_hour, detection_end_hour). Returns: DataFrame with stay_id column for stays meeting the criterion. @@ -198,11 +213,12 @@ def _check_absolute_criterion_vectorized( ) # Join on stay_id and filter: j > i, within 48h, rise >= threshold, - # and the later measurement is in the detection window + # 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/tasks/aki_kdigo.yaml b/src/slices/data/tasks/aki_kdigo.yaml index 5edbe2d..b772186 100644 --- a/src/slices/data/tasks/aki_kdigo.yaml +++ b/src/slices/data/tasks/aki_kdigo.yaml @@ -1,14 +1,14 @@ task_name: aki_kdigo task_type: binary observation_window_hours: 48 -prediction_window_hours: null +prediction_window_hours: 24 # Look ahead 24h (hours 48-72) for AKI label gap_hours: 0 label_sources: - stays - timeseries label_params: creatinine_col: crea - baseline_window_hours: 24 + baseline_window_hours: 48 # 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 diff --git a/tests/test_aki_builder.py b/tests/test_aki_builder.py index 185c550..ec8f50a 100644 --- a/tests/test_aki_builder.py +++ b/tests/test_aki_builder.py @@ -8,15 +8,16 @@ 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 +38,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 +118,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 +135,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 +161,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 +187,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 +225,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 +255,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 +292,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.""" From f5c30aa95d76b53f794d08bc604d748bf15e536c Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 12 Mar 2026 15:51:33 -0400 Subject: [PATCH 028/121] feat: add --revision and --reason flags to experiment runner Allow rerunning sprint experiments with changed settings via a post-processing transform that rewrites run IDs, output dirs, and injects W&B tags. Includes validation guards and revision-aware status display. --- configs/finetune.yaml | 4 + configs/pretrain.yaml | 4 + configs/supervised.yaml | 4 + scripts/run_experiments.py | 142 +++++++++++++++++++++++++++++++---- src/slices/training/utils.py | 4 + 5 files changed, 142 insertions(+), 16 deletions(-) diff --git a/configs/finetune.yaml b/configs/finetune.yaml index 499e680..18d01cb 100644 --- a/configs/finetune.yaml +++ b/configs/finetune.yaml @@ -38,6 +38,10 @@ paradigm: mae # Experiment sprint (for W&B filtering, e.g. sprint=1) sprint: null +# Revision tracking (for rerunning sprints with changed settings) +revision: null +rerun_reason: null + # Random seed for reproducibility seed: 42 diff --git a/configs/pretrain.yaml b/configs/pretrain.yaml index 515d4b2..d935e3b 100644 --- a/configs/pretrain.yaml +++ b/configs/pretrain.yaml @@ -25,6 +25,10 @@ dataset: miiv # miiv | eicu | combined # Experiment sprint (for W&B filtering, e.g. sprint=1) sprint: null +# Revision tracking (for rerunning sprints with changed settings) +revision: null +rerun_reason: null + # Random seed for reproducibility seed: 42 diff --git a/configs/supervised.yaml b/configs/supervised.yaml index 5af387d..eceb732 100644 --- a/configs/supervised.yaml +++ b/configs/supervised.yaml @@ -36,6 +36,10 @@ paradigm: supervised # Experiment sprint (for W&B filtering, e.g. sprint=1) sprint: null +# Revision tracking (for rerunning sprints with changed settings) +revision: null +rerun_reason: null + # Random seed for reproducibility seed: 42 diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index 1701afc..9809cbd 100644 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -9,15 +9,18 @@ uv run python scripts/run_experiments.py warmup --sprint 1 uv run python scripts/run_experiments.py run --sprint 1 --parallel 4 uv run python scripts/run_experiments.py run --sprint 1 2 3 --parallel 6 --dry-run + uv run python scripts/run_experiments.py run --sprint 1 --revision v2 --reason "fix LR" uv run python scripts/run_experiments.py status uv run python scripts/run_experiments.py status --sprint 1 uv run python scripts/run_experiments.py retry --failed --parallel 4 + uv run python scripts/run_experiments.py retry --failed --sprint 1 --revision v2 --parallel 4 """ from __future__ import annotations import argparse import json import os +import re import signal import subprocess import sys @@ -434,6 +437,44 @@ def generate_all_runs() -> list[Run]: return builder.build_all() +def apply_revision(runs: list[Run], revision: str, reason: str | None = None) -> list[Run]: + """Post-process runs to inject revision into IDs, output dirs, and overrides. + + Rewrites: + - Run IDs: s1_pretrain_... → s1_rev-v2_pretrain_... + - Output dirs: outputs/sprint1/... → outputs/sprint1_rev-v2/... + - depends_on: remap references within the revised set + - extra_overrides: inject revision= and rerun_reason= + """ + old_to_new: dict[str, str] = {} + + for r in runs: + old_id = r.id + # Insert rev- after the sprint prefix (e.g. s1_ → s1_rev-v2_) + prefix = f"s{r.sprint}_" + if old_id.startswith(prefix): + new_id = f"{prefix}rev-{revision}_{old_id[len(prefix):]}" + else: + new_id = f"rev-{revision}_{old_id}" + old_to_new[old_id] = new_id + + r.id = new_id + # Rewrite output dir: sprint1/ → sprint1_rev-v2/ + r.output_dir = r.output_dir.replace( + f"outputs/sprint{r.sprint}/", + f"outputs/sprint{r.sprint}_rev-{revision}/", + ) + r.extra_overrides["revision"] = revision + if reason: + r.extra_overrides["rerun_reason"] = reason + + # Remap depends_on references within the revised set + for r in runs: + r.depends_on = [old_to_new.get(d, d) for d in r.depends_on] + + return runs + + # --------------------------------------------------------------------------- # State Management # --------------------------------------------------------------------------- @@ -676,54 +717,85 @@ def _print_dry_run(runs: list[Run], runs_by_id: dict[str, Run]): # --------------------------------------------------------------------------- # Status / Retry # --------------------------------------------------------------------------- +def _extract_revision_from_id(run_id: str) -> str | None: + """Extract revision name from a run ID (e.g. 's1_rev-v2_pretrain_...' → 'v2').""" + m = re.search(r"_rev-([^_]+)_", run_id) + return m.group(1) if m else None + + +def _sprint_display(sprint: str, revision: str | None) -> str: + """Format sprint column, e.g. '1' or '1/v2'.""" + if revision: + return f"{sprint}/{revision}" + return sprint + + +def _extract_sprint_from_id(run_id: str) -> str | None: + """Extract sprint from a run ID (e.g. 's1_rev-v2_pretrain_...' → '1').""" + m = re.match(r"s([^_]+)_", run_id) + return m.group(1) if m else None + + def print_status(sprint_filter: list[str] | None = None): all_runs = generate_all_runs() state = load_state() + generated_ids = {r.id for r in all_runs} - # Group by sprint - sprints: dict[str, list[Run]] = {} + # Group generated runs by (sprint, revision=None) + groups: dict[tuple[str, str | None], list[str]] = {} for r in all_runs: - sprints.setdefault(r.sprint, []).append(r) - - sprint_keys = sorted(sprints.keys(), key=_sprint_sort_key) + key = (r.sprint, None) + groups.setdefault(key, []).append(r.id) + + # Add revised runs from state (not in generated set) by parsing their IDs + for run_id in state.get("runs", {}): + if run_id not in generated_ids and "_rev-" in run_id: + sprint = _extract_sprint_from_id(run_id) + rev = _extract_revision_from_id(run_id) + if sprint and rev: + key = (sprint, rev) + groups.setdefault(key, []).append(run_id) + + group_keys = sorted(groups.keys(), key=lambda k: (_sprint_sort_key(k[0]), k[1] or "")) if sprint_filter: - sprint_keys = [s for s in sprint_keys if s in sprint_filter] + group_keys = [k for k in group_keys if k[0] in sprint_filter] print( - f"{'Sprint':>6} | {'Total':>5} | {'Done':>4} | {'Run':>3} | " + f"{'Sprint':>8} | {'Total':>5} | {'Done':>4} | {'Run':>3} | " f"{'Fail':>4} | {'Pend':>4} | {'Skip':>4}" ) - print("-" * 52) + print("-" * 54) totals = {"total": 0, "completed": 0, "running": 0, "failed": 0, "pending": 0, "skipped": 0} - for s in sprint_keys: - runs = sprints[s] + for sprint, rev in group_keys: + run_ids = groups[(sprint, rev)] + label = _sprint_display(sprint, rev) counts = { - "total": len(runs), + "total": len(run_ids), "completed": 0, "running": 0, "failed": 0, "pending": 0, "skipped": 0, } - for r in runs: - status = get_run_status(state, r.id) + for rid in run_ids: + status = get_run_status(state, rid) if status in counts: counts[status] += 1 else: counts["pending"] += 1 # unknown → pending print( - f"{s:>6} | {counts['total']:>5} | {counts['completed']:>4} | " + f"{label:>8} | {counts['total']:>5} | {counts['completed']:>4} | " f"{counts['running']:>3} | {counts['failed']:>4} | " f"{counts['pending']:>4} | {counts['skipped']:>4}" ) for k in totals: totals[k] += counts[k] - print("-" * 52) + print("-" * 54) print( - f"{'TOTAL':>6} | {totals['total']:>5} | {totals['completed']:>4} | " + f"{'TOTAL':>8} | {totals['total']:>5} | {totals['completed']:>4} | " f"{totals['running']:>3} | {totals['failed']:>4} | " f"{totals['pending']:>4} | {totals['skipped']:>4}" ) @@ -785,7 +857,13 @@ def cmd_run(args): print(f"No runs found for sprint(s): {', '.join(sprints)}") return + # Apply revision transform if requested + if args.revision: + runs = apply_revision(runs, args.revision, args.reason) + print(f"Sprint(s) {', '.join(sprints)}: {len(runs)} runs") + if args.revision: + print(f"Revision: {args.revision}" + (f" ({args.reason})" if args.reason else "")) run_scheduler(runs, state, args.parallel, args.dry_run) @@ -796,6 +874,18 @@ def cmd_status(args): def cmd_retry(args): all_runs = generate_all_runs() + + # Apply revision if specified (must happen before ID matching against state) + if args.revision: + sprint_filter = [str(s) for s in args.sprint] if args.sprint else None + if sprint_filter: + revised = [r for r in all_runs if r.sprint in sprint_filter] + rest = [r for r in all_runs if r.sprint not in sprint_filter] + revised = apply_revision(revised, args.revision, args.reason) + all_runs = rest + revised + else: + all_runs = apply_revision(all_runs, args.revision, args.reason) + state = load_state() recover_stale_running(state) @@ -837,6 +927,9 @@ def cmd_warmup(args): experiment runs can load preprocessed data from disk instead of each process independently loading and converting raw parquet files (which causes OOM). """ + if args.revision: + print(f"Note: --revision={args.revision} ignored (warmup is revision-independent)") + all_runs = generate_all_runs() # Filter to requested sprints @@ -905,6 +998,8 @@ def main(): p_run.add_argument("--sprint", nargs="+", required=True, help="Sprint(s) to run (e.g. 1 1b 2)") p_run.add_argument("--parallel", type=int, default=4, help="Max parallel jobs (default: 4)") 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 (e.g. v2)") + p_run.add_argument("--reason", type=str, default=None, help="Reason for rerun") # status p_status = sub.add_parser("status", help="Show experiment status") @@ -915,6 +1010,9 @@ def main(): p_retry.add_argument("--failed", action="store_true", help="Retry failed runs") p_retry.add_argument("--skipped", action="store_true", help="Retry skipped runs") p_retry.add_argument("--parallel", type=int, default=4, help="Max parallel jobs (default: 4)") + p_retry.add_argument("--sprint", nargs="+", default=None, help="Scope retry to sprint(s)") + 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") # warmup p_warmup = sub.add_parser( @@ -923,9 +1021,21 @@ def main(): p_warmup.add_argument( "--sprint", nargs="+", required=True, help="Sprint(s) to warmup (e.g. 1 1b 2)" ) + 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() + # Validate --reason requires --revision + if getattr(args, "reason", None) and not getattr(args, "revision", None): + parser.error("--reason requires --revision") + + # Validate retry --revision requires --sprint + if args.command == "retry" and args.revision and not args.sprint: + parser.error("--revision requires --sprint to scope which sprints to revise") + if args.command == "run": cmd_run(args) elif args.command == "status": diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index e7175a7..5ee67eb 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -242,6 +242,10 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: tags = list(cfg.logging.get("wandb_tags", [])) if cfg.get("sprint") is not None: tags.append(f"sprint:{cfg.sprint}") + if cfg.get("revision") is not None: + tags.append(f"revision:{cfg.revision}") + if cfg.get("rerun_reason") is not None: + tags.append(f"rerun-reason:{cfg.rerun_reason}") tags = tags or None logger = WandbLogger( From 82604fceeae3065f450081b82f08356dfe2ae514 Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 12 Mar 2026 15:55:30 -0400 Subject: [PATCH 029/121] feat: improve wandb tagging and add run query script Add query_runs.py to wandb-analysis skill for filtering runs by tags, config, state, and group with aggregation support. Add sprint prefix to run names, protocol/mask_ratio/label_fraction/lr tags, and dynamic probe vs finetune run name distinction. --- configs/finetune.yaml | 3 ++- configs/pretrain.yaml | 3 ++- configs/supervised.yaml | 3 ++- src/slices/training/utils.py | 24 +++++++++++++++++++++++- 4 files changed, 29 insertions(+), 4 deletions(-) diff --git a/configs/finetune.yaml b/configs/finetune.yaml index 18d01cb..5dd38bb 100644 --- a/configs/finetune.yaml +++ b/configs/finetune.yaml @@ -136,7 +136,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: s${sprint}_finetune_${dataset}_${paradigm}_${task.task_name}_seed${seed} wandb_group: finetune_${dataset}_${paradigm}_${task.task_name} wandb_tags: - "phase:finetune" @@ -144,6 +144,7 @@ logging: - "paradigm:${paradigm}" - "task:${task.task_name}" - "seed:${seed}" + - "lr:${optimizer.lr}" log_every_n_steps: 10 # Hydra configuration diff --git a/configs/pretrain.yaml b/configs/pretrain.yaml index d935e3b..dc8ed4e 100644 --- a/configs/pretrain.yaml +++ b/configs/pretrain.yaml @@ -94,13 +94,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: s${sprint}_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/supervised.yaml b/configs/supervised.yaml index eceb732..e111190 100644 --- a/configs/supervised.yaml +++ b/configs/supervised.yaml @@ -124,7 +124,7 @@ logging: use_wandb: true wandb_project: ${project_name} wandb_entity: null - run_name: supervised_${dataset}_${task.task_name}_seed${seed} + run_name: s${sprint}_supervised_${dataset}_${task.task_name}_seed${seed} wandb_group: supervised_${dataset}_${task.task_name} wandb_tags: - "phase:supervised" @@ -132,6 +132,7 @@ logging: - "paradigm:supervised" - "task:${task.task_name}" - "seed:${seed}" + - "lr:${optimizer.lr}" log_every_n_steps: 10 # Hydra configuration diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 5ee67eb..f9e6378 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -246,12 +246,34 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: tags.append(f"revision:{cfg.revision}") if cfg.get("rerun_reason") is not None: tags.append(f"rerun-reason:{cfg.rerun_reason}") + + # Add protocol tag for finetune runs (Protocol A = frozen, Protocol B = unfrozen) + freeze_encoder = cfg.get("training", {}).get("freeze_encoder") + if freeze_encoder is not None: + protocol = "A" if freeze_encoder else "B" + tags.append(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: + tags.append(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: + tags.append(f"label_fraction:{label_fraction}") + 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) + logger = WandbLogger( project=cfg.logging.wandb_project, entity=cfg.logging.get("wandb_entity", None), - name=cfg.logging.get("run_name", None), + name=run_name, group=cfg.logging.get("wandb_group", None), tags=tags, save_dir=cfg.output_dir, From 726d74a3b6eab909dbe40e14868b241e8dc5ce82 Mon Sep 17 00:00:00 2001 From: hill Date: Fri, 13 Mar 2026 13:24:02 -0400 Subject: [PATCH 030/121] feat: import only essential RICU tables to speed up extraction Skip unused tables (emar_detail, pharmacy, poe, prescriptions, etc.) that take 20+ min to import but aren't needed for concept extraction. Also remove unnecessary intermediate gc() calls in batch merging. --- scripts/preprocessing/extract_with_ricu.R | 46 +++++++++++++++++++---- 1 file changed, 38 insertions(+), 8 deletions(-) diff --git a/scripts/preprocessing/extract_with_ricu.R b/scripts/preprocessing/extract_with_ricu.R index 04130e2..501b2ba 100644 --- a/scripts/preprocessing/extract_with_ricu.R +++ b/scripts/preprocessing/extract_with_ricu.R @@ -57,6 +57,28 @@ VALID_DATASETS <- c("miiv", "eicu", "hirid", "aumc", "mimic", "sic", # Batch size for concept loading — controls peak memory during extraction. CONCEPT_BATCH_SIZE <- 8L +# 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 = c("patient", "vitalAperiodic", "vitalPeriodic", "lab", + "infusionDrug", "respiratoryCharting", "nurseCharting", + "intakeOutput", "diagnosis", "medication"), + hirid = c("observations", "pharma", "general"), + aumc = c("numericitems", "drugitems", "procedureorderitems", + "admissions", "listitems") +) + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -192,8 +214,7 @@ extract_timeseries <- function(dataset, concepts, seq_length_hours, dict) { for (b in 2:n_batches) { merged <- merge(merged, batch_results[[b]], by = c("stay_id", "hour"), all = TRUE) - batch_results[[b]] <- NULL - gc() + batch_results[[b]] <- NULL # allow GC to reclaim } } rm(batch_results, wins) @@ -857,12 +878,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) { From f5528cd7fbcfe5443137b97a3e6b4b7b4bd1908a Mon Sep 17 00:00:00 2001 From: hill Date: Fri, 13 Mar 2026 15:33:40 -0400 Subject: [PATCH 031/121] feat: add sprint baseline tagging and update experiment plan with results MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `tag` subcommand to run_experiments.py for propagating sprint tags to inherited baseline runs in W&B, enabling single-query filtering of all runs relevant to a sprint. Update EXPERIMENT_PLAN.md with Sprint 1–2 outcomes, rev2 methodological corrections (AKI forward-looking window, complementary contrastive masks), and baseline inheritance documentation. --- scripts/run_experiments.py | 112 +++++++++++++++++++++++++++++++++++++ 1 file changed, 112 insertions(+) diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index 9809cbd..4f7e1fb 100644 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -14,6 +14,8 @@ uv run python scripts/run_experiments.py status --sprint 1 uv run python scripts/run_experiments.py retry --failed --parallel 4 uv run python scripts/run_experiments.py retry --failed --sprint 1 --revision v2 --parallel 4 + uv run python scripts/run_experiments.py tag --sprint 2 --dry-run + uv run python scripts/run_experiments.py tag --sprint 2 3 5 """ from __future__ import annotations @@ -43,6 +45,24 @@ MASK_RATIO_ABLATION = [0.3, 0.75] # 0.5 reused from Phase 1 TRANSFER_PAIRS = [("miiv", "eicu"), ("eicu", "miiv")] +# Baseline inheritance: which earlier sprints' runs should be tagged as baselines +# for later sprints. See docs/EXPERIMENT_PLAN.md § "Baseline Inheritance Across Sprints". +# Intentionally coarse-grained: ALL runs from source sprints are tagged, even when +# only a subset is relevant (e.g. 1b only needs mortality_24h from Sprint 1). Use +# secondary W&B tags (task:, protocol:, phase:, etc.) to filter down when querying. +BASELINE_SPRINTS: dict[str, list[str]] = { + "1": [], + "1b": ["1"], + "1c": ["1"], + "2": ["1"], + "3": [], + "4": [], + "5": ["1", "2", "3", "4"], + "6": ["1", "2", "3", "4", "5"], + "7": ["1", "3", "5"], + "8": ["1", "1b", "1c", "5"], +} + STATE_FILE = Path("outputs/experiment_state.json") LOG_DIR = Path("logs/runner") @@ -989,6 +1009,77 @@ def cmd_warmup(args): print("You can now run experiments in parallel without OOM.") +def cmd_tag(args): + """Add sprint tags to inherited baseline runs in W&B. + + For each requested sprint, finds runs from inherited baseline sprints + (per BASELINE_SPRINTS mapping) and adds the target sprint tag via the + W&B API. Idempotent — already-tagged runs are skipped. + + Usage: + uv run python scripts/run_experiments.py tag --sprint 2 + uv run python scripts/run_experiments.py tag --sprint 2 3 --dry-run + uv run python scripts/run_experiments.py tag --sprint 2 --project slices --entity myteam + """ + try: + import wandb # noqa: F811 + except ImportError: + print("Error: wandb is required for tagging. Install with: pip install wandb") + sys.exit(1) + + api = wandb.Api() + + project = args.project + entity = args.entity + # Resolve entity: CLI flag > WANDB_ENTITY env var > wandb default entity + if entity is None: + entity = os.environ.get("WANDB_ENTITY") or api.default_entity + if entity is None: + print("Error: Could not determine W&B entity. Set --entity or WANDB_ENTITY env var.") + sys.exit(1) + + sprints = [str(s) for s in args.sprint] + dry_run = args.dry_run + + for target_sprint in sprints: + source_sprints = BASELINE_SPRINTS.get(target_sprint, []) + if not source_sprints: + print(f"Sprint {target_sprint}: no baseline sprints to inherit from — skipping.") + continue + + target_tag = f"sprint:{target_sprint}" + sources = ", ".join(source_sprints) + print(f"\nSprint {target_sprint}: inheriting baselines from sprint(s) {sources}") + + # Query runs from each source sprint + tagged = 0 + skipped = 0 + for source_sprint in source_sprints: + source_tag = f"sprint:{source_sprint}" + runs = api.runs( + f"{entity}/{project}", + filters={"tags": {"$in": [source_tag]}}, + ) + + for run in runs: + if target_tag in run.tags: + skipped += 1 + continue + + if dry_run: + print(f" [dry-run] Would tag: {run.name} ({run.id}) with {target_tag}") + tagged += 1 + else: + run.tags = run.tags + [target_tag] + run.update() + tagged += 1 + + action = "would tag" if dry_run else "tagged" + print(f" {action} {tagged} runs, skipped {skipped} (already tagged)") + + print("\nDone.") + + def main(): parser = argparse.ArgumentParser(description="SLICES parallel experiment runner") sub = parser.add_subparsers(dest="command", required=True) @@ -1026,6 +1117,25 @@ def main(): ) p_warmup.add_argument("--reason", type=str, default=None, help="Ignored for warmup") + # tag + p_tag = sub.add_parser("tag", help="Tag inherited baseline runs in W&B with target sprint tags") + p_tag.add_argument( + "--sprint", + nargs="+", + required=True, + help="Sprint(s) whose baselines to tag (e.g. 2 or 2 3 5)", + ) + p_tag.add_argument("--dry-run", action="store_true", help="Show what would be tagged") + p_tag.add_argument( + "--project", type=str, default="slices", help="W&B project name (default: slices)" + ) + p_tag.add_argument( + "--entity", + type=str, + default=None, + help="W&B entity (default: WANDB_ENTITY env var or wandb default)", + ) + args = parser.parse_args() # Validate --reason requires --revision @@ -1044,6 +1154,8 @@ def main(): cmd_retry(args) elif args.command == "warmup": cmd_warmup(args) + elif args.command == "tag": + cmd_tag(args) if __name__ == "__main__": From db005fa4941485131e246940c13846061f7c2115 Mon Sep 17 00:00:00 2001 From: Hannes Ill Date: Fri, 13 Mar 2026 19:48:25 +0000 Subject: [PATCH 032/121] fix: truncate long rerun-reason wandb tags to 64 characters --- src/slices/training/utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index f9e6378..9788170 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -245,7 +245,10 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: if cfg.get("revision") is not None: tags.append(f"revision:{cfg.revision}") if cfg.get("rerun_reason") is not None: - tags.append(f"rerun-reason:{cfg.rerun_reason}") + tag = f"rerun-reason:{cfg.rerun_reason}" + if len(tag) > 64: + tag = tag[:61] + "..." + tags.append(tag) # Add protocol tag for finetune runs (Protocol A = frozen, Protocol B = unfrozen) freeze_encoder = cfg.get("training", {}).get("freeze_encoder") From 942c45ddcf62c32395e3f18b4b21f9bba060bf81 Mon Sep 17 00:00:00 2001 From: Hannes Ill Date: Wed, 18 Mar 2026 00:05:37 +0000 Subject: [PATCH 033/121] fix: include run_experiments and prepare_dataset in auto-shutdown activity check MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pgrep pattern only matched training scripts, so warmup and experiment orchestration were treated as idle — causing the VM to shut down mid-run after 59 minutes. --- scripts/auto_shutdown.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/auto_shutdown.sh b/scripts/auto_shutdown.sh index dffcb6a..69d6ac5 100755 --- a/scripts/auto_shutdown.sh +++ b/scripts/auto_shutdown.sh @@ -9,7 +9,7 @@ IDLE_THRESHOLD_MIN=59 STAMP_FILE="/tmp/slices_last_training_activity" # Check if any training process is running -if pgrep -f "scripts/training/(pretrain|finetune|supervised)\.py" > /dev/null 2>&1; then +if pgrep -f "scripts/(training/(pretrain|finetune|supervised)|run_experiments|preprocessing/prepare_dataset)\.py" > /dev/null 2>&1; then # Training active — update timestamp and exit date +%s > "$STAMP_FILE" exit 0 From 32d426644ccb0c5b02a5fca2d878f364171da28c Mon Sep 17 00:00:00 2001 From: Hannes Ill Date: Wed, 18 Mar 2026 01:25:33 +0000 Subject: [PATCH 034/121] fix: reduce peak memory in tensor preprocessing pipeline Pre-allocate output arrays instead of batch-and-stack to avoid 2x peak memory. Free raw DataFrames and Python lists earlier in the pipeline. --- src/slices/data/dataset.py | 13 +++-- src/slices/data/tensor_preprocessing.py | 68 +++++++++---------------- 2 files changed, 32 insertions(+), 49 deletions(-) diff --git a/src/slices/data/dataset.py b/src/slices/data/dataset.py index 3b805bd..74104ad 100644 --- a/src/slices/data/dataset.py +++ b/src/slices/data/dataset.py @@ -248,11 +248,17 @@ def _load_data(self) -> None: self.feature_means = cached_tensors["feature_means"] self.feature_stds = cached_tensors["feature_stds"] logger.info("Using cached preprocessed tensors") + # Free raw DataFrame — not needed when using cached tensors + del self.timeseries_df else: # Pre-extract raw arrays raw_timeseries = self.timeseries_df["timeseries"].to_list() raw_masks = self.timeseries_df["mask"].to_list() + # Free raw DataFrame immediately to reduce peak memory. + # get_preprocessing_stages() reloads from parquet on demand. + del self.timeseries_df + # Try to load existing normalization stats (for reproducibility) cached_stats = load_normalization_stats( self.data_dir, self.train_indices, self.normalize @@ -275,10 +281,6 @@ def _load_data(self) -> None: self._excluded_stay_ids, ) - # Free raw DataFrame — only _timeseries_tensor and _mask_tensor needed from here. - # get_preprocessing_stages() reloads from parquet on demand. - del self.timeseries_df - # Pre-compute labels and static features self._precompute_labels_and_static() @@ -308,6 +310,9 @@ def _precompute_tensors( raw_timeseries, raw_masks, self.seq_length, self.n_features ) + # Free raw Python lists to reduce peak memory + del raw_timeseries, raw_masks + # Step 2: Compute normalization statistics self.feature_means = torch.zeros(self.n_features) self.feature_stds = torch.ones(self.n_features) diff --git a/src/slices/data/tensor_preprocessing.py b/src/slices/data/tensor_preprocessing.py index 8151541..816e0df 100644 --- a/src/slices/data/tensor_preprocessing.py +++ b/src/slices/data/tensor_preprocessing.py @@ -43,54 +43,32 @@ def convert_raw_to_tensors( n_samples = len(raw_timeseries) logger.debug("[1/3] Converting to tensors...") - 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 >= seq_length: - # Truncate - ts_arr = np.array(ts_data[:seq_length], dtype=np.float32) - mask_arr = np.array(mask_data[:seq_length], dtype=bool) - else: - # Pad with NaN / False - ts_arr = np.full((seq_length, n_features), np.nan, dtype=np.float32) - mask_arr = np.zeros((seq_length, 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 + # 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) - # Free numpy arrays - del timeseries_np, masks_np, all_timeseries, all_masks - return timeseries_tensor, masks_tensor From 3da6cfc4a077f668efd90245688fc5601642bb14 Mon Sep 17 00:00:00 2001 From: Hannes Ill Date: Wed, 18 Mar 2026 21:34:13 +0000 Subject: [PATCH 035/121] fix: avoid OOM in warmup and tensor extraction for large datasets Warmup now runs each cache build in a subprocess so Python's allocator releases memory between iterations. Tensor extraction uses Polars explode() for zero-copy conversion instead of to_list(), avoiding ~3-4x memory overhead from Python object creation. --- scripts/run_experiments.py | 48 ++++++++++------ src/slices/data/dataset.py | 45 +++++++++------ src/slices/data/tensor_preprocessing.py | 75 +++++++++++++++++++++++++ 3 files changed, 135 insertions(+), 33 deletions(-) diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index 4f7e1fb..a716002 100644 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -980,8 +980,8 @@ def cmd_warmup(args): print(f"Warming up tensor caches for sprint(s) {', '.join(sprints)}") print(f" {len(combos)} unique (dataset, task, seed, label_fraction) combinations\n") - # Import here to avoid loading heavy deps when not needed - from slices.data.datamodule import ICUDataModule + import subprocess + import sys for i, (dataset, task, seed, label_fraction) in enumerate(combos, 1): processed_dir = f"data/processed/{dataset}" @@ -989,21 +989,35 @@ def cmd_warmup(args): frac_str = f", frac={label_fraction}" if label_fraction < 1.0 else "" print(f"[{i}/{len(combos)}] {dataset} / {task_str} / seed={seed}{frac_str}") - try: - dm = ICUDataModule( - processed_dir=processed_dir, - task_name=task, - batch_size=1, # doesn't matter, we just need setup() - num_workers=0, - seed=seed, - label_fraction=label_fraction, - ) - dm.setup() - print(f" -> Cached ({len(dm.dataset):,} samples)") - # Free memory before next combo - del dm - except Exception as e: - print(f" -> ERROR: {e}") + # Run each warmup in a subprocess so that memory is fully returned to + # the OS between iterations. Python's allocator retains freed memory, + # so in-process iteration causes cumulative RSS growth and OOM on the + # combined dataset (~23 GB peak per run). + task_arg = task or "" + result = subprocess.run( + [ + sys.executable, "-c", + "from slices.data.datamodule import ICUDataModule; " + f"dm = ICUDataModule(" + f" processed_dir={processed_dir!r}," + f" task_name={task!r}," + f" batch_size=1," + f" num_workers=0," + f" seed={seed}," + f" label_fraction={label_fraction}," + f"); " + "dm.setup(); " + f"print(len(dm.dataset))", + ], + 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)}") print("\nWarmup complete. Tensor caches saved to data/processed//.tensor_cache/") print("You can now run experiments in parallel without OOM.") diff --git a/src/slices/data/dataset.py b/src/slices/data/dataset.py index 74104ad..8bee404 100644 --- a/src/slices/data/dataset.py +++ b/src/slices/data/dataset.py @@ -26,6 +26,7 @@ compute_normalization_stats, compute_single_sample_stages, convert_raw_to_tensors, + extract_tensors_from_dataframe, ) # Module-level logger @@ -251,9 +252,12 @@ def _load_data(self) -> None: # Free raw DataFrame — not needed when using cached tensors del self.timeseries_df else: - # Pre-extract raw arrays - raw_timeseries = self.timeseries_df["timeseries"].to_list() - raw_masks = self.timeseries_df["mask"].to_list() + # Extract tensors directly from Polars DataFrame. + # Uses explode() for zero-copy extraction instead of to_list(), + # which creates Python objects and can OOM on large datasets. + timeseries_tensor, masks_tensor = extract_tensors_from_dataframe( + self.timeseries_df, self.seq_length, self.n_features + ) # Free raw DataFrame immediately to reduce peak memory. # get_preprocessing_stages() reloads from parquet on demand. @@ -264,8 +268,12 @@ def _load_data(self) -> None: self.data_dir, self.train_indices, self.normalize ) - # Pre-compute tensors for all samples (much faster __getitem__) - self._precompute_tensors(raw_timeseries, raw_masks, self.train_indices, cached_stats) + # Apply normalization and imputation (steps 2 and 3) + self._precompute_tensors( + precomputed_tensors=(timeseries_tensor, masks_tensor), + train_indices=self.train_indices, + cached_stats=cached_stats, + ) # Save preprocessed tensors to cache for next run save_cached_tensors( @@ -286,10 +294,11 @@ def _load_data(self) -> None: 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__. @@ -298,20 +307,24 @@ def _precompute_tensors( 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). + 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: Convert nested lists to tensors - timeseries_tensor, masks_tensor = convert_raw_to_tensors( - raw_timeseries, raw_masks, self.seq_length, self.n_features - ) + 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 - # Free raw Python lists to reduce peak memory - del raw_timeseries, raw_masks + n_samples = timeseries_tensor.shape[0] + logger.info(f"Preprocessing {n_samples:,} samples...") # Step 2: Compute normalization statistics self.feature_means = torch.zeros(self.n_features) diff --git a/src/slices/data/tensor_preprocessing.py b/src/slices/data/tensor_preprocessing.py index 816e0df..4591246 100644 --- a/src/slices/data/tensor_preprocessing.py +++ b/src/slices/data/tensor_preprocessing.py @@ -18,6 +18,78 @@ 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) + + # 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]]], @@ -30,6 +102,9 @@ def convert_raw_to_tensors( 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). From ad63aa3ca00eaf2155c69dce5ad1d6acb8c4a176 Mon Sep 17 00:00:00 2001 From: Hannes Ill Date: Wed, 18 Mar 2026 21:34:16 +0000 Subject: [PATCH 036/121] fix: derive eICU date_of_death from hospitaldischargeoffset eICU extraction set date_of_death=NA for all patients, causing windowed mortality tasks (mortality_24h) to produce 0 positive labels. For expired patients, approximate death time from hospitaldischargeoffset using the same synthetic epoch as stays extraction. --- scripts/preprocessing/extract_with_ricu.R | 26 +++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/scripts/preprocessing/extract_with_ricu.R b/scripts/preprocessing/extract_with_ricu.R index 501b2ba..df71c70 100644 --- a/scripts/preprocessing/extract_with_ricu.R +++ b/scripts/preprocessing/extract_with_ricu.R @@ -605,11 +605,33 @@ extract_mortality_eicu <- function(dataset, dict) { NA_integer_ } + # Derive date_of_death from hospitaldischargeoffset for expired patients. + # eICU does not store death timestamps directly, but hospitaldischargeoffset + # gives minutes from unit admission to hospital discharge. For patients who + # expired (hospitaldischargestatus == "Expired"), this is the approximate + # time of death. 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") + has_offset <- "hospitaldischargeoffset" %in% names(pat) + expired <- !is.na(hosp_flag) & hosp_flag == 1L + + dod <- as.POSIXct(rep(NA, nrow(pat)), tz = "UTC") + if (has_offset) { + dod[expired] <- epoch + pat$hospitaldischargeoffset[expired] * 60 + } + + # Also derive dischtime from hospitaldischargeoffset for all patients + dischtime <- as.POSIXct(rep(NA, nrow(pat)), tz = "UTC") + if (has_offset) { + dischtime <- epoch + pat$hospitaldischargeoffset * 60 + } + df <- data.frame( stay_id = pat[[id_col]], - date_of_death = NA, + date_of_death = dod, hospital_expire_flag = hosp_flag, - dischtime = NA, + dischtime = dischtime, discharge_location = if ("unitdischargelocation" %in% names(pat)) { pat$unitdischargelocation } else { From c2572edc0619711e51eb5df513e8f92e45fcfdc9 Mon Sep 17 00:00:00 2001 From: Hannes Ill Date: Wed, 18 Mar 2026 21:34:17 +0000 Subject: [PATCH 037/121] fix: detect interactive sessions in auto-shutdown activity check --- scripts/auto_shutdown.sh | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/scripts/auto_shutdown.sh b/scripts/auto_shutdown.sh index 69d6ac5..8608bec 100755 --- a/scripts/auto_shutdown.sh +++ b/scripts/auto_shutdown.sh @@ -8,9 +8,10 @@ set -euo pipefail IDLE_THRESHOLD_MIN=59 STAMP_FILE="/tmp/slices_last_training_activity" -# Check if any training process is running -if pgrep -f "scripts/(training/(pretrain|finetune|supervised)|run_experiments|preprocessing/prepare_dataset)\.py" > /dev/null 2>&1; then - # Training active — update timestamp and exit +# Check if any training process or Claude Code session is running +if pgrep -f "scripts/(training/(pretrain|finetune|supervised)|run_experiments|preprocessing/prepare_dataset)\.py" > /dev/null 2>&1 \ + || pgrep -f "claude" > /dev/null 2>&1; then + # Activity detected — update timestamp and exit date +%s > "$STAMP_FILE" exit 0 fi From bd8119ee3030db18d7f1c9090ff29d9349567cf0 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 17 Mar 2026 00:13:16 -0400 Subject: [PATCH 038/121] fix: improve eICU extraction robustness and compatibility - Add per-concept retry fallback when batch loading fails in RICU extraction, skipping unsupported concepts instead of crashing - Remove hardcoded eICU essential tables list (eICU needs all tables due to load_dictionary() check) - Reduce concept batch size from 8 to 4 for lower peak memory - Add dur_var to feature blocklist (internal ricu helper, not clinical) - Handle all-null date_of_death column dtype in mortality label builder (eICU has no date_of_death, causing Null dtype and downstream errors) --- scripts/preprocessing/extract_with_ricu.R | 129 +++++++++++++++------- src/slices/constants.py | 9 +- src/slices/data/labels/mortality.py | 15 ++- 3 files changed, 106 insertions(+), 47 deletions(-) diff --git a/scripts/preprocessing/extract_with_ricu.R b/scripts/preprocessing/extract_with_ricu.R index df71c70..a6e18ac 100644 --- a/scripts/preprocessing/extract_with_ricu.R +++ b/scripts/preprocessing/extract_with_ricu.R @@ -55,7 +55,7 @@ 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, @@ -71,9 +71,9 @@ ESSENTIAL_TABLES <- list( "procedureevents", "ingredientevents", "datetimeevents", "icustays", "patients", "admissions", "transfers", "d_labitems", "diagnoses_icd", "d_icd_diagnoses"), - eicu = c("patient", "vitalAperiodic", "vitalPeriodic", "lab", - "infusionDrug", "respiratoryCharting", "nurseCharting", - "intakeOutput", "diagnosis", "medication"), + # 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") @@ -161,60 +161,105 @@ extract_timeseries <- function(dataset, concepts, seq_length_hours, dict) { 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 + }) + } + + skipped_concepts <- character(0) + for (b in seq_len(n_batches)) { batch_concepts <- concepts[batch_indices[[b]]] message(sprintf(" Batch %d/%d (%d concepts): %s", 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)] - - # 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]])) + 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) } - # Convert to data.table for efficient merging - dt <- as.data.table(ts_capped) - - # 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] - batch_results[[b]] <- dt + rm(dt); gc() + } - # Free intermediate objects - rm(ts_data, ts_dense, ts_capped, dt) - gc() + if (length(skipped_concepts) > 0) { + message(sprintf(" Skipped %d concepts: %s", + length(skipped_concepts), + paste(skipped_concepts, collapse = ", "))) } - # Merge all batches by (stay_id, hour) + # Merge all batches by (stay_id, hour), skipping NULLs message(" Merging batches...") + batch_results <- Filter(Negate(is.null), batch_results) + if (length(batch_results) == 0) { + stop("All concept batches failed. Cannot continue.") + } merged <- batch_results[[1]] - if (n_batches > 1) { - for (b in 2:n_batches) { + if (length(batch_results) > 1) { + for (b in 2:length(batch_results)) { merged <- merge(merged, batch_results[[b]], by = c("stay_id", "hour"), all = TRUE) - batch_results[[b]] <- NULL # allow GC to reclaim } } rm(batch_results, wins) diff --git a/src/slices/constants.py b/src/slices/constants.py index fbe193a..7b67a9b 100644 --- a/src/slices/constants.py +++ b/src/slices/constants.py @@ -38,7 +38,8 @@ # ============================================================================= # 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: +# - los_hosp / los_icu: leak downstream task labels (updated hourly by RICU, +# directly reveal the length-of-stay answer) +# - dur_var: internal ricu helper variable, not a clinical concept +FEATURE_BLOCKLIST: frozenset = frozenset({"los_hosp", "los_icu", "dur_var"}) diff --git a/src/slices/data/labels/mortality.py b/src/slices/data/labels/mortality.py index e24aec4..6c11efd 100644 --- a/src/slices/data/labels/mortality.py +++ b/src/slices/data/labels/mortality.py @@ -97,7 +97,20 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: # 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 + # When all-null (e.g., eICU), polars may infer Null/Boolean dtype — cast to + # Datetime so downstream comparisons don't fail on type mismatch. + date_of_death_dtype = merged["date_of_death"].dtype + if date_of_death_dtype not in ( + pl.Date, + pl.Datetime, + pl.Datetime("us"), + pl.Datetime("ns"), + pl.Datetime("ms"), + ): + merged = merged.with_columns( + pl.col("date_of_death").cast(pl.Datetime("us", "UTC")).alias("date_of_death") + ) + date_of_death_dtype = pl.Datetime("us", "UTC") is_date_type = date_of_death_dtype == pl.Date if window_hours is None and obs_hours is None: From 4a179ef80174ca3ced9da35efb8e338d7270ac30 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 17 Mar 2026 00:48:46 -0400 Subject: [PATCH 039/121] fix: harmonize column dtypes before concatenating datasets The patient_id column has different types across datasets (Int32 in MIIV vs String in eICU), causing polars concat to fail. Detect and cast mismatched columns to String before merging. --- scripts/preprocessing/create_combined_dataset.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/scripts/preprocessing/create_combined_dataset.py b/scripts/preprocessing/create_combined_dataset.py index 4eb49c9..8626fdc 100644 --- a/scripts/preprocessing/create_combined_dataset.py +++ b/scripts/preprocessing/create_combined_dataset.py @@ -189,6 +189,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"]]) From 318f883eb72b50c70477000c69569f409d7b8632 Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 18 Mar 2026 21:31:46 -0400 Subject: [PATCH 040/121] fix: derive eICU date_of_death from hospitaldischargeoffset for mortality_24h MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit eICU does not provide explicit death timestamps, so date_of_death was always NA — causing mortality_24h to produce zero positive labels. Derive date_of_death for expired patients from hospitaldischargeoffset using the same synthetic epoch as the stays extraction. Also add the missing observation/prediction window parameters to the mortality_24h task config, without which the label builder was falling through to the hospital mortality path and never using date_of_death at all. Verified end-to-end: produces 1,703 positives (0.88% rate) from 200k stays. --- configs/tasks/mortality_24h.yaml | 6 +++++ scripts/preprocessing/extract_with_ricu.R | 5 +++- src/slices/data/labels/mortality.py | 2 +- src/slices/data/tensor_preprocessing.py | 28 ++++++++--------------- 4 files changed, 20 insertions(+), 21 deletions(-) diff --git a/configs/tasks/mortality_24h.yaml b/configs/tasks/mortality_24h.yaml index f0c0cfa..8b03a47 100644 --- a/configs/tasks/mortality_24h.yaml +++ b/configs/tasks/mortality_24h.yaml @@ -1,10 +1,16 @@ # @package task # 24-hour mortality prediction (binary classification) # Predict if patient will die within 24h after the observation window +# +# Timeline: |---- observation (48h) ----|---- prediction (24h) ----| +# intime obs_end pred_end task_name: mortality_24h task_type: binary n_classes: null +observation_window_hours: 48 # must match seq_length_hours +prediction_window_hours: 24 +gap_hours: 0 head_type: mlp hidden_dims: [64] dropout: 0.3 diff --git a/scripts/preprocessing/extract_with_ricu.R b/scripts/preprocessing/extract_with_ricu.R index a6e18ac..636c09f 100644 --- a/scripts/preprocessing/extract_with_ricu.R +++ b/scripts/preprocessing/extract_with_ricu.R @@ -643,7 +643,10 @@ extract_mortality_eicu <- function(dataset, dict) { names(pat)[1] } - # eICU uses status strings and offsets rather than timestamps + # eICU uses status strings and offsets rather than timestamps. + # We derive date_of_death from hospitaldischargeoffset for expired patients, + # using the same synthetic epoch as the stays extraction so timelines are + # consistent (intime = epoch, outtime = epoch + unitdischargeoffset*60). hosp_flag <- if ("hospitaldischargestatus" %in% names(pat)) { as.integer(tolower(pat$hospitaldischargestatus) == "expired") } else { diff --git a/src/slices/data/labels/mortality.py b/src/slices/data/labels/mortality.py index 6c11efd..3b2fb75 100644 --- a/src/slices/data/labels/mortality.py +++ b/src/slices/data/labels/mortality.py @@ -90,7 +90,7 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: 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. " + "This dataset may not provide death timestamps. " "All windowed mortality labels will be null (excluded). " "Consider using hospital-mortality (hospital_expire_flag) instead." ) diff --git a/src/slices/data/tensor_preprocessing.py b/src/slices/data/tensor_preprocessing.py index 4591246..ced51dd 100644 --- a/src/slices/data/tensor_preprocessing.py +++ b/src/slices/data/tensor_preprocessing.py @@ -6,9 +6,12 @@ """ import logging -from typing import Dict, List +from typing import TYPE_CHECKING, Dict, List import numpy as np + +if TYPE_CHECKING: + import polars as pl import torch from tqdm import tqdm @@ -52,9 +55,7 @@ def extract_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" - ) + 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) @@ -77,12 +78,7 @@ def extract_tensors_from_dataframe( 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) - ) + 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 @@ -119,9 +115,7 @@ def convert_raw_to_tensors( 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 - ) + 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) @@ -133,12 +127,8 @@ def convert_raw_to_tensors( 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 - ) + 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) From 50d08637dba216d652366f53c1466dda35d8b4cd Mon Sep 17 00:00:00 2001 From: hill Date: Sat, 21 Mar 2026 19:08:33 -0400 Subject: [PATCH 041/121] feat: chunk RICU timeseries output by stay to reduce memory usage R script now writes timeseries as a directory of part files (chunks of 10k stays) instead of a single monolithic parquet file. Python extractor and tests updated to handle both single-file and partitioned directory formats. --- scripts/preprocessing/extract_with_ricu.R | 132 +++++++++++++++++----- src/slices/data/extractors/ricu.py | 28 ++++- tests/test_ricu_extractor.py | 24 ++++ 3 files changed, 148 insertions(+), 36 deletions(-) diff --git a/scripts/preprocessing/extract_with_ricu.R b/scripts/preprocessing/extract_with_ricu.R index 636c09f..caa4db6 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 @@ -26,6 +27,7 @@ if (length(missing) > 0) { install.packages(missing, repos = "https://cloud.r-project.org") } + suppressPackageStartupMessages({ library(ricu) library(arrow) @@ -144,7 +146,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) @@ -159,10 +162,8 @@ 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({ @@ -197,7 +198,34 @@ extract_timeseries <- function(dataset, concepts, seq_length_hours, dict) { }) } + # --- 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]]] @@ -239,41 +267,86 @@ extract_timeseries <- function(dataset, concepts, seq_length_hours, dict) { skipped_concepts <- c(skipped_concepts, batch_concepts) } - batch_results[[b]] <- dt + 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() } + rm(wins); gc() + if (length(skipped_concepts) > 0) { message(sprintf(" Skipped %d concepts: %s", length(skipped_concepts), paste(skipped_concepts, collapse = ", "))) } - # Merge all batches by (stay_id, hour), skipping NULLs - message(" Merging batches...") - batch_results <- Filter(Negate(is.null), batch_results) - if (length(batch_results) == 0) { + if (batches_written == 0) { stop("All concept batches failed. Cannot continue.") } - merged <- batch_results[[1]] - if (length(batch_results) > 1) { - for (b in 2:length(batch_results)) { - merged <- merge(merged, batch_results[[b]], - by = c("stay_id", "hour"), all = TRUE) + + # --- 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) } - } - rm(batch_results, wins) - gc() + 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() - # 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) + if (ci %% 5 == 0 || ci == n_chunks) { + message(sprintf(" Chunk %d/%d done.", ci, n_chunks)) + } + } 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) } # --------------------------------------------------------------------------- @@ -1000,13 +1073,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, seq_length_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) diff --git a/src/slices/data/extractors/ricu.py b/src/slices/data/extractors/ricu.py index 1c3f020..2972169 100644 --- a/src/slices/data/extractors/ricu.py +++ b/src/slices/data/extractors/ricu.py @@ -14,7 +14,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 @@ -30,6 +31,21 @@ 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 +53,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 @@ -85,7 +101,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,7 +157,7 @@ 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) diff --git a/tests/test_ricu_extractor.py b/tests/test_ricu_extractor.py index bcda728..1757244 100644 --- a/tests/test_ricu_extractor.py +++ b/tests/test_ricu_extractor.py @@ -228,6 +228,30 @@ 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"), + ) + 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().""" From 47d1c7a8919a6c0a0220cff1283af2a482de14b1 Mon Sep 17 00:00:00 2001 From: hill Date: Mon, 23 Mar 2026 00:29:56 -0400 Subject: [PATCH 042/121] feat: add Sprint 7p model capacity pilot for SSL label-efficiency ablation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add medium (d_model=128, 4L, ~816K params) and large (d_model=256, 4L, ~3.2M params) encoder configs to test whether the weak SSL advantage observed in Sprint 6 is due to insufficient model capacity. Sprint 7p runs MAE + supervised on MIIV/seed42/mortality_24h at label_fractions 0.01/0.1/0.5 with both protocols — 20 runs total. The hypothesis: larger models should widen the SSL-supervised gap because pretrained initialization regularizes more effectively than random init when capacity exceeds what few labels can teach. --- configs/model/transformer_large.yaml | 28 ++++++ configs/model/transformer_medium.yaml | 28 ++++++ scripts/run_experiments.py | 127 +++++++++++++++++++++++++- 3 files changed, 181 insertions(+), 2 deletions(-) create mode 100644 configs/model/transformer_large.yaml create mode 100644 configs/model/transformer_medium.yaml diff --git a/configs/model/transformer_large.yaml b/configs/model/transformer_large.yaml new file mode 100644 index 0000000..7ba5ca4 --- /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 capacity ablation (Sprint 7 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..d5f2551 --- /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 capacity ablation (Sprint 7 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/scripts/run_experiments.py b/scripts/run_experiments.py index a716002..721ccf8 100644 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -41,6 +41,55 @@ SEEDS = [42, 123, 456] LABEL_FRACTIONS_FULL = [0.01, 0.05, 0.1, 0.25, 0.5] LABEL_FRACTIONS_TREND = [0.1] +LABEL_FRACTIONS_PILOT = [0.01, 0.1, 0.5] + +# Model capacity variants for Sprint 7 pilot (capacity ablation) +MODEL_SIZES = { + "medium": { + "model": "transformer_medium", + "encoder": {"d_model": 128, "n_layers": 4, "n_heads": 8, "d_ff": 512}, + "ssl_scale": { + "mae": { + "ssl.decoder_d_model": 128, + "ssl.decoder_n_layers": 2, + "ssl.decoder_n_heads": 8, + "ssl.decoder_d_ff": 512, + }, + "jepa": { + "ssl.predictor_d_model": 64, + "ssl.predictor_n_layers": 2, + "ssl.predictor_n_heads": 4, + "ssl.predictor_d_ff": 256, + }, + "contrastive": { + "ssl.proj_hidden_dim": 512, + "ssl.proj_output_dim": 128, + }, + }, + }, + "large": { + "model": "transformer_large", + "encoder": {"d_model": 256, "n_layers": 4, "n_heads": 8, "d_ff": 1024}, + "ssl_scale": { + "mae": { + "ssl.decoder_d_model": 256, + "ssl.decoder_n_layers": 2, + "ssl.decoder_n_heads": 8, + "ssl.decoder_d_ff": 1024, + }, + "jepa": { + "ssl.predictor_d_model": 128, + "ssl.predictor_n_layers": 2, + "ssl.predictor_n_heads": 8, + "ssl.predictor_d_ff": 512, + }, + "contrastive": { + "ssl.proj_hidden_dim": 1024, + "ssl.proj_output_dim": 256, + }, + }, + }, +} LR_ABLATION = [2e-4, 5e-4, 2e-3] # 1e-3 reused from Phase 1 MASK_RATIO_ABLATION = [0.3, 0.75] # 0.5 reused from Phase 1 TRANSFER_PAIRS = [("miiv", "eicu"), ("eicu", "miiv")] @@ -60,6 +109,7 @@ "5": ["1", "2", "3", "4"], "6": ["1", "2", "3", "4", "5"], "7": ["1", "3", "5"], + "7p": ["6"], "8": ["1", "1b", "1c", "5"], } @@ -405,6 +455,78 @@ def build_sprint6(self): for frac in LABEL_FRACTIONS_TREND: self._add_supervised(sprint, ds, seed, task, frac) + def build_sprint7p(self): + """Model capacity pilot — test whether bigger models widen SSL-supervised gap. + + MIIV only, seed 42, mortality_24h, MAE + supervised. + Two model sizes (medium=128d/4L, large=256d/4L) at 3 label fractions. + Compares against Sprint 6 baseline (64d/2L) — no need to rerun baseline. + + Total: 2 pretrain + 2×3 finetune + 2×3 probe + 2×3 supervised = 20 runs. + """ + sprint = "7p" + ds, seed, task = "miiv", 42, "mortality_24h" + + for size_name, size_cfg in MODEL_SIZES.items(): + # Common model override for all runs at this size + model_extra = {"model": size_cfg["model"]} + + # --- MAE pretrain (new encoder size, full data) --- + pretrain_extra = { + **model_extra, + **size_cfg["ssl_scale"]["mae"], + } + pt = self._add_pretrain(sprint, "mae", ds, seed, extra=pretrain_extra) + + # --- MAE finetune (Protocol B) + probe (Protocol A) --- + finetune_extra = {**model_extra} + for frac in LABEL_FRACTIONS_PILOT: + self._add_finetune( + sprint, + "mae", + ds, + seed, + task, + False, + pt, + label_fraction=frac, + extra=finetune_extra, + name_extra={"size": size_name}, + ) + self._add_finetune( + sprint, + "mae", + ds, + seed, + task, + True, + pt, + label_fraction=frac, + extra=finetune_extra, + name_extra={"size": size_name}, + ) + + # --- Supervised baseline (same larger encoder, from scratch) --- + sup_extra = {**model_extra} + for frac in LABEL_FRACTIONS_PILOT: + sup_name = f"supervised_{task}_{ds}_seed{seed}_{size_name}" + if frac < 1.0: + frac_str = str(frac).replace(".", "") + sup_name += f"_frac{frac_str}" + run = Run( + id=f"s{sprint}_{sup_name}", + sprint=sprint, + run_type="supervised", + paradigm="supervised", + dataset=ds, + seed=seed, + output_dir=_output_dir(sprint, sup_name), + task=task, + label_fraction=frac, + extra_overrides=sup_extra, + ) + self.runs.append(run) + def build_sprint7(self): """Cross-dataset transfer — Protocol B only (full finetune).""" sprint = "7" @@ -447,6 +569,7 @@ def build_all(self) -> list[Run]: self.build_sprint4() self.build_sprint5() self.build_sprint6() + self.build_sprint7p() self.build_sprint7() self.build_sprint8() return self.runs @@ -993,10 +1116,10 @@ def cmd_warmup(args): # the OS between iterations. Python's allocator retains freed memory, # so in-process iteration causes cumulative RSS growth and OOM on the # combined dataset (~23 GB peak per run). - task_arg = task or "" result = subprocess.run( [ - sys.executable, "-c", + sys.executable, + "-c", "from slices.data.datamodule import ICUDataModule; " f"dm = ICUDataModule(" f" processed_dir={processed_dir!r}," From 76d3541e64b256a4eec9ff16b539fc6ef4ba86d4 Mon Sep 17 00:00:00 2001 From: Hannes Ill Date: Sun, 22 Mar 2026 15:34:29 +0000 Subject: [PATCH 043/121] fix: remove extraction-only fields from training task config Remove observation_window_hours, prediction_window_hours, and gap_hours from the Hydra training config for mortality_24h. These fields are only needed in the data extraction configs (src/slices/data/tasks/) and were causing pydantic validation errors during training due to extra="forbid" on TaskConfig. --- configs/tasks/mortality_24h.yaml | 3 --- 1 file changed, 3 deletions(-) diff --git a/configs/tasks/mortality_24h.yaml b/configs/tasks/mortality_24h.yaml index 8b03a47..fd7db19 100644 --- a/configs/tasks/mortality_24h.yaml +++ b/configs/tasks/mortality_24h.yaml @@ -8,9 +8,6 @@ task_name: mortality_24h task_type: binary n_classes: null -observation_window_hours: 48 # must match seq_length_hours -prediction_window_hours: 24 -gap_hours: 0 head_type: mlp hidden_dims: [64] dropout: 0.3 From 65dbf6122b8160ff20a5f71bb3d231d89d237786 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 24 Mar 2026 14:46:37 -0400 Subject: [PATCH 044/121] feat: add Sprints 10-12 for 5-seed expansion, classical baselines, and SMART Sprint 10: Seeds 789/1011 for full Sprints 1-8 scope (SSL + supervised, 630 runs) Sprint 11: XGBoost + GRU-D baselines with 5 seeds incl. label efficiency (360 runs) Sprint 12: SMART (NeurIPS 2024) as external SSL SOTA reference (135 runs) Sprint 9 (fairness) stays as eval-only, runs last after all training. Also fix W&B group names to include protocol (A/B) and label_fraction, so the Group view correctly aggregates seed replicates. --- scripts/run_experiments.py | 247 ++++++++++++++++++++++++++++++++++- src/slices/training/utils.py | 12 +- 2 files changed, 253 insertions(+), 6 deletions(-) diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index 721ccf8..b4b3217 100644 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -39,6 +39,7 @@ DATASETS = ["miiv", "eicu", "combined"] TASKS = ["mortality_24h", "mortality_hospital", "aki_kdigo", "los_remaining"] SEEDS = [42, 123, 456] +SEEDS_EXTENDED = [42, 123, 456, 789, 1011] # 5 seeds for cheap baselines LABEL_FRACTIONS_FULL = [0.01, 0.05, 0.1, 0.25, 0.5] LABEL_FRACTIONS_TREND = [0.1] LABEL_FRACTIONS_PILOT = [0.01, 0.1, 0.5] @@ -128,8 +129,8 @@ class Run: id: str sprint: str - run_type: str # "pretrain" | "finetune" | "supervised" - paradigm: str # "mae" | "jepa" | "contrastive" | "supervised" + run_type: str # "pretrain" | "finetune" | "supervised" | "gru_d" | "xgboost" + paradigm: str # "mae" | "jepa" | "contrastive" | "supervised" | "gru_d" | "xgboost" dataset: str seed: int output_dir: str @@ -149,6 +150,10 @@ def build_command(self, runs_by_id: dict[str, Run]) -> list[str]: return self._finetune_cmd(runs_by_id) elif self.run_type == "supervised": return self._supervised_cmd() + elif self.run_type == "gru_d": + return self._gru_d_cmd() + elif self.run_type == "xgboost": + return self._xgboost_cmd() else: raise ValueError(f"Unknown run_type: {self.run_type}") @@ -225,6 +230,44 @@ def _supervised_cmd(self) -> list[str]: cmd.append(f"{k}={v}") 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"sprint={self.sprint}", + f"hydra.run.dir={self.output_dir}", + ] + if self.label_fraction < 1.0: + cmd.append(f"label_fraction={self.label_fraction}") + for k, v in self.extra_overrides.items(): + cmd.append(f"{k}={v}") + 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"sprint={self.sprint}", + f"hydra.run.dir={self.output_dir}", + ] + if self.label_fraction < 1.0: + cmd.append(f"label_fraction={self.label_fraction}") + for k, v in self.extra_overrides.items(): + cmd.append(f"{k}={v}") + return cmd + # --------------------------------------------------------------------------- # Matrix Generation @@ -359,10 +402,52 @@ def _add_supervised( self.runs.append(run) return run + def _add_gru_d( + self, sprint: str, dataset: str, seed: int, task: str, label_fraction: float = 1.0 + ) -> Run: + name = f"gru_d_{task}_{dataset}_seed{seed}" + if label_fraction < 1.0: + frac_str = str(label_fraction).replace(".", "") + name += f"_frac{frac_str}" + run = Run( + id=f"s{sprint}_{name}", + sprint=sprint, + run_type="gru_d", + paradigm="gru_d", + dataset=dataset, + seed=seed, + output_dir=_output_dir(sprint, name), + task=task, + label_fraction=label_fraction, + ) + self.runs.append(run) + return run + + def _add_xgboost( + self, sprint: str, dataset: str, seed: int, task: str, label_fraction: float = 1.0 + ) -> Run: + name = f"xgboost_{task}_{dataset}_seed{seed}" + if label_fraction < 1.0: + frac_str = str(label_fraction).replace(".", "") + name += f"_frac{frac_str}" + run = Run( + id=f"s{sprint}_{name}", + sprint=sprint, + run_type="xgboost", + paradigm="xgboost", + dataset=dataset, + seed=seed, + output_dir=_output_dir(sprint, name), + task=task, + label_fraction=label_fraction, + ) + self.runs.append(run) + return run + # --- Sprint builders --- def build_sprint1(self): - """MIMIC, all tasks, Protocol B + supervised, seed=42.""" + """MIMIC, all tasks, Protocol B + supervised + baselines, seed=42.""" ds, seed, sprint = "miiv", 42, "1" for p in SSL_PARADIGMS: pt = self._add_pretrain(sprint, p, ds, seed) @@ -370,6 +455,8 @@ def build_sprint1(self): self._add_finetune(sprint, p, ds, seed, task, False, pt) for task in TASKS: self._add_supervised(sprint, ds, seed, task) + self._add_gru_d(sprint, ds, seed, task) + self._add_xgboost(sprint, ds, seed, task) def build_sprint1b(self): """LR sensitivity, MIMIC, mortality_24h, seed=42.""" @@ -398,7 +485,7 @@ def build_sprint2(self): self._add_finetune(sprint, p, ds, seed, task, True, pt) def build_sprint3(self): - """eICU, both protocols + supervised, seed=42.""" + """eICU, both protocols + supervised + baselines, seed=42.""" ds, seed, sprint = "eicu", 42, "3" for p in SSL_PARADIGMS: pt = self._add_pretrain(sprint, p, ds, seed) @@ -407,9 +494,11 @@ def build_sprint3(self): self._add_finetune(sprint, p, ds, seed, task, True, pt) for task in TASKS: self._add_supervised(sprint, ds, seed, task) + self._add_gru_d(sprint, ds, seed, task) + self._add_xgboost(sprint, ds, seed, task) def build_sprint4(self): - """Combined, both protocols + supervised, seed=42.""" + """Combined, both protocols + supervised + baselines, seed=42.""" ds, seed, sprint = "combined", 42, "4" for p in SSL_PARADIGMS: pt = self._add_pretrain(sprint, p, ds, seed) @@ -418,6 +507,8 @@ def build_sprint4(self): self._add_finetune(sprint, p, ds, seed, task, True, pt) for task in TASKS: self._add_supervised(sprint, ds, seed, task) + self._add_gru_d(sprint, ds, seed, task) + self._add_xgboost(sprint, ds, seed, task) def build_sprint5(self): """Seeds 123, 456 for datasets miiv, eicu, combined.""" @@ -431,6 +522,8 @@ def build_sprint5(self): self._add_finetune(sprint, p, ds, seed, task, True, pt) for task in TASKS: self._add_supervised(sprint, ds, seed, task) + self._add_gru_d(sprint, ds, seed, task) + self._add_xgboost(sprint, ds, seed, task) def build_sprint6(self): """Label efficiency ablation — reuses Phase 1 encoders.""" @@ -451,9 +544,13 @@ def build_sprint6(self): # Supervised label efficiency for frac in LABEL_FRACTIONS_FULL: self._add_supervised(sprint, ds, seed, "mortality_24h", frac) + self._add_gru_d(sprint, ds, seed, "mortality_24h", frac) + self._add_xgboost(sprint, ds, seed, "mortality_24h", frac) for task in TASKS[1:]: for frac in LABEL_FRACTIONS_TREND: self._add_supervised(sprint, ds, seed, task, frac) + self._add_gru_d(sprint, ds, seed, task, frac) + self._add_xgboost(sprint, ds, seed, task, frac) def build_sprint7p(self): """Model capacity pilot — test whether bigger models widen SSL-supervised gap. @@ -559,6 +656,143 @@ def build_sprint8(self): sprint, p, "miiv", seed, "mortality_24h", False, pt, name_extra=extra ) + def build_sprint10(self): + """Seeds 789, 1011 for Sprints 1-8 scope (SSL + supervised).""" + sprint = "10" + new_seeds = [789, 1011] + + # --- Core experiments (Sprint 1-4 scope) --- + for seed in new_seeds: + for ds in DATASETS: + for p in SSL_PARADIGMS: + pt = self._add_pretrain(sprint, p, ds, seed) + for task in TASKS: + self._add_finetune(sprint, p, ds, seed, task, False, pt) + self._add_finetune(sprint, p, ds, seed, task, True, pt) + for task in TASKS: + self._add_supervised(sprint, ds, seed, task) + + # --- Label efficiency (Sprint 6 scope) --- + for seed in new_seeds: + for ds in DATASETS: + for p in SSL_PARADIGMS: + pt = self._get_pretrain(p, ds, seed) + for frac in LABEL_FRACTIONS_FULL: + self._add_finetune(sprint, p, ds, seed, "mortality_24h", False, pt, frac) + self._add_finetune(sprint, p, ds, seed, "mortality_24h", True, pt, frac) + for task in TASKS[1:]: + for frac in LABEL_FRACTIONS_TREND: + self._add_finetune(sprint, p, ds, seed, task, False, pt, frac) + self._add_finetune(sprint, p, ds, seed, task, True, pt, frac) + for frac in LABEL_FRACTIONS_FULL: + self._add_supervised(sprint, ds, seed, "mortality_24h", frac) + for task in TASKS[1:]: + for frac in LABEL_FRACTIONS_TREND: + self._add_supervised(sprint, ds, seed, task, frac) + + # --- Cross-dataset transfer (Sprint 7 scope) --- + for seed in new_seeds: + for source_ds, target_ds in TRANSFER_PAIRS: + for p in SSL_PARADIGMS: + pt = self._get_pretrain(p, source_ds, seed) + for task in TASKS: + self._add_finetune( + sprint, + p, + target_ds, + seed, + task, + False, + pt, + source_dataset=source_ds, + ) + + # --- HP ablations (Sprint 8 scope) --- + for seed in new_seeds: + for p in SSL_PARADIGMS: + for lr in LR_ABLATION: + extra = {"optimizer.lr": lr} + pt = self._add_pretrain(sprint, p, "miiv", seed, extra) + self._add_finetune( + sprint, + p, + "miiv", + seed, + "mortality_24h", + False, + pt, + name_extra=extra, + ) + for mr in MASK_RATIO_ABLATION: + extra = {"ssl.mask_ratio": mr} + pt = self._add_pretrain(sprint, p, "miiv", seed, extra) + self._add_finetune( + sprint, + p, + "miiv", + seed, + "mortality_24h", + False, + pt, + name_extra=extra, + ) + + def build_sprint11(self): + """Classical baselines (XGBoost + GRU-D), all datasets/tasks, 5 seeds.""" + sprint = "11" + for seed in SEEDS_EXTENDED: + for ds in DATASETS: + for task in TASKS: + self._add_xgboost(sprint, ds, seed, task) + self._add_gru_d(sprint, ds, seed, task) + # Baseline label-efficiency (mirrors Sprint 6 baseline scope) + for frac in LABEL_FRACTIONS_FULL: + self._add_gru_d(sprint, ds, seed, "mortality_24h", frac) + self._add_xgboost(sprint, ds, seed, "mortality_24h", frac) + for task_name in TASKS[1:]: + for frac in LABEL_FRACTIONS_TREND: + self._add_gru_d(sprint, ds, seed, task_name, frac) + self._add_xgboost(sprint, ds, seed, task_name, frac) + + def build_sprint12(self): + """SMART (NeurIPS 2024) as external SSL reference, 5 seeds. + + SMART uses its own encoder (MART architecture, d_model=32) and element-wise + masking — NOT part of the controlled comparison (different architecture). + Included as an external SSL SOTA reference point in the appendix. + """ + sprint = "12" + # model=smart selects the MART encoder; ssl=smart is set via paradigm arg + pretrain_extra = {"model": "smart"} + finetune_extra = {"model": "smart"} + + for seed in SEEDS_EXTENDED: + for ds in DATASETS: + # Pretrain (ssl=smart is set by _pretrain_cmd via self.paradigm) + pt = self._add_pretrain(sprint, "smart", ds, seed, pretrain_extra) + # Both protocols + all tasks + for task in TASKS: + self._add_finetune( + sprint, + "smart", + ds, + seed, + task, + False, + pt, + extra=finetune_extra, + ) + self._add_finetune( + sprint, + "smart", + ds, + seed, + task, + True, + pt, + extra=finetune_extra, + ) + def build_all(self) -> list[Run]: """Build full experiment matrix. Order matters for dedup.""" self.build_sprint1() @@ -572,6 +806,9 @@ def build_all(self) -> list[Run]: self.build_sprint7p() self.build_sprint7() self.build_sprint8() + self.build_sprint10() + self.build_sprint11() + self.build_sprint12() return self.runs diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 9788170..2c7fe1b 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -273,11 +273,21 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: if run_name and freeze_encoder is True: run_name = run_name.replace("_finetune_", "_probe_", 1) + # 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) + if label_fraction is not None and label_fraction < 1.0: + frac_str = str(label_fraction).replace(".", "") + group += f"_frac{frac_str}" + logger = WandbLogger( project=cfg.logging.wandb_project, entity=cfg.logging.get("wandb_entity", None), name=run_name, - group=cfg.logging.get("wandb_group", None), + group=group, tags=tags, save_dir=cfg.output_dir, log_model=False, From 8ad2651f46a5d085c40a3a7d288a2c2441b83644 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 24 Mar 2026 14:52:38 -0400 Subject: [PATCH 045/121] feat: add GRU-D encoder and XGBoost baseline Add GRU-D encoder with decay mechanism for handling irregular time series, register it in encoder factory and checkpoint loading. Add XGBoost baseline training script and configs for classical baseline comparison. Add xgboost dependency. --- configs/gru_d.yaml | 105 ++++++++ configs/model/gru_d.yaml | 13 + configs/xgboost.yaml | 81 ++++++ pyproject.toml | 1 + scripts/training/xgboost_baseline.py | 310 ++++++++++++++++++++++ src/slices/models/encoders/__init__.py | 3 + src/slices/models/encoders/factory.py | 3 + src/slices/models/encoders/gru_d.py | 146 ++++++++++ src/slices/training/checkpoint_loading.py | 7 +- uv.lock | 21 ++ 10 files changed, 689 insertions(+), 1 deletion(-) create mode 100644 configs/gru_d.yaml create mode 100644 configs/model/gru_d.yaml create mode 100644 configs/xgboost.yaml create mode 100644 scripts/training/xgboost_baseline.py create mode 100644 src/slices/models/encoders/gru_d.py diff --git a/configs/gru_d.yaml b/configs/gru_d.yaml new file mode 100644 index 0000000..41e63ff --- /dev/null +++ b/configs/gru_d.yaml @@ -0,0 +1,105 @@ +# 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 + +# Experiment sprint +sprint: 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: s${sprint}_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/xgboost.yaml b/configs/xgboost.yaml new file mode 100644 index 0000000..7d989e2 --- /dev/null +++ b/configs/xgboost.yaml @@ -0,0 +1,81 @@ +# 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 + +# Experiment sprint +sprint: 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 + scale_pos_weight: null # set to "balanced" to auto-compute from label distribution + +# 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: s${sprint}_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/pyproject.toml b/pyproject.toml index 1f23e88..2730cea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,7 @@ dependencies = [ "pre-commit>=4.5.0", "portalocker>=2.8.0", "m4-infra>=0.0.0.dev0", + "xgboost>=2.0", ] [project.optional-dependencies] diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py new file mode 100644 index 0000000..55fbd83 --- /dev/null +++ b/scripts/training/xgboost_baseline.py @@ -0,0 +1,310 @@ +"""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, + mean_absolute_error, + mean_squared_error, + r2_score, + roc_auc_score, +) +from slices.data.datamodule import ICUDataModule +from xgboost import XGBClassifier, XGBRegressor + + +def extract_tabular_features(dataset, indices: list[int]) -> tuple[np.ndarray, np.ndarray]: + """Extract per-feature summary statistics from time-series data. + + For each of the D features, computes: mean, std, min, max, first, last, + obs_count, obs_fraction over observed values. Total = D * 8 features. + Vectorized over the full batch using the dataset's stacked tensors. + + Args: + dataset: ICUDataset instance. + indices: List of sample indices. + + Returns: + X: Feature matrix (N, D*8). + y: Label array (N,). + """ + 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, T, D = ts.shape + mask_float = mask.float() + + # obs_count, obs_fraction: (N, D) + obs_count = mask_float.sum(dim=1) # (N, D) + obs_frac = obs_count / T + + # Masked sum for mean: sum observed values / count + masked_ts = ts * mask_float # zero out unobserved + feat_sum = masked_ts.sum(dim=1) # (N, D) + safe_count = obs_count.clamp(min=1) + feat_mean = feat_sum / safe_count # (N, D) + + # Masked std + 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) + + # Min/max: fill unobserved with +inf/-inf before taking min/max + zeros = ts.new_zeros(N, D) + 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 and last observed value per feature + # Use argmax on mask along time axis for first; flip for last + # For features with no observations, falls back to 0 + first_idx = mask.float().argmax(dim=1) # (N, D) — first True index + last_idx = T - 1 - mask.flip(dims=[1]).float().argmax(dim=1) # (N, D) + + feat_first = ts.gather(1, first_idx.unsqueeze(1)).squeeze(1) # (N, D) + feat_last = ts.gather(1, last_idx.unsqueeze(1)).squeeze(1) # (N, D) + + # Zero out features with no observations + 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 + + # Stack: (N, D, 8) -> (N, D*8) + features = torch.stack( + [feat_mean, feat_std, feat_min, feat_max, feat_first, feat_last, obs_count, obs_frac], + dim=-1, + ) # (N, D, 8) + X = features.reshape(N, D * 8).numpy() + + # Labels + if dataset._labels_tensor is not None: + y = dataset._labels_tensor[idx_tensor].numpy() + else: + y = np.zeros(N, dtype=np.float32) + + return X, y + + +@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) + + task_name = cfg.task.get("task_name", "mortality_24h") + task_type = cfg.task.get("task_type", "binary") + + 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") + + # ========================================================================= + # 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) + 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": -1, + } + + if task_type == "regression": + model = XGBRegressor(**common_params) + else: + # Auto-compute scale_pos_weight if requested + scale_pos_weight = xgb_cfg.get("scale_pos_weight", None) + if scale_pos_weight == "balanced": + n_pos = y_train.sum() + n_neg = len(y_train) - n_pos + scale_pos_weight = n_neg / max(n_pos, 1) + print(f" scale_pos_weight (balanced): {scale_pos_weight:.2f}") + model = XGBClassifier( + **common_params, + scale_pos_weight=scale_pos_weight, + eval_metric="logloss", + ) + + 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] + y_pred_class = (y_pred_proba >= 0.5).astype(int) + + metrics["test/auroc"] = roc_auc_score(y_test, y_pred_proba) + metrics["test/auprc"] = average_precision_score(y_test, y_pred_proba) + metrics["test/accuracy"] = accuracy_score(y_test, y_pred_class) + + 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 + + # 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, + ) + + 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. Log to W&B + # ========================================================================= + if cfg.logging.get("use_wandb", False): + import wandb + + tags = list(cfg.logging.get("wandb_tags", [])) + if cfg.get("sprint"): + tags.append(f"sprint:{cfg.sprint}") + if cfg.get("label_fraction", 1.0) < 1.0: + tags.append(f"label_fraction:{cfg.label_fraction}") + + run = wandb.init( + project=cfg.logging.wandb_project, + entity=cfg.logging.get("wandb_entity", None), + name=cfg.logging.get("run_name", f"xgboost_{cfg.dataset}_{task_name}"), + group=cfg.logging.get("wandb_group", f"xgboost_{cfg.dataset}_{task_name}"), + tags=tags, + config=OmegaConf.to_container(cfg, resolve=True), + ) + run.summary.update(metrics) + if fairness_report: + run.summary.update( + { + f"fairness/{k}": v + for k, v in fairness_report.items() + if isinstance(v, (int, float)) + } + ) + wandb.finish() + + # ========================================================================= + # 7. 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}") + + print("\n" + "=" * 80) + print("XGBoost Baseline Complete!") + print("=" * 80) + + +if __name__ == "__main__": + main() diff --git a/src/slices/models/encoders/__init__.py b/src/slices/models/encoders/__init__.py index 26802a4..956d1df 100644 --- a/src/slices/models/encoders/__init__.py +++ b/src/slices/models/encoders/__init__.py @@ -2,6 +2,7 @@ from .base import BaseEncoder, EncoderConfig 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,6 +13,8 @@ "BaseEncoder", "EncoderConfig", "EncoderWithMissingToken", + "GRUDConfig", + "GRUDEncoder", "LinearConfig", "LinearEncoder", "ObservationTransformerConfig", 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..838b40d --- /dev/null +++ b/src/slices/models/encoders/gru_d.py @@ -0,0 +1,146 @@ +"""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 + self.W_gamma_h = nn.Parameter(torch.Tensor(d_model)) + 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 _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_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_prev = self.x_mean.expand(B, -1).clone() + + for t in range(T): + # Input decay + gamma_x = torch.exp(-torch.relu(self.W_gamma_x * delta[:, t, :] + self.b_gamma_x)) + x_decayed = mask_float[:, t, :] * x[:, t, :] + (1 - mask_float[:, t, :]) * ( + gamma_x * x_prev + (1 - gamma_x) * self.x_mean + ) + + # Hidden state decay + delta_h = delta[:, t, :].mean(dim=-1, keepdim=True).expand(-1, self.config.d_model) + gamma_h = torch.exp(-torch.relu(self.W_gamma_h * delta_h + self.b_gamma_h)) + 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_prev = x_decayed + + h = self.dropout(h) + return h + + def get_output_dim(self) -> int: + return self.config.d_model diff --git a/src/slices/training/checkpoint_loading.py b/src/slices/training/checkpoint_loading.py index 761b7fa..79ac6b5 100644 --- a/src/slices/training/checkpoint_loading.py +++ b/src/slices/training/checkpoint_loading.py @@ -16,6 +16,7 @@ from slices.models.encoders import ( EncoderWithMissingToken, + GRUDEncoder, ObservationTransformerEncoder, SMARTEncoder, TransformerEncoder, @@ -38,6 +39,10 @@ def infer_encoder_type(state_dict: Dict[str, Any]) -> Optional[str]: """ 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" @@ -73,7 +78,7 @@ def wrap_encoder_with_missing_token( The encoder, possibly wrapped with EncoderWithMissingToken. """ # Skip for encoders that handle missingness intrinsically - if isinstance(encoder, (ObservationTransformerEncoder, SMARTEncoder)): + if isinstance(encoder, (ObservationTransformerEncoder, SMARTEncoder, GRUDEncoder)): logger.info( "Skipping EncoderWithMissingToken wrapper for %s " "(handles missingness intrinsically)", diff --git a/uv.lock b/uv.lock index 91b78d3..79711f7 100644 --- a/uv.lock +++ b/uv.lock @@ -2296,6 +2296,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" }, ] @@ -3626,6 +3627,7 @@ dependencies = [ { name = "seaborn" }, { name = "torch" }, { name = "wandb" }, + { name = "xgboost" }, ] [package.optional-dependencies] @@ -3669,6 +3671,7 @@ requires-dist = [ { name = "seaborn", specifier = ">=0.13.2" }, { name = "torch", specifier = ">=2.0" }, { name = "wandb", specifier = ">=0.15" }, + { name = "xgboost", specifier = ">=2.0" }, ] provides-extras = ["dev"] @@ -4092,6 +4095,24 @@ wheels = [ { 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" }, ] +[[package]] +name = "xgboost" +version = "3.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "nvidia-nccl-cu12", marker = "sys_platform == 'linux'" }, + { name = "scipy" }, +] +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/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]] name = "xlsxwriter" version = "3.2.9" From 328a7319ee7a45e16c0a790c9773736ddd7eccb6 Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 25 Mar 2026 21:24:12 -0400 Subject: [PATCH 046/121] fix: cache raw tensors per dataset to eliminate 140GB+ cache bloat The tensor cache was storing normalized tensors keyed by (task, seed, label_fraction), creating 100+ files totaling 140GB+ and filling the 194GB VM disk during warmup. Now caches raw (unnormalized) tensors once per dataset (~300MB each). Normalization and stay filtering are applied at runtime (<1s). Reduces total cache from ~140GB to ~0.8GB (3 files). Warmup simplified from ~180 sequential DataModule setups to 3 (one per dataset). --- scripts/run_experiments.py | 55 +++++------- src/slices/data/dataset.py | 149 +++++++++++++++----------------- src/slices/data/tensor_cache.py | 111 ++++++------------------ 3 files changed, 119 insertions(+), 196 deletions(-) diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index b4b3217..62ee2c2 100644 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -1300,12 +1300,11 @@ def cmd_retry(args): def cmd_warmup(args): - """Pre-build tensor caches for requested sprints. + """Pre-build raw tensor caches for requested sprints. - Instantiates ICUDataModule for each unique (dataset, task, seed, label_fraction) - combination sequentially. This populates the tensor cache so that parallel - experiment runs can load preprocessed data from disk instead of each process - independently loading and converting raw parquet files (which causes OOM). + Creates one raw tensor cache per dataset (not per task/seed/fraction). + The cache stores unnormalized tensors; normalization is applied at runtime. + This means only ~0.8GB total cache instead of 140GB+. """ if args.revision: print(f"Note: --revision={args.revision} ignored (warmup is revision-independent)") @@ -1327,47 +1326,35 @@ def cmd_warmup(args): all_by_id = {r.id: r for r in all_runs} runs = [all_by_id[d] for d in deps_needed if d in all_by_id] + runs - # Collect unique (dataset, task, seed, label_fraction) combos - # task=None for pretrain, task= for finetune/supervised - combos: dict[tuple, None] = {} # ordered set via dict + # Collect unique datasets — raw tensor cache is per-dataset only + datasets: dict[str, None] = {} # ordered set via dict for r in runs: - if r.run_type == "pretrain": - key = (r.dataset, None, r.seed, 1.0) - else: - key = (r.dataset, r.task, r.seed, r.label_fraction) - combos[key] = None + datasets[r.dataset] = None print(f"Warming up tensor caches for sprint(s) {', '.join(sprints)}") - print(f" {len(combos)} unique (dataset, task, seed, label_fraction) combinations\n") + print(f" {len(datasets)} unique datasets to cache\n") import subprocess import sys - for i, (dataset, task, seed, label_fraction) in enumerate(combos, 1): + for i, dataset in enumerate(datasets, 1): processed_dir = f"data/processed/{dataset}" - task_str = task or "(pretrain/no task)" - frac_str = f", frac={label_fraction}" if label_fraction < 1.0 else "" - print(f"[{i}/{len(combos)}] {dataset} / {task_str} / seed={seed}{frac_str}") - - # Run each warmup in a subprocess so that memory is fully returned to - # the OS between iterations. Python's allocator retains freed memory, - # so in-process iteration causes cumulative RSS growth and OOM on the - # combined dataset (~23 GB peak per run). + print(f"[{i}/{len(datasets)}] {dataset}") + + # Run in a subprocess so that memory is fully returned to the OS. + # Uses task_name=None (unsupervised) to load ALL stays without + # any task-specific filtering — this populates the raw cache. result = subprocess.run( [ sys.executable, "-c", - "from slices.data.datamodule import ICUDataModule; " - f"dm = ICUDataModule(" - f" processed_dir={processed_dir!r}," - f" task_name={task!r}," - f" batch_size=1," - f" num_workers=0," - f" seed={seed}," - f" label_fraction={label_fraction}," + "from slices.data.dataset import ICUDataset; " + f"ds = ICUDataset(" + f" data_dir={processed_dir!r}," + f" task_name=None," + f" normalize=False," f"); " - "dm.setup(); " - f"print(len(dm.dataset))", + f"print(len(ds))", ], capture_output=True, text=True, @@ -1379,7 +1366,7 @@ def cmd_warmup(args): stderr_tail = result.stderr.strip().split("\n")[-3:] print(f" -> ERROR: {' '.join(stderr_tail)}") - print("\nWarmup complete. Tensor caches saved to data/processed//.tensor_cache/") + print("\nWarmup complete. Raw tensor caches saved to data/processed//.tensor_cache/") print("You can now run experiments in parallel without OOM.") diff --git a/src/slices/data/dataset.py b/src/slices/data/dataset.py index 8bee404..7ab78a9 100644 --- a/src/slices/data/dataset.py +++ b/src/slices/data/dataset.py @@ -187,108 +187,103 @@ def __init__( ) 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. Filter excluded stays, then normalize using run-specific train_indices + + 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)) - ) - self.labels_df = self.labels_df.filter( - ~pl.col("stay_id").is_in(list(self._excluded_stay_ids)) - ) - logger.debug(f"Filtered down to {len(self.timeseries_df):,} stays") - - # Create 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 = load_cached_tensors( - self.data_dir, - self.normalize, - self.seq_length, - self.n_features, - self.train_indices, - 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: - # 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") - # Free raw DataFrame — not needed when using cached tensors - del self.timeseries_df + 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 tensors directly from Polars DataFrame. - # Uses explode() for zero-copy extraction instead of to_list(), - # which creates Python objects and can OOM on large datasets. - timeseries_tensor, masks_tensor = extract_tensors_from_dataframe( - self.timeseries_df, self.seq_length, self.n_features - ) + # Extract from parquet and save raw cache + logger.debug(f"Loading timeseries from {timeseries_path.name}") + timeseries_df = pl.read_parquet(timeseries_path) - # Free raw DataFrame immediately to reduce peak memory. - # get_preprocessing_stages() reloads from parquet on demand. - del self.timeseries_df + 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." + ) - # Try to load existing normalization stats (for reproducibility) - cached_stats = load_normalization_stats( - self.data_dir, self.train_indices, self.normalize - ) + all_stay_ids = timeseries_df["stay_id"].to_list() - # Apply normalization and imputation (steps 2 and 3) - self._precompute_tensors( - precomputed_tensors=(timeseries_tensor, masks_tensor), - train_indices=self.train_indices, - cached_stats=cached_stats, + timeseries_tensor, masks_tensor = extract_tensors_from_dataframe( + timeseries_df, self.seq_length, self.n_features ) + del timeseries_df - # Save preprocessed tensors to cache for next run + # Save raw tensors to cache (shared across all runs for this dataset) save_cached_tensors( self.data_dir, - self._timeseries_tensor, - self._mask_tensor, - self.feature_means, - self.feature_stds, - self.normalize, + timeseries_tensor, + masks_tensor, self.seq_length, self.n_features, - self.train_indices, - self._excluded_stay_ids, ) + # Phase 2: Filter excluded stays, then normalize + + # Apply stay exclusion filtering (task-specific, e.g. missing labels) + 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) + timeseries_tensor = timeseries_tensor[idx_tensor] + masks_tensor = masks_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") + else: + self.stay_ids = all_stay_ids + + # Build stay_id -> index mapping + logger.debug("Building stay_id index mapping") + 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 existing normalization stats (for reproducibility) + cached_stats = load_normalization_stats(self.data_dir, self.train_indices, self.normalize) + + # Apply normalization and imputation + self._precompute_tensors( + precomputed_tensors=(timeseries_tensor, masks_tensor), + train_indices=self.train_indices, + cached_stats=cached_stats, + ) + # Pre-compute labels and static features self._precompute_labels_and_static() diff --git a/src/slices/data/tensor_cache.py b/src/slices/data/tensor_cache.py index 82b3679..1a97eda 100644 --- a/src/slices/data/tensor_cache.py +++ b/src/slices/data/tensor_cache.py @@ -9,7 +9,7 @@ import logging import warnings from pathlib import Path -from typing import Any, Dict, List, Optional, Set +from typing import Any, Dict, List, Optional import torch import yaml @@ -112,106 +112,67 @@ def save_normalization_stats( def get_tensor_cache_key( - normalize: bool, seq_length: int, n_features: int, - train_indices: Optional[List[int]], - excluded_stay_ids: Optional[Set[int]], ) -> str: - """Generate a hash key for tensor caching based on preprocessing parameters. + """Generate a hash key for raw tensor caching based on shape 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) + 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: - normalize: Whether normalization is enabled. seq_length: Sequence length. n_features: Number of features. - train_indices: Optional list of training indices. - excluded_stay_ids: Optional set of excluded stay IDs. Returns: Hash string to use as cache identifier. """ - # Create a deterministic string representation of parameters params = { - "normalize": normalize, "seq_length": seq_length, "n_features": 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 excluded_stay_ids: - excluded_str = ",".join(map(str, sorted(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( data_dir: Path, - normalize: bool, seq_length: int, n_features: int, - train_indices: Optional[List[int]], - excluded_stay_ids: Optional[Set[int]], ) -> Path: - """Get the path to the tensor cache file. + """Get the path to the raw tensor cache file. Args: data_dir: Path to data directory. - normalize: Whether normalization is enabled. seq_length: Sequence length. n_features: Number of features. - train_indices: Optional list of training indices. - excluded_stay_ids: Optional set of excluded stay IDs. Returns: Path to the tensor cache file. """ - cache_key = get_tensor_cache_key( - normalize, seq_length, n_features, train_indices, excluded_stay_ids - ) + 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, normalize: bool -) -> Optional[Dict[str, Any]]: - """Try to load and validate a tensor cache file. +def _try_load_cache(cache_path: Path, n_features: int, seq_length: int) -> 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. - normalize: Expected normalize flag. Returns: - Dictionary with cached tensors if valid, None otherwise. + 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, mmap=True) + cached = torch.load(cache_path, weights_only=True) # Validate cache metadata if cached.get("n_features") != n_features: @@ -220,9 +181,6 @@ def _try_load_cache( if cached.get("seq_length") != seq_length: logger.debug("Cached tensors have different sequence length, recomputing") return None - if cached.get("normalize") != normalize: - logger.debug("Cached tensors have different normalize flag, recomputing") - 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 @@ -246,7 +204,7 @@ def _try_load_cache( cached["mask_tensor"] = masks n_samples = timeseries.shape[0] if hasattr(timeseries, "shape") else len(timeseries) - logger.info(f"Loaded {n_samples:,} cached samples") + logger.info(f"Loaded {n_samples:,} cached raw samples") return cached except Exception as e: @@ -256,75 +214,58 @@ def _try_load_cache( def load_cached_tensors( data_dir: Path, - normalize: bool, seq_length: int, n_features: int, - train_indices: Optional[List[int]], - excluded_stay_ids: Optional[Set[int]], ) -> Optional[Dict[str, Any]]: - """Load cached preprocessed tensors if they exist and are valid. + """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. - normalize: Whether normalization is enabled. seq_length: Sequence length. n_features: Number of features. - train_indices: Optional list of training indices. - excluded_stay_ids: Optional set of excluded stay IDs. Returns: - Dictionary with cached tensors and metadata if valid, None otherwise. + Dictionary with cached raw tensors and metadata if valid, None otherwise. """ - cache_path = get_tensor_cache_path( - data_dir, normalize, seq_length, n_features, train_indices, excluded_stay_ids - ) - return _try_load_cache(cache_path, n_features, seq_length, normalize) + cache_path = get_tensor_cache_path(data_dir, seq_length, n_features) + return _try_load_cache(cache_path, n_features, seq_length) def save_cached_tensors( data_dir: Path, timeseries_tensor: torch.Tensor, mask_tensor: torch.Tensor, - feature_means: torch.Tensor, - feature_stds: torch.Tensor, - normalize: bool, seq_length: int, n_features: int, - train_indices: Optional[List[int]], - excluded_stay_ids: Optional[Set[int]], ) -> None: - """Save preprocessed tensors to cache file. + """Save raw tensors to cache file. + + 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: Preprocessed timeseries tensor. + timeseries_tensor: Raw timeseries tensor (unnormalized). mask_tensor: Observation mask tensor. - feature_means: Per-feature means. - feature_stds: Per-feature stds. - normalize: Whether normalization is enabled. seq_length: Sequence length. n_features: Number of features. - train_indices: Optional list of training indices. - excluded_stay_ids: Optional set of excluded stay IDs. """ - cache_path = get_tensor_cache_path( - data_dir, normalize, seq_length, n_features, train_indices, excluded_stay_ids - ) + cache_path = get_tensor_cache_path(data_dir, seq_length, n_features) cache_dir = cache_path.parent cache_dir.mkdir(exist_ok=True) cache_data = { "timeseries_tensor": timeseries_tensor, "mask_tensor": mask_tensor, - "feature_means": feature_means, - "feature_stds": feature_stds, "n_features": n_features, "seq_length": seq_length, - "normalize": normalize, } try: - logger.debug(f"Saving tensors to cache: {cache_path.name}") + logger.debug(f"Saving raw 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}") From 582e0e479c48e5c35338dbe85ff54c2fe1c8d7b7 Mon Sep 17 00:00:00 2001 From: hill Date: Fri, 27 Mar 2026 00:13:14 -0400 Subject: [PATCH 047/121] feat: add canonical results export pipeline from W&B Pulls all experiment runs from W&B and produces two structured parquet files: per_seed_results.parquet (one row per run with wandb IDs for traceability) and aggregated_results.parquet (mean/std/min/max across seeds per unique config). Key features: - Cross-sprint seed merging for core experiments (Sprints 1-5, 10) - Conditional fingerprinting: core runs group without sprint, ablation runs retain sprint as a grouping key - Experiment type classification (core, lr_ablation, label_efficiency, transfer, capacity_pilot, etc.) - Deduplication keeps only the most recent run per config+seed - Retry with exponential backoff for W&B API timeouts - Smart validation that only warns about core configs with low seed counts --- results/.gitkeep | 0 scripts/export_results.py | 690 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 690 insertions(+) create mode 100644 results/.gitkeep create mode 100644 scripts/export_results.py diff --git a/results/.gitkeep b/results/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/scripts/export_results.py b/scripts/export_results.py new file mode 100644 index 0000000..a3da2db --- /dev/null +++ b/scripts/export_results.py @@ -0,0 +1,690 @@ +#!/usr/bin/env python3 +""" +Canonical results export pipeline for SLICES. + +Pulls all experiment runs from W&B, extracts config + test metrics, and produces +two structured parquet files: + + - results/per_seed_results.parquet — one row per W&B run (~3000 rows) + - results/aggregated_results.parquet — one row per unique config (~600 rows), + with mean/std/min/max across seeds + +Both files include wandb run IDs for traceability back to W&B. + +Usage: + uv run python scripts/export_results.py + uv run python scripts/export_results.py --sprint 1 --validate-seeds 3 + uv run python scripts/export_results.py --paradigm mae jepa --dataset miiv + uv run python scripts/export_results.py --output-dir results/sprint1 --sprint 1 +""" +from __future__ import annotations + +import argparse +import json +import math +import os +import sys +import time +from pathlib import Path + +import pandas as pd +import wandb + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +# Sprint → experiment type mapping. +# Core sprints share the same config and only add seeds; they merge across sprints. +# Non-core sprints have ablation-specific pretrain configs that aren't visible +# in the finetune W&B config, so sprint must remain a grouping key. +EXPERIMENT_TYPES = { + "1": "core", + "2": "core", + "3": "core", + "4": "core", + "5": "core", + "10": "core", + "1b": "lr_ablation", + "1c": "mask_ablation", + "8": "hp_ablation", + "6": "label_efficiency", + "7": "transfer", + "7p": "capacity_pilot", + "11": "classical_baselines", + "12": "smart_reference", +} + +# For core runs: sprint is noise — same config across sprints, different seeds. +# lr/mask_ratio excluded because they're determined by protocol (redundant) or +# only in the pretrain config (NaN for finetune runs). +CORE_FINGERPRINT = [ + "experiment_type", + "paradigm", + "dataset", + "task", + "protocol", + "label_fraction", + "model_size", + "source_dataset", + "phase", +] + +# For non-core runs: sprint distinguishes pretrain configs that aren't captured +# in the finetune W&B config (e.g., which LR/mask_ratio the pretrain used). +ABLATION_FINGERPRINT = CORE_FINGERPRINT + ["sprint"] + +# Legacy fingerprint used for deduplication (includes sprint + all config dims). +DEDUP_FINGERPRINT = [ + "sprint", + "paradigm", + "dataset", + "task", + "protocol", + "label_fraction", + "lr", + "mask_ratio", + "model_size", + "source_dataset", + "phase", +] + +# Metrics to extract from run summaries. +TEST_METRICS = [ + # Binary classification + "test/auroc", + "test/auprc", + "test/accuracy", + "test/f1", + "test/precision", + "test/recall", + "test/specificity", + "test/brier_score", + "test/ece", + # Regression (LOS) + "test/mse", + "test/mae", + "test/r2", + # Universal + "test/loss", +] + +VAL_METRICS = [ + "val/auroc", + "val/auprc", + "val/loss", +] + +ALL_METRICS = TEST_METRICS + VAL_METRICS + +# d_model values that indicate non-default capacity (Sprint 7p pilot). +# Default d_model=128 maps to "default"; other values are named here. +SIZED_MODELS = {128: "default", 256: "large"} + +# Phases that correspond to evaluation runs (not pretraining). +EVAL_PHASES = ["finetune", "supervised", "gru_d", "xgboost"] + + +# --------------------------------------------------------------------------- +# W&B Fetching +# --------------------------------------------------------------------------- + + +def fetch_all_runs( + project: str, + entity: str | None = None, + state: str = "finished", + sprint: list[str] | None = None, + paradigm: list[str] | None = None, + dataset: list[str] | None = None, + phase: list[str] | None = None, +) -> list: + """Fetch runs from W&B with server-side filtering. + + Returns a list of wandb.Run objects. + """ + api = wandb.Api(timeout=300) + path = f"{entity}/{project}" if entity else project + + # Build server-side filters + filters: dict = {} + if state: + filters["state"] = state + + # Tag-based filters — use $all so runs must match ALL specified tags. + tag_filters: list[str] = [] + + if sprint: + if len(sprint) == 1: + tag_filters.append(f"sprint:{sprint[0]}") + # Multiple sprints: fetch all, filter client-side (W&B $all requires ALL tags) + + if paradigm and len(paradigm) == 1: + tag_filters.append(f"paradigm:{paradigm[0]}") + + if dataset and len(dataset) == 1: + tag_filters.append(f"dataset:{dataset[0]}") + + if phase and len(phase) == 1: + tag_filters.append(f"phase:{phase[0]}") + + if tag_filters: + filters["tags"] = {"$all": tag_filters} + + 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") + + # Client-side filtering for multi-value filters + sprint_set = set(sprint) if sprint and len(sprint) > 1 else None + paradigm_set = set(paradigm) if paradigm and len(paradigm) > 1 else None + dataset_set = set(dataset) if dataset and len(dataset) > 1 else None + phase_set = set(phase) if phase and len(phase) > 1 else None + + runs = [] + for run in runs_iter: + tags = set(run.tags) + + if sprint_set: + if not any(f"sprint:{s}" in tags for s in sprint_set): + continue + if paradigm_set: + if not any(f"paradigm:{p}" in tags for p in paradigm_set): + continue + if dataset_set: + if not any(f"dataset:{d}" in tags for d in dataset_set): + continue + if phase_set: + if not any(f"phase:{p}" in tags for p in phase_set): + 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): + """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 _retry(fn, max_retries=3, base_delay=5): + """Retry a function with exponential backoff on timeout/connection 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) + print( + f" Retry {attempt + 1}/{max_retries} after {delay}s: {e!r}", + file=sys.stderr, + ) + time.sleep(delay) + + +def _load_run_data(run) -> tuple[dict, dict, str, str, str, list[str], str, str]: + """Load all data from a W&B run in one call (minimizes lazy-load API hits). + + Returns (config, summary_dict, run_id, run_url, run_name, tags, group, created_at). + """ + # Access all lazy-loaded properties together so retries cover them all + config = dict(run.config) + summary = dict(run.summary._json_dict) + return ( + config, + summary, + run.id, + run.url, + run.name or "", + list(run.tags), + run.group, + run.created_at or "", + ) + + +def extract_run(run, metric_keys: list[str]) -> dict: + """Extract config + metrics from a W&B run in a single retry-protected call.""" + config, summary, run_id, run_url, run_name, tags, group, created_at = _retry( + lambda: _load_run_data(run) + ) + + # Derive protocol from freeze_encoder + freeze = _get_nested(config, "training.freeze_encoder") + if freeze is True: + protocol = "A" + elif freeze is False: + protocol = "B" + else: + protocol = None + + # Derive model_size from encoder.d_model + d_model = _get_nested(config, "encoder.d_model") + if d_model is not None and d_model != 128: + # Sprint 7p capacity pilot: medium=128, large=256 + model_size = SIZED_MODELS.get(d_model, f"d{d_model}") + else: + model_size = "default" + + # Detect source_dataset for transfer runs from run name or config + source_dataset = config.get("source_dataset", None) + if source_dataset is None: + if "_from_" in run_name: + parts = run_name.split("_from_") + if len(parts) >= 2: + source_part = parts[1].split("_")[0] + if source_part in ("miiv", "eicu", "combined"): + source_dataset = source_part + + # Detect phase from tags + phase = None + for tag in tags: + if tag.startswith("phase:"): + phase = tag.split(":", 1)[1] + break + if phase is None: + p = config.get("paradigm", "") + if p in ("supervised", "gru_d", "xgboost"): + phase = p + else: + phase = "finetune" + + # Extract metrics from summary + metrics = {} + for key in metric_keys: + val = summary.get(key, None) + if val is not None and isinstance(val, (int, float)): + if math.isnan(val) or math.isinf(val): + metrics[key] = float("nan") + else: + metrics[key] = float(val) + else: + metrics[key] = float("nan") + + sprint_str = str(config.get("sprint", "")) + experiment_type = EXPERIMENT_TYPES.get(sprint_str, "unknown") + + row = { + "wandb_run_id": run_id, + "wandb_run_url": run_url, + "wandb_run_name": run_name, + "wandb_group": group, + "created_at": created_at, + "experiment_type": experiment_type, + "sprint": sprint_str, + "paradigm": config.get("paradigm", None), + "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": _get_nested(config, "optimizer.lr", default=None), + "mask_ratio": _get_nested(config, "ssl.mask_ratio", default=None), + "model_size": model_size, + "source_dataset": source_dataset, + "revision": config.get("revision", None), + "phase": phase, + } + row.update(metrics) + return row + + +# --------------------------------------------------------------------------- +# DataFrame Construction +# --------------------------------------------------------------------------- + + +def build_per_seed_df(runs: list) -> pd.DataFrame: + """Build the per-seed DataFrame from raw W&B runs. + + One row per run with config columns + metric columns + wandb IDs. + """ + rows = [] + failed = [] + for i, run in enumerate(runs): + if (i + 1) % 100 == 0: + print(f" Processing run {i + 1}/{len(runs)}...", file=sys.stderr) + try: + row = extract_run(run, ALL_METRICS) + rows.append(row) + except Exception as e: + run_id = getattr(run, "id", "unknown") + failed.append(run_id) + print(f" FAILED to extract run {run_id}: {e!r}", file=sys.stderr) + if failed: + print( + f" WARNING: {len(failed)} runs failed extraction and were skipped.", + file=sys.stderr, + ) + + df = pd.DataFrame(rows) + + # Ensure correct dtypes + 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") + df["label_fraction"] = df["label_fraction"].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") + + # Deduplicate: keep only the most recent run per (fingerprint + seed). + # This handles reruns/revisions — the latest finished run is canonical. + # Uses conditional fingerprints: core runs dedup by CORE_FINGERPRINT (which + # excludes sprint/lr/mask_ratio), non-core by ABLATION_FINGERPRINT. + df = df.sort_values("created_at", ascending=False) + before = len(df) + + core_mask = df["experiment_type"] == "core" + core_dedup_cols = [c for c in CORE_FINGERPRINT + ["seed"] if c in df.columns] + ablation_dedup_cols = [c for c in ABLATION_FINGERPRINT + ["seed"] if c in df.columns] + + core_df = df[core_mask].drop_duplicates(subset=core_dedup_cols, keep="first") + non_core_df = df[~core_mask].drop_duplicates(subset=ablation_dedup_cols, keep="first") + df = pd.concat([core_df, non_core_df], ignore_index=True) + + after = len(df) + if before != after: + print( + f" Deduplicated: {before} -> {after} rows " + f"({before - after} older duplicate runs removed).", + file=sys.stderr, + ) + + # Sort for reproducibility + sort_cols = [c for c in ["sprint", "paradigm", "dataset", "task", "seed"] if c in df.columns] + if sort_cols: + df = df.sort_values(sort_cols, na_position="last").reset_index(drop=True) + + return df + + +def _aggregate_group(df: pd.DataFrame, fingerprint_cols: list[str]) -> pd.DataFrame: + """Aggregate a subset of runs by the given fingerprint columns. + + Returns a DataFrame with one row per unique fingerprint, containing + mean/std/min/max of metrics and lists of run IDs/seeds/sprints. + """ + metric_cols = [c for c in ALL_METRICS if c in df.columns] + + # Sentinel for NaN in groupby + work = df.copy() + for col in fingerprint_cols: + if col in work.columns: + work[col] = work[col].fillna("__none__") + + grouped = work.groupby(fingerprint_cols, dropna=False) + agg_rows = [] + + for fingerprint, group in grouped: + if isinstance(fingerprint, str): + fingerprint = (fingerprint,) + row = dict(zip(fingerprint_cols, fingerprint)) + + # Restore None from sentinel + for col in fingerprint_cols: + if row.get(col) == "__none__": + row[col] = None + + # Run metadata + 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["sprint_list"] = json.dumps(sorted(group["sprint"].dropna().unique().tolist())) + + # Metric aggregation + for metric in metric_cols: + values = group[metric].dropna() + 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() + 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") + + agg_rows.append(row) + + return pd.DataFrame(agg_rows) + + +def build_aggregated_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: + """Group by fingerprint, compute mean/std/min/max across seeds. + + Uses conditional grouping: + - Core runs (Sprints 1-5, 10): group WITHOUT sprint to merge seeds across sprints + - Non-core runs (ablations, label efficiency, etc.): group WITH sprint + """ + core_mask = per_seed_df["experiment_type"] == "core" + core = per_seed_df[core_mask] + non_core = per_seed_df[~core_mask] + + parts = [] + if len(core) > 0: + print(f" Aggregating {len(core)} core runs (cross-sprint)...", file=sys.stderr) + parts.append(_aggregate_group(core, CORE_FINGERPRINT)) + if len(non_core) > 0: + print(f" Aggregating {len(non_core)} non-core runs (per-sprint)...", file=sys.stderr) + parts.append(_aggregate_group(non_core, ABLATION_FINGERPRINT)) + + if not parts: + return pd.DataFrame() + + agg_df = pd.concat(parts, ignore_index=True) + + # Sort + sort_cols = [ + c for c in ["experiment_type", "paradigm", "dataset", "task"] if c 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 + + +# --------------------------------------------------------------------------- +# Validation +# --------------------------------------------------------------------------- + + +def validate( + per_seed_df: pd.DataFrame, + aggregated_df: pd.DataFrame, + expected_core_seeds: int = 5, +) -> list[str]: + """Validate results and return warning strings. + + Only warns about core configs with fewer seeds than expected. + Non-core configs have variable seed counts by design. + """ + warnings = [] + + # Check core configs for expected seed count + if "experiment_type" in aggregated_df.columns: + core = aggregated_df[aggregated_df["experiment_type"] == "core"] + low_seed = core[core["n_seeds"] < expected_core_seeds] + if len(low_seed) > 0: + warnings.append( + f"WARNING: {len(low_seed)}/{len(core)} core configs have fewer " + f"than {expected_core_seeds} seeds:" + ) + for _, row in low_seed.iterrows(): + desc = ", ".join( + f"{c}={row[c]}" + for c in ["paradigm", "dataset", "task", "protocol"] + if c in row and row[c] is not None + ) + seeds = json.loads(row["seed_list"]) if pd.notna(row.get("seed_list")) else [] + warnings.append(f" {desc} — n_seeds={row['n_seeds']}, seeds={seeds}") + + # Check for runs with no test metrics at all + test_cols = [c for c in TEST_METRICS if c in per_seed_df.columns] + if test_cols: + all_nan = per_seed_df[test_cols].isna().all(axis=1) + if all_nan.any(): + n_empty = all_nan.sum() + warnings.append( + f"WARNING: {n_empty} runs have no test metrics at all. " + "These may be pretraining runs or crashed evaluations." + ) + + # Summary by experiment type + n_runs = len(per_seed_df) + n_configs = len(aggregated_df) + print("\nValidation summary:", file=sys.stderr) + print(f" Total runs: {n_runs}", file=sys.stderr) + print(f" Unique configs: {n_configs}", file=sys.stderr) + if "experiment_type" in aggregated_df.columns: + for etype, group in aggregated_df.groupby("experiment_type"): + seed_dist = dict(group["n_seeds"].value_counts().sort_index()) + print(f" {etype}: {len(group)} configs, seeds: {seed_dist}", file=sys.stderr) + print(f" Warnings: {len(warnings)}", file=sys.stderr) + + return warnings + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main(): + 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"), + help="W&B project name (default: WANDB_PROJECT env var or 'slices')", + ) + parser.add_argument( + "--entity", + default=os.environ.get("WANDB_ENTITY"), + help="W&B entity name (default: WANDB_ENTITY env var)", + ) + parser.add_argument( + "--sprint", + nargs="+", + help="Filter to specific sprint(s), e.g. --sprint 1 1b 2", + ) + parser.add_argument( + "--paradigm", + nargs="+", + help="Filter to specific paradigm(s), e.g. --paradigm mae jepa", + ) + parser.add_argument( + "--dataset", + nargs="+", + help="Filter to specific dataset(s), e.g. --dataset miiv eicu", + ) + parser.add_argument( + "--phase", + nargs="+", + default=EVAL_PHASES, + help=f"Filter to specific phase(s) (default: {EVAL_PHASES})", + ) + parser.add_argument( + "--state", + default="finished", + help="Run state filter (default: finished)", + ) + parser.add_argument( + "--output-dir", + default="results", + help="Output directory (default: results/)", + ) + parser.add_argument( + "--expected-core-seeds", + type=int, + default=5, + help="Expected seed count for core configs (default: 5)", + ) + + args = parser.parse_args() + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + # Fetch runs + runs = fetch_all_runs( + project=args.project, + entity=args.entity, + state=args.state, + sprint=args.sprint, + paradigm=args.paradigm, + dataset=args.dataset, + phase=args.phase, + ) + + if not runs: + print("No runs found matching filters. Exiting.", file=sys.stderr) + sys.exit(0) + + # Build DataFrames + print(f"\nBuilding per-seed DataFrame from {len(runs)} runs...", file=sys.stderr) + per_seed_df = build_per_seed_df(runs) + print(f" Shape: {per_seed_df.shape}", file=sys.stderr) + + print("\nBuilding aggregated DataFrame...", file=sys.stderr) + aggregated_df = build_aggregated_df(per_seed_df) + print(f" Shape: {aggregated_df.shape}", file=sys.stderr) + + # Validate + warnings = validate(per_seed_df, aggregated_df, expected_core_seeds=args.expected_core_seeds) + for w in warnings: + print(w, file=sys.stderr) + + # Save + per_seed_path = output_dir / "per_seed_results.parquet" + aggregated_path = output_dir / "aggregated_results.parquet" + + per_seed_df.to_parquet(per_seed_path, index=False) + aggregated_df.to_parquet(aggregated_path, index=False) + + print("\nSaved:", file=sys.stderr) + print(f" {per_seed_path} ({len(per_seed_df)} rows)", file=sys.stderr) + print(f" {aggregated_path} ({len(aggregated_df)} rows)", file=sys.stderr) + + # Print quick summary to stdout for piping + print("\n--- Quick Summary ---") + print(f"Runs: {len(per_seed_df)}, Configs: {len(aggregated_df)}") + if "experiment_type" in aggregated_df.columns: + for etype, group in aggregated_df.groupby("experiment_type"): + seed_dist = dict(group["n_seeds"].value_counts().sort_index()) + print(f" {etype}: {len(group)} configs, seeds: {seed_dist}") + if "paradigm" in aggregated_df.columns: + print(f"Paradigms: {sorted(aggregated_df['paradigm'].dropna().unique())}") + if "dataset" in aggregated_df.columns: + print(f"Datasets: {sorted(aggregated_df['dataset'].dropna().unique())}") + if "task" in aggregated_df.columns: + print(f"Tasks: {sorted(aggregated_df['task'].dropna().unique())}") + + +if __name__ == "__main__": + main() From d4b585b64a2a5dc626dbff276023864bfc798e9c Mon Sep 17 00:00:00 2001 From: hill Date: Fri, 27 Mar 2026 12:07:49 -0400 Subject: [PATCH 048/121] feat: add post-run fairness evaluation script with W&B write-back Add standalone scripts/eval/evaluate_fairness.py that batch-evaluates fairness metrics on existing W&B runs without retraining. Queries W&B for core experiment runs, loads best checkpoints, runs inference, and writes fairness/* metrics back to the same run summaries. Key components: - scripts/eval/evaluate_fairness.py: argparse CLI with --sprint/--dataset filters, --dry-run, --skip-existing for resumability, datamodule reuse - src/slices/eval/inference.py: shared run_inference() extracted from run_fairness_evaluation() to avoid duplication - export_results.py: add 21 fairness metric keys to TEST_METRICS No existing code behavior is changed. The utils.py refactor replaces the inline inference loop with a call to the shared function. --- configs/eval/default.yaml | 4 +- scripts/eval/evaluate_fairness.py | 603 ++++++++++++++++++++++++++++++ scripts/export_results.py | 22 ++ src/slices/eval/__init__.py | 2 + src/slices/eval/inference.py | 61 +++ src/slices/training/utils.py | 28 +- 6 files changed, 697 insertions(+), 23 deletions(-) create mode 100644 scripts/eval/evaluate_fairness.py create mode 100644 src/slices/eval/inference.py diff --git a/configs/eval/default.yaml b/configs/eval/default.yaml index 293f114..a042712 100644 --- a/configs/eval/default.yaml +++ b/configs/eval/default.yaml @@ -27,7 +27,9 @@ metrics: # 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/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py new file mode 100644 index 0000000..d56db1c --- /dev/null +++ b/scripts/eval/evaluate_fairness.py @@ -0,0 +1,603 @@ +#!/usr/bin/env python3 +"""Post-run fairness evaluation for SLICES experiment runs. + +Queries W&B for finished finetune/supervised runs, loads their best checkpoints, +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 core experiment runs (~480 runs). Supports +resumability via --skip-existing (default) and scoping via --sprint/--paradigm/ +--dataset filters. + +Usage: + # Evaluate all core runs (default: sprints 1-5, 10) + uv run python scripts/eval/evaluate_fairness.py + + # Scope to specific sprint/dataset + uv run python scripts/eval/evaluate_fairness.py --sprint 1 --dataset miiv + + # Preview which runs would be evaluated + uv run python scripts/eval/evaluate_fairness.py --dry-run + + # Override paths (e.g., different machine than training) + uv run python scripts/eval/evaluate_fairness.py \ + --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 --force + + # Debug with a single run + uv run python scripts/eval/evaluate_fairness.py --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 + +log = logging.getLogger("evaluate_fairness") + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +CORE_SPRINTS = ["1", "2", "3", "4", "5", "10"] +DEFAULT_PHASES = ["finetune", "supervised"] +DEFAULT_PROTECTED_ATTRIBUTES = ["gender", "age_group", "race"] + + +# --------------------------------------------------------------------------- +# 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 fetch_eval_runs( + project: str, + entity: Optional[str], + sprints: Optional[list[str]], + paradigms: Optional[list[str]], + datasets: Optional[list[str]], + phases: 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"} + tag_filters: list[str] = [] + + if sprints and len(sprints) == 1: + tag_filters.append(f"sprint:{sprints[0]}") + if paradigms and len(paradigms) == 1: + tag_filters.append(f"paradigm:{paradigms[0]}") + if datasets and len(datasets) == 1: + tag_filters.append(f"dataset:{datasets[0]}") + if phases and len(phases) == 1: + tag_filters.append(f"phase:{phases[0]}") + + if tag_filters: + filters["tags"] = {"$all": tag_filters} + + 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") + + # Client-side filtering for multi-value filters + sprint_set = set(sprints) if sprints and len(sprints) > 1 else None + paradigm_set = set(paradigms) if paradigms and len(paradigms) > 1 else None + dataset_set = set(datasets) if datasets and len(datasets) > 1 else None + phase_set = set(phases) if phases and len(phases) > 1 else None + + runs = [] + for run in runs_iter: + tags = set(run.tags) + + if sprint_set and not any(f"sprint:{s}" in tags for s in sprint_set): + continue + if paradigm_set and not any(f"paradigm:{p}" in tags for p in paradigm_set): + continue + if dataset_set and not any(f"dataset:{d}" in tags for d in dataset_set): + continue + if phase_set and not any(f"phase:{p}" in tags for p in phase_set): + continue + + runs.append(run) + + log.info("Fetched %d runs.", len(runs)) + return runs + + +def has_fairness_metrics(run) -> bool: + """Check if run summary already contains fairness/* keys.""" + try: + summary = dict(run.summary._json_dict) + return any(k.startswith("fairness/") for k in summary) + except Exception: + return False + + +def write_fairness_to_wandb( + run_path: str, + fairness_flat: dict[str, Any], + dry_run: bool = False, +) -> None: + """Write fairness metrics to W&B run summary (additive only).""" + import wandb + + if dry_run: + log.info(" [DRY RUN] Would write %d fairness keys to %s", len(fairness_flat), run_path) + return + + def _do_update(): + api = wandb.Api(timeout=120) + run = api.run(run_path) + 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.""" + 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 / "checkpoints" + return Path(output_dir) / "checkpoints" + + +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. + + Globs for finetune-*.ckpt or supervised-*.ckpt (excluding last.ckpt), + parses the metric value from the filename, and returns the best one. + Falls back to the most recently modified non-last checkpoint. + + 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 + + # Glob for all .ckpt files, exclude last.ckpt + ckpts = [p for p in ckpt_dir.glob("*.ckpt") if p.name != "last.ckpt"] + if not ckpts: + # Fall back to last.ckpt if nothing else exists + last = ckpt_dir / "last.ckpt" + if last.exists(): + log.info(" Only last.ckpt found, using it") + return last + log.warning(" No checkpoints found in %s", ckpt_dir) + return None + + # Try to parse metric value from filename: prefix-{epoch}-{metric}.ckpt + # e.g. finetune-012-val/auprc=0.8432.ckpt or finetune-012-0.8432.ckpt + 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 metric 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 + + # Fallback: most recently modified + ckpts.sort(key=lambda p: p.stat().st_mtime, reverse=True) + log.info(" Using most recent checkpoint: %s", ckpts[0].name) + return ckpts[0] + + +# --------------------------------------------------------------------------- +# 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 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.fairness_evaluator import FairnessEvaluator + 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") + evaluator = FairnessEvaluator( + static_df=datamodule.dataset.static_df, + protected_attributes=protected_attributes, + min_subgroup_size=min_subgroup_size, + task_type=task_type, + ) + report = evaluator.evaluate(predictions, labels, stay_ids) + return report + + +def flatten_fairness_report(report: dict[str, Any]) -> dict[str, Any]: + """Flatten nested fairness report into flat fairness/* keys for W&B. + + Mirrors the flattening in run_fairness_evaluation() (training/utils.py). + """ + flat: dict[str, Any] = {} + for attr, metrics in report.items(): + for metric_name, value in metrics.items(): + if isinstance(value, (int, float)): + flat[f"fairness/{attr}/{metric_name}"] = value + elif isinstance(value, dict): + for sub_key, sub_val in value.items(): + if isinstance(sub_val, (int, float)): + flat[f"fairness/{attr}/{metric_name}/{sub_key}"] = sub_val + return flat + + +# --------------------------------------------------------------------------- +# 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"), + help="W&B project name (default: $WANDB_PROJECT or 'slices')", + ) + parser.add_argument( + "--entity", + default=os.environ.get("WANDB_ENTITY"), + help="W&B entity (default: $WANDB_ENTITY)", + ) + parser.add_argument("--sprint", nargs="+", help="Filter to sprint(s). Default: core sprints") + parser.add_argument("--paradigm", nargs="+", help="Filter to paradigm(s)") + parser.add_argument("--dataset", nargs="+", help="Filter to dataset(s)") + 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=50, help="Min samples per subgroup" + ) + parser.add_argument("--dry-run", action="store_true", help="List runs without processing") + 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") + return parser.parse_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 core sprints if not specified + sprints = args.sprint or CORE_SPRINTS + log.info("Sprints: %s", sprints) + + # Fetch runs from W&B + runs = _retry( + lambda: fetch_eval_runs( + project=args.project, + entity=args.entity, + sprints=sprints, + paradigms=args.paradigm, + datasets=args.dataset, + phases=args.phase, + ) + ) + + if not runs: + print("No runs found matching filters.", file=sys.stderr) + sys.exit(0) + + # 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)] + skipped = before - len(runs) + if skipped: + log.info("Skipped %d runs with existing fairness metrics.", skipped) + + 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 + output_dir = cfg.get("output_dir", "?") + task_type = _get_nested(cfg, "task.task_type", "binary") + ckpt = find_best_checkpoint(output_dir, args.outputs_root, task_type) + ckpt_str = str(ckpt) if ckpt else "NOT FOUND" + 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. Find checkpoint + output_dir = cfg.get("output_dir", "") + task_type = _get_nested(cfg, "task.task_type", "binary") + ckpt_path = find_best_checkpoint(output_dir, args.outputs_root, task_type) + if ckpt_path is None: + results["skipped"] += 1 + results["errors"].append((run.id, run_desc, "no checkpoint found")) + continue + + # 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 + + model = build_model(cfg, ckpt_path, datamodule) + + # 3. Evaluate fairness + report = evaluate_run_fairness( + model, + datamodule, + args.protected_attributes, + args.min_subgroup_size, + device, + ) + + if not report: + log.warning(" No fairness results (no valid attribute groups)") + results["skipped"] += 1 + continue + + # 4. Flatten and write back + fairness_flat = flatten_fairness_report(report) + 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 + 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 __name__ == "__main__": + main() diff --git a/scripts/export_results.py b/scripts/export_results.py index a3da2db..554d4d9 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -107,6 +107,28 @@ "test/r2", # Universal "test/loss", + # Fairness (populated by scripts/eval/evaluate_fairness.py) + "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/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/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", ] VAL_METRICS = [ diff --git a/src/slices/eval/__init__.py b/src/slices/eval/__init__.py index ec8cc6a..a14efea 100644 --- a/src/slices/eval/__init__.py +++ b/src/slices/eval/__init__.py @@ -9,6 +9,7 @@ from slices.eval.fairness_evaluator import FairnessEvaluator from slices.eval.imputation import ImputationEvaluator +from slices.eval.inference import run_inference from slices.eval.metrics import ( MetricConfig, build_metrics, @@ -24,4 +25,5 @@ "ImputationEvaluator", "bootstrap_ci", "paired_bootstrap_test", + "run_inference", ] diff --git a/src/slices/eval/inference.py b/src/slices/eval/inference.py new file mode 100644 index 0000000..1f18b41 --- /dev/null +++ b/src/slices/eval/inference.py @@ -0,0 +1,61 @@ +"""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 typing import List, Tuple, Union + +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.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 = torch.cat(all_labels) + return predictions, labels, all_stay_ids diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 2c7fe1b..8020dd2 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -322,29 +322,13 @@ def run_fairness_evaluation( print("=" * 80) from slices.eval.fairness_evaluator import FairnessEvaluator + 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, + ) evaluator = FairnessEvaluator( static_df=datamodule.dataset.static_df, From c5667a5c689dfbcd1e33bd952bf916e18e189ff2 Mon Sep 17 00:00:00 2001 From: hill Date: Fri, 27 Mar 2026 12:37:09 -0400 Subject: [PATCH 049/121] fix: handle Lightning checkpoint dir issue in fairness eval Lightning's ModelCheckpoint creates subdirectories instead of files when the monitor metric contains '/' (e.g. val/auprc). This left only last.ckpt as usable checkpoints. Updated find_best_checkpoint to: - Exclude all last*.ckpt variants (not just last.ckpt) - Search subdirectories for .ckpt files (Strategy 2) - Fall back to last.ckpt when no best checkpoint exists Using last.ckpt is consistent with existing test metrics in W&B, which were also computed from the final model due to the same issue. --- scripts/eval/evaluate_fairness.py | 75 +++++++++++++++++-------------- 1 file changed, 41 insertions(+), 34 deletions(-) diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index d56db1c..a65517e 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -189,9 +189,11 @@ def find_best_checkpoint( ) -> Optional[Path]: """Find the best .ckpt file in the run's checkpoint directory. - Globs for finetune-*.ckpt or supervised-*.ckpt (excluding last.ckpt), - parses the metric value from the filename, and returns the best one. - Falls back to the most recently modified non-last checkpoint. + 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). @@ -206,38 +208,43 @@ def find_best_checkpoint( log.warning(" Checkpoint dir not found: %s", ckpt_dir) return None - # Glob for all .ckpt files, exclude last.ckpt - ckpts = [p for p in ckpt_dir.glob("*.ckpt") if p.name != "last.ckpt"] - if not ckpts: - # Fall back to last.ckpt if nothing else exists - last = ckpt_dir / "last.ckpt" - if last.exists(): - log.info(" Only last.ckpt found, using it") - return last - log.warning(" No checkpoints found in %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")] - # Try to parse metric value from filename: prefix-{epoch}-{metric}.ckpt - # e.g. finetune-012-val/auprc=0.8432.ckpt or finetune-012-0.8432.ckpt - 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 metric 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 - - # Fallback: most recently modified - 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 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 # --------------------------------------------------------------------------- From 6242014033927a266fe20c7af007dba7502fe95e Mon Sep 17 00:00:00 2001 From: hill Date: Fri, 27 Mar 2026 12:40:27 -0400 Subject: [PATCH 050/121] fix: replace '/' with '_' in checkpoint filenames to prevent silent save failure Lightning's ModelCheckpoint interprets '/' in filenames as a directory separator, causing checkpoints like 'finetune-005-val/auprc=0.8432.ckpt' to create empty directories instead of saving checkpoint files. This meant all test metrics were computed from the last-epoch model rather than the best-epoch model. Fix: use safe_monitor (val_auprc, val_mse, val_loss) in the filename template and set auto_insert_metric_name=False. New checkpoints will save as e.g. 'finetune-005-val_auprc=0.8432.ckpt'. Affected runs must be retrained to get proper best-checkpoint metrics. --- src/slices/training/utils.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 8020dd2..0a302ed 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -160,12 +160,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}-{val_loss:.4f}", monitor="val/loss", mode="min", save_top_k=3, save_last=True, verbose=True, + auto_insert_metric_name=False, ) callbacks.append(checkpoint_callback) @@ -200,14 +201,19 @@ def setup_finetune_callbacks(cfg: DictConfig, checkpoint_prefix: str = "finetune monitor = cfg.training.get("early_stopping_monitor", default_monitor) mode = cfg.training.get("early_stopping_mode", default_mode) + # Replace '/' with '_' in monitor name for safe filenames. + # Lightning interprets '/' as a directory separator in checkpoint filenames, + # which causes best checkpoints to silently not save (empty directories). + safe_monitor = monitor.replace("/", "_") checkpoint_callback = ModelCheckpoint( dirpath=cfg.get("checkpoint_dir", "checkpoints"), - filename=f"{checkpoint_prefix}-{{epoch:03d}}-{{{monitor}:.4f}}", + filename=f"{checkpoint_prefix}-{{epoch:03d}}-{{{safe_monitor}:.4f}}", monitor=monitor, mode=mode, save_top_k=3, save_last=True, verbose=True, + auto_insert_metric_name=False, ) callbacks.append(checkpoint_callback) From 7c5c742091d2e94f1cf8d172c0ddc3245d13da81 Mon Sep 17 00:00:00 2001 From: hill Date: Sat, 28 Mar 2026 18:54:14 -0400 Subject: [PATCH 051/121] fix: use summary_metrics to avoid N+1 W&B API calls summary._json_dict triggers a per-run GraphQL reload, causing ~886 sequential HTTP requests when filtering runs. summary_metrics uses data already fetched in the batch query. --- scripts/eval/evaluate_fairness.py | 10 +++++++--- scripts/export_results.py | 2 +- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index a65517e..36ea5e5 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -135,10 +135,14 @@ def fetch_eval_runs( def has_fairness_metrics(run) -> bool: - """Check if run summary already contains fairness/* keys.""" + """Check if run summary already contains fairness/* keys. + + Uses ``summary_metrics`` (populated from the batch query) instead of + ``summary._json_dict`` which triggers a per-run GraphQL reload. + """ try: - summary = dict(run.summary._json_dict) - return any(k.startswith("fairness/") for k in summary) + sm = run.summary_metrics or {} + return any(k.startswith("fairness/") for k in sm) except Exception: return False diff --git a/scripts/export_results.py b/scripts/export_results.py index 554d4d9..a2cf243 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -272,7 +272,7 @@ def _load_run_data(run) -> tuple[dict, dict, str, str, str, list[str], str, str] """ # Access all lazy-loaded properties together so retries cover them all config = dict(run.config) - summary = dict(run.summary._json_dict) + summary = dict(run.summary_metrics or {}) return ( config, summary, From bb6f01646bc67be2f2e1e37c093919398e49d254 Mon Sep 17 00:00:00 2001 From: hill Date: Mon, 30 Mar 2026 10:37:13 -0400 Subject: [PATCH 052/121] fix: align SMART implementation with original paper code Compared against github.com/yzhHoward/SMART and fixed 4 deviations: - Remove extra final_norm LayerNorm after MART blocks (not in original) - Use per-step momentum schedule via estimated_stepping_batches instead of per-epoch granularity, matching the original linear step schedule - Remove spurious attention weight dropout in VarAttention (original only applies proj_drop on the output, not on attention weights) - Set PE dropout=0.0 for SMART encoder (original has no PE dropout) --- src/slices/models/encoders/smart.py | 10 ++-------- src/slices/training/pretrain_module.py | 11 ++++------- tests/test_smart_encoder.py | 2 +- 3 files changed, 7 insertions(+), 16 deletions(-) diff --git a/src/slices/models/encoders/smart.py b/src/slices/models/encoders/smart.py index 0ed4a34..15cb058 100644 --- a/src/slices/models/encoders/smart.py +++ b/src/slices/models/encoders/smart.py @@ -266,7 +266,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 +411,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,9 +431,6 @@ def __init__(self, config: SMARTEncoderConfig) -> None: ] ) - # Final layer norm - self.final_norm = nn.LayerNorm(config.d_model) - def _validate_config(self) -> None: """Validate configuration parameters.""" if self.config.d_model % self.config.n_heads != 0: @@ -493,9 +490,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/training/pretrain_module.py b/src/slices/training/pretrain_module.py index 80ed825..c1705e9 100644 --- a/src/slices/training/pretrain_module.py +++ b/src/slices/training/pretrain_module.py @@ -175,13 +175,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/tests/test_smart_encoder.py b/tests/test_smart_encoder.py index 6779f5b..bcb5abe 100644 --- a/tests/test_smart_encoder.py +++ b/tests/test_smart_encoder.py @@ -337,7 +337,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.""" From 23e2dd86d57b2a7f2cba8aa5ac4ddd47d8ec7eec Mon Sep 17 00:00:00 2001 From: hill Date: Mon, 30 Mar 2026 11:26:44 -0400 Subject: [PATCH 053/121] feat: add TS2Vec temporal contrastive SSL objective MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address the thesis vulnerability that the standard contrastive objective may be disadvantaged by masking-only augmentation. TS2Vec uses input-level Gaussian noise + independent random masks + hierarchical temporal contrastive loss — augmentations natural to contrastive learning. Same encoder (TransformerEncoder, obs_aware=True) and projection head dims as standard contrastive for direct comparison. Sprint 13 runs 15 pretrain + 60 finetune across 3 datasets × 5 seeds × 4 tasks (Protocol B). --- configs/ssl/ts2vec.yaml | 35 ++ scripts/run_experiments.py | 16 + scripts/training/pretrain.py | 1 + src/slices/models/pretraining/__init__.py | 3 + src/slices/models/pretraining/factory.py | 3 + src/slices/models/pretraining/ts2vec.py | 486 ++++++++++++++++++++++ 6 files changed, 544 insertions(+) create mode 100644 configs/ssl/ts2vec.yaml create mode 100644 src/slices/models/pretraining/ts2vec.py 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/scripts/run_experiments.py b/scripts/run_experiments.py index 62ee2c2..32662f0 100644 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -793,6 +793,21 @@ def build_sprint12(self): extra=finetune_extra, ) + def build_sprint13(self): + """TS2Vec temporal contrastive variant, 5 seeds. + + Addresses the "contrastive was set up to fail" vulnerability by giving + the contrastive paradigm its natural augmentations (noise + masking) + and a temporal contrastive loss. Same encoder, same training budget. + Protocol B (full finetune) only — matches the primary evaluation protocol. + """ + sprint = "13" + for seed in SEEDS_EXTENDED: + for ds in DATASETS: + pt = self._add_pretrain(sprint, "ts2vec", ds, seed) + for task in TASKS: + self._add_finetune(sprint, "ts2vec", ds, seed, task, False, pt) + def build_all(self) -> list[Run]: """Build full experiment matrix. Order matters for dedup.""" self.build_sprint1() @@ -809,6 +824,7 @@ def build_all(self) -> list[Run]: self.build_sprint10() self.build_sprint11() self.build_sprint12() + self.build_sprint13() return self.runs diff --git a/scripts/training/pretrain.py b/scripts/training/pretrain.py index eeecf2c..4602a81 100644 --- a/scripts/training/pretrain.py +++ b/scripts/training/pretrain.py @@ -33,6 +33,7 @@ "mae": {"transformer"}, "jepa": {"transformer"}, "contrastive": {"transformer"}, + "ts2vec": {"transformer"}, "smart": {"smart"}, } 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/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/ts2vec.py b/src/slices/models/pretraining/ts2vec.py new file mode 100644 index 0000000..ab74e9e --- /dev/null +++ b/src/slices/models/pretraining/ts2vec.py @@ -0,0 +1,486 @@ +"""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 +from .masking import create_timestep_mask, extract_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 + + # Validate encoder + if not getattr(getattr(encoder, "config", None), "obs_aware", False): + raise ValueError( + "TS2Vec requires an encoder with obs_aware=True. " f"Got: {type(encoder).__name__}" + ) + encoder_pooling = getattr(encoder.config, "pooling", "none") + if encoder_pooling != "none": + raise ValueError( + "TS2Vec requires encoder with pooling='none' for per-token " + f"representations, but got pooling='{encoder_pooling}'" + ) + + 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, _, _ = self.encoder.tokenize(x1, obs_mask) + tokens_2, _, _ = self.encoder.tokenize(x2, obs_mask) + + # 4. Independent timestep masks + ssl_mask_1 = create_timestep_mask(B, T, self.config.mask_ratio, device) + ssl_mask_2 = create_timestep_mask(B, T, self.config.mask_ratio, device) + + # 5. Encode visible tokens for each view + vis_1, vp_1 = extract_visible_timesteps(tokens_1, ssl_mask_1) + enc_1 = self.encoder.encode(vis_1, vp_1) + + vis_2, vp_2 = extract_visible_timesteps(tokens_2, ssl_mask_2) + 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, ssl_mask_1, T) + full_2 = self._scatter_to_full(enc_2, ssl_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 = ssl_mask_1 & ssl_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, ssl_mask_1, ssl_mask_2) + + # Add masking statistics + with torch.no_grad(): + metrics.update( + { + "ts2vec_n_timesteps": torch.tensor(T), + "ts2vec_n_visible_view1": torch.tensor(ssl_mask_1.sum().item() / B), + "ts2vec_n_visible_view2": torch.tensor(ssl_mask_2.sum().item() / B), + "ts2vec_n_overlap_per_sample": torch.tensor(overlap.sum().item() / B), + "ts2vec_overlap_ratio": overlap.float().mean(), + } + ) + + 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. + """ + B, n_vis, d_enc = encoded.shape + device = encoded.device + + full = torch.zeros(B, n_timesteps, d_enc, device=device, dtype=encoded.dtype) + vis_indices = ssl_mask.float().argsort(dim=1, descending=True, stable=True) + scatter_idx = vis_indices[:, :n_vis].unsqueeze(-1).expand(-1, -1, d_enc) + full.scatter_(1, scatter_idx, encoded) + + return full + + @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, + } From ded458b34da662fd0df424a18ea7a337135114d2 Mon Sep 17 00:00:00 2001 From: hill Date: Mon, 30 Mar 2026 15:04:45 -0400 Subject: [PATCH 054/121] feat: add offline FLOPs measurement script for SSL paradigm comparison Computes per-step and total training FLOPs for MAE, JEPA, Contrastive, and TS2Vec using PyTorch's FlopCounterMode. Optionally queries W&B for gradient step counts and wall-clock time to compute totals. --- scripts/measure_flops.py | 340 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 340 insertions(+) create mode 100644 scripts/measure_flops.py diff --git a/scripts/measure_flops.py b/scripts/measure_flops.py new file mode 100644 index 0000000..73d0442 --- /dev/null +++ b/scripts/measure_flops.py @@ -0,0 +1,340 @@ +"""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 35 # override feature dim + uv run python scripts/measure_flops.py --sparsity 0.7 # fraction missing +""" + +import argparse +import os +from typing import Dict, Tuple + +import torch +import torch.nn as nn +from slices.models.encoders.factory import build_encoder +from slices.models.pretraining.factory import build_ssl_objective, get_ssl_config_class +from torch.utils.flop_counter import FlopCounterMode + +# --------------------------------------------------------------------------- +# Paradigm configs — mirror the YAML defaults used in actual runs +# --------------------------------------------------------------------------- + +ENCODER_CONFIG = { + "d_input": 35, + "d_model": 64, + "max_seq_length": 168, + "n_layers": 2, + "n_heads": 4, + "d_ff": 256, + "dropout": 0.1, + "activation": "gelu", + "prenorm": True, + "layer_norm_eps": 1e-5, + "use_positional_encoding": True, + "pooling": "none", # Required for SSL pretraining + "obs_aware": True, # All SSL paradigms use obs-aware tokenization +} + +# SSL configs — exact values from configs/ssl/*.yaml +SSL_CONFIGS: Dict[str, dict] = { + "mae": { + "mask_ratio": 0.5, + "decoder_d_model": 64, + "decoder_n_layers": 2, + "decoder_n_heads": 4, + "decoder_d_ff": 256, + "decoder_dropout": 0.1, + }, + "jepa": { + "mask_ratio": 0.5, + "mask_strategy": "block", + "mask_n_blocks": 3, + "predictor_d_model": 32, + "predictor_n_layers": 2, + "predictor_n_heads": 4, + "predictor_d_ff": 128, + "predictor_dropout": 0.1, + "momentum_base": 0.999, + "momentum_final": 1.0, + "loss_type": "mse", + }, + "contrastive": { + "mode": "instance", + "mask_ratio": 0.5, + "proj_hidden_dim": 256, + "proj_output_dim": 64, + "temperature": 0.07, + "complementary_masks": True, + }, + "ts2vec": { + "mask_ratio": 0.5, + "noise_scale": 0.01, + "crop_ratio": 1.0, + "proj_hidden_dim": 256, + "proj_output_dim": 64, + "temperature": 0.05, + "n_hierarchical_scales": 4, + }, +} + + +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) -> 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 pretrain runs (tagged phase:pretrain) + filters = {"tags": {"$all": ["phase:pretrain"]}, "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=35, help="Number of input features (d_input)" + ) + parser.add_argument("--seq-length", type=int, default=48, 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"), + help="W&B project", + ) + args = parser.parse_args() + + # Build encoder config with CLI overrides + encoder_config = {**ENCODER_CONFIG, "d_input": args.n_features} + + # 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 + + paradigms = ["mae", "jepa", "contrastive", "ts2vec"] + + 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: + wandb_data = query_wandb_runs(args.wandb_entity, args.wandb_project) + 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() From b67a82a2ac67337910fed612dfd161720623ff7c Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 31 Mar 2026 21:53:37 -0400 Subject: [PATCH 055/121] fix: export pipeline handles sprint 10/11 correctly - Add "baseline" to EVAL_PHASES so sprint 11 GRU-D/XGBoost runs (tagged phase:baseline) get picked up - Reclassify sprint 10 sub-1.0 label_fraction runs as experiment_type=label_efficiency so they merge with sprint 6 seeds during aggregation - Use cross-sprint dedup/aggregation for both core and label_efficiency types, giving 5 seeds per config --- scripts/export_results.py | 44 ++++++++++++++++++++++++++------------- 1 file changed, 30 insertions(+), 14 deletions(-) diff --git a/scripts/export_results.py b/scripts/export_results.py index a2cf243..a63eb2a 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -144,7 +144,7 @@ SIZED_MODELS = {128: "default", 256: "large"} # Phases that correspond to evaluation runs (not pretraining). -EVAL_PHASES = ["finetune", "supervised", "gru_d", "xgboost"] +EVAL_PHASES = ["finetune", "supervised", "gru_d", "xgboost", "baseline"] # --------------------------------------------------------------------------- @@ -346,6 +346,13 @@ def extract_run(run, metric_keys: list[str]) -> dict: sprint_str = str(config.get("sprint", "")) experiment_type = EXPERIMENT_TYPES.get(sprint_str, "unknown") + # Sprint 10 contains both core (label_fraction=1.0) and label_efficiency + # (label_fraction<1.0) runs. Reclassify the sub-1.0 runs so they merge + # with Sprint 6 label_efficiency during aggregation. + label_frac = config.get("label_fraction", 1.0) + if experiment_type == "core" and label_frac is not None and float(label_frac) < 1.0: + experiment_type = "label_efficiency" + row = { "wandb_run_id": run_id, "wandb_run_url": run_url, @@ -419,12 +426,13 @@ def build_per_seed_df(runs: list) -> pd.DataFrame: df = df.sort_values("created_at", ascending=False) before = len(df) - core_mask = df["experiment_type"] == "core" - core_dedup_cols = [c for c in CORE_FINGERPRINT + ["seed"] if c in df.columns] + cross_sprint_types = {"core", "label_efficiency"} + cross_sprint_mask = df["experiment_type"].isin(cross_sprint_types) + cross_sprint_dedup_cols = [c for c in CORE_FINGERPRINT + ["seed"] if c in df.columns] ablation_dedup_cols = [c for c in ABLATION_FINGERPRINT + ["seed"] if c in df.columns] - core_df = df[core_mask].drop_duplicates(subset=core_dedup_cols, keep="first") - non_core_df = df[~core_mask].drop_duplicates(subset=ablation_dedup_cols, keep="first") + core_df = df[cross_sprint_mask].drop_duplicates(subset=cross_sprint_dedup_cols, keep="first") + non_core_df = df[~cross_sprint_mask].drop_duplicates(subset=ablation_dedup_cols, keep="first") df = pd.concat([core_df, non_core_df], ignore_index=True) after = len(df) @@ -502,17 +510,25 @@ def build_aggregated_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: - Core runs (Sprints 1-5, 10): group WITHOUT sprint to merge seeds across sprints - Non-core runs (ablations, label efficiency, etc.): group WITH sprint """ - core_mask = per_seed_df["experiment_type"] == "core" - core = per_seed_df[core_mask] - non_core = per_seed_df[~core_mask] + # Cross-sprint types: seeds were added across sprints for the same config, + # so aggregate WITHOUT sprint to merge them. + cross_sprint_types = {"core", "label_efficiency"} + cross_sprint_mask = per_seed_df["experiment_type"].isin(cross_sprint_types) + cross_sprint = per_seed_df[cross_sprint_mask] + per_sprint = per_seed_df[~cross_sprint_mask] parts = [] - if len(core) > 0: - print(f" Aggregating {len(core)} core runs (cross-sprint)...", file=sys.stderr) - parts.append(_aggregate_group(core, CORE_FINGERPRINT)) - if len(non_core) > 0: - print(f" Aggregating {len(non_core)} non-core runs (per-sprint)...", file=sys.stderr) - parts.append(_aggregate_group(non_core, ABLATION_FINGERPRINT)) + if len(cross_sprint) > 0: + print( + f" Aggregating {len(cross_sprint)} cross-sprint runs (core + label_efficiency)...", + file=sys.stderr, + ) + parts.append(_aggregate_group(cross_sprint, CORE_FINGERPRINT)) + if len(per_sprint) > 0: + print( + f" Aggregating {len(per_sprint)} per-sprint runs (ablations, etc.)...", file=sys.stderr + ) + parts.append(_aggregate_group(per_sprint, ABLATION_FINGERPRINT)) if not parts: return pd.DataFrame() From cf103d76d22a0f76792a114e4cad861edefa6397 Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 1 Apr 2026 23:35:21 -0400 Subject: [PATCH 056/121] fix: implement FIXES.md Issues 2-9 for experiment integrity recovery Issue 2: Precision-aware mortality labeling. R extraction now fetches admissions.deathtime (exact timestamp) alongside patients.dod (date-only). New schema columns: death_time, death_date, death_time_precision, death_source. Label builder uses per-row precision: exact comparison for timestamps, interval-based conservative labeling for date-only deaths (boundary overlaps produce null). Backward-compatible migration for legacy data. SEMANTIC_VERSION bumped to 2.0.0. Issue 3: Label manifest freshness checking. Each label builder declares SEMANTIC_VERSION. Extraction writes label_manifest (builder_version + config_hash) to metadata.yaml. Training aborts on version/config mismatch, missing manifest, or missing task entry. Combined dataset combiner cross-checks both source manifests and aborts on disagreement. XGBoost baseline validates freshness. Issue 4: Sequence-length override fix. extract_tensors_from_dataframe detects stored vs requested seq_length mismatch and falls back to convert_raw_to_tensors for correct truncation/padding. Issue 5: Full-split normalization for low-label runs. DataModule saves full train indices before subsampling and passes them to ICUDataset for normalization, so label_fraction<1.0 runs use the same stats as full runs. Issue 6: Exporter dedup fix. Runner propagates upstream_pretrain_lr, upstream_pretrain_mask_ratio, experiment_subtype to finetune W&B config. Exporter recovers historical metadata from output_dir patterns. HP ablation finetunes reclassified from "core" to "hp_ablation". ABLATION_FINGERPRINT includes upstream pretrain dims. Row-count-by-type validation added after dedup. Issue 8: Checkpoint fail-fast. finetune.py and supervised.py default to RuntimeError on corrupted best checkpoint (configurable fallback). Both log _eval_checkpoint_source, _best_ckpt_path, _best_ckpt_load_ok, _best_ckpt_error to W&B summary. Exporter extracts all four fields. Issue 9: Hash-keyed normalization stats cache. Stats files keyed by hash of train_indices + normalize flag. Atomic writes via tempfile + os.replace. Legacy normalization_stats.yaml fallback preserved. 30 regression tests added in tests/test_fixes.py. --- scripts/export_results.py | 77 ++- .../preprocessing/create_combined_dataset.py | 63 ++ scripts/preprocessing/extract_with_ricu.R | 91 ++- scripts/run_experiments.py | 72 ++- scripts/training/finetune.py | 44 +- scripts/training/supervised.py | 46 +- scripts/training/xgboost_baseline.py | 5 + src/slices/data/datamodule.py | 14 +- src/slices/data/extractors/ricu.py | 60 ++ src/slices/data/labels/aki.py | 2 + src/slices/data/labels/base.py | 22 + src/slices/data/labels/los.py | 2 + src/slices/data/labels/mortality.py | 542 ++++++++++------- src/slices/data/tensor_cache.py | 80 ++- src/slices/data/tensor_preprocessing.py | 14 + src/slices/training/utils.py | 86 ++- tests/test_dataset_datamodule.py | 16 +- tests/test_fixes.py | 571 ++++++++++++++++++ 18 files changed, 1526 insertions(+), 281 deletions(-) create mode 100644 tests/test_fixes.py diff --git a/scripts/export_results.py b/scripts/export_results.py index a63eb2a..628e954 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -23,6 +23,7 @@ import json import math import os +import re import sys import time from pathlib import Path @@ -72,7 +73,13 @@ # For non-core runs: sprint distinguishes pretrain configs that aren't captured # in the finetune W&B config (e.g., which LR/mask_ratio the pretrain used). -ABLATION_FINGERPRINT = CORE_FINGERPRINT + ["sprint"] +# upstream_pretrain_lr/mask_ratio ensure HP ablation rows with different pretrain +# configs are not collapsed (fixes Sprint 1b, 1c, 8, 10 dedup collisions). +ABLATION_FINGERPRINT = CORE_FINGERPRINT + [ + "sprint", + "upstream_pretrain_lr", + "upstream_pretrain_mask_ratio", +] # Legacy fingerprint used for deduplication (includes sprint + all config dims). DEDUP_FINGERPRINT = [ @@ -285,6 +292,52 @@ def _load_run_data(run) -> tuple[dict, dict, str, str, str, list[str], str, str] ) +# Lookup tables for decoding run-name-encoded HP values. +# run_experiments.py encodes: str(v).replace(".", "") +# str(2e-4)="0.0002" -> "00002", str(5e-4)="0.0005" -> "00005", +# str(2e-3)="0.002" -> "0002", str(0.3)="0.3" -> "03", str(0.75)="0.75" -> "075" +_LR_DECODE = { + "00002": 2e-4, + "00005": 5e-4, + "0002": 2e-3, +} +_MR_DECODE = {"03": 0.3, "075": 0.75} + + +def _recover_pretrain_metadata( + run_name: str, config: dict +) -> tuple[float | None, float | None, str | None]: + """Recover upstream pretrain metadata from config (new runs) or run name (historical). + + Returns (upstream_pretrain_lr, upstream_pretrain_mask_ratio, experiment_subtype). + """ + # New runs have explicit fields + up_lr = config.get("upstream_pretrain_lr") + up_mr = config.get("upstream_pretrain_mask_ratio") + subtype = config.get("experiment_subtype") + if up_lr is not None or up_mr is not None: + return up_lr, up_mr, subtype + + # Historical recovery: parse from run name / output_dir + output_dir = config.get("output_dir", run_name) + + lr_match = re.search(r"_lr([^_]+)", output_dir) + if lr_match: + lr_str = lr_match.group(1) + up_lr = _LR_DECODE.get(lr_str) + if up_lr is not None: + subtype = "lr_ablation" + + mr_match = re.search(r"_mask_ratio([^_]+)", output_dir) + if mr_match: + mr_str = mr_match.group(1) + up_mr = _MR_DECODE.get(mr_str) + if up_mr is not None: + subtype = subtype or "mask_ablation" + + return up_lr, up_mr, subtype + + def extract_run(run, metric_keys: list[str]) -> dict: """Extract config + metrics from a W&B run in a single retry-protected call.""" config, summary, run_id, run_url, run_name, tags, group, created_at = _retry( @@ -353,6 +406,13 @@ def extract_run(run, metric_keys: list[str]) -> dict: if experiment_type == "core" and label_frac is not None and float(label_frac) < 1.0: experiment_type = "label_efficiency" + # Recover upstream pretrain metadata (from config for new runs, run name for historical) + up_lr, up_mr, experiment_subtype = _recover_pretrain_metadata(run_name, config) + + # Reclassify HP ablation runs that would otherwise stay as "core" + if experiment_type == "core" and experiment_subtype in ("lr_ablation", "mask_ablation"): + experiment_type = "hp_ablation" + row = { "wandb_run_id": run_id, "wandb_run_url": run_url, @@ -373,6 +433,13 @@ def extract_run(run, metric_keys: list[str]) -> dict: "source_dataset": source_dataset, "revision": config.get("revision", None), "phase": phase, + "upstream_pretrain_lr": up_lr, + "upstream_pretrain_mask_ratio": up_mr, + "experiment_subtype": experiment_subtype, + "_eval_checkpoint_source": summary.get("_eval_checkpoint_source", 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), } row.update(metrics) return row @@ -448,6 +515,14 @@ def build_per_seed_df(runs: list) -> pd.DataFrame: if sort_cols: df = df.sort_values(sort_cols, na_position="last").reset_index(drop=True) + # Row count validation by experiment_type for auditability + if "experiment_type" in df.columns: + counts = df["experiment_type"].value_counts().sort_index() + print("\n Row counts by experiment_type after dedup:", file=sys.stderr) + for etype, count in counts.items(): + print(f" {etype}: {count}", file=sys.stderr) + print(f" TOTAL: {len(df)}", file=sys.stderr) + return df diff --git a/scripts/preprocessing/create_combined_dataset.py b/scripts/preprocessing/create_combined_dataset.py index 8626fdc..dbf2dc5 100644 --- a/scripts/preprocessing/create_combined_dataset.py +++ b/scripts/preprocessing/create_combined_dataset.py @@ -217,6 +217,68 @@ def main(): tasks_b = set(meta_b.get("task_names", [])) common_tasks = sorted(tasks_a & tasks_b) + # 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", "source_datasets": names, @@ -227,6 +289,7 @@ def main(): "seq_length_hours": meta_a.get("seq_length_hours", 48), "min_stay_hours": meta_a.get("min_stay_hours", 48), "task_names": common_tasks, + "label_manifest": combined_manifest, "n_stays": len(static_combined), "stays_per_dataset": { names[0]: len(data_a["static"]), diff --git a/scripts/preprocessing/extract_with_ricu.R b/scripts/preprocessing/extract_with_ricu.R index caa4db6..4639606 100644 --- a/scripts/preprocessing/extract_with_ricu.R +++ b/scripts/preprocessing/extract_with_ricu.R @@ -615,9 +615,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 } @@ -636,6 +637,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] @@ -646,10 +648,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) { @@ -662,6 +665,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 } @@ -689,7 +724,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) { @@ -703,6 +738,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 } @@ -758,6 +815,12 @@ extract_mortality_eicu <- function(dataset, dict) { } else { NA_character_ }, + # Precision-aware schema: eICU timestamps are derived from offsets + # (minute-level precision) — these are true timestamps, not date-only. + death_time = dod, + death_date = as.Date(rep(NA, nrow(pat))), + death_time_precision = ifelse(expired, "timestamp", NA_character_), + death_source = ifelse(expired, "derived", NA_character_), stringsAsFactors = FALSE ) df @@ -769,12 +832,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 )) } @@ -782,13 +850,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 diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index 32662f0..1ca2134 100644 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -141,6 +141,10 @@ class Run: extra_overrides: dict = field(default_factory=dict) # For transfer learning: dataset the encoder was pretrained on source_dataset: str | None = None + # Upstream pretrain metadata for provenance tracking in finetune runs + upstream_pretrain_lr: float | None = None + upstream_pretrain_mask_ratio: float | None = None + experiment_subtype: str | None = None # "lr_ablation" | "mask_ablation" | None def build_command(self, runs_by_id: dict[str, Run]) -> list[str]: """Build the subprocess command for this run.""" @@ -210,6 +214,13 @@ def _finetune_cmd(self, runs_by_id: dict[str, Run]) -> list[str]: cmd.append(f"label_fraction={self.label_fraction}") for k, v in self.extra_overrides.items(): cmd.append(f"{k}={v}") + # Propagate upstream pretrain metadata so it lands in W&B config + 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}") + if self.experiment_subtype is not None: + cmd.append(f"+experiment_subtype={self.experiment_subtype}") return cmd def _supervised_cmd(self) -> list[str]: @@ -340,6 +351,9 @@ def _add_finetune( 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, ) -> Run: """Add a finetune run. @@ -347,6 +361,9 @@ def _add_finetune( extra: Hydra overrides added to the finetune command AND name. name_extra: Dict used only for directory/ID disambiguation (e.g. pretrain ablation params). Not passed to command. + upstream_pretrain_lr: LR used in the upstream pretrain run (for provenance). + upstream_pretrain_mask_ratio: Mask ratio used in upstream pretrain run. + experiment_subtype: "lr_ablation", "mask_ablation", or None. """ prefix = "probe" if freeze else "finetune" name = f"{prefix}_{paradigm}_{task}_{dataset}_seed{seed}" @@ -377,6 +394,9 @@ def _add_finetune( 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, ) self.runs.append(run) return run @@ -465,7 +485,18 @@ def build_sprint1b(self): for lr in LR_ABLATION: extra = {"optimizer.lr": lr} pt = self._add_pretrain(sprint, p, ds, seed, extra) - self._add_finetune(sprint, p, ds, seed, task, False, pt, name_extra=extra) + self._add_finetune( + sprint, + p, + ds, + seed, + task, + False, + pt, + name_extra=extra, + upstream_pretrain_lr=lr, + experiment_subtype="lr_ablation", + ) def build_sprint1c(self): """Mask ratio sensitivity, MIMIC, mortality_24h, seed=42.""" @@ -474,7 +505,18 @@ def build_sprint1c(self): for mr in MASK_RATIO_ABLATION: extra = {"ssl.mask_ratio": mr} pt = self._add_pretrain(sprint, p, ds, seed, extra) - self._add_finetune(sprint, p, ds, seed, task, False, pt, name_extra=extra) + self._add_finetune( + sprint, + p, + ds, + seed, + task, + False, + pt, + name_extra=extra, + upstream_pretrain_mask_ratio=mr, + experiment_subtype="mask_ablation", + ) def build_sprint2(self): """MIMIC Protocol A, seed=42 — reuses Sprint 1 pretrains.""" @@ -646,14 +688,32 @@ def build_sprint8(self): extra = {"optimizer.lr": lr} pt = self._add_pretrain(sprint, p, "miiv", seed, extra) self._add_finetune( - sprint, p, "miiv", seed, "mortality_24h", False, pt, name_extra=extra + sprint, + p, + "miiv", + seed, + "mortality_24h", + False, + pt, + name_extra=extra, + upstream_pretrain_lr=lr, + experiment_subtype="lr_ablation", ) # Mask ratio ablation for mr in MASK_RATIO_ABLATION: extra = {"ssl.mask_ratio": mr} pt = self._add_pretrain(sprint, p, "miiv", seed, extra) self._add_finetune( - sprint, p, "miiv", seed, "mortality_24h", False, pt, name_extra=extra + sprint, + p, + "miiv", + seed, + "mortality_24h", + False, + pt, + name_extra=extra, + upstream_pretrain_mask_ratio=mr, + experiment_subtype="mask_ablation", ) def build_sprint10(self): @@ -722,6 +782,8 @@ def build_sprint10(self): False, pt, name_extra=extra, + upstream_pretrain_lr=lr, + experiment_subtype="lr_ablation", ) for mr in MASK_RATIO_ABLATION: extra = {"ssl.mask_ratio": mr} @@ -735,6 +797,8 @@ def build_sprint10(self): False, pt, name_extra=extra, + upstream_pretrain_mask_ratio=mr, + experiment_subtype="mask_ablation", ) def build_sprint11(self): diff --git a/scripts/training/finetune.py b/scripts/training/finetune.py index 041e9ba..72d1b2b 100644 --- a/scripts/training/finetune.py +++ b/scripts/training/finetune.py @@ -72,7 +72,8 @@ def main(cfg: DictConfig) -> None: del ckpt # 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]) # Set random seed for reproducibility pl.seed_everything(cfg.seed, workers=True) @@ -84,8 +85,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, @@ -230,6 +229,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}") @@ -237,11 +239,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) # ========================================================================= @@ -267,6 +295,14 @@ def main(cfg: DictConfig) -> None: if logger: logger.experiment.summary.update(test_results[0]) + logger.experiment.summary.update( + { + "_eval_checkpoint_source": eval_checkpoint_source, + "_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/supervised.py b/scripts/training/supervised.py index 73b607d..4325337 100644 --- a/scripts/training/supervised.py +++ b/scripts/training/supervised.py @@ -138,7 +138,10 @@ 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] + ) # Set random seed for reproducibility pl.seed_everything(cfg.seed, workers=True) @@ -299,6 +302,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}") @@ -306,8 +312,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}") @@ -392,10 +424,18 @@ 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, + "_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 index 55fbd83..7e76498 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -125,9 +125,14 @@ def main(cfg: DictConfig) -> None: print("1. Setting up DataModule") print("=" * 80) + from slices.training.utils import validate_data_prerequisites + task_name = cfg.task.get("task_name", "mortality_24h") task_type = cfg.task.get("task_type", "binary") + # Validate data prerequisites including label freshness + validate_data_prerequisites(cfg.data.processed_dir, cfg.dataset, task_names=[task_name]) + datamodule = ICUDataModule( processed_dir=cfg.data.processed_dir, task_name=task_name, diff --git a/src/slices/data/datamodule.py b/src/slices/data/datamodule.py index 0054b6d..b753ba6 100644 --- a/src/slices/data/datamodule.py +++ b/src/slices/data/datamodule.py @@ -210,23 +210,27 @@ def setup(self, stage: Optional[str] = None) -> None: self.test_ratio, ) + # Save full train indices for normalization BEFORE subsampling. + # Normalization must always use the full training split to avoid + # noisy stats at low label fractions and pretrain/finetune mismatch. + normalization_train_indices = list(self.train_indices) + # Subsample training indices for label-efficiency ablations if self.label_fraction < 1.0: 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 compute_patient_level_splits. If any are still missing, - # it indicates a bug. + # Create dataset with FULL training indices for normalization. + # self.train_indices (possibly subsampled) is used by train_dataloader + # to create a Subset for optimization. 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=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 diff --git a/src/slices/data/extractors/ricu.py b/src/slices/data/extractors/ricu.py index 2972169..6cfac7b 100644 --- a/src/slices/data/extractors/ricu.py +++ b/src/slices/data/extractors/ricu.py @@ -28,6 +28,7 @@ from slices.constants import FEATURE_BLOCKLIST, LABEL_HORIZON_HOURS from slices.data.config_schemas import TimeSeriesConceptConfig +from slices.data.labels import LabelBuilder, LabelBuilderFactory from .base import BaseExtractor, ExtractorConfig @@ -161,6 +162,55 @@ def extract_data_source(self, source_name: str, stay_ids: List[int]) -> pl.DataF 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: + # Datetime column — could be true timestamp or midnight-cast date. + # For MIIV legacy data this is midnight-cast dod; for eICU it's + # a true offset-derived timestamp. Mark as "timestamp" since that + # was the prior behavior; the new R extraction fixes MIIV properly. + df = df.with_columns( + pl.col("date_of_death").cast(pl.Datetime("us", "UTC")).alias("death_time"), + 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("legacy")) + .otherwise(pl.lit(None)) + .alias("death_source"), + ) return df def run(self) -> None: @@ -317,6 +367,15 @@ def run(self) -> None: self._validate_labels(labels, stay_ids) + # Build label manifest for freshness checking at training time + label_manifest = {} + for tc in task_configs: + builder = LabelBuilderFactory.create(tc) + label_manifest[tc.task_name] = { + "builder_version": builder.SEMANTIC_VERSION, + "config_hash": LabelBuilder.config_hash(tc), + } + metadata = { "dataset": self._get_dataset_name(), "feature_set": self.config.feature_set, @@ -328,6 +387,7 @@ def run(self) -> None: "min_stay_hours": self.config.min_stay_hours, "task_names": task_names, "n_stays": len(stays_filtered), + "label_manifest": label_manifest, "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 d7b756f..8951a81 100644 --- a/src/slices/data/labels/aki.py +++ b/src/slices/data/labels/aki.py @@ -23,6 +23,8 @@ class AKILabelBuilder(LabelBuilder): Label = null if no creatinine in the prediction window. """ + SEMANTIC_VERSION = "1.0.0" + def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: """Build AKI labels from stays and timeseries data. diff --git a/src/slices/data/labels/base.py b/src/slices/data/labels/base.py index db40af9..5810914 100644 --- a/src/slices/data/labels/base.py +++ b/src/slices/data/labels/base.py @@ -1,5 +1,7 @@ """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 @@ -47,6 +49,26 @@ 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. + """ + 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_params": config.label_params, + } + 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. 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 3b2fb75..3a1f7f8 100644 --- a/src/slices/data/labels/mortality.py +++ b/src/slices/data/labels/mortality.py @@ -27,22 +27,34 @@ 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: falls back to hospital_expire_flag """ + SEMANTIC_VERSION = "2.0.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 +73,9 @@ 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 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 +86,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,36 +101,8 @@ 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 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 - # When all-null (e.g., eICU), polars may infer Null/Boolean dtype — cast to - # Datetime so downstream comparisons don't fail on type mismatch. - date_of_death_dtype = merged["date_of_death"].dtype - if date_of_death_dtype not in ( - pl.Date, - pl.Datetime, - pl.Datetime("us"), - pl.Datetime("ns"), - pl.Datetime("ms"), - ): - merged = merged.with_columns( - pl.col("date_of_death").cast(pl.Datetime("us", "UTC")).alias("date_of_death") - ) - date_of_death_dtype = pl.Datetime("us", "UTC") - 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", @@ -126,87 +112,199 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: 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) - .then(1) - .otherwise(0) - .cast(pl.Int32) - .alias("label"), - ] - ) + 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) - else: - # For DATETIME type, compare full datetime for precision - comparison = pl.col("date_of_death") <= ( - pl.col("intime") + pl.duration(hours=window_hours) + # 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.""" + if "death_time_precision" in merged.columns: + # Ensure death_time column is tz-naive datetime for comparison with + # intime/outtime (which are always tz-naive in SLICES) + 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") + ) + 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"), + ) + + # 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: + # Datetime — treat as timestamp (legacy behavior) + 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.col("date_of_death").cast(pl.Datetime("us")).alias("death_time"), + 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.lit("legacy").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"), - ] + 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) + left_icu_during_obs = pl.col("outtime") < obs_end + + # For timestamp precision: exact comparison + ts_died_during_obs = ( + pl.col("death_time").is_not_null() + & (pl.col("death_time") <= obs_end) + & left_icu_during_obs + ) + + # For date precision: conservative — use outtime check + date_died_during_obs = left_icu_during_obs & pl.col("death_date").is_not_null() + + died_during_obs = ( + pl.when(pl.col("death_time_precision") == "timestamp") + .then(ts_died_during_obs) + .when(pl.col("death_time_precision") == "date") + .then(date_died_during_obs) + .otherwise(pl.lit(False)) + ) + + return merged.select( + [ + "stay_id", + pl.when(died_during_obs) + .then(None) + .when(pl.col("hospital_expire_flag") == 1) + .then(1) + .otherwise(0) + .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.""" + ts_died = pl.col("death_time").is_not_null() & (pl.col("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(pl.col("death_time_precision") == "timestamp") + .then(ts_died) + .when(pl.col("death_time_precision") == "date") + .then(date_died) + .otherwise(pl.lit(False)) + ) + + return merged.select( + [ + "stay_id", + pl.when(died_in_icu).then(1).otherwise(0).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) + + ts_died = pl.col("death_time").is_not_null() & (pl.col("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) + ) ) + ) - return labels + died = ( + pl.when(pl.col("death_time_precision") == "timestamp") + .then(ts_died) + .when(pl.col("death_time_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(False)) + ) + + return merged.select( + [ + "stay_id", + pl.when(died.is_null()) + .then(None) + .when(died) + .then(1) + .otherwise(0) + .cast(pl.Int32) + .alias("label"), + ] + ) def _build_windowed_mortality_labels( self, @@ -214,7 +312,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. @@ -222,148 +319,135 @@ 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: null for windowed tasks """ - # 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 - # 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) + left_icu_during_obs = pl.col("outtime") < obs_end - # 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 = ( + pl.col("death_time").is_not_null() + & (pl.col("death_time") <= obs_end) + & left_icu_during_obs + ) + ts_died_during_gap = ( + ( + pl.col("death_time").is_not_null() + & (pl.col("death_time") > obs_end) + & (pl.col("death_time") < pred_start) ) + if gap_hours > 0 + else pl.lit(False) + ) + ts_died_during_pred = ( + pl.col("death_time").is_not_null() + & (pl.col("death_time") >= pred_start) + & (pl.col("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 <= obs_end + date_died_during_obs = ( + pl.col("death_date").is_not_null() & (date_end <= obs_end) & left_icu_during_obs + ) - # 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) + & left_icu_during_obs + ) - # 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) + ) + + # 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) + ) - 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) + # 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(pl.col("death_time_precision") == "timestamp") + .then(ts_label) + .when(pl.col("death_time_precision") == "date") + .then(date_label) + .otherwise( + # Unknown precision or no death info → fall back + pl.when(pl.col("hospital_expire_flag") == 1) + .then(None) # can't resolve timing → exclude + .otherwise(0) # not flagged as dead → survived ) + ) - # 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/tensor_cache.py b/src/slices/data/tensor_cache.py index 1a97eda..73eead9 100644 --- a/src/slices/data/tensor_cache.py +++ b/src/slices/data/tensor_cache.py @@ -7,6 +7,8 @@ import hashlib import logging +import os +import tempfile import warnings from pathlib import Path from typing import Any, Dict, List, Optional @@ -17,6 +19,21 @@ logger = logging.getLogger(__name__) +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"{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]], @@ -24,11 +41,11 @@ def load_normalization_stats( ) -> 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. + 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.yaml. + 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. @@ -40,15 +57,33 @@ def load_normalization_stats( if not normalize: return None - stats_path = data_dir / "normalization_stats.yaml" - if not stats_path.exists(): + # 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) + 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(stats_path) as f: + with open(legacy_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 @@ -66,7 +101,7 @@ def load_normalization_stats( return stats except Exception as e: warnings.warn( - f"Failed to load normalization stats from {stats_path}: {e}. " + f"Failed to load normalization stats from {legacy_path}: {e}. " "Will recompute statistics.", UserWarning, ) @@ -81,7 +116,11 @@ def save_normalization_stats( train_indices: Optional[List[int]], normalize: bool, ) -> None: - """Save computed normalization statistics to file for reproducibility. + """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. @@ -91,19 +130,34 @@ def save_normalization_stats( train_indices: Optional list of training indices used to compute stats. normalize: Whether normalization is enabled. """ - stats_path = data_dir / "normalization_stats.yaml" + 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, - "train_indices": train_indices, + "split_hash": split_hash, + "train_indices_count": len(train_indices) if train_indices else 0, "normalize": normalize, } try: - with open(stats_path, "w") as f: - yaml.dump(stats, f, default_flow_style=False) + 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}", diff --git a/src/slices/data/tensor_preprocessing.py b/src/slices/data/tensor_preprocessing.py index ced51dd..ac941e4 100644 --- a/src/slices/data/tensor_preprocessing.py +++ b/src/slices/data/tensor_preprocessing.py @@ -60,6 +60,20 @@ def extract_tensors_from_dataframe( 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 diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 0a302ed..37c37fc 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -366,12 +366,26 @@ def run_fairness_evaluation( # ============================================================================= -def validate_data_prerequisites(processed_dir: str, dataset: str) -> None: +def validate_data_prerequisites( + processed_dir: str, + dataset: str, + task_names: Optional[List[str]] = None, +) -> None: """Validate that required data files exist before training. + Checks file existence and, if task_names are provided, validates the label + manifest in metadata.yaml to ensure labels were built with the current + builder version and task config. + Raises: FileNotFoundError: If required files are missing. + RuntimeError: If label manifest indicates stale labels. """ + + import yaml + + from slices.data.labels import LabelBuilder, LabelBuilderFactory, LabelConfig + path = Path(processed_dir) if not path.exists(): @@ -387,12 +401,70 @@ 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}" - ) + # Validate label manifest if task_names are provided + if task_names: + metadata_path = path / "metadata.yaml" + if not metadata_path.exists(): + 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}" + ) + + with open(metadata_path) as f: + metadata = yaml.safe_load(f) + + 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}" + ) + + # Load current task configs and compare against manifest + tasks_path = Path(__file__).parent.parent / "data" / "tasks" + for task_name in task_names: + config_file = tasks_path / f"{task_name}.yaml" + if not config_file.exists(): + continue + + with open(config_file) as f: + config_dict = yaml.safe_load(f) + current_config = LabelConfig(**config_dict) + 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/test_dataset_datamodule.py b/tests/test_dataset_datamodule.py index ea10b6f..811bea6 100644 --- a/tests/test_dataset_datamodule.py +++ b/tests/test_dataset_datamodule.py @@ -1019,8 +1019,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() @@ -1254,10 +1255,13 @@ def test_no_inline_import_warnings(self): 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_fixes.py b/tests/test_fixes.py new file mode 100644 index 0000000..9b63c35 --- /dev/null +++ b/tests/test_fixes.py @@ -0,0 +1,571 @@ +"""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. +""" + +from datetime import datetime, timedelta + +import polars as pl +import pytest +import torch +import yaml +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, + load_normalization_stats, + save_normalization_stats, +) +from slices.data.tensor_preprocessing import ( + extract_tensors_from_dataframe, +) + +# ============================================================================ +# 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=48, + ) + config_48 = LabelConfig( + task_name="mortality_48h", + task_type="binary", + prediction_window_hours=48, + observation_window_hours=48, + ) + 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_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.training.utils import validate_data_prerequisites + + 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 = 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 — but we need the task YAML to exist for validation. + # Since validate_data_prerequisites looks for configs/tasks/*.yaml, + # it will skip tasks where the config file doesn't exist. + # So this test verifies the no-crash path. + validate_data_prerequisites(str(tmp_path), "miiv", task_names=["mortality_24h"]) + + +# ============================================================================ +# 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, + } + 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] + + +# ============================================================================ +# 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_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_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 data with only date_of_death should be migrated correctly.""" + 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] == "timestamp" + + +# ============================================================================ +# Issue 6: Exporter LR/mask_ratio decode and historical recovery +# ============================================================================ + + +class TestExporterHistoricalRecovery: + """Tests for _recover_pretrain_metadata run-name parsing.""" + + 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): + """Historical LR ablation finetunes recovered from output_dir pattern.""" + from scripts.export_results import _recover_pretrain_metadata + + config = {"output_dir": "outputs/sprint10/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_ablation" + + def test_recover_mask_ratio_from_output_dir(self): + """Historical mask_ratio ablation finetunes recovered from output_dir.""" + from scripts.export_results import _recover_pretrain_metadata + + config = { + "output_dir": "outputs/sprint1c/finetune_jepa_mortality_24h_miiv_seed42_mask_ratio03" + } + up_lr, up_mr, subtype = _recover_pretrain_metadata("some_run", config) + assert up_mr == 0.3 + assert subtype == "mask_ablation" + + 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_ablation"} + up_lr, up_mr, subtype = _recover_pretrain_metadata("irrelevant_name", config) + assert up_lr == 5e-4 + assert subtype == "lr_ablation" + + 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/sprint1/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 From c05c28e9589ff235e51a3ca5fb02234361cfe391 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 7 Apr 2026 12:13:49 -0400 Subject: [PATCH 057/121] Add agent files to gitignore --- .gitignore | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.gitignore b/.gitignore index 2634ff6..d353468 100644 --- a/.gitignore +++ b/.gitignore @@ -195,6 +195,8 @@ cython_debug/ .pypirc # AI +.agents/ +.codex/ .cursorignore .cursorindexingignore .cursorrules From 930ce5eb9a0c21fb99363967659a5a87a830ab07 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 7 Apr 2026 12:49:10 -0400 Subject: [PATCH 058/121] fix: harden preprocessing cache freshness and experiment cleanup --- README.md | 2 +- pyproject.toml | 4 + .../preprocessing/create_combined_dataset.py | 70 ++++++++- scripts/preprocessing/prepare_dataset.py | 8 ++ src/slices/data/tensor_cache.py | 134 +++++++++++++++++- src/slices/debug/plots.py | 8 +- src/slices/models/pretraining/README.md | 14 +- src/slices/training/README.md | 4 +- src/slices/training/utils.py | 5 +- tests/README.md | 14 +- tests/test_fixes.py | 117 ++++++++++++++- uv.lock | 17 +++ 12 files changed, 361 insertions(+), 36 deletions(-) diff --git a/README.md b/README.md index 8229a0d..3d32719 100644 --- a/README.md +++ b/README.md @@ -173,7 +173,7 @@ SLICES/ ├── scripts/ # Entry point scripts │ ├── preprocessing/ │ └── training/ -└── tests/ # pytest test suite (940+ tests) +└── tests/ # pytest test suite (900+ tests) ``` ## Configuration diff --git a/pyproject.toml b/pyproject.toml index 2730cea..08fb326 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,7 +78,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/scripts/preprocessing/create_combined_dataset.py b/scripts/preprocessing/create_combined_dataset.py index dbf2dc5..8f5f8cf 100644 --- a/scripts/preprocessing/create_combined_dataset.py +++ b/scripts/preprocessing/create_combined_dataset.py @@ -75,13 +75,22 @@ def validate_no_id_collision(static_a: pl.DataFrame, static_b: pl.DataFrame) -> 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 +98,56 @@ 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 merge_labels(labels_a: pl.DataFrame, labels_b: pl.DataFrame) -> pl.DataFrame: """Merge label DataFrames, handling column mismatches.""" @@ -164,6 +223,7 @@ def main(): # 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 diff --git a/scripts/preprocessing/prepare_dataset.py b/scripts/preprocessing/prepare_dataset.py index b3ff47c..ea97402 100644 --- a/scripts/preprocessing/prepare_dataset.py +++ b/scripts/preprocessing/prepare_dataset.py @@ -22,6 +22,10 @@ import yaml from omegaconf import DictConfig 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: @@ -269,6 +273,10 @@ def main(cfg: DictConfig) -> None: feature_names=feature_names, seq_length=seq_length, ) + stats["normalize"] = True + stats["split_hash"] = None + stats["data_fingerprint"] = get_data_fingerprint(processed_dir) + stats["preprocessing_fingerprint"] = get_preprocessing_fingerprint() # Save normalization stats stats_path = processed_dir / "normalization_stats.yaml" diff --git a/src/slices/data/tensor_cache.py b/src/slices/data/tensor_cache.py index 73eead9..ace6bee 100644 --- a/src/slices/data/tensor_cache.py +++ b/src/slices/data/tensor_cache.py @@ -6,10 +6,13 @@ """ 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 @@ -17,6 +20,90 @@ import yaml logger = logging.getLogger(__name__) +_CACHE_FINGERPRINT_VERSION = "2026-04-07" + + +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: @@ -57,6 +144,9 @@ def load_normalization_stats( if not normalize: return None + 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) @@ -65,6 +155,13 @@ def load_normalization_stats( 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: @@ -98,6 +195,14 @@ def load_normalization_stats( ) 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( @@ -143,7 +248,10 @@ def save_normalization_stats( "feature_names": feature_names, "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: @@ -210,7 +318,14 @@ def get_tensor_cache_path( return cache_dir / f"tensors_{cache_key}.pt" -def _try_load_cache(cache_path: Path, n_features: int, seq_length: int) -> Optional[Dict[str, Any]]: +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: @@ -235,6 +350,13 @@ def _try_load_cache(cache_path: Path, n_features: int, seq_length: int) -> Optio 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 @@ -285,7 +407,13 @@ def load_cached_tensors( 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) + 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( @@ -316,6 +444,8 @@ def save_cached_tensors( "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: 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/models/pretraining/README.md b/src/slices/models/pretraining/README.md index 619b362..4644fec 100644 --- a/src/slices/models/pretraining/README.md +++ b/src/slices/models/pretraining/README.md @@ -29,7 +29,7 @@ mae = MAEObjective(encoder, mae_config) loss, metrics = mae(x, obs_mask) ``` -**See**: `docs/MAE_IMPLEMENTATION.md` for detailed documentation. +**See**: the objective source files in this package and the Hydra configs under `configs/ssl/`. ## Factory Pattern @@ -58,6 +58,9 @@ pretraining/ ├── base.py # Abstract base classes ├── factory.py # Factory functions and registry ├── mae.py # MAE implementation +├── jepa.py # JEPA implementation +├── contrastive.py # Contrastive implementation +├── smart.py # SMART sanity-check objective └── README.md # This file ``` @@ -68,14 +71,9 @@ All SSL objectives require: - Input shape: `(B, T, D)` - Observation mask shape: `(B, T, D)` -## Future Objectives +## Implemented Objectives -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 +The active codebase already includes MAE, JEPA, contrastive, and SMART objectives. ## References diff --git a/src/slices/training/README.md b/src/slices/training/README.md index 710faa3..e927ab9 100644 --- a/src/slices/training/README.md +++ b/src/slices/training/README.md @@ -197,7 +197,7 @@ hidden_size: 256 4. Use in pretraining: ```bash -uv run python scripts/pretrain.py encoder=my_encoder +uv run python scripts/training/pretrain.py encoder=my_encoder ``` ### Adding a New SSL Objective @@ -241,7 +241,7 @@ name: my_ssl 4. Use in pretraining: ```bash -uv run python scripts/pretrain.py ssl=my_ssl +uv run python scripts/training/pretrain.py ssl=my_ssl ``` ### `FineTuneModule` diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 37c37fc..e898fe1 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -428,7 +428,10 @@ def validate_data_prerequisites( for task_name in task_names: config_file = tasks_path / f"{task_name}.yaml" if not config_file.exists(): - continue + 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) diff --git a/tests/README.md b/tests/README.md index 3c5f913..2f8304f 100644 --- a/tests/README.md +++ b/tests/README.md @@ -16,17 +16,17 @@ 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 +uv run pytest tests/test_dataset_datamodule.py -v ``` ### Run specific test class ```bash -uv run pytest tests/test_data_io.py::TestCsvToParquetAll -v +uv run pytest tests/test_dataset_datamodule.py::TestICUDataset -v ``` ### Run specific test function ```bash -uv run pytest tests/test_data_io.py::TestCsvToParquetAll::test_successful_conversion -v +uv run pytest tests/test_fixes.py::TestNormalizationStatsCache::test_legacy_stats_without_fingerprints_are_ignored -v ``` ### Run tests matching a pattern @@ -47,18 +47,18 @@ uv run pytest tests/ -vv | `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 | +| `test_ricu_extractor.py` | RICU extractor compatibility | Chunked parquet ingestion, schema migration, extraction edge cases | +| `test_fixes.py` | Regression coverage for experiment-integrity fixes | Label freshness, cache invalidation, combined-dataset safety | ## Test Categories ### Unit Tests - `test_extractor_config.py` - Configuration validation - `test_base_extractor.py` - Core extractor methods -- `test_data_io.py` - Data I/O utilities +- `test_transforms.py` - SSL/data augmentation utilities - `test_task_builders.py` - Task builders ### Integration Tests @@ -66,7 +66,7 @@ uv run pytest tests/ -vv - `test_dataset_datamodule.py` - Data loading pipeline ### Edge Case Tests -- `test_timeseries_extraction.py::TestTimeSeriesEdgeCases` - Empty data, extreme values +- `test_ricu_extractor.py::TestRicuExtractorEdgeCases` - Missing schema pieces, chunked input, legacy migration - `test_task_builders.py::TestMortalityBoundaryConditions` - Boundary conditions ## Writing New Tests diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 9b63c35..e69a0d7 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -17,12 +17,14 @@ 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, -) +from slices.data.tensor_preprocessing import extract_tensors_from_dataframe # ============================================================================ # Issue 3: Label manifest freshness checking @@ -171,12 +173,32 @@ def test_validate_data_prerequisites_matching_passes(self, tmp_path): with open(tmp_path / "metadata.yaml", "w") as f: yaml.dump(metadata, f) - # Should not raise — but we need the task YAML to exist for validation. - # Since validate_data_prerequisites looks for configs/tasks/*.yaml, - # it will skip tasks where the config file doesn't exist. - # So this test verifies the no-crash path. + # 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_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 @@ -250,6 +272,8 @@ def test_legacy_fallback(self, tmp_path): "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) @@ -258,6 +282,85 @@ def test_legacy_fallback(self, tmp_path): 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 + + +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) + # ============================================================================ # Issue 4: Sequence-length override diff --git a/uv.lock b/uv.lock index 79711f7..6f232dc 100644 --- a/uv.lock +++ b/uv.lock @@ -3642,9 +3642,13 @@ dev = [ [package.dev-dependencies] dev = [ + { name = "black" }, + { name = "mypy" }, { name = "pytest" }, + { name = "pytest-cov" }, { name = "python-pptx" }, { name = "ruff" }, + { name = "types-pyyaml" }, ] [package.metadata] @@ -3677,9 +3681,13 @@ provides-extras = ["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]] @@ -3906,6 +3914,15 @@ 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" }, ] +[[package]] +name = "types-pyyaml" +version = "6.0.12.20250915" +source = { registry = "https://pypi.org/simple" } +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/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]] name = "typing-extensions" version = "4.15.0" From 2f5231635f48b033bb06128ec387a275216b5e69 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 7 Apr 2026 13:21:33 -0400 Subject: [PATCH 059/121] feat: add thesis rerun launcher and wandb targeting --- scripts/export_results.py | 69 +++++++++++++--- scripts/launch_thesis_tmux.sh | 148 ++++++++++++++++++++++++++++++++++ scripts/run_experiments.py | 57 ++++++++++++- tests/test_fixes.py | 80 ++++++++++++++++++ 4 files changed, 341 insertions(+), 13 deletions(-) create mode 100755 scripts/launch_thesis_tmux.sh diff --git a/scripts/export_results.py b/scripts/export_results.py index 628e954..807a0b2 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -16,6 +16,7 @@ uv run python scripts/export_results.py --sprint 1 --validate-seeds 3 uv run python scripts/export_results.py --paradigm mae jepa --dataset miiv uv run python scripts/export_results.py --output-dir results/sprint1 --sprint 1 + uv run python scripts/export_results.py --project slices-thesis --revision thesis-v1 """ from __future__ import annotations @@ -54,8 +55,11 @@ "7p": "capacity_pilot", "11": "classical_baselines", "12": "smart_reference", + "13": "temporal_contrastive", } +CROSS_SPRINT_TYPES = {"core", "label_efficiency", "transfer", "hp_ablation"} + # For core runs: sprint is noise — same config across sprints, different seeds. # lr/mask_ratio excluded because they're determined by protocol (redundant) or # only in the pretrain config (NaN for finetune runs). @@ -146,9 +150,12 @@ ALL_METRICS = TEST_METRICS + VAL_METRICS -# d_model values that indicate non-default capacity (Sprint 7p pilot). -# Default d_model=128 maps to "default"; other values are named here. -SIZED_MODELS = {128: "default", 256: "large"} +# Canonical model variants used in the thesis matrix. +MODEL_VARIANTS = { + (64, 2): "default", + (128, 4): "medium", + (256, 4): "large", +} # Phases that correspond to evaluation runs (not pretraining). EVAL_PHASES = ["finetune", "supervised", "gru_d", "xgboost", "baseline"] @@ -167,6 +174,7 @@ def fetch_all_runs( 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 filtering. @@ -197,6 +205,9 @@ def fetch_all_runs( if phase and len(phase) == 1: tag_filters.append(f"phase:{phase[0]}") + if revision and len(revision) == 1: + tag_filters.append(f"revision:{revision[0]}") + if tag_filters: filters["tags"] = {"$all": tag_filters} @@ -210,6 +221,7 @@ def fetch_all_runs( paradigm_set = set(paradigm) if paradigm and len(paradigm) > 1 else None dataset_set = set(dataset) if dataset and len(dataset) > 1 else None phase_set = set(phase) if phase and len(phase) > 1 else None + revision_set = set(revision) if revision and len(revision) > 1 else None runs = [] for run in runs_iter: @@ -227,6 +239,9 @@ def fetch_all_runs( if phase_set: if not any(f"phase:{p}" in tags for p in phase_set): continue + if revision_set: + if not any(f"revision:{r}" in tags for r in revision_set): + continue runs.append(run) @@ -338,6 +353,32 @@ def _recover_pretrain_metadata( return up_lr, up_mr, subtype +def _infer_model_size(config: dict) -> str: + """Infer model-size label from the encoder config.""" + d_model = _get_nested(config, "encoder.d_model") + n_layers = _get_nested(config, "encoder.n_layers") + sprint = str(config.get("sprint", "")) + + if d_model is None: + return "default" + + variant = MODEL_VARIANTS.get((d_model, n_layers)) + if variant is not None: + return variant + + if sprint == "7p": + if d_model == 128: + return "medium" + if d_model == 256: + return "large" + + 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 extract_run(run, metric_keys: list[str]) -> dict: """Extract config + metrics from a W&B run in a single retry-protected call.""" config, summary, run_id, run_url, run_name, tags, group, created_at = _retry( @@ -353,13 +394,7 @@ def extract_run(run, metric_keys: list[str]) -> dict: else: protocol = None - # Derive model_size from encoder.d_model - d_model = _get_nested(config, "encoder.d_model") - if d_model is not None and d_model != 128: - # Sprint 7p capacity pilot: medium=128, large=256 - model_size = SIZED_MODELS.get(d_model, f"d{d_model}") - else: - model_size = "default" + model_size = _infer_model_size(config) # Detect source_dataset for transfer runs from run name or config source_dataset = config.get("source_dataset", None) @@ -413,6 +448,10 @@ def extract_run(run, metric_keys: list[str]) -> dict: if experiment_type == "core" and experiment_subtype in ("lr_ablation", "mask_ablation"): experiment_type = "hp_ablation" + # Sprint 10 also adds transfer seeds for Sprint 7 scope. + if experiment_type == "core" and source_dataset is not None: + experiment_type = "transfer" + row = { "wandb_run_id": run_id, "wandb_run_url": run_url, @@ -493,7 +532,7 @@ def build_per_seed_df(runs: list) -> pd.DataFrame: df = df.sort_values("created_at", ascending=False) before = len(df) - cross_sprint_types = {"core", "label_efficiency"} + cross_sprint_types = CROSS_SPRINT_TYPES cross_sprint_mask = df["experiment_type"].isin(cross_sprint_types) cross_sprint_dedup_cols = [c for c in CORE_FINGERPRINT + ["seed"] if c in df.columns] ablation_dedup_cols = [c for c in ABLATION_FINGERPRINT + ["seed"] if c in df.columns] @@ -587,7 +626,7 @@ def build_aggregated_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: """ # Cross-sprint types: seeds were added across sprints for the same config, # so aggregate WITHOUT sprint to merge them. - cross_sprint_types = {"core", "label_efficiency"} + cross_sprint_types = CROSS_SPRINT_TYPES cross_sprint_mask = per_seed_df["experiment_type"].isin(cross_sprint_types) cross_sprint = per_seed_df[cross_sprint_mask] per_sprint = per_seed_df[~cross_sprint_mask] @@ -723,6 +762,11 @@ def main(): 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), e.g. --revision thesis-v1", + ) parser.add_argument( "--state", default="finished", @@ -753,6 +797,7 @@ def main(): paradigm=args.paradigm, dataset=args.dataset, phase=args.phase, + revision=args.revision, ) if not runs: diff --git a/scripts/launch_thesis_tmux.sh b/scripts/launch_thesis_tmux.sh new file mode 100755 index 0000000..32a5332 --- /dev/null +++ b/scripts/launch_thesis_tmux.sh @@ -0,0 +1,148 @@ +#!/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:-2026-04-07 post-cache-freshness final rerun}" +PARALLEL_MAIN="${PARALLEL_MAIN:-4}" +PARALLEL_APPENDIX="${PARALLEL_APPENDIX:-4}" +BATCH_SIZE_FAIRNESS="${BATCH_SIZE_FAIRNESS:-64}" +DEVICE_FAIRNESS="${DEVICE_FAIRNESS:-auto}" +INCLUDE_SPRINT12="${INCLUDE_SPRINT12:-0}" +INCLUDE_SPRINT13="${INCLUDE_SPRINT13:-1}" +INCLUDE_SPRINT7P="${INCLUDE_SPRINT7P:-0}" +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 + +mkdir -p "$LOG_DIR" + +main_sprints=(1 1b 1c 2 3 4 5 6 7 8 10 11) +fairness_sprints=(1 2 3 4 5 10) +tag_sprints=(1b 1c 2 5 6 7 8) +warmup_sprints=("${main_sprints[@]}") +appendix_sprints=() + +if [[ "$INCLUDE_SPRINT13" == "1" ]]; then + main_sprints+=(13) + fairness_sprints+=(13) +fi + +if [[ "$INCLUDE_SPRINT7P" == "1" ]]; then + main_sprints+=(7p) + fairness_sprints+=(7p) + tag_sprints+=(7p) +fi + +if [[ "$INCLUDE_SPRINT12" == "1" ]]; then + appendix_sprints+=(12) + fairness_sprints+=(12) + warmup_sprints+=(12) +fi + +quote_cmd() { + printf "%q " "$@" +} + +run_args=(--project "$WANDB_PROJECT") +export_args=(--project "$WANDB_PROJECT") +fairness_args=(--project "$WANDB_PROJECT") +tag_args=(--project "$WANDB_PROJECT") + +if [[ -n "$WANDB_ENTITY" ]]; then + run_args+=(--entity "$WANDB_ENTITY") + export_args+=(--entity "$WANDB_ENTITY") + fairness_args+=(--entity "$WANDB_ENTITY") + tag_args+=(--entity "$WANDB_ENTITY") +fi + +run_args+=(--revision "$REVISION" --reason "$REASON") +export_args+=(--revision "$REVISION" --output-dir "$RESULTS_DIR") +fairness_args+=(--sprint "${fairness_sprints[@]}" --batch-size "$BATCH_SIZE_FAIRNESS" --device "$DEVICE_FAIRNESS") + +warmup_cmd=(uv run python scripts/run_experiments.py warmup --sprint "${warmup_sprints[@]}") +main_cmd=(uv run python scripts/run_experiments.py run --sprint "${main_sprints[@]}" --parallel "$PARALLEL_MAIN" "${run_args[@]}") +tag_cmd=(uv run python scripts/run_experiments.py tag --sprint "${tag_sprints[@]}" "${tag_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_sprints[@]} > 0)); then + appendix_cmd=(uv run python scripts/run_experiments.py run --sprint "${appendix_sprints[@]}" --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 tag_line "%q " "${tag_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" <&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/run_experiments.py b/scripts/run_experiments.py index 1ca2134..a25ddc5 100644 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -10,6 +10,8 @@ uv run python scripts/run_experiments.py run --sprint 1 --parallel 4 uv run python scripts/run_experiments.py run --sprint 1 2 3 --parallel 6 --dry-run uv run python scripts/run_experiments.py run --sprint 1 --revision v2 --reason "fix LR" + uv run python scripts/run_experiments.py run --sprint 1 2 \\ + --project slices-thesis --entity myteam uv run python scripts/run_experiments.py status uv run python scripts/run_experiments.py status --sprint 1 uv run python scripts/run_experiments.py retry --failed --parallel 4 @@ -935,6 +937,24 @@ def apply_revision(runs: list[Run], revision: str, reason: str | None = None) -> 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 + + # --------------------------------------------------------------------------- # State Management # --------------------------------------------------------------------------- @@ -1320,10 +1340,14 @@ def cmd_run(args): # Apply revision transform if requested if args.revision: runs = apply_revision(runs, args.revision, args.reason) + runs = apply_wandb_target(runs, args.project, args.entity) print(f"Sprint(s) {', '.join(sprints)}: {len(runs)} runs") if args.revision: print(f"Revision: {args.revision}" + (f" ({args.reason})" if args.reason else "")) + if args.project or args.entity: + target = f"{args.entity}/{args.project}" if args.entity else args.project + print(f"W&B target: {target}") run_scheduler(runs, state, args.parallel, args.dry_run) @@ -1345,6 +1369,7 @@ def cmd_retry(args): all_runs = rest + revised else: all_runs = apply_revision(all_runs, args.revision, args.reason) + all_runs = apply_wandb_target(all_runs, args.project, args.entity) state = load_state() recover_stale_running(state) @@ -1376,6 +1401,9 @@ def cmd_retry(args): save_state(state) print(f"Retrying {len(runs_to_retry)} 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}") run_scheduler(runs_to_retry, state, args.parallel, dry_run=False) @@ -1532,6 +1560,18 @@ def main(): 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 (e.g. v2)") p_run.add_argument("--reason", type=str, default=None, help="Reason for rerun") + 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)", + ) # status p_status = sub.add_parser("status", help="Show experiment status") @@ -1545,6 +1585,18 @@ def main(): p_retry.add_argument("--sprint", nargs="+", default=None, help="Scope retry to sprint(s)") 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( + "--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)", + ) # warmup p_warmup = sub.add_parser( @@ -1568,7 +1620,10 @@ def main(): ) p_tag.add_argument("--dry-run", action="store_true", help="Show what would be tagged") p_tag.add_argument( - "--project", type=str, default="slices", help="W&B project name (default: slices)" + "--project", + type=str, + default=os.environ.get("WANDB_PROJECT", "slices"), + help="W&B project name (default: WANDB_PROJECT env var or 'slices')", ) p_tag.add_argument( "--entity", diff --git a/tests/test_fixes.py b/tests/test_fixes.py index e69a0d7..a9b98aa 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -672,3 +672,83 @@ def test_core_run_returns_none(self): 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 Sprint 7p 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({"sprint": "7p", "encoder": {"d_model": 128, "n_layers": 4}}) == ( + "medium" + ) + assert _infer_model_size({"sprint": "7p", "encoder": {"d_model": 256, "n_layers": 4}}) == ( + "large" + ) + + def test_extract_run_reclassifies_transfer_and_sprint13(self): + """Sprint-10 transfer seeds and Sprint 13 should export with stable experiment types.""" + from scripts.export_results import extract_run + + class DummyRun: + def __init__(self, config, tags, name): + self.config = config + self.summary_metrics = {} + self.id = "dummy" + self.url = "https://example.com" + self.name = name + self.tags = tags + self.group = "g" + self.created_at = "2026-04-07T00:00:00" + + transfer = DummyRun( + config={ + "sprint": "10", + "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_type"] == "transfer" + + ts2vec = DummyRun( + config={ + "sprint": "13", + "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_type"] == "temporal_contrastive" + + +class TestExperimentRunnerWandbOverrides: + """Tests for clean project/entity overrides in the experiment runner.""" + + def test_apply_wandb_target_injects_project_and_entity(self): + from scripts.run_experiments import Run, apply_wandb_target + + run = Run( + id="s1_supervised_mortality_24h_miiv_seed42", + sprint="1", + run_type="supervised", + paradigm="supervised", + dataset="miiv", + seed=42, + output_dir="outputs/sprint1/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" From 4a329c7ef10ae7ab14130dd74e49ac7908145426 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 7 Apr 2026 14:48:27 -0400 Subject: [PATCH 060/121] feat: expand thesis sprint coverage --- scripts/eval/evaluate_fairness.py | 14 +-- scripts/export_results.py | 30 +++++-- scripts/launch_thesis_tmux.sh | 7 +- scripts/run_experiments.py | 141 ++++++++++++++++-------------- tests/test_fixes.py | 28 ++++++ 5 files changed, 137 insertions(+), 83 deletions(-) diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index 36ea5e5..63960e3 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -5,12 +5,12 @@ 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 core experiment runs (~480 runs). Supports +Designed for batch evaluation of the thesis fairness corpus. Supports resumability via --skip-existing (default) and scoping via --sprint/--paradigm/ --dataset filters. Usage: - # Evaluate all core runs (default: sprints 1-5, 10) + # Evaluate the default thesis fairness corpus uv run python scripts/eval/evaluate_fairness.py # Scope to specific sprint/dataset @@ -49,7 +49,7 @@ # Constants # --------------------------------------------------------------------------- -CORE_SPRINTS = ["1", "2", "3", "4", "5", "10"] +CORE_SPRINTS = ["1", "2", "3", "4", "5", "7p", "10", "13"] DEFAULT_PHASES = ["finetune", "supervised"] DEFAULT_PROTECTED_ATTRIBUTES = ["gender", "age_group", "race"] @@ -408,7 +408,11 @@ def parse_args() -> argparse.Namespace: default=os.environ.get("WANDB_ENTITY"), help="W&B entity (default: $WANDB_ENTITY)", ) - parser.add_argument("--sprint", nargs="+", help="Filter to sprint(s). Default: core sprints") + parser.add_argument( + "--sprint", + nargs="+", + help="Filter to sprint(s). Default: thesis fairness sprints", + ) parser.add_argument("--paradigm", nargs="+", help="Filter to paradigm(s)") parser.add_argument("--dataset", nargs="+", help="Filter to dataset(s)") parser.add_argument( @@ -467,7 +471,7 @@ def main() -> None: device = _resolve_device(args.device) log.info("Device: %s", device) - # Default to core sprints if not specified + # Default to the thesis fairness sprint set if not specified sprints = args.sprint or CORE_SPRINTS log.info("Sprints: %s", sprints) diff --git a/scripts/export_results.py b/scripts/export_results.py index 807a0b2..e8abb30 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -59,6 +59,16 @@ } CROSS_SPRINT_TYPES = {"core", "label_efficiency", "transfer", "hp_ablation"} +FIXED_SEED_EXPERIMENT_TYPES = { + "core", + "label_efficiency", + "transfer", + "hp_ablation", + "capacity_pilot", + "classical_baselines", + "smart_reference", + "temporal_contrastive", +} # For core runs: sprint is noise — same config across sprints, different seeds. # lr/mask_ratio excluded because they're determined by protocol (redundant) or @@ -634,7 +644,7 @@ def build_aggregated_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: parts = [] if len(cross_sprint) > 0: print( - f" Aggregating {len(cross_sprint)} cross-sprint runs (core + label_efficiency)...", + f" Aggregating {len(cross_sprint)} cross-sprint runs (merged experiment families)...", file=sys.stderr, ) parts.append(_aggregate_group(cross_sprint, CORE_FINGERPRINT)) @@ -671,18 +681,19 @@ def validate( ) -> list[str]: """Validate results and return warning strings. - Only warns about core configs with fewer seeds than expected. - Non-core configs have variable seed counts by design. + Warns when fixed-seed experiment families have fewer seeds than expected. """ warnings = [] - # Check core configs for expected seed count + # Check fixed-seed experiment families for expected seed count if "experiment_type" in aggregated_df.columns: - core = aggregated_df[aggregated_df["experiment_type"] == "core"] - low_seed = core[core["n_seeds"] < expected_core_seeds] + fixed_seed = aggregated_df[ + aggregated_df["experiment_type"].isin(FIXED_SEED_EXPERIMENT_TYPES) + ] + low_seed = fixed_seed[fixed_seed["n_seeds"] < expected_core_seeds] if len(low_seed) > 0: warnings.append( - f"WARNING: {len(low_seed)}/{len(core)} core configs have fewer " + f"WARNING: {len(low_seed)}/{len(fixed_seed)} fixed-seed configs have fewer " f"than {expected_core_seeds} seeds:" ) for _, row in low_seed.iterrows(): @@ -692,7 +703,10 @@ def validate( if c in row and row[c] is not None ) seeds = json.loads(row["seed_list"]) if pd.notna(row.get("seed_list")) else [] - warnings.append(f" {desc} — n_seeds={row['n_seeds']}, seeds={seeds}") + warnings.append( + f" experiment_type={row['experiment_type']}, {desc} — " + f"n_seeds={row['n_seeds']}, seeds={seeds}" + ) # Check for runs with no test metrics at all test_cols = [c for c in TEST_METRICS if c in per_seed_df.columns] diff --git a/scripts/launch_thesis_tmux.sh b/scripts/launch_thesis_tmux.sh index 32a5332..e65a36a 100755 --- a/scripts/launch_thesis_tmux.sh +++ b/scripts/launch_thesis_tmux.sh @@ -13,7 +13,7 @@ BATCH_SIZE_FAIRNESS="${BATCH_SIZE_FAIRNESS:-64}" DEVICE_FAIRNESS="${DEVICE_FAIRNESS:-auto}" INCLUDE_SPRINT12="${INCLUDE_SPRINT12:-0}" INCLUDE_SPRINT13="${INCLUDE_SPRINT13:-1}" -INCLUDE_SPRINT7P="${INCLUDE_SPRINT7P:-0}" +INCLUDE_SPRINT7P="${INCLUDE_SPRINT7P:-1}" RUN_EXPORT="${RUN_EXPORT:-1}" STATUS_INTERVAL="${STATUS_INTERVAL:-60}" RESULTS_DIR="${RESULTS_DIR:-results/${WANDB_PROJECT}_${REVISION}}" @@ -34,7 +34,6 @@ mkdir -p "$LOG_DIR" main_sprints=(1 1b 1c 2 3 4 5 6 7 8 10 11) fairness_sprints=(1 2 3 4 5 10) tag_sprints=(1b 1c 2 5 6 7 8) -warmup_sprints=("${main_sprints[@]}") appendix_sprints=() if [[ "$INCLUDE_SPRINT13" == "1" ]]; then @@ -51,6 +50,10 @@ fi if [[ "$INCLUDE_SPRINT12" == "1" ]]; then appendix_sprints+=(12) fairness_sprints+=(12) +fi + +warmup_sprints=("${main_sprints[@]}") +if [[ "$INCLUDE_SPRINT12" == "1" ]]; then warmup_sprints+=(12) fi diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index a25ddc5..b117530 100644 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -112,7 +112,7 @@ "5": ["1", "2", "3", "4"], "6": ["1", "2", "3", "4", "5"], "7": ["1", "3", "5"], - "7p": ["6"], + "7p": ["6", "10"], "8": ["1", "1b", "1c", "5"], } @@ -597,76 +597,79 @@ def build_sprint6(self): self._add_xgboost(sprint, ds, seed, task, frac) def build_sprint7p(self): - """Model capacity pilot — test whether bigger models widen SSL-supervised gap. + """Focused model-capacity study for the thesis. - MIIV only, seed 42, mortality_24h, MAE + supervised. - Two model sizes (medium=128d/4L, large=256d/4L) at 3 label fractions. - Compares against Sprint 6 baseline (64d/2L) — no need to rerun baseline. + MIIV only, mortality_24h only, MAE + supervised. + Two larger model sizes (medium=128d/4L, large=256d/4L) at 3 label fractions + across all five thesis seeds. Baseline comparisons come from Sprint 6 + (seeds 42/123/456) and Sprint 10 (seeds 789/1011), so the default-size + baseline is inherited rather than rerun here. - Total: 2 pretrain + 2×3 finetune + 2×3 probe + 2×3 supervised = 20 runs. + Total: 5 seeds × [2 pretrain + 2×3 finetune + 2×3 probe + 2×3 supervised] = 100 runs. """ sprint = "7p" - ds, seed, task = "miiv", 42, "mortality_24h" - - for size_name, size_cfg in MODEL_SIZES.items(): - # Common model override for all runs at this size - model_extra = {"model": size_cfg["model"]} - - # --- MAE pretrain (new encoder size, full data) --- - pretrain_extra = { - **model_extra, - **size_cfg["ssl_scale"]["mae"], - } - pt = self._add_pretrain(sprint, "mae", ds, seed, extra=pretrain_extra) - - # --- MAE finetune (Protocol B) + probe (Protocol A) --- - finetune_extra = {**model_extra} - for frac in LABEL_FRACTIONS_PILOT: - self._add_finetune( - sprint, - "mae", - ds, - seed, - task, - False, - pt, - label_fraction=frac, - extra=finetune_extra, - name_extra={"size": size_name}, - ) - self._add_finetune( - sprint, - "mae", - ds, - seed, - task, - True, - pt, - label_fraction=frac, - extra=finetune_extra, - name_extra={"size": size_name}, - ) + ds, task = "miiv", "mortality_24h" - # --- Supervised baseline (same larger encoder, from scratch) --- - sup_extra = {**model_extra} - for frac in LABEL_FRACTIONS_PILOT: - sup_name = f"supervised_{task}_{ds}_seed{seed}_{size_name}" - if frac < 1.0: - frac_str = str(frac).replace(".", "") - sup_name += f"_frac{frac_str}" - run = Run( - id=f"s{sprint}_{sup_name}", - sprint=sprint, - run_type="supervised", - paradigm="supervised", - dataset=ds, - seed=seed, - output_dir=_output_dir(sprint, sup_name), - task=task, - label_fraction=frac, - extra_overrides=sup_extra, - ) - self.runs.append(run) + for seed in SEEDS_EXTENDED: + for size_name, size_cfg in MODEL_SIZES.items(): + # Common model override for all runs at this size + model_extra = {"model": size_cfg["model"]} + + # --- MAE pretrain (new encoder size, full data) --- + pretrain_extra = { + **model_extra, + **size_cfg["ssl_scale"]["mae"], + } + pt = self._add_pretrain(sprint, "mae", ds, seed, extra=pretrain_extra) + + # --- MAE finetune (Protocol B) + probe (Protocol A) --- + finetune_extra = {**model_extra} + for frac in LABEL_FRACTIONS_PILOT: + self._add_finetune( + sprint, + "mae", + ds, + seed, + task, + False, + pt, + label_fraction=frac, + extra=finetune_extra, + name_extra={"size": size_name}, + ) + self._add_finetune( + sprint, + "mae", + ds, + seed, + task, + True, + pt, + label_fraction=frac, + extra=finetune_extra, + name_extra={"size": size_name}, + ) + + # --- Supervised baseline (same larger encoder, from scratch) --- + sup_extra = {**model_extra} + for frac in LABEL_FRACTIONS_PILOT: + sup_name = f"supervised_{task}_{ds}_seed{seed}_{size_name}" + if frac < 1.0: + frac_str = str(frac).replace(".", "") + sup_name += f"_frac{frac_str}" + run = Run( + id=f"s{sprint}_{sup_name}", + sprint=sprint, + run_type="supervised", + paradigm="supervised", + dataset=ds, + seed=seed, + output_dir=_output_dir(sprint, sup_name), + task=task, + label_fraction=frac, + extra_overrides=sup_extra, + ) + self.runs.append(run) def build_sprint7(self): """Cross-dataset transfer — Protocol B only (full finetune).""" @@ -860,12 +863,13 @@ def build_sprint12(self): ) def build_sprint13(self): - """TS2Vec temporal contrastive variant, 5 seeds. + """TS2Vec temporal contrastive variant, 5 seeds, both protocols. Addresses the "contrastive was set up to fail" vulnerability by giving the contrastive paradigm its natural augmentations (noise + masking) and a temporal contrastive loss. Same encoder, same training budget. - Protocol B (full finetune) only — matches the primary evaluation protocol. + Evaluated under both Protocol A (probe) and Protocol B (full finetune) + so it fits the main thesis A/B framing. """ sprint = "13" for seed in SEEDS_EXTENDED: @@ -873,6 +877,7 @@ def build_sprint13(self): pt = self._add_pretrain(sprint, "ts2vec", ds, seed) for task in TASKS: self._add_finetune(sprint, "ts2vec", ds, seed, task, False, pt) + self._add_finetune(sprint, "ts2vec", ds, seed, task, True, pt) def build_all(self) -> list[Run]: """Build full experiment matrix. Order matters for dedup.""" diff --git a/tests/test_fixes.py b/tests/test_fixes.py index a9b98aa..b912142 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -752,3 +752,31 @@ def test_apply_wandb_target_injects_project_and_entity(self): 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" + + +class TestExperimentRunnerMatrix: + """Tests for thesis-final sprint matrix coverage.""" + + def test_sprint7p_expanded_to_five_seeds(self): + from scripts.run_experiments import BASELINE_SPRINTS, SEEDS_EXTENDED, MatrixBuilder + + builder = MatrixBuilder() + builder.build_sprint7p() + + assert len(builder.runs) == 100 + assert sorted({run.seed for run in builder.runs}) == SEEDS_EXTENDED + assert BASELINE_SPRINTS["7p"] == ["6", "10"] + + def test_sprint13_includes_both_protocols(self): + from scripts.run_experiments import MatrixBuilder + + builder = MatrixBuilder() + builder.build_sprint13() + + 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 From 3cb462f8d8cba763756e75e00204ea62f5f178af Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 7 Apr 2026 14:57:54 -0400 Subject: [PATCH 061/121] docs: align thesis plan with formal sprint scope --- README.md | 29 +++++++++++++++++++++++------ 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 3d32719..f86eb14 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,11 @@ 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 formal thesis corpus now also includes two targeted extensions: + +- `Sprint 7p`: a focused capacity study that scales MAE and supervised encoders on MIIV `mortality_24h` +- `Sprint 13`: a TS2Vec temporal-contrastive extension that tests whether a stronger contrastive-family instantiation changes the conclusion + ### The Comparison Triangle | Comparison | What Varies | What It Tests | @@ -37,7 +42,7 @@ 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 thesis objectives share the same timestep-level obs-aware Transformer encoder and differ only in the SSL objective and masking logic: | Objective | Predicts | Target | Loss | |---|---|---|---| @@ -45,7 +50,7 @@ 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 formal temporal-contrastive thesis extension (Sprint `13`). **SMART** (NeurIPS 2024) remains an appendix-only external reference because it swaps in its own MART encoder and element-wise masking, so it is not part of the controlled thesis comparison. ## Pipeline @@ -111,6 +116,9 @@ 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 +# TS2Vec (formal thesis extension for the contrastive family) +uv run python scripts/training/pretrain.py dataset=miiv ssl=ts2vec + # SMART (sanity check / extensibility demo — not part of thesis experiments) uv run python scripts/training/pretrain.py dataset=miiv ssl=smart model=smart ``` @@ -127,6 +135,14 @@ uv run python scripts/training/finetune.py dataset=miiv checkpoint=outputs/.../e uv run python scripts/training/supervised.py dataset=miiv ``` +### 5. Thesis Reruns + +```bash +WANDB_PROJECT=slices-thesis \ +WANDB_ENTITY=your-wandb-entity \ +bash scripts/launch_thesis_tmux.sh +``` + ## Project Structure ``` @@ -151,7 +167,8 @@ SLICES/ │ │ │ ├── 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/ @@ -163,12 +180,12 @@ SLICES/ │ ├── fairness.py # Per-group AUROC, demographic parity │ └── 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/ @@ -195,7 +212,7 @@ uv run python scripts/training/pretrain.py dataset=miiv ssl=jepa training.overfi | Group | Options | Purpose | |---|---|---| -| `ssl/` | `mae`, `jepa`, `contrastive`, `smart` | SSL objective hyperparameters | +| `ssl/` | `mae`, `jepa`, `contrastive`, `ts2vec`, `smart` | SSL objective hyperparameters | | `model/` | `transformer`, `observation_transformer`, `smart` | Encoder architecture | | `data/` | `ricu` | Dataset and paths | | `tasks/` | `mortality`, `los`, `aki`, ... | Downstream task definitions | From 97f30e7021cbe1401ed54b305d332186fd8153b3 Mon Sep 17 00:00:00 2001 From: hill Date: Fri, 10 Apr 2026 16:35:56 -0400 Subject: [PATCH 062/121] Fix audit issues across training, extraction, and export --- scripts/eval/evaluate_imputation.py | 20 ++- scripts/export_results.py | 141 ++++++++++++--- scripts/run_experiments.py | 5 +- scripts/training/finetune.py | 11 +- scripts/training/supervised.py | 11 +- src/slices/data/datamodule.py | 27 ++- src/slices/data/dataset.py | 27 ++- src/slices/data/extractors/base.py | 17 ++ src/slices/data/extractors/ricu.py | 62 ++++++- src/slices/data/labels/mortality.py | 13 +- src/slices/data/splits.py | 20 +++ src/slices/training/finetune_module.py | 2 + tests/test_dataset_datamodule.py | 50 ++++++ tests/test_finetune_module.py | 54 ++++++ tests/test_fixes.py | 240 +++++++++++++++++++++++-- tests/test_imputation_eval.py | 37 ++++ tests/test_ricu_extractor.py | 56 ++++++ 17 files changed, 717 insertions(+), 76 deletions(-) diff --git a/scripts/eval/evaluate_imputation.py b/scripts/eval/evaluate_imputation.py index 3ee71a5..2150c27 100644 --- a/scripts/eval/evaluate_imputation.py +++ b/scripts/eval/evaluate_imputation.py @@ -24,8 +24,9 @@ import lightning.pytorch as pl import torch from omegaconf import DictConfig, OmegaConf -from slices.data.datamodule import ICUDataModule +from slices.data.datamodule import ICUDataModule, icu_collate_fn from slices.eval.imputation import ImputationEvaluator +from torch.utils.data import DataLoader, Subset @hydra.main( @@ -124,6 +125,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 +157,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}") diff --git a/scripts/export_results.py b/scripts/export_results.py index e8abb30..43692ef 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -58,7 +58,7 @@ "13": "temporal_contrastive", } -CROSS_SPRINT_TYPES = {"core", "label_efficiency", "transfer", "hp_ablation"} +CROSS_SPRINT_TYPES = {"core", "label_efficiency", "transfer"} FIXED_SEED_EXPERIMENT_TYPES = { "core", "label_efficiency", @@ -85,19 +85,22 @@ "phase", ] -# For non-core runs: sprint distinguishes pretrain configs that aren't captured -# in the finetune W&B config (e.g., which LR/mask_ratio the pretrain used). -# upstream_pretrain_lr/mask_ratio ensure HP ablation rows with different pretrain -# configs are not collapsed (fixes Sprint 1b, 1c, 8, 10 dedup collisions). -ABLATION_FINGERPRINT = CORE_FINGERPRINT + [ - "sprint", +# HP ablations merge across sprints, but must preserve the upstream pretrain +# config and subtype so LR and mask-ratio families cannot collapse together. +HP_ABLATION_FINGERPRINT = CORE_FINGERPRINT + [ + "experiment_subtype", "upstream_pretrain_lr", "upstream_pretrain_mask_ratio", ] -# Legacy fingerprint used for deduplication (includes sprint + all config dims). -DEDUP_FINGERPRINT = [ +# For per-sprint experiment families, sprint remains part of the identity. +ABLATION_FINGERPRINT = HP_ABLATION_FINGERPRINT + [ "sprint", +] + +FINGERPRINT_AUDIT_COLUMNS = [ + "experiment_type", + "experiment_subtype", "paradigm", "dataset", "task", @@ -108,6 +111,9 @@ "model_size", "source_dataset", "phase", + "sprint", + "upstream_pretrain_lr", + "upstream_pretrain_mask_ratio", ] # Metrics to extract from run summaries. @@ -454,6 +460,10 @@ def extract_run(run, metric_keys: list[str]) -> dict: # Recover upstream pretrain metadata (from config for new runs, run name for historical) up_lr, up_mr, experiment_subtype = _recover_pretrain_metadata(run_name, config) + if experiment_type in ("lr_ablation", "mask_ablation"): + experiment_subtype = experiment_subtype or experiment_type + experiment_type = "hp_ablation" + # Reclassify HP ablation runs that would otherwise stay as "core" if experiment_type == "core" and experiment_subtype in ("lr_ablation", "mask_ablation"): experiment_type = "hp_ablation" @@ -499,6 +509,73 @@ def extract_run(run, metric_keys: list[str]) -> dict: # --------------------------------------------------------------------------- +def _fingerprint_for_experiment_type(experiment_type: str) -> list[str]: + """Return the canonical fingerprint for a given experiment family.""" + if experiment_type in CROSS_SPRINT_TYPES: + return CORE_FINGERPRINT + if experiment_type == "hp_ablation": + return HP_ABLATION_FINGERPRINT + return ABLATION_FINGERPRINT + + +def _allowed_varying_columns(experiment_type: str) -> set[str]: + """Columns allowed to vary within one canonical experiment family.""" + if experiment_type in CROSS_SPRINT_TYPES or experiment_type == "hp_ablation": + return {"sprint"} + return set() + + +def _fillna_for_grouping(df: pd.DataFrame, columns: list[str]) -> pd.DataFrame: + """Fill NaNs so pandas groupby/drop_duplicates treat missing values deterministically.""" + 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: + """Fail if the chosen fingerprint would collapse non-equivalent configurations.""" + collisions = [] + + for experiment_type in sorted(df["experiment_type"].dropna().unique()): + subset = df[df["experiment_type"] == experiment_type] + if subset.empty: + continue + + fingerprint = [*(_fingerprint_for_experiment_type(experiment_type)), "seed"] + fingerprint = [c for c in fingerprint if c in subset.columns] + if not fingerprint: + continue + + allowed_vary = _allowed_varying_columns(experiment_type) + identity_cols = [ + c for c in FINGERPRINT_AUDIT_COLUMNS if c in subset.columns and c not in allowed_vary + ] + + work = _fillna_for_grouping(subset, list(dict.fromkeys(fingerprint + identity_cols))) + grouped = work.groupby(fingerprint, dropna=False) + + for key, group in grouped: + if len(group) <= 1: + continue + + unique_identities = group[identity_cols].drop_duplicates() + if len(unique_identities) > 1: + if not isinstance(key, tuple): + key = (key,) + collisions.append((experiment_type, dict(zip(fingerprint, key)), len(group))) + + if collisions: + preview = [] + for experiment_type, fingerprint, count in collisions[:5]: + desc = ", ".join(f"{k}={v}" for k, v in fingerprint.items()) + preview.append(f"{experiment_type}: {desc} ({count} rows)") + raise RuntimeError( + "Ambiguous export fingerprint would collapse distinct runs:\n " + "\n ".join(preview) + ) + + def build_per_seed_df(runs: list) -> pd.DataFrame: """Build the per-seed DataFrame from raw W&B runs. @@ -537,19 +614,18 @@ def build_per_seed_df(runs: list) -> pd.DataFrame: # Deduplicate: keep only the most recent run per (fingerprint + seed). # This handles reruns/revisions — the latest finished run is canonical. - # Uses conditional fingerprints: core runs dedup by CORE_FINGERPRINT (which - # excludes sprint/lr/mask_ratio), non-core by ABLATION_FINGERPRINT. + # Different experiment families use different fingerprints. df = df.sort_values("created_at", ascending=False) before = len(df) + _assert_no_ambiguous_fingerprint_collisions(df) - cross_sprint_types = CROSS_SPRINT_TYPES - cross_sprint_mask = df["experiment_type"].isin(cross_sprint_types) - cross_sprint_dedup_cols = [c for c in CORE_FINGERPRINT + ["seed"] if c in df.columns] - ablation_dedup_cols = [c for c in ABLATION_FINGERPRINT + ["seed"] if c in df.columns] - - core_df = df[cross_sprint_mask].drop_duplicates(subset=cross_sprint_dedup_cols, keep="first") - non_core_df = df[~cross_sprint_mask].drop_duplicates(subset=ablation_dedup_cols, keep="first") - df = pd.concat([core_df, non_core_df], ignore_index=True) + parts = [] + for experiment_type in sorted(df["experiment_type"].dropna().unique()): + subset = df[df["experiment_type"] == experiment_type] + dedup_cols = [*(_fingerprint_for_experiment_type(experiment_type)), "seed"] + dedup_cols = [c for c in dedup_cols if c in subset.columns] + parts.append(subset.drop_duplicates(subset=dedup_cols, keep="first")) + df = pd.concat(parts, ignore_index=True) if parts else df.iloc[0:0] after = len(df) if before != after: @@ -631,23 +707,32 @@ def build_aggregated_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: """Group by fingerprint, compute mean/std/min/max across seeds. Uses conditional grouping: - - Core runs (Sprints 1-5, 10): group WITHOUT sprint to merge seeds across sprints - - Non-core runs (ablations, label efficiency, etc.): group WITH sprint + - Core / label-efficiency / transfer runs: group WITHOUT sprint + - HP ablations: group WITHOUT sprint, but include subtype + upstream pretrain config + - Other experiment families: group WITH sprint """ - # Cross-sprint types: seeds were added across sprints for the same config, - # so aggregate WITHOUT sprint to merge them. - cross_sprint_types = CROSS_SPRINT_TYPES - cross_sprint_mask = per_seed_df["experiment_type"].isin(cross_sprint_types) - cross_sprint = per_seed_df[cross_sprint_mask] - per_sprint = per_seed_df[~cross_sprint_mask] - parts = [] + cross_sprint = per_seed_df[per_seed_df["experiment_type"].isin(CROSS_SPRINT_TYPES)] if len(cross_sprint) > 0: print( f" Aggregating {len(cross_sprint)} cross-sprint runs (merged experiment families)...", file=sys.stderr, ) parts.append(_aggregate_group(cross_sprint, CORE_FINGERPRINT)) + + hp_ablation = per_seed_df[per_seed_df["experiment_type"] == "hp_ablation"] + if len(hp_ablation) > 0: + print( + " Aggregating " + f"{len(hp_ablation)} HP-ablation runs " + "(merged by subtype + upstream config)...", + file=sys.stderr, + ) + parts.append(_aggregate_group(hp_ablation, HP_ABLATION_FINGERPRINT)) + + per_sprint = per_seed_df[ + ~per_seed_df["experiment_type"].isin(CROSS_SPRINT_TYPES | {"hp_ablation"}) + ] if len(per_sprint) > 0: print( f" Aggregating {len(per_sprint)} per-sprint runs (ablations, etc.)...", file=sys.stderr diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index b117530..90742b6 100644 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -1363,10 +1363,10 @@ def cmd_status(args): def cmd_retry(args): all_runs = generate_all_runs() + sprint_filter = {str(s) for s in args.sprint} if args.sprint else None # Apply revision if specified (must happen before ID matching against state) if args.revision: - sprint_filter = [str(s) for s in args.sprint] if args.sprint else None if sprint_filter: revised = [r for r in all_runs if r.sprint in sprint_filter] rest = [r for r in all_runs if r.sprint not in sprint_filter] @@ -1379,8 +1379,9 @@ def cmd_retry(args): state = load_state() recover_stale_running(state) + candidate_runs = [r for r in all_runs if sprint_filter is None or r.sprint in sprint_filter] runs_to_retry = [] - for r in all_runs: + for r in candidate_runs: status = get_run_status(state, r.id) if args.failed and status == "failed": set_run_status(state, r.id, "pending") diff --git a/scripts/training/finetune.py b/scripts/training/finetune.py index 72d1b2b..562028a 100644 --- a/scripts/training/finetune.py +++ b/scripts/training/finetune.py @@ -113,6 +113,7 @@ def main(cfg: DictConfig) -> None: ) label_stats = datamodule.get_label_statistics() + train_label_stats = datamodule.get_train_label_statistics() if task_name in label_stats: stats = label_stats[task_name] print(f"\n Label distribution for '{task_name}':") @@ -139,8 +140,8 @@ 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] + if task_name in train_label_stats: + stats = train_label_stats[task_name] n_pos = stats.get("positive", 0) n_neg = stats.get("negative", 0) n_total = n_pos + n_neg @@ -151,9 +152,13 @@ def main(cfg: DictConfig) -> None: ) raw = [n_total / (2 * n_neg), n_total / (2 * n_pos)] cfg.training.class_weight = [w**0.5 for w in raw] + 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") + 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) diff --git a/scripts/training/supervised.py b/scripts/training/supervised.py index 4325337..7b63caf 100644 --- a/scripts/training/supervised.py +++ b/scripts/training/supervised.py @@ -183,6 +183,7 @@ def main(cfg: DictConfig) -> None: ) label_stats = datamodule.get_label_statistics() + train_label_stats = datamodule.get_train_label_statistics() if task_name in label_stats: stats = label_stats[task_name] print(f"\n Label distribution for '{task_name}':") @@ -209,8 +210,8 @@ 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] + if task_name in train_label_stats: + stats = train_label_stats[task_name] n_pos = stats.get("positive", 0) n_neg = stats.get("negative", 0) n_total = n_pos + n_neg @@ -221,9 +222,13 @@ def main(cfg: DictConfig) -> None: ) raw = [n_total / (2 * n_neg), n_total / (2 * n_pos)] cfg.training.class_weight = [w**0.5 for w in raw] + 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") + 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) diff --git a/src/slices/data/datamodule.py b/src/slices/data/datamodule.py index b753ba6..9bb14b0 100644 --- a/src/slices/data/datamodule.py +++ b/src/slices/data/datamodule.py @@ -169,6 +169,7 @@ def __init__( # Will be set in setup() self.dataset: Optional[ICUDataset] = None + self.full_train_indices: List[int] = [] self.train_indices: List[int] = [] self.val_indices: List[int] = [] self.test_indices: List[int] = [] @@ -213,7 +214,8 @@ def setup(self, stage: Optional[str] = None) -> None: # Save full train indices for normalization BEFORE subsampling. # Normalization must always use the full training split to avoid # noisy stats at low label fractions and pretrain/finetune mismatch. - normalization_train_indices = list(self.train_indices) + self.full_train_indices = list(self.train_indices) + normalization_train_indices = list(self.full_train_indices) # Subsample training indices for label-efficiency ablations if self.label_fraction < 1.0: @@ -253,13 +255,15 @@ def setup(self, stage: Optional[str] = None) -> None: save_split_info( self.processed_dir, self.dataset, - self.train_indices, + self.full_train_indices, self.val_indices, self.test_indices, self.seed, self.train_ratio, self.val_ratio, self.test_ratio, + train_subset_indices=self.train_indices, + label_fraction=self.label_fraction, ) # Free temporary data used only during setup — Dataset holds its own copies @@ -401,6 +405,7 @@ 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), "val_stays": len(self.val_indices), "test_stays": len(self.test_indices), "total_stays": total_stays, @@ -411,6 +416,9 @@ 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_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, } @@ -420,3 +428,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 7ab78a9..0a3126d 100644 --- a/src/slices/data/dataset.py +++ b/src/slices/data/dataset.py @@ -534,20 +534,31 @@ 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. 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, ...} """ + 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() + if task_name in labels_df.columns: + labels = labels_df[task_name].drop_nulls() positive = (labels == 1).sum() total = len(labels) stats[task_name] = { @@ -558,14 +569,12 @@ def get_label_statistics(self) -> Dict[str, Dict[str, Any]]: } 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}_") - ] + multilabel_cols = [c for c in 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() + 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 @@ -577,7 +586,7 @@ def get_label_statistics(self) -> Dict[str, Dict[str, Any]]: } prevalences.append(prev) stats[task_name] = { - "total": len(self.labels_df), + "total": len(labels_df), "n_labels": len(multilabel_cols), "mean_prevalence": ( sum(prevalences) / len(prevalences) if prevalences else 0.0 diff --git a/src/slices/data/extractors/base.py b/src/slices/data/extractors/base.py index ca9601f..2e15732 100644 --- a/src/slices/data/extractors/base.py +++ b/src/slices/data/extractors/base.py @@ -186,6 +186,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). @@ -572,6 +580,15 @@ 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 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 6cfac7b..18c1295 100644 --- a/src/slices/data/extractors/ricu.py +++ b/src/slices/data/extractors/ricu.py @@ -84,6 +84,54 @@ 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 _validate_ricu_horizon(self) -> None: + """Validate that the Python extraction horizon does not exceed the R export.""" + ricu_seq_length = self._metadata.get("seq_length_hours") + if ricu_seq_length is None: + raise ValueError( + "ricu_metadata.yaml is missing seq_length_hours; " + "cannot validate the export horizon safely." + ) + + ricu_seq_length = int(ricu_seq_length) + if self.config.seq_length_hours > ricu_seq_length: + raise ValueError( + "Python extraction requests seq_length_hours=" + f"{self.config.seq_length_hours}, but the upstream RICU export only " + f"contains {ricu_seq_length} hours. Re-run the R export with a longer " + "horizon or lower the Python extraction seq_length_hours." + ) + + 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"]), + "files": files, + } def _get_dataset_name(self) -> str: return self._metadata["dataset"] @@ -195,15 +243,14 @@ def _migrate_mortality_schema(df: pl.DataFrame) -> pl.DataFrame: .alias("death_source"), ) else: - # Datetime column — could be true timestamp or midnight-cast date. - # For MIIV legacy data this is midnight-cast dod; for eICU it's - # a true offset-derived timestamp. Mark as "timestamp" since that - # was the prior behavior; the new R extraction fixes MIIV properly. + # 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.col("date_of_death").cast(pl.Datetime("us", "UTC")).alias("death_time"), - pl.lit(None).cast(pl.Date).alias("death_date"), + 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("timestamp")) + .then(pl.lit("date")) .otherwise(pl.lit(None)) .alias("death_time_precision"), pl.when(pl.col("date_of_death").is_not_null()) @@ -388,6 +435,7 @@ def run(self) -> None: "task_names": task_names, "n_stays": len(stays_filtered), "label_manifest": label_manifest, + "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/mortality.py b/src/slices/data/labels/mortality.py index 3a1f7f8..b59cf05 100644 --- a/src/slices/data/labels/mortality.py +++ b/src/slices/data/labels/mortality.py @@ -37,7 +37,7 @@ class MortalityLabelBuilder(LabelBuilder): - "unknown" or missing: falls back to hospital_expire_flag """ - SEMANTIC_VERSION = "2.0.0" + SEMANTIC_VERSION = "2.1.0" def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: """Build mortality labels from stay and mortality data. @@ -173,7 +173,10 @@ def _ensure_precision_columns(merged: pl.DataFrame) -> pl.DataFrame: pl.lit("legacy").alias("death_source"), ) else: - # Datetime — treat as timestamp (legacy behavior) + # 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"), @@ -185,10 +188,10 @@ def _ensure_precision_columns(merged: pl.DataFrame) -> pl.DataFrame: pl.col("date_of_death").cast(pl.Datetime("us")).alias("date_of_death") ) return merged.with_columns( - pl.col("date_of_death").cast(pl.Datetime("us")).alias("death_time"), - pl.lit(None).cast(pl.Date).alias("death_date"), + 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("timestamp")) + .then(pl.lit("date")) .otherwise(pl.lit(None)) .alias("death_time_precision"), pl.lit("legacy").alias("death_source"), diff --git a/src/slices/data/splits.py b/src/slices/data/splits.py index bd14b08..631681c 100644 --- a/src/slices/data/splits.py +++ b/src/slices/data/splits.py @@ -420,6 +420,8 @@ def save_split_info( 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. @@ -433,6 +435,10 @@ def save_split_info( 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())) @@ -446,6 +452,7 @@ def save_split_info( "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, @@ -454,6 +461,19 @@ def save_split_info( "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/training/finetune_module.py b/src/slices/training/finetune_module.py index 0e61a03..37a41c9 100644 --- a/src/slices/training/finetune_module.py +++ b/src/slices/training/finetune_module.py @@ -237,11 +237,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 diff --git a/tests/test_dataset_datamodule.py b/tests/test_dataset_datamodule.py index 811bea6..653bef5 100644 --- a/tests/test_dataset_datamodule.py +++ b/tests/test_dataset_datamodule.py @@ -770,6 +770,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 +897,34 @@ 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_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 + class TestCollateFn: """Tests for the collate function.""" diff --git a/tests/test_finetune_module.py b/tests/test_finetune_module.py index abf3e80..d3f2a65 100644 --- a/tests/test_finetune_module.py +++ b/tests/test_finetune_module.py @@ -1297,3 +1297,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 index b912142..c7ffb06 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -6,7 +6,9 @@ """ from datetime import datetime, timedelta +from types import SimpleNamespace +import pandas as pd import polars as pl import pytest import torch @@ -26,6 +28,21 @@ ) 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 # ============================================================================ @@ -587,7 +604,7 @@ def test_survived_patient(self): assert labels["label"][0] == 0 def test_legacy_schema_migration(self): - """Legacy data with only date_of_death should be migrated correctly.""" + """Legacy datetime date_of_death should be migrated conservatively.""" merged = pl.DataFrame( { "stay_id": [1], @@ -600,7 +617,9 @@ def test_legacy_schema_migration(self): 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] == "timestamp" + assert result["death_time_precision"][0] == "date" + assert result["death_time"][0] is None + assert result["death_date"][0].isoformat() == "2020-01-03" # ============================================================================ @@ -689,18 +708,8 @@ def test_extract_run_reclassifies_transfer_and_sprint13(self): """Sprint-10 transfer seeds and Sprint 13 should export with stable experiment types.""" from scripts.export_results import extract_run - class DummyRun: - def __init__(self, config, tags, name): - self.config = config - self.summary_metrics = {} - self.id = "dummy" - self.url = "https://example.com" - self.name = name - self.tags = tags - self.group = "g" - self.created_at = "2026-04-07T00:00:00" - - transfer = DummyRun( + transfer = DummyWandbRun( + run_id="transfer", config={ "sprint": "10", "dataset": "eicu", @@ -716,7 +725,8 @@ def __init__(self, config, tags, name): ) assert extract_run(transfer, [])["experiment_type"] == "transfer" - ts2vec = DummyRun( + ts2vec = DummyWandbRun( + run_id="ts2vec", config={ "sprint": "13", "dataset": "miiv", @@ -731,6 +741,130 @@ def __init__(self, config, tags, name): ) assert extract_run(ts2vec, [])["experiment_type"] == "temporal_contrastive" + 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={ + "sprint": "10", + "dataset": "miiv", + "paradigm": "mae", + "seed": 42, + "output_dir": "outputs/sprint10/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={ + "sprint": "10", + "dataset": "miiv", + "paradigm": "mae", + "seed": 42, + "output_dir": "outputs/sprint10/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_type"]) == {"hp_ablation"} + assert sorted(df["upstream_pretrain_lr"].tolist()) == [2e-4, 5e-4] + + def test_build_aggregated_df_merges_hp_ablation_across_sprints(self): + """Sprint 10 historical seeds should merge with Sprint 1b/1c HP-ablation families.""" + from scripts.export_results import build_aggregated_df + + per_seed_df = pd.DataFrame( + [ + { + "wandb_run_id": f"run_{i}", + "experiment_type": "hp_ablation", + "experiment_subtype": "lr_ablation", + "sprint": sprint, + "paradigm": "mae", + "dataset": "miiv", + "task": "mortality_24h", + "seed": seed, + "protocol": "B", + "label_fraction": 1.0, + "model_size": "default", + "source_dataset": None, + "phase": "finetune", + "upstream_pretrain_lr": 2e-4, + "upstream_pretrain_mask_ratio": None, + } + for i, (sprint, seed) in enumerate( + [("1b", 42), ("1b", 123), ("10", 456), ("10", 789), ("10", 1011)] + ) + ] + ) + + agg = build_aggregated_df(per_seed_df) + + assert len(agg) == 1 + assert agg.iloc[0]["experiment_type"] == "hp_ablation" + assert agg.iloc[0]["experiment_subtype"] == "lr_ablation" + assert agg.iloc[0]["n_seeds"] == 5 + assert set(yaml.safe_load(agg.iloc[0]["sprint_list"])) == {"1b", "10"} + + 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={ + "sprint": "1", + "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={ + "sprint": "2", + "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.""" @@ -754,6 +888,82 @@ def test_apply_wandb_target_injects_project_and_entity(self): assert result[0].extra_overrides["logging.wandb_entity"] == "hannes-ill" +class TestExperimentRunnerRetry: + """Tests for retry scoping and dependency preservation.""" + + def test_cmd_retry_respects_sprint_filter_but_keeps_dependencies(self, monkeypatch): + import scripts.run_experiments as runner + + dependency = runner.Run( + id="dep_pretrain", + sprint="0", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/sprint0/pretrain_mae_miiv_seed42", + ) + target = runner.Run( + id="target_finetune", + sprint="1", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/sprint1/finetune_mae_miiv_seed42", + depends_on=[dependency.id], + task="mortality_24h", + ) + other = runner.Run( + id="other_failed", + sprint="2", + run_type="supervised", + paradigm="supervised", + dataset="miiv", + seed=42, + output_dir="outputs/sprint2/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( + sprint=["1"], + 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 + + class TestExperimentRunnerMatrix: """Tests for thesis-final sprint matrix coverage.""" diff --git a/tests/test_imputation_eval.py b/tests/test_imputation_eval.py index 8ef3cb9..efa127e 100644 --- a/tests/test_imputation_eval.py +++ b/tests/test_imputation_eval.py @@ -248,6 +248,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.""" diff --git a/tests/test_ricu_extractor.py b/tests/test_ricu_extractor.py index 1757244..332c02f 100644 --- a/tests/test_ricu_extractor.py +++ b/tests/test_ricu_extractor.py @@ -135,6 +135,7 @@ 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, ) extractor = RicuExtractor(config) @@ -164,6 +165,18 @@ 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) + # --------------------------------------------------------------------------- # Core extraction methods @@ -244,6 +257,7 @@ def test_reads_from_partitioned_directory(self, ricu_output_dir: Path, tmp_path: config = ExtractorConfig( parquet_root=str(ricu_output_dir), output_dir=str(tmp_path / "processed_partitioned"), + seq_length_hours=6, ) extractor = RicuExtractor(config) ts = extractor.extract_timeseries([100, 200, 300]) @@ -289,6 +303,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.""" @@ -423,6 +447,8 @@ def test_run_metadata_content(self, ricu_output_dir: Path, tmp_path: Path) -> No assert meta["n_stays"] == 3 assert meta["feature_names"] == ["hr", "sbp", "crea"] assert "ricu_metadata" in meta + assert "upstream_source_signature" in meta + assert len(meta["upstream_source_signature"]["files"]) >= 4 def test_run_filters_short_stays(self, ricu_output_dir: Path, tmp_path: Path) -> None: output_dir = tmp_path / "processed" @@ -523,3 +549,33 @@ 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_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 From c4c372e5af37f7c291ae9bc3d79b1cc0c3e1de61 Mon Sep 17 00:00:00 2001 From: hill Date: Fri, 10 Apr 2026 20:45:16 -0400 Subject: [PATCH 063/121] Canonicalize thesis baseline matrix --- README.md | 3 +- docs/EXPERIMENT_PLAN.md | 522 ++++++++++++++++++++++++++++++++++ scripts/launch_thesis_tmux.sh | 4 +- scripts/run_experiments.py | 49 ++-- tests/test_fixes.py | 21 ++ 5 files changed, 578 insertions(+), 21 deletions(-) create mode 100644 docs/EXPERIMENT_PLAN.md diff --git a/README.md b/README.md index f86eb14..4a8ac30 100644 --- a/README.md +++ b/README.md @@ -21,8 +21,9 @@ 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 formal thesis corpus now also includes two targeted extensions: +The formal thesis corpus now also includes a contextual baseline family plus two targeted extensions: +- `Sprint 11`: canonical GRU-D and XGBoost baselines, reported as contextual ICU references alongside the controlled benchmark - `Sprint 7p`: a focused capacity study that scales MAE and supervised encoders on MIIV `mortality_24h` - `Sprint 13`: a TS2Vec temporal-contrastive extension that tests whether a stronger contrastive-family instantiation changes the conclusion diff --git a/docs/EXPERIMENT_PLAN.md b/docs/EXPERIMENT_PLAN.md new file mode 100644 index 0000000..2153eb4 --- /dev/null +++ b/docs/EXPERIMENT_PLAN.md @@ -0,0 +1,522 @@ +# 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? +4. Does increasing model capacity widen the SSL-supervised gap on low-label mortality prediction? +5. Does a temporal contrastive objective (TS2Vec) materially change the contrastive-family conclusion? + +--- + +## 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; final thesis reporting also includes GRU-D and XGBoost as contextual baselines | +| **Downstream Task** | mortality_24h, mortality_hospital, aki_kdigo, los_remaining | 3 classification + 1 regression | +| **Seed** | 42, 123, 456, 789, 1011 | Core matrix starts with 3 seeds; Sprint `10` extends the thesis corpus to 5 seeds, and Sprints `7p`, `11`, `12`, and `13` use the same 5-seed footing | + +**Formal thesis scope**: The controlled comparison remains the three-way MAE/JEPA/Contrastive benchmark plus supervised. The final thesis matrix also includes GRU-D and XGBoost as contextual baselines, canonicalized in Sprint `11` rather than duplicated across the early SSL/supervised sprints. Two additional thesis experiments are now tracked explicitly: Sprint `7p` is a focused capacity study (larger MAE + supervised encoders on MIIV mortality_24h), and Sprint `13` is a temporal-contrastive extension (TS2Vec) run under the same downstream A/B protocol. Sprint `12` (SMART) remains appendix-only because it changes the encoder family. + +### 1.2 Controlled Variables (Held Constant) + +| Variable | Value | Rationale | +|----------|-------|-----------| +| Observation window | 48 hours | Standard in ICU benchmarks | +| AKI prediction window | Hours 48–72 (forward-looking) | Prevents leakage from creatinine values used in label construction | +| Min stay | 48 hours (72 hours for AKI) | Ensures full observation + prediction 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 | +| Precision | fp32 | Avoids bf16 numerical issues with small model | +| Weight decay | 0.05 | Same regularization across all downstream runs | + +### 1.3 Shared Encoder Architecture + +All three SSL paradigms share the same encoder architecture, removing tokenization as a confounding variable: + +| Paradigm | Tokenization | Encoder | Masking | Rationale | +|----------|-------------|---------|---------|-----------| +| MAE | Timestep-level (obs-aware) | Transformer (d=64, L=2, H=4, obs_aware=True) | Random timestep masking | Reconstruct masked timestep features | +| JEPA | Timestep-level (obs-aware) | Transformer (d=64, L=2, H=4, obs_aware=True) | Block masking (3 contiguous segments) | Predict in latent space; block masking prevents trivial interpolation | +| Contrastive | Timestep-level (obs-aware) | Transformer (d=64, L=2, H=4, obs_aware=True) | Two complementary masked views (instance mode, zero overlap) | Instance-level NT-Xent: mean-pool each view → align sequence embeddings | +| Supervised | Timestep-level (obs-aware) | Transformer (d=64, L=2, H=4, obs_aware=True) | N/A | Same encoder as SSL for fair comparison | + +**Obs-aware tokenization**: Each timestep token is produced by an MLP projection of `concat(values, obs_mask)` → `d_model`. This encodes both the observed values and the missingness pattern, avoiding the "mostly zeros" problem of naively feeding sparse timestep vectors through a linear projection. With ~22 observed features per timestep on average, the MLP maps ~44 non-zero inputs (22 values + 22 mask bits) to 64-dim tokens — no information bottleneck. + +**Design rationale**: An earlier observation-level design (1 token per observed value, using `ObservationTransformerEncoder`) produced ~1660 tokens per sample at 20% observation density, causing O(N²) attention to exceed L4 GPU memory. Timestep-level tokenization reduces this to 48 tokens (~35× reduction) while preserving all information through the obs-aware MLP. This also strengthens the controlled comparison by making all SSL paradigms share identical tokenization. + +--- + +## 2. Training Protocol + +### 2.1 Pretraining + +| Setting | MAE | JEPA | Contrastive | +|---------|-----|------|-------------| +| Epochs | 100 | 100 | 100 | +| Batch size | 256 | 256 | 256 | +| Learning rate | 1e-3 | 1e-3 | 1e-3 | +| Scheduler | Warmup cosine (10 warmup) | Warmup cosine (10 warmup) | Warmup cosine (10 warmup) | +| Mask ratio | 0.5 (random masking) | 0.5 (block masking) | 0.5 (complementary masks, zero overlap) | +| Gradient clipping | 1.0 | 1.0 | 1.0 | +| Early stopping | None (fixed schedule) | None (fixed schedule) | None (fixed schedule) | +| Checkpoint | Last epoch + best val loss | Last epoch + best val loss | Last epoch + best val loss | + +**Fixed-schedule training**: All paradigms train for the full epoch budget without early stopping, following the convention in self-distillation literature (I-JEPA, DINO, BYOL) where validation loss is unreliable as a stopping criterion — JEPA's val loss rises as representations get richer, not worse. Fixed schedules ensure identical gradient step budgets across paradigms. The cosine LR schedule decays to eta_min=1e-6, providing natural convergence. Both last-epoch and best-val-loss encoders are saved; the last-epoch encoder is used for all primary results, with best-val-loss as a robustness check in the appendix. + +### 2.2 Downstream Evaluation (Two Protocols) + +SSL encoders are evaluated under **two** protocols to answer different questions: + +#### Protocol A: Linear Probing (primary SSL comparison) + +Isolates representation quality — the only variable is the SSL pretraining objective. + +| 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 AUPRC (classification) or val MAE (regression) | +| Head | MLP, hidden_dims=[64], dropout=0.3, ReLU | +| Gradient clipping | 1.0 | +| Label smoothing | 0.1 | +| Weight decay | 0.05 | +| Class weighting | sqrt(balanced) — square root of inverse frequency | + +**Rationale**: Linear probing answers "Which SSL objective produces the best representations?" by preventing the encoder from adapting to the downstream task. + +#### Protocol B: Full Finetuning (practical utility) + +Measures how useful SSL pretraining is as weight initialization. + +| Setting | Value | +|---------|-------| +| Protocol | Full finetuning (freeze_encoder=false) | +| Epochs | 100 | +| Batch size | 64 | +| Learning rate | 3e-4 | +| Scheduler | Cosine decay (eta_min=1e-6) | +| Early stopping | Patience=10 on val AUPRC (classification) or val MAE (regression) | +| Head | MLP, hidden_dims=[64], dropout=0.3, ReLU | +| Gradient clipping | 1.0 | +| Label smoothing | 0.1 | +| Weight decay | 0.05 | +| Class weighting | sqrt(balanced) — square root of inverse frequency | + +**Rationale**: Full finetuning answers "Which SSL objective is most useful in practice as initialization?" Results may diverge from linear probing — an encoder with mediocre frozen representations can still provide excellent initialization. + +**Regularization rationale**: Early Sprint 1 runs showed severe overfitting across all paradigms (train AUROC ~0.98 by epoch 1–2, monotonically degrading val metrics). The regularization suite — reduced LR (1e-3 → 3e-4), increased dropout (0.1 → 0.3), higher weight decay (0.01 → 0.05), sqrt class weights, and label smoothing — was introduced to address this. These settings apply uniformly to all paradigms and the supervised baseline, preserving the controlled comparison. + +### 2.3 Supervised Baseline + +| Setting | Value | +|---------|-------| +| Epochs | 100 | +| Batch size | 64 | +| Learning rate | 3e-4 | +| Scheduler | Cosine decay (eta_min=1e-6) | +| Early stopping | Patience=10 on val AUPRC (classification) or val MAE (regression) | +| Encoder | Transformer (d=64, L=2, H=4, obs_aware=True), trained end-to-end | +| Head | MLP, hidden_dims=[64], dropout=0.3, ReLU | +| Gradient clipping | 1.0 | +| Label smoothing | 0.1 | +| Weight decay | 0.05 | +| Class weighting | sqrt(balanced) — square root of inverse frequency | + +**Same-architecture rationale**: The supervised baseline uses the identical encoder architecture and tokenization as the SSL paradigms to eliminate model capacity as a confounding variable. The only difference is the training procedure: supervised trains end-to-end on labeled data from random initialization, while SSL paradigms pretrain on unlabeled data and then evaluate via Protocol A (linear probe) and Protocol B (full finetune). Comparing supervised vs. Protocol B isolates the effect of SSL initialization; comparing SSL paradigms via Protocol A isolates representation quality. + +--- + +## 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 5 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 | Transformer (obs_aware) | ssl=mae | +| P2 | MIMIC-IV | JEPA | Transformer (obs_aware) | ssl=jepa | +| P3 | MIMIC-IV | Contrastive | Transformer (obs_aware) | ssl=contrastive | +| P4 | eICU | MAE | Transformer (obs_aware) | ssl=mae dataset=eicu | +| P5 | eICU | JEPA | Transformer (obs_aware) | ssl=jepa dataset=eicu | +| P6 | eICU | Contrastive | Transformer (obs_aware) | ssl=contrastive dataset=eicu | +| P7 | Combined | MAE | Transformer (obs_aware) | ssl=mae dataset=combined | +| P8 | Combined | JEPA | Transformer (obs_aware) | ssl=jepa dataset=combined | +| P9 | Combined | Contrastive | Transformer (obs_aware) | ssl=contrastive dataset=combined | + +### 4.2 Phase 2: Downstream Evaluation (84 runs × 3 seeds = 252 runs) + +Each of the 9 pretrained encoders is evaluated on 4 downstream tasks under **both** protocols: +- Protocol A (linear probe): 9 encoders × 4 tasks = 36 runs +- Protocol B (full finetune): 9 encoders × 4 tasks = 36 runs +- Supervised baseline: 3 datasets × 4 tasks = 12 runs + +**Result table template** (one per dataset, two sub-tables per dataset): + +| | mortality_24h | mortality_hospital | aki_kdigo | los_remaining | +|---|---|---|---|---| +| **MAE** | AUPRC ± std | AUPRC ± std | AUPRC ± std | MAE ± std | +| **JEPA** | | | | | +| **Contrastive** | | | | | +| **Supervised** | (Protocol B only) | | | | + +### 4.3 Phase 3: Fairness Evaluation (on all downstream runs) + +No additional training. Compute fairness metrics on test predictions from Phase 2. + +--- + +## 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**: Evaluate SSL under both downstream protocols plus the supervised Transformer baseline at low label fractions {1%, 5%, 10%, 25%, 50%}. The 100% endpoints are inherited from the main full-label runs rather than relaunched here. + +| Scope | Tasks | Label Fractions | Runs | +|-------|-------|----------------|------| +| Full sweep | mortality_24h (anchor task) | 1%, 5%, 10%, 25%, 50% | (3 SSL paradigms × 2 protocols + supervised) × 3 datasets × 5 fractions = 105 runs/seed | +| Trend check | Remaining 3 tasks | 10% | (3 SSL paradigms × 2 protocols + supervised) × 3 datasets × 3 tasks = 63 runs/seed | +| **Total** | | | **168 runs/seed × 3 seeds = 504 runs** | + +**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** (from main experiments, no extra runs): in-domain pretraining, supervised from scratch. + +### 5.3 Learning Rate Sensitivity + +**Question**: Are paradigm rankings robust to the shared learning rate? + +**Design**: Pretrain with LR ∈ {2e-4, 5e-4, 1e-3, 2e-3} on MIMIC-IV, finetune on mortality_24h. LR=1e-3 reused from main experiments. + +**Total**: 9 new pretraining + 9 new finetuning = 18 runs × 3 seeds = 54 runs + +### 5.4 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. 0.5 reused from main experiments. + +**Total**: 6 new pretraining + 6 new finetuning = 12 runs × 3 seeds = 36 runs + +### 5.5 Focused Capacity Study (Sprint 7p) + +**Question**: Are the modest SSL gains on low-label mortality prediction capacity-limited, or does supervised scale similarly when model size increases? + +**Design**: MIIV only, mortality_24h only, MAE + supervised only. Compare two larger encoder sizes (128d/4L and 256d/4L) against the inherited default-size baseline (64d/2L) from Sprints `6` and `10`, at label fractions {0.01, 0.1, 0.5}, across 5 seeds. + +**Total**: 2 sizes × [1 pretrain + 3 Protocol B finetunes + 3 Protocol A probes + 3 supervised] = 20 runs/seed × 5 seeds = **100 runs** + +### 5.6 Temporal Contrastive Extension (Sprint 13) + +**Question**: Does giving the contrastive family its natural temporal objective and augmentations overturn the conclusion from the instance-level contrastive baseline? + +**Design**: TS2Vec on all 3 datasets, all 4 downstream tasks, both Protocol A and Protocol B, 5 seeds. Same base encoder family and matched training budget as the controlled comparison. + +**Total**: 1 paradigm × 3 datasets × 5 seeds = 15 pretrains + 120 downstream probe/finetune runs = **135 runs** + +--- + +## 6. Run Summary + +| Phase | Description | Unit Runs | Planned Runs | Cumulative | +|-------|-------------|-----------|--------------|------------| +| Phase 1 | Pretraining | 9 | 27 | 27 | +| Phase 2 | Downstream (linear probe + full finetune + supervised) | 84 | 252 | 279 | +| Phase 3 | Fairness evaluation | 0 (eval only) | 0 | 279 | +| Ablation 5.1 | Label efficiency | 168 | 504 | 783 | +| Ablation 5.2 | Cross-dataset transfer | 24 | 72 | 855 | +| Ablation 5.3 | LR sensitivity | 18 | 54 | 909 | +| Ablation 5.4 | Mask ratio sensitivity | 12 | 36 | 945 | +| Sprint 7p | Focused capacity study (MAE + supervised, 5 seeds) | — | 100 | 1045 | +| Sprint 10 | Seeds 789, 1011 for Sprints 1–8 scope (SSL + supervised) | — | 630 | 1675 | +| Sprint 11 | Classical baselines (XGBoost + GRU-D), 5 seeds | — | 360 | 2035 | +| Sprint 13 | TS2Vec temporal contrastive extension, 5 seeds | — | 135 | 2170 | +| Sprint 12 | SMART external SSL reference, 5 seeds | — | 135 | 2305 | +| **Total** | | | **2305** | | + +**Planning note**: The formal thesis corpus through Sprint `13` totals **2170** runs. Sprint `12` adds **135** appendix-only SMART reference runs, bringing the full run budget to **2305**. + +--- + +## 7. Sprint Execution Order + +| Sprint | Description | Runs | Status | +|--------|-------------|------|--------| +| 1 | Pipeline validation (MIMIC, Protocol B + supervised) | 19 | COMPLETE | +| 1b | LR sensitivity (MIMIC, mortality_24h) | 18 | COMPLETE | +| 1c | Mask ratio sensitivity (MIMIC, mortality_24h) | 12 | COMPLETE | +| 2 | MIMIC Protocol A (completes MIMIC results table) | 12 | COMPLETE | +| 3 | eICU (both protocols + supervised) | 31 | COMPLETE | +| 4 | Combined (both protocols + supervised) | 31 | COMPLETE | +| 5 | Seeds 123, 456 for Sprints 1–4 | 186 | COMPLETE | +| 6 | Label efficiency ablation | 504 | COMPLETE | +| 7 | Cross-dataset transfer | 72 | COMPLETE | +| 8 | LR + mask ratio extra seeds | 60 | COMPLETE | +| 10 | Seeds 789, 1011 for Sprints 1–8 scope (SSL + supervised) | 630 | TODO | +| 7p | Focused capacity study (MIIV mortality_24h, MAE + supervised, 5 seeds) | 100 | TODO | +| 11 | Classical baselines (XGBoost + GRU-D), 5 seeds | 360 | TODO | +| 13 | TS2Vec temporal contrastive extension, both protocols, 5 seeds | 135 | TODO | +| 12 | SMART external SSL reference, 5 seeds | 135 | TODO | +| 9 | Fairness (eval only, runs last after all training) | 0 | TODO | + +For detailed sprint outcomes, see `EXPERIMENT_RESULTS.md`. + +### Upcoming: Sprint 10 — Extra Seeds 789, 1011 (630 runs) + +**Motivation**: Go from 3 seeds to 5 seeds for all SSL and supervised experiments, improving statistical power (4 d.f. instead of 2 for variance estimates). + +**Scope**: Seeds 789 and 1011 for the full Sprints 1–8 scope. Classical baselines are not duplicated here; Sprint `11` remains their single canonical family. + +| Scope | Equivalent Sprint | Runs | +|-------|-------------------|------| +| Core experiments (pretrain + probe + finetune + supervised) | Sprints 1–4 via Sprint 5 | 186 | +| Label efficiency ablation (SSL + supervised) | Sprint 6 | 336 | +| Cross-dataset transfer | Sprint 7 | 48 | +| LR + mask ratio HP ablations | Sprint 8 | 60 | +| **Total** | | **630** | + +**Breakdown**: 48 pretrain (~30 min each, the bottleneck) + 510 finetune/probe (~3–5 min each) + 72 supervised (~3–5 min each). + +**Execution**: `uv run python scripts/run_experiments.py run --sprint 10 --parallel 4` (limited by GPU-bound pretraining) + +### Upcoming: Sprint 7p — Focused Capacity Study (100 runs) + +**Motivation**: The main benchmark suggests SSL gains are real but modest. Sprint `7p` tests whether that gap widens once model capacity increases, while keeping the question narrow enough to stay interpretable: MAE vs supervised, MIIV only, mortality_24h only, and low-label fractions where representation quality matters most. + +**Matrix**: 5 seeds × 2 model sizes × [1 pretrain + 3 Protocol A probes + 3 Protocol B finetunes + 3 supervised baselines] = **100 runs** + +**Comparison set**: Default-size baselines are inherited from Sprint `6` (seeds 42/123/456) and Sprint `10` (seeds 789/1011), so Sprint `7p` only launches the larger-capacity runs. + +**Execution**: `uv run python scripts/run_experiments.py run --sprint 7p --parallel 4` + +### Upcoming: Sprint 11 — Classical Baselines (360 runs) + +**Motivation**: Provide non-neural (XGBoost) and missingness-native neural (GRU-D) reference points. Without these, SSL AUROC numbers exist in a vacuum. XGBoost with engineered features is the perennially competitive clinical baseline (BAT 2025); GRU-D (Che et al. 2018) handles missing data natively via learnable decay. + +**Canonicalization note**: These baselines are part of the final thesis matrix, but they are launched only here. Earlier sprints intentionally contain only the controlled SSL + supervised benchmark so the classical baselines are reported once as a separate contextual comparison family. + +**Matrix**: 2 baselines × 3 datasets × 4 tasks × 5 seeds = 120 full-label runs + 240 label-efficiency runs = **360 runs** + +**XGBoost**: 8 summary statistics per input feature (mean, std, min, max, first, last, obs_count, obs_fraction) → 896 tabular features. Uses `ICUDataModule` for identical data splits. + +**GRU-D**: d_model=64, 1 GRU layer, same downstream training protocol as supervised Transformer (lr=3e-4, wd=0.05, patience=10, label smoothing, class weights). Trains on GPU via `supervised.py --config-name gru_d`. + +**Key comparisons enabled**: +1. XGBoost/GRU-D vs SSL Protocol B — absolute performance context +2. XGBoost/GRU-D vs supervised Transformer — is the Transformer architecture justified? +3. XGBoost cannot transfer across datasets — highlights SSL's cross-dataset transfer capability (Sprint 7) +4. XGBoost/GRU-D on label efficiency curves — direct comparison with SSL learning curves + +**Execution**: `uv run python scripts/run_experiments.py run --sprint 11 --parallel 12` (all runs independent, no pretraining dependency) + +### Upcoming: Sprint 13 — TS2Vec Temporal Contrastive Extension (135 runs) + +**Motivation**: Sprint `13` addresses the strongest foreseeable reviewer objection to the contrastive-family result: "instance-level contrastive was set up to fail." TS2Vec gives contrastive learning its natural temporal objective and augmentations while preserving the rest of the benchmark machinery. + +**Scope**: 3 datasets × 4 tasks × 2 protocols × 5 seeds, with matched downstream Protocol A/B evaluation and the same base encoder family. + +**Matrix**: 15 pretrains + 120 downstream probe/finetune runs = **135 runs** + +**Interpretation**: Sprint `13` is a formal thesis extension, not a replacement for the controlled MAE/JEPA/Contrastive triangle. It tests whether a better-instantiated contrastive family changes the conclusion. + +**Execution**: `uv run python scripts/run_experiments.py run --sprint 13 --parallel 4` + +### Upcoming: Sprint 12 — SMART External SSL Reference (135 runs) + +**Motivation**: SMART (NeurIPS 2024) is the most recent published SSL method for ICU time series. Including it provides an external SOTA reference point beyond the controlled three-way comparison. Addresses reviewer concern #7 ("no comparison to published clinical SSL methods"). + +**Important**: SMART uses a different encoder architecture (MART, d_model=32, per-variable query tokens, element-wise masking) and is therefore NOT part of the controlled comparison. Results go in the appendix as an external reference, not in the main results table. + +**Matrix**: 1 paradigm × 3 datasets × 4 tasks × 2 protocols × 5 seeds = 120 finetune + 15 pretrain = **135 runs** + +**Execution**: `uv run python scripts/run_experiments.py run --sprint 12 --parallel 4` + +### Upcoming: Sprint 9 — Fairness Analysis (0 extra runs) + +**Runs last**, after all training sprints are complete. Zero additional training — pure evaluation on existing test predictions. + +1. Compute fairness metrics on all finetune/supervised test predictions from Sprints 1–5, 7p, 10, 12, and 13 +2. Enable `eval.fairness.enabled=true` and rerun evaluation +3. Generate fairness tables and disparity plots +4. Protected attributes: sex (all datasets), age group (all), race/ethnicity (MIMIC-IV only) +5. Sprint `11` classical baselines are not part of the default standalone fairness sweep because `evaluate_fairness.py` currently targets `finetune` and `supervised` phases only + +--- + +## 8. W&B Tracking + +### Project Structure + +Legacy exploratory runs remain in W&B project `slices`. Final thesis reruns should log to a separate project such as `slices-thesis` and use a single revision tag such as `thesis-v1`. + +### 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` / `paradigm:xgboost` / `paradigm:gru_d` | Paradigm | +| `task:mortality_24h` / etc. | Downstream task | +| `ablation:label-efficiency` / `ablation:transfer` | Ablation type | +| `sprint:N` | Execution sprint | +| `seed:42` / `seed:123` / `seed:456` / `seed:789` / `seed:1011` | 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 | + +### Baseline Inheritance Across Sprints + +Later sprints reuse runs from earlier sprints as comparison baselines. To enable filtering all relevant runs for a sprint in a single W&B query, baseline runs are tagged with later sprint tags using `run_experiments.py tag --sprint N`. + +| Sprint | Inherits From | What's Inherited | +|--------|--------------|------------------| +| **1** | — | Self-contained (root sprint) | +| **1b** | 1 | Default-LR pretrain+finetune (mortality_24h) + supervised baseline | +| **1c** | 1 | Default-mask-ratio pretrain+finetune (mortality_24h) + supervised baseline | +| **2** | 1 | All Sprint 1 runs — completes the MIMIC results table | +| **3** | — | Self-contained (eICU has own pretrains + supervised) | +| **4** | — | Self-contained (Combined has own pretrains + supervised) | +| **5** | 1, 2, 3, 4 | All seed=42 runs (to aggregate with seeds 123, 456) | +| **6** | 1, 2, 3, 4, 5 | All fraction=1.0 finetune + supervised (learning curve right endpoints) | +| **7** | 1, 3, 5 | In-domain finetunes + supervised for MIMIC and eICU | +| **7p** | 6, 10 | Default-size label-efficiency baselines for seeds 42/123/456/789/1011 | +| **8** | 1, 1b, 1c, 5 | Seed=42 ablation runs + default-HP data points for seeds 123, 456 | +| **10** | — | Self-contained (new pretrains for seeds 789, 1011) | +| **11** | — | Self-contained (classical baselines, 5 seeds) | +| **13** | — | Self-contained (TS2Vec pretrains + both protocols, 5 seeds) | +| **12** | — | Self-contained (SMART pretrains + finetune, 5 seeds) | +| **9** | All | Evaluates test predictions from all training sprints (runs last) | + +**Usage**: `uv run python scripts/run_experiments.py tag --sprint N` (idempotent). + +--- + +## 9. Expected Outputs + +### Tables for Paper + +1. **Table 1**: Main results — AUPRC (mean ± std) for each paradigm × dataset × task +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 + +- Per-seed results for transparency +- Training curves (loss vs steps) for all pretraining runs +- Hyperparameter sensitivity (if reviewer requests) +- Literature comparison table (see `docs/LIT_COMP.md`) + +--- + +## 10. Risk Mitigation + +| Risk | Mitigation | +|------|-----------| +| eICU lacks race/ethnicity data | Fairness analysis for race is MIMIC-only; document as limitation | +| Unequal training budgets across paradigms | Track gradient steps; normalize if >2x difference | +| Class imbalance affects metrics | Use AUPRC as primary metric; sqrt(balanced) class weighting; early stopping on val AUPRC (not val loss); report class ratios | +| Downstream overfitting | Addressed with regularization suite: dropout 0.3, LR 3e-4, weight decay 0.05, sqrt class weights, label smoothing 0.1 — applied uniformly to all paradigms | +| Compute budget exceeded | Sprint ordering ensures usable results at each checkpoint | +| Shared hyperparameters unfair to one paradigm | LR sensitivity (Sprint 1b) and mask ratio sensitivity (Sprint 1c) validate that rankings are robust to shared hyperparameter choices | +| Reviewer requests more seeds | Sprint `10` expands the main matrix to 5 seeds, and Sprints `7p`/`11`/`12`/`13` keep the same 4 d.f. footing for variance estimates | +| AKI label leakage | Fixed: forward-looking prediction window (hours 48–72) prevents model from seeing creatinine values used for KDIGO label construction | diff --git a/scripts/launch_thesis_tmux.sh b/scripts/launch_thesis_tmux.sh index e65a36a..a382f7f 100755 --- a/scripts/launch_thesis_tmux.sh +++ b/scripts/launch_thesis_tmux.sh @@ -6,12 +6,12 @@ SESSION_NAME="${SESSION_NAME:-slices-thesis}" WANDB_PROJECT="${WANDB_PROJECT:-slices-thesis}" WANDB_ENTITY="${WANDB_ENTITY:-}" REVISION="${REVISION:-thesis-v1}" -REASON="${REASON:-2026-04-07 post-cache-freshness final rerun}" +REASON="${REASON:-thesis-final canonical-sprint11-baselines}" PARALLEL_MAIN="${PARALLEL_MAIN:-4}" PARALLEL_APPENDIX="${PARALLEL_APPENDIX:-4}" BATCH_SIZE_FAIRNESS="${BATCH_SIZE_FAIRNESS:-64}" DEVICE_FAIRNESS="${DEVICE_FAIRNESS:-auto}" -INCLUDE_SPRINT12="${INCLUDE_SPRINT12:-0}" +INCLUDE_SPRINT12="${INCLUDE_SPRINT12:-1}" INCLUDE_SPRINT13="${INCLUDE_SPRINT13:-1}" INCLUDE_SPRINT7P="${INCLUDE_SPRINT7P:-1}" RUN_EXPORT="${RUN_EXPORT:-1}" diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index 90742b6..61ddad4 100644 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -469,7 +469,12 @@ def _add_xgboost( # --- Sprint builders --- def build_sprint1(self): - """MIMIC, all tasks, Protocol B + supervised + baselines, seed=42.""" + """MIMIC, all tasks, Protocol B + supervised, seed=42. + + Classical baselines (GRU-D, XGBoost) are canonicalized in Sprint 11 + so the final thesis matrix reports them once rather than duplicating + them across the early SSL/supervised sprints. + """ ds, seed, sprint = "miiv", 42, "1" for p in SSL_PARADIGMS: pt = self._add_pretrain(sprint, p, ds, seed) @@ -477,8 +482,6 @@ def build_sprint1(self): self._add_finetune(sprint, p, ds, seed, task, False, pt) for task in TASKS: self._add_supervised(sprint, ds, seed, task) - self._add_gru_d(sprint, ds, seed, task) - self._add_xgboost(sprint, ds, seed, task) def build_sprint1b(self): """LR sensitivity, MIMIC, mortality_24h, seed=42.""" @@ -529,7 +532,10 @@ def build_sprint2(self): self._add_finetune(sprint, p, ds, seed, task, True, pt) def build_sprint3(self): - """eICU, both protocols + supervised + baselines, seed=42.""" + """eICU, both protocols + supervised, seed=42. + + Classical baselines are launched in Sprint 11 only. + """ ds, seed, sprint = "eicu", 42, "3" for p in SSL_PARADIGMS: pt = self._add_pretrain(sprint, p, ds, seed) @@ -538,11 +544,12 @@ def build_sprint3(self): self._add_finetune(sprint, p, ds, seed, task, True, pt) for task in TASKS: self._add_supervised(sprint, ds, seed, task) - self._add_gru_d(sprint, ds, seed, task) - self._add_xgboost(sprint, ds, seed, task) def build_sprint4(self): - """Combined, both protocols + supervised + baselines, seed=42.""" + """Combined, both protocols + supervised, seed=42. + + Classical baselines are launched in Sprint 11 only. + """ ds, seed, sprint = "combined", 42, "4" for p in SSL_PARADIGMS: pt = self._add_pretrain(sprint, p, ds, seed) @@ -551,11 +558,13 @@ def build_sprint4(self): self._add_finetune(sprint, p, ds, seed, task, True, pt) for task in TASKS: self._add_supervised(sprint, ds, seed, task) - self._add_gru_d(sprint, ds, seed, task) - self._add_xgboost(sprint, ds, seed, task) def build_sprint5(self): - """Seeds 123, 456 for datasets miiv, eicu, combined.""" + """Seeds 123, 456 for datasets miiv, eicu, combined. + + Extends the SSL + supervised matrix to three seeds. Classical baselines + are canonicalized in Sprint 11. + """ sprint = "5" for seed in [123, 456]: for ds in DATASETS: @@ -566,11 +575,13 @@ def build_sprint5(self): self._add_finetune(sprint, p, ds, seed, task, True, pt) for task in TASKS: self._add_supervised(sprint, ds, seed, task) - self._add_gru_d(sprint, ds, seed, task) - self._add_xgboost(sprint, ds, seed, task) def build_sprint6(self): - """Label efficiency ablation — reuses Phase 1 encoders.""" + """Label efficiency ablation — reuses Phase 1 encoders. + + Includes SSL Protocol A/B plus the supervised Transformer baseline. + Classical baselines are canonicalized separately in Sprint 11. + """ sprint = "6" for seed in SEEDS: for ds in DATASETS: @@ -588,13 +599,9 @@ def build_sprint6(self): # Supervised label efficiency for frac in LABEL_FRACTIONS_FULL: self._add_supervised(sprint, ds, seed, "mortality_24h", frac) - self._add_gru_d(sprint, ds, seed, "mortality_24h", frac) - self._add_xgboost(sprint, ds, seed, "mortality_24h", frac) for task in TASKS[1:]: for frac in LABEL_FRACTIONS_TREND: self._add_supervised(sprint, ds, seed, task, frac) - self._add_gru_d(sprint, ds, seed, task, frac) - self._add_xgboost(sprint, ds, seed, task, frac) def build_sprint7p(self): """Focused model-capacity study for the thesis. @@ -807,7 +814,13 @@ def build_sprint10(self): ) def build_sprint11(self): - """Classical baselines (XGBoost + GRU-D), all datasets/tasks, 5 seeds.""" + """Canonical classical baselines (XGBoost + GRU-D), 5 seeds. + + These baselines are intentionally centralized here rather than being + duplicated across earlier sprints. That keeps the final thesis matrix + aligned with the docs and export logic: one canonical experiment family + for contextual non-SSL baselines. + """ sprint = "11" for seed in SEEDS_EXTENDED: for ds in DATASETS: diff --git a/tests/test_fixes.py b/tests/test_fixes.py index c7ffb06..ca93fd6 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -967,6 +967,27 @@ def test_cmd_retry_respects_sprint_filter_but_keeps_dependencies(self, monkeypat class TestExperimentRunnerMatrix: """Tests for thesis-final sprint matrix coverage.""" + def test_classical_baselines_are_canonicalized_to_sprint11(self): + from collections import Counter + + from scripts.run_experiments import generate_all_runs + + runs = generate_all_runs() + by_sprint = Counter(run.sprint for run in runs) + + assert len(runs) == 2305 + assert by_sprint["1"] == 19 + assert by_sprint["3"] == 31 + assert by_sprint["4"] == 31 + assert by_sprint["5"] == 186 + assert by_sprint["6"] == 504 + assert by_sprint["11"] == 360 + + non_s11_classical = [ + run for run in runs if run.run_type in ("gru_d", "xgboost") and run.sprint != "11" + ] + assert non_s11_classical == [] + def test_sprint7p_expanded_to_five_seeds(self): from scripts.run_experiments import BASELINE_SPRINTS, SEEDS_EXTENDED, MatrixBuilder From e10aeec06832042f55f02707c158cb2673f1356d Mon Sep 17 00:00:00 2001 From: hill Date: Sun, 12 Apr 2026 10:53:05 -0400 Subject: [PATCH 064/121] Use 24h observation window defaults --- configs/data/ricu.yaml | 4 ++++ configs/model/observation_transformer.yaml | 2 +- configs/model/transformer.yaml | 2 +- configs/tasks/aki_kdigo.yaml | 2 +- configs/tasks/los_remaining.yaml | 2 +- configs/tasks/mortality_24h.yaml | 2 +- configs/tasks/mortality_hospital.yaml | 2 +- src/slices/constants.py | 10 +++++----- src/slices/data/tasks/aki_kdigo.yaml | 6 +++--- src/slices/data/tasks/los_remaining.yaml | 2 +- src/slices/data/tasks/mortality.yaml | 10 +++++----- src/slices/data/tasks/mortality_24h.yaml | 12 ++++++------ src/slices/data/tasks/mortality_hospital.yaml | 10 +++++----- tests/test_extractor_config.py | 4 ++-- 14 files changed, 37 insertions(+), 33 deletions(-) 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/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 b20d29b..2854cb1 100644 --- a/configs/model/transformer.yaml +++ b/configs/model/transformer.yaml @@ -14,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 diff --git a/configs/tasks/aki_kdigo.yaml b/configs/tasks/aki_kdigo.yaml index 0aff0a8..e0f2c08 100644 --- a/configs/tasks/aki_kdigo.yaml +++ b/configs/tasks/aki_kdigo.yaml @@ -1,6 +1,6 @@ # @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 diff --git a/configs/tasks/los_remaining.yaml b/configs/tasks/los_remaining.yaml index ce2562c..52d7718 100644 --- a/configs/tasks/los_remaining.yaml +++ b/configs/tasks/los_remaining.yaml @@ -1,6 +1,6 @@ # @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 diff --git a/configs/tasks/mortality_24h.yaml b/configs/tasks/mortality_24h.yaml index fd7db19..235eae2 100644 --- a/configs/tasks/mortality_24h.yaml +++ b/configs/tasks/mortality_24h.yaml @@ -2,7 +2,7 @@ # 24-hour mortality prediction (binary classification) # Predict if patient will die within 24h after the observation window # -# Timeline: |---- observation (48h) ----|---- prediction (24h) ----| +# Timeline: |---- observation (24h) ----|---- prediction (24h) ----| # intime obs_end pred_end task_name: mortality_24h diff --git a/configs/tasks/mortality_hospital.yaml b/configs/tasks/mortality_hospital.yaml index 6a1a5c9..21ba4ba 100644 --- a/configs/tasks/mortality_hospital.yaml +++ b/configs/tasks/mortality_hospital.yaml @@ -1,6 +1,6 @@ # @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 diff --git a/src/slices/constants.py b/src/slices/constants.py index 7b67a9b..19249de 100644 --- a/src/slices/constants.py +++ b/src/slices/constants.py @@ -11,11 +11,11 @@ # ============================================================================= # Observation Window # ============================================================================= -SEQ_LENGTH_HOURS: int = 48 # Model input window -MIN_STAY_HOURS: int = 48 -# RICU extraction horizon — longer than SEQ_LENGTH_HOURS to enable -# forward-looking labels (e.g. AKI prediction from hours 48-72). -LABEL_HORIZON_HOURS: int = 72 +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 # ============================================================================= # Extraction diff --git a/src/slices/data/tasks/aki_kdigo.yaml b/src/slices/data/tasks/aki_kdigo.yaml index b772186..7ff13b0 100644 --- a/src/slices/data/tasks/aki_kdigo.yaml +++ b/src/slices/data/tasks/aki_kdigo.yaml @@ -1,14 +1,14 @@ task_name: aki_kdigo task_type: binary -observation_window_hours: 48 -prediction_window_hours: 24 # Look ahead 24h (hours 48-72) for AKI label +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: 48 # Use full observation window for baseline + 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 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/tests/test_extractor_config.py b/tests/test_extractor_config.py index 7a951cd..64bac23 100644 --- a/tests/test_extractor_config.py +++ b/tests/test_extractor_config.py @@ -24,10 +24,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): From d5f754bc34fab19f688f5361ddf61be2e2685c5e Mon Sep 17 00:00:00 2001 From: hill Date: Sun, 12 Apr 2026 12:10:27 -0400 Subject: [PATCH 065/121] Align 24h benchmark defaults and add regressions --- configs/debug.yaml | 2 +- configs/pretrain.yaml | 4 +- scripts/debug/explore_missingness.py | 3 +- .../preprocessing/create_combined_dataset.py | 5 +- src/slices/data/labels/aki.py | 9 +- src/slices/debug/snapshots.py | 6 +- tests/test_dataset_datamodule.py | 16 +-- tests/test_fixes.py | 130 ++++++++++++++++-- tests/test_ricu_extractor.py | 42 ++++++ 9 files changed, 190 insertions(+), 27 deletions(-) diff --git a/configs/debug.yaml b/configs/debug.yaml index ac1c379..4e0f9c8 100644 --- a/configs/debug.yaml +++ b/configs/debug.yaml @@ -41,7 +41,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/pretrain.yaml b/configs/pretrain.yaml index dc8ed4e..1cd3112 100644 --- a/configs/pretrain.yaml +++ b/configs/pretrain.yaml @@ -45,8 +45,8 @@ 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 for SSL pretraining # Architecture (d_model, n_layers, etc.) inherited from configs/model/transformer.yaml diff --git a/scripts/debug/explore_missingness.py b/scripts/debug/explore_missingness.py index d869be6..04c80fd 100644 --- a/scripts/debug/explore_missingness.py +++ b/scripts/debug/explore_missingness.py @@ -20,6 +20,7 @@ import numpy as np 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]: @@ -349,7 +350,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/preprocessing/create_combined_dataset.py b/scripts/preprocessing/create_combined_dataset.py index 8f5f8cf..59335cb 100644 --- a/scripts/preprocessing/create_combined_dataset.py +++ b/scripts/preprocessing/create_combined_dataset.py @@ -22,6 +22,7 @@ import polars as pl import yaml +from slices.constants import MIN_STAY_HOURS, SEQ_LENGTH_HOURS # Offset applied to stay_id and patient_id for the second dataset # to avoid collisions. Large enough to exceed any real ID range. @@ -346,8 +347,8 @@ 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), + "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), diff --git a/src/slices/data/labels/aki.py b/src/slices/data/labels/aki.py index 8951a81..571f8d1 100644 --- a/src/slices/data/labels/aki.py +++ b/src/slices/data/labels/aki.py @@ -5,6 +5,8 @@ import polars as pl +from slices.constants import SEQ_LENGTH_HOURS + from .base import LabelBuilder logger = logging.getLogger(__name__) @@ -17,13 +19,14 @@ 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 baseline_window_hours (default 48h). + 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.0.0" + SEMANTIC_VERSION = "1.1.0" def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: """Build AKI labels from stays and timeseries data. @@ -47,7 +50,7 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: ) crea_col = self.config.label_params.get("creatinine_col", "crea") - baseline_hours = self.config.label_params.get("baseline_window_hours", 48) + 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) diff --git a/src/slices/debug/snapshots.py b/src/slices/debug/snapshots.py index 70d9903..ef77524 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. @@ -59,7 +61,7 @@ class SnapshotConfig: output_dir: Union[str, Path] = "debug_snapshots" stages: List[PipelineStage] = field(default_factory=lambda: list(PipelineStage)) stay_ids: Optional[List[int]] = None - max_hours: int = 48 + max_hours: int = SEQ_LENGTH_HOURS include_masks: bool = True flatten_arrays: bool = True @@ -305,7 +307,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. diff --git a/tests/test_dataset_datamodule.py b/tests/test_dataset_datamodule.py index 653bef5..acec995 100644 --- a/tests/test_dataset_datamodule.py +++ b/tests/test_dataset_datamodule.py @@ -24,7 +24,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 +57,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 +107,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 +150,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 @@ -624,7 +624,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): @@ -752,7 +752,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.""" diff --git a/tests/test_fixes.py b/tests/test_fixes.py index ca93fd6..79f6dca 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -5,6 +5,7 @@ and exporter dedup logic. """ +import importlib from datetime import datetime, timedelta from types import SimpleNamespace @@ -13,6 +14,7 @@ import pytest import torch import yaml +from omegaconf import OmegaConf from slices.data.labels import LabelBuilder, LabelBuilderFactory, LabelConfig from slices.data.labels.aki import AKILabelBuilder from slices.data.labels.los import LOSLabelBuilder @@ -167,16 +169,11 @@ def test_validate_data_prerequisites_missing_task_in_manifest_raises(self, tmp_p 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 - 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"], - ) + 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": { @@ -964,6 +961,123 @@ def test_cmd_retry_respects_sprint_filter_but_keeps_dependencies(self, monkeypat assert other.id not in scheduled_ids +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): + 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) + + 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 + + class TestExperimentRunnerMatrix: """Tests for thesis-final sprint matrix coverage.""" diff --git a/tests/test_ricu_extractor.py b/tests/test_ricu_extractor.py index 332c02f..e2a240b 100644 --- a/tests/test_ricu_extractor.py +++ b/tests/test_ricu_extractor.py @@ -177,6 +177,48 @@ def test_init_raises_when_python_horizon_exceeds_ricu_export( with pytest.raises(ValueError, match="upstream RICU export only contains 6 hours"): RicuExtractor(config) + 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 + # --------------------------------------------------------------------------- # Core extraction methods From 614846218c89c01a9b09c3d6418fcacbead03e1c Mon Sep 17 00:00:00 2001 From: hill Date: Sun, 12 Apr 2026 12:17:10 -0400 Subject: [PATCH 066/121] test: derive conftest fixture seq_length from benchmark constant Shared sample_batch and sample_config fixtures were still hardcoding 48h, so tests that consumed them ran against a sequence length that no longer matches the 24h benchmark. Import SEQ_LENGTH_HOURS so fixtures track the constant. --- tests/conftest.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 88dc8bf..ebe7f70 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -9,6 +9,7 @@ import polars as pl import pytest import torch +from slices.constants import SEQ_LENGTH_HOURS @pytest.fixture @@ -19,7 +20,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 +41,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", } From 8b866877c354007642961368708303d8ca3b3bfd Mon Sep 17 00:00:00 2001 From: hill Date: Sun, 12 Apr 2026 16:23:45 -0400 Subject: [PATCH 067/121] Surface AKI label quality and low-positive train subsets --- scripts/analyze_labels.py | 38 ++++++++++- scripts/training/finetune.py | 11 ++++ scripts/training/supervised.py | 11 ++++ scripts/training/xgboost_baseline.py | 15 ++++- src/slices/data/labels/base.py | 3 + src/slices/data/tasks/aki_kdigo.yaml | 2 + src/slices/training/utils.py | 84 +++++++++++++++++++++++++ tests/test_analyze_labels.py | 43 +++++++++++++ tests/test_training_utils.py | 94 ++++++++++++++++++++++++++++ 9 files changed, 299 insertions(+), 2 deletions(-) create mode 100644 tests/test_analyze_labels.py diff --git a/scripts/analyze_labels.py b/scripts/analyze_labels.py index 588d5c6..a5090f0 100644 --- a/scripts/analyze_labels.py +++ b/scripts/analyze_labels.py @@ -30,6 +30,7 @@ import numpy as np import polars as pl import yaml +from slices.data.utils import get_package_data_dir def load_labels(processed_dir: Path) -> pl.DataFrame: @@ -68,6 +69,18 @@ def load_metadata(processed_dir: Path) -> Dict[str, Any]: 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.""" + config_path = get_package_data_dir() / "tasks" / f"{task_name}.yaml" + if not config_path.exists(): + return {} + + with open(config_path) as f: + config = yaml.safe_load(f) or {} + + return config.get("quality_checks", {}) or {} + + def get_task_columns(labels_df: pl.DataFrame) -> List[str]: """Get task column names (excluding stay_id). @@ -182,6 +195,7 @@ def compute_auroc_ci_width( def analyze_task( labels_df: pl.DataFrame, task_name: str, + quality_checks: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """Analyze a single task's label distribution. @@ -211,6 +225,8 @@ 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: @@ -254,6 +270,17 @@ 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 @@ -387,6 +414,11 @@ def print_task_stats(stats: Dict[str, Any]) -> None: print(f" Missing labels: {stats['missing_labels']:,} ({stats['missing_percentage']:.1f}%)") 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["class_distribution"]: print("\n Class Distribution:") for class_val, class_stats in stats["class_distribution"].items(): @@ -598,7 +630,11 @@ def main(): print("=" * 60) for task_name in task_columns: - stats = analyze_task(labels_df, task_name) + stats = analyze_task( + labels_df, + task_name, + quality_checks=load_task_quality_checks(task_name), + ) all_stats[task_name] = stats print_task_stats(stats) diff --git a/scripts/training/finetune.py b/scripts/training/finetune.py index 562028a..82a34c4 100644 --- a/scripts/training/finetune.py +++ b/scripts/training/finetune.py @@ -35,6 +35,7 @@ from slices.data.datamodule import ICUDataModule from slices.training import FineTuneModule from slices.training.utils import ( + report_and_validate_train_label_support, run_fairness_evaluation, setup_finetune_callbacks, setup_wandb_logger, @@ -127,6 +128,16 @@ def main(cfg: DictConfig) -> None: f"({(1 - stats.get('prevalence', 0))*100:.1f}%)" ) + report_and_validate_train_label_support( + datamodule=datamodule, + task_name=task_name, + task_type=cfg.task.get("task_type", "binary"), + 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 # ========================================================================= diff --git a/scripts/training/supervised.py b/scripts/training/supervised.py index 7b63caf..8e84317 100644 --- a/scripts/training/supervised.py +++ b/scripts/training/supervised.py @@ -34,6 +34,7 @@ from slices.models.encoders import EncoderWithMissingToken from slices.training import FineTuneModule from slices.training.utils import ( + report_and_validate_train_label_support, run_fairness_evaluation, save_encoder_checkpoint, setup_finetune_callbacks, @@ -197,6 +198,16 @@ def main(cfg: DictConfig) -> None: f"({(1 - stats.get('prevalence', 0))*100:.1f}%)" ) + report_and_validate_train_label_support( + datamodule=datamodule, + task_name=task_name, + task_type=cfg.task.get("task_type", "binary"), + 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) # ========================================================================= diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index 7e76498..fb1aad3 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -125,7 +125,10 @@ def main(cfg: DictConfig) -> None: print("1. Setting up DataModule") print("=" * 80) - from slices.training.utils import validate_data_prerequisites + from slices.training.utils import ( + report_and_validate_train_label_support, + validate_data_prerequisites, + ) task_name = cfg.task.get("task_name", "mortality_24h") task_type = cfg.task.get("task_type", "binary") @@ -148,6 +151,16 @@ def main(cfg: DictConfig) -> None: print(f" Val: {len(datamodule.val_indices)} stays") print(f" Test: {len(datamodule.test_indices)} stays") + 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 # ========================================================================= diff --git a/src/slices/data/labels/base.py b/src/slices/data/labels/base.py index 5810914..139e2b6 100644 --- a/src/slices/data/labels/base.py +++ b/src/slices/data/labels/base.py @@ -28,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" @@ -58,6 +59,8 @@ def config_hash(config: LabelConfig) -> str: 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. hashable = { "task_name": config.task_name, "task_type": config.task_type, diff --git a/src/slices/data/tasks/aki_kdigo.yaml b/src/slices/data/tasks/aki_kdigo.yaml index 7ff13b0..ac06651 100644 --- a/src/slices/data/tasks/aki_kdigo.yaml +++ b/src/slices/data/tasks/aki_kdigo.yaml @@ -12,6 +12,8 @@ label_params: absolute_rise_threshold: 0.3 relative_rise_threshold: 1.5 relative_window_hours: 168 # 7 days per KDIGO 2012 specification +quality_checks: + max_missing_percentage: 28.0 # Warn if 24h baseline pushes exclusions above prior range primary_metric: auprc additional_metrics: - auroc diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index e898fe1..7485594 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -366,6 +366,90 @@ def run_fairness_evaluation( # ============================================================================= +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( + " - Dataset / task / seed / fraction: " + f"{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 validate_data_prerequisites( processed_dir: str, dataset: str, diff --git a/tests/test_analyze_labels.py b/tests/test_analyze_labels.py new file mode 100644 index 0000000..7e7e0d9 --- /dev/null +++ b/tests/test_analyze_labels.py @@ -0,0 +1,43 @@ +"""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"] == [] diff --git a/tests/test_training_utils.py b/tests/test_training_utils.py index b94232e..fd101c3 100644 --- a/tests/test_training_utils.py +++ b/tests/test_training_utils.py @@ -15,6 +15,7 @@ from slices.training.utils import ( build_optimizer, build_scheduler, + report_and_validate_train_label_support, save_encoder_checkpoint, ) @@ -177,6 +178,99 @@ 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 TestSupervisedCheckpointFormat: """Test that save_encoder_weights from supervised script produces v3 format.""" From d72e7966580b1291862d8996602e99c54afd9b20 Mon Sep 17 00:00:00 2001 From: hill Date: Mon, 13 Apr 2026 13:20:35 -0400 Subject: [PATCH 068/121] Validate raw label horizon for 24h benchmark tasks --- scripts/preprocessing/extract_with_ricu.R | 13 ++-- src/slices/data/extractors/base.py | 30 ++++++++ src/slices/data/extractors/ricu.py | 85 ++++++++++++++++++++--- src/slices/data/labels/aki.py | 24 ++++++- src/slices/data/labels/base.py | 16 +++++ tests/test_ricu_extractor.py | 76 ++++++++++++++++++++ tests/test_task_builders.py | 30 ++++++++ 7 files changed, 256 insertions(+), 18 deletions(-) diff --git a/scripts/preprocessing/extract_with_ricu.R b/scripts/preprocessing/extract_with_ricu.R index 4639606..50634b7 100644 --- a/scripts/preprocessing/extract_with_ricu.R +++ b/scripts/preprocessing/extract_with_ricu.R @@ -999,12 +999,13 @@ write_metadata <- function(output_dir, dataset, concept_cols, seq_length_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 = seq_length_hours, + raw_export_horizon_hours = seq_length_hours, + n_stays = n_stays, + ricu_version = as.character(packageVersion("ricu")) ) write_yaml(metadata, file.path(output_dir, "ricu_metadata.yaml")) } diff --git a/src/slices/data/extractors/base.py b/src/slices/data/extractors/base.py index 2e15732..7e8850d 100644 --- a/src/slices/data/extractors/base.py +++ b/src/slices/data/extractors/base.py @@ -330,6 +330,35 @@ def _validate_labels(self, labels: pl.DataFrame, stay_ids: List[int]) -> None: 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. @@ -389,6 +418,7 @@ 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) # For single-label tasks, the builder returns a 'label' column # that we rename to the task name. diff --git a/src/slices/data/extractors/ricu.py b/src/slices/data/extractors/ricu.py index 18c1295..3baa9ad 100644 --- a/src/slices/data/extractors/ricu.py +++ b/src/slices/data/extractors/ricu.py @@ -28,7 +28,7 @@ from slices.constants import FEATURE_BLOCKLIST, LABEL_HORIZON_HOURS from slices.data.config_schemas import TimeSeriesConceptConfig -from slices.data.labels import LabelBuilder, LabelBuilderFactory +from slices.data.labels import LabelBuilder, LabelBuilderFactory, LabelConfig from .base import BaseExtractor, ExtractorConfig @@ -74,6 +74,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}") @@ -86,24 +87,80 @@ def __init__(self, config: ExtractorConfig) -> None: self._metadata = yaml.safe_load(f) self._validate_ricu_horizon() - def _validate_ricu_horizon(self) -> None: - """Validate that the Python extraction horizon does not exceed the R export.""" - ricu_seq_length = self._metadata.get("seq_length_hours") - if ricu_seq_length is None: + 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 seq_length_hours; " - "cannot validate the export horizon safely." + "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()), ) - ricu_seq_length = int(ricu_seq_length) - if self.config.seq_length_hours > ricu_seq_length: + 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 {ricu_seq_length} hours. Re-run the R export with a longer " + 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] = [] @@ -130,6 +187,7 @@ def _get_upstream_source_signature(self) -> dict: 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(), "files": files, } @@ -359,7 +417,7 @@ def run(self) -> None: # ----------------------------------------------------------------- # 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: @@ -430,7 +488,12 @@ 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), diff --git a/src/slices/data/labels/aki.py b/src/slices/data/labels/aki.py index 571f8d1..af6ecc1 100644 --- a/src/slices/data/labels/aki.py +++ b/src/slices/data/labels/aki.py @@ -28,6 +28,18 @@ class AKILabelBuilder(LabelBuilder): 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_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: """Build AKI labels from stays and timeseries data. @@ -149,6 +161,15 @@ 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 + # Build label: # - null if no creatinine data, no baseline, or no data in prediction window # - 1 if either AKI criterion met @@ -175,7 +196,8 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: n_null = result_df["label"].null_count() logger.info( f"AKI labels: {n_aki} AKI, {n_no_aki} no AKI, " - f"{n_null} excluded (no creatinine/baseline or no data in prediction window)" + f"{n_null} excluded ({missing_crea_or_baseline} no creatinine/baseline, " + f"{missing_post_obs} no creatinine in prediction window)" ) return result_df diff --git a/src/slices/data/labels/base.py b/src/slices/data/labels/base.py index 139e2b6..917f941 100644 --- a/src/slices/data/labels/base.py +++ b/src/slices/data/labels/base.py @@ -80,6 +80,22 @@ def __init__(self, config: LabelConfig) -> None: """ self.config = config + 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) + @abstractmethod def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: """Build task labels from raw extracted data. diff --git a/tests/test_ricu_extractor.py b/tests/test_ricu_extractor.py index e2a240b..33bf9e7 100644 --- a/tests/test_ricu_extractor.py +++ b/tests/test_ricu_extractor.py @@ -27,6 +27,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", } @@ -136,6 +137,7 @@ def test_init_success(self, ricu_output_dir: Path, tmp_path: Path) -> None: parquet_root=str(ricu_output_dir), output_dir=str(tmp_path / "processed"), seq_length_hours=6, + tasks=[], ) extractor = RicuExtractor(config) @@ -177,6 +179,76 @@ def test_init_raises_when_python_horizon_exceeds_ricu_export( 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: @@ -300,6 +372,7 @@ def test_reads_from_partitioned_directory(self, ricu_output_dir: Path, tmp_path: 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]) @@ -486,6 +559,9 @@ 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 "ricu_metadata" in meta diff --git a/tests/test_task_builders.py b/tests/test_task_builders.py index 4f3b1cd..b8afbc5 100644 --- a/tests/test_task_builders.py +++ b/tests/test_task_builders.py @@ -997,3 +997,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 From ee4e9d82130ff289e5b944932466926d8060097f Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 14 Apr 2026 11:22:45 -0400 Subject: [PATCH 069/121] Align configs and utilities with 24h benchmark --- .gitignore | 2 + README.md | 10 +- configs/finetune.yaml | 3 +- configs/supervised.yaml | 3 +- configs/tasks/aki_kdigo.yaml | 1 + configs/tasks/los_remaining.yaml | 1 + configs/tasks/mortality.yaml | 14 + configs/tasks/mortality_24h.yaml | 1 + configs/tasks/mortality_hospital.yaml | 1 + docs/EXPERIMENT_PLAN.md | 12 +- examples/transformer_encoder_example.py | 6 +- .../preprocessing/create_combined_dataset.py | 13 + src/slices/data/sliding_window.py | 10 +- src/slices/debug/sampling.py | 2 +- src/slices/models/encoders/README.md | 2 +- src/slices/training/config_schemas.py | 1 + src/slices/training/utils.py | 37 +- tests/README.md | 2 +- tests/test_fixes.py | 8 +- tests/test_task_builders.py | 436 ++++++++++-------- tests/test_training_utils.py | 67 +++ tests/test_transformer_encoder.py | 6 +- 22 files changed, 415 insertions(+), 223 deletions(-) create mode 100644 configs/tasks/mortality.yaml diff --git a/.gitignore b/.gitignore index d353468..dbad839 100644 --- a/.gitignore +++ b/.gitignore @@ -215,6 +215,8 @@ docs/ # SLICES-specific /data/ outputs/ +/results/* +!/results/.gitkeep wandb/ .tensor_cache/ diff --git a/README.md b/README.md index 4a8ac30..ca2018e 100644 --- a/README.md +++ b/README.md @@ -57,9 +57,9 @@ The controlled thesis objectives share the same timestep-level obs-aware Transfo ``` 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 @@ -216,7 +216,7 @@ uv run python scripts/training/pretrain.py dataset=miiv ssl=jepa training.overfi | `ssl/` | `mae`, `jepa`, `contrastive`, `ts2vec`, `smart` | SSL objective hyperparameters | | `model/` | `transformer`, `observation_transformer`, `smart` | 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 | ## Data Format @@ -224,7 +224,7 @@ 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 `ICUDataset` returns batches with: diff --git a/configs/finetune.yaml b/configs/finetune.yaml index 5dd38bb..4121ef1 100644 --- a/configs/finetune.yaml +++ b/configs/finetune.yaml @@ -100,7 +100,8 @@ 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 diff --git a/configs/supervised.yaml b/configs/supervised.yaml index e111190..7367681 100644 --- a/configs/supervised.yaml +++ b/configs/supervised.yaml @@ -88,7 +88,8 @@ 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 diff --git a/configs/tasks/aki_kdigo.yaml b/configs/tasks/aki_kdigo.yaml index e0f2c08..f964ca1 100644 --- a/configs/tasks/aki_kdigo.yaml +++ b/configs/tasks/aki_kdigo.yaml @@ -4,6 +4,7 @@ task_name: aki_kdigo task_type: binary +primary_metric: auprc n_classes: null head_type: mlp hidden_dims: [64] diff --git a/configs/tasks/los_remaining.yaml b/configs/tasks/los_remaining.yaml index 52d7718..d6d602c 100644 --- a/configs/tasks/los_remaining.yaml +++ b/configs/tasks/los_remaining.yaml @@ -4,6 +4,7 @@ task_name: los_remaining task_type: regression +primary_metric: mae n_classes: null head_type: mlp hidden_dims: [64] diff --git a/configs/tasks/mortality.yaml b/configs/tasks/mortality.yaml new file mode 100644 index 0000000..5cce9e8 --- /dev/null +++ b/configs/tasks/mortality.yaml @@ -0,0 +1,14 @@ +# @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 +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 235eae2..4fc3a55 100644 --- a/configs/tasks/mortality_24h.yaml +++ b/configs/tasks/mortality_24h.yaml @@ -7,6 +7,7 @@ task_name: mortality_24h task_type: binary +primary_metric: auprc n_classes: null head_type: mlp hidden_dims: [64] diff --git a/configs/tasks/mortality_hospital.yaml b/configs/tasks/mortality_hospital.yaml index 21ba4ba..eacb43c 100644 --- a/configs/tasks/mortality_hospital.yaml +++ b/configs/tasks/mortality_hospital.yaml @@ -4,6 +4,7 @@ task_name: mortality_hospital task_type: binary +primary_metric: auprc n_classes: null head_type: mlp hidden_dims: [64] diff --git a/docs/EXPERIMENT_PLAN.md b/docs/EXPERIMENT_PLAN.md index 2153eb4..0360889 100644 --- a/docs/EXPERIMENT_PLAN.md +++ b/docs/EXPERIMENT_PLAN.md @@ -33,9 +33,9 @@ Secondary questions addressed through ablations: | Variable | Value | Rationale | |----------|-------|-----------| -| Observation window | 48 hours | Standard in ICU benchmarks | -| AKI prediction window | Hours 48–72 (forward-looking) | Prevents leakage from creatinine values used in label construction | -| Min stay | 48 hours (72 hours for AKI) | Ensures full observation + prediction window | +| Observation window | 24 hours | Current benchmark input window | +| AKI prediction window | Hours 24–48 (forward-looking) | Prevents leakage from creatinine values used in label construction | +| Min stay | 24 hours | Ensures a full observation window before downstream labeling | | 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 | @@ -204,11 +204,13 @@ Min subgroup size: 50 patients. Groups below threshold are excluded from fairnes ### 4.2 Phase 2: Downstream Evaluation (84 runs × 3 seeds = 252 runs) -Each of the 9 pretrained encoders is evaluated on 4 downstream tasks under **both** protocols: +Each of the 9 pretrained encoders is evaluated on the **4-task primary thesis matrix** under **both** protocols: - Protocol A (linear probe): 9 encoders × 4 tasks = 36 runs - Protocol B (full finetune): 9 encoders × 4 tasks = 36 runs - Supervised baseline: 3 datasets × 4 tasks = 12 runs +The repo now also exposes an optional `mortality` ICU-mortality task, but it is not included in these run totals unless explicitly added as an extension. + **Result table template** (one per dataset, two sub-tables per dataset): | | mortality_24h | mortality_hospital | aki_kdigo | los_remaining | @@ -282,7 +284,7 @@ No additional training. Compute fairness metrics on test predictions from Phase **Question**: Does giving the contrastive family its natural temporal objective and augmentations overturn the conclusion from the instance-level contrastive baseline? -**Design**: TS2Vec on all 3 datasets, all 4 downstream tasks, both Protocol A and Protocol B, 5 seeds. Same base encoder family and matched training budget as the controlled comparison. +**Design**: TS2Vec on all 3 datasets, the same 4-task primary thesis matrix, both Protocol A and Protocol B, 5 seeds. Same base encoder family and matched training budget as the controlled comparison. **Total**: 1 paradigm × 3 datasets × 5 seeds = 15 pretrains + 120 downstream probe/finetune runs = **135 runs** diff --git a/examples/transformer_encoder_example.py b/examples/transformer_encoder_example.py index 0af8d0c..72293ca 100644 --- a/examples/transformer_encoder_example.py +++ b/examples/transformer_encoder_example.py @@ -31,9 +31,9 @@ def basic_usage(): 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 + # Simulate benchmark-style ICU data: batch_size=16, seq_length=24 hours, features=35 batch_size = 16 - seq_length = 48 + seq_length = 24 x = torch.randn(batch_size, seq_length, config.d_input) # Forward pass @@ -57,7 +57,7 @@ def with_observation_mask(): encoder = TransformerEncoder(config) batch_size = 8 - seq_length = 48 + seq_length = 24 x = torch.randn(batch_size, seq_length, config.d_input) # Create observation mask (True = observed, False = missing) diff --git a/scripts/preprocessing/create_combined_dataset.py b/scripts/preprocessing/create_combined_dataset.py index 59335cb..7fbfe70 100644 --- a/scripts/preprocessing/create_combined_dataset.py +++ b/scripts/preprocessing/create_combined_dataset.py @@ -348,6 +348,19 @@ def main(): "feature_names": meta_a.get("feature_names", []), "n_features": meta_a.get("n_features", 0), "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"), + }, "min_stay_hours": meta_a.get("min_stay_hours", MIN_STAY_HOURS), "task_names": common_tasks, "label_manifest": combined_manifest, diff --git a/src/slices/data/sliding_window.py b/src/slices/data/sliding_window.py index e2e7395..5fbce54 100644 --- a/src/slices/data/sliding_window.py +++ b/src/slices/data/sliding_window.py @@ -31,15 +31,15 @@ class SlidingWindowDataset(Dataset): >>> base_dataset = ICUDataset("data/processed/mimic-iv-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/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/models/encoders/README.md b/src/slices/models/encoders/README.md index 3f41368..5340603 100644 --- a/src/slices/models/encoders/README.md +++ b/src/slices/models/encoders/README.md @@ -262,7 +262,7 @@ pretrained_encoder = ssl_model.get_encoder() ### Model Size vs. Performance -For ICU data (typical: 35 features, 48-hour windows): +For ICU data (typical benchmark setup: 35 features, 24-hour observation 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 diff --git a/src/slices/training/config_schemas.py b/src/slices/training/config_schemas.py index 6a216a2..d2bd07e 100644 --- a/src/slices/training/config_schemas.py +++ b/src/slices/training/config_schemas.py @@ -30,6 +30,7 @@ class TaskConfig(BaseModel): task_name: str task_type: str = "binary" + primary_metric: Optional[str] = None head_type: str = "mlp" hidden_dims: List[int] = [64] dropout: float = 0.1 diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 7485594..eb59385 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -188,17 +188,44 @@ 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) + 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." + ) + mode = cfg.training.get("early_stopping_mode", default_mode) # Replace '/' with '_' in monitor name for safe filenames. diff --git a/tests/README.md b/tests/README.md index 2f8304f..ac0f993 100644 --- a/tests/README.md +++ b/tests/README.md @@ -146,7 +146,7 @@ open htmlcov/index.html # macOS - Observation mask handling ### Task Building -- Mortality prediction windows (24h, 48h, hospital, ICU) +- Mortality task variants (24h, hospital, ICU) plus window-boundary cases - Boundary conditions (exact boundaries) - Empty data handling - Multiple tasks extraction diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 79f6dca..67eda1b 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -79,13 +79,13 @@ def test_config_hash_changes_on_window_change(self): task_name="mortality_24h", task_type="binary", prediction_window_hours=24, - observation_window_hours=48, + observation_window_hours=24, ) config_48 = LabelConfig( - task_name="mortality_48h", + task_name="mortality_24h", task_type="binary", prediction_window_hours=48, - observation_window_hours=48, + observation_window_hours=24, ) assert LabelBuilder.config_hash(config_24) != LabelBuilder.config_hash(config_48) @@ -1006,7 +1006,7 @@ def get_label_statistics(self): } } - def get_train_label_statistics(self): + def get_train_label_statistics(self, use_full_train: bool = False): return { "mortality_24h": { "total": 10, diff --git a/tests/test_task_builders.py b/tests/test_task_builders.py index b8afbc5..83059a3 100644 --- a/tests/test_task_builders.py +++ b/tests/test_task_builders.py @@ -11,6 +11,28 @@ 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 +114,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): @@ -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( @@ -554,22 +592,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 +671,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 +733,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 +786,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 +849,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 +909,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 +968,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( diff --git a/tests/test_training_utils.py b/tests/test_training_utils.py index fd101c3..31bdf30 100644 --- a/tests/test_training_utils.py +++ b/tests/test_training_utils.py @@ -7,9 +7,12 @@ - Supervised script save_encoder_weights: v3 format, loadable by FineTuneModule """ +from pathlib import Path + import pytest import torch import torch.nn as nn +import yaml from omegaconf import OmegaConf from slices.training import FineTuneModule from slices.training.utils import ( @@ -17,6 +20,7 @@ build_scheduler, report_and_validate_train_label_support, save_encoder_checkpoint, + setup_finetune_callbacks, ) @@ -271,6 +275,69 @@ def test_raises_when_binary_train_subset_loses_a_class(self): ) +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") + + 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 + + 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 35a0f04..fcdb82a 100644 --- a/tests/test_transformer_encoder.py +++ b/tests/test_transformer_encoder.py @@ -637,7 +637,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 +651,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 +661,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) From 400f6ee51ec1cfc8d06453bbe99ca776c386d62a Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 14 Apr 2026 12:07:29 -0400 Subject: [PATCH 070/121] Persist AKI label quality stats and clarify raw export horizon Add an explicit raw_export_horizon_hours interface to the RICU export pipeline while keeping seq_length_hours as a deprecated compatibility alias. Persist task-level label quality statistics into processed metadata so extractions retain positive/negative/null counts and AKI-specific null reasons. Add an end-to-end AKI extractor test for the intended benchmark setup: 24-hour model input with a 48-hour upstream raw export and mixed post-observation creatinine follow-up. --- .../preprocessing/create_combined_dataset.py | 4 + scripts/preprocessing/extract_ricu.py | 5 +- scripts/preprocessing/extract_with_ricu.R | 50 ++++-- src/slices/data/extractors/base.py | 8 +- src/slices/data/extractors/ricu.py | 4 +- src/slices/data/labels/aki.py | 29 +++- src/slices/data/labels/base.py | 44 +++++- tests/test_ricu_extractor.py | 147 ++++++++++++++++++ 8 files changed, 272 insertions(+), 19 deletions(-) diff --git a/scripts/preprocessing/create_combined_dataset.py b/scripts/preprocessing/create_combined_dataset.py index 7fbfe70..cd7579f 100644 --- a/scripts/preprocessing/create_combined_dataset.py +++ b/scripts/preprocessing/create_combined_dataset.py @@ -361,6 +361,10 @@ def main(): 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, diff --git a/scripts/preprocessing/extract_ricu.py b/scripts/preprocessing/extract_ricu.py index 4567354..771fe20 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 @@ -45,7 +46,7 @@ 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) diff --git a/scripts/preprocessing/extract_with_ricu.R b/scripts/preprocessing/extract_with_ricu.R index 50634b7..2bbda39 100644 --- a/scripts/preprocessing/extract_with_ricu.R +++ b/scripts/preprocessing/extract_with_ricu.R @@ -16,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") @@ -45,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 = 72L, - 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)") ) @@ -996,14 +998,14 @@ 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, - raw_export_horizon_hours = seq_length_hours, + 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")) ) @@ -1015,14 +1017,36 @@ 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( + "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 = ", "))) @@ -1134,8 +1158,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(). @@ -1149,7 +1173,7 @@ main <- function(opts) { # 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, seq_length_hours, dict, + ts_result <- extract_timeseries(dataset, concepts, raw_export_horizon_hours, dict, output_path = ts_path) concept_cols <- ts_result$concept_cols n_stays <- ts_result$n_stays @@ -1171,7 +1195,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/src/slices/data/extractors/base.py b/src/slices/data/extractors/base.py index 7e8850d..6212222 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 @@ -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(): @@ -378,6 +379,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) @@ -419,6 +421,10 @@ 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. diff --git a/src/slices/data/extractors/ricu.py b/src/slices/data/extractors/ricu.py index 3baa9ad..c245cda 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 \ @@ -498,6 +499,7 @@ def run(self) -> None: "task_names": task_names, "n_stays": len(stays_filtered), "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), diff --git a/src/slices/data/labels/aki.py b/src/slices/data/labels/aki.py index af6ecc1..2f46f68 100644 --- a/src/slices/data/labels/aki.py +++ b/src/slices/data/labels/aki.py @@ -1,7 +1,7 @@ """AKI (Acute Kidney Injury) label builder using KDIGO criteria.""" import logging -from typing import Dict +from typing import Any, Dict import polars as pl @@ -40,6 +40,21 @@ def required_raw_timeseries_horizon_hours(self) -> int: 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. @@ -54,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), @@ -83,6 +102,10 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: # 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), @@ -169,6 +192,10 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: & 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, no baseline, or no data in prediction window diff --git a/src/slices/data/labels/base.py b/src/slices/data/labels/base.py index 917f941..2965e5a 100644 --- a/src/slices/data/labels/base.py +++ b/src/slices/data/labels/base.py @@ -4,7 +4,7 @@ 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 @@ -79,6 +79,7 @@ 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. @@ -96,6 +97,47 @@ def required_raw_timeseries_horizon_hours(self) -> int: 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: """Build task labels from raw extracted data. diff --git a/tests/test_ricu_extractor.py b/tests/test_ricu_extractor.py index 33bf9e7..a17a538 100644 --- a/tests/test_ricu_extractor.py +++ b/tests/test_ricu_extractor.py @@ -111,6 +111,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.""" @@ -564,6 +646,7 @@ def test_run_metadata_content(self, ricu_output_dir: Path, tmp_path: Path) -> No 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 @@ -645,6 +728,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" From 09a0176ba93690bbbae23d5a99c615f52dd4a959 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 14 Apr 2026 14:22:29 -0400 Subject: [PATCH 071/121] Unblock combined prep, SMART teachers, and fairness revisions Route dataset preparation through a shared helper, run combined dataset prep from the supported setup path, and make combined creation emit the required split and normalization artifacts by default. Keep the SMART EMA target encoder in eval mode so teacher targets stay deterministic, and add revision filtering to standalone fairness evaluation plus the tmux launcher so reruns do not mix fairness outputs across revisions. Behavioral implications: fresh setup now builds data/processed/combined when requested or by default, SMART sprint 12 uses stable EMA targets, and fairness sweeps respect revision tags. --- scripts/eval/evaluate_fairness.py | 8 + scripts/launch_thesis_tmux.sh | 3 + .../preprocessing/create_combined_dataset.py | 25 +- scripts/preprocessing/prepare_dataset.py | 289 +----------------- scripts/setup_and_extract.sh | 70 ++++- src/slices/data/preparation.py | 244 +++++++++++++++ src/slices/models/pretraining/smart.py | 8 + tests/test_fixes.py | 284 +++++++++++++++++ tests/test_smart_objective.py | 16 + 9 files changed, 654 insertions(+), 293 deletions(-) create mode 100644 src/slices/data/preparation.py diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index 63960e3..77a55be 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -84,6 +84,7 @@ def fetch_eval_runs( 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 @@ -102,6 +103,8 @@ def fetch_eval_runs( tag_filters.append(f"dataset:{datasets[0]}") if phases and len(phases) == 1: tag_filters.append(f"phase:{phases[0]}") + if revisions and len(revisions) == 1: + tag_filters.append(f"revision:{revisions[0]}") if tag_filters: filters["tags"] = {"$all": tag_filters} @@ -114,6 +117,7 @@ def fetch_eval_runs( paradigm_set = set(paradigms) if paradigms and len(paradigms) > 1 else None dataset_set = set(datasets) if datasets and len(datasets) > 1 else None phase_set = set(phases) if phases and len(phases) > 1 else None + revision_set = set(revisions) if revisions and len(revisions) > 1 else None runs = [] for run in runs_iter: @@ -127,6 +131,8 @@ def fetch_eval_runs( continue if phase_set and not any(f"phase:{p}" in tags for p in phase_set): continue + if revision_set and not any(f"revision:{r}" in tags for r in revision_set): + continue runs.append(run) @@ -415,6 +421,7 @@ def parse_args() -> argparse.Namespace: ) 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)") parser.add_argument( "--phase", nargs="+", @@ -484,6 +491,7 @@ def main() -> None: paradigms=args.paradigm, datasets=args.dataset, phases=args.phase, + revisions=args.revision, ) ) diff --git a/scripts/launch_thesis_tmux.sh b/scripts/launch_thesis_tmux.sh index a382f7f..b560963 100755 --- a/scripts/launch_thesis_tmux.sh +++ b/scripts/launch_thesis_tmux.sh @@ -75,6 +75,9 @@ fi run_args+=(--revision "$REVISION" --reason "$REASON") export_args+=(--revision "$REVISION" --output-dir "$RESULTS_DIR") +if [[ -n "$REVISION" ]]; then + fairness_args+=(--revision "$REVISION") +fi fairness_args+=(--sprint "${fairness_sprints[@]}" --batch-size "$BATCH_SIZE_FAIRNESS" --device "$DEVICE_FAIRNESS") warmup_cmd=(uv run python scripts/run_experiments.py warmup --sprint "${warmup_sprints[@]}") diff --git a/scripts/preprocessing/create_combined_dataset.py b/scripts/preprocessing/create_combined_dataset.py index cd7579f..33092fe 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,6 +24,7 @@ import polars as pl import yaml from slices.constants import MIN_STAY_HOURS, SEQ_LENGTH_HOURS +from slices.data.preparation import prepare_processed_dataset # Offset applied to stay_id and patient_id for the second dataset # to avoid collisions. Large enough to exceed any real ID range. @@ -192,6 +194,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]) @@ -393,6 +406,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/prepare_dataset.py b/scripts/preprocessing/prepare_dataset.py index ea97402..c787d7a 100644 --- a/scripts/preprocessing/prepare_dataset.py +++ b/scripts/preprocessing/prepare_dataset.py @@ -1,298 +1,21 @@ -"""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 -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 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, - ) - stats["normalize"] = True - stats["split_hash"] = None - stats["data_fingerprint"] = get_data_fingerprint(processed_dir) - stats["preprocessing_fingerprint"] = get_preprocessing_fingerprint() - - # 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/setup_and_extract.sh b/scripts/setup_and_extract.sh index 715b1e7..a3855ef 100755 --- a/scripts/setup_and_extract.sh +++ b/scripts/setup_and_extract.sh @@ -5,9 +5,10 @@ # 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 set -euo pipefail @@ -31,21 +32,48 @@ error() { echo -e "${RED}[ERROR]${NC} $1"; } # Parse arguments # --------------------------------------------------------------------------- SKIP_DEPS=false -DATASETS=() +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") ;; + miiv|eicu|combined) REQUESTED_DATASETS+=("$arg") ;; *) error "Unknown argument: $arg" - echo "Usage: $0 [--skip-deps] [miiv] [eicu]" + echo "Usage: $0 [--skip-deps] [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) @@ -70,7 +98,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" @@ -136,7 +164,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" @@ -178,12 +206,36 @@ for ds in "${DATASETS[@]}"; do info "Dataset $ds ready!" done +if [ "$BUILD_COMBINED" = true ]; then + section "Processing dataset: combined" + + processed_dir="data/processed/combined" + + if [ -f "$processed_dir/timeseries.parquet" ] && \ + [ -f "$processed_dir/static.parquet" ] && \ + [ -f "$processed_dir/labels.parquet" ] && \ + [ -f "$processed_dir/metadata.yaml" ] && \ + [ -f "$processed_dir/splits.yaml" ] && \ + [ -f "$processed_dir/normalization_stats.yaml" ]; then + info "Combined dataset already exists and is prepared: $processed_dir (skipping)" + else + 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" + info "Combined dataset creation complete: $processed_dir" + fi + + 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/src/slices/data/preparation.py b/src/slices/data/preparation.py new file mode 100644 index 0000000..021f0cc --- /dev/null +++ b/src/slices/data/preparation.py @@ -0,0 +1,244 @@ +"""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}\n" "Run 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["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/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/tests/test_fixes.py b/tests/test_fixes.py index 67eda1b..bb5d807 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -6,7 +6,12 @@ """ import importlib +import os +import shutil +import subprocess +import sys from datetime import datetime, timedelta +from pathlib import Path from types import SimpleNamespace import pandas as pd @@ -376,6 +381,285 @@ def test_invariant_mismatch_raises(self): mod.validate_feature_compatibility(meta_a, meta_b) +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], + } + ) + labels_df.write_parquet(processed_dir / "labels.parquet") + + 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": ["mortality_24h"], + "label_manifest": { + "mortality_24h": { + "builder_version": "1.0.0", + "config_hash": "abc123", + } + }, + } + 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").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 + + +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.setattr( + sys, + "argv", + ["evaluate_fairness.py", "--revision", "thesis-v1"], + ) + args = mod.parse_args() + + assert args.revision == ["thesis-v1"] + + def test_fetch_eval_runs_single_revision_adds_server_side_filter(self, monkeypatch): + """A single revision should be sent as a W&B tag filter.""" + mod = importlib.import_module("scripts.eval.evaluate_fairness") + captured = {} + + 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 [] + + monkeypatch.setitem(sys.modules, "wandb", SimpleNamespace(Api=DummyApi)) + + mod.fetch_eval_runs( + project="proj", + entity="entity", + sprints=["1"], + paradigms=["mae"], + datasets=["miiv"], + phases=["finetune"], + revisions=["thesis-v1"], + ) + + assert captured["path"] == "entity/proj" + assert captured["filters"]["tags"]["$all"] == [ + "sprint:1", + "paradigm:mae", + "dataset:miiv", + "phase:finetune", + "revision:thesis-v1", + ] + + 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, + sprints=None, + paradigms=None, + datasets=None, + phases=["finetune", "supervised"], + revisions=["v1", "v2"], + ) + + assert [run.id for run in filtered] == ["keep_v1", "keep_v2"] + + 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").mkdir(parents=True, exist_ok=True) + shutil.copy2(repo_root / "scripts" / "launch_thesis_tmux.sh", temp_repo / "scripts") + + 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["REVISION"] = "thesis-v2" + env["RUN_EXPORT"] = "0" + + result = subprocess.run( + ["bash", "scripts/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 + + # ============================================================================ # Issue 4: Sequence-length override # ============================================================================ diff --git a/tests/test_smart_objective.py b/tests/test_smart_objective.py index d32aa82..9cfd6fb 100644 --- a/tests/test_smart_objective.py +++ b/tests/test_smart_objective.py @@ -131,6 +131,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 +154,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.""" From bb018dd24b5918c11b32775173248d119351d0fa Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 14 Apr 2026 16:39:57 -0400 Subject: [PATCH 072/121] Finalize experiment readiness fixes across fairness, export, and SSL --- scripts/eval/evaluate_fairness.py | 125 +++++++++- scripts/export_results.py | 198 ++++++++++++++- .../preprocessing/create_combined_dataset.py | 46 +++- scripts/training/xgboost_baseline.py | 1 + src/slices/data/datamodule.py | 36 +-- src/slices/eval/__init__.py | 11 +- src/slices/eval/fairness_evaluator.py | 208 ++++++++++++---- src/slices/eval/statistical.py | 227 ++++++++++++++++-- src/slices/models/encoders/transformer.py | 5 +- src/slices/models/pretraining/contrastive.py | 45 +++- src/slices/models/pretraining/jepa.py | 41 +++- src/slices/models/pretraining/masking.py | 54 ++++- src/slices/models/pretraining/ts2vec.py | 54 ++++- src/slices/training/utils.py | 1 + tests/test_contrastive_objective.py | 15 ++ tests/test_dataset_datamodule.py | 22 ++ tests/test_export_results.py | 43 ++++ tests/test_fairness_evaluator.py | 120 +++++++++ tests/test_fixes.py | 115 +++++++++ tests/test_jepa_objective.py | 17 ++ tests/test_statistical.py | 38 ++- tests/test_transformer_encoder.py | 3 +- tests/test_ts2vec_objective.py | 47 ++++ 23 files changed, 1334 insertions(+), 138 deletions(-) create mode 100644 tests/test_export_results.py create mode 100644 tests/test_ts2vec_objective.py diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index 77a55be..6ed56b2 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -10,24 +10,25 @@ --dataset filters. Usage: - # Evaluate the default thesis fairness corpus - uv run python scripts/eval/evaluate_fairness.py + # Evaluate the thesis fairness corpus for one explicit revision + uv run python scripts/eval/evaluate_fairness.py --revision thesis-v1 # Scope to specific sprint/dataset - uv run python scripts/eval/evaluate_fairness.py --sprint 1 --dataset miiv + uv run python scripts/eval/evaluate_fairness.py --revision thesis-v1 --sprint 1 --dataset miiv # Preview which runs would be evaluated - uv run python scripts/eval/evaluate_fairness.py --dry-run + uv run python scripts/eval/evaluate_fairness.py --revision thesis-v1 --dry-run # Override paths (e.g., different machine than training) uv run python scripts/eval/evaluate_fairness.py \ + --revision thesis-v1 \ --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 --force + uv run python scripts/eval/evaluate_fairness.py --revision thesis-v1 --force # Debug with a single run - uv run python scripts/eval/evaluate_fairness.py --max-runs 1 + uv run python scripts/eval/evaluate_fairness.py --revision thesis-v1 --max-runs 1 """ from __future__ import annotations @@ -192,6 +193,25 @@ def _resolve_ckpt_dir(output_dir: str, outputs_root: Optional[str] = None) -> Pa return Path(output_dir) / "checkpoints" +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, @@ -257,6 +277,61 @@ def find_best_checkpoint( 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." + ) + + # --------------------------------------------------------------------------- # Model + data reconstruction # --------------------------------------------------------------------------- @@ -361,6 +436,7 @@ def evaluate_run_fairness( protected_attributes=protected_attributes, min_subgroup_size=min_subgroup_size, task_type=task_type, + dataset_name=getattr(getattr(datamodule, "processed_dir", None), "name", None), ) report = evaluator.evaluate(predictions, labels, stay_ids) return report @@ -421,7 +497,14 @@ def parse_args() -> argparse.Namespace: ) 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)") + 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="+", @@ -463,7 +546,19 @@ def parse_args() -> argparse.Namespace: ) 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") - return parser.parse_args() + 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: @@ -526,10 +621,12 @@ def _sort_key(r): print("\n[DRY RUN] Listing runs:\n") for i, r in enumerate(runs): cfg = r.config - output_dir = cfg.get("output_dir", "?") task_type = _get_nested(cfg, "task.task_type", "binary") - ckpt = find_best_checkpoint(output_dir, args.outputs_root, task_type) - ckpt_str = str(ckpt) if ckpt else "NOT FOUND" + try: + ckpt, ckpt_source = resolve_evaluation_checkpoint(r, args.outputs_root, task_type) + ckpt_str = f"{ckpt} [{ckpt_source}]" if ckpt else "NOT FOUND" + 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', '?')}/" @@ -554,13 +651,15 @@ def _sort_key(r): try: # 1. Find checkpoint - output_dir = cfg.get("output_dir", "") task_type = _get_nested(cfg, "task.task_type", "binary") - ckpt_path = find_best_checkpoint(output_dir, args.outputs_root, task_type) + ckpt_path, ckpt_source = resolve_evaluation_checkpoint( + run, args.outputs_root, task_type + ) if ckpt_path is None: results["skipped"] += 1 results["errors"].append((run.id, run_desc, "no checkpoint found")) continue + log.info(" Checkpoint: %s (%s)", ckpt_path, ckpt_source) # 2. Reconstruct model + data (reuse datamodule if same dataset/task/seed) dm_key = (ds, task, seed, cfg.get("label_fraction", 1.0)) diff --git a/scripts/export_results.py b/scripts/export_results.py index 43692ef..657c4d2 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -3,11 +3,13 @@ Canonical results export pipeline for SLICES. Pulls all experiment runs from W&B, extracts config + test metrics, and produces -two structured parquet files: +three structured parquet files: - results/per_seed_results.parquet — one row per W&B run (~3000 rows) - results/aggregated_results.parquet — one row per unique config (~600 rows), with mean/std/min/max across seeds + - results/statistical_tests.parquet — pairwise Wilcoxon + Bonferroni + + Cohen's d significance table Both files include wandb run IDs for traceability back to W&B. @@ -27,10 +29,16 @@ import re import sys import time +from itertools import combinations from pathlib import Path import pandas as pd import wandb +from slices.eval.statistical import ( + bonferroni_correction, + cohens_d, + paired_wilcoxon_signed_rank, +) # --------------------------------------------------------------------------- # Constants @@ -173,6 +181,20 @@ (256, 4): "large", } +PRIMARY_TEST_METRIC_BY_TASK = { + "mortality_24h": "test/auprc", + "mortality_hospital": "test/auprc", + "mortality": "test/auprc", + "aki_kdigo": "test/auprc", + "los_remaining": "test/mae", +} + +LOWER_IS_BETTER_METRICS = { + "test/loss", + "test/mae", + "test/mse", +} + # Phases that correspond to evaluation runs (not pretraining). EVAL_PHASES = ["finetune", "supervised", "gru_d", "xgboost", "baseline"] @@ -754,6 +776,168 @@ def build_aggregated_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: return agg_df +def _primary_metric_for_task(task_name: str | None) -> str | None: + """Return the thesis-primary test metric for a downstream task.""" + 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 build_statistical_tests_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: + """Build pairwise paradigm significance tables from per-seed results. + + Comparisons are scoped by the canonical experiment fingerprint with + `paradigm` and `task` removed, then pooled across matched `(task, seed)` + pairs for tasks that share the same primary metric. This matches the thesis + plan's "paired across seeds and tasks" intent without mixing incompatible + scales such as AUPRC and MAE in the same test. + """ + if per_seed_df.empty or "experiment_type" not in per_seed_df.columns: + return pd.DataFrame() + + rows = [] + + for experiment_type in sorted(per_seed_df["experiment_type"].dropna().unique()): + subset = per_seed_df[per_seed_df["experiment_type"] == experiment_type].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.index, + subset.to_dict("records"), + subset["primary_metric_name"], + strict=True, + ) + ] + + fingerprint = _fingerprint_for_experiment_type(experiment_type) + scope_cols = [ + c for c in fingerprint if c in subset.columns and c not in {"paradigm", "task"} + ] + if "primary_metric_name" not in scope_cols: + scope_cols.append("primary_metric_name") + + if not scope_cols: + continue + + 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 + + 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): + pairs_a = scope_group[scope_group["paradigm"] == paradigm_a][ + ["task", "seed", "primary_metric_value"] + ].rename(columns={"primary_metric_value": "value_a"}) + pairs_b = scope_group[scope_group["paradigm"] == paradigm_b][ + ["task", "seed", "primary_metric_value"] + ].rename(columns={"primary_metric_value": "value_b"}) + + paired = pairs_a.merge(pairs_b, on=["task", "seed"], how="inner") + paired = paired.dropna(subset=["value_a", "value_b"]) + if paired.empty: + continue + + 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) + median_improvement = float(pd.Series(improvement).median()) + better_paradigm = ( + paradigm_a + if mean_improvement > 0 + else paradigm_b if mean_improvement < 0 else None + ) + + row = { + **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()) + ), + "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": median_improvement, + "better_paradigm": better_paradigm, + "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 + ), + } + family_rows.append(row) + + if not family_rows: + continue + + 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_type", + "dataset", + "protocol", + "phase", + "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 + + # --------------------------------------------------------------------------- # Validation # --------------------------------------------------------------------------- @@ -912,6 +1096,10 @@ def main(): aggregated_df = build_aggregated_df(per_seed_df) print(f" Shape: {aggregated_df.shape}", file=sys.stderr) + 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) + # Validate warnings = validate(per_seed_df, aggregated_df, expected_core_seeds=args.expected_core_seeds) for w in warnings: @@ -920,17 +1108,23 @@ def main(): # Save per_seed_path = output_dir / "per_seed_results.parquet" aggregated_path = output_dir / "aggregated_results.parquet" + statistical_path = output_dir / "statistical_tests.parquet" per_seed_df.to_parquet(per_seed_path, index=False) aggregated_df.to_parquet(aggregated_path, index=False) + statistical_df.to_parquet(statistical_path, index=False) print("\nSaved:", file=sys.stderr) print(f" {per_seed_path} ({len(per_seed_df)} rows)", file=sys.stderr) print(f" {aggregated_path} ({len(aggregated_df)} rows)", file=sys.stderr) + print(f" {statistical_path} ({len(statistical_df)} rows)", file=sys.stderr) # Print quick summary to stdout for piping print("\n--- Quick Summary ---") - print(f"Runs: {len(per_seed_df)}, Configs: {len(aggregated_df)}") + print( + f"Runs: {len(per_seed_df)}, Configs: {len(aggregated_df)}, " + f"StatTests: {len(statistical_df)}" + ) if "experiment_type" in aggregated_df.columns: for etype, group in aggregated_df.groupby("experiment_type"): seed_dist = dict(group["n_seeds"].value_counts().sort_index()) diff --git a/scripts/preprocessing/create_combined_dataset.py b/scripts/preprocessing/create_combined_dataset.py index 33092fe..8fd3f4c 100644 --- a/scripts/preprocessing/create_combined_dataset.py +++ b/scripts/preprocessing/create_combined_dataset.py @@ -26,8 +26,8 @@ from slices.constants import MIN_STAY_HOURS, SEQ_LENGTH_HOURS from slices.data.preparation import prepare_processed_dataset -# Offset applied to stay_id and patient_id for the second dataset -# to avoid collisions. Large enough to exceed any real ID range. +# 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 @@ -53,19 +53,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 @@ -76,6 +88,16 @@ 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 share the same preprocessing contract.""" @@ -234,6 +256,13 @@ 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"]) @@ -387,6 +416,11 @@ def main(): 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 diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index fb1aad3..9891648 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -268,6 +268,7 @@ def main(cfg: DictConfig) -> None: ), 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": diff --git a/src/slices/data/datamodule.py b/src/slices/data/datamodule.py index 9bb14b0..cabd80e 100644 --- a/src/slices/data/datamodule.py +++ b/src/slices/data/datamodule.py @@ -250,21 +250,27 @@ 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") - save_split_info( - self.processed_dir, - self.dataset, - self.full_train_indices, - self.val_indices, - self.test_indices, - self.seed, - self.train_ratio, - self.val_ratio, - self.test_ratio, - train_subset_indices=self.train_indices, - label_fraction=self.label_fraction, - ) + # 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_split_info( + self.processed_dir, + self.dataset, + self.full_train_indices, + self.val_indices, + self.test_indices, + self.seed, + self.train_ratio, + self.val_ratio, + self.test_ratio, + train_subset_indices=self.train_indices, + 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 diff --git a/src/slices/eval/__init__.py b/src/slices/eval/__init__.py index a14efea..be74974 100644 --- a/src/slices/eval/__init__.py +++ b/src/slices/eval/__init__.py @@ -15,7 +15,13 @@ 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", @@ -23,7 +29,10 @@ "get_default_metrics", "FairnessEvaluator", "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..f61396c 100644 --- a/src/slices/eval/fairness_evaluator.py +++ b/src/slices/eval/fairness_evaluator.py @@ -3,17 +3,15 @@ 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 thesis plan: +- 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 typing import Any, Dict, List, Optional, Tuple import polars as pl @@ -47,6 +45,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 +53,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 +62,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 +89,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. @@ -105,7 +125,7 @@ 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,15 +133,21 @@ 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) @@ -131,14 +157,13 @@ def _encode_attribute( group_ids[ages_tensor < 0] = -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, {}) + for row in rows: val = row.get("gender") raw_vals.append(val) if val is not None: @@ -150,29 +175,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, @@ -196,35 +261,62 @@ def evaluate( 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 patient-count threshold from the thesis plan 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,6 +329,7 @@ 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] = {} @@ -295,6 +388,7 @@ def _evaluate_binary( "disparate_impact_ratio": di_ratio, "n_valid_groups": len(valid_groups), "group_sizes": group_sizes, + "group_sample_sizes": group_sample_sizes, } def _evaluate_regression( @@ -305,6 +399,7 @@ 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] = {} @@ -341,6 +436,7 @@ def _evaluate_regression( "worst_group_mse": worst_group_mse, "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 +457,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 +519,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/statistical.py b/src/slices/eval/statistical.py index b4056c2..fd22058 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 @@ -165,6 +160,141 @@ 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 with tie correction. + + This implementation uses the standard large-sample normal approximation + with tie correction. It is sufficient for the thesis export pipeline where + comparisons pool multiple tasks x seeds and are primarily used for + significance tables rather than exact small-sample inference. + + 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, + } + + abs_diffs = [abs(diff) for diff in nonzero_diffs] + ranks, tie_counts = _average_ranks(abs_diffs) + + w_plus = sum(rank for rank, diff in zip(ranks, nonzero_diffs) if diff > 0.0) + w_minus = sum(rank for rank, diff in zip(ranks, nonzero_diffs) if diff < 0.0) + statistic = min(w_plus, w_minus) + + expected = n_nonzero * (n_nonzero + 1) / 4.0 + variance = n_nonzero * (n_nonzero + 1) * (2 * n_nonzero + 1) / 24.0 + tie_correction = sum(t * (t + 1) * (2 * t + 1) for t in tie_counts) / 48.0 + variance -= tie_correction + + if variance <= 0.0: + z_score = 0.0 + p_value = 1.0 + else: + numerator = abs(w_plus - expected) + if correction: + numerator = max(numerator - 0.5, 0.0) + z_score = numerator / math.sqrt(variance) + p_value = min(2.0 * _normal_sf(z_score), 1.0) + + return { + "statistic": float(statistic), + "z_score": float(z_score), + "p_value": float(p_value), + "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 +307,70 @@ def _compute_metric( return metric_fn.compute().item() else: return metric_fn(preds, targets).item() + + +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 _average_ranks(values: Sequence[float]) -> tuple[list[float], list[int]]: + indexed = sorted(enumerate(values), key=lambda item: item[1]) + ranks = [0.0] * len(values) + tie_counts: list[int] = [] + + i = 0 + while i < len(indexed): + j = i + 1 + while j < len(indexed) and indexed[j][1] == indexed[i][1]: + j += 1 + + start_rank = i + 1 + end_rank = j + average_rank = (start_rank + end_rank) / 2.0 + for k in range(i, j): + original_index = indexed[k][0] + ranks[original_index] = average_rank + + tie_size = j - i + if tie_size > 1: + tie_counts.append(tie_size) + i = j + + return ranks, tie_counts + + +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 _normal_sf(z_score: float) -> float: + return 0.5 * math.erfc(z_score / math.sqrt(2.0)) + + +def _is_nan(value: float) -> bool: + return isinstance(value, float) and math.isnan(value) diff --git a/src/slices/models/encoders/transformer.py b/src/slices/models/encoders/transformer.py index 0a26252..9f2f983 100644 --- a/src/slices/models/encoders/transformer.py +++ b/src/slices/models/encoders/transformer.py @@ -287,13 +287,14 @@ def tokenize( # Add sinusoidal time PE tokens = tokens + self.time_pe[:T].unsqueeze(0) # (B, T, d_model) - # Fixed-length: all timesteps valid - padding_mask = torch.ones(B, T, dtype=torch.bool, device=device) + # 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 diff --git a/src/slices/models/pretraining/contrastive.py b/src/slices/models/pretraining/contrastive.py index 6bb2c64..d590b50 100644 --- a/src/slices/models/pretraining/contrastive.py +++ b/src/slices/models/pretraining/contrastive.py @@ -167,28 +167,57 @@ def forward( # 1. Tokenize (shared between both views) tokens, padding_mask, token_info = self.encoder.tokenize(x, obs_mask) + 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) + 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_2 = ~ssl_mask_1 else: - ssl_mask_2 = create_timestep_mask(B, T, self.config.mask_ratio, device) + 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 # 3. Encode both views - vis_tokens_1, vis_padding_1 = extract_visible_timesteps(tokens, ssl_mask_1) + 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) - vis_tokens_2, vis_padding_2 = extract_visible_timesteps(tokens, ssl_mask_2) + 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) if self.config.mode == "temporal": # 4a. Scatter encoded tokens back to full (B, T, d) grid - full_1 = self._scatter_to_full(encoded_1, ssl_mask_1, T) - full_2 = self._scatter_to_full(encoded_2, ssl_mask_2, T) + 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, ssl_mask_1, ssl_mask_2) + loss, metrics = self._temporal_nt_xent_loss( + full_1, + full_2, + effective_mask_1, + effective_mask_2, + ) else: # 4b. Instance mode: mean-pool -> project -> instance NT-Xent pooled_1 = self._mean_pool(encoded_1, vis_padding_1) @@ -197,7 +226,7 @@ def forward( 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, ssl_mask_1, ssl_mask_2) + loss, metrics = self._nt_xent_loss(z1, z2, effective_mask_1, effective_mask_2) return loss, metrics diff --git a/src/slices/models/pretraining/jepa.py b/src/slices/models/pretraining/jepa.py index c8cd383..cad7d64 100644 --- a/src/slices/models/pretraining/jepa.py +++ b/src/slices/models/pretraining/jepa.py @@ -218,17 +218,34 @@ def forward( # 1. Tokenize timesteps (online encoder) tokens, padding_mask, token_info = self.encoder.tokenize(x, obs_mask) + 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 + 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) + 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_timesteps(tokens, ssl_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) @@ -244,13 +261,18 @@ def forward( # 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, n_timesteps=T, ) # (B, T, d_encoder) # 7. Compute loss on masked positions - loss, metrics = self._compute_loss(predicted_repr, target_repr, ssl_mask) + loss, metrics = self._compute_loss( + predicted_repr, + target_repr, + ssl_mask, + valid_timestep_mask, + ) return loss, metrics @@ -259,6 +281,7 @@ def _compute_loss( predicted: torch.Tensor, target: torch.Tensor, ssl_mask: torch.Tensor, + valid_timestep_mask: torch.Tensor, ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: """Compute loss on masked timestep representations. @@ -270,14 +293,14 @@ def _compute_loss( Returns: (loss, metrics_dict) """ - # Loss only on masked timesteps - loss_mask = ~ssl_mask # (B, T) + # 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.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) @@ -318,7 +341,7 @@ def _compute_loss( metrics = { "jepa_loss": loss.detach(), "ssl_loss": loss.detach(), - "jepa_mask_ratio_actual": n_masked / max(B * T, 1), + "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, diff --git a/src/slices/models/pretraining/masking.py b/src/slices/models/pretraining/masking.py index 00e9b4a..528eeed 100644 --- a/src/slices/models/pretraining/masking.py +++ b/src/slices/models/pretraining/masking.py @@ -86,6 +86,7 @@ def create_timestep_mask( n_timesteps: int, mask_ratio: float, device: torch.device, + valid_timestep_mask: torch.Tensor | None = None, ) -> torch.Tensor: """Create random mask at timestep level. @@ -94,6 +95,9 @@ def create_timestep_mask( 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. @@ -101,13 +105,24 @@ def create_timestep_mask( rand_vals = torch.rand(batch_size, n_timesteps, device=device) ssl_mask = rand_vals >= mask_ratio # True = visible - # Ensure at least 1 visible timestep per sample - n_visible = ssl_mask.sum(dim=1) # (B,) - needs_fix = n_visible == 0 + 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) + + # Empty timesteps are not eligible SSL tokens; mark them visible so they + # are excluded from masked-token accounting and downstream losses. + ssl_mask = ssl_mask | (~valid_timestep_mask) + + # Ensure at least 1 visible eligible timestep per sample + n_visible = (ssl_mask & valid_timestep_mask).sum(dim=1) # (B,) + has_valid = valid_timestep_mask.any(dim=1) + needs_fix = (n_visible == 0) & has_valid if needs_fix.any(): for b in range(batch_size): if needs_fix[b]: - ssl_mask[b, 0] = True + first_valid = valid_timestep_mask[b].nonzero(as_tuple=True)[0][0] + ssl_mask[b, first_valid] = True return ssl_mask @@ -118,6 +133,7 @@ def create_block_timestep_mask( 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. @@ -136,6 +152,8 @@ def create_block_timestep_mask( 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. @@ -188,18 +206,37 @@ def create_block_timestep_mask( masked = in_block.any(dim=1) # (B, T) ssl_mask = (~masked).to(device=device) + 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 = ssl_mask | (~valid_timestep_mask) + + n_visible = (ssl_mask & valid_timestep_mask).sum(dim=1) + has_valid = valid_timestep_mask.any(dim=1) + needs_fix = (n_visible == 0) & has_valid + if needs_fix.any(): + for b in range(batch_size): + if needs_fix[b]: + first_valid = valid_timestep_mask[b].nonzero(as_tuple=True)[0][0] + ssl_mask[b, first_valid] = True + 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) @@ -207,11 +244,16 @@ def extract_visible_timesteps( """ B, T, d_model = tokens.shape - n_visible = ssl_mask.sum(dim=1) # (B,) + 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 = ssl_mask.float().argsort(dim=1, descending=True, stable=True) + 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) diff --git a/src/slices/models/pretraining/ts2vec.py b/src/slices/models/pretraining/ts2vec.py index ab74e9e..df7b991 100644 --- a/src/slices/models/pretraining/ts2vec.py +++ b/src/slices/models/pretraining/ts2vec.py @@ -161,44 +161,74 @@ def forward( T = crop_len # 3. Tokenize each view independently (noise means different MLP inputs) - tokens_1, _, _ = self.encoder.tokenize(x1, obs_mask) + 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) - ssl_mask_2 = create_timestep_mask(B, T, self.config.mask_ratio, device) + 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) + 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) + 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, ssl_mask_1, T) - full_2 = self._scatter_to_full(enc_2, ssl_mask_2, T) + 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 = ssl_mask_1 & ssl_mask_2 # (B, T) + 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, ssl_mask_1, ssl_mask_2) + 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(ssl_mask_1.sum().item() / B), - "ts2vec_n_visible_view2": torch.tensor(ssl_mask_2.sum().item() / B), + "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.float().mean(), + "ts2vec_overlap_ratio": overlap.sum().float() + / valid_timestep_mask.sum().clamp(min=1).float(), } ) diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index eb59385..6756a2b 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -369,6 +369,7 @@ def run_fairness_evaluation( fairness_cfg.get("protected_attributes", ["gender", "age_group"]) ), min_subgroup_size=fairness_cfg.get("min_subgroup_size", 50), + 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) diff --git a/tests/test_contrastive_objective.py b/tests/test_contrastive_objective.py index 65b32f1..45a468e 100644 --- a/tests/test_contrastive_objective.py +++ b/tests/test_contrastive_objective.py @@ -187,6 +187,21 @@ def test_two_views_encoded(self, encoder, contrastive_config): 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 views 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 + # ============================================================================= # NT-Xent loss tests diff --git a/tests/test_dataset_datamodule.py b/tests/test_dataset_datamodule.py index acec995..0eb893f 100644 --- a/tests/test_dataset_datamodule.py +++ b/tests/test_dataset_datamodule.py @@ -925,6 +925,28 @@ def test_label_efficiency_splits_preserve_full_split_provenance(self, mock_extra 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.""" diff --git a/tests/test_export_results.py b/tests/test_export_results.py new file mode 100644 index 0000000..bb018e9 --- /dev/null +++ b/tests/test_export_results.py @@ -0,0 +1,43 @@ +"""Tests for the results export pipeline.""" + +import importlib + +import pandas as pd + + +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)]: + rows.append( + { + "experiment_type": "core", + "sprint": "1", + "paradigm": paradigm, + "dataset": "miiv", + "task": task, + "seed": seed, + "protocol": "A", + "label_fraction": 1.0, + "model_size": "default", + "source_dataset": None, + "phase": "finetune", + "test/auprc": base + offset + (seed / 100000.0), + } + ) + + per_seed_df = pd.DataFrame(rows) + stats_df = mod.build_statistical_tests_df(per_seed_df) + + assert len(stats_df) == 1 + row = stats_df.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["n_tasks"] == 2 + assert row["better_paradigm"] == "mae" + assert row["p_value_bonferroni"] >= row["p_value"] diff --git a/tests/test_fairness_evaluator.py b/tests/test_fairness_evaluator.py index d1466fd..77daf67 100644 --- a/tests/test_fairness_evaluator.py +++ b/tests/test_fairness_evaluator.py @@ -232,6 +232,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.""" @@ -293,6 +336,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().""" diff --git a/tests/test_fixes.py b/tests/test_fixes.py index bb5d807..5c1116a 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -380,6 +380,22 @@ def test_invariant_mismatch_raises(self): 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" + class TestCombinedSetupPath: """Regression tests for the combined-dataset setup flow.""" @@ -537,6 +553,8 @@ class TestFairnessRevisionScoping: 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, @@ -547,6 +565,29 @@ def test_parse_args_accepts_revision_filter(self, monkeypatch): 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_fetch_eval_runs_single_revision_adds_server_side_filter(self, monkeypatch): """A single revision should be sent as a W&B tag filter.""" mod = importlib.import_module("scripts.eval.evaluate_fairness") @@ -615,6 +656,80 @@ def runs(self, path, filters, order): 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_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] diff --git a/tests/test_jepa_objective.py b/tests/test_jepa_objective.py index bd7328b..ee21a9e 100644 --- a/tests/test_jepa_objective.py +++ b/tests/test_jepa_objective.py @@ -246,6 +246,23 @@ def test_encoder_sees_fewer_timesteps(self, encoder, jepa_config): 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) + # ============================================================================= # Momentum tests diff --git a/tests/test_statistical.py b/tests/test_statistical.py index 10a8d92..9e5d518 100644 --- a/tests/test_statistical.py +++ b/tests/test_statistical.py @@ -4,11 +4,19 @@ - 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 slices.eval.statistical import ( + bonferroni_correction, + bootstrap_ci, + cohens_d, + paired_bootstrap_test, + paired_wilcoxon_signed_rank, +) # ── Helpers ────────────────────────────────────────────────────────────────── @@ -206,3 +214,31 @@ 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 + + +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_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_transformer_encoder.py b/tests/test_transformer_encoder.py index fcdb82a..d7fc29b 100644 --- a/tests/test_transformer_encoder.py +++ b/tests/test_transformer_encoder.py @@ -891,10 +891,11 @@ def test_tokenize_shapes(self, encoder): assert tokens.shape == (B, T, 32) assert padding_mask.shape == (B, T) - assert padding_mask.all() # All True for fixed 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): diff --git a/tests/test_ts2vec_objective.py b/tests/test_ts2vec_objective.py new file mode 100644 index 0000000..e92027f --- /dev/null +++ b/tests/test_ts2vec_objective.py @@ -0,0 +1,47 @@ +"""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 From ab696e84390962de2aa06be6e295e9350dbeaaf4 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 14 Apr 2026 16:53:32 -0400 Subject: [PATCH 073/121] Align medium-priority experiment fixes Update stale AKI and SSL documentation, include Sprint 12 in the standalone fairness defaults, and run the tmux status pane through the uv-managed project environment. Unify fairness report flattening across training, standalone evaluation, and XGBoost logging, and make normalize=false impute feature means using train-scoped stats when available instead of effectively zero-filling. Behavioral implications: XGBoost fairness summaries now retain nested subgroup metrics; non-normalized dataset loading now uses mean imputation for missing values; the full test suite passes (989 passed). --- configs/ssl/contrastive.yaml | 6 ++-- configs/ssl/jepa.yaml | 6 ++-- configs/ssl/mae.yaml | 8 ++--- docs/EXPERIMENT_PLAN.md | 2 +- scripts/eval/evaluate_fairness.py | 20 ++--------- scripts/launch_thesis_tmux.sh | 2 +- scripts/training/xgboost_baseline.py | 10 ++---- src/slices/data/dataset.py | 18 ++++++++-- src/slices/data/tensor_cache.py | 5 +-- src/slices/data/tensor_preprocessing.py | 9 +++-- src/slices/eval/__init__.py | 3 +- src/slices/eval/fairness_evaluator.py | 19 +++++++++++ src/slices/models/pretraining/README.md | 12 +++---- src/slices/training/utils.py | 14 ++------ tests/test_dataset_datamodule.py | 31 ++++++++++++----- tests/test_fairness_evaluator.py | 23 ++++++++++++- tests/test_fixes.py | 44 +++++++++++++++++++++++++ 17 files changed, 156 insertions(+), 76 deletions(-) diff --git a/configs/ssl/contrastive.yaml b/configs/ssl/contrastive.yaml index 4a94bc8..b7b96ba 100644 --- a/configs/ssl/contrastive.yaml +++ b/configs/ssl/contrastive.yaml @@ -1,9 +1,9 @@ # @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) diff --git a/configs/ssl/jepa.yaml b/configs/ssl/jepa.yaml index d39dbb2..53ca21b 100644 --- a/configs/ssl/jepa.yaml +++ b/configs/ssl/jepa.yaml @@ -1,9 +1,9 @@ # @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 masking interface 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) diff --git a/configs/ssl/mae.yaml b/configs/ssl/mae.yaml index 45ac05d..c6fb5d0 100644 --- a/configs/ssl/mae.yaml +++ b/configs/ssl/mae.yaml @@ -1,10 +1,10 @@ # @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 diff --git a/docs/EXPERIMENT_PLAN.md b/docs/EXPERIMENT_PLAN.md index 0360889..3fbbf2d 100644 --- a/docs/EXPERIMENT_PLAN.md +++ b/docs/EXPERIMENT_PLAN.md @@ -521,4 +521,4 @@ Later sprints reuse runs from earlier sprints as comparison baselines. To enable | Compute budget exceeded | Sprint ordering ensures usable results at each checkpoint | | Shared hyperparameters unfair to one paradigm | LR sensitivity (Sprint 1b) and mask ratio sensitivity (Sprint 1c) validate that rankings are robust to shared hyperparameter choices | | Reviewer requests more seeds | Sprint `10` expands the main matrix to 5 seeds, and Sprints `7p`/`11`/`12`/`13` keep the same 4 d.f. footing for variance estimates | -| AKI label leakage | Fixed: forward-looking prediction window (hours 48–72) prevents model from seeing creatinine values used for KDIGO label construction | +| AKI label leakage | Fixed: forward-looking prediction window (hours 24–48) prevents model from seeing creatinine values used for KDIGO label construction | diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index 6ed56b2..aecf6f6 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -43,6 +43,7 @@ from typing import Any, Optional import torch +from slices.eval.fairness_evaluator import flatten_fairness_report log = logging.getLogger("evaluate_fairness") @@ -50,7 +51,7 @@ # Constants # --------------------------------------------------------------------------- -CORE_SPRINTS = ["1", "2", "3", "4", "5", "7p", "10", "13"] +CORE_SPRINTS = ["1", "2", "3", "4", "5", "7p", "10", "12", "13"] DEFAULT_PHASES = ["finetune", "supervised"] DEFAULT_PROTECTED_ATTRIBUTES = ["gender", "age_group", "race"] @@ -442,23 +443,6 @@ def evaluate_run_fairness( return report -def flatten_fairness_report(report: dict[str, Any]) -> dict[str, Any]: - """Flatten nested fairness report into flat fairness/* keys for W&B. - - Mirrors the flattening in run_fairness_evaluation() (training/utils.py). - """ - flat: dict[str, Any] = {} - for attr, metrics in report.items(): - for metric_name, value in metrics.items(): - if isinstance(value, (int, float)): - flat[f"fairness/{attr}/{metric_name}"] = value - elif isinstance(value, dict): - for sub_key, sub_val in value.items(): - if isinstance(sub_val, (int, float)): - flat[f"fairness/{attr}/{metric_name}/{sub_key}"] = sub_val - return flat - - # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- diff --git a/scripts/launch_thesis_tmux.sh b/scripts/launch_thesis_tmux.sh index b560963..bc1a485 100755 --- a/scripts/launch_thesis_tmux.sh +++ b/scripts/launch_thesis_tmux.sh @@ -140,7 +140,7 @@ chmod +x "$runner_script" status_cmd=( bash -lc - "cd $(printf "%q" "$REPO_ROOT"); while true; do clear; date; echo; python scripts/run_experiments.py status; sleep $(printf "%q" "$STATUS_INTERVAL"); done" + "cd $(printf "%q" "$REPO_ROOT"); while true; do clear; date; echo; uv run python scripts/run_experiments.py status; sleep $(printf "%q" "$STATUS_INTERVAL"); done" ) printf -v status_line "%q " "${status_cmd[@]}" diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index 9891648..38fad3d 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -256,7 +256,7 @@ def main(cfg: DictConfig) -> None: print("5. Fairness Evaluation") print("=" * 80) - from slices.eval.fairness_evaluator import FairnessEvaluator + 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] @@ -302,13 +302,7 @@ def main(cfg: DictConfig) -> None: ) run.summary.update(metrics) if fairness_report: - run.summary.update( - { - f"fairness/{k}": v - for k, v in fairness_report.items() - if isinstance(v, (int, float)) - } - ) + run.summary.update(flatten_fairness_report(fairness_report)) wandb.finish() # ========================================================================= diff --git a/src/slices/data/dataset.py b/src/slices/data/dataset.py index 0a3126d..2679fdd 100644 --- a/src/slices/data/dataset.py +++ b/src/slices/data/dataset.py @@ -122,7 +122,9 @@ 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. @@ -321,7 +323,7 @@ def _precompute_tensors( n_samples = timeseries_tensor.shape[0] logger.info(f"Preprocessing {n_samples:,} samples...") - # Step 2: Compute normalization statistics + # 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) @@ -348,6 +350,18 @@ def _precompute_tensors( "train_indices must be provided when normalize=True to prevent data leakage. " "Pass train_indices from your data splits, or set normalize=False." ) + else: + 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( diff --git a/src/slices/data/tensor_cache.py b/src/slices/data/tensor_cache.py index ace6bee..09c068c 100644 --- a/src/slices/data/tensor_cache.py +++ b/src/slices/data/tensor_cache.py @@ -126,7 +126,7 @@ def load_normalization_stats( current_train_indices: Optional[List[int]], normalize: bool, ) -> Optional[Dict[str, Any]]: - """Load cached normalization statistics from file if they exist and match current split. + """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. @@ -141,9 +141,6 @@ def load_normalization_stats( Dictionary with 'feature_means' and 'feature_stds' if file exists and split matches, None otherwise. """ - if not normalize: - return None - expected_data_fingerprint = get_data_fingerprint(data_dir) expected_preprocessing_fingerprint = get_preprocessing_fingerprint() diff --git a/src/slices/data/tensor_preprocessing.py b/src/slices/data/tensor_preprocessing.py index ac941e4..9b12693 100644 --- a/src/slices/data/tensor_preprocessing.py +++ b/src/slices/data/tensor_preprocessing.py @@ -235,8 +235,8 @@ def apply_normalization_and_imputation( """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 (avoids physiologically - impossible values like 0 heart rate). + 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). @@ -258,9 +258,8 @@ def apply_normalization_and_imputation( # 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. + # 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] diff --git a/src/slices/eval/__init__.py b/src/slices/eval/__init__.py index be74974..29293d8 100644 --- a/src/slices/eval/__init__.py +++ b/src/slices/eval/__init__.py @@ -7,7 +7,7 @@ - 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 ( @@ -28,6 +28,7 @@ "build_metrics", "get_default_metrics", "FairnessEvaluator", + "flatten_fairness_report", "ImputationEvaluator", "bonferroni_correction", "bootstrap_ci", diff --git a/src/slices/eval/fairness_evaluator.py b/src/slices/eval/fairness_evaluator.py index f61396c..89d785e 100644 --- a/src/slices/eval/fairness_evaluator.py +++ b/src/slices/eval/fairness_evaluator.py @@ -12,6 +12,7 @@ import logging import re +from numbers import Real from typing import Any, Dict, List, Optional, Tuple import polars as pl @@ -28,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. diff --git a/src/slices/models/pretraining/README.md b/src/slices/models/pretraining/README.md index 4644fec..49088d2 100644 --- a/src/slices/models/pretraining/README.md +++ b/src/slices/models/pretraining/README.md @@ -6,11 +6,12 @@ This module implements various SSL objectives for pretraining on unlabeled ICU t ### MAE (Masked Autoencoder) -Learns representations by masking portions of the input and reconstructing them. +Learns representations by masking timestep tokens and reconstructing the masked +features from the remaining context. **Key Features**: -- 4 masking strategies: random, block, timestep, feature -- Respects observation mask (ICU data-aware) +- Timestep-level masking aligned with the benchmark encoder +- Obs-aware tokenization that embeds values together with the observation mask - Lightweight transformer decoder - Configurable mask ratio and decoder architecture @@ -19,8 +20,7 @@ Learns representations by masking portions of the input and reconstructing them. from slices.models.pretraining import MAEConfig, MAEObjective mae_config = MAEConfig( - mask_ratio=0.15, - mask_strategy="block", + mask_ratio=0.5, decoder_d_model=64, decoder_n_layers=2, ) @@ -67,7 +67,7 @@ pretraining/ ## Requirements All SSL objectives require: -- Encoder with `pooling="none"` (for per-timestep outputs) +- Encoder with `pooling="none"` and `obs_aware=True` - Input shape: `(B, T, D)` - Observation mask shape: `(B, T, D)` diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 6756a2b..73754ee 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -354,7 +354,7 @@ 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 predictions, labels_tensor, all_stay_ids = run_inference( @@ -375,16 +375,8 @@ def run_fairness_evaluation( 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 diff --git a/tests/test_dataset_datamodule.py b/tests/test_dataset_datamodule.py index 0eb893f..55d79c8 100644 --- a/tests/test_dataset_datamodule.py +++ b/tests/test_dataset_datamodule.py @@ -159,8 +159,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 +169,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.""" diff --git a/tests/test_fairness_evaluator.py b/tests/test_fairness_evaluator.py index 77daf67..dfd37a7 100644 --- a/tests/test_fairness_evaluator.py +++ b/tests/test_fairness_evaluator.py @@ -13,7 +13,7 @@ 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( @@ -435,3 +435,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_fixes.py b/tests/test_fixes.py index 5c1116a..42ab689 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -588,6 +588,12 @@ def test_parse_args_requires_revision_when_no_scope_is_available(self, monkeypat assert excinfo.value.code == 2 + def test_default_core_sprints_include_sprint12(self): + """The standalone thesis fairness corpus should include Sprint 12.""" + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + assert "12" in mod.CORE_SPRINTS + def test_fetch_eval_runs_single_revision_adds_server_side_filter(self, monkeypatch): """A single revision should be sent as a W&B tag filter.""" mod = importlib.import_module("scripts.eval.evaluate_fairness") @@ -774,6 +780,44 @@ def test_tmux_launcher_passes_revision_to_fairness(self, tmp_path): assert "scripts/eval/evaluate_fairness.py" in runner_text assert "--revision thesis-v2" in runner_text + 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").mkdir(parents=True, exist_ok=True) + shutil.copy2(repo_root / "scripts" / "launch_thesis_tmux.sh", temp_repo / "scripts") + + 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["RUN_EXPORT"] = "0" + + result = subprocess.run( + ["bash", "scripts/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/run_experiments.py\\ status" in tmux_log.read_text() + # ============================================================================ # Issue 4: Sequence-length override From 7bd520ad6957184ba7b76917b93f63ee74c282fe Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 14 Apr 2026 17:27:52 -0400 Subject: [PATCH 074/121] Refresh tracked documentation for the current thesis workflow Update the tracked repository docs to match the latest experiment design, fairness evaluation flow, and high-level testing guidance. This keeps the main markdown entry points aligned with the current code and analysis process. Behavioral implications: documentation only; no runtime behavior, benchmark results, or migration steps change in this commit. --- README.md | 26 +- docs/EXPERIMENT_PLAN.md | 39 +- src/slices/eval/README.md | 305 ++++----------- src/slices/models/encoders/README.md | 336 +++++------------ src/slices/models/pretraining/README.md | 160 +++++--- src/slices/training/README.md | 470 +++++------------------- tests/README.md | 173 +++------ 7 files changed, 454 insertions(+), 1055 deletions(-) diff --git a/README.md b/README.md index ca2018e..8d65a8d 100644 --- a/README.md +++ b/README.md @@ -160,8 +160,9 @@ 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 @@ -177,8 +178,10 @@ 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/ts2vec/smart) @@ -191,7 +194,7 @@ SLICES/ ├── scripts/ # Entry point scripts │ ├── preprocessing/ │ └── training/ -└── tests/ # pytest test suite (900+ tests) +└── tests/ # pytest test suite ``` ## Configuration @@ -242,9 +245,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 @@ -258,9 +261,12 @@ mypy src/ ## 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 task definitions ## References diff --git a/docs/EXPERIMENT_PLAN.md b/docs/EXPERIMENT_PLAN.md index 3fbbf2d..50d8da3 100644 --- a/docs/EXPERIMENT_PLAN.md +++ b/docs/EXPERIMENT_PLAN.md @@ -222,7 +222,8 @@ The repo now also exposes an optional `mortality` ICU-mortality task, but it is ### 4.3 Phase 3: Fairness Evaluation (on all downstream runs) -No additional training. Compute fairness metrics on test predictions from Phase 2. +No additional training. Compute fairness metrics by rerunning inference from the +same evaluation checkpoint provenance recorded for each downstream run. --- @@ -333,7 +334,8 @@ No additional training. Compute fairness metrics on test predictions from Phase | 12 | SMART external SSL reference, 5 seeds | 135 | TODO | | 9 | Fairness (eval only, runs last after all training) | 0 | TODO | -For detailed sprint outcomes, see `EXPERIMENT_RESULTS.md`. +For concrete rerun outputs, use the exported result artifacts under `results/` +and the thesis W&B project summaries. ### Upcoming: Sprint 10 — Extra Seeds 789, 1011 (630 runs) @@ -410,7 +412,7 @@ For detailed sprint outcomes, see `EXPERIMENT_RESULTS.md`. **Runs last**, after all training sprints are complete. Zero additional training — pure evaluation on existing test predictions. 1. Compute fairness metrics on all finetune/supervised test predictions from Sprints 1–5, 7p, 10, 12, and 13 -2. Enable `eval.fairness.enabled=true` and rerun evaluation +2. Run `scripts/eval/evaluate_fairness.py` with an explicit `--revision` tag for the thesis rerun corpus 3. Generate fairness tables and disparity plots 4. Protected attributes: sex (all datasets), age group (all), race/ethnicity (MIMIC-IV only) 5. Sprint `11` classical baselines are not part of the default standalone fairness sweep because `evaluate_fairness.py` currently targets `finetune` and `supervised` phases only @@ -425,37 +427,40 @@ Legacy exploratory runs remain in W&B project `slices`. Final thesis reruns shou ### Run Naming Convention -``` -{phase}_{dataset}_{paradigm}_{task}_seed{seed} -``` +Run names are phase-specific and generated from the Hydra configs rather than a +single universal template. The current defaults use the `s${sprint}_...` +pattern from `configs/*.yaml`. 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` +- `s10_pretrain_miiv_mae_seed42` +- `s10_finetune_eicu_jepa_mortality_24h_seed123` +- `s10_supervised_combined_aki_kdigo_seed456` +- `s11_xgboost_miiv_mortality_hospital_seed42` + +Downstream tooling should rely on tags and structured config fields rather than +parsing run names. ### Tags | Tag | Purpose | |-----|---------| -| `phase:pretrain` / `phase:finetune` / `phase:supervised` | Training phase | -| `dataset:mimic` / `dataset:eicu` / `dataset:combined` | Dataset | +| `phase:pretrain` / `phase:finetune` / `phase:supervised` / `phase:baseline` | Training phase | +| `dataset:miiv` / `dataset:eicu` / `dataset:combined` | Dataset | | `paradigm:mae` / `paradigm:jepa` / `paradigm:contrastive` / `paradigm:supervised` / `paradigm:xgboost` / `paradigm:gru_d` | Paradigm | | `task:mortality_24h` / etc. | Downstream task | | `ablation:label-efficiency` / `ablation:transfer` | Ablation type | | `sprint:N` | Execution sprint | +| `revision:` | Thesis rerun revision for fairness/export filtering | | `seed:42` / `seed:123` / `seed:456` / `seed:789` / `seed:1011` | 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 | +| Pretraining | `train/loss`, `val/loss`, `train/gradient_steps`, `train/wall_clock_seconds` | +| Classification downstream | `val/auroc`, `val/auprc`, `test/auroc`, `test/auprc`, `test/brier_score`, `test/ece` | +| Regression downstream | `val/loss`, `test/mae`, `test/mse`, `test/r2` | +| Fairness | `fairness/gender/*`, `fairness/age_group/*`, `fairness/race/*` | ### Baseline Inheritance Across Sprints diff --git a/src/slices/eval/README.md b/src/slices/eval/README.md index aff5355..848131b 100644 --- a/src/slices/eval/README.md +++ b/src/slices/eval/README.md @@ -1,275 +1,122 @@ # 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"], + metrics=["auroc", "auprc", "brier_score", "ece"], + 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 +Supported task types: -# Get default metrics for task type -val_metrics = get_default_metrics("binary", prefix="val") -# Uses default: [auroc, auprc] -``` +- `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` -### 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 -``` +Default metric sets are intentionally small: -## Available Metrics +- Binary: `auroc`, `auprc`, `brier_score`, `ece` +- Multiclass: `auroc`, `accuracy` +- Multilabel: `auroc` +- Regression: `mse`, `mae`, `r2` -### Binary Classification -- `auroc` - Area Under ROC Curve (default) -- `auprc` - Area Under Precision-Recall Curve (default) -- `accuracy` - Accuracy at 0.5 threshold -- `f1` - F1 Score +## Fairness -### Multiclass Classification -- `auroc` - One-vs-Rest AUROC (default) -- `auprc` - One-vs-Rest AUPRC -- `accuracy` - Multiclass accuracy (default) -- `f1` - Macro F1 Score +There are two layers: -### Multilabel Classification -- `auroc` - Per-label AUROC (default) -- `auprc` - Per-label AUPRC -- `accuracy` - Subset accuracy -- `f1` - Macro F1 Score +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. -### Regression -- Not yet implemented (placeholder) -- Planned: MSE, MAE, R², RMSE +Current evaluator behavior matches the thesis pipeline: -## Integration with Training +- 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()` -The `FineTuneModule` automatically uses eval metrics: +Typical usage: ```python -from slices.training import FineTuneModule +from slices.eval import FairnessEvaluator, flatten_fairness_report -# Metrics are built from config -module = FineTuneModule(config) - -# During training, metrics are logged automatically -# val/auroc, val/auprc, test/auroc, test/auprc, etc. -``` - -Metrics are read from `config.eval.metrics.names` if specified, otherwise defaults are used. - -## Extending with New Metrics - -### Adding a New Metric - -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) - # ... -``` - -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 - -- **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 +For batch thesis 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. -### Available Demographics (MIMIC-IV) +## Statistical Utilities -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 +`statistical.py` provides the utilities used by the export pipeline: -## Design Principles +- `bootstrap_ci` +- `paired_bootstrap_test` +- `paired_wilcoxon_signed_rank` +- `bonferroni_correction` +- `cohens_d` -### Minimal but Extendable -- Start with essential metrics (AUROC, AUPRC) -- Easy to add more without breaking existing code -- No unnecessary complexity +These are exported from `slices.eval` for direct reuse. -### Configuration-Driven -- Metrics specified in YAML config files -- Sensible defaults that "just work" -- Override when needed for experiments +## Inference Helper -### Separation of Concerns -- Evaluation logic separate from training logic -- Metrics independent of Lightning module -- Can be used standalone or in training +`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. -### Clinical Focus -- Metrics relevant to clinical prediction (AUROC, calibration) -- Demographics available for fairness analysis -- Designed for high-stakes decision support +## Imputation Evaluation -## Examples +`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 thesis +headline results. -### 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/models/encoders/README.md b/src/slices/models/encoders/README.md index 5340603..9d51078 100644 --- a/src/slices/models/encoders/README.md +++ b/src/slices/models/encoders/README.md @@ -1,288 +1,138 @@ -# 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 thesis 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) +encoder = TransformerEncoder(config) +embeddings = encoder(x, 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: +For SSL: ```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) +ssl_config = TransformerConfig( + d_input=35, + d_model=64, + n_layers=2, + n_heads=4, + d_ff=256, + obs_aware=True, + pooling="none", +) ``` -**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. +## Other Encoders In This Package -### Why Sinusoidal Positional Encoding? +- `ObservationTransformerEncoder` in `observation.py` + - observation-level tokenization + - kept as an alternate architecture, not the canonical thesis 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 -- **Extrapolation**: Can handle sequences longer than seen during training -- **Deterministic**: No learned parameters, more interpretable -- **Standard**: Works well for time-series (timestamps are ordered) +## Which Encoder Goes With Which Objective -Alternative: Learned positional embeddings (not yet implemented). +- `mae`, `jepa`, `contrastive`, `ts2vec` + - use `TransformerEncoder` + - require `obs_aware=True` + - require `pooling="none"` during pretraining +- `smart` + - requires `SMARTEncoder` + - requires `pooling="none"` -### Why Mean Pooling by Default? +The compatibility checks are enforced in `scripts/training/pretrain.py`. -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 +## Notes On Missingness And Padding -For supervised tasks, CLS pooling may work better (requires fine-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()`. -## 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() -``` - -After pretraining, extract the encoder: -```python -pretrained_encoder = ssl_model.get_encoder() -``` - -## Performance Tips - -### Memory Efficiency - -- Use gradient checkpointing for very deep models (not yet implemented) -- Use `torch.no_grad()` for inference -- Reduce batch size if OOM - -### Training Speed - -- 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 +## Testing -### Model Size vs. Performance +Relevant test coverage lives in: -For ICU data (typical benchmark setup: 35 features, 24-hour observation 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 +- `tests/test_transformer_encoder.py` +- `tests/test_smart_encoder.py` +- `tests/test_factories.py` -## Testing +Example script: -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 ``` - -## References - -- **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 diff --git a/src/slices/models/pretraining/README.md b/src/slices/models/pretraining/README.md index 49088d2..c762bd5 100644 --- a/src/slices/models/pretraining/README.md +++ b/src/slices/models/pretraining/README.md @@ -1,81 +1,131 @@ # 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 - -### MAE (Masked Autoencoder) +## Implemented Objectives -Learns representations by masking timestep tokens and reconstructing the masked -features from the remaining context. +- `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` -**Key Features**: -- Timestep-level masking aligned with the benchmark encoder -- Obs-aware tokenization that embeds values together with the observation mask -- Lightweight transformer decoder -- Configurable mask ratio and decoder architecture +The registry lives in `factory.py`. -**Usage**: -```python -from slices.models.pretraining import MAEConfig, MAEObjective +## Encoder Requirements -mae_config = MAEConfig( - mask_ratio=0.5, - 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**: the objective source files in this package and the Hydra configs under `configs/ssl/`. +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 controlled thesis objectives 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 stronger temporal objective +- 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 -├── jepa.py # JEPA implementation -├── contrastive.py # Contrastive implementation -├── smart.py # SMART sanity-check objective -└── 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 -All SSL objectives require: -- Encoder with `pooling="none"` and `obs_aware=True` -- Input shape: `(B, T, D)` -- Observation mask shape: `(B, T, D)` +### MAE -## Implemented Objectives +- 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 + +- uses the same timestep-token interface as MAE +- target encoder is an EMA copy kept in eval mode +- supports `mse` and `cosine` latent losses + +### 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 + +### 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 +``` -The active codebase already includes MAE, JEPA, contrastive, and SMART objectives. +## 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/training/README.md b/src/slices/training/README.md index e927ab9..b1cfde6 100644 --- a/src/slices/training/README.md +++ b/src/slices/training/README.md @@ -1,430 +1,152 @@ # 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 +```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 ``` -#### 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 -``` +### `FineTuneModule` -## Adding New Components +Composes: -### Adding a New Encoder +- an encoder +- an optional missing-token wrapper +- a task head +- metric collections for train, val, and test -1. Create encoder class in `src/slices/models/encoders/`: +Supported task types: -```python -from .base import BaseEncoder, EncoderConfig +- `binary` +- `multiclass` +- `multilabel` +- `regression` -@dataclass -class MyEncoderConfig(EncoderConfig): - # Add encoder-specific params - hidden_size: int = 256 +Checkpoint inputs: -class MyEncoder(BaseEncoder): - def __init__(self, config: MyEncoderConfig): - super().__init__(config) - # ... build encoder +- encoder checkpoint from `encoder.pt` +- full Lightning pretrain checkpoint via `pretrain_checkpoint` +- no checkpoint for supervised-from-scratch runs - def forward(self, x, mask=None, padding_mask=None): - # ... encode input - return encoded -``` +What it handles: -2. Register in `src/slices/models/encoders/factory.py`: +- 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 -```python -from .my_encoder import MyEncoder, MyEncoderConfig +The main entrypoints are: -ENCODER_REGISTRY["my_encoder"] = MyEncoder -ENCODER_CONFIG_REGISTRY["my_encoder"] = MyEncoderConfig -``` +- `scripts/training/finetune.py` +- `scripts/training/supervised.py` -3. Create config file `configs/encoder/my_encoder.yaml`: - -```yaml -name: my_encoder -hidden_size: 256 -# ... other params -``` - -4. Use in pretraining: +Examples: ```bash -uv run python scripts/training/pretrain.py encoder=my_encoder +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 ``` -### Adding a New SSL Objective +## Configuration Layout -1. Create objective class in `src/slices/models/pretraining/`: +The training scripts use Hydra config groups: -```python -from .base import BaseSSLObjective, SSLConfig +- `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 -@dataclass -class MySSLConfig(SSLConfig): - name: str = "my_ssl" - # Add objective-specific params +Note that custom encoders belong in `configs/model/`, not `configs/encoder/`. -class MySSLObjective(BaseSSLObjective): - def __init__(self, encoder, config: MySSLConfig): - super().__init__(encoder, config) - # ... build objective-specific components +## Checkpoint Formats - def forward(self, x, mask): - # ... compute SSL loss - return loss, metrics -``` +The downstream path supports two checkpoint styles: -2. Register in `src/slices/models/pretraining/factory.py`: +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 -```python -from .my_ssl import MySSLObjective, MySSLConfig +`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. -SSL_REGISTRY["my_ssl"] = MySSLObjective -CONFIG_REGISTRY["my_ssl"] = MySSLConfig -``` +## Utilities -3. Create config file `configs/ssl/my_ssl.yaml`: +`utils.py` contains shared helpers for: -```yaml -name: my_ssl -# ... ssl-specific params -``` - -4. Use in pretraining: - -```bash -uv run python scripts/training/pretrain.py ssl=my_ssl -``` +- optimizer and scheduler construction +- callback setup +- W&B logger setup +- checkpoint export +- label-support validation +- fairness evaluation hooks +- data prerequisite checks -### `FineTuneModule` +## Extending -Lightning module for downstream task finetuning. This module composes a pretrained encoder with a task head for supervised learning on clinical prediction tasks. +To add a new encoder: -#### Key Features +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 -- **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 +To add a new SSL objective: -#### Usage +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 -```python -from omegaconf import OmegaConf -from slices.training import FineTuneModule +To add a new task head: -# Load config (typically from Hydra) -config = OmegaConf.load("configs/finetune.yaml") - -# Create module with pretrained encoder -module = FineTuneModule( - config=config, - checkpoint_path="outputs/encoder.pt", # From pretraining -) - -# Train with Lightning Trainer -import lightning.pytorch as L -trainer = L.Trainer(max_epochs=50, devices=1) -trainer.fit(module, datamodule=datamodule) -``` - -#### 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 - -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/tests/README.md b/tests/README.md index ac0f993..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_dataset_datamodule.py -v -``` +Inspect the current collected tests: -### Run specific test class ```bash -uv run pytest tests/test_dataset_datamodule.py::TestICUDataset -v +uv run pytest --collect-only -q tests ``` -### Run specific test function -```bash -uv run pytest tests/test_fixes.py::TestNormalizationStatsCache::test_legacy_stats_without_fingerprints_are_ignored -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_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_ricu_extractor.py` | RICU extractor compatibility | Chunked parquet ingestion, schema migration, extraction edge cases | -| `test_fixes.py` | Regression coverage for experiment-integrity fixes | Label freshness, cache invalidation, combined-dataset safety | +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_transforms.py` - SSL/data augmentation 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_ricu_extractor.py::TestRicuExtractorEdgeCases` - Missing schema pieces, chunked input, legacy migration -- `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 task variants (24h, hospital, ICU) plus window-boundary cases -- 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. From 323a1faa3aecc154072bc8e034b6a9c2cd8860cd Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 14 Apr 2026 17:28:54 -0400 Subject: [PATCH 075/121] Guard sparse contrastive views and clarify fairness thresholds Prevent instance-contrastive complementary masks from producing an empty eligible view after empty-timestep filtering by adding a sparse-sample fallback and regression coverage. Clarify the standalone fairness CLI help text so the subgroup threshold is described in patients, matching the evaluator's behavior. Behavioral implication: sparse contrastive batches no longer produce NaN losses under the default complementary-mask setting. No migration or benchmark config changes are required. --- scripts/eval/evaluate_fairness.py | 2 +- src/slices/models/pretraining/contrastive.py | 15 +++++-- src/slices/models/pretraining/masking.py | 44 ++++++++++++++++++++ tests/test_contrastive_objective.py | 24 +++++++++++ 4 files changed, 81 insertions(+), 4 deletions(-) diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index aecf6f6..ccf2e27 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -514,7 +514,7 @@ def parse_args() -> argparse.Namespace: help=f"Attributes to evaluate (default: {DEFAULT_PROTECTED_ATTRIBUTES})", ) parser.add_argument( - "--min-subgroup-size", type=int, default=50, help="Min samples per subgroup" + "--min-subgroup-size", type=int, default=50, help="Min patients per subgroup" ) parser.add_argument("--dry-run", action="store_true", help="List runs without processing") parser.add_argument( diff --git a/src/slices/models/pretraining/contrastive.py b/src/slices/models/pretraining/contrastive.py index d590b50..96e41c1 100644 --- a/src/slices/models/pretraining/contrastive.py +++ b/src/slices/models/pretraining/contrastive.py @@ -38,7 +38,11 @@ import torch.nn.functional as F from .base import BaseSSLObjective, SSLConfig -from .masking import create_timestep_mask, extract_visible_timesteps +from .masking import ( + create_complementary_timestep_masks, + create_timestep_mask, + extract_visible_timesteps, +) @dataclass @@ -60,7 +64,9 @@ class ContrastiveConfig(SSLConfig): # Temperature for NT-Xent temperature: float = 0.07 - # Use complementary (non-overlapping) masks for views + # 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: @@ -178,7 +184,10 @@ def forward( valid_timestep_mask=valid_timestep_mask, ) if self.config.complementary_masks: - ssl_mask_2 = ~ssl_mask_1 + ssl_mask_1, ssl_mask_2 = create_complementary_timestep_masks( + ssl_mask_1, + valid_timestep_mask, + ) else: ssl_mask_2 = create_timestep_mask( B, diff --git a/src/slices/models/pretraining/masking.py b/src/slices/models/pretraining/masking.py index 528eeed..c57757a 100644 --- a/src/slices/models/pretraining/masking.py +++ b/src/slices/models/pretraining/masking.py @@ -127,6 +127,50 @@ def create_timestep_mask( 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 create_block_timestep_mask( batch_size: int, n_timesteps: int, diff --git a/tests/test_contrastive_objective.py b/tests/test_contrastive_objective.py index 45a468e..f3987b1 100644 --- a/tests/test_contrastive_objective.py +++ b/tests/test_contrastive_objective.py @@ -202,6 +202,30 @@ def test_empty_timesteps_are_excluded_from_visible_counts(self, encoder, contras assert metrics["contrastive_n_visible_view1"] <= 2 assert metrics["contrastive_n_visible_view2"] <= 2 + 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 + # ============================================================================= # NT-Xent loss tests From c2ef1c6676772e7b4305545b0186183f7cd61580 Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 16 Apr 2026 08:16:07 +0100 Subject: [PATCH 076/121] Normalize RICU mortality timestamps before label comparisons Strip timezone metadata from ICU stay and mortality timestamp columns before the precision-aware mortality label logic so UTC-aware RICU exports compare cleanly against date-only death fields. Add a regression test covering UTC-aware stay times and bump the mortality label semantic version so stale processed labels are flagged for regeneration. --- src/slices/data/labels/mortality.py | 29 ++++++++++++++++++++++++++--- tests/test_fixes.py | 27 ++++++++++++++++++++++++++- 2 files changed, 52 insertions(+), 4 deletions(-) diff --git a/src/slices/data/labels/mortality.py b/src/slices/data/labels/mortality.py index b59cf05..a191a67 100644 --- a/src/slices/data/labels/mortality.py +++ b/src/slices/data/labels/mortality.py @@ -37,7 +37,7 @@ class MortalityLabelBuilder(LabelBuilder): - "unknown" or missing: falls back to hospital_expire_flag """ - SEMANTIC_VERSION = "2.1.0" + SEMANTIC_VERSION = "2.1.1" def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: """Build mortality labels from stay and mortality data. @@ -133,9 +133,10 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: @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) + if "death_time_precision" in merged.columns: - # Ensure death_time column is tz-naive datetime for comparison with - # intime/outtime (which are always tz-naive in SLICES) + # 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 ( @@ -197,6 +198,28 @@ def _ensure_precision_columns(merged: pl.DataFrame) -> pl.DataFrame: pl.lit("legacy").alias("death_source"), ) + @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 + def _build_hospital_mortality_with_obs( self, merged: pl.DataFrame, obs_hours: int ) -> pl.DataFrame: diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 42ab689..aed9bc9 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -10,7 +10,7 @@ import shutil import subprocess import sys -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from pathlib import Path from types import SimpleNamespace @@ -1023,6 +1023,31 @@ def test_date_only_before_observation_end(self): 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) From f4afef89df8b0118a3a233945800adf71c56a701 Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 16 Apr 2026 08:28:19 +0100 Subject: [PATCH 077/121] Mask fully unobserved timesteps in JEPA predictor Pass the valid timestep mask into the JEPA predictor transformer so fully unobserved hours cannot participate as attention context for masked valid targets. Add regression coverage that inspects the predictor call and verifies the expected padding mask is applied for sparse hourly inputs.\n\nBehavioral implication: JEPA no longer learns from empty-hour positional structure through predictor attention. No migration or benchmark configuration changes are required. --- src/slices/models/pretraining/jepa.py | 16 +++++++++++--- tests/test_jepa_objective.py | 31 +++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 3 deletions(-) diff --git a/src/slices/models/pretraining/jepa.py b/src/slices/models/pretraining/jepa.py index cad7d64..9655659 100644 --- a/src/slices/models/pretraining/jepa.py +++ b/src/slices/models/pretraining/jepa.py @@ -110,7 +110,9 @@ def forward( Args: encoded_visible: (B, n_vis, d_encoder) encoded visible tokens. ssl_mask: (B, T) bool, True = visible, False = masked. - token_info: dict with timestep_idx (B, T). + 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: @@ -135,8 +137,16 @@ def forward( full_tokens = full_tokens + self.time_pe[timestep_idx] full_tokens = self.embed_dropout(full_tokens) - # Run predictor transformer (no padding mask, all T valid) - decoded = self.predictor(full_tokens) + valid_timestep_mask = token_info.get("valid_timestep_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, + ) + + # 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, T, d_encoder) diff --git a/tests/test_jepa_objective.py b/tests/test_jepa_objective.py index ee21a9e..1597981 100644 --- a/tests/test_jepa_objective.py +++ b/tests/test_jepa_objective.py @@ -263,6 +263,37 @@ def test_empty_timesteps_are_excluded_from_ssl_counts(self, encoder, jepa_config "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) + # ============================================================================= # Momentum tests From 1b14b88e71d8c38b21b27d7a3550160a0a100447 Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 16 Apr 2026 14:20:25 +0100 Subject: [PATCH 078/121] Fix experiment audit gaps in export and fairness flows Require explicit revision scoping for result exports, let supervised and SSL runs pair across phase boundaries in statistical tables, and extend standalone fairness reruns to support baseline XGBoost artifacts while preserving recorded checkpoint provenance. Also fix regression fairness shape handling, validate training.allow_best_ckpt_fallback, and exclude fully unobserved timesteps from MAE masking. Behavioral implications: exports now fail closed without a revision, fairness reruns cover baseline runs, and LOS fairness metrics use the correct samplewise residuals. --- configs/finetune.yaml | 4 + configs/supervised.yaml | 4 + scripts/eval/evaluate_fairness.py | 121 ++++++++++++++++++++------ scripts/export_results.py | 21 ++++- scripts/training/xgboost_baseline.py | 82 +---------------- src/slices/eval/fairness_evaluator.py | 9 +- src/slices/eval/inference.py | 99 ++++++++++++++++++++- src/slices/models/pretraining/mae.py | 8 +- src/slices/training/config_schemas.py | 1 + tests/test_config_schemas.py | 27 ++++++ tests/test_evaluate_fairness.py | 37 ++++++++ tests/test_export_results.py | 35 +++++++- tests/test_fairness_evaluator.py | 30 +++++++ tests/test_inference.py | 30 +++++++ tests/test_mae_objective.py | 48 ++++++++++ 15 files changed, 437 insertions(+), 119 deletions(-) create mode 100644 tests/test_evaluate_fairness.py create mode 100644 tests/test_inference.py diff --git a/configs/finetune.yaml b/configs/finetune.yaml index 4121ef1..440d4e2 100644 --- a/configs/finetune.yaml +++ b/configs/finetune.yaml @@ -84,6 +84,10 @@ 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 diff --git a/configs/supervised.yaml b/configs/supervised.yaml index 7367681..8265ae1 100644 --- a/configs/supervised.yaml +++ b/configs/supervised.yaml @@ -72,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 diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index ccf2e27..7260cd4 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -1,13 +1,14 @@ #!/usr/bin/env python3 """Post-run fairness evaluation for SLICES experiment runs. -Queries W&B for finished finetune/supervised runs, loads their best checkpoints, -runs inference on the test set, computes fairness metrics via FairnessEvaluator, -and writes results back to the same W&B run's summary. +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 thesis fairness corpus. Supports -resumability via --skip-existing (default) and scoping via --sprint/--paradigm/ ---dataset filters. +Designed for batch evaluation of the thesis fairness corpus across finetune, +supervised, and classical baseline runs. Supports resumability via +--skip-existing (default) and scoping via --sprint/--paradigm/--dataset filters. Usage: # Evaluate the thesis fairness corpus for one explicit revision @@ -52,7 +53,7 @@ # --------------------------------------------------------------------------- CORE_SPRINTS = ["1", "2", "3", "4", "5", "7p", "10", "12", "13"] -DEFAULT_PHASES = ["finetune", "supervised"] +DEFAULT_PHASES = ["finetune", "supervised", "baseline"] DEFAULT_PROTECTED_ATTRIBUTES = ["gender", "age_group", "race"] @@ -184,14 +185,19 @@ def _do_update(): 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 / "checkpoints" - return Path(output_dir) / "checkpoints" + return Path(outputs_root) / rel + return Path(output_dir) def _resolve_logged_checkpoint_path( @@ -333,6 +339,26 @@ def resolve_evaluation_checkpoint( ) +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 + + # --------------------------------------------------------------------------- # Model + data reconstruction # --------------------------------------------------------------------------- @@ -421,7 +447,6 @@ def evaluate_run_fairness( device: str, ) -> dict[str, Any]: """Run fairness evaluation on a single run.""" - from slices.eval.fairness_evaluator import FairnessEvaluator from slices.eval.inference import run_inference model = model.to(device) @@ -432,6 +457,30 @@ def evaluate_run_fairness( ) 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, @@ -439,8 +488,7 @@ def evaluate_run_fairness( task_type=task_type, dataset_name=getattr(getattr(datamodule, "processed_dir", None), "name", None), ) - report = evaluator.evaluate(predictions, labels, stay_ids) - return report + return evaluator.evaluate(predictions, labels, stay_ids) # --------------------------------------------------------------------------- @@ -607,8 +655,10 @@ def _sort_key(r): cfg = r.config task_type = _get_nested(cfg, "task.task_type", "binary") try: - ckpt, ckpt_source = resolve_evaluation_checkpoint(r, args.outputs_root, task_type) - ckpt_str = f"{ckpt} [{ckpt_source}]" if ckpt else "NOT FOUND" + 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( @@ -634,16 +684,12 @@ def _sort_key(r): log.info("[%d/%d] %s (%s/%s/seed%s)", i + 1, len(runs), run_desc, ds, task, seed) try: - # 1. Find checkpoint + # 1. Resolve the saved evaluation artifact task_type = _get_nested(cfg, "task.task_type", "binary") - ckpt_path, ckpt_source = resolve_evaluation_checkpoint( + artifact_path, artifact_source = resolve_evaluation_artifact( run, args.outputs_root, task_type ) - if ckpt_path is None: - results["skipped"] += 1 - results["errors"].append((run.id, run_desc, "no checkpoint found")) - continue - log.info(" Checkpoint: %s (%s)", ckpt_path, ckpt_source) + 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)) @@ -653,15 +699,37 @@ def _sort_key(r): datamodule = build_datamodule(cfg, args.batch_size, args.data_root) prev_dm_key = dm_key - model = build_model(cfg, ckpt_path, datamodule) + 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_run_fairness( - model, + report = evaluate_predictions_fairness( + predictions, + labels, + stay_ids, datamodule, args.protected_attributes, args.min_subgroup_size, - device, + task_type, ) if not report: @@ -680,7 +748,8 @@ def _sort_key(r): results["processed"] += 1 # Free model memory - del model + if model is not None: + del model if torch.cuda.is_available(): torch.cuda.empty_cache() diff --git a/scripts/export_results.py b/scripts/export_results.py index 657c4d2..fc23916 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -819,7 +819,7 @@ def build_statistical_tests_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: fingerprint = _fingerprint_for_experiment_type(experiment_type) scope_cols = [ - c for c in fingerprint if c in subset.columns and c not in {"paradigm", "task"} + c for c in fingerprint if c in subset.columns and c not in {"paradigm", "task", "phase"} ] if "primary_metric_name" not in scope_cols: scope_cols.append("primary_metric_name") @@ -1008,7 +1008,7 @@ def validate( # --------------------------------------------------------------------------- -def main(): +def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Export SLICES experiment results from W&B to parquet files.", formatter_class=argparse.RawDescriptionHelpFormatter, @@ -1066,8 +1066,23 @@ def main(): default=5, help="Expected seed count for core configs (default: 5)", ) - 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(): + args = parse_args() output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index 38fad3d..d249304 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -13,7 +13,6 @@ from pathlib import Path import hydra -import numpy as np import torch from omegaconf import DictConfig, OmegaConf from sklearn.metrics import ( @@ -25,89 +24,10 @@ roc_auc_score, ) from slices.data.datamodule import ICUDataModule +from slices.eval.inference import extract_tabular_features from xgboost import XGBClassifier, XGBRegressor -def extract_tabular_features(dataset, indices: list[int]) -> tuple[np.ndarray, np.ndarray]: - """Extract per-feature summary statistics from time-series data. - - For each of the D features, computes: mean, std, min, max, first, last, - obs_count, obs_fraction over observed values. Total = D * 8 features. - Vectorized over the full batch using the dataset's stacked tensors. - - Args: - dataset: ICUDataset instance. - indices: List of sample indices. - - Returns: - X: Feature matrix (N, D*8). - y: Label array (N,). - """ - 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, T, D = ts.shape - mask_float = mask.float() - - # obs_count, obs_fraction: (N, D) - obs_count = mask_float.sum(dim=1) # (N, D) - obs_frac = obs_count / T - - # Masked sum for mean: sum observed values / count - masked_ts = ts * mask_float # zero out unobserved - feat_sum = masked_ts.sum(dim=1) # (N, D) - safe_count = obs_count.clamp(min=1) - feat_mean = feat_sum / safe_count # (N, D) - - # Masked std - 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) - - # Min/max: fill unobserved with +inf/-inf before taking min/max - zeros = ts.new_zeros(N, D) - 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 and last observed value per feature - # Use argmax on mask along time axis for first; flip for last - # For features with no observations, falls back to 0 - first_idx = mask.float().argmax(dim=1) # (N, D) — first True index - last_idx = T - 1 - mask.flip(dims=[1]).float().argmax(dim=1) # (N, D) - - feat_first = ts.gather(1, first_idx.unsqueeze(1)).squeeze(1) # (N, D) - feat_last = ts.gather(1, last_idx.unsqueeze(1)).squeeze(1) # (N, D) - - # Zero out features with no observations - 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 - - # Stack: (N, D, 8) -> (N, D*8) - features = torch.stack( - [feat_mean, feat_std, feat_min, feat_max, feat_first, feat_last, obs_count, obs_frac], - dim=-1, - ) # (N, D, 8) - X = features.reshape(N, D * 8).numpy() - - # Labels - if dataset._labels_tensor is not None: - y = dataset._labels_tensor[idx_tensor].numpy() - else: - y = np.zeros(N, dtype=np.float32) - - return X, y - - @hydra.main(version_base=None, config_path="../../configs", config_name="xgboost") def main(cfg: DictConfig) -> None: """Train XGBoost baseline.""" diff --git a/src/slices/eval/fairness_evaluator.py b/src/slices/eval/fairness_evaluator.py index 89d785e..c63b45f 100644 --- a/src/slices/eval/fairness_evaluator.py +++ b/src/slices/eval/fairness_evaluator.py @@ -277,6 +277,11 @@ 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: @@ -428,8 +433,8 @@ def _evaluate_regression( 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() diff --git a/src/slices/eval/inference.py b/src/slices/eval/inference.py index 1f18b41..d39264f 100644 --- a/src/slices/eval/inference.py +++ b/src/slices/eval/inference.py @@ -5,8 +5,10 @@ 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 @@ -48,14 +50,103 @@ def run_inference( 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_preds.append(probs.squeeze(-1).cpu()) + all_labels.append(batch["label"].squeeze(-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) - labels = torch.cat(all_labels) + predictions = torch.cat(all_preds).squeeze(-1) + labels = torch.cat(all_labels).squeeze(-1) 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/models/pretraining/mae.py b/src/slices/models/pretraining/mae.py index 9b9012f..fe68335 100644 --- a/src/slices/models/pretraining/mae.py +++ b/src/slices/models/pretraining/mae.py @@ -212,7 +212,13 @@ def forward( # tokens: (B, T, d_model) # 2. Create SSL mask on timesteps - ssl_mask = create_timestep_mask(B, T, self.config.mask_ratio, device) + 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_timesteps(tokens, ssl_mask) diff --git a/src/slices/training/config_schemas.py b/src/slices/training/config_schemas.py index d2bd07e..1760f40 100644 --- a/src/slices/training/config_schemas.py +++ b/src/slices/training/config_schemas.py @@ -89,6 +89,7 @@ class TrainingConfig(BaseModel): early_stopping_mode: Optional[str] = None label_smoothing: float = 0.0 overfit_batches: Union[int, float] = 0 + allow_best_ckpt_fallback: bool = False class OptimizerConfig(BaseModel): diff --git a/tests/test_config_schemas.py b/tests/test_config_schemas.py index a5c455e..0f5c12f 100644 --- a/tests/test_config_schemas.py +++ b/tests/test_config_schemas.py @@ -126,6 +126,11 @@ 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 TestOptimizerConfigValidation: """Tests that OptimizerConfig catches invalid configs.""" @@ -176,3 +181,25 @@ 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 diff --git a/tests/test_evaluate_fairness.py b/tests/test_evaluate_fairness.py new file mode 100644 index 0000000..de11cd5 --- /dev/null +++ b/tests/test_evaluate_fairness.py @@ -0,0 +1,37 @@ +"""Focused tests for the standalone fairness rerun script.""" + +import importlib +from types import SimpleNamespace + + +def test_default_phases_include_baseline(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + assert "baseline" in mod.DEFAULT_PHASES + + +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" diff --git a/tests/test_export_results.py b/tests/test_export_results.py index bb018e9..9c0e921 100644 --- a/tests/test_export_results.py +++ b/tests/test_export_results.py @@ -1,6 +1,7 @@ """Tests for the results export pipeline.""" import importlib +import sys import pandas as pd @@ -9,7 +10,10 @@ 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 paradigm, phase, offset in [ + ("mae", "finetune", 0.10), + ("supervised", "supervised", 0.0), + ]: for seed in [42, 123, 456]: for task, base in [("mortality_24h", 0.30), ("aki_kdigo", 0.25)]: rows.append( @@ -24,7 +28,7 @@ def test_build_statistical_tests_df_produces_pairwise_significance_rows(): "label_fraction": 1.0, "model_size": "default", "source_dataset": None, - "phase": "finetune", + "phase": phase, "test/auprc": base + offset + (seed / 100000.0), } ) @@ -41,3 +45,30 @@ def test_build_statistical_tests_df_produces_pairwise_significance_rows(): assert row["n_tasks"] == 2 assert row["better_paradigm"] == "mae" assert row["p_value_bonferroni"] >= row["p_value"] + + +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.setattr(sys, "argv", ["export_results.py"]) + + args = mod.parse_args() + + assert args.revision == ["thesis-v1"] + + +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"]) + + try: + mod.parse_args() + except SystemExit as exc: + assert exc.code == 2 + else: + raise AssertionError("parse_args() should require an explicit revision scope") diff --git a/tests/test_fairness_evaluator.py b/tests/test_fairness_evaluator.py index dfd37a7..d312958 100644 --- a/tests/test_fairness_evaluator.py +++ b/tests/test_fairness_evaluator.py @@ -123,6 +123,36 @@ def test_selects_minimum(self): assert worst == pytest.approx(min(valid_aurocs)) +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) + + class TestAgeBinning: """Tests for age binning.""" diff --git a/tests/test_inference.py b/tests/test_inference.py new file mode 100644 index 0000000..22f549b --- /dev/null +++ b/tests/test_inference.py @@ -0,0 +1,30 @@ +"""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] diff --git a/tests/test_mae_objective.py b/tests/test_mae_objective.py index c7a93cd..5c062e5 100644 --- a/tests/test_mae_objective.py +++ b/tests/test_mae_objective.py @@ -297,6 +297,54 @@ def test_mae_respects_obs_mask(self, encoder, mae_config): n_observed = obs_mask.sum().item() assert n_loss <= n_observed + 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["valid_timestep_mask"] = valid_timestep_mask.clone() + return torch.ones(batch_size, n_timesteps, dtype=torch.bool, device=device) + + monkeypatch.setattr( + "slices.models.pretraining.mae.create_timestep_mask", + fake_create_timestep_mask, + ) + + 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["valid_timestep_mask"], expected) + def test_mae_get_encoder(self, encoder, mae_config): mae = MAEObjective(encoder, mae_config) assert mae.get_encoder() is encoder From ebf9a093ad08cd098ddc7b50000d293162ae2f9c Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 16 Apr 2026 14:49:30 +0100 Subject: [PATCH 079/121] Harden label freshness and extraction safety checks Validate label manifests against the active Hydra task config, auto-detect finetune paradigms from full pretrain checkpoints, and align training task schemas with the label semantics they depend on. This changes extractor behavior to fail closed on missing task configs, invalid LOS, missing timeseries coverage, and missing labels, and to overwrite rather than resume-merge when task configs or builder versions drift. Regression tests cover the new paths and the full pytest suite passes. --- configs/tasks/aki_kdigo.yaml | 14 ++ configs/tasks/los_remaining.yaml | 6 + configs/tasks/mortality.yaml | 7 + configs/tasks/mortality_24h.yaml | 7 + configs/tasks/mortality_hospital.yaml | 7 + scripts/training/finetune.py | 70 ++++++++-- scripts/training/supervised.py | 5 +- scripts/training/xgboost_baseline.py | 7 +- src/slices/data/extractors/base.py | 46 +++++-- src/slices/data/extractors/ricu.py | 14 +- src/slices/training/config_schemas.py | 12 +- src/slices/training/utils.py | 92 +++++++++---- tests/test_base_extractor.py | 42 +++++- tests/test_config_schemas.py | 20 +++ tests/test_fixes.py | 177 +++++++++++++++++++++++++ tests/test_ricu_extractor.py | 183 ++++++++++++++++++++++++++ tests/test_training_utils.py | 15 +++ 17 files changed, 665 insertions(+), 59 deletions(-) diff --git a/configs/tasks/aki_kdigo.yaml b/configs/tasks/aki_kdigo.yaml index f964ca1..409957b 100644 --- a/configs/tasks/aki_kdigo.yaml +++ b/configs/tasks/aki_kdigo.yaml @@ -4,6 +4,20 @@ 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 n_classes: null head_type: mlp diff --git a/configs/tasks/los_remaining.yaml b/configs/tasks/los_remaining.yaml index d6d602c..464b817 100644 --- a/configs/tasks/los_remaining.yaml +++ b/configs/tasks/los_remaining.yaml @@ -4,6 +4,12 @@ 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 diff --git a/configs/tasks/mortality.yaml b/configs/tasks/mortality.yaml index 5cce9e8..44a23d9 100644 --- a/configs/tasks/mortality.yaml +++ b/configs/tasks/mortality.yaml @@ -4,6 +4,13 @@ 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 diff --git a/configs/tasks/mortality_24h.yaml b/configs/tasks/mortality_24h.yaml index 4fc3a55..37d64c8 100644 --- a/configs/tasks/mortality_24h.yaml +++ b/configs/tasks/mortality_24h.yaml @@ -7,6 +7,13 @@ 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 diff --git a/configs/tasks/mortality_hospital.yaml b/configs/tasks/mortality_hospital.yaml index eacb43c..0f3847e 100644 --- a/configs/tasks/mortality_hospital.yaml +++ b/configs/tasks/mortality_hospital.yaml @@ -4,6 +4,13 @@ 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 diff --git a/scripts/training/finetune.py b/scripts/training/finetune.py index 82a34c4..1fa9325 100644 --- a/scripts/training/finetune.py +++ b/scripts/training/finetune.py @@ -43,6 +43,42 @@ ) +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.""" @@ -60,21 +96,33 @@ def main(cfg: DictConfig) -> None: "'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 task_name = cfg.task.get("task_name", "mortality_24h") - validate_data_prerequisites(cfg.data.processed_dir, cfg.dataset, task_names=[task_name]) + 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) diff --git a/scripts/training/supervised.py b/scripts/training/supervised.py index 8e84317..1934208 100644 --- a/scripts/training/supervised.py +++ b/scripts/training/supervised.py @@ -141,7 +141,10 @@ def main(cfg: DictConfig) -> None: # Validate data prerequisites 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] + cfg.data.processed_dir, + cfg.dataset, + task_names=[task_name_for_validation], + task_configs=[cfg.task], ) # Set random seed for reproducibility diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index d249304..e0376e9 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -54,7 +54,12 @@ def main(cfg: DictConfig) -> None: task_type = cfg.task.get("task_type", "binary") # Validate data prerequisites including label freshness - validate_data_prerequisites(cfg.data.processed_dir, cfg.dataset, task_names=[task_name]) + 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, diff --git a/src/slices/data/extractors/base.py b/src/slices/data/extractors/base.py index 6212222..7fa975d 100644 --- a/src/slices/data/extractors/base.py +++ b/src/slices/data/extractors/base.py @@ -133,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) @@ -148,6 +147,19 @@ 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), + } + return manifest + def _validate_observation_window(self, task_config: LabelConfig) -> None: """Validate that task's observation_window_hours matches extraction seq_length_hours. @@ -268,8 +280,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( @@ -291,7 +304,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 = { @@ -301,8 +317,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: @@ -324,7 +342,10 @@ 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 @@ -625,6 +646,15 @@ def _check_existing_extraction(self) -> Optional[Dict[str, pl.DataFrame]]: "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 c245cda..3c71533 100644 --- a/src/slices/data/extractors/ricu.py +++ b/src/slices/data/extractors/ricu.py @@ -29,7 +29,7 @@ from slices.constants import FEATURE_BLOCKLIST, LABEL_HORIZON_HOURS from slices.data.config_schemas import TimeSeriesConceptConfig -from slices.data.labels import LabelBuilder, LabelBuilderFactory, LabelConfig +from slices.data.labels import LabelBuilderFactory, LabelConfig from .base import BaseExtractor, ExtractorConfig @@ -355,6 +355,8 @@ def run(self) -> None: stays = self.extract_stays() progress.update(task, completed=True) + self._validate_stays(stays) + # 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) @@ -383,7 +385,6 @@ def run(self) -> None: console.print("[red]Error: No stays remaining after filtering![/red]") return - self._validate_stays(stays_filtered) self._stays_cache = stays_filtered # ----------------------------------------------------------------- @@ -473,14 +474,7 @@ def run(self) -> None: self._validate_labels(labels, stay_ids) - # Build label manifest for freshness checking at training time - label_manifest = {} - for tc in task_configs: - builder = LabelBuilderFactory.create(tc) - label_manifest[tc.task_name] = { - "builder_version": builder.SEMANTIC_VERSION, - "config_hash": LabelBuilder.config_hash(tc), - } + label_manifest = self._build_label_manifest(task_configs) metadata = { "dataset": self._get_dataset_name(), diff --git a/src/slices/training/config_schemas.py b/src/slices/training/config_schemas.py index 1760f40..b937d38 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,12 +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 diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 73754ee..95b1ea8 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -5,11 +5,13 @@ """ 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 yaml from lightning.pytorch.callbacks import ( EarlyStopping, LearningRateMonitor, @@ -18,6 +20,8 @@ from lightning.pytorch.loggers import WandbLogger from omegaconf import DictConfig, OmegaConf +from slices.data.labels import LabelBuilder, LabelBuilderFactory, LabelConfig + # ============================================================================= # Optimizer / Scheduler # ============================================================================= @@ -474,22 +478,23 @@ 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_names are provided, validates the label - manifest in metadata.yaml to ensure labels were built with the current - builder version and task config. + 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 ``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. """ - import yaml - - from slices.data.labels import LabelBuilder, LabelBuilderFactory, LabelConfig - path = Path(processed_dir) if not path.exists(): @@ -505,8 +510,62 @@ def validate_data_prerequisites( f"Run first: uv run python scripts/preprocessing/prepare_dataset.py dataset={dataset}" ) - # Validate label manifest if task_names are provided - if task_names: + 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: metadata_path = path / "metadata.yaml" if not metadata_path.exists(): raise FileNotFoundError( @@ -527,19 +586,8 @@ def validate_data_prerequisites( f"extract_ricu.py dataset={dataset}" ) - # Load current task configs and compare against manifest - tasks_path = Path(__file__).parent.parent / "data" / "tasks" - 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) - current_config = LabelConfig(**config_dict) + 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) diff --git a/tests/test_base_extractor.py b/tests/test_base_extractor.py index 3392a11..c1c505b 100644 --- a/tests/test_base_extractor.py +++ b/tests/test_base_extractor.py @@ -699,8 +699,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 +712,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 +830,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.""" diff --git a/tests/test_config_schemas.py b/tests/test_config_schemas.py index 0f5c12f..712f667 100644 --- a/tests/test_config_schemas.py +++ b/tests/test_config_schemas.py @@ -77,6 +77,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 +89,21 @@ 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 TestTrainingConfigValidation: """Tests that TrainingConfig catches invalid/misspelled keys.""" diff --git a/tests/test_fixes.py b/tests/test_fixes.py index aed9bc9..1fac6ad 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -195,6 +195,52 @@ def test_validate_data_prerequisites_matching_passes(self, tmp_path): # 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_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 @@ -1497,6 +1543,8 @@ def fake_finetune_module(config, checkpoint_path=None, pretrain_checkpoint_path= 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( { @@ -1546,6 +1594,135 @@ def fake_finetune_module(config, checkpoint_path=None, pretrain_checkpoint_path= assert captured["max_seq_length"] == 24 +class TestFinetuneCheckpointParadigmDetection: + """Regression tests for checkpoint-driven finetune metadata.""" + + 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 sprint matrix coverage.""" diff --git a/tests/test_ricu_extractor.py b/tests/test_ricu_extractor.py index a17a538..f39d9eb 100644 --- a/tests/test_ricu_extractor.py +++ b/tests/test_ricu_extractor.py @@ -815,6 +815,189 @@ 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_fails_closed_on_invalid_los(self, ricu_output_dir: Path, tmp_path: Path) -> None: + """Abort when stay metadata is incomplete instead of emitting a partial dataset.""" + 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=[], + ) + + with pytest.raises(ValueError, match="invalid LOS"): + RicuExtractor(config).run() + + 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: diff --git a/tests/test_training_utils.py b/tests/test_training_utils.py index 31bdf30..1465f25 100644 --- a/tests/test_training_utils.py +++ b/tests/test_training_utils.py @@ -7,6 +7,7 @@ - Supervised script save_encoder_weights: v3 format, loadable by FineTuneModule """ +from dataclasses import fields from pathlib import Path import pytest @@ -14,6 +15,7 @@ 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, @@ -326,6 +328,7 @@ def test_explicit_monitor_override_still_wins(self): 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"] @@ -337,6 +340,18 @@ def test_training_task_configs_cover_package_tasks(self): 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 TestSupervisedCheckpointFormat: """Test that save_encoder_weights from supervised script produces v3 format.""" From 6b9d4eb0d4201a46b65885c27ff525c7b81bbbcb Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 16 Apr 2026 16:03:31 +0100 Subject: [PATCH 080/121] Fix evaluation defaults and downstream audit regressions Include Sprint 11 baselines in the default fairness sweep and train the imputation probe decoder even when loading MAE pretrain checkpoints. Harden downstream training and evaluation by making label statistics task-aware, disabling binary class-weight derivation for regression, fitting trivial supervised baselines on train labels only, and preserving singleton batches during inference. Behavioral implications: - Default fairness runs now cover Sprint 11 baseline experiments. - MAE imputation evaluation now uses a trained probe instead of a random decoder. - Regression runs no longer emit bogus binary stats or leaked trivial baselines during reporting. --- configs/evaluate_imputation.yaml | 4 +- docs/EVAL_FAIRNESS.md | 135 ++++++++++++++++ docs/EXPERIMENT_PLAN.md | 4 +- scripts/eval/evaluate_fairness.py | 2 +- scripts/eval/evaluate_imputation.py | 27 ++-- scripts/launch_thesis_tmux.sh | 2 +- scripts/training/finetune.py | 57 ++++--- scripts/training/supervised.py | 113 ++++++++------ src/slices/data/dataset.py | 177 ++++++++++++++++----- src/slices/eval/imputation.py | 17 ++- src/slices/eval/inference.py | 8 +- src/slices/training/utils.py | 44 ++++++ tests/test_dataset_datamodule.py | 65 ++++++++ tests/test_evaluate_fairness.py | 6 + tests/test_fixes.py | 228 ++++++++++++++++++++++++++++ tests/test_imputation_eval.py | 120 +++++++++++++++ tests/test_inference.py | 26 ++++ 17 files changed, 900 insertions(+), 135 deletions(-) create mode 100644 docs/EVAL_FAIRNESS.md diff --git a/configs/evaluate_imputation.yaml b/configs/evaluate_imputation.yaml index 7658c87..1a9a10f 100644 --- a/configs/evaluate_imputation.yaml +++ b/configs/evaluate_imputation.yaml @@ -13,14 +13,14 @@ defaults: # 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/docs/EVAL_FAIRNESS.md b/docs/EVAL_FAIRNESS.md new file mode 100644 index 0000000..cab2a3d --- /dev/null +++ b/docs/EVAL_FAIRNESS.md @@ -0,0 +1,135 @@ +# Fairness Evaluation + +This document describes the current fairness pipeline in SLICES. + +It intentionally focuses on implementation status and data semantics rather than +freezing a set of thesis numbers that may change as revisions are rerun. 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 thesis fairness corpus: + - Sprints `1`, `2`, `3`, `4`, `5`, `7p`, `10`, `11`, `12`, `13` + - 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. + +## 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 thesis 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` +- `group_sizes` +- `group_sample_sizes` + +For regression tasks it logs: + +- `per_group_mse` +- `per_group_mae` +- `per_group_r2` +- `worst_group_mse` +- `group_sizes` +- `group_sample_sizes` + +### 3. Post-run batch evaluation + +`scripts/eval/evaluate_fairness.py`: + +1. Fetches finished W&B runs matching the requested sprint/paradigm/dataset + 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. + +### 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. + +## Design Decisions + +1. Threshold-dependent classification fairness metrics use a fixed threshold of + `0.5`. +2. Race fairness follows the thesis plan'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 thesis downstream corpus, including + the canonical Sprint `11` classical baselines. + +## How To Refresh Fairness Results + +1. Run the fairness sweep for a specific revision: + +```bash +uv run python scripts/eval/evaluate_fairness.py --revision thesis-v1 +``` + +2. Re-export the result tables: + +```bash +uv run python scripts/export_results.py --revision thesis-v1 +``` + +## 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 should not be used as the frozen numeric appendix for the thesis. diff --git a/docs/EXPERIMENT_PLAN.md b/docs/EXPERIMENT_PLAN.md index 50d8da3..cc6d87d 100644 --- a/docs/EXPERIMENT_PLAN.md +++ b/docs/EXPERIMENT_PLAN.md @@ -411,11 +411,11 @@ and the thesis W&B project summaries. **Runs last**, after all training sprints are complete. Zero additional training — pure evaluation on existing test predictions. -1. Compute fairness metrics on all finetune/supervised test predictions from Sprints 1–5, 7p, 10, 12, and 13 +1. Compute fairness metrics on all downstream test predictions from Sprints 1–5, 7p, 10, 11, 12, and 13 2. Run `scripts/eval/evaluate_fairness.py` with an explicit `--revision` tag for the thesis rerun corpus 3. Generate fairness tables and disparity plots 4. Protected attributes: sex (all datasets), age group (all), race/ethnicity (MIMIC-IV only) -5. Sprint `11` classical baselines are not part of the default standalone fairness sweep because `evaluate_fairness.py` currently targets `finetune` and `supervised` phases only +5. The default standalone fairness sweep includes Sprint `11` classical baselines via `phase:baseline` --- diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index 7260cd4..40e5743 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -52,7 +52,7 @@ # Constants # --------------------------------------------------------------------------- -CORE_SPRINTS = ["1", "2", "3", "4", "5", "7p", "10", "12", "13"] +CORE_SPRINTS = ["1", "2", "3", "4", "5", "7p", "10", "11", "12", "13"] DEFAULT_PHASES = ["finetune", "supervised", "baseline"] DEFAULT_PROTECTED_ATTRIBUTES = ["gender", "age_group", "race"] diff --git a/scripts/eval/evaluate_imputation.py b/scripts/eval/evaluate_imputation.py index 2150c27..b362034 100644 --- a/scripts/eval/evaluate_imputation.py +++ b/scripts/eval/evaluate_imputation.py @@ -9,7 +9,7 @@ checkpoint=outputs/pretrain/encoder.pt \ data.processed_dir=data/processed/mimic-iv - # 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 @@ -101,19 +101,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 diff --git a/scripts/launch_thesis_tmux.sh b/scripts/launch_thesis_tmux.sh index bc1a485..a0e90b8 100755 --- a/scripts/launch_thesis_tmux.sh +++ b/scripts/launch_thesis_tmux.sh @@ -32,7 +32,7 @@ fi mkdir -p "$LOG_DIR" main_sprints=(1 1b 1c 2 3 4 5 6 7 8 10 11) -fairness_sprints=(1 2 3 4 5 10) +fairness_sprints=(1 2 3 4 5 10 11) tag_sprints=(1b 1c 2 5 6 7 8) appendix_sprints=() diff --git a/scripts/training/finetune.py b/scripts/training/finetune.py index 1fa9325..12c61bf 100644 --- a/scripts/training/finetune.py +++ b/scripts/training/finetune.py @@ -36,6 +36,7 @@ from slices.training import FineTuneModule from slices.training.utils import ( report_and_validate_train_label_support, + resolve_balanced_class_weights, run_fairness_evaluation, setup_finetune_callbacks, setup_wandb_logger, @@ -163,23 +164,30 @@ 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}%)" + ) report_and_validate_train_label_support( datamodule=datamodule, task_name=task_name, - task_type=cfg.task.get("task_type", "binary"), + task_type=task_type, dataset=cfg.dataset, seed=cfg.seed, label_fraction=cfg.get("label_fraction", 1.0), @@ -199,25 +207,24 @@ def main(cfg: DictConfig) -> None: # Resolve "balanced" class weights from label distribution if cfg.training.get("class_weight") == "balanced": - if task_name in train_label_stats: - stats = train_label_stats[task_name] - n_pos = stats.get("positive", 0) - n_neg = 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)] - cfg.training.class_weight = [w**0.5 for w in raw] + 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 train-split 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) diff --git a/scripts/training/supervised.py b/scripts/training/supervised.py index 1934208..e0142c2 100644 --- a/scripts/training/supervised.py +++ b/scripts/training/supervised.py @@ -35,6 +35,7 @@ from slices.training import FineTuneModule from slices.training.utils import ( report_and_validate_train_label_support, + resolve_balanced_class_weights, run_fairness_evaluation, save_encoder_checkpoint, setup_finetune_callbacks, @@ -43,48 +44,68 @@ ) +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 @@ -188,23 +209,30 @@ 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}%)" + ) report_and_validate_train_label_support( datamodule=datamodule, task_name=task_name, - task_type=cfg.task.get("task_type", "binary"), + task_type=task_type, dataset=cfg.dataset, seed=cfg.seed, label_fraction=cfg.get("label_fraction", 1.0), @@ -224,25 +252,24 @@ def main(cfg: DictConfig) -> None: # Resolve "balanced" class weights from label distribution if cfg.training.get("class_weight") == "balanced": - if task_name in train_label_stats: - stats = train_label_stats[task_name] - n_pos = stats.get("positive", 0) - n_neg = 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)] - cfg.training.class_weight = [w**0.5 for w in raw] + 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 train-split 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) diff --git a/src/slices/data/dataset.py b/src/slices/data/dataset.py index 2679fdd..d36d88e 100644 --- a/src/slices/data/dataset.py +++ b/src/slices/data/dataset.py @@ -5,6 +5,7 @@ """ import logging +from collections import Counter from pathlib import Path from typing import Any, Dict, List, Optional, Set, Tuple, Union @@ -169,6 +170,7 @@ 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() # Validate task_name if provided if task_name is not None and task_name not in self.task_names: @@ -188,6 +190,57 @@ def __init__( 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. @@ -553,16 +606,18 @@ def get_label_statistics( ) -> Dict[str, Dict[str, Any]]: """Compute label statistics for each task or subset. - 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. + 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 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: @@ -571,42 +626,92 @@ def get_label_statistics( stats: Dict[str, Dict[str, Any]] = {} for task_name in self.task_names: - if task_name in labels_df.columns: - labels = 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 labels_df.columns if c.startswith(f"{task_name}_")] - if 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] = { - "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(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]]: diff --git a/src/slices/eval/imputation.py b/src/slices/eval/imputation.py index 10b5404..fe9b2f5 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. """ @@ -152,10 +152,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(). @@ -168,7 +168,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 @@ -328,10 +328,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 index d39264f..0adb290 100644 --- a/src/slices/eval/inference.py +++ b/src/slices/eval/inference.py @@ -50,16 +50,16 @@ def run_inference( if probs.dim() > 1 and probs.shape[1] == 2: all_preds.append(probs[:, 1].cpu()) else: - all_preds.append(probs.squeeze(-1).cpu()) - all_labels.append(batch["label"].squeeze(-1).cpu()) + 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).squeeze(-1) - labels = torch.cat(all_labels).squeeze(-1) + predictions = torch.cat(all_preds, dim=0) + labels = torch.cat(all_labels, dim=0) return predictions, labels, all_stay_ids diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 95b1ea8..f46fc1d 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -474,6 +474,50 @@ def report_and_validate_train_label_support( } +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_data_prerequisites( processed_dir: str, dataset: str, diff --git a/tests/test_dataset_datamodule.py b/tests/test_dataset_datamodule.py index 55d79c8..bc453a4 100644 --- a/tests/test_dataset_datamodule.py +++ b/tests/test_dataset_datamodule.py @@ -233,6 +233,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( diff --git a/tests/test_evaluate_fairness.py b/tests/test_evaluate_fairness.py index de11cd5..3d7bb03 100644 --- a/tests/test_evaluate_fairness.py +++ b/tests/test_evaluate_fairness.py @@ -10,6 +10,12 @@ def test_default_phases_include_baseline(): assert "baseline" in mod.DEFAULT_PHASES +def test_default_core_sprints_include_sprint11(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + assert "11" in mod.CORE_SPRINTS + + def test_resolve_evaluation_artifact_supports_xgboost(tmp_path): mod = importlib.import_module("scripts.eval.evaluate_fairness") diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 1fac6ad..833d7d8 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -826,6 +826,48 @@ def test_tmux_launcher_passes_revision_to_fairness(self, tmp_path): assert "scripts/eval/evaluate_fairness.py" in runner_text assert "--revision thesis-v2" in runner_text + def test_tmux_launcher_includes_sprint11_in_fairness(self, tmp_path): + """The fairness sweep should include the canonical baseline sprint by default.""" + repo_root = Path(__file__).resolve().parents[1] + temp_repo = tmp_path / "repo" + (temp_repo / "scripts").mkdir(parents=True, exist_ok=True) + shutil.copy2(repo_root / "scripts" / "launch_thesis_tmux.sh", temp_repo / "scripts") + + 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["RUN_EXPORT"] = "0" + + result = subprocess.run( + ["bash", "scripts/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 "--sprint 1 2 3 4 5 10 11" in runner_text + 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] @@ -1593,6 +1635,192 @@ def fake_finetune_module(config, checkpoint_path=None, pretrain_checkpoint_path= 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.""" diff --git a/tests/test_imputation_eval.py b/tests/test_imputation_eval.py index efa127e..db962db 100644 --- a/tests/test_imputation_eval.py +++ b/tests/test_imputation_eval.py @@ -9,9 +9,12 @@ - from_encoder_checkpoint() creates linear decoder """ +import importlib + import pytest import torch import torch.nn as nn +from omegaconf import OmegaConf from slices.eval.imputation import ImputationEvaluator from torch.utils.data import DataLoader @@ -406,3 +409,120 @@ 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): + """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, "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( + { + "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": "outputs/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" diff --git a/tests/test_inference.py b/tests/test_inference.py index 22f549b..5cb3b3e 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -28,3 +28,29 @@ def test_run_inference_squeezes_singleton_regression_outputs(): 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] From 726f8843b036088c107eaf0e3e0b64f5365b30ca Mon Sep 17 00:00:00 2001 From: hill Date: Mon, 20 Apr 2026 10:18:33 -0400 Subject: [PATCH 081/121] Validate cached splits against task cohorts Reject cached patient-level splits when the active task filters the cohort to a different patient set than the prepared full dataset. This keeps seed-specific task runs consistent instead of reusing the prep-time full-cohort split for missing-label tasks such as AKI. Behavioral impact: downstream runs now recompute splits for task-filtered cohorts when the canonical cache was built on a different patient universe. Added a regression test covering the prepared-cache plus missing-label task case. --- src/slices/data/splits.py | 17 ++++++++++-- tests/test_dataset_datamodule.py | 47 ++++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+), 3 deletions(-) diff --git a/src/slices/data/splits.py b/src/slices/data/splits.py index 631681c..971d802 100644 --- a/src/slices/data/splits.py +++ b/src/slices/data/splits.py @@ -27,7 +27,7 @@ def load_cached_splits( """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. + that the patient lists are consistent with the current task-filtered cohort. Args: processed_dir: Path to processed data directory containing splits.yaml. @@ -83,8 +83,19 @@ def load_cached_splits( 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()) + # Validate cached patient lists against the current cohort represented by + # stay_ids. For supervised tasks this may be a label-filtered subset, so + # validating against the full static table would incorrectly reuse + # full-cohort cached splits for a smaller task-specific cohort. + 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) cached_patients = train_patients | val_patients | test_patients if current_patients != cached_patients: diff --git a/tests/test_dataset_datamodule.py b/tests/test_dataset_datamodule.py index bc453a4..431ec78 100644 --- a/tests/test_dataset_datamodule.py +++ b/tests/test_dataset_datamodule.py @@ -10,6 +10,7 @@ import yaml from slices.data.datamodule import ICUDataModule, icu_collate_fn from slices.data.dataset import ICUDataset +from slices.data.splits import load_cached_splits @pytest.fixture @@ -977,6 +978,52 @@ 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_splits_are_invalidated_when_task_filter_removes_patients( + self, mock_extracted_data + ): + """Task-filtered cohorts must not reuse a full-cohort cached split.""" + dm = ICUDataModule( + processed_dir=mock_extracted_data, + task_name="mortality_24h", + 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"], + ) + 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, + ) + + assert cached is None + 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( From a3cda16ac65a4a2b3829885185b105520eb2fe82 Mon Sep 17 00:00:00 2001 From: hill Date: Mon, 20 Apr 2026 10:29:12 -0400 Subject: [PATCH 082/121] Preserve experiment provenance and scheduler safeguards Recover transfer source datasets from output directories, write source_dataset into future transfer finetune configs, and align XGBoost exports with the Protocol-B comparison family so Sprint 10 transfer rows and Sprint 11 baseline comparisons are not misclassified or dropped. Behavioral impact: --parallel now acts as a slot budget where pretrains consume four slots to avoid mixing them with downstream GPU jobs by default. XGBoost early stopping now tracks AUPRC for classification and MAE for regression, matching the benchmark primary metrics. Added regression coverage for exporter recovery, runner command generation, scheduler slot selection, and classical-baseline significance tables. --- scripts/export_results.py | 55 ++++---- scripts/run_experiments.py | 71 +++++++++- scripts/training/xgboost_baseline.py | 12 +- tests/test_fixes.py | 187 +++++++++++++++++++++++++++ 4 files changed, 295 insertions(+), 30 deletions(-) diff --git a/scripts/export_results.py b/scripts/export_results.py index fc23916..9b8a141 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -33,13 +33,14 @@ from pathlib import Path import pandas as pd -import wandb from slices.eval.statistical import ( bonferroni_correction, cohens_d, paired_wilcoxon_signed_rank, ) +import wandb + # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- @@ -357,6 +358,20 @@ def _load_run_data(run) -> tuple[dict, dict, str, str, str, list[str], str, str] _MR_DECODE = {"03": 0.3, "075": 0.75} +def _recover_source_dataset(run_name: str, config: dict) -> str | None: + """Recover transfer provenance from explicit config or historical names.""" + 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]: @@ -423,27 +438,6 @@ def extract_run(run, metric_keys: list[str]) -> dict: lambda: _load_run_data(run) ) - # Derive protocol from freeze_encoder - freeze = _get_nested(config, "training.freeze_encoder") - if freeze is True: - protocol = "A" - elif freeze is False: - protocol = "B" - else: - protocol = None - - model_size = _infer_model_size(config) - - # Detect source_dataset for transfer runs from run name or config - source_dataset = config.get("source_dataset", None) - if source_dataset is None: - if "_from_" in run_name: - parts = run_name.split("_from_") - if len(parts) >= 2: - source_part = parts[1].split("_")[0] - if source_part in ("miiv", "eicu", "combined"): - source_dataset = source_part - # Detect phase from tags phase = None for tag in tags: @@ -457,6 +451,23 @@ def extract_run(run, metric_keys: list[str]) -> dict: else: phase = "finetune" + # Derive protocol from freeze_encoder. XGBoost runs do not use this field, + # but they belong with the full-training / Protocol-B comparison family. + freeze = _get_nested(config, "training.freeze_encoder") + if freeze is True: + protocol = "A" + elif freeze is False: + protocol = "B" + elif config.get("paradigm") == "xgboost" or phase == "baseline": + protocol = "B" + else: + protocol = None + + model_size = _infer_model_size(config) + + # Detect source_dataset for transfer runs from config, output_dir, or run name. + source_dataset = _recover_source_dataset(run_name, config) + # Extract metrics from summary metrics = {} for key in metric_keys: diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index 61ddad4..6d4b9cb 100644 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -216,6 +216,8 @@ def _finetune_cmd(self, runs_by_id: dict[str, Run]) -> list[str]: cmd.append(f"label_fraction={self.label_fraction}") for k, v in self.extra_overrides.items(): cmd.append(f"{k}={v}") + if self.source_dataset is not None: + cmd.append(f"+source_dataset={self.source_dataset}") # Propagate upstream pretrain metadata so it lands in W&B config if self.upstream_pretrain_lr is not None: cmd.append(f"+upstream_pretrain_lr={self.upstream_pretrain_lr}") @@ -1029,10 +1031,57 @@ def recover_stale_running(state: dict): print(f" Recovered stale run: {run_id}") +# Scheduler slot weights. Pretraining occupies the full default budget of 4 +# slots so `--parallel 4` does not mix heavyweight pretrains with downstream +# GPU jobs on a single device. +RUN_SLOT_COSTS = { + "pretrain": 4, + "finetune": 1, + "supervised": 1, + "gru_d": 1, + "xgboost": 1, +} + + +def _slot_cost(run: Run, slot_budget: int) -> int: + """Return the scheduler slot cost for a run under the active budget.""" + 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]: + """Choose a launch batch that fits the available scheduler slots. + + Heavier runs are prioritized first so pretrains claim the budget before + lighter downstream jobs, avoiding accidental co-scheduling on one GPU. + """ + active_slots = sum(_slot_cost(runs_by_id[rid], slot_budget) for rid 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): + if parallel < 1: + raise ValueError(f"--parallel must be >= 1, got {parallel}") + runs_by_id = {r.id: r for r in runs} active: dict[str, subprocess.Popen] = {} # run_id -> Popen shutting_down = False @@ -1072,7 +1121,7 @@ def handle_signal(signum, frame): _print_dry_run(runs, runs_by_id) return - print(f"Scheduler started (parallel={parallel})") + print(f"Scheduler started (slot_budget={parallel})") print(f"State file: {STATE_FILE}") print() @@ -1124,9 +1173,9 @@ def handle_signal(signum, frame): if deps_ok: ready.append(r) - # 4. Launch runs up to parallel limit - slots = parallel - len(active) - for r in ready[:slots]: + # 4. Launch runs that fit the remaining scheduler slots + launch_batch = _select_ready_runs(ready, set(active), runs_by_id, parallel) + for r in launch_batch: cmd = r.build_command(runs_by_id) log_file = LOG_DIR / f"{r.id}.log" now = datetime.now(timezone.utc).isoformat() @@ -1575,7 +1624,12 @@ def main(): # run p_run = sub.add_parser("run", help="Run experiments for given sprints") p_run.add_argument("--sprint", nargs="+", required=True, help="Sprint(s) to run (e.g. 1 1b 2)") - p_run.add_argument("--parallel", type=int, default=4, help="Max parallel jobs (default: 4)") + 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 (e.g. v2)") p_run.add_argument("--reason", type=str, default=None, help="Reason for rerun") @@ -1600,7 +1654,12 @@ def main(): 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 skipped runs") - p_retry.add_argument("--parallel", type=int, default=4, help="Max parallel jobs (default: 4)") + 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("--sprint", nargs="+", default=None, help="Scope retry to sprint(s)") 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") diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index e0376e9..8586326 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -28,6 +28,11 @@ from xgboost import XGBClassifier, XGBRegressor +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" + + @hydra.main(version_base=None, config_path="../../configs", config_name="xgboost") def main(cfg: DictConfig) -> None: """Train XGBoost baseline.""" @@ -123,7 +128,10 @@ def main(cfg: DictConfig) -> None: } if task_type == "regression": - model = XGBRegressor(**common_params) + model = XGBRegressor( + **common_params, + eval_metric=_xgboost_eval_metric(task_type), + ) else: # Auto-compute scale_pos_weight if requested scale_pos_weight = xgb_cfg.get("scale_pos_weight", None) @@ -135,7 +143,7 @@ def main(cfg: DictConfig) -> None: model = XGBClassifier( **common_params, scale_pos_weight=scale_pos_weight, - eval_metric="logloss", + eval_metric=_xgboost_eval_metric(task_type), ) model.fit( diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 833d7d8..bd172c9 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -1294,6 +1294,49 @@ def test_extract_run_reclassifies_transfer_and_sprint13(self): ) assert extract_run(ts2vec, [])["experiment_type"] == "temporal_contrastive" + def test_extract_run_recovers_transfer_from_output_dir(self): + """Historical transfer runs should recover source_dataset from output_dir.""" + from scripts.export_results import extract_run + + transfer = DummyWandbRun( + run_id="transfer_output_dir", + config={ + "sprint": "10", + "dataset": "eicu", + "paradigm": "mae", + "seed": 789, + "output_dir": "outputs/sprint10/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_type"] == "transfer" + + def test_extract_run_assigns_protocol_b_to_xgboost_baselines(self): + """XGBoost baselines must align with the Protocol-B comparison family.""" + from scripts.export_results import extract_run + + run = DummyWandbRun( + run_id="xgb", + config={ + "sprint": "11", + "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"] == "B" + 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 @@ -1440,6 +1483,35 @@ def test_apply_wandb_target_injects_project_and_entity(self): assert result[0].extra_overrides["logging.wandb_project"] == "slices-thesis" assert result[0].extra_overrides["logging.wandb_entity"] == "hannes-ill" + def test_transfer_finetune_command_propagates_source_dataset(self): + from scripts.run_experiments import Run + + pretrain = Run( + id="pretrain_mae_miiv_seed42", + sprint="7", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/sprint7/pretrain_mae_miiv_seed42", + ) + finetune = Run( + id="finetune_mae_eicu_seed42", + sprint="7", + run_type="finetune", + paradigm="mae", + dataset="eicu", + seed=42, + output_dir="outputs/sprint7/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 + class TestExperimentRunnerRetry: """Tests for retry scoping and dependency preservation.""" @@ -1516,6 +1588,121 @@ def test_cmd_retry_respects_sprint_filter_but_keeps_dependencies(self, monkeypat assert dependency.id in scheduled_ids assert other.id not in scheduled_ids + def test_select_ready_runs_prioritizes_pretrains_with_slot_budget(self): + import scripts.run_experiments as runner + + pretrain = runner.Run( + id="pretrain", + sprint="10", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/sprint10/pretrain_mae_miiv_seed42", + ) + finetune = runner.Run( + id="finetune", + sprint="10", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/sprint10/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.run_experiments as runner + + pretrain = runner.Run( + id="pretrain", + sprint="10", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/sprint10/pretrain_mae_miiv_seed42", + ) + finetune = runner.Run( + id="finetune", + sprint="10", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/sprint10/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 == [] + + +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", "B", 0.08), + ("gru_d", "baseline", "B", 0.03), + ("supervised", "supervised", "B", 0.0), + ]: + for seed in [42, 123]: + for task, base in [("mortality_24h", 0.30), ("aki_kdigo", 0.25)]: + rows.append( + { + "experiment_type": "classical_baselines", + "sprint": "11", + "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 + + +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" + class TestTrainingScriptClassWeighting: """Regression tests for entrypoint-level class weight resolution.""" From 6717f292ce35127e2840d6b436cbac1e1184b5a8 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 21 Apr 2026 09:33:37 -0400 Subject: [PATCH 083/121] Fix SSL timestep reassembly for sparse batches Add a padding-safe visible-token scatter helper and use it when rebuilding full timestep grids for MAE, JEPA, temporal contrastive, and TS2Vec objectives. Pass MAE's valid timestep mask into visible-token extraction so fully unobserved hours are excluded consistently with the other SSL objectives. This changes SSL pretraining behavior for sparse mixed-density batches: padded visible tokens no longer overwrite masked timesteps, and MAE masking metrics now use observed timesteps as the denominator. Existing SSL checkpoints produced before this fix should not be mixed with new canonical runs. Regression tests cover padded-token scatter and MAE valid-timestep propagation. --- src/slices/models/pretraining/contrastive.py | 17 +---- src/slices/models/pretraining/jepa.py | 32 +++++--- src/slices/models/pretraining/mae.py | 53 +++++++++----- src/slices/models/pretraining/masking.py | 51 +++++++++++++ src/slices/models/pretraining/ts2vec.py | 12 +-- tests/test_contrastive_objective.py | 23 ++++++ tests/test_jepa_objective.py | 40 ++++++++++ tests/test_mae_objective.py | 77 +++++++++++++++++++- tests/test_ts2vec_objective.py | 22 ++++++ 9 files changed, 272 insertions(+), 55 deletions(-) diff --git a/src/slices/models/pretraining/contrastive.py b/src/slices/models/pretraining/contrastive.py index 96e41c1..42e5b70 100644 --- a/src/slices/models/pretraining/contrastive.py +++ b/src/slices/models/pretraining/contrastive.py @@ -42,6 +42,7 @@ create_complementary_timestep_masks, create_timestep_mask, extract_visible_timesteps, + scatter_visible_timesteps, ) @@ -266,9 +267,8 @@ def _scatter_to_full( ) -> torch.Tensor: """Scatter visible encoded tokens back to full (B, T, d) tensor. - Uses the same argsort-scatter pattern as the MAE decoder and JEPA - predictor to place encoded visible tokens at their original temporal - positions, with zeros at masked positions. + Places encoded visible tokens at their original temporal positions, + with zeros at masked positions. Args: encoded: (B, n_vis, d_enc) encoded visible tokens. @@ -278,16 +278,7 @@ def _scatter_to_full( Returns: (B, T, d_enc) with encoded tokens at visible positions, zeros elsewhere. """ - B, n_vis, d_enc = encoded.shape - device = encoded.device - - full = torch.zeros(B, n_timesteps, d_enc, device=device, dtype=encoded.dtype) - - vis_indices = ssl_mask.float().argsort(dim=1, descending=True, stable=True) - scatter_idx = vis_indices[:, :n_vis].unsqueeze(-1).expand(-1, -1, d_enc) - full.scatter_(1, scatter_idx, encoded) - - return full + return scatter_visible_timesteps(encoded, ssl_mask, n_timesteps) def _temporal_nt_xent_loss( self, diff --git a/src/slices/models/pretraining/jepa.py b/src/slices/models/pretraining/jepa.py index 9655659..0f68da0 100644 --- a/src/slices/models/pretraining/jepa.py +++ b/src/slices/models/pretraining/jepa.py @@ -26,7 +26,12 @@ from slices.models.common import build_sinusoidal_pe from .base import BaseSSLObjective, SSLConfig -from .masking import create_block_timestep_mask, create_timestep_mask, extract_visible_timesteps +from .masking import ( + create_block_timestep_mask, + create_timestep_mask, + extract_visible_timesteps, + scatter_visible_timesteps, +) @dataclass @@ -118,26 +123,29 @@ def forward( Returns: (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, n_timesteps, 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 - 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.to(full_tokens.dtype)) + full_tokens = scatter_visible_timesteps( + vis_proj, + visible_mask, + n_timesteps, + fill_value=self.mask_token, + ) # Add time PE timestep_idx = token_info["timestep_idx"] full_tokens = full_tokens + self.time_pe[timestep_idx] full_tokens = self.embed_dropout(full_tokens) - valid_timestep_mask = token_info.get("valid_timestep_mask") predictor_padding_mask = None if valid_timestep_mask is not None: predictor_padding_mask = ~valid_timestep_mask.to( diff --git a/src/slices/models/pretraining/mae.py b/src/slices/models/pretraining/mae.py index fe68335..93739fd 100644 --- a/src/slices/models/pretraining/mae.py +++ b/src/slices/models/pretraining/mae.py @@ -20,7 +20,7 @@ from slices.models.common import build_sinusoidal_pe from .base import BaseSSLObjective, SSLConfig -from .masking import create_timestep_mask, extract_visible_timesteps +from .masking import create_timestep_mask, extract_visible_timesteps, scatter_visible_timesteps @dataclass @@ -113,28 +113,31 @@ def forward( Returns: (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: mask tokens everywhere, then scatter visible - full_tokens = self.mask_token.expand(B, n_timesteps, d_dec).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 their original positions - vis_indices = ssl_mask.float().argsort(dim=1, descending=True, stable=True) - 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) - full_tokens.scatter_(1, scatter_idx_expanded, vis_proj.to(full_tokens.dtype)) + 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) - # Run decoder transformer (no padding mask needed, all T positions valid) + # Run decoder transformer over the full reconstruction grid. decoded = self.decoder(full_tokens) # Predict D features per timestep @@ -221,7 +224,12 @@ def forward( ) # 3. Extract visible tokens - visible_tokens, vis_padding = extract_visible_timesteps(tokens, ssl_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) @@ -235,7 +243,13 @@ def forward( ) # (B, T, D) # 6. Compute loss on observed features at masked timesteps - loss, metrics = self._compute_loss(predictions, x, ssl_mask, obs_mask) + loss, metrics = self._compute_loss( + predictions, + x, + ssl_mask, + obs_mask, + valid_timestep_mask, + ) return loss, metrics @@ -245,6 +259,7 @@ def _compute_loss( true_values: torch.Tensor, ssl_mask: torch.Tensor, obs_mask: torch.Tensor, + valid_timestep_mask: torch.Tensor, ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: """Compute MSE loss on observed features at masked timesteps. @@ -253,6 +268,7 @@ def _compute_loss( 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) @@ -266,8 +282,9 @@ def _compute_loss( with torch.no_grad(): B, T, D = true_values.shape n_timesteps = T - n_masked_timesteps = (~ssl_mask).sum().item() - n_visible_timesteps = ssl_mask.sum().item() + 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 @@ -284,7 +301,7 @@ 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_timesteps / max(B * n_timesteps, 1), + "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, diff --git a/src/slices/models/pretraining/masking.py b/src/slices/models/pretraining/masking.py index c57757a..86817d0 100644 --- a/src/slices/models/pretraining/masking.py +++ b/src/slices/models/pretraining/masking.py @@ -307,3 +307,54 @@ def extract_visible_timesteps( 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/ts2vec.py b/src/slices/models/pretraining/ts2vec.py index df7b991..ea661bd 100644 --- a/src/slices/models/pretraining/ts2vec.py +++ b/src/slices/models/pretraining/ts2vec.py @@ -37,7 +37,7 @@ import torch.nn.functional as F from .base import BaseSSLObjective, SSLConfig -from .masking import create_timestep_mask, extract_visible_timesteps +from .masking import create_timestep_mask, extract_visible_timesteps, scatter_visible_timesteps @dataclass @@ -276,15 +276,7 @@ def _scatter_to_full( Returns: (B, T, d_enc) with encoded tokens at visible positions, zeros elsewhere. """ - B, n_vis, d_enc = encoded.shape - device = encoded.device - - full = torch.zeros(B, n_timesteps, d_enc, device=device, dtype=encoded.dtype) - vis_indices = ssl_mask.float().argsort(dim=1, descending=True, stable=True) - scatter_idx = vis_indices[:, :n_vis].unsqueeze(-1).expand(-1, -1, d_enc) - full.scatter_(1, scatter_idx, encoded) - - return full + return scatter_visible_timesteps(encoded, ssl_mask, n_timesteps) @staticmethod def _masked_max_pool1d( diff --git a/tests/test_contrastive_objective.py b/tests/test_contrastive_objective.py index f3987b1..f1c7cc4 100644 --- a/tests/test_contrastive_objective.py +++ b/tests/test_contrastive_objective.py @@ -531,6 +531,29 @@ def test_scatter_to_full_shape_and_zeros(self, encoder, temporal_config): # 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 diff --git a/tests/test_jepa_objective.py b/tests/test_jepa_objective.py index 1597981..ff06be5 100644 --- a/tests/test_jepa_objective.py +++ b/tests/test_jepa_objective.py @@ -81,6 +81,46 @@ def test_predictor_output_in_encoder_space(self): 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 diff --git a/tests/test_mae_objective.py b/tests/test_mae_objective.py index 5c062e5..f32b458 100644 --- a/tests/test_mae_objective.py +++ b/tests/test_mae_objective.py @@ -138,6 +138,42 @@ 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])) + # ============================================================================= # MAE Objective tests @@ -297,6 +333,23 @@ def test_mae_respects_obs_mask(self, encoder, mae_config): 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) @@ -313,13 +366,32 @@ def fake_create_timestep_mask( captured["n_timesteps"] = n_timesteps captured["mask_ratio"] = mask_ratio captured["device"] = device - captured["valid_timestep_mask"] = valid_timestep_mask.clone() + 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( @@ -343,7 +415,8 @@ def fake_create_timestep_mask( mae(x, obs_mask) expected = obs_mask.any(dim=-1) - assert torch.equal(captured["valid_timestep_mask"], expected) + 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) diff --git a/tests/test_ts2vec_objective.py b/tests/test_ts2vec_objective.py index e92027f..91c394a 100644 --- a/tests/test_ts2vec_objective.py +++ b/tests/test_ts2vec_objective.py @@ -45,3 +45,25 @@ def test_empty_timesteps_are_excluded_from_ts2vec_masks(encoder): 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])) From 8e7440392700f4789ffce8627dc3748e00e4de00 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 21 Apr 2026 12:38:41 -0400 Subject: [PATCH 084/121] Prepare repository for public benchmark release Move internal experiment automation and stale planning docs under internal paths, update public docs/examples/configs to current benchmark conventions, and restore debug snapshot export support. This keeps auto-shutdown available under scripts/internal, treats docs/internal as private material, and updates dependency metadata for cleaner benchmark installation. Validation: ruff, black, mypy, pytest. --- .gitignore | 3 +- README.md | 28 +- configs/debug.yaml | 4 +- configs/evaluate_imputation.yaml | 4 +- configs/finetune.yaml | 4 +- configs/gru_d.yaml | 2 +- configs/model/transformer_large.yaml | 2 +- configs/model/transformer_medium.yaml | 2 +- configs/pretrain.yaml | 6 +- configs/supervised.yaml | 4 +- configs/xgboost.yaml | 2 +- docs/EVAL_FAIRNESS.md | 27 +- docs/{ => internal}/EXPERIMENT_PLAN.md | 12 +- examples/mae_pretraining_example.py | 6 +- examples/transformer_with_dataset_example.py | 28 +- pyproject.toml | 19 +- scripts/analyze_labels.py | 8 +- scripts/debug/explore_missingness.py | 4 +- scripts/debug/export_snapshots.py | 59 + scripts/debug/inspect_embeddings.py | 4 +- scripts/eval/evaluate_fairness.py | 23 +- scripts/eval/evaluate_imputation.py | 4 +- scripts/export_results.py | 43 +- scripts/{ => internal}/auto_shutdown.sh | 4 +- scripts/{ => internal}/launch_thesis_tmux.sh | 12 +- scripts/{ => internal}/run_experiments.py | 32 +- scripts/measure_flops.py | 13 +- scripts/sanity_checks/sc_mae_pretraining.py | 24 +- .../sanity_checks/sc_supervised_learning.py | 12 +- src/slices/data/datamodule.py | 2 +- src/slices/data/dataset.py | 4 +- src/slices/data/sliding_window.py | 2 +- src/slices/debug/snapshots.py | 27 +- src/slices/eval/README.md | 6 +- src/slices/eval/fairness_evaluator.py | 4 +- src/slices/eval/statistical.py | 2 +- src/slices/models/encoders/README.md | 4 +- src/slices/models/encoders/wrapper.py | 4 +- src/slices/models/pretraining/README.md | 2 +- src/slices/training/utils.py | 2 +- tests/test_fixes.py | 50 +- tests/test_snapshots.py | 41 + uv.lock | 1714 +---------------- 43 files changed, 434 insertions(+), 1825 deletions(-) rename docs/{ => internal}/EXPERIMENT_PLAN.md (97%) create mode 100644 scripts/debug/export_snapshots.py rename scripts/{ => internal}/auto_shutdown.sh (84%) rename scripts/{ => internal}/launch_thesis_tmux.sh (86%) rename scripts/{ => internal}/run_experiments.py (98%) create mode 100644 tests/test_snapshots.py diff --git a/.gitignore b/.gitignore index dbad839..90778d6 100644 --- a/.gitignore +++ b/.gitignore @@ -210,7 +210,7 @@ __marimo__/ # Notes notes/ -docs/ +docs/internal/ # SLICES-specific /data/ @@ -219,6 +219,7 @@ outputs/ !/results/.gitkeep wandb/ .tensor_cache/ +logs/ # PyTorch Lightning lightning_logs/ diff --git a/README.md b/README.md index 8d65a8d..349817a 100644 --- a/README.md +++ b/README.md @@ -21,11 +21,9 @@ 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 formal thesis corpus now also includes a contextual baseline family plus two targeted extensions: - -- `Sprint 11`: canonical GRU-D and XGBoost baselines, reported as contextual ICU references alongside the controlled benchmark -- `Sprint 7p`: a focused capacity study that scales MAE and supervised encoders on MIIV `mortality_24h` -- `Sprint 13`: a TS2Vec temporal-contrastive extension that tests whether a stronger contrastive-family instantiation changes the conclusion +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 @@ -43,7 +41,8 @@ Fair comparison of SSL objectives for clinical time series is currently impossib ## SSL Paradigms -The controlled thesis objectives share the same timestep-level obs-aware Transformer encoder and differ only in the SSL objective and masking logic: +The controlled SSL objectives share the same timestep-level obs-aware Transformer +encoder and differ only in the SSL objective and masking logic: | Objective | Predicts | Target | Loss | |---|---|---|---| @@ -51,7 +50,9 @@ The controlled thesis objectives share the same timestep-level obs-aware Transfo | **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 | -**TS2Vec** is included as a formal temporal-contrastive thesis extension (Sprint `13`). **SMART** (NeurIPS 2024) remains an appendix-only external reference because it swaps in its own MART encoder and element-wise masking, so it is not part of the controlled thesis comparison. +**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 @@ -105,7 +106,8 @@ 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) @@ -117,10 +119,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 -# TS2Vec (formal thesis extension for the contrastive family) +# TS2Vec (temporal contrastive extension) uv run python scripts/training/pretrain.py dataset=miiv ssl=ts2vec -# SMART (sanity check / extensibility demo — not part of thesis experiments) +# SMART (external reference with a different encoder/tokenization contract) uv run python scripts/training/pretrain.py dataset=miiv ssl=smart model=smart ``` @@ -136,12 +138,10 @@ uv run python scripts/training/finetune.py dataset=miiv checkpoint=outputs/.../e uv run python scripts/training/supervised.py dataset=miiv ``` -### 5. Thesis Reruns +### 5. Validate ```bash -WANDB_PROJECT=slices-thesis \ -WANDB_ENTITY=your-wandb-entity \ -bash scripts/launch_thesis_tmux.sh +uv run pytest tests/ -q ``` ## Project Structure diff --git a/configs/debug.yaml b/configs/debug.yaml index 4e0f9c8..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 diff --git a/configs/evaluate_imputation.yaml b/configs/evaluate_imputation.yaml index 1a9a10f..9c19b5e 100644 --- a/configs/evaluate_imputation.yaml +++ b/configs/evaluate_imputation.yaml @@ -5,12 +5,14 @@ # 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; evaluation still trains a probe decoder diff --git a/configs/finetune.yaml b/configs/finetune.yaml index 440d4e2..050f30f 100644 --- a/configs/finetune.yaml +++ b/configs/finetune.yaml @@ -35,10 +35,10 @@ 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) +# Experiment group (stored as a sprint tag for historical W&B compatibility) sprint: null -# Revision tracking (for rerunning sprints with changed settings) +# Revision tracking (for rerunning experiment groups with changed settings) revision: null rerun_reason: null diff --git a/configs/gru_d.yaml b/configs/gru_d.yaml index 41e63ff..2c265aa 100644 --- a/configs/gru_d.yaml +++ b/configs/gru_d.yaml @@ -25,7 +25,7 @@ dataset: miiv # Paradigm paradigm: gru_d -# Experiment sprint +# Experiment group (stored as a sprint tag for historical W&B compatibility) sprint: null revision: null rerun_reason: null diff --git a/configs/model/transformer_large.yaml b/configs/model/transformer_large.yaml index 7ba5ca4..b1fa545 100644 --- a/configs/model/transformer_large.yaml +++ b/configs/model/transformer_large.yaml @@ -1,7 +1,7 @@ # @package encoder # Transformer Encoder — Large (256-dim, 4-layer) # -# Scaled-up variant for capacity ablation (Sprint 7 pilot). +# Scaled-up variant for the capacity ablation pilot. # ~3.2M params vs ~112K for the base transformer. # All other settings identical to transformer.yaml. diff --git a/configs/model/transformer_medium.yaml b/configs/model/transformer_medium.yaml index d5f2551..d4aa114 100644 --- a/configs/model/transformer_medium.yaml +++ b/configs/model/transformer_medium.yaml @@ -1,7 +1,7 @@ # @package encoder # Transformer Encoder — Medium (128-dim, 4-layer) # -# Scaled-up variant for capacity ablation (Sprint 7 pilot). +# Scaled-up variant for the capacity ablation pilot. # ~816K params vs ~112K for the base transformer. # All other settings identical to transformer.yaml. diff --git a/configs/pretrain.yaml b/configs/pretrain.yaml index 1cd3112..ae137c4 100644 --- a/configs/pretrain.yaml +++ b/configs/pretrain.yaml @@ -22,10 +22,10 @@ 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) +# Experiment group (stored as a sprint tag for historical W&B compatibility) sprint: null -# Revision tracking (for rerunning sprints with changed settings) +# Revision tracking (for rerunning experiment groups with changed settings) revision: null rerun_reason: null @@ -53,7 +53,7 @@ data: 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: 100 batch_size: 256 diff --git a/configs/supervised.yaml b/configs/supervised.yaml index 8265ae1..367cfb8 100644 --- a/configs/supervised.yaml +++ b/configs/supervised.yaml @@ -33,10 +33,10 @@ dataset: miiv # miiv | eicu | combined # Paradigm — always "supervised" for this config paradigm: supervised -# Experiment sprint (for W&B filtering, e.g. sprint=1) +# Experiment group (stored as a sprint tag for historical W&B compatibility) sprint: null -# Revision tracking (for rerunning sprints with changed settings) +# Revision tracking (for rerunning experiment groups with changed settings) revision: null rerun_reason: null diff --git a/configs/xgboost.yaml b/configs/xgboost.yaml index 7d989e2..2a47a45 100644 --- a/configs/xgboost.yaml +++ b/configs/xgboost.yaml @@ -22,7 +22,7 @@ dataset: miiv # Paradigm paradigm: xgboost -# Experiment sprint +# Experiment group (stored as a sprint tag for historical W&B compatibility) sprint: null revision: null rerun_reason: null diff --git a/docs/EVAL_FAIRNESS.md b/docs/EVAL_FAIRNESS.md index cab2a3d..8923b97 100644 --- a/docs/EVAL_FAIRNESS.md +++ b/docs/EVAL_FAIRNESS.md @@ -2,18 +2,18 @@ This document describes the current fairness pipeline in SLICES. -It intentionally focuses on implementation status and data semantics rather than -freezing a set of thesis numbers that may change as revisions are rerun. The -canonical outputs are the `fairness/*` keys written to W&B summaries and the -parquet exports produced by `scripts/export_results.py`. +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 thesis fairness corpus: - - Sprints `1`, `2`, `3`, `4`, `5`, `7p`, `10`, `11`, `12`, `13` +- Default benchmark fairness corpus: + - Internal experiment groups `1`, `2`, `3`, `4`, `5`, `7p`, `10`, `11`, + `12`, `13` - 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. @@ -28,7 +28,7 @@ parquet exports produced by `scripts/export_results.py`. Important details: -- eICU race is excluded from the thesis fairness sweep. +- 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 @@ -103,25 +103,25 @@ summary metrics. 1. Threshold-dependent classification fairness metrics use a fixed threshold of `0.5`. -2. Race fairness follows the thesis plan's canonical five-bin schema rather than +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 thesis downstream corpus, including - the canonical Sprint `11` classical baselines. +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: ```bash -uv run python scripts/eval/evaluate_fairness.py --revision thesis-v1 +uv run python scripts/eval/evaluate_fairness.py --revision benchmark-v1 ``` 2. Re-export the result tables: ```bash -uv run python scripts/export_results.py --revision thesis-v1 +uv run python scripts/export_results.py --revision benchmark-v1 ``` ## Source Of Truth @@ -132,4 +132,5 @@ Use these in order: 2. `results/*.parquet` generated by `scripts/export_results.py` 3. This document for implementation semantics only -This file should not be used as the frozen numeric appendix for the thesis. +This file documents implementation semantics; use exported result tables for +numeric reporting. diff --git a/docs/EXPERIMENT_PLAN.md b/docs/internal/EXPERIMENT_PLAN.md similarity index 97% rename from docs/EXPERIMENT_PLAN.md rename to docs/internal/EXPERIMENT_PLAN.md index cc6d87d..1628cca 100644 --- a/docs/EXPERIMENT_PLAN.md +++ b/docs/internal/EXPERIMENT_PLAN.md @@ -353,7 +353,7 @@ and the thesis W&B project summaries. **Breakdown**: 48 pretrain (~30 min each, the bottleneck) + 510 finetune/probe (~3–5 min each) + 72 supervised (~3–5 min each). -**Execution**: `uv run python scripts/run_experiments.py run --sprint 10 --parallel 4` (limited by GPU-bound pretraining) +**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 10 --parallel 4` (limited by GPU-bound pretraining) ### Upcoming: Sprint 7p — Focused Capacity Study (100 runs) @@ -363,7 +363,7 @@ and the thesis W&B project summaries. **Comparison set**: Default-size baselines are inherited from Sprint `6` (seeds 42/123/456) and Sprint `10` (seeds 789/1011), so Sprint `7p` only launches the larger-capacity runs. -**Execution**: `uv run python scripts/run_experiments.py run --sprint 7p --parallel 4` +**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 7p --parallel 4` ### Upcoming: Sprint 11 — Classical Baselines (360 runs) @@ -383,7 +383,7 @@ and the thesis W&B project summaries. 3. XGBoost cannot transfer across datasets — highlights SSL's cross-dataset transfer capability (Sprint 7) 4. XGBoost/GRU-D on label efficiency curves — direct comparison with SSL learning curves -**Execution**: `uv run python scripts/run_experiments.py run --sprint 11 --parallel 12` (all runs independent, no pretraining dependency) +**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 11 --parallel 12` (all runs independent, no pretraining dependency) ### Upcoming: Sprint 13 — TS2Vec Temporal Contrastive Extension (135 runs) @@ -395,7 +395,7 @@ and the thesis W&B project summaries. **Interpretation**: Sprint `13` is a formal thesis extension, not a replacement for the controlled MAE/JEPA/Contrastive triangle. It tests whether a better-instantiated contrastive family changes the conclusion. -**Execution**: `uv run python scripts/run_experiments.py run --sprint 13 --parallel 4` +**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 13 --parallel 4` ### Upcoming: Sprint 12 — SMART External SSL Reference (135 runs) @@ -405,7 +405,7 @@ and the thesis W&B project summaries. **Matrix**: 1 paradigm × 3 datasets × 4 tasks × 2 protocols × 5 seeds = 120 finetune + 15 pretrain = **135 runs** -**Execution**: `uv run python scripts/run_experiments.py run --sprint 12 --parallel 4` +**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 12 --parallel 4` ### Upcoming: Sprint 9 — Fairness Analysis (0 extra runs) @@ -485,7 +485,7 @@ Later sprints reuse runs from earlier sprints as comparison baselines. To enable | **12** | — | Self-contained (SMART pretrains + finetune, 5 seeds) | | **9** | All | Evaluates test predictions from all training sprints (runs last) | -**Usage**: `uv run python scripts/run_experiments.py tag --sprint N` (idempotent). +**Usage**: `uv run python scripts/internal/run_experiments.py tag --sprint N` (idempotent). --- diff --git a/examples/mae_pretraining_example.py b/examples/mae_pretraining_example.py index 2acd205..214ad90 100644 --- a/examples/mae_pretraining_example.py +++ b/examples/mae_pretraining_example.py @@ -142,7 +142,7 @@ def main(): """Main example demonstrating MAE pretraining.""" # Configuration - data_dir = Path("data/processed/mimic-iv-demo") + data_dir = Path("data/processed/miiv") device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("=" * 80) @@ -152,7 +152,7 @@ def main(): # 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.") + print("Run ./scripts/setup_and_extract.sh miiv first to extract data.") return # Load dataset @@ -397,7 +397,7 @@ def main(): 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(" - Add to scripts/training/pretrain.py for full training") print(" - Evaluate learned representations on downstream tasks") print("=" * 80) diff --git a/examples/transformer_with_dataset_example.py b/examples/transformer_with_dataset_example.py index 2dfd77c..8bcda9f 100644 --- a/examples/transformer_with_dataset_example.py +++ b/examples/transformer_with_dataset_example.py @@ -12,8 +12,8 @@ def check_data_availability() -> bool: - """Check if extracted MIMIC-IV demo data is available.""" - data_dir = Path("data/processed/mimic-iv-demo") + """Check if processed MIIV data is available.""" + data_dir = Path("data/processed/miiv") return data_dir.exists() and (data_dir / "timeseries.parquet").exists() @@ -22,8 +22,8 @@ def basic_integration(): 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("⚠️ Processed MIIV data not found.") + print(" Run: ./scripts/setup_and_extract.sh miiv") print(" This example requires extracted ICU data.\n") return False @@ -31,7 +31,7 @@ def basic_integration(): # Load dataset dataset = ICUDataset( - "data/processed/mimic-iv-demo", + "data/processed/miiv", task_name="mortality_24h", normalize=True, ) @@ -81,7 +81,7 @@ def batch_processing(): print("=== Batch Processing ===") if not check_data_availability(): - print("⚠️ MIMIC-IV demo data not found. Skipping.\n") + print("⚠️ Processed MIIV data not found. Skipping.\n") return False from slices.data import ICUDataset @@ -89,7 +89,7 @@ def batch_processing(): # Load dataset dataset = ICUDataset( - "data/processed/mimic-iv-demo", + "data/processed/miiv", task_name="mortality_24h", normalize=True, ) @@ -136,14 +136,14 @@ def with_datamodule(): print("=== Integration with ICUDataModule ===") if not check_data_availability(): - print("⚠️ MIMIC-IV demo data not found. Skipping.\n") + print("⚠️ Processed MIIV data not found. Skipping.\n") return False from slices.data import ICUDataModule # Create datamodule datamodule = ICUDataModule( - processed_dir="data/processed/mimic-iv-demo", + processed_dir="data/processed/miiv", task_name="mortality_24h", batch_size=32, num_workers=0, @@ -208,14 +208,14 @@ def different_sequence_lengths(): print("=== Variable-Length Sequences ===") if not check_data_availability(): - print("⚠️ MIMIC-IV demo data not found. Skipping.\n") + print("⚠️ Processed MIIV data not found. Skipping.\n") return False from slices.data import ICUDataset # Load dataset with shorter sequences dataset = ICUDataset( - "data/processed/mimic-iv-demo", + "data/processed/miiv", seq_length=24, # Use only first 24 hours normalize=True, ) @@ -270,10 +270,10 @@ def main(): if any(success): print("✓ Examples completed successfully!") else: - print("⚠️ No examples ran (MIMIC-IV demo data not found)") + print("⚠️ No examples ran (processed MIIV 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("1. Prepare the MIIV source data") + print("2. Run: ./scripts/setup_and_extract.sh miiv") print("=" * 70) diff --git a/pyproject.toml b/pyproject.toml index 08fb326..15101ae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,20 +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", @@ -38,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 diff --git a/scripts/analyze_labels.py b/scripts/analyze_labels.py index a5090f0..8223476 100644 --- a/scripts/analyze_labels.py +++ b/scripts/analyze_labels.py @@ -10,16 +10,16 @@ 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 diff --git a/scripts/debug/explore_missingness.py b/scripts/debug/explore_missingness.py index 04c80fd..a4d1b8c 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 """ @@ -334,7 +334,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)", ) diff --git a/scripts/debug/export_snapshots.py b/scripts/debug/export_snapshots.py new file mode 100644 index 0000000..a96fb09 --- /dev/null +++ b/scripts/debug/export_snapshots.py @@ -0,0 +1,59 @@ +#!/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 fd97461..44bccc2 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 """ diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index 40e5743..f996bbd 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -6,30 +6,31 @@ metrics via FairnessEvaluator, and writes results back to the same W&B run's summary. -Designed for batch evaluation of the thesis fairness corpus across finetune, +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 --sprint/--paradigm/--dataset filters. Usage: - # Evaluate the thesis fairness corpus for one explicit revision - uv run python scripts/eval/evaluate_fairness.py --revision thesis-v1 + # Evaluate the benchmark fairness corpus for one explicit revision + uv run python scripts/eval/evaluate_fairness.py --revision benchmark-v1 # Scope to specific sprint/dataset - uv run python scripts/eval/evaluate_fairness.py --revision thesis-v1 --sprint 1 --dataset miiv + uv run python scripts/eval/evaluate_fairness.py \ + --revision benchmark-v1 --sprint 1 --dataset miiv # Preview which runs would be evaluated - uv run python scripts/eval/evaluate_fairness.py --revision thesis-v1 --dry-run + uv run python scripts/eval/evaluate_fairness.py --revision benchmark-v1 --dry-run # Override paths (e.g., different machine than training) uv run python scripts/eval/evaluate_fairness.py \ - --revision thesis-v1 \ + --revision benchmark-v1 \ --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 --revision thesis-v1 --force + uv run python scripts/eval/evaluate_fairness.py --revision benchmark-v1 --force # Debug with a single run - uv run python scripts/eval/evaluate_fairness.py --revision thesis-v1 --max-runs 1 + uv run python scripts/eval/evaluate_fairness.py --revision benchmark-v1 --max-runs 1 """ from __future__ import annotations @@ -525,7 +526,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--sprint", nargs="+", - help="Filter to sprint(s). Default: thesis fairness sprints", + help="Filter to experiment group(s). Default: benchmark fairness groups", ) parser.add_argument("--paradigm", nargs="+", help="Filter to paradigm(s)") parser.add_argument("--dataset", nargs="+", help="Filter to dataset(s)") @@ -605,9 +606,9 @@ def main() -> None: device = _resolve_device(args.device) log.info("Device: %s", device) - # Default to the thesis fairness sprint set if not specified + # Default to the benchmark fairness groups if not specified. sprints = args.sprint or CORE_SPRINTS - log.info("Sprints: %s", sprints) + log.info("Experiment groups: %s", sprints) # Fetch runs from W&B runs = _retry( diff --git a/scripts/eval/evaluate_imputation.py b/scripts/eval/evaluate_imputation.py index b362034..451e89b 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 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 \ diff --git a/scripts/export_results.py b/scripts/export_results.py index 9b8a141..6513e0e 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -18,7 +18,7 @@ uv run python scripts/export_results.py --sprint 1 --validate-seeds 3 uv run python scripts/export_results.py --paradigm mae jepa --dataset miiv uv run python scripts/export_results.py --output-dir results/sprint1 --sprint 1 - uv run python scripts/export_results.py --project slices-thesis --revision thesis-v1 + uv run python scripts/export_results.py --project slices-benchmark --revision benchmark-v1 """ from __future__ import annotations @@ -45,10 +45,10 @@ # Constants # --------------------------------------------------------------------------- -# Sprint → experiment type mapping. -# Core sprints share the same config and only add seeds; they merge across sprints. -# Non-core sprints have ablation-specific pretrain configs that aren't visible -# in the finetune W&B config, so sprint must remain a grouping key. +# Historical sprint tag to benchmark experiment type mapping. +# Core groups share the same config and only add seeds; they merge across groups. +# Non-core groups have ablation-specific pretrain configs that are not visible +# in the finetune W&B config, so the sprint tag remains a grouping key. EXPERIMENT_TYPES = { "1": "core", "2": "core", @@ -79,7 +79,7 @@ "temporal_contrastive", } -# For core runs: sprint is noise — same config across sprints, different seeds. +# For core runs: the sprint tag is not part of the config identity. # lr/mask_ratio excluded because they're determined by protocol (redundant) or # only in the pretrain config (NaN for finetune runs). CORE_FINGERPRINT = [ @@ -94,7 +94,7 @@ "phase", ] -# HP ablations merge across sprints, but must preserve the upstream pretrain +# HP ablations merge across groups, but must preserve the upstream pretrain # config and subtype so LR and mask-ratio families cannot collapse together. HP_ABLATION_FINGERPRINT = CORE_FINGERPRINT + [ "experiment_subtype", @@ -102,7 +102,7 @@ "upstream_pretrain_mask_ratio", ] -# For per-sprint experiment families, sprint remains part of the identity. +# For per-group experiment families, the sprint tag remains part of the identity. ABLATION_FINGERPRINT = HP_ABLATION_FINGERPRINT + [ "sprint", ] @@ -175,7 +175,7 @@ ALL_METRICS = TEST_METRICS + VAL_METRICS -# Canonical model variants used in the thesis matrix. +# Canonical model variants used in the benchmark matrix. MODEL_VARIANTS = { (64, 2): "default", (128, 4): "medium", @@ -483,9 +483,8 @@ def extract_run(run, metric_keys: list[str]) -> dict: sprint_str = str(config.get("sprint", "")) experiment_type = EXPERIMENT_TYPES.get(sprint_str, "unknown") - # Sprint 10 contains both core (label_fraction=1.0) and label_efficiency - # (label_fraction<1.0) runs. Reclassify the sub-1.0 runs so they merge - # with Sprint 6 label_efficiency during aggregation. + # Historical group 10 contains both core (label_fraction=1.0) and + # label-efficiency runs that merge with group 6 during aggregation. label_frac = config.get("label_fraction", 1.0) if experiment_type == "core" and label_frac is not None and float(label_frac) < 1.0: experiment_type = "label_efficiency" @@ -501,7 +500,7 @@ def extract_run(run, metric_keys: list[str]) -> dict: if experiment_type == "core" and experiment_subtype in ("lr_ablation", "mask_ablation"): experiment_type = "hp_ablation" - # Sprint 10 also adds transfer seeds for Sprint 7 scope. + # Historical group 10 also adds transfer seeds for group 7 scope. if experiment_type == "core" and source_dataset is not None: experiment_type = "transfer" @@ -688,7 +687,7 @@ def _aggregate_group(df: pd.DataFrame, fingerprint_cols: list[str]) -> pd.DataFr """Aggregate a subset of runs by the given fingerprint columns. Returns a DataFrame with one row per unique fingerprint, containing - mean/std/min/max of metrics and lists of run IDs/seeds/sprints. + mean/std/min/max of metrics and lists of run IDs/seeds/groups. """ metric_cols = [c for c in ALL_METRICS if c in df.columns] @@ -740,9 +739,9 @@ def build_aggregated_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: """Group by fingerprint, compute mean/std/min/max across seeds. Uses conditional grouping: - - Core / label-efficiency / transfer runs: group WITHOUT sprint - - HP ablations: group WITHOUT sprint, but include subtype + upstream pretrain config - - Other experiment families: group WITH sprint + - Core / label-efficiency / transfer runs: group without the sprint tag + - HP ablations: group without the sprint tag, but include subtype + upstream pretrain config + - Other experiment families: group with the sprint tag """ parts = [] cross_sprint = per_seed_df[per_seed_df["experiment_type"].isin(CROSS_SPRINT_TYPES)] @@ -788,7 +787,7 @@ def build_aggregated_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: def _primary_metric_for_task(task_name: str | None) -> str | None: - """Return the thesis-primary test metric for a downstream task.""" + """Return the benchmark-primary test metric for a downstream task.""" if task_name is None: return None if task_name in PRIMARY_TEST_METRIC_BY_TASK: @@ -803,8 +802,8 @@ def build_statistical_tests_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: Comparisons are scoped by the canonical experiment fingerprint with `paradigm` and `task` removed, then pooled across matched `(task, seed)` - pairs for tasks that share the same primary metric. This matches the thesis - plan's "paired across seeds and tasks" intent without mixing incompatible + pairs for tasks that share the same primary metric. This matches the benchmark + plan's paired-across-seeds-and-tasks intent without mixing incompatible scales such as AUPRC and MAE in the same test. """ if per_seed_df.empty or "experiment_type" not in per_seed_df.columns: @@ -1038,7 +1037,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--sprint", nargs="+", - help="Filter to specific sprint(s), e.g. --sprint 1 1b 2", + help="Filter to specific experiment group(s), e.g. --sprint 1 1b 2", ) parser.add_argument( "--paradigm", @@ -1059,7 +1058,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--revision", nargs="+", - help="Filter to specific revision tag(s), e.g. --revision thesis-v1", + help="Filter to specific revision tag(s), e.g. --revision benchmark-v1", ) parser.add_argument( "--state", diff --git a/scripts/auto_shutdown.sh b/scripts/internal/auto_shutdown.sh similarity index 84% rename from scripts/auto_shutdown.sh rename to scripts/internal/auto_shutdown.sh index 8608bec..c30c24e 100755 --- a/scripts/auto_shutdown.sh +++ b/scripts/internal/auto_shutdown.sh @@ -1,6 +1,6 @@ #!/usr/bin/env bash # ============================================================================= -# Auto-shutdown: powers off the VM if no training has run for 30 minutes. +# 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 @@ -9,7 +9,7 @@ IDLE_THRESHOLD_MIN=59 STAMP_FILE="/tmp/slices_last_training_activity" # Check if any training process or Claude Code session is running -if pgrep -f "scripts/(training/(pretrain|finetune|supervised)|run_experiments|preprocessing/prepare_dataset)\.py" > /dev/null 2>&1 \ +if pgrep -f "scripts/(training/(pretrain|finetune|supervised)|internal/run_experiments|preprocessing/prepare_dataset)\.py" > /dev/null 2>&1 \ || pgrep -f "claude" > /dev/null 2>&1; then # Activity detected — update timestamp and exit date +%s > "$STAMP_FILE" diff --git a/scripts/launch_thesis_tmux.sh b/scripts/internal/launch_thesis_tmux.sh similarity index 86% rename from scripts/launch_thesis_tmux.sh rename to scripts/internal/launch_thesis_tmux.sh index a0e90b8..e1ba1fe 100755 --- a/scripts/launch_thesis_tmux.sh +++ b/scripts/internal/launch_thesis_tmux.sh @@ -1,7 +1,7 @@ #!/usr/bin/env bash set -euo pipefail -REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" SESSION_NAME="${SESSION_NAME:-slices-thesis}" WANDB_PROJECT="${WANDB_PROJECT:-slices-thesis}" WANDB_ENTITY="${WANDB_ENTITY:-}" @@ -80,14 +80,14 @@ if [[ -n "$REVISION" ]]; then fi fairness_args+=(--sprint "${fairness_sprints[@]}" --batch-size "$BATCH_SIZE_FAIRNESS" --device "$DEVICE_FAIRNESS") -warmup_cmd=(uv run python scripts/run_experiments.py warmup --sprint "${warmup_sprints[@]}") -main_cmd=(uv run python scripts/run_experiments.py run --sprint "${main_sprints[@]}" --parallel "$PARALLEL_MAIN" "${run_args[@]}") -tag_cmd=(uv run python scripts/run_experiments.py tag --sprint "${tag_sprints[@]}" "${tag_args[@]}") +warmup_cmd=(uv run python scripts/internal/run_experiments.py warmup --sprint "${warmup_sprints[@]}") +main_cmd=(uv run python scripts/internal/run_experiments.py run --sprint "${main_sprints[@]}" --parallel "$PARALLEL_MAIN" "${run_args[@]}") +tag_cmd=(uv run python scripts/internal/run_experiments.py tag --sprint "${tag_sprints[@]}" "${tag_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_sprints[@]} > 0)); then - appendix_cmd=(uv run python scripts/run_experiments.py run --sprint "${appendix_sprints[@]}" --parallel "$PARALLEL_APPENDIX" "${run_args[@]}") + appendix_cmd=(uv run python scripts/internal/run_experiments.py run --sprint "${appendix_sprints[@]}" --parallel "$PARALLEL_APPENDIX" "${run_args[@]}") else appendix_cmd=() fi @@ -140,7 +140,7 @@ chmod +x "$runner_script" status_cmd=( bash -lc - "cd $(printf "%q" "$REPO_ROOT"); while true; do clear; date; echo; uv run python scripts/run_experiments.py status; sleep $(printf "%q" "$STATUS_INTERVAL"); done" + "cd $(printf "%q" "$REPO_ROOT"); while true; do clear; date; echo; uv run python scripts/internal/run_experiments.py status; sleep $(printf "%q" "$STATUS_INTERVAL"); done" ) printf -v status_line "%q " "${status_cmd[@]}" diff --git a/scripts/run_experiments.py b/scripts/internal/run_experiments.py similarity index 98% rename from scripts/run_experiments.py rename to scripts/internal/run_experiments.py index 6d4b9cb..8b3900f 100644 --- a/scripts/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -6,18 +6,19 @@ and executes them in parallel with crash recovery and state persistence. Usage: - uv run python scripts/run_experiments.py warmup --sprint 1 - uv run python scripts/run_experiments.py run --sprint 1 --parallel 4 - uv run python scripts/run_experiments.py run --sprint 1 2 3 --parallel 6 --dry-run - uv run python scripts/run_experiments.py run --sprint 1 --revision v2 --reason "fix LR" - uv run python scripts/run_experiments.py run --sprint 1 2 \\ + uv run python scripts/internal/run_experiments.py warmup --sprint 1 + uv run python scripts/internal/run_experiments.py run --sprint 1 --parallel 4 + uv run python scripts/internal/run_experiments.py run --sprint 1 2 3 --parallel 6 --dry-run + uv run python scripts/internal/run_experiments.py run --sprint 1 --revision v2 --reason "fix LR" + uv run python scripts/internal/run_experiments.py run --sprint 1 2 \\ --project slices-thesis --entity myteam - uv run python scripts/run_experiments.py status - uv run python scripts/run_experiments.py status --sprint 1 - uv run python scripts/run_experiments.py retry --failed --parallel 4 - uv run python scripts/run_experiments.py retry --failed --sprint 1 --revision v2 --parallel 4 - uv run python scripts/run_experiments.py tag --sprint 2 --dry-run - uv run python scripts/run_experiments.py tag --sprint 2 3 5 + uv run python scripts/internal/run_experiments.py status + uv run python scripts/internal/run_experiments.py status --sprint 1 + uv run python scripts/internal/run_experiments.py retry --failed --parallel 4 + uv run python scripts/internal/run_experiments.py retry --failed \ + --sprint 1 --revision v2 --parallel 4 + uv run python scripts/internal/run_experiments.py tag --sprint 2 --dry-run + uv run python scripts/internal/run_experiments.py tag --sprint 2 3 5 """ from __future__ import annotations @@ -98,7 +99,7 @@ TRANSFER_PAIRS = [("miiv", "eicu"), ("eicu", "miiv")] # Baseline inheritance: which earlier sprints' runs should be tagged as baselines -# for later sprints. See docs/EXPERIMENT_PLAN.md § "Baseline Inheritance Across Sprints". +# for later sprints. See docs/internal/EXPERIMENT_PLAN.md § "Baseline Inheritance Across Sprints". # Intentionally coarse-grained: ALL runs from source sprints are tagged, even when # only a subset is relevant (e.g. 1b only needs mortality_24h from Sprint 1). Use # secondary W&B tags (task:, protocol:, phase:, etc.) to filter down when querying. @@ -1554,9 +1555,10 @@ def cmd_tag(args): W&B API. Idempotent — already-tagged runs are skipped. Usage: - uv run python scripts/run_experiments.py tag --sprint 2 - uv run python scripts/run_experiments.py tag --sprint 2 3 --dry-run - uv run python scripts/run_experiments.py tag --sprint 2 --project slices --entity myteam + uv run python scripts/internal/run_experiments.py tag --sprint 2 + uv run python scripts/internal/run_experiments.py tag --sprint 2 3 --dry-run + uv run python scripts/internal/run_experiments.py tag --sprint 2 \ + --project slices --entity myteam """ try: import wandb # noqa: F811 diff --git a/scripts/measure_flops.py b/scripts/measure_flops.py index 73d0442..5a700f8 100644 --- a/scripts/measure_flops.py +++ b/scripts/measure_flops.py @@ -17,6 +17,7 @@ import torch import torch.nn as nn +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 from torch.utils.flop_counter import FlopCounterMode @@ -28,7 +29,7 @@ ENCODER_CONFIG = { "d_input": 35, "d_model": 64, - "max_seq_length": 168, + "max_seq_length": SEQ_LENGTH_HOURS, "n_layers": 2, "n_heads": 4, "d_ff": 256, @@ -191,7 +192,9 @@ def main(): parser.add_argument( "--n-features", type=int, default=35, help="Number of input features (d_input)" ) - parser.add_argument("--seq-length", type=int, default=48, help="Sequence length (T)") + 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)" @@ -209,7 +212,11 @@ def main(): args = parser.parse_args() # Build encoder config with CLI overrides - encoder_config = {**ENCODER_CONFIG, "d_input": args.n_features} + 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 diff --git a/scripts/sanity_checks/sc_mae_pretraining.py b/scripts/sanity_checks/sc_mae_pretraining.py index 33ee8db..5ea9068 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 @@ -62,7 +62,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 +115,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 +479,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..99cac8c 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 @@ -193,8 +193,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/src/slices/data/datamodule.py b/src/slices/data/datamodule.py index cabd80e..19e612c 100644 --- a/src/slices/data/datamodule.py +++ b/src/slices/data/datamodule.py @@ -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, ... ) diff --git a/src/slices/data/dataset.py b/src/slices/data/dataset.py index d36d88e..0fead51 100644 --- a/src/slices/data/dataset.py +++ b/src/slices/data/dataset.py @@ -98,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]) @@ -733,7 +733,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]) diff --git a/src/slices/data/sliding_window.py b/src/slices/data/sliding_window.py index 5fbce54..754f1de 100644 --- a/src/slices/data/sliding_window.py +++ b/src/slices/data/sliding_window.py @@ -28,7 +28,7 @@ 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=24, # 24h benchmark windows diff --git a/src/slices/debug/snapshots.py b/src/slices/debug/snapshots.py index ef77524..a627c1f 100644 --- a/src/slices/debug/snapshots.py +++ b/src/slices/debug/snapshots.py @@ -45,6 +45,9 @@ class LegacyPipelineStage(str, Enum): FINAL = "final" +SnapshotStage = Union[PipelineStage, LegacyPipelineStage] + + @dataclass class SnapshotConfig: """Configuration for pipeline snapshots. @@ -59,7 +62,9 @@ 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 = SEQ_LENGTH_HOURS include_masks: bool = True @@ -77,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 @@ -111,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), @@ -149,7 +154,7 @@ def capture_labels_snapshot( } return PipelineSnapshot( - stage=PipelineStage.LABELS, + stage=LegacyPipelineStage.LABELS, data=df, metadata={ "n_stays": len(df), @@ -184,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": ( @@ -395,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: @@ -526,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. @@ -539,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: @@ -568,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: @@ -576,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 848131b..7ee5505 100644 --- a/src/slices/eval/README.md +++ b/src/slices/eval/README.md @@ -52,7 +52,7 @@ There are two layers: 2. `FairnessEvaluator` in `fairness_evaluator.py` builds full subgroup reports from predictions, labels, stay IDs, and `static.parquet` demographics. -Current evaluator behavior matches the thesis pipeline: +Current evaluator behavior matches the benchmark pipeline: - sex and age-group analysis for all supported datasets when columns exist - race/ethnicity fairness only on MIMIC-IV rows @@ -77,7 +77,7 @@ report = evaluator.evaluate(predictions, labels, stay_ids) flat_report = flatten_fairness_report(report) ``` -For batch thesis sweeps, the repository uses +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. @@ -103,7 +103,7 @@ export code can consume. `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 thesis +metrics and is intended for representation diagnostics rather than benchmark headline results. ## Integration Points diff --git a/src/slices/eval/fairness_evaluator.py b/src/slices/eval/fairness_evaluator.py index c63b45f..3d9c816 100644 --- a/src/slices/eval/fairness_evaluator.py +++ b/src/slices/eval/fairness_evaluator.py @@ -3,7 +3,7 @@ Computes per-group AUROC, worst-group AUROC, demographic parity, and equalized odds across protected attributes (gender, age group, race). -Race handling follows the thesis plan: +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) @@ -290,7 +290,7 @@ def evaluate( # Get unique valid groups (exclude -1 = unknown) unique_groups = [g for g in group_ids.unique().tolist() if g >= 0] - # Filter groups below the patient-count threshold from the thesis plan + # Filter groups below the benchmark patient-count threshold. valid_groups = [] group_sizes = {} group_sample_sizes = {} diff --git a/src/slices/eval/statistical.py b/src/slices/eval/statistical.py index fd22058..a80d4d8 100644 --- a/src/slices/eval/statistical.py +++ b/src/slices/eval/statistical.py @@ -168,7 +168,7 @@ def paired_wilcoxon_signed_rank( """Paired Wilcoxon signed-rank test with tie correction. This implementation uses the standard large-sample normal approximation - with tie correction. It is sufficient for the thesis export pipeline where + with tie correction. It is sufficient for the benchmark export pipeline where comparisons pool multiple tasks x seeds and are primarily used for significance tables rather than exact small-sample inference. diff --git a/src/slices/models/encoders/README.md b/src/slices/models/encoders/README.md index 9d51078..721fe05 100644 --- a/src/slices/models/encoders/README.md +++ b/src/slices/models/encoders/README.md @@ -5,7 +5,7 @@ downstream finetuning, and baseline experiments. ## Canonical Benchmark Encoder -The main thesis benchmark uses `TransformerEncoder` from `transformer.py`. +The main benchmark uses `TransformerEncoder` from `transformer.py`. Current canonical config: @@ -95,7 +95,7 @@ ssl_config = TransformerConfig( - `ObservationTransformerEncoder` in `observation.py` - observation-level tokenization - - kept as an alternate architecture, not the canonical thesis encoder + - 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` 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 c762bd5..95fe90d 100644 --- a/src/slices/models/pretraining/README.md +++ b/src/slices/models/pretraining/README.md @@ -35,7 +35,7 @@ These constraints are enforced in the objective constructors and by ## Benchmark Design -The controlled thesis objectives share the same timestep-level obs-aware +The controlled benchmark objectives share the same timestep-level obs-aware transformer encoder and differ only in the SSL objective and masking logic. Highlights: diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index f46fc1d..333ef33 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -557,7 +557,7 @@ def validate_data_prerequisites( label_config_fields = {field.name for field in fields(LabelConfig)} def coerce_label_config( - task_config: Union[LabelConfig, DictConfig, Dict[str, Any]] + task_config: Union[LabelConfig, DictConfig, Dict[str, Any]], ) -> LabelConfig: if isinstance(task_config, LabelConfig): return task_config diff --git a/tests/test_fixes.py b/tests/test_fixes.py index bd172c9..39e7972 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -543,7 +543,7 @@ 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").mkdir(parents=True, exist_ok=True) + (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"): @@ -786,8 +786,11 @@ 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").mkdir(parents=True, exist_ok=True) - shutil.copy2(repo_root / "scripts" / "launch_thesis_tmux.sh", temp_repo / "scripts") + (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() @@ -811,7 +814,7 @@ def test_tmux_launcher_passes_revision_to_fairness(self, tmp_path): env["RUN_EXPORT"] = "0" result = subprocess.run( - ["bash", "scripts/launch_thesis_tmux.sh"], + ["bash", "scripts/internal/launch_thesis_tmux.sh"], cwd=temp_repo, env=env, capture_output=True, @@ -830,8 +833,11 @@ def test_tmux_launcher_includes_sprint11_in_fairness(self, tmp_path): """The fairness sweep should include the canonical baseline sprint by default.""" repo_root = Path(__file__).resolve().parents[1] temp_repo = tmp_path / "repo" - (temp_repo / "scripts").mkdir(parents=True, exist_ok=True) - shutil.copy2(repo_root / "scripts" / "launch_thesis_tmux.sh", temp_repo / "scripts") + (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() @@ -854,7 +860,7 @@ def test_tmux_launcher_includes_sprint11_in_fairness(self, tmp_path): env["RUN_EXPORT"] = "0" result = subprocess.run( - ["bash", "scripts/launch_thesis_tmux.sh"], + ["bash", "scripts/internal/launch_thesis_tmux.sh"], cwd=temp_repo, env=env, capture_output=True, @@ -872,8 +878,11 @@ 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").mkdir(parents=True, exist_ok=True) - shutil.copy2(repo_root / "scripts" / "launch_thesis_tmux.sh", temp_repo / "scripts") + (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() @@ -896,7 +905,7 @@ def test_tmux_status_pane_uses_uv_run_python(self, tmp_path): env["RUN_EXPORT"] = "0" result = subprocess.run( - ["bash", "scripts/launch_thesis_tmux.sh"], + ["bash", "scripts/internal/launch_thesis_tmux.sh"], cwd=temp_repo, env=env, capture_output=True, @@ -904,7 +913,10 @@ def test_tmux_status_pane_uses_uv_run_python(self, tmp_path): ) assert result.returncode == 0, result.stdout + result.stderr - assert "uv\\ run\\ python\\ scripts/run_experiments.py\\ status" in tmux_log.read_text() + assert ( + "uv\\ run\\ python\\ scripts/internal/run_experiments.py\\ status" + in tmux_log.read_text() + ) # ============================================================================ @@ -1466,7 +1478,7 @@ class TestExperimentRunnerWandbOverrides: """Tests for clean project/entity overrides in the experiment runner.""" def test_apply_wandb_target_injects_project_and_entity(self): - from scripts.run_experiments import Run, apply_wandb_target + from scripts.internal.run_experiments import Run, apply_wandb_target run = Run( id="s1_supervised_mortality_24h_miiv_seed42", @@ -1484,7 +1496,7 @@ def test_apply_wandb_target_injects_project_and_entity(self): assert result[0].extra_overrides["logging.wandb_entity"] == "hannes-ill" def test_transfer_finetune_command_propagates_source_dataset(self): - from scripts.run_experiments import Run + from scripts.internal.run_experiments import Run pretrain = Run( id="pretrain_mae_miiv_seed42", @@ -1517,7 +1529,7 @@ class TestExperimentRunnerRetry: """Tests for retry scoping and dependency preservation.""" def test_cmd_retry_respects_sprint_filter_but_keeps_dependencies(self, monkeypatch): - import scripts.run_experiments as runner + import scripts.internal.run_experiments as runner dependency = runner.Run( id="dep_pretrain", @@ -1589,7 +1601,7 @@ def test_cmd_retry_respects_sprint_filter_but_keeps_dependencies(self, monkeypat assert other.id not in scheduled_ids def test_select_ready_runs_prioritizes_pretrains_with_slot_budget(self): - import scripts.run_experiments as runner + import scripts.internal.run_experiments as runner pretrain = runner.Run( id="pretrain", @@ -1621,7 +1633,7 @@ def test_select_ready_runs_prioritizes_pretrains_with_slot_budget(self): 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.run_experiments as runner + import scripts.internal.run_experiments as runner pretrain = runner.Run( id="pretrain", @@ -2144,7 +2156,7 @@ class TestExperimentRunnerMatrix: def test_classical_baselines_are_canonicalized_to_sprint11(self): from collections import Counter - from scripts.run_experiments import generate_all_runs + from scripts.internal.run_experiments import generate_all_runs runs = generate_all_runs() by_sprint = Counter(run.sprint for run in runs) @@ -2163,7 +2175,7 @@ def test_classical_baselines_are_canonicalized_to_sprint11(self): assert non_s11_classical == [] def test_sprint7p_expanded_to_five_seeds(self): - from scripts.run_experiments import BASELINE_SPRINTS, SEEDS_EXTENDED, MatrixBuilder + from scripts.internal.run_experiments import BASELINE_SPRINTS, SEEDS_EXTENDED, MatrixBuilder builder = MatrixBuilder() builder.build_sprint7p() @@ -2173,7 +2185,7 @@ def test_sprint7p_expanded_to_five_seeds(self): assert BASELINE_SPRINTS["7p"] == ["6", "10"] def test_sprint13_includes_both_protocols(self): - from scripts.run_experiments import MatrixBuilder + from scripts.internal.run_experiments import MatrixBuilder builder = MatrixBuilder() builder.build_sprint13() diff --git a/tests/test_snapshots.py b/tests/test_snapshots.py new file mode 100644 index 0000000..8520e2a --- /dev/null +++ b/tests/test_snapshots.py @@ -0,0 +1,41 @@ +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/uv.lock b/uv.lock index 6f232dc..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" @@ -2337,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" @@ -2474,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" @@ -2492,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" @@ -2600,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" @@ -2647,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" @@ -2737,19 +1757,7 @@ wheels = [ { 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/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]] @@ -2767,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" @@ -2851,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" @@ -2896,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" @@ -2972,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" @@ -3019,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]] @@ -3041,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" @@ -3092,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" @@ -3206,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" @@ -3261,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" @@ -3312,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" @@ -3539,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" @@ -3588,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" @@ -3611,21 +2279,22 @@ 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" }, ] @@ -3639,6 +2308,10 @@ dev = [ { name = "pytest-cov" }, { name = "ruff" }, ] +viz = [ + { name = "matplotlib" }, + { name = "umap-learn" }, +] [package.dev-dependencies] dev = [ @@ -3654,30 +2327,33 @@ dev = [ [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 = [ @@ -3699,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" @@ -3899,21 +2487,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/16/b5/b0d3d8b901b6a04ca38df5e24c27e53afb15b93624d7fd7d658c7cd9352a/triton-3.5.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bac7f7d959ad0f48c0e97d6643a1cc0fd5786fe61cb1f83b537c6b2d54776478", size = 170582192, upload-time = "2025-11-11T17:41:23.963Z" }, ] -[[package]] -name = "typer" -version = "0.21.1" -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" } -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" }, -] - [[package]] name = "types-pyyaml" version = "6.0.12.20250915" @@ -3954,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]] @@ -4018,100 +2595,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/30/20/6c091d451e2a07689bfbfaeb7592d488011420e721de170884fedd68c644/wandb-0.23.1-py3-none-win_arm64.whl", hash = "sha256:8aee7f3bb573f2c0acf860f497ca9c684f9b35f2ca51011ba65af3d4592b77c1", size = 20137463, upload-time = "2025-12-03T02:25:08.317Z" }, ] -[[package]] -name = "websockets" -version = "16.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" }, -] - -[[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" } -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" }, -] - [[package]] name = "xgboost" version = "3.2.0" @@ -4232,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" }, -] From dbf9aa96a687cb203446105af5c6176a488efefc Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 21 Apr 2026 14:19:35 -0400 Subject: [PATCH 085/121] Document AKI label filtering in README Explain that supervised task evaluation excludes missing labels and that AKI requires creatinine observations in both the baseline and prediction windows. This clarifies that AKI results use a task-specific labelable cohort rather than the full ICU stay cohort. --- README.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/README.md b/README.md index 349817a..5d8fcea 100644 --- a/README.md +++ b/README.md @@ -230,6 +230,10 @@ Extracted ICU stays are stored as Parquet files: - `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 From 866cd654ca6054ce682718f8844c7f49a9cc716b Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 21 Apr 2026 14:21:33 -0400 Subject: [PATCH 086/121] Tighten thesis experiment evaluation scope Use strict linear heads for Protocol A probes, include ablation and transfer sprints in fairness evaluation, and preserve revision scoping across retry dependencies and XGBoost W&B tags. This changes exported statistical tests to add contextual classical-vs-neural comparisons and updates the experiment plan to document the AKI labelable cohort and protocol head behavior. --- docs/internal/EXPERIMENT_PLAN.md | 22 +++- scripts/eval/evaluate_fairness.py | 2 +- scripts/export_results.py | 51 +++++++- scripts/internal/launch_thesis_tmux.sh | 2 +- scripts/internal/run_experiments.py | 18 ++- scripts/training/xgboost_baseline.py | 25 ++-- tests/test_evaluate_fairness.py | 6 + tests/test_export_results.py | 43 +++++++ tests/test_fixes.py | 159 ++++++++++++++++++++++++- 9 files changed, 307 insertions(+), 21 deletions(-) diff --git a/docs/internal/EXPERIMENT_PLAN.md b/docs/internal/EXPERIMENT_PLAN.md index 1628cca..1465b4a 100644 --- a/docs/internal/EXPERIMENT_PLAN.md +++ b/docs/internal/EXPERIMENT_PLAN.md @@ -38,10 +38,19 @@ Secondary questions addressed through ablations: | Min stay | 24 hours | Ensures a full observation window before downstream labeling | | 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 | +| Downstream head | Protocol A: linear; Protocol B/supervised: MLP, hidden_dims=[64] | Same head within each protocol across all paradigms | | Precision | fp32 | Avoids bf16 numerical issues with small model | | Weight decay | 0.05 | Same regularization across all downstream runs | +**AKI labelable cohort**: `aki_kdigo` is evaluated on a task-specific subset +with ascertainable creatinine-based KDIGO labels. A stay requires non-null +creatinine in the 0-24h baseline window and at least one creatinine measurement +in the 24-48h prediction window; stays without an ascertainable AKI label are +excluded from AKI supervised training/evaluation but remain available for SSL +pretraining and non-AKI tasks. This leaves 56,403/74,829 MIIV stays, +88,239/132,900 eICU stays, and 144,642/207,729 combined stays. Most exclusions +are due to no 24-48h creatinine measurement. + ### 1.3 Shared Encoder Architecture All three SSL paradigms share the same encoder architecture, removing tokenization as a confounding variable: @@ -92,13 +101,13 @@ Isolates representation quality — the only variable is the SSL pretraining obj | Learning rate | 1e-4 | | Scheduler | Cosine decay (eta_min=1e-6) | | Early stopping | Patience=10 on val AUPRC (classification) or val MAE (regression) | -| Head | MLP, hidden_dims=[64], dropout=0.3, ReLU | +| Head | Single linear layer, dropout=0.0 | | Gradient clipping | 1.0 | | Label smoothing | 0.1 | | Weight decay | 0.05 | | Class weighting | sqrt(balanced) — square root of inverse frequency | -**Rationale**: Linear probing answers "Which SSL objective produces the best representations?" by preventing the encoder from adapting to the downstream task. +**Rationale**: Linear probing answers "Which SSL objective produces the best representations?" by preventing the encoder from adapting to the downstream task and by removing nonlinear head capacity as a confound. #### Protocol B: Full Finetuning (practical utility) @@ -138,7 +147,7 @@ Measures how useful SSL pretraining is as weight initialization. | Weight decay | 0.05 | | Class weighting | sqrt(balanced) — square root of inverse frequency | -**Same-architecture rationale**: The supervised baseline uses the identical encoder architecture and tokenization as the SSL paradigms to eliminate model capacity as a confounding variable. The only difference is the training procedure: supervised trains end-to-end on labeled data from random initialization, while SSL paradigms pretrain on unlabeled data and then evaluate via Protocol A (linear probe) and Protocol B (full finetune). Comparing supervised vs. Protocol B isolates the effect of SSL initialization; comparing SSL paradigms via Protocol A isolates representation quality. +**Same-architecture rationale**: The supervised baseline uses the identical encoder architecture and tokenization as the SSL paradigms to eliminate model capacity as a confounding variable. The only difference is the training procedure: supervised trains end-to-end on labeled data from random initialization, while SSL paradigms pretrain on unlabeled data and then evaluate via Protocol A (strict linear probe) and Protocol B (full finetune). Comparing supervised vs. Protocol B isolates the effect of SSL initialization; comparing SSL paradigms via Protocol A isolates representation quality. --- @@ -155,6 +164,9 @@ Measures how useful SSL pretraining is as weight initialization. **AUPRC as primary** for classification: ICU tasks are heavily imbalanced; AUPRC is more informative than AUROC in this regime. +For `aki_kdigo`, prevalence and sample counts are reported on the labelable +cohort defined in Section 1.2, not on the full >=24h ICU stay cohort. + ### 3.2 Fairness Metrics Computed on test set predictions for all classification tasks: @@ -411,7 +423,7 @@ and the thesis W&B project summaries. **Runs last**, after all training sprints are complete. Zero additional training — pure evaluation on existing test predictions. -1. Compute fairness metrics on all downstream test predictions from Sprints 1–5, 7p, 10, 11, 12, and 13 +1. Compute fairness metrics on all downstream test predictions from Sprints 1–8, 7p, 10, 11, 12, and 13 2. Run `scripts/eval/evaluate_fairness.py` with an explicit `--revision` tag for the thesis rerun corpus 3. Generate fairness tables and disparity plots 4. Protected attributes: sex (all datasets), age group (all), race/ethnicity (MIMIC-IV only) diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index f996bbd..b61a71a 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -53,7 +53,7 @@ # Constants # --------------------------------------------------------------------------- -CORE_SPRINTS = ["1", "2", "3", "4", "5", "7p", "10", "11", "12", "13"] +CORE_SPRINTS = ["1", "2", "3", "4", "5", "6", "7", "8", "7p", "10", "11", "12", "13"] DEFAULT_PHASES = ["finetune", "supervised", "baseline"] DEFAULT_PROTECTED_ATTRIBUTES = ["gender", "age_group", "race"] diff --git a/scripts/export_results.py b/scripts/export_results.py index 6513e0e..c48dcfd 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -68,6 +68,10 @@ } CROSS_SPRINT_TYPES = {"core", "label_efficiency", "transfer"} +CONTEXTUAL_STAT_TYPES = { + "classical_context_full", + "classical_context_label_efficiency", +} FIXED_SEED_EXPERIMENT_TYPES = { "core", "label_efficiency", @@ -543,7 +547,7 @@ def extract_run(run, metric_keys: list[str]) -> dict: def _fingerprint_for_experiment_type(experiment_type: str) -> list[str]: """Return the canonical fingerprint for a given experiment family.""" - if experiment_type in CROSS_SPRINT_TYPES: + if experiment_type in CROSS_SPRINT_TYPES or experiment_type in CONTEXTUAL_STAT_TYPES: return CORE_FINGERPRINT if experiment_type == "hp_ablation": return HP_ABLATION_FINGERPRINT @@ -797,6 +801,44 @@ def _primary_metric_for_task(task_name: str | None) -> str | None: return "test/auprc" +def _add_contextual_classical_stat_rows(per_seed_df: pd.DataFrame) -> pd.DataFrame: + """Add synthetic statistic scopes for classical-vs-neural comparisons.""" + required = {"experiment_type", "protocol", "label_fraction", "paradigm"} + if per_seed_df.empty or not required.issubset(per_seed_df.columns): + return per_seed_df + + work = per_seed_df.copy() + label_fraction = pd.to_numeric(work["label_fraction"], errors="coerce").fillna(1.0) + protocol_b = work["protocol"] == "B" + classical = work["experiment_type"] == "classical_baselines" + + parts = [work] + + full_context_mask = ( + protocol_b & (label_fraction == 1.0) & ((work["experiment_type"] == "core") | classical) + ) + label_efficiency_context_mask = ( + protocol_b + & (label_fraction < 1.0) + & ((work["experiment_type"] == "label_efficiency") | classical) + ) + + for mask, experiment_type in [ + (full_context_mask, "classical_context_full"), + (label_efficiency_context_mask, "classical_context_label_efficiency"), + ]: + context = work.loc[mask].copy() + if context.empty: + continue + paradigms = set(context["paradigm"].dropna()) + if not (paradigms & {"xgboost", "gru_d"}) or paradigms <= {"xgboost", "gru_d"}: + continue + context["experiment_type"] = experiment_type + 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: """Build pairwise paradigm significance tables from per-seed results. @@ -809,6 +851,7 @@ def build_statistical_tests_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: if per_seed_df.empty or "experiment_type" not in per_seed_df.columns: return pd.DataFrame() + per_seed_df = _add_contextual_classical_stat_rows(per_seed_df) rows = [] for experiment_type in sorted(per_seed_df["experiment_type"].dropna().unique()): @@ -859,6 +902,12 @@ def build_statistical_tests_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: family_rows = [] for paradigm_a, paradigm_b in combinations(paradigms, 2): + if experiment_type in CONTEXTUAL_STAT_TYPES: + a_classical = paradigm_a in {"gru_d", "xgboost"} + b_classical = paradigm_b in {"gru_d", "xgboost"} + if a_classical == b_classical: + continue + pairs_a = scope_group[scope_group["paradigm"] == paradigm_a][ ["task", "seed", "primary_metric_value"] ].rename(columns={"primary_metric_value": "value_a"}) diff --git a/scripts/internal/launch_thesis_tmux.sh b/scripts/internal/launch_thesis_tmux.sh index e1ba1fe..0f929a7 100755 --- a/scripts/internal/launch_thesis_tmux.sh +++ b/scripts/internal/launch_thesis_tmux.sh @@ -32,7 +32,7 @@ fi mkdir -p "$LOG_DIR" main_sprints=(1 1b 1c 2 3 4 5 6 7 8 10 11) -fairness_sprints=(1 2 3 4 5 10 11) +fairness_sprints=(1 2 3 4 5 6 7 8 10 11) tag_sprints=(1b 1c 2 5 6 7 8) appendix_sprints=() diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index 8b3900f..d22b564 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -26,6 +26,7 @@ import json import os import re +import shlex import signal import subprocess import sys @@ -205,6 +206,9 @@ def _finetune_cmd(self, runs_by_id: dict[str, Run]) -> list[str]: f"training.max_epochs={PROTO_A['max_epochs']}", f"training.early_stopping_patience={PROTO_A['patience']}", f"optimizer.lr={PROTO_A['lr']}", + "task.head_type=linear", + "task.hidden_dims=[]", + "task.dropout=0.0", ] elif self.freeze_encoder is False: cmd += [ @@ -1254,7 +1258,7 @@ def _print_dry_run(runs: list[Run], runs_by_id: dict[str, Run]): print(f"DRY RUN: {len(runs)} runs\n") for r in runs: deps = ", ".join(r.depends_on) if r.depends_on else "(none)" - cmd = " ".join(r.build_command(runs_by_id)) + cmd = shlex.join(r.build_command(runs_by_id)) print(f"[{r.id}]") print(f" sprint={r.sprint} type={r.run_type} deps={deps}") print(f" dir={r.output_dir}") @@ -1432,7 +1436,17 @@ def cmd_retry(args): if args.revision: if sprint_filter: revised = [r for r in all_runs if r.sprint in sprint_filter] - rest = [r for r in all_runs if r.sprint not in sprint_filter] + revised_ids = {r.id for r in revised} + all_by_id = {r.id: r for r in all_runs} + deps_needed = { + dep_id + for r in revised + for dep_id in r.depends_on + if dep_id not in revised_ids and dep_id in all_by_id + } + revised = [all_by_id[dep_id] for dep_id in deps_needed] + revised + revised_ids = {r.id for r in revised} + rest = [r for r in all_runs if r.id not in revised_ids] revised = apply_revision(revised, args.revision, args.reason) all_runs = rest + revised else: diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index 8586326..b26fec8 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -33,6 +33,23 @@ def _xgboost_eval_metric(task_type: str) -> str: return "mae" if task_type == "regression" else "aucpr" +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("sprint") is not None: + tags.append(f"sprint:{cfg.sprint}") + if cfg.get("revision") is not None: + tags.append(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] + "..." + tags.append(tag) + if cfg.get("label_fraction", 1.0) < 1.0: + tags.append(f"label_fraction:{cfg.label_fraction}") + return tags or None + + @hydra.main(version_base=None, config_path="../../configs", config_name="xgboost") def main(cfg: DictConfig) -> None: """Train XGBoost baseline.""" @@ -219,18 +236,12 @@ def main(cfg: DictConfig) -> None: if cfg.logging.get("use_wandb", False): import wandb - tags = list(cfg.logging.get("wandb_tags", [])) - if cfg.get("sprint"): - tags.append(f"sprint:{cfg.sprint}") - if cfg.get("label_fraction", 1.0) < 1.0: - tags.append(f"label_fraction:{cfg.label_fraction}") - run = wandb.init( project=cfg.logging.wandb_project, entity=cfg.logging.get("wandb_entity", None), name=cfg.logging.get("run_name", f"xgboost_{cfg.dataset}_{task_name}"), group=cfg.logging.get("wandb_group", f"xgboost_{cfg.dataset}_{task_name}"), - tags=tags, + tags=_build_wandb_tags(cfg), config=OmegaConf.to_container(cfg, resolve=True), ) run.summary.update(metrics) diff --git a/tests/test_evaluate_fairness.py b/tests/test_evaluate_fairness.py index 3d7bb03..fa92851 100644 --- a/tests/test_evaluate_fairness.py +++ b/tests/test_evaluate_fairness.py @@ -16,6 +16,12 @@ def test_default_core_sprints_include_sprint11(): assert "11" in mod.CORE_SPRINTS +def test_default_core_sprints_include_ablation_and_transfer_sprints(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + assert {"6", "7", "8"}.issubset(set(mod.CORE_SPRINTS)) + + def test_resolve_evaluation_artifact_supports_xgboost(tmp_path): mod = importlib.import_module("scripts.eval.evaluate_fairness") diff --git a/tests/test_export_results.py b/tests/test_export_results.py index 9c0e921..f006a22 100644 --- a/tests/test_export_results.py +++ b/tests/test_export_results.py @@ -47,6 +47,49 @@ def test_build_statistical_tests_df_produces_pairwise_significance_rows(): assert row["p_value_bonferroni"] >= row["p_value"] +def test_build_statistical_tests_df_adds_classical_context_rows(): + mod = importlib.import_module("scripts.export_results") + + rows = [] + for experiment_type, paradigm, phase, offset in [ + ("core", "mae", "finetune", 0.10), + ("core", "supervised", "supervised", 0.0), + ("classical_baselines", "xgboost", "baseline", 0.03), + ("classical_baselines", "gru_d", "baseline", 0.05), + ]: + for seed in [42, 123]: + for task, base in [("mortality_24h", 0.30), ("aki_kdigo", 0.25)]: + rows.append( + { + "experiment_type": experiment_type, + "sprint": "11" if experiment_type == "classical_baselines" else "1", + "paradigm": paradigm, + "dataset": "miiv", + "task": task, + "seed": seed, + "protocol": "B", + "label_fraction": 1.0, + "model_size": "default", + "source_dataset": None, + "phase": phase, + "test/auprc": base + offset + (seed / 100000.0), + } + ) + + stats_df = mod.build_statistical_tests_df(pd.DataFrame(rows)) + contextual = stats_df[stats_df["experiment_type"] == "classical_context_full"] + pairs = { + tuple(sorted((row["paradigm_a"], row["paradigm_b"]))) for _, row in contextual.iterrows() + } + + assert ("supervised", "xgboost") in pairs + assert ("mae", "xgboost") in pairs + assert ("gru_d", "supervised") in pairs + assert ("gru_d", "mae") in pairs + assert ("gru_d", "xgboost") not in pairs + assert ("mae", "supervised") not in pairs + + def test_parse_args_uses_revision_env_when_cli_omits_it(monkeypatch): mod = importlib.import_module("scripts.export_results") diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 39e7972..ea9b3f9 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -634,11 +634,13 @@ def test_parse_args_requires_revision_when_no_scope_is_available(self, monkeypat assert excinfo.value.code == 2 - def test_default_core_sprints_include_sprint12(self): - """The standalone thesis fairness corpus should include Sprint 12.""" + def test_default_core_sprints_include_all_downstream_training_sprints(self): + """The standalone thesis fairness corpus should include every downstream sprint.""" mod = importlib.import_module("scripts.eval.evaluate_fairness") - assert "12" in mod.CORE_SPRINTS + assert set(["1", "2", "3", "4", "5", "6", "7", "8", "7p", "10", "11", "12", "13"]).issubset( + set(mod.CORE_SPRINTS) + ) def test_fetch_eval_runs_single_revision_adds_server_side_filter(self, monkeypatch): """A single revision should be sent as a W&B tag filter.""" @@ -872,7 +874,7 @@ def test_tmux_launcher_includes_sprint11_in_fairness(self, tmp_path): assert len(runner_scripts) == 1 runner_text = runner_scripts[0].read_text() - assert "--sprint 1 2 3 4 5 10 11" in runner_text + assert "--sprint 1 2 3 4 5 6 7 8 10 11" in runner_text def test_tmux_status_pane_uses_uv_run_python(self, tmp_path): """The status pane should use uv-managed Python.""" @@ -1349,6 +1351,28 @@ def test_extract_run_assigns_protocol_b_to_xgboost_baselines(self): assert extract_run(run, [])["protocol"] == "B" + 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"]}, + "sprint": "11", + "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 "sprint:11" in tags + assert "revision:thesis-v1" in tags + assert "label_fraction:0.1" in tags + assert any(tag.startswith("rerun-reason:") and len(tag) <= 64 for tag in tags) + 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 @@ -1524,6 +1548,68 @@ def test_transfer_finetune_command_propagates_source_dataset(self): cmd = finetune.build_command({pretrain.id: pretrain, finetune.id: finetune}) assert "+source_dataset=miiv" in cmd + def test_protocol_a_finetune_command_uses_strict_linear_probe(self): + from scripts.internal.run_experiments import Run + + pretrain = Run( + id="pretrain_mae_miiv_seed42", + sprint="2", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/sprint1/pretrain_mae_miiv_seed42", + ) + probe = Run( + id="probe_mae_miiv_seed42", + sprint="2", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/sprint2/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 "training.freeze_encoder=true" in cmd + assert "task.head_type=linear" in cmd + assert "task.hidden_dims=[]" in cmd + assert "task.dropout=0.0" in cmd + + def test_protocol_b_finetune_command_keeps_task_mlp_head(self): + from scripts.internal.run_experiments import Run + + pretrain = Run( + id="pretrain_mae_miiv_seed42", + sprint="1", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/sprint1/pretrain_mae_miiv_seed42", + ) + finetune = Run( + id="finetune_mae_miiv_seed42", + sprint="1", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/sprint1/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 "training.freeze_encoder=false" in cmd + assert "task.head_type=linear" not in cmd + class TestExperimentRunnerRetry: """Tests for retry scoping and dependency preservation.""" @@ -1600,6 +1686,71 @@ def test_cmd_retry_respects_sprint_filter_but_keeps_dependencies(self, monkeypat assert dependency.id in scheduled_ids assert other.id not in scheduled_ids + def test_cmd_retry_revises_dependencies_for_revision_scoped_retry(self, monkeypatch): + import scripts.internal.run_experiments as runner + + dependency = runner.Run( + id="s1_pretrain_mae_miiv_seed42", + sprint="1", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/sprint1/pretrain_mae_miiv_seed42", + ) + target = runner.Run( + id="s2_finetune_mae_mortality_24h_miiv_seed42", + sprint="2", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/sprint2/finetune_mae_mortality_24h_miiv_seed42", + depends_on=[dependency.id], + task="mortality_24h", + ) + + revised_dep_id = "s1_rev-thesis-v1_pretrain_mae_miiv_seed42" + revised_target_id = "s2_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( + sprint=["2"], + 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_select_ready_runs_prioritizes_pretrains_with_slot_budget(self): import scripts.internal.run_experiments as runner From cd3a34de9cd5254555a3c4141866aa1850e1946d Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 21 Apr 2026 15:03:14 -0400 Subject: [PATCH 087/121] Harden benchmark validation gates Add CI for linting, type checking, and tests; make fairness and result export fail closed unless --allow-incomplete is set; ignore regenerated rulesync state. Behavioral implications: empty W&B result scopes and incomplete validation now fail by default, so benchmark exports should stop before silently producing incomplete artifacts. The stale example scripts are removed from the public surface. --- .github/workflows/ci.yml | 36 ++ .gitignore | 1 + examples/eval_metrics_example.py | 226 ----------- examples/mae_pretraining_example.py | 406 ------------------- examples/transformer_encoder_example.py | 273 ------------- examples/transformer_with_dataset_example.py | 282 ------------- scripts/eval/evaluate_fairness.py | 13 +- scripts/export_results.py | 10 +- tests/test_evaluate_fairness.py | 29 ++ tests/test_export_results.py | 31 ++ 10 files changed, 118 insertions(+), 1189 deletions(-) create mode 100644 .github/workflows/ci.yml delete mode 100644 examples/eval_metrics_example.py delete mode 100644 examples/mae_pretraining_example.py delete mode 100644 examples/transformer_encoder_example.py delete mode 100644 examples/transformer_with_dataset_example.py 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 90778d6..089124c 100644 --- a/.gitignore +++ b/.gitignore @@ -202,6 +202,7 @@ cython_debug/ .cursorrules .claude/ .cursor/ +.rulesync/ # Marimo marimo/_static/ 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 214ad90..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/miiv") - 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("Run ./scripts/setup_and_extract.sh miiv 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/training/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 72293ca..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 benchmark-style ICU data: batch_size=16, seq_length=24 hours, features=35 - batch_size = 16 - seq_length = 24 - 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 = 24 - 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 8bcda9f..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 processed MIIV data is available.""" - data_dir = Path("data/processed/miiv") - 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("⚠️ Processed MIIV data not found.") - print(" Run: ./scripts/setup_and_extract.sh miiv") - print(" This example requires extracted ICU data.\n") - return False - - from slices.data import ICUDataset - - # Load dataset - dataset = ICUDataset( - "data/processed/miiv", - 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("⚠️ Processed MIIV data not found. Skipping.\n") - return False - - from slices.data import ICUDataset - from torch.utils.data import DataLoader - - # Load dataset - dataset = ICUDataset( - "data/processed/miiv", - 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("⚠️ Processed MIIV data not found. Skipping.\n") - return False - - from slices.data import ICUDataModule - - # Create datamodule - datamodule = ICUDataModule( - processed_dir="data/processed/miiv", - 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("⚠️ Processed MIIV data not found. Skipping.\n") - return False - - from slices.data import ICUDataset - - # Load dataset with shorter sequences - dataset = ICUDataset( - "data/processed/miiv", - 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 (processed MIIV data not found)") - print("\nTo run these examples:") - print("1. Prepare the MIIV source data") - print("2. Run: ./scripts/setup_and_extract.sh miiv") - - print("=" * 70) - - -if __name__ == "__main__": - main() diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index b61a71a..1061c16 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -566,6 +566,11 @@ def parse_args() -> argparse.Namespace: "--min-subgroup-size", type=int, default=50, help="Min patients per subgroup" ) parser.add_argument("--dry-run", action="store_true", help="List runs without processing") + 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", @@ -625,7 +630,7 @@ def main() -> None: if not runs: print("No runs found matching filters.", file=sys.stderr) - sys.exit(0) + sys.exit(0 if args.allow_incomplete else 1) # Filter out runs that already have fairness metrics (unless --force) if args.skip_existing and not args.force: @@ -634,6 +639,9 @@ def main() -> None: skipped = before - len(runs) if skipped: log.info("Skipped %d runs with 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] @@ -774,6 +782,9 @@ def _sort_key(r): 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/export_results.py b/scripts/export_results.py index c48dcfd..292816a 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -1125,6 +1125,11 @@ def parse_args() -> argparse.Namespace: default=5, help="Expected seed count for core configs (default: 5)", ) + parser.add_argument( + "--allow-incomplete", + action="store_true", + help="Exit successfully even if no runs match or validation warnings are emitted.", + ) args = parser.parse_args() if not args.revision: @@ -1159,7 +1164,7 @@ def main(): if not runs: print("No runs found matching filters. Exiting.", file=sys.stderr) - sys.exit(0) + sys.exit(0 if args.allow_incomplete else 1) # Build DataFrames print(f"\nBuilding per-seed DataFrame from {len(runs)} runs...", file=sys.stderr) @@ -1203,6 +1208,9 @@ def main(): for etype, group in aggregated_df.groupby("experiment_type"): seed_dist = dict(group["n_seeds"].value_counts().sort_index()) print(f" {etype}: {len(group)} configs, seeds: {seed_dist}") + + if warnings and not args.allow_incomplete: + sys.exit(1) if "paradigm" in aggregated_df.columns: print(f"Paradigms: {sorted(aggregated_df['paradigm'].dropna().unique())}") if "dataset" in aggregated_df.columns: diff --git a/tests/test_evaluate_fairness.py b/tests/test_evaluate_fairness.py index fa92851..90c4786 100644 --- a/tests/test_evaluate_fairness.py +++ b/tests/test_evaluate_fairness.py @@ -1,8 +1,11 @@ """Focused tests for the standalone fairness rerun script.""" import importlib +import sys from types import SimpleNamespace +import pytest + def test_default_phases_include_baseline(): mod = importlib.import_module("scripts.eval.evaluate_fairness") @@ -22,6 +25,32 @@ def test_default_core_sprints_include_ablation_and_transfer_sprints(): assert {"6", "7", "8"}.issubset(set(mod.CORE_SPRINTS)) +def test_parse_args_exposes_allow_incomplete_escape_hatch(monkeypatch): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + monkeypatch.setattr( + sys, + "argv", + ["evaluate_fairness.py", "--revision", "thesis-v1", "--allow-incomplete"], + ) + + args = mod.parse_args() + + assert args.allow_incomplete is True + + +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") diff --git a/tests/test_export_results.py b/tests/test_export_results.py index f006a22..c618b50 100644 --- a/tests/test_export_results.py +++ b/tests/test_export_results.py @@ -4,6 +4,7 @@ import sys import pandas as pd +import pytest def test_build_statistical_tests_df_produces_pairwise_significance_rows(): @@ -102,6 +103,36 @@ def test_parse_args_uses_revision_env_when_cli_omits_it(monkeypatch): assert args.revision == ["thesis-v1"] +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"], + ) + + args = mod.parse_args() + + assert args.allow_incomplete is True + + +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_parse_args_requires_revision_without_cli_or_env(monkeypatch): mod = importlib.import_module("scripts.export_results") From 87eff39862ebd620a760b7150ffc034501110900 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 21 Apr 2026 15:24:54 -0400 Subject: [PATCH 088/121] Add processed data rebuild wrapper Introduce a lightweight Python orchestrator for rebuilding data/processed from existing RICU outputs, including optional combined dataset creation. Document the wrapper in the quick start and add regression coverage for the generated command sequence. This does not change extraction semantics or migrate existing artifacts. --- README.md | 13 +- scripts/preprocessing/build_processed_data.py | 137 ++++++++++++++++++ tests/test_fixes.py | 45 ++++++ 3 files changed, 193 insertions(+), 2 deletions(-) create mode 100644 scripts/preprocessing/build_processed_data.py diff --git a/README.md b/README.md index 5d8fcea..eb7c043 100644 --- a/README.md +++ b/README.md @@ -97,10 +97,19 @@ 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 ``` 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/tests/test_fixes.py b/tests/test_fixes.py index ea9b3f9..a6e0b73 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -592,6 +592,51 @@ def test_setup_and_extract_default_path_includes_combined(self, tmp_path): 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.""" From 2120ca4d747679db4f8533d9106aaccd21a48461 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 21 Apr 2026 15:39:57 -0400 Subject: [PATCH 089/121] Filter invalid RICU stays before cohort validation Exclude RICU stay rows with missing or negative LOS before applying the benchmark minimum-stay filter, then validate the extracted cohort afterward. This lets known-bad raw source rows outside the cohort avoid blocking processed-data rebuilds while preserving fail-closed validation for rows that enter the dataset. No migration is required; regenerated processed artifacts retain the same benchmark cohort sizes. --- src/slices/data/extractors/ricu.py | 17 +++++++++++++++-- tests/test_ricu_extractor.py | 12 ++++++++---- 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/src/slices/data/extractors/ricu.py b/src/slices/data/extractors/ricu.py index 3c71533..0af8e1b 100644 --- a/src/slices/data/extractors/ricu.py +++ b/src/slices/data/extractors/ricu.py @@ -355,11 +355,23 @@ def run(self) -> None: stays = self.extract_stays() progress.update(task, completed=True) - self._validate_stays(stays) + # 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( @@ -385,6 +397,7 @@ def run(self) -> None: console.print("[red]Error: No stays remaining after filtering![/red]") return + self._validate_stays(stays_filtered) self._stays_cache = stays_filtered # ----------------------------------------------------------------- diff --git a/tests/test_ricu_extractor.py b/tests/test_ricu_extractor.py index f39d9eb..73065b0 100644 --- a/tests/test_ricu_extractor.py +++ b/tests/test_ricu_extractor.py @@ -912,8 +912,10 @@ def test_run_rebuilds_when_builder_version_changes( updated_metadata["label_manifest"]["mortality_hospital"]["builder_version"] == "9.9.9" ) - def test_run_fails_closed_on_invalid_los(self, ricu_output_dir: Path, tmp_path: Path) -> None: - """Abort when stay metadata is incomplete instead of emitting a partial dataset.""" + 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" @@ -933,8 +935,10 @@ def test_run_fails_closed_on_invalid_los(self, ricu_output_dir: Path, tmp_path: tasks=[], ) - with pytest.raises(ValueError, match="invalid LOS"): - RicuExtractor(config).run() + 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 From 94bb249edda62050c73f42292cab6f14984d23fe Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 21 Apr 2026 18:48:16 -0400 Subject: [PATCH 090/121] Fix thesis preflight training blockers Compute SSL and downstream splits from one full-cohort patient partition before task filtering to prevent split leakage. Mask fully unobserved timesteps in the MAE decoder, make the experiment scheduler fail nonzero for failed or dependency-skipped runs, and document explicit thesis W&B project, revision, and entity usage. Behavioral implications: new splits.yaml writes preserve canonical full-cohort patient sets; thesis launch now stops before tagging, fairness, and export when training is incomplete. --- docs/EVAL_FAIRNESS.md | 10 +- docs/internal/EXPERIMENT_PLAN.md | 28 ++- scripts/eval/evaluate_fairness.py | 17 +- scripts/export_results.py | 19 +- scripts/internal/launch_thesis_tmux.sh | 20 +- scripts/internal/run_experiments.py | 77 ++++--- src/slices/data/datamodule.py | 23 +-- src/slices/data/splits.py | 267 +++++++++++++++++++------ src/slices/models/pretraining/mae.py | 14 +- tests/test_dataset_datamodule.py | 165 ++++++++++++++- tests/test_fixes.py | 44 ++++ tests/test_mae_objective.py | 34 ++++ 12 files changed, 574 insertions(+), 144 deletions(-) diff --git a/docs/EVAL_FAIRNESS.md b/docs/EVAL_FAIRNESS.md index 8923b97..8cb1965 100644 --- a/docs/EVAL_FAIRNESS.md +++ b/docs/EVAL_FAIRNESS.md @@ -115,13 +115,19 @@ summary metrics. 1. Run the fairness sweep for a specific revision: ```bash -uv run python scripts/eval/evaluate_fairness.py --revision benchmark-v1 +uv run python scripts/eval/evaluate_fairness.py \ + --project slices-thesis \ + --revision thesis-v1 \ + --entity ``` 2. Re-export the result tables: ```bash -uv run python scripts/export_results.py --revision benchmark-v1 +uv run python scripts/export_results.py \ + --project slices-thesis \ + --revision thesis-v1 \ + --entity ``` ## Source Of Truth diff --git a/docs/internal/EXPERIMENT_PLAN.md b/docs/internal/EXPERIMENT_PLAN.md index 1465b4a..0ef2452 100644 --- a/docs/internal/EXPERIMENT_PLAN.md +++ b/docs/internal/EXPERIMENT_PLAN.md @@ -349,6 +349,16 @@ same evaluation checkpoint provenance recorded for each downstream run. For concrete rerun outputs, use the exported result artifacts under `results/` and the thesis W&B project summaries. +Final thesis reruns should use the tmux launcher so W&B scoping stays +consistent: + +`WANDB_ENTITY= bash scripts/internal/launch_thesis_tmux.sh` + +If launching a sprint manually, include the thesis W&B target on every training +command: + +`--project slices-thesis --revision thesis-v1 --entity ` + ### Upcoming: Sprint 10 — Extra Seeds 789, 1011 (630 runs) **Motivation**: Go from 3 seeds to 5 seeds for all SSL and supervised experiments, improving statistical power (4 d.f. instead of 2 for variance estimates). @@ -365,7 +375,7 @@ and the thesis W&B project summaries. **Breakdown**: 48 pretrain (~30 min each, the bottleneck) + 510 finetune/probe (~3–5 min each) + 72 supervised (~3–5 min each). -**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 10 --parallel 4` (limited by GPU-bound pretraining) +**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 10 --parallel 4 --project slices-thesis --revision thesis-v1 --entity ` (limited by GPU-bound pretraining) ### Upcoming: Sprint 7p — Focused Capacity Study (100 runs) @@ -375,7 +385,7 @@ and the thesis W&B project summaries. **Comparison set**: Default-size baselines are inherited from Sprint `6` (seeds 42/123/456) and Sprint `10` (seeds 789/1011), so Sprint `7p` only launches the larger-capacity runs. -**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 7p --parallel 4` +**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 7p --parallel 4 --project slices-thesis --revision thesis-v1 --entity ` ### Upcoming: Sprint 11 — Classical Baselines (360 runs) @@ -395,7 +405,7 @@ and the thesis W&B project summaries. 3. XGBoost cannot transfer across datasets — highlights SSL's cross-dataset transfer capability (Sprint 7) 4. XGBoost/GRU-D on label efficiency curves — direct comparison with SSL learning curves -**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 11 --parallel 12` (all runs independent, no pretraining dependency) +**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 11 --parallel 12 --project slices-thesis --revision thesis-v1 --entity ` (all runs independent, no pretraining dependency) ### Upcoming: Sprint 13 — TS2Vec Temporal Contrastive Extension (135 runs) @@ -407,7 +417,7 @@ and the thesis W&B project summaries. **Interpretation**: Sprint `13` is a formal thesis extension, not a replacement for the controlled MAE/JEPA/Contrastive triangle. It tests whether a better-instantiated contrastive family changes the conclusion. -**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 13 --parallel 4` +**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 13 --parallel 4 --project slices-thesis --revision thesis-v1 --entity ` ### Upcoming: Sprint 12 — SMART External SSL Reference (135 runs) @@ -417,14 +427,14 @@ and the thesis W&B project summaries. **Matrix**: 1 paradigm × 3 datasets × 4 tasks × 2 protocols × 5 seeds = 120 finetune + 15 pretrain = **135 runs** -**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 12 --parallel 4` +**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 12 --parallel 4 --project slices-thesis --revision thesis-v1 --entity ` ### Upcoming: Sprint 9 — Fairness Analysis (0 extra runs) **Runs last**, after all training sprints are complete. Zero additional training — pure evaluation on existing test predictions. 1. Compute fairness metrics on all downstream test predictions from Sprints 1–8, 7p, 10, 11, 12, and 13 -2. Run `scripts/eval/evaluate_fairness.py` with an explicit `--revision` tag for the thesis rerun corpus +2. Run `scripts/eval/evaluate_fairness.py --project slices-thesis --revision thesis-v1 --entity ` for the thesis rerun corpus 3. Generate fairness tables and disparity plots 4. Protected attributes: sex (all datasets), age group (all), race/ethnicity (MIMIC-IV only) 5. The default standalone fairness sweep includes Sprint `11` classical baselines via `phase:baseline` @@ -435,7 +445,7 @@ and the thesis W&B project summaries. ### Project Structure -Legacy exploratory runs remain in W&B project `slices`. Final thesis reruns should log to a separate project such as `slices-thesis` and use a single revision tag such as `thesis-v1`. +Legacy exploratory runs remain in W&B project `slices`. Final thesis reruns should log to `slices-thesis`, use the single revision tag `thesis-v1`, and set the W&B entity explicitly. ### Run Naming Convention @@ -476,7 +486,7 @@ parsing run names. ### Baseline Inheritance Across Sprints -Later sprints reuse runs from earlier sprints as comparison baselines. To enable filtering all relevant runs for a sprint in a single W&B query, baseline runs are tagged with later sprint tags using `run_experiments.py tag --sprint N`. +Later sprints reuse runs from earlier sprints as comparison baselines. To enable filtering all relevant runs for a sprint in a single W&B query, baseline runs are tagged with later sprint tags using `run_experiments.py tag --sprint N --project slices-thesis --entity `. | Sprint | Inherits From | What's Inherited | |--------|--------------|------------------| @@ -497,7 +507,7 @@ Later sprints reuse runs from earlier sprints as comparison baselines. To enable | **12** | — | Self-contained (SMART pretrains + finetune, 5 seeds) | | **9** | All | Evaluates test predictions from all training sprints (runs last) | -**Usage**: `uv run python scripts/internal/run_experiments.py tag --sprint N` (idempotent). +**Usage**: `uv run python scripts/internal/run_experiments.py tag --sprint N --project slices-thesis --entity ` (idempotent). --- diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index 1061c16..9828305 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -12,25 +12,30 @@ Usage: # Evaluate the benchmark fairness corpus for one explicit revision - uv run python scripts/eval/evaluate_fairness.py --revision benchmark-v1 + uv run python scripts/eval/evaluate_fairness.py \ + --project slices-thesis --revision thesis-v1 --entity # Scope to specific sprint/dataset uv run python scripts/eval/evaluate_fairness.py \ - --revision benchmark-v1 --sprint 1 --dataset miiv + --project slices-thesis --revision thesis-v1 --entity \ + --sprint 1 --dataset miiv # Preview which runs would be evaluated - uv run python scripts/eval/evaluate_fairness.py --revision benchmark-v1 --dry-run + 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 \ - --revision benchmark-v1 \ + --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 --revision benchmark-v1 --force + 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 --revision benchmark-v1 --max-runs 1 + uv run python scripts/eval/evaluate_fairness.py \ + --project slices-thesis --revision thesis-v1 --entity --max-runs 1 """ from __future__ import annotations diff --git a/scripts/export_results.py b/scripts/export_results.py index 292816a..113db14 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -14,11 +14,17 @@ Both files include wandb run IDs for traceability back to W&B. Usage: - uv run python scripts/export_results.py - uv run python scripts/export_results.py --sprint 1 --validate-seeds 3 - uv run python scripts/export_results.py --paradigm mae jepa --dataset miiv - uv run python scripts/export_results.py --output-dir results/sprint1 --sprint 1 - uv run python scripts/export_results.py --project slices-benchmark --revision benchmark-v1 + 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 \ + --sprint 1 --validate-seeds 3 + uv run python scripts/export_results.py \ + --project slices-thesis --revision thesis-v1 --entity \ + --paradigm mae jepa --dataset miiv + uv run python scripts/export_results.py \ + --project slices-thesis --revision thesis-v1 --entity \ + --output-dir results/sprint1 --sprint 1 """ from __future__ import annotations @@ -33,14 +39,13 @@ from pathlib import Path import pandas as pd +import wandb from slices.eval.statistical import ( bonferroni_correction, cohens_d, paired_wilcoxon_signed_rank, ) -import wandb - # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- diff --git a/scripts/internal/launch_thesis_tmux.sh b/scripts/internal/launch_thesis_tmux.sh index 0f929a7..2c11152 100755 --- a/scripts/internal/launch_thesis_tmux.sh +++ b/scripts/internal/launch_thesis_tmux.sh @@ -29,6 +29,11 @@ if tmux has-session -t "$SESSION_NAME" 2>/dev/null; then exit 1 fi +if [[ -z "$WANDB_ENTITY" ]]; then + echo "WANDB_ENTITY must be set for thesis runs." >&2 + exit 1 +fi + mkdir -p "$LOG_DIR" main_sprints=(1 1b 1c 2 3 4 5 6 7 8 10 11) @@ -61,17 +66,10 @@ quote_cmd() { printf "%q " "$@" } -run_args=(--project "$WANDB_PROJECT") -export_args=(--project "$WANDB_PROJECT") -fairness_args=(--project "$WANDB_PROJECT") -tag_args=(--project "$WANDB_PROJECT") - -if [[ -n "$WANDB_ENTITY" ]]; then - run_args+=(--entity "$WANDB_ENTITY") - export_args+=(--entity "$WANDB_ENTITY") - fairness_args+=(--entity "$WANDB_ENTITY") - tag_args+=(--entity "$WANDB_ENTITY") -fi +run_args=(--project "$WANDB_PROJECT" --entity "$WANDB_ENTITY") +export_args=(--project "$WANDB_PROJECT" --entity "$WANDB_ENTITY") +fairness_args=(--project "$WANDB_PROJECT" --entity "$WANDB_ENTITY") +tag_args=(--project "$WANDB_PROJECT" --entity "$WANDB_ENTITY") run_args+=(--revision "$REVISION" --reason "$REASON") export_args+=(--revision "$REVISION" --output-dir "$RESULTS_DIR") diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index d22b564..fbd5c85 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -7,18 +7,25 @@ Usage: uv run python scripts/internal/run_experiments.py warmup --sprint 1 - uv run python scripts/internal/run_experiments.py run --sprint 1 --parallel 4 - uv run python scripts/internal/run_experiments.py run --sprint 1 2 3 --parallel 6 --dry-run - uv run python scripts/internal/run_experiments.py run --sprint 1 --revision v2 --reason "fix LR" + uv run python scripts/internal/run_experiments.py run --sprint 1 --parallel 4 \\ + --project slices-thesis --revision thesis-v1 --entity + uv run python scripts/internal/run_experiments.py run --sprint 1 2 3 --parallel 6 \\ + --dry-run --project slices-thesis --revision thesis-v1 --entity + uv run python scripts/internal/run_experiments.py run --sprint 1 \\ + --project slices-thesis --revision thesis-v1 --entity --reason "fix LR" uv run python scripts/internal/run_experiments.py run --sprint 1 2 \\ - --project slices-thesis --entity myteam + --project slices-thesis --revision thesis-v1 --entity uv run python scripts/internal/run_experiments.py status uv run python scripts/internal/run_experiments.py status --sprint 1 - uv run python scripts/internal/run_experiments.py retry --failed --parallel 4 - uv run python scripts/internal/run_experiments.py retry --failed \ - --sprint 1 --revision v2 --parallel 4 - uv run python scripts/internal/run_experiments.py tag --sprint 2 --dry-run - uv run python scripts/internal/run_experiments.py tag --sprint 2 3 5 + uv run python scripts/internal/run_experiments.py retry --failed --parallel 4 \\ + --project slices-thesis --revision thesis-v1 --entity + uv run python scripts/internal/run_experiments.py retry --failed \\ + --sprint 1 --revision thesis-v1 --parallel 4 \\ + --project slices-thesis --entity + uv run python scripts/internal/run_experiments.py tag --sprint 2 --dry-run \\ + --project slices-thesis --entity + uv run python scripts/internal/run_experiments.py tag --sprint 2 3 5 \\ + --project slices-thesis --entity """ from __future__ import annotations @@ -1083,7 +1090,7 @@ def _select_ready_runs( # --------------------------------------------------------------------------- # Scheduler # --------------------------------------------------------------------------- -def run_scheduler(runs: list[Run], state: dict, parallel: int, dry_run: bool): +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}") @@ -1124,7 +1131,7 @@ def handle_signal(signum, frame): if dry_run: _print_dry_run(runs, runs_by_id) - return + return 0 print(f"Scheduler started (slot_budget={parallel})") print(f"State file: {STATE_FILE}") @@ -1218,6 +1225,7 @@ def handle_signal(signum, frame): # Print summary print() _print_summary(runs, state) + return _scheduler_exit_code(runs, state) def _format_elapsed(state: dict, run_id: str, now_iso: str) -> str: @@ -1367,6 +1375,21 @@ def _print_summary(runs: list[Run], state: dict): ) +def _scheduler_exit_code(runs: list[Run], state: dict) -> int: + """Return nonzero when the requested run set did not complete cleanly.""" + failed = [r.id for r in runs if get_run_status(state, r.id) == "failed"] + skipped = [r.id for r in runs if get_run_status(state, r.id) == "skipped"] + + if failed or skipped: + print( + "Scheduler incomplete: " + f"{len(failed)} failed, {len(skipped)} skipped due to dependency failure" + ) + return 1 + + return 0 + + def _sprint_sort_key(s: str) -> tuple[int, str]: """Sort sprints: 1, 1b, 2, 3, ...""" num = "" @@ -1407,7 +1430,7 @@ def cmd_run(args): if not runs: print(f"No runs found for sprint(s): {', '.join(sprints)}") - return + return 0 # Apply revision transform if requested if args.revision: @@ -1420,7 +1443,7 @@ def cmd_run(args): if args.project or args.entity: target = f"{args.entity}/{args.project}" if args.entity else args.project print(f"W&B target: {target}") - run_scheduler(runs, state, args.parallel, args.dry_run) + return run_scheduler(runs, state, args.parallel, args.dry_run) def cmd_status(args): @@ -1469,7 +1492,7 @@ def cmd_retry(args): if not runs_to_retry: print("No runs to retry.") - return + return 0 # Include their dependencies that are completed (already fine) or pending all_by_id = {r.id: r for r in all_runs} @@ -1487,7 +1510,7 @@ def cmd_retry(args): if args.project or args.entity: target = f"{args.entity}/{args.project}" if args.entity else args.project print(f"W&B target: {target}") - run_scheduler(runs_to_retry, state, args.parallel, dry_run=False) + return run_scheduler(runs_to_retry, state, args.parallel, dry_run=False) def cmd_warmup(args): @@ -1569,10 +1592,12 @@ def cmd_tag(args): W&B API. Idempotent — already-tagged runs are skipped. Usage: - uv run python scripts/internal/run_experiments.py tag --sprint 2 - uv run python scripts/internal/run_experiments.py tag --sprint 2 3 --dry-run - uv run python scripts/internal/run_experiments.py tag --sprint 2 \ - --project slices --entity myteam + uv run python scripts/internal/run_experiments.py tag --sprint 2 \\ + --project slices-thesis --entity + uv run python scripts/internal/run_experiments.py tag --sprint 2 3 --dry-run \\ + --project slices-thesis --entity + uv run python scripts/internal/run_experiments.py tag --sprint 2 \\ + --project slices-thesis --entity """ try: import wandb # noqa: F811 @@ -1737,15 +1762,19 @@ def main(): parser.error("--revision requires --sprint to scope which sprints to revise") if args.command == "run": - cmd_run(args) + exit_code = cmd_run(args) elif args.command == "status": - cmd_status(args) + exit_code = cmd_status(args) elif args.command == "retry": - cmd_retry(args) + exit_code = cmd_retry(args) elif args.command == "warmup": - cmd_warmup(args) + exit_code = cmd_warmup(args) elif args.command == "tag": - cmd_tag(args) + exit_code = cmd_tag(args) + else: + exit_code = 0 + + sys.exit(exit_code or 0) if __name__ == "__main__": diff --git a/src/slices/data/datamodule.py b/src/slices/data/datamodule.py index 19e612c..9c68c25 100644 --- a/src/slices/data/datamodule.py +++ b/src/slices/data/datamodule.py @@ -25,7 +25,7 @@ from slices.data.sliding_window import SlidingWindowDataset from slices.data.splits import ( compute_patient_level_splits, - save_split_info, + save_global_split_info, subsample_train_indices, ) @@ -256,17 +256,16 @@ def setup(self, stage: Optional[str] = None) -> None: logger.debug("[Step 3/3] Preserving canonical split information") splits_path = self.processed_dir / "splits.yaml" if not splits_path.exists(): - save_split_info( - self.processed_dir, - self.dataset, - self.full_train_indices, - self.val_indices, - self.test_indices, - self.seed, - self.train_ratio, - self.val_ratio, - self.test_ratio, - train_subset_indices=self.train_indices, + 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: diff --git a/src/slices/data/splits.py b/src/slices/data/splits.py index 971d802..9f2d880 100644 --- a/src/slices/data/splits.py +++ b/src/slices/data/splits.py @@ -23,20 +23,27 @@ def load_cached_splits( 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 current task-filtered cohort. + 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 in order from timeseries parquet. + 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, @@ -83,10 +90,31 @@ def load_cached_splits( zip(static_df["stay_id"].to_list(), static_df["patient_id"].to_list()) ) - # Validate cached patient lists against the current cohort represented by - # stay_ids. For supervised tasks this may be a label-filtered subset, so - # validating against the full static table would incorrectly reuse - # full-cohort cached splits for a smaller task-specific cohort. + # 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) @@ -96,11 +124,9 @@ def load_cached_splits( ) return None current_patients.add(patient_id) - cached_patients = train_patients | val_patients | test_patients - - if current_patients != cached_patients: + if not current_patients.issubset(cached_patients): logger.debug( - f"Cached splits have different patients " + f"Cached splits do not cover current patients " f"(cached: {len(cached_patients)}, current: {len(current_patients)}), " f"recomputing" ) @@ -141,8 +167,9 @@ def filter_stays_with_missing_labels( ) -> 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. + 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. @@ -190,6 +217,126 @@ def filter_stays_with_missing_labels( 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], @@ -202,15 +349,17 @@ def compute_patient_level_splits( ]: """Compute 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. + 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: Filters out stays with missing labels BEFORE computing splits - to ensure indices remain consistent with the Dataset after it loads. + 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. @@ -245,8 +394,9 @@ def compute_patient_level_splits( logger.debug(f"Loading labels from {labels_path.name}") 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 + # 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 ) @@ -254,7 +404,14 @@ def compute_patient_level_splits( # 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 + 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 @@ -269,44 +426,39 @@ def compute_patient_level_splits( excluded_stay_ids, ) - logger.info("Computing patient-level splits...") + logger.info("Computing full-cohort patient-level splits...") - # Get stay_id -> patient_id mapping (only for filtered stays) + # Get stay_id -> patient_id mapping logger.debug("Building stay-to-patient 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) - filtered_patient_ids = {stay_to_patient[sid] for sid in stay_ids if sid in stay_to_patient} - unique_patients = sorted(filtered_patient_ids) + # 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(stay_ids) and n_patients > 0: - filtered_stay_set = set(stay_ids) - if filtered_patient_ids == filtered_stay_set: + 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 patients deterministically using seed + # Shuffle full-cohort patients deterministically using seed logger.debug(f"Shuffling patients (seed={seed})") - 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 :]) + 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" @@ -338,30 +490,17 @@ def compute_patient_level_splits( 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." - ) + # 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, + ) # Final validation: all stays should be accounted for total_assigned = len(train_indices) + len(val_indices) + len(test_indices) diff --git a/src/slices/models/pretraining/mae.py b/src/slices/models/pretraining/mae.py index 93739fd..2926d0d 100644 --- a/src/slices/models/pretraining/mae.py +++ b/src/slices/models/pretraining/mae.py @@ -137,8 +137,18 @@ def forward( full_tokens = full_tokens + self.time_pe[timestep_idx] full_tokens = self.embed_dropout(full_tokens) - # Run decoder transformer over the full reconstruction grid. - decoded = self.decoder(full_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, + ) + + # 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 D features per timestep predictions = self.output_proj(decoded) # (B, T, D) diff --git a/tests/test_dataset_datamodule.py b/tests/test_dataset_datamodule.py index 431ec78..26c4cd6 100644 --- a/tests/test_dataset_datamodule.py +++ b/tests/test_dataset_datamodule.py @@ -10,7 +10,70 @@ import yaml from slices.data.datamodule import ICUDataModule, icu_collate_fn from slices.data.dataset import ICUDataset -from slices.data.splits import load_cached_splits +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 @@ -978,13 +1041,11 @@ 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_splits_are_invalidated_when_task_filter_removes_patients( - self, mock_extracted_data - ): - """Task-filtered cohorts must not reuse a full-cohort cached split.""" + 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="mortality_24h", + task_name=None, seed=42, ) dm.setup() @@ -1010,6 +1071,10 @@ def test_cached_splits_are_invalidated_when_task_filter_removes_patients( 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( @@ -1020,9 +1085,95 @@ def test_cached_splits_are_invalidated_when_task_filter_removes_patients( train_ratio=0.7, val_ratio=0.15, test_ratio=0.15, + cohort_stay_ids=all_stay_ids, ) - assert cached is None + 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_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.""" diff --git a/tests/test_fixes.py b/tests/test_fixes.py index a6e0b73..bf8eb44 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -857,6 +857,7 @@ def test_tmux_launcher_passes_revision_to_fairness(self, tmp_path): 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" @@ -875,6 +876,8 @@ def test_tmux_launcher_passes_revision_to_fairness(self, tmp_path): 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 def test_tmux_launcher_includes_sprint11_in_fairness(self, tmp_path): """The fairness sweep should include the canonical baseline sprint by default.""" @@ -904,6 +907,7 @@ def test_tmux_launcher_includes_sprint11_in_fairness(self, tmp_path): 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" result = subprocess.run( @@ -949,6 +953,7 @@ def test_tmux_status_pane_uses_uv_run_python(self, tmp_path): 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" result = subprocess.run( @@ -1860,6 +1865,45 @@ def test_select_ready_runs_does_not_mix_pretrain_with_active_gpu_work(self): 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", + sprint="1", + run_type="supervised", + paradigm="supervised", + dataset="miiv", + seed=42, + output_dir="outputs/sprint1/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_exit_code_nonzero_for_dependency_skips(self): + import scripts.internal.run_experiments as runner + + skipped_run = runner.Run( + id="skipped", + sprint="1", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/sprint1/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.""" diff --git a/tests/test_mae_objective.py b/tests/test_mae_objective.py index f32b458..6356fdd 100644 --- a/tests/test_mae_objective.py +++ b/tests/test_mae_objective.py @@ -174,6 +174,40 @@ def test_decoder_ignores_padded_visible_tokens(self): 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 From f28f09849d3b19147f516cd4e98f393f323b93bf Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 21 Apr 2026 18:59:56 -0400 Subject: [PATCH 091/121] Fix thesis export statistical safeguards Replace the custom Wilcoxon approximation with SciPy-backed inference, classify inherited exports from requested sprint tags, and add exact seed plus paired task/seed coverage validation. Add TS2Vec-vs-core contrastive export coverage, XGBoost calibration metrics, canonical GRU-D decay behavior, and downstream obs-aware Transformer empty-timestep masking. GRU-D hidden decay parameters now map feature-wise deltas to hidden units, so existing GRU-D checkpoints should be rerun before reporting. --- scripts/export_results.py | 407 ++++++++++++++++++++----- scripts/training/xgboost_baseline.py | 26 ++ src/slices/eval/statistical.py | 82 ++--- src/slices/models/encoders/gru_d.py | 54 +++- src/slices/training/finetune_module.py | 17 +- tests/test_export_results.py | 137 +++++++++ tests/test_finetune_module.py | 24 ++ tests/test_fixes.py | 8 + tests/test_gru_d_encoder.py | 36 +++ tests/test_statistical.py | 7 + 10 files changed, 658 insertions(+), 140 deletions(-) create mode 100644 tests/test_gru_d_encoder.py diff --git a/scripts/export_results.py b/scripts/export_results.py index 113db14..41d8c78 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -3,13 +3,15 @@ Canonical results export pipeline for SLICES. Pulls all experiment runs from W&B, extracts config + test metrics, and produces -three structured parquet files: +structured parquet files: - results/per_seed_results.parquet — one row per W&B run (~3000 rows) - results/aggregated_results.parquet — one row per unique config (~600 rows), with mean/std/min/max across seeds - results/statistical_tests.parquet — pairwise Wilcoxon + Bonferroni + Cohen's d significance table + - results/ts2vec_vs_core_contrastive.parquet — direct TS2Vec vs core + contrastive comparisons Both files include wandb run IDs for traceability back to W&B. @@ -18,7 +20,7 @@ --project slices-thesis --revision thesis-v1 --entity uv run python scripts/export_results.py \ --project slices-thesis --revision thesis-v1 --entity \ - --sprint 1 --validate-seeds 3 + --sprint 1 --expected-seeds 42 123 456 uv run python scripts/export_results.py \ --project slices-thesis --revision thesis-v1 --entity \ --paradigm mae jepa --dataset miiv @@ -40,6 +42,7 @@ import pandas as pd import wandb + from slices.eval.statistical import ( bonferroni_correction, cohens_d, @@ -87,6 +90,7 @@ "smart_reference", "temporal_contrastive", } +EXPECTED_FIXED_SEEDS = {42, 123, 456, 789, 1011} # For core runs: the sprint tag is not part of the config identity. # lr/mask_ratio excluded because they're determined by protocol (redundant) or @@ -381,6 +385,36 @@ def _recover_source_dataset(run_name: str, config: dict) -> str | None: return None +def _sprint_tags(tags: list[str]) -> list[str]: + """Return sprint values present in W&B tags, preserving tag order.""" + return [tag.split(":", 1)[1] for tag in tags if tag.startswith("sprint:")] + + +def _select_export_sprint( + config_sprint: str, + tags: list[str], + requested_sprints: list[str] | None = None, +) -> str: + """Choose the sprint family used for export classification. + + Historical inherited baselines keep their original `config.sprint`, but the + runner adds target sprint tags. When the export is scoped to target sprint + tags, classify the run in that requested target family. + """ + if not requested_sprints: + return config_sprint + + tagged_sprints = set(_sprint_tags(tags)) + requested_matches = [ + str(sprint) for sprint in requested_sprints if str(sprint) in tagged_sprints + ] + if not requested_matches: + return config_sprint + if config_sprint in requested_matches: + return config_sprint + return requested_matches[0] + + def _recover_pretrain_metadata( run_name: str, config: dict ) -> tuple[float | None, float | None, str | None]: @@ -441,7 +475,11 @@ def _infer_model_size(config: dict) -> str: return f"d{d_model}_L{n_layers}" -def extract_run(run, metric_keys: list[str]) -> dict: +def extract_run( + run, + metric_keys: list[str], + requested_sprints: list[str] | None = None, +) -> dict: """Extract config + metrics from a W&B run in a single retry-protected call.""" config, summary, run_id, run_url, run_name, tags, group, created_at = _retry( lambda: _load_run_data(run) @@ -489,7 +527,8 @@ def extract_run(run, metric_keys: list[str]) -> dict: else: metrics[key] = float("nan") - sprint_str = str(config.get("sprint", "")) + config_sprint = str(config.get("sprint", "")) + sprint_str = _select_export_sprint(config_sprint, tags, requested_sprints) experiment_type = EXPERIMENT_TYPES.get(sprint_str, "unknown") # Historical group 10 contains both core (label_fraction=1.0) and @@ -521,6 +560,8 @@ def extract_run(run, metric_keys: list[str]) -> dict: "created_at": created_at, "experiment_type": experiment_type, "sprint": sprint_str, + "config_sprint": config_sprint, + "sprint_tags": json.dumps(_sprint_tags(tags)), "paradigm": config.get("paradigm", None), "dataset": config.get("dataset", None), "task": _get_nested(config, "task.task_name", default=None), @@ -617,7 +658,10 @@ def _assert_no_ambiguous_fingerprint_collisions(df: pd.DataFrame) -> None: ) -def build_per_seed_df(runs: list) -> pd.DataFrame: +def build_per_seed_df( + runs: list, + requested_sprints: list[str] | None = None, +) -> pd.DataFrame: """Build the per-seed DataFrame from raw W&B runs. One row per run with config columns + metric columns + wandb IDs. @@ -628,7 +672,7 @@ def build_per_seed_df(runs: list) -> pd.DataFrame: if (i + 1) % 100 == 0: print(f" Processing run {i + 1}/{len(runs)}...", file=sys.stderr) try: - row = extract_run(run, ALL_METRICS) + row = extract_run(run, ALL_METRICS, requested_sprints=requested_sprints) rows.append(row) except Exception as e: run_id = getattr(run, "id", "unknown") @@ -806,6 +850,109 @@ def _primary_metric_for_task(task_name: str | None) -> str | None: return "test/auprc" +def _format_task_seed_pairs(keys: set[tuple[str, int]]) -> str: + """Serialize task/seed coverage gaps for audit columns.""" + formatted = [{"task": task, "seed": int(seed)} for task, seed in sorted(keys)] + return json.dumps(formatted) + + +def _metric_pairs_for_paradigm( + scope_group: pd.DataFrame, + paradigm: str, + value_col: str = "primary_metric_value", +) -> pd.DataFrame: + """Return finite task/seed metric rows for one paradigm.""" + 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]]: + """Build paired metric frame and task/seed coverage summary.""" + 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: + """Compute paired Wilcoxon/effect-size row for a comparison.""" + 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) + median_improvement = float(pd.Series(improvement).median()) + better_paradigm = ( + paradigm_a if mean_improvement > 0 else paradigm_b if mean_improvement < 0 else None + ) + + 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": median_improvement, + "better_paradigm": better_paradigm, + "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: """Add synthetic statistic scopes for classical-vs-neural comparisons.""" required = {"experiment_type", "protocol", "label_fraction", "paradigm"} @@ -913,59 +1060,18 @@ def build_statistical_tests_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: if a_classical == b_classical: continue - pairs_a = scope_group[scope_group["paradigm"] == paradigm_a][ - ["task", "seed", "primary_metric_value"] - ].rename(columns={"primary_metric_value": "value_a"}) - pairs_b = scope_group[scope_group["paradigm"] == paradigm_b][ - ["task", "seed", "primary_metric_value"] - ].rename(columns={"primary_metric_value": "value_b"}) - - paired = pairs_a.merge(pairs_b, on=["task", "seed"], how="inner") - paired = paired.dropna(subset=["value_a", "value_b"]) + paired, coverage = _paired_metric_frame(scope_group, paradigm_a, paradigm_b) if paired.empty: continue - 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) - median_improvement = float(pd.Series(improvement).median()) - better_paradigm = ( - paradigm_a - if mean_improvement > 0 - else paradigm_b if mean_improvement < 0 else None + row = _paired_stat_row( + scope=scope, + paired=paired, + coverage=coverage, + paradigm_a=paradigm_a, + paradigm_b=paradigm_b, + higher_is_better=higher_is_better, ) - - row = { - **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()) - ), - "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": median_improvement, - "better_paradigm": better_paradigm, - "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 - ), - } family_rows.append(row) if not family_rows: @@ -1002,6 +1108,105 @@ def build_statistical_tests_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: return stats_df +def build_ts2vec_vs_core_contrastive_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: + """Compare Sprint-13 TS2Vec directly against core contrastive runs.""" + required = {"experiment_type", "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_type"] == "temporal_contrastive") + & (per_seed_df["paradigm"] == "ts2vec") + ) + | ((per_seed_df["experiment_type"] == "core") & (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 + ] + work = _fillna_for_grouping(subset, scope_cols) + rows = [] + + 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 + + paradigms = set(scope_group["paradigm"].dropna()) + if not {"contrastive", "ts2vec"}.issubset(paradigms): + 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 + row = _paired_stat_row( + scope={ + **scope, + "comparison_type": "ts2vec_vs_core_contrastive", + "experiment_type_a": "temporal_contrastive", + "experiment_type_b": "core", + }, + paired=paired, + coverage=coverage, + paradigm_a="ts2vec", + paradigm_b="contrastive", + higher_is_better=higher_is_better, + ) + rows.append(row) + + 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) + + result = pd.DataFrame(rows) + sort_cols = [ + col + for col in [ + "dataset", + "protocol", + "label_fraction", + "model_size", + "primary_metric_name", + ] + if col in result.columns + ] + if sort_cols: + result = result.sort_values(sort_cols, na_position="last").reset_index(drop=True) + return result + + # --------------------------------------------------------------------------- # Validation # --------------------------------------------------------------------------- @@ -1010,36 +1215,75 @@ def build_statistical_tests_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: def validate( per_seed_df: pd.DataFrame, aggregated_df: pd.DataFrame, - expected_core_seeds: int = 5, + statistical_df: pd.DataFrame | None = None, + expected_seeds: set[int] | None = None, ) -> list[str]: """Validate results and return warning strings. - Warns when fixed-seed experiment families have fewer seeds than expected. + Warns when fixed-seed experiment families do not have the exact expected + seed set, and when statistical comparisons have incomplete task/seed pairing. """ warnings = [] + expected_seeds = set(EXPECTED_FIXED_SEEDS if expected_seeds is None else expected_seeds) - # Check fixed-seed experiment families for expected seed count + # Check fixed-seed experiment families for exact expected seed set. if "experiment_type" in aggregated_df.columns: fixed_seed = aggregated_df[ aggregated_df["experiment_type"].isin(FIXED_SEED_EXPERIMENT_TYPES) ] - low_seed = fixed_seed[fixed_seed["n_seeds"] < expected_core_seeds] - if len(low_seed) > 0: + 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(low_seed)}/{len(fixed_seed)} fixed-seed configs have fewer " - f"than {expected_core_seeds} seeds:" + 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 in low_seed.iterrows(): + for row, seeds in bad_seed_rows: desc = ", ".join( f"{c}={row[c]}" for c in ["paradigm", "dataset", "task", "protocol"] if c in row and row[c] is not None ) - seeds = json.loads(row["seed_list"]) if pd.notna(row.get("seed_list")) else [] warnings.append( f" experiment_type={row['experiment_type']}, {desc} — " - f"n_seeds={row['n_seeds']}, seeds={seeds}" + 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:" ) + for _, row in incomplete.head(10).iterrows(): + scope = ", ".join( + f"{c}={row[c]}" + for c in [ + "experiment_type", + "dataset", + "protocol", + "label_fraction", + "primary_metric_name", + "paradigm_a", + "paradigm_b", + ] + if c in row and row[c] is not None + ) + warnings.append( + f" {scope} — shared={row['n_shared_task_seed_pairs']}, " + f"union={row['n_union_task_seed_pairs']}" + ) # Check for runs with no test metrics at all test_cols = [c for c in TEST_METRICS if c in per_seed_df.columns] @@ -1125,10 +1369,14 @@ def parse_args() -> argparse.Namespace: help="Output directory (default: results/)", ) parser.add_argument( - "--expected-core-seeds", + "--expected-seeds", + nargs="+", type=int, - default=5, - help="Expected seed count for core configs (default: 5)", + 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", @@ -1173,7 +1421,7 @@ def main(): # Build DataFrames print(f"\nBuilding per-seed DataFrame from {len(runs)} runs...", file=sys.stderr) - per_seed_df = build_per_seed_df(runs) + per_seed_df = build_per_seed_df(runs, requested_sprints=args.sprint) print(f" Shape: {per_seed_df.shape}", file=sys.stderr) print("\nBuilding aggregated DataFrame...", file=sys.stderr) @@ -1184,8 +1432,21 @@ def main(): 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) + # Validate - warnings = validate(per_seed_df, aggregated_df, expected_core_seeds=args.expected_core_seeds) + 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), + ) for w in warnings: print(w, file=sys.stderr) @@ -1193,21 +1454,25 @@ def main(): per_seed_path = output_dir / "per_seed_results.parquet" aggregated_path = output_dir / "aggregated_results.parquet" statistical_path = output_dir / "statistical_tests.parquet" + ts2vec_contrastive_path = output_dir / "ts2vec_vs_core_contrastive.parquet" per_seed_df.to_parquet(per_seed_path, index=False) aggregated_df.to_parquet(aggregated_path, index=False) statistical_df.to_parquet(statistical_path, index=False) + ts2vec_contrastive_df.to_parquet(ts2vec_contrastive_path, index=False) print("\nSaved:", file=sys.stderr) print(f" {per_seed_path} ({len(per_seed_df)} rows)", file=sys.stderr) print(f" {aggregated_path} ({len(aggregated_df)} rows)", file=sys.stderr) print(f" {statistical_path} ({len(statistical_df)} rows)", file=sys.stderr) + print(f" {ts2vec_contrastive_path} ({len(ts2vec_contrastive_df)} rows)", file=sys.stderr) # Print quick summary to stdout for piping print("\n--- Quick Summary ---") print( f"Runs: {len(per_seed_df)}, Configs: {len(aggregated_df)}, " - f"StatTests: {len(statistical_df)}" + f"StatTests: {len(statistical_df)}, " + f"TS2VecContrastiveTests: {len(ts2vec_contrastive_df)}" ) if "experiment_type" in aggregated_df.columns: for etype, group in aggregated_df.groupby("experiment_type"): diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index b26fec8..884762f 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -13,11 +13,13 @@ 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, mean_absolute_error, mean_squared_error, r2_score, @@ -33,6 +35,28 @@ def _xgboost_eval_metric(task_type: str) -> str: return "mae" if task_type == "regression" else "aucpr" +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 _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", [])) @@ -191,6 +215,8 @@ def main(cfg: DictConfig) -> None: metrics["test/auroc"] = roc_auc_score(y_test, y_pred_proba) metrics["test/auprc"] = average_precision_score(y_test, y_pred_proba) metrics["test/accuracy"] = accuracy_score(y_test, y_pred_class) + metrics["test/brier_score"] = brier_score_loss(y_test, y_pred_proba) + metrics["test/ece"] = _binary_ece(y_test, y_pred_proba) print("\n Test Results:") for key, value in metrics.items(): diff --git a/src/slices/eval/statistical.py b/src/slices/eval/statistical.py index a80d4d8..98c3ac3 100644 --- a/src/slices/eval/statistical.py +++ b/src/slices/eval/statistical.py @@ -165,12 +165,10 @@ def paired_wilcoxon_signed_rank( values_b: Sequence[float], correction: bool = True, ) -> Dict[str, float]: - """Paired Wilcoxon signed-rank test with tie correction. + """Paired Wilcoxon signed-rank test. - This implementation uses the standard large-sample normal approximation - with tie correction. It is sufficient for the benchmark export pipeline where - comparisons pool multiple tasks x seeds and are primarily used for - significance tables rather than exact small-sample inference. + 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. @@ -204,32 +202,30 @@ def paired_wilcoxon_signed_rank( "n_nonzero_pairs": 0.0, } - abs_diffs = [abs(diff) for diff in nonzero_diffs] - ranks, tie_counts = _average_ranks(abs_diffs) - - w_plus = sum(rank for rank, diff in zip(ranks, nonzero_diffs) if diff > 0.0) - w_minus = sum(rank for rank, diff in zip(ranks, nonzero_diffs) if diff < 0.0) - statistic = min(w_plus, w_minus) - - expected = n_nonzero * (n_nonzero + 1) / 4.0 - variance = n_nonzero * (n_nonzero + 1) * (2 * n_nonzero + 1) / 24.0 - tie_correction = sum(t * (t + 1) * (2 * t + 1) for t in tie_counts) / 48.0 - variance -= tie_correction - - if variance <= 0.0: - z_score = 0.0 - p_value = 1.0 - else: - numerator = abs(w_plus - expected) - if correction: - numerator = max(numerator - 0.5, 0.0) - z_score = numerator / math.sqrt(variance) - p_value = min(2.0 * _normal_sf(z_score), 1.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(statistic), + "statistic": float(result.statistic), "z_score": float(z_score), - "p_value": float(p_value), + "p_value": float(result.pvalue), "n_pairs": float(n_pairs), "n_nonzero_pairs": float(n_nonzero), } @@ -322,32 +318,6 @@ def _finite_pairs( return pairs -def _average_ranks(values: Sequence[float]) -> tuple[list[float], list[int]]: - indexed = sorted(enumerate(values), key=lambda item: item[1]) - ranks = [0.0] * len(values) - tie_counts: list[int] = [] - - i = 0 - while i < len(indexed): - j = i + 1 - while j < len(indexed) and indexed[j][1] == indexed[i][1]: - j += 1 - - start_rank = i + 1 - end_rank = j - average_rank = (start_rank + end_rank) / 2.0 - for k in range(i, j): - original_index = indexed[k][0] - ranks[original_index] = average_rank - - tie_size = j - i - if tie_size > 1: - tie_counts.append(tie_size) - i = j - - return ranks, tie_counts - - def _sample_variance(values: Sequence[float]) -> float: if len(values) < 2: return float("nan") @@ -368,9 +338,5 @@ def _cohens_d_from_differences(differences: Sequence[float]) -> float: return mean_diff / std_diff -def _normal_sf(z_score: float) -> float: - return 0.5 * math.erfc(z_score / math.sqrt(2.0)) - - def _is_nan(value: float) -> bool: return isinstance(value, float) and math.isnan(value) diff --git a/src/slices/models/encoders/gru_d.py b/src/slices/models/encoders/gru_d.py index 838b40d..5a7e225 100644 --- a/src/slices/models/encoders/gru_d.py +++ b/src/slices/models/encoders/gru_d.py @@ -55,8 +55,8 @@ def __init__(self, config: GRUDConfig) -> None: self.W_gamma_x = nn.Parameter(torch.Tensor(d_input)) self.b_gamma_x = nn.Parameter(torch.Tensor(d_input)) - # Hidden state decay parameters - self.W_gamma_h = nn.Parameter(torch.Tensor(d_model)) + # 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) @@ -72,6 +72,36 @@ def _reset_parameters(self) -> None: 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. @@ -119,25 +149,31 @@ def forward( delta = self._compute_time_deltas(mask) h = torch.zeros(B, self.config.d_model, device=device) - x_prev = self.x_mean.expand(B, -1).clone() + x_last_observed = self.x_mean.expand(B, -1).clone() for t in range(T): # Input decay - gamma_x = torch.exp(-torch.relu(self.W_gamma_x * delta[:, t, :] + self.b_gamma_x)) - x_decayed = mask_float[:, t, :] * x[:, t, :] + (1 - mask_float[:, t, :]) * ( - gamma_x * x_prev + (1 - gamma_x) * self.x_mean + 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 - delta_h = delta[:, t, :].mean(dim=-1, keepdim=True).expand(-1, self.config.d_model) - gamma_h = torch.exp(-torch.relu(self.W_gamma_h * delta_h + self.b_gamma_h)) + 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_prev = x_decayed + x_last_observed = self._update_last_observed( + x[:, t, :], + mask[:, t, :], + x_last_observed, + ) h = self.dropout(h) return h diff --git a/src/slices/training/finetune_module.py b/src/slices/training/finetune_module.py index 37a41c9..96f308a 100644 --- a/src/slices/training/finetune_module.py +++ b/src/slices/training/finetune_module.py @@ -19,7 +19,7 @@ from omegaconf import DictConfig, OmegaConf from slices.eval import MetricConfig, build_metrics -from slices.models.encoders import 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, @@ -266,7 +266,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: @@ -275,6 +280,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/tests/test_export_results.py b/tests/test_export_results.py index c618b50..3a63fab 100644 --- a/tests/test_export_results.py +++ b/tests/test_export_results.py @@ -7,6 +7,18 @@ 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 test_build_statistical_tests_df_produces_pairwise_significance_rows(): mod = importlib.import_module("scripts.export_results") @@ -46,6 +58,35 @@ def test_build_statistical_tests_df_produces_pairwise_significance_rows(): assert row["n_tasks"] == 2 assert row["better_paradigm"] == "mae" assert row["p_value_bonferroni"] >= row["p_value"] + assert row["n_shared_task_seed_pairs"] == 6 + assert row["n_union_task_seed_pairs"] == 6 + + +def test_extract_run_uses_requested_inherited_sprint_tag_for_family(): + mod = importlib.import_module("scripts.export_results") + + run = DummyRun( + config={ + "sprint": "6", + "dataset": "miiv", + "paradigm": "mae", + "seed": 42, + "encoder": {"d_model": 64, "n_layers": 2}, + "training": {"freeze_encoder": False}, + "task": {"task_name": "mortality_24h"}, + "label_fraction": 0.1, + }, + tags=["sprint:6", "sprint:7p", "phase:finetune"], + name="s6_inherited_capacity_baseline", + ) + + default_row = mod.extract_run(run, []) + requested_row = mod.extract_run(run, [], requested_sprints=["7p"]) + + assert default_row["experiment_type"] == "label_efficiency" + assert requested_row["experiment_type"] == "capacity_pilot" + assert requested_row["sprint"] == "7p" + assert requested_row["config_sprint"] == "6" def test_build_statistical_tests_df_adds_classical_context_rows(): @@ -91,6 +132,102 @@ def test_build_statistical_tests_df_adds_classical_context_rows(): assert ("mae", "supervised") not in pairs +def test_build_statistical_tests_df_reports_incomplete_pair_coverage(): + mod = importlib.import_module("scripts.export_results") + + rows = [] + for paradigm, seeds in [("mae", [42, 123]), ("supervised", [42])]: + for seed in seeds: + rows.append( + { + "experiment_type": "core", + "sprint": "1", + "paradigm": paradigm, + "dataset": "miiv", + "task": "mortality_24h", + "seed": seed, + "protocol": "B", + "label_fraction": 1.0, + "model_size": "default", + "source_dataset": None, + "phase": "finetune", + "test/auprc": 0.3 + (0.1 if paradigm == "mae" else 0.0), + } + ) + + stats_df = mod.build_statistical_tests_df(pd.DataFrame(rows)) + row = stats_df.iloc[0] + + assert row["n_pairs"] == 1 + assert row["n_task_seed_pairs_a"] == 2 + assert row["n_task_seed_pairs_b"] == 1 + assert row["n_shared_task_seed_pairs"] == 1 + assert row["n_union_task_seed_pairs"] == 2 + assert '"seed": 123' in row["missing_task_seed_pairs_b"] + + +def test_build_ts2vec_vs_core_contrastive_df_compares_across_experiment_types(): + mod = importlib.import_module("scripts.export_results") + + rows = [] + for experiment_type, paradigm, offset in [ + ("temporal_contrastive", "ts2vec", 0.04), + ("core", "contrastive", 0.0), + ]: + for seed in [42, 123]: + for task, base in [("mortality_24h", 0.30), ("aki_kdigo", 0.25)]: + rows.append( + { + "experiment_type": experiment_type, + "sprint": "13" if paradigm == "ts2vec" else "1", + "paradigm": paradigm, + "dataset": "miiv", + "task": task, + "seed": seed, + "protocol": "B", + "label_fraction": 1.0, + "model_size": "default", + "source_dataset": None, + "phase": "finetune", + "test/auprc": base + offset + (seed / 100000.0), + } + ) + + 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 + assert row["better_paradigm"] == "ts2vec" + + +def test_validate_flags_wrong_fixed_seed_set_with_correct_count(): + mod = importlib.import_module("scripts.export_results") + + aggregated_df = pd.DataFrame( + [ + { + "experiment_type": "core", + "paradigm": "mae", + "dataset": "miiv", + "task": "mortality_24h", + "protocol": "B", + "n_seeds": 5, + "seed_list": "[1, 2, 3, 4, 5]", + } + ] + ) + per_seed_df = pd.DataFrame([{"test/auprc": 0.5}]) + + warnings = mod.validate(per_seed_df, 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_parse_args_uses_revision_env_when_cli_omits_it(monkeypatch): mod = importlib.import_module("scripts.export_results") diff --git a/tests/test_finetune_module.py b/tests/test_finetune_module.py index d3f2a65..b7b0f0e 100644 --- a/tests/test_finetune_module.py +++ b/tests/test_finetune_module.py @@ -412,6 +412,30 @@ 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) diff --git a/tests/test_fixes.py b/tests/test_fixes.py index bf8eb44..d227e91 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -1955,6 +1955,14 @@ def test_xgboost_uses_primary_metric_for_early_stopping(self): 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) + class TestTrainingScriptClassWeighting: """Regression tests for entrypoint-level class weight resolution.""" diff --git a/tests/test_gru_d_encoder.py b/tests/test_gru_d_encoder.py new file mode 100644 index 0000000..47d9512 --- /dev/null +++ b/tests/test_gru_d_encoder.py @@ -0,0 +1,36 @@ +"""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_statistical.py b/tests/test_statistical.py index 9e5d518..5fe995d 100644 --- a/tests/test_statistical.py +++ b/tests/test_statistical.py @@ -228,6 +228,13 @@ def test_paired_wilcoxon_detects_consistent_improvement(self): 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) From ea58cb30971d5cfcd09e17dd07cf355a9251ab3f Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 21 Apr 2026 19:03:51 -0400 Subject: [PATCH 092/121] Add operational experiment guardrails Tighten retry handling for skipped runs with failed dependencies, cap XGBoost CPU parallelism, make raw tensor cache writes atomic, and validate split ratio bounds. Refresh label freshness hashing, Sprint 7p W&B model-size metadata, and fairness scope docs so operational reporting stays unambiguous. Behavioral implications: retry --skipped now fails fast unless blocked failed dependencies are also requested with --failed; Sprint 11 XGBoost runs default to n_jobs=4 instead of using all CPU threads per process; label config changes involving sources or supported datasets now invalidate stale extracted labels. --- configs/xgboost.yaml | 1 + docs/EVAL_FAIRNESS.md | 4 +- scripts/internal/run_experiments.py | 185 ++++++++++++++++++---- scripts/training/xgboost_baseline.py | 8 +- src/slices/data/config_schemas.py | 6 +- src/slices/data/datamodule.py | 9 ++ src/slices/data/extractors/base.py | 6 + src/slices/data/labels/base.py | 5 + src/slices/data/tensor_cache.py | 16 +- src/slices/training/utils.py | 7 + tests/test_config_schemas.py | 24 +++ tests/test_dataset_datamodule.py | 9 ++ tests/test_fixes.py | 229 +++++++++++++++++++++++++++ 13 files changed, 466 insertions(+), 43 deletions(-) diff --git a/configs/xgboost.yaml b/configs/xgboost.yaml index 2a47a45..7d86d85 100644 --- a/configs/xgboost.yaml +++ b/configs/xgboost.yaml @@ -46,6 +46,7 @@ xgboost: min_child_weight: 1 early_stopping_rounds: 20 scale_pos_weight: null # set to "balanced" to auto-compute from label distribution + n_jobs: 4 # cap per-process CPU use when Sprint 11 runs many baselines in parallel # Data loading training: diff --git a/docs/EVAL_FAIRNESS.md b/docs/EVAL_FAIRNESS.md index 8cb1965..cab93cd 100644 --- a/docs/EVAL_FAIRNESS.md +++ b/docs/EVAL_FAIRNESS.md @@ -12,8 +12,8 @@ to W&B summaries and the parquet exports produced by `scripts/export_results.py` - The standalone script is revision-scoped. Pass `--revision ` so reruns do not get mixed together. - Default benchmark fairness corpus: - - Internal experiment groups `1`, `2`, `3`, `4`, `5`, `7p`, `10`, `11`, - `12`, `13` + - Internal experiment groups `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`, `7p`, + `10`, `11`, `12`, `13` - 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. diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index fbd5c85..ed7f900 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -19,6 +19,8 @@ uv run python scripts/internal/run_experiments.py status --sprint 1 uv run python scripts/internal/run_experiments.py retry --failed --parallel 4 \\ --project slices-thesis --revision thesis-v1 --entity + uv run python scripts/internal/run_experiments.py retry --skipped --failed --parallel 4 \\ + --project slices-thesis --revision thesis-v1 --entity uv run python scripts/internal/run_experiments.py retry --failed \\ --sprint 1 --revision thesis-v1 --parallel 4 \\ --project slices-thesis --entity @@ -480,6 +482,12 @@ def _add_xgboost( self.runs.append(run) return run + @staticmethod + def _add_model_size_metadata(run: Run, size_name: str) -> Run: + """Add Hydra metadata used by W&B naming/tagging without changing run IDs.""" + run.extra_overrides = {**run.extra_overrides, "+model_size": size_name} + return run + # --- Sprint builders --- def build_sprint1(self): @@ -642,33 +650,40 @@ def build_sprint7p(self): **size_cfg["ssl_scale"]["mae"], } pt = self._add_pretrain(sprint, "mae", ds, seed, extra=pretrain_extra) + self._add_model_size_metadata(pt, size_name) # --- MAE finetune (Protocol B) + probe (Protocol A) --- finetune_extra = {**model_extra} for frac in LABEL_FRACTIONS_PILOT: - self._add_finetune( - sprint, - "mae", - ds, - seed, - task, - False, - pt, - label_fraction=frac, - extra=finetune_extra, - name_extra={"size": size_name}, + self._add_model_size_metadata( + self._add_finetune( + sprint, + "mae", + ds, + seed, + task, + False, + pt, + label_fraction=frac, + extra=finetune_extra, + name_extra={"size": size_name}, + ), + size_name, ) - self._add_finetune( - sprint, - "mae", - ds, - seed, - task, - True, - pt, - label_fraction=frac, - extra=finetune_extra, - name_extra={"size": size_name}, + self._add_model_size_metadata( + self._add_finetune( + sprint, + "mae", + ds, + seed, + task, + True, + pt, + label_fraction=frac, + extra=finetune_extra, + name_extra={"size": size_name}, + ), + size_name, ) # --- Supervised baseline (same larger encoder, from scratch) --- @@ -690,6 +705,7 @@ def build_sprint7p(self): label_fraction=frac, extra_overrides=sup_extra, ) + self._add_model_size_metadata(run, size_name) self.runs.append(run) def build_sprint7(self): @@ -1451,6 +1467,89 @@ def cmd_status(args): print_status(sprint_filter) +def _collect_dependency_closure( + runs: list[Run], + all_by_id: dict[str, Run], +) -> tuple[list[Run], dict[str, set[str]], set[str]]: + """Return transitive dependencies for runs, with provenance for reporting.""" + deps: list[Run] = [] + required_by: dict[str, set[str]] = {} + missing: set[str] = set() + seen: set[str] = set() + queue: list[tuple[str, str]] = [(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 _retry_command_suggestion(args, *, include_failed: bool, include_skipped: bool) -> str: + """Build a retry command that preserves the user's current retry scope.""" + cmd = ["uv", "run", "python", "scripts/internal/run_experiments.py", "retry"] + if include_failed: + cmd.append("--failed") + if include_skipped: + cmd.append("--skipped") + if args.sprint: + cmd.append("--sprint") + cmd.extend(str(s) for s in args.sprint) + 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)]) + 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 a dependency report for retry requests that cannot make progress.""" + 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 cmd_retry(args): all_runs = generate_all_runs() sprint_filter = {str(s) for s in args.sprint} if args.sprint else None @@ -1484,26 +1583,46 @@ def cmd_retry(args): for r in candidate_runs: status = get_run_status(state, r.id) if args.failed and status == "failed": - set_run_status(state, r.id, "pending") runs_to_retry.append(r) elif args.skipped and status == "skipped": - set_run_status(state, r.id, "pending") runs_to_retry.append(r) if not runs_to_retry: print("No runs to retry.") return 0 - # Include their dependencies that are completed (already fine) or pending all_by_id = {r.id: r for r in all_runs} - retry_ids = {r.id for r in runs_to_retry} - deps_needed = set() - for r in runs_to_retry: - for d in r.depends_on: - if d not in retry_ids: - deps_needed.add(d) - extra = [all_by_id[d] for d in deps_needed if d in all_by_id] - runs_to_retry = extra + runs_to_retry + dependencies, required_by, missing = _collect_dependency_closure(runs_to_retry, all_by_id) + + # Failed/skipped dependencies cannot satisfy downstream runs. Treat --failed + # and --skipped as the explicit request to reset those dependency states too. + 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, + ) + + runs_by_retry_id: dict[str, Run] = {} + for run in dependencies + runs_to_retry: + runs_by_retry_id.setdefault(run.id, run) + + for run in runs_by_retry_id.values(): + if get_run_status(state, run.id) in {"failed", "skipped"}: + set_run_status(state, run.id, "pending") + + runs_to_retry = list(runs_by_retry_id.values()) save_state(state) print(f"Retrying {len(runs_to_retry)} runs") diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index 884762f..7e66fdf 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -25,9 +25,10 @@ r2_score, roc_auc_score, ) +from xgboost import XGBClassifier, XGBRegressor + from slices.data.datamodule import ICUDataModule from slices.eval.inference import extract_tabular_features -from xgboost import XGBClassifier, XGBRegressor def _xgboost_eval_metric(task_type: str) -> str: @@ -156,6 +157,9 @@ def main(cfg: DictConfig) -> None: 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, @@ -165,7 +169,7 @@ def main(cfg: DictConfig) -> None: "min_child_weight": xgb_cfg.min_child_weight, "early_stopping_rounds": early_stopping_rounds, "random_state": cfg.seed, - "n_jobs": -1, + "n_jobs": n_jobs, } if task_type == "regression": diff --git a/src/slices/data/config_schemas.py b/src/slices/data/config_schemas.py index fe6ffdf..0696af2 100644 --- a/src/slices/data/config_schemas.py +++ b/src/slices/data/config_schemas.py @@ -203,9 +203,9 @@ 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 diff --git a/src/slices/data/datamodule.py b/src/slices/data/datamodule.py index 9c68c25..f02d2ff 100644 --- a/src/slices/data/datamodule.py +++ b/src/slices/data/datamodule.py @@ -163,6 +163,15 @@ 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}") diff --git a/src/slices/data/extractors/base.py b/src/slices/data/extractors/base.py index 7fa975d..cb9546c 100644 --- a/src/slices/data/extractors/base.py +++ b/src/slices/data/extractors/base.py @@ -157,6 +157,12 @@ def _build_label_manifest( 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 diff --git a/src/slices/data/labels/base.py b/src/slices/data/labels/base.py index 2965e5a..830ea19 100644 --- a/src/slices/data/labels/base.py +++ b/src/slices/data/labels/base.py @@ -61,13 +61,18 @@ def config_hash(config: LabelConfig) -> str: """ # 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] diff --git a/src/slices/data/tensor_cache.py b/src/slices/data/tensor_cache.py index 09c068c..6dccb75 100644 --- a/src/slices/data/tensor_cache.py +++ b/src/slices/data/tensor_cache.py @@ -420,7 +420,7 @@ def save_cached_tensors( seq_length: int, n_features: int, ) -> None: - """Save raw tensors to cache file. + """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. @@ -434,7 +434,7 @@ def save_cached_tensors( """ cache_path = get_tensor_cache_path(data_dir, seq_length, n_features) cache_dir = cache_path.parent - cache_dir.mkdir(exist_ok=True) + cache_dir.mkdir(parents=True, exist_ok=True) cache_data = { "timeseries_tensor": timeseries_tensor, @@ -447,7 +447,17 @@ def save_cached_tensors( try: logger.debug(f"Saving raw tensors to cache: {cache_path.name}") - torch.save(cache_data, cache_path) + 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}") diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 333ef33..66a0e41 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -286,6 +286,9 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: if len(tag) > 64: tag = tag[:61] + "..." tags.append(tag) + model_size = cfg.get("model_size") + if model_size is not None: + tags.append(f"model_size:{model_size}") # Add protocol tag for finetune runs (Protocol A = frozen, Protocol B = unfrozen) freeze_encoder = cfg.get("training", {}).get("freeze_encoder") @@ -309,6 +312,8 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: run_name = cfg.logging.get("run_name", None) if run_name and freeze_encoder is True: run_name = run_name.replace("_finetune_", "_probe_", 1) + if run_name and model_size is not None: + run_name += f"_{model_size}" # Adjust group to include protocol and label_fraction so that W&B "Group" view # aggregates exactly the runs that differ only by seed. @@ -316,6 +321,8 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: if group: if freeze_encoder is True: group = group.replace("finetune_", "probe_", 1) + 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}" diff --git a/tests/test_config_schemas.py b/tests/test_config_schemas.py index 712f667..bebf302 100644 --- a/tests/test_config_schemas.py +++ b/tests/test_config_schemas.py @@ -9,6 +9,8 @@ import pytest import yaml from pydantic import ValidationError + +from slices.data.config_schemas import DataConfig from slices.training.config_schemas import ( OptimizerConfig, SchedulerConfig, @@ -105,6 +107,28 @@ def test_label_definition_fields_are_accepted(self): assert cfg.label_sources == ["stays", "mortality_info"] +class TestDataConfigValidation: + """Tests for data loading config bounds.""" + + 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, + ) + + class TestTrainingConfigValidation: """Tests that TrainingConfig catches invalid/misspelled keys.""" diff --git a/tests/test_dataset_datamodule.py b/tests/test_dataset_datamodule.py index 26c4cd6..24fa414 100644 --- a/tests/test_dataset_datamodule.py +++ b/tests/test_dataset_datamodule.py @@ -8,6 +8,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.splits import compute_patient_level_splits, load_cached_splits @@ -848,6 +849,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( diff --git a/tests/test_fixes.py b/tests/test_fixes.py index d227e91..9ccb0f4 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -20,6 +20,7 @@ import torch import yaml from omegaconf import OmegaConf + from slices.data.labels import LabelBuilder, LabelBuilderFactory, LabelConfig from slices.data.labels.aki import AKILabelBuilder from slices.data.labels.los import LOSLabelBuilder @@ -110,6 +111,36 @@ def test_config_hash_changes_on_gap_change(self): ) 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 @@ -398,6 +429,26 @@ def test_tensor_cache_invalidated_on_data_fingerprint_mismatch(self, tmp_path, m 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.""" @@ -1569,6 +1620,50 @@ def test_apply_wandb_target_injects_project_and_entity(self): assert result[0].extra_overrides["logging.wandb_project"] == "slices-thesis" assert result[0].extra_overrides["logging.wandb_entity"] == "hannes-ill" + 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/sprint7p/example", + "sprint": "7p", + "model_size": "medium", + "label_fraction": 0.1, + "training": {"freeze_encoder": False}, + "logging": { + "use_wandb": True, + "wandb_project": "slices", + "wandb_entity": None, + "run_name": "s7p_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"] == "s7p_finetune_miiv_mae_mortality_24h_seed42_medium" + assert captured["kwargs"]["group"] == "finetune_miiv_mae_mortality_24h_medium_frac01" + assert "model_size:medium" in captured["kwargs"]["tags"] + def test_transfer_finetune_command_propagates_source_dataset(self): from scripts.internal.run_experiments import Run @@ -1736,6 +1831,131 @@ def test_cmd_retry_respects_sprint_filter_but_keeps_dependencies(self, monkeypat 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", + sprint="0", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/sprint0/pretrain_mae_miiv_seed42", + ) + target = runner.Run( + id="target_finetune", + sprint="1", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/sprint1/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( + sprint=["1"], + 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", + sprint="0", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/sprint0/pretrain_mae_miiv_seed42", + ) + target = runner.Run( + id="target_finetune", + sprint="1", + run_type="finetune", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/sprint1/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( + sprint=["1"], + 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_revises_dependencies_for_revision_scoped_retry(self, monkeypatch): import scripts.internal.run_experiments as runner @@ -1963,6 +2183,11 @@ def test_xgboost_binary_ece_matches_uniform_bin_definition(self): 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 + class TestTrainingScriptClassWeighting: """Regression tests for entrypoint-level class weight resolution.""" @@ -2431,6 +2656,10 @@ def test_sprint7p_expanded_to_five_seeds(self): assert len(builder.runs) == 100 assert sorted({run.seed for run in builder.runs}) == SEEDS_EXTENDED assert BASELINE_SPRINTS["7p"] == ["6", "10"] + assert {run.extra_overrides["+model_size"] for run in builder.runs} == { + "medium", + "large", + } def test_sprint13_includes_both_protocols(self): from scripts.internal.run_experiments import MatrixBuilder From cdf032f56a7af2d44461018e6c712b37dfbd96bc Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 21 Apr 2026 20:05:27 -0400 Subject: [PATCH 093/121] Honor inherited capacity baselines during export Add explicit export membership rows for Sprint 7p inherited default-size baselines so revision-wide and multi-sprint exports include them in the capacity_pilot family while preserving their source label-efficiency rows. This changes result export classification only; existing W&B runs and training artifacts do not need migration. --- scripts/export_results.py | 67 ++++++++++++++++++++++++++++++++++-- tests/test_export_results.py | 53 ++++++++++++++++++++++++++++ 2 files changed, 117 insertions(+), 3 deletions(-) diff --git a/scripts/export_results.py b/scripts/export_results.py index 41d8c78..ca50404 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -5,7 +5,7 @@ Pulls all experiment runs from W&B, extracts config + test metrics, and produces structured parquet files: - - results/per_seed_results.parquet — one row per W&B run (~3000 rows) + - results/per_seed_results.parquet — one row per W&B run/export membership - results/aggregated_results.parquet — one row per unique config (~600 rows), with mean/std/min/max across seeds - results/statistical_tests.parquet — pairwise Wilcoxon + Bonferroni + @@ -42,7 +42,6 @@ import pandas as pd import wandb - from slices.eval.statistical import ( bonferroni_correction, cohens_d, @@ -91,6 +90,8 @@ "temporal_contrastive", } EXPECTED_FIXED_SEEDS = {42, 123, 456, 789, 1011} +CAPACITY_PILOT_LABEL_FRACTIONS = {0.01, 0.1, 0.5} +CAPACITY_PILOT_PARADIGMS = {"mae", "supervised"} # For core runs: the sprint tag is not part of the config identity. # lr/mask_ratio excluded because they're determined by protocol (redundant) or @@ -415,6 +416,53 @@ def _select_export_sprint( return requested_matches[0] +def _is_capacity_pilot_inherited_row(row: dict) -> bool: + """Return True for default-size rows inherited into Sprint 7p comparisons.""" + try: + label_fraction = float(row.get("label_fraction", 1.0)) + except (TypeError, ValueError): + return False + + return ( + row.get("dataset") == "miiv" + and row.get("task") == "mortality_24h" + and row.get("paradigm") in CAPACITY_PILOT_PARADIGMS + and row.get("model_size") == "default" + and any( + math.isclose(label_fraction, allowed, rel_tol=0.0, abs_tol=1e-9) + for allowed in CAPACITY_PILOT_LABEL_FRACTIONS + ) + ) + + +def _additional_inherited_export_sprints( + row: dict, + tags: list[str], + requested_sprints: list[str] | None = None, +) -> list[str]: + """Return inherited target sprints that need explicit export rows. + + Unscoped and multi-sprint thesis exports need both the source-family row and + target-family comparison row for inherited baselines. Sprint tags are + intentionally coarse, so target-specific eligibility keeps unrelated source + rows out of the target family. + """ + config_sprint = str(row.get("config_sprint", "")) + current_sprint = str(row.get("sprint", "")) + requested = {str(sprint) for sprint in requested_sprints} if requested_sprints else None + inherited: list[str] = [] + + for sprint in _sprint_tags(tags): + if requested is not None and sprint not in requested: + continue + if sprint in {config_sprint, current_sprint}: + continue + if sprint == "7p" and _is_capacity_pilot_inherited_row(row): + inherited.append(sprint) + + return list(dict.fromkeys(inherited)) + + def _recover_pretrain_metadata( run_name: str, config: dict ) -> tuple[float | None, float | None, str | None]: @@ -664,7 +712,7 @@ def build_per_seed_df( ) -> pd.DataFrame: """Build the per-seed DataFrame from raw W&B runs. - One row per run with config columns + metric columns + wandb IDs. + One row per run/export membership with config columns, metrics, and W&B IDs. """ rows = [] failed = [] @@ -674,6 +722,19 @@ def build_per_seed_df( try: row = extract_run(run, ALL_METRICS, requested_sprints=requested_sprints) rows.append(row) + tags = [f"sprint:{s}" for s in json.loads(row.get("sprint_tags", "[]"))] + for inherited_sprint in _additional_inherited_export_sprints( + row, + tags, + requested_sprints=requested_sprints, + ): + rows.append( + extract_run( + run, + ALL_METRICS, + requested_sprints=[inherited_sprint], + ) + ) except Exception as e: run_id = getattr(run, "id", "unknown") failed.append(run_id) diff --git a/tests/test_export_results.py b/tests/test_export_results.py index 3a63fab..35ed734 100644 --- a/tests/test_export_results.py +++ b/tests/test_export_results.py @@ -89,6 +89,59 @@ def test_extract_run_uses_requested_inherited_sprint_tag_for_family(): assert requested_row["config_sprint"] == "6" +def test_build_per_seed_df_adds_revision_wide_capacity_membership_for_inherited_row(): + mod = importlib.import_module("scripts.export_results") + + run = DummyRun( + config={ + "sprint": "6", + "dataset": "miiv", + "paradigm": "mae", + "seed": 42, + "encoder": {"d_model": 64, "n_layers": 2}, + "training": {"freeze_encoder": False}, + "task": {"task_name": "mortality_24h"}, + "label_fraction": 0.1, + }, + tags=["sprint:6", "sprint:7p", "phase:finetune"], + name="s6_inherited_capacity_baseline", + ) + + df = mod.build_per_seed_df([run]) + + assert set(df["experiment_type"]) == {"label_efficiency", "capacity_pilot"} + capacity = df[df["experiment_type"] == "capacity_pilot"].iloc[0] + assert capacity["sprint"] == "7p" + assert capacity["config_sprint"] == "6" + assert capacity["wandb_run_id"] == "dummy-id" + + scoped_df = mod.build_per_seed_df([run], requested_sprints=["6", "7p"]) + assert set(scoped_df["experiment_type"]) == {"label_efficiency", "capacity_pilot"} + + +def test_build_per_seed_df_does_not_add_capacity_membership_for_irrelevant_tagged_rows(): + mod = importlib.import_module("scripts.export_results") + + run = DummyRun( + config={ + "sprint": "6", + "dataset": "eicu", + "paradigm": "jepa", + "seed": 42, + "encoder": {"d_model": 64, "n_layers": 2}, + "training": {"freeze_encoder": False}, + "task": {"task_name": "mortality_hospital"}, + "label_fraction": 0.05, + }, + tags=["sprint:6", "sprint:7p", "phase:finetune"], + name="s6_coarsely_tagged_but_not_capacity_scope", + ) + + df = mod.build_per_seed_df([run]) + + assert df["experiment_type"].tolist() == ["label_efficiency"] + + def test_build_statistical_tests_df_adds_classical_context_rows(): mod = importlib.import_module("scripts.export_results") From 0f02f7923ffa25dfe2f864638d1cc9a1d2eecf43 Mon Sep 17 00:00:00 2001 From: hill Date: Tue, 21 Apr 2026 21:22:49 -0400 Subject: [PATCH 094/121] Fix benchmark leakage and experiment safeguards Correct mortality label exclusion for deaths inside the observation window and filter zero-observation ICU stays while making SSL objectives robust to all-masked samples. Harden result export, experiment resumption, fairness reruns, and XGBoost artifact handling so publication benchmark runs fail visibly instead of silently dropping or misclassifying results. Behavioral impact: regenerated processed datasets use mortality label version 2.1.2; benchmark exports now fail closed on extraction errors unless explicitly overridden. --- configs/finetune.yaml | 3 + configs/pretrain.yaml | 3 + configs/xgboost.yaml | 2 +- scripts/eval/evaluate_fairness.py | 66 ++++++- scripts/export_results.py | 121 ++++++++++++- scripts/internal/launch_thesis_tmux.sh | 3 + scripts/internal/run_experiments.py | 95 ++++++++-- scripts/training/finetune.py | 5 +- scripts/training/pretrain.py | 5 +- scripts/training/xgboost_baseline.py | 23 ++- src/slices/data/extractors/ricu.py | 23 +++ src/slices/data/labels/mortality.py | 44 ++--- src/slices/models/encoders/transformer.py | 21 ++- src/slices/models/pretraining/contrastive.py | 38 ++++ src/slices/models/pretraining/jepa.py | 4 + src/slices/models/pretraining/mae.py | 4 + src/slices/training/utils.py | 8 +- tests/test_contrastive_objective.py | 16 ++ tests/test_evaluate_fairness.py | 33 ++++ tests/test_export_results.py | 127 ++++++++++++++ tests/test_fixes.py | 175 ++++++++++++++++++- tests/test_jepa_objective.py | 14 ++ tests/test_task_builders.py | 5 +- tests/test_transformer_encoder.py | 13 ++ 24 files changed, 777 insertions(+), 74 deletions(-) diff --git a/configs/finetune.yaml b/configs/finetune.yaml index 050f30f..49e0c18 100644 --- a/configs/finetune.yaml +++ b/configs/finetune.yaml @@ -52,6 +52,9 @@ seed: 42 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 diff --git a/configs/pretrain.yaml b/configs/pretrain.yaml index ae137c4..d49ac99 100644 --- a/configs/pretrain.yaml +++ b/configs/pretrain.yaml @@ -36,6 +36,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: diff --git a/configs/xgboost.yaml b/configs/xgboost.yaml index 7d86d85..9d1e5c2 100644 --- a/configs/xgboost.yaml +++ b/configs/xgboost.yaml @@ -45,7 +45,7 @@ xgboost: colsample_bytree: 0.8 min_child_weight: 1 early_stopping_rounds: 20 - scale_pos_weight: null # set to "balanced" to auto-compute from label distribution + scale_pos_weight: balanced # auto-compute from label distribution for binary tasks n_jobs: 4 # cap per-process CPU use when Sprint 11 runs many baselines in parallel # Data loading diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index 9828305..579edb2 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -42,6 +42,7 @@ import argparse import json import logging +import math import os import re import sys @@ -149,15 +150,46 @@ def fetch_eval_runs( return runs -def has_fairness_metrics(run) -> bool: - """Check if run summary already contains fairness/* keys. +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_finite_summary_value(summary: dict[str, Any], key: str) -> bool: + value = summary.get(key) + if value is None: + return False + if isinstance(value, float) and math.isnan(value): + return False + return True + + +def has_fairness_metrics(run, protected_attributes: list[str]) -> 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. """ try: sm = run.summary_metrics or {} - return any(k.startswith("fairness/") for k in sm) + expected_attrs = _expected_fairness_attributes(run, protected_attributes) + if not expected_attrs: + return False + for attr in expected_attrs: + prefix = f"fairness/{attr}/" + if not _has_finite_summary_value(sm, f"{prefix}n_valid_groups"): + return False + has_task_metric = _has_finite_summary_value( + sm, + f"{prefix}worst_group_auroc", + ) or _has_finite_summary_value(sm, f"{prefix}worst_group_mse") + if not has_task_metric: + return False + return True except Exception: return False @@ -166,17 +198,32 @@ def write_fairness_to_wandb( run_path: str, fairness_flat: dict[str, Any], dry_run: bool = False, + clear_existing: bool = False, ) -> None: - """Write fairness metrics to W&B run summary (additive only).""" + """Write fairness metrics to W&B run summary.""" import wandb if dry_run: - log.info(" [DRY RUN] Would write %d fairness keys to %s", len(fairness_flat), run_path) + 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/") and key not in fairness_flat: + run.summary[key] = None run.summary.update(fairness_flat) run.summary.save() @@ -640,7 +687,7 @@ def main() -> None: # 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)] + runs = [r for r in runs if not has_fairness_metrics(r, args.protected_attributes)] skipped = before - len(runs) if skipped: log.info("Skipped %d runs with existing fairness metrics.", skipped) @@ -758,7 +805,12 @@ def _sort_key(r): if args.entity else f"{args.project}/{run.id}" ) - write_fairness_to_wandb(run_path, fairness_flat, args.dry_run) + write_fairness_to_wandb( + run_path, + fairness_flat, + args.dry_run, + clear_existing=args.force, + ) results["processed"] += 1 # Free model memory diff --git a/scripts/export_results.py b/scripts/export_results.py index ca50404..14b7d76 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -435,6 +435,72 @@ def _is_capacity_pilot_inherited_row(row: dict) -> bool: ) +def _is_label_efficiency_inherited_row(row: dict) -> bool: + """Return True for full-label rows inherited as 100% endpoints.""" + try: + label_fraction = float(row.get("label_fraction", 1.0)) + except (TypeError, ValueError): + return False + + return ( + math.isclose(label_fraction, 1.0, rel_tol=0.0, abs_tol=1e-9) + and row.get("phase") in {"finetune", "supervised"} + and row.get("paradigm") in {"mae", "jepa", "contrastive", "supervised"} + and row.get("model_size") == "default" + and row.get("source_dataset") is None + ) + + +def _is_transfer_inherited_row(row: dict) -> bool: + """Return True for in-domain SSL rows inherited into transfer comparisons.""" + try: + label_fraction = float(row.get("label_fraction", 1.0)) + except (TypeError, ValueError): + return False + + return ( + row.get("phase") == "finetune" + and row.get("paradigm") in {"mae", "jepa", "contrastive"} + and row.get("dataset") in {"miiv", "eicu"} + and row.get("protocol") == "B" + and math.isclose(label_fraction, 1.0, rel_tol=0.0, abs_tol=1e-9) + and row.get("model_size") == "default" + and row.get("source_dataset") is None + ) + + +def _is_hp_anchor_inherited_row(row: dict) -> bool: + """Return True for default LR/mask rows inherited into HP ablations.""" + try: + label_fraction = float(row.get("label_fraction", 1.0)) + except (TypeError, ValueError): + return False + + return ( + row.get("phase") == "finetune" + and row.get("paradigm") in {"mae", "jepa", "contrastive"} + and row.get("dataset") == "miiv" + and row.get("task") == "mortality_24h" + and row.get("protocol") == "B" + and math.isclose(label_fraction, 1.0, rel_tol=0.0, abs_tol=1e-9) + and row.get("model_size") == "default" + and row.get("source_dataset") is None + ) + + +def _is_inherited_export_row_for_target(row: dict, sprint: str) -> bool: + """Return whether a coarsely sprint-tagged source row belongs in target export.""" + if sprint == "6": + return _is_label_efficiency_inherited_row(row) + if sprint == "7": + return _is_transfer_inherited_row(row) + if sprint in {"1b", "1c", "8"}: + return _is_hp_anchor_inherited_row(row) + if sprint == "7p": + return _is_capacity_pilot_inherited_row(row) + return False + + def _additional_inherited_export_sprints( row: dict, tags: list[str], @@ -457,7 +523,7 @@ def _additional_inherited_export_sprints( continue if sprint in {config_sprint, current_sprint}: continue - if sprint == "7p" and _is_capacity_pilot_inherited_row(row): + if _is_inherited_export_row_for_target(row, sprint): inherited.append(sprint) return list(dict.fromkeys(inherited)) @@ -709,6 +775,7 @@ def _assert_no_ambiguous_fingerprint_collisions(df: pd.DataFrame) -> None: def build_per_seed_df( runs: list, requested_sprints: list[str] | None = None, + allow_extraction_failures: bool = False, ) -> pd.DataFrame: """Build the per-seed DataFrame from raw W&B runs. @@ -744,8 +811,16 @@ def build_per_seed_df( 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 # Ensure correct dtypes if "seed" in df.columns: @@ -833,6 +908,7 @@ def _aggregate_group(df: pd.DataFrame, fingerprint_cols: list[str]) -> pd.DataFr # Metric aggregation 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 @@ -1357,6 +1433,35 @@ def validate( "These may be pretraining runs or crashed evaluations." ) + # Check primary metric presence for every evaluation row. + if not per_seed_df.empty and "task" in per_seed_df.columns: + missing_primary = [] + for _, row in per_seed_df.iterrows(): + phase = row.get("phase") + if 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." + ) + for row, primary_metric in missing_primary[:10]: + warnings.append( + " " + + ", ".join( + f"{c}={row.get(c)}" + for c in ["wandb_run_id", "paradigm", "dataset", "task", "seed", "phase"] + if c in row + ) + + f" — missing {primary_metric}" + ) + # Summary by experiment type n_runs = len(per_seed_df) n_configs = len(aggregated_df) @@ -1444,6 +1549,14 @@ def parse_args() -> argparse.Namespace: 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." + ), + ) args = parser.parse_args() if not args.revision: @@ -1482,7 +1595,11 @@ def main(): # Build DataFrames print(f"\nBuilding per-seed DataFrame from {len(runs)} runs...", file=sys.stderr) - per_seed_df = build_per_seed_df(runs, requested_sprints=args.sprint) + per_seed_df = build_per_seed_df( + runs, + requested_sprints=args.sprint, + allow_extraction_failures=args.allow_extraction_failures, + ) print(f" Shape: {per_seed_df.shape}", file=sys.stderr) print("\nBuilding aggregated DataFrame...", file=sys.stderr) diff --git a/scripts/internal/launch_thesis_tmux.sh b/scripts/internal/launch_thesis_tmux.sh index 2c11152..f159cc2 100755 --- a/scripts/internal/launch_thesis_tmux.sh +++ b/scripts/internal/launch_thesis_tmux.sh @@ -73,6 +73,9 @@ tag_args=(--project "$WANDB_PROJECT" --entity "$WANDB_ENTITY") run_args+=(--revision "$REVISION" --reason "$REASON") export_args+=(--revision "$REVISION" --output-dir "$RESULTS_DIR") +if [[ -n "$REVISION" ]]; then + tag_args+=(--revision "$REVISION") +fi if [[ -n "$REVISION" ]]; then fairness_args+=(--revision "$REVISION") fi diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index ed7f900..dd198c5 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -121,10 +121,10 @@ "3": [], "4": [], "5": ["1", "2", "3", "4"], - "6": ["1", "2", "3", "4", "5"], - "7": ["1", "3", "5"], + "6": ["1", "2", "3", "4", "5", "10"], + "7": ["1", "3", "5", "10"], "7p": ["6", "10"], - "8": ["1", "1b", "1c", "5"], + "8": ["1", "1b", "1c", "5", "10"], } STATE_FILE = Path("outputs/experiment_state.json") @@ -186,6 +186,9 @@ def _pretrain_cmd(self) -> list[str]: f"sprint={self.sprint}", f"hydra.run.dir={self.output_dir}", ] + last_ckpt = Path(self.output_dir) / "checkpoints" / "last.ckpt" + if last_ckpt.exists(): + cmd.append(f"ckpt_path={last_ckpt}") for k, v in self.extra_overrides.items(): cmd.append(f"{k}={v}") return cmd @@ -209,6 +212,9 @@ def _finetune_cmd(self, runs_by_id: dict[str, Run]) -> list[str]: f"sprint={self.sprint}", f"hydra.run.dir={self.output_dir}", ] + last_ckpt = Path(self.output_dir) / "checkpoints" / "last.ckpt" + if last_ckpt.exists(): + cmd.append(f"ckpt_path={last_ckpt}") if self.freeze_encoder is True: cmd += [ "training.freeze_encoder=true", @@ -253,6 +259,9 @@ def _supervised_cmd(self) -> list[str]: f"sprint={self.sprint}", f"hydra.run.dir={self.output_dir}", ] + last_ckpt = Path(self.output_dir) / "checkpoints" / "last.ckpt" + if last_ckpt.exists(): + cmd.append(f"ckpt_path={last_ckpt}") if self.label_fraction < 1.0: cmd.append(f"label_fraction={self.label_fraction}") for k, v in self.extra_overrides.items(): @@ -273,6 +282,9 @@ def _gru_d_cmd(self) -> list[str]: f"sprint={self.sprint}", f"hydra.run.dir={self.output_dir}", ] + last_ckpt = Path(self.output_dir) / "checkpoints" / "last.ckpt" + if last_ckpt.exists(): + cmd.append(f"ckpt_path={last_ckpt}") if self.label_fraction < 1.0: cmd.append(f"label_fraction={self.label_fraction}") for k, v in self.extra_overrides.items(): @@ -1048,14 +1060,43 @@ def is_pid_alive(pid: int) -> bool: return False +def _pid_matches_command(pid: int, expected_command: str | None) -> bool: + """Best-effort guard against PID reuse before trusting a running state.""" + 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() + if not actual_command: + return False + return expected_command == actual_command or expected_command in actual_command + + def recover_stale_running(state: dict): - """Reset running entries whose PIDs are dead back to pending.""" + """Reset running entries whose PIDs are dead or no longer match the run.""" 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(pid): + if ( + pid is None + or not is_pid_alive(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}") @@ -1220,7 +1261,13 @@ def handle_signal(signum, frame): proc._log_fh = log_fh # type: ignore[attr-defined] active[r.id] = proc set_run_status( - state, r.id, "running", started_at=now, pid=proc.pid, log_file=str(log_file) + state, + r.id, + "running", + started_at=now, + pid=proc.pid, + log_file=str(log_file), + command=shlex.join(cmd), ) # 5. Save state @@ -1395,11 +1442,14 @@ def _scheduler_exit_code(runs: list[Run], state: dict) -> int: """Return nonzero when the requested run set did not complete cleanly.""" failed = [r.id for r in runs if get_run_status(state, r.id) == "failed"] skipped = [r.id for r in runs if get_run_status(state, r.id) == "skipped"] + running = [r.id for r in runs if get_run_status(state, r.id) == "running"] + pending = [r.id for r in runs if get_run_status(state, r.id) == "pending"] - if failed or skipped: + if failed or skipped or running or pending: print( "Scheduler incomplete: " - f"{len(failed)} failed, {len(skipped)} skipped due to dependency failure" + f"{len(failed)} failed, {len(skipped)} skipped due to dependency failure, " + f"{len(running)} running, {len(pending)} pending" ) return 1 @@ -1667,9 +1717,7 @@ def cmd_warmup(args): print(f"Warming up tensor caches for sprint(s) {', '.join(sprints)}") print(f" {len(datasets)} unique datasets to cache\n") - import subprocess - import sys - + failed = [] for i, dataset in enumerate(datasets, 1): processed_dir = f"data/processed/{dataset}" print(f"[{i}/{len(datasets)}] {dataset}") @@ -1679,7 +1727,9 @@ def cmd_warmup(args): # any task-specific filtering — this populates the raw cache. result = subprocess.run( [ - sys.executable, + "uv", + "run", + "python", "-c", "from slices.data.dataset import ICUDataset; " f"ds = ICUDataset(" @@ -1698,9 +1748,15 @@ def cmd_warmup(args): 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/") print("You can now run experiments in parallel without OOM.") + return 0 def cmd_tag(args): @@ -1737,6 +1793,7 @@ def cmd_tag(args): sprints = [str(s) for s in args.sprint] dry_run = args.dry_run + revision_tag = f"revision:{args.revision}" if args.revision else None for target_sprint in sprints: source_sprints = BASELINE_SPRINTS.get(target_sprint, []) @@ -1747,10 +1804,13 @@ def cmd_tag(args): target_tag = f"sprint:{target_sprint}" sources = ", ".join(source_sprints) print(f"\nSprint {target_sprint}: inheriting baselines from sprint(s) {sources}") + if revision_tag: + print(f" Filtering source runs to {revision_tag}") # Query runs from each source sprint tagged = 0 skipped = 0 + skipped_revision = 0 for source_sprint in source_sprints: source_tag = f"sprint:{source_sprint}" runs = api.runs( @@ -1759,6 +1819,9 @@ def cmd_tag(args): ) for run in runs: + if revision_tag and revision_tag not in run.tags: + skipped_revision += 1 + continue if target_tag in run.tags: skipped += 1 continue @@ -1773,6 +1836,8 @@ def cmd_tag(args): action = "would tag" if dry_run else "tagged" print(f" {action} {tagged} runs, skipped {skipped} (already tagged)") + if skipped_revision: + print(f" skipped {skipped_revision} source runs outside {revision_tag}") print("\nDone.") @@ -1857,6 +1922,12 @@ def main(): help="Sprint(s) whose baselines to tag (e.g. 2 or 2 3 5)", ) p_tag.add_argument("--dry-run", action="store_true", help="Show what would be tagged") + p_tag.add_argument( + "--revision", + type=str, + default=os.environ.get("REVISION") or os.environ.get("WANDB_REVISION"), + help="Only inherit source runs with this revision tag (default: REVISION/WANDB_REVISION)", + ) p_tag.add_argument( "--project", type=str, diff --git a/scripts/training/finetune.py b/scripts/training/finetune.py index 12c61bf..fede201 100644 --- a/scripts/training/finetune.py +++ b/scripts/training/finetune.py @@ -290,7 +290,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 diff --git a/scripts/training/pretrain.py b/scripts/training/pretrain.py index 4602a81..babb3e3 100644 --- a/scripts/training/pretrain.py +++ b/scripts/training/pretrain.py @@ -186,7 +186,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. Save Encoder diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index 7e66fdf..55796e5 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -25,10 +25,9 @@ r2_score, roc_auc_score, ) -from xgboost import XGBClassifier, XGBRegressor - from slices.data.datamodule import ICUDataModule from slices.eval.inference import extract_tabular_features +from xgboost import XGBClassifier, XGBRegressor def _xgboost_eval_metric(task_type: str) -> str: @@ -261,7 +260,16 @@ def main(cfg: DictConfig) -> None: evaluator.print_report(fairness_report) # ========================================================================= - # 6. Log to W&B + # 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 cfg.logging.get("use_wandb", False): import wandb @@ -279,15 +287,6 @@ def main(cfg: DictConfig) -> None: run.summary.update(flatten_fairness_report(fairness_report)) wandb.finish() - # ========================================================================= - # 7. 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}") - print("\n" + "=" * 80) print("XGBoost Baseline Complete!") print("=" * 80) diff --git a/src/slices/data/extractors/ricu.py b/src/slices/data/extractors/ricu.py index 0af8e1b..b57545f 100644 --- a/src/slices/data/extractors/ricu.py +++ b/src/slices/data/extractors/ricu.py @@ -429,6 +429,27 @@ 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 # ----------------------------------------------------------------- @@ -505,6 +526,8 @@ def run(self) -> None: "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(), diff --git a/src/slices/data/labels/mortality.py b/src/slices/data/labels/mortality.py index a191a67..30371cb 100644 --- a/src/slices/data/labels/mortality.py +++ b/src/slices/data/labels/mortality.py @@ -37,7 +37,7 @@ class MortalityLabelBuilder(LabelBuilder): - "unknown" or missing: falls back to hospital_expire_flag """ - SEMANTIC_VERSION = "2.1.1" + SEMANTIC_VERSION = "2.1.2" def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: """Build mortality labels from stay and mortality data. @@ -225,23 +225,26 @@ def _build_hospital_mortality_with_obs( ) -> pl.DataFrame: """Hospital mortality with observation window exclusion.""" obs_end = pl.col("intime") + pl.duration(hours=obs_hours) - left_icu_during_obs = pl.col("outtime") < obs_end - # For timestamp precision: exact comparison - ts_died_during_obs = ( - pl.col("death_time").is_not_null() - & (pl.col("death_time") <= obs_end) - & left_icu_during_obs - ) + # 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 = pl.col("death_time").is_not_null() & (pl.col("death_time") < obs_end) - # For date precision: conservative — use outtime check - date_died_during_obs = left_icu_during_obs & pl.col("death_date").is_not_null() + # 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(pl.col("death_time_precision") == "timestamp") .then(ts_died_during_obs) .when(pl.col("death_time_precision") == "date") - .then(date_died_during_obs) + .then(date_died_during_obs | date_obs_boundary) .otherwise(pl.lit(False)) ) @@ -356,7 +359,6 @@ def _build_windowed_mortality_labels( obs_end = pl.col("intime") + pl.duration(hours=obs_hours) pred_start = pl.col("intime") + pl.duration(hours=prediction_start_hours) - left_icu_during_obs = pl.col("outtime") < obs_end if until_icu_discharge: pred_end = pl.col("outtime") @@ -365,11 +367,7 @@ def _build_windowed_mortality_labels( pred_end = pl.col("intime") + pl.duration(hours=prediction_end_hours) # --- Timestamp precision: exact comparisons --- - ts_died_during_obs = ( - pl.col("death_time").is_not_null() - & (pl.col("death_time") <= obs_end) - & left_icu_during_obs - ) + ts_died_during_obs = pl.col("death_time").is_not_null() & (pl.col("death_time") < obs_end) ts_died_during_gap = ( ( pl.col("death_time").is_not_null() @@ -392,17 +390,13 @@ def _build_windowed_mortality_labels( date_start = pl.col("death_date").cast(pl.Datetime("us")) date_end = date_start + pl.duration(hours=23, minutes=59, seconds=59) - # Died before obs_end? The entire date interval is <= obs_end - date_died_during_obs = ( - pl.col("death_date").is_not_null() & (date_end <= obs_end) & left_icu_during_obs - ) + # 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) # 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) - & left_icu_during_obs + pl.col("death_date").is_not_null() & (date_start < obs_end) & (date_end >= obs_end) ) # Definite during prediction: entire interval within [pred_start, pred_end] diff --git a/src/slices/models/encoders/transformer.py b/src/slices/models/encoders/transformer.py index 9f2f983..4520d4d 100644 --- a/src/slices/models/encoders/transformer.py +++ b/src/slices/models/encoders/transformer.py @@ -313,8 +313,7 @@ def encode( Returns: (B, N, d_model) encoded tokens. """ - # Convert to PyTorch convention: True = ignore - key_padding_mask = ~padding_mask + key_padding_mask = self._key_padding_mask_for_attention(padding_mask) x = tokens for layer in self.layers: @@ -378,7 +377,7 @@ def forward( # 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 @@ -409,6 +408,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/pretraining/contrastive.py b/src/slices/models/pretraining/contrastive.py index 42e5b70..7d6ed03 100644 --- a/src/slices/models/pretraining/contrastive.py +++ b/src/slices/models/pretraining/contrastive.py @@ -394,6 +394,42 @@ def _nt_xent_loss( 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] + 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_visible_view1": valid_mask_1.sum().item() + / max(n_valid_samples, 1), + "contrastive_n_visible_view2": valid_mask_2.sum().item() + / max(n_valid_samples, 1), + "contrastive_n_masked_view1": (~valid_mask_1).sum().item() + / max(n_valid_samples, 1), + "contrastive_n_masked_view2": (~valid_mask_2).sum().item() + / max(n_valid_samples, 1), + "contrastive_n_samples_used": n_valid_samples, + "contrastive_n_samples_skipped": z1.shape[0] - n_valid_samples, + } + 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] + B = z1.shape[0] temperature = self.config.temperature @@ -462,6 +498,8 @@ def _nt_xent_loss( "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_samples_used": B, + "contrastive_n_samples_skipped": valid_samples.numel() - B, } return loss, metrics diff --git a/src/slices/models/pretraining/jepa.py b/src/slices/models/pretraining/jepa.py index 0f68da0..135a033 100644 --- a/src/slices/models/pretraining/jepa.py +++ b/src/slices/models/pretraining/jepa.py @@ -152,6 +152,10 @@ def forward( 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 # Fully unobserved hours should not be available as predictor context. decoded = self.predictor(full_tokens, src_key_padding_mask=predictor_padding_mask) diff --git a/src/slices/models/pretraining/mae.py b/src/slices/models/pretraining/mae.py index 2926d0d..75750cc 100644 --- a/src/slices/models/pretraining/mae.py +++ b/src/slices/models/pretraining/mae.py @@ -143,6 +143,10 @@ def forward( 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 # Fully unobserved hours should not be available as decoder context. if decoder_padding_mask is None: diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 66a0e41..0830391 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -164,7 +164,7 @@ 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, @@ -232,13 +232,9 @@ def setup_finetune_callbacks(cfg: DictConfig, checkpoint_prefix: str = "finetune mode = cfg.training.get("early_stopping_mode", default_mode) - # Replace '/' with '_' in monitor name for safe filenames. - # Lightning interprets '/' as a directory separator in checkpoint filenames, - # which causes best checkpoints to silently not save (empty directories). - safe_monitor = monitor.replace("/", "_") checkpoint_callback = ModelCheckpoint( dirpath=cfg.get("checkpoint_dir", "checkpoints"), - filename=f"{checkpoint_prefix}-{{epoch:03d}}-{{{safe_monitor}:.4f}}", + filename=f"{checkpoint_prefix}-{{epoch:03d}}", monitor=monitor, mode=mode, save_top_k=3, diff --git a/tests/test_contrastive_objective.py b/tests/test_contrastive_objective.py index f1c7cc4..67bef93 100644 --- a/tests/test_contrastive_objective.py +++ b/tests/test_contrastive_objective.py @@ -226,6 +226,22 @@ def test_single_eligible_timestep_falls_back_to_shared_view(self, encoder): 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 diff --git a/tests/test_evaluate_fairness.py b/tests/test_evaluate_fairness.py index 90c4786..87512d9 100644 --- a/tests/test_evaluate_fairness.py +++ b/tests/test_evaluate_fairness.py @@ -76,3 +76,36 @@ def test_resolve_evaluation_artifact_supports_xgboost(tmp_path): assert artifact_path == model_path assert source == "xgboost_model" + + +def test_has_fairness_metrics_requires_requested_attribute_completeness(): + mod = importlib.import_module("scripts.eval.evaluate_fairness") + + run = SimpleNamespace( + config={"dataset": "miiv"}, + summary_metrics={ + "fairness/gender/n_valid_groups": 2, + "fairness/gender/worst_group_auroc": 0.71, + "fairness/age_group/n_valid_groups": 4, + "fairness/age_group/worst_group_auroc": 0.69, + }, + ) + + 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"}, + summary_metrics={ + "fairness/gender/n_valid_groups": 2, + "fairness/gender/worst_group_auroc": 0.71, + "fairness/age_group/n_valid_groups": 4, + "fairness/age_group/worst_group_auroc": 0.69, + }, + ) + + assert mod.has_fairness_metrics(run, ["gender", "age_group", "race"]) is True diff --git a/tests/test_export_results.py b/tests/test_export_results.py index 35ed734..aeff9e6 100644 --- a/tests/test_export_results.py +++ b/tests/test_export_results.py @@ -142,6 +142,49 @@ def test_build_per_seed_df_does_not_add_capacity_membership_for_irrelevant_tagge assert df["experiment_type"].tolist() == ["label_efficiency"] +def test_build_per_seed_df_adds_unscoped_inherited_memberships_for_thesis_families(): + mod = importlib.import_module("scripts.export_results") + + base_config = { + "sprint": "1", + "dataset": "miiv", + "paradigm": "mae", + "seed": 42, + "encoder": {"d_model": 64, "n_layers": 2}, + "training": {"freeze_encoder": False}, + "task": {"task_name": "mortality_24h"}, + "label_fraction": 1.0, + } + runs = [ + DummyRun(base_config, ["sprint:1", "sprint:6", "phase:finetune"], name="core_for_s6"), + DummyRun(base_config, ["sprint:1", "sprint:7", "phase:finetune"], name="core_for_s7"), + DummyRun(base_config, ["sprint:1", "sprint:8", "phase:finetune"], name="core_for_s8"), + ] + for i, run in enumerate(runs): + run.id = f"run-{i}" + run.url = f"https://wandb.test/run-{i}" + run.created_at = f"2026-04-21T00:00:0{i}" + + df = mod.build_per_seed_df(runs) + memberships = set(zip(df["wandb_run_id"], df["sprint"], df["experiment_type"], strict=True)) + + assert ("run-0", "6", "label_efficiency") in memberships + assert ("run-1", "7", "transfer") in memberships + assert ("run-2", "8", "hp_ablation") in memberships + + +def test_build_per_seed_df_fails_closed_on_extraction_errors(): + mod = importlib.import_module("scripts.export_results") + + bad_run = object() + + with pytest.raises(RuntimeError, match="failed extraction"): + mod.build_per_seed_df([bad_run]) + + df = mod.build_per_seed_df([bad_run], allow_extraction_failures=True) + assert df.empty + + def test_build_statistical_tests_df_adds_classical_context_rows(): mod = importlib.import_module("scripts.export_results") @@ -281,6 +324,76 @@ def test_validate_flags_wrong_fixed_seed_set_with_correct_count(): assert any("unexpected=[1, 2, 3, 4, 5]" in warning for warning in warnings) +def test_build_aggregated_df_records_per_metric_non_null_counts(): + mod = importlib.import_module("scripts.export_results") + + per_seed_df = pd.DataFrame( + [ + { + "experiment_type": "core", + "sprint": "1", + "paradigm": "mae", + "dataset": "miiv", + "task": "mortality_24h", + "seed": 42, + "protocol": "B", + "label_fraction": 1.0, + "model_size": "default", + "source_dataset": None, + "phase": "finetune", + "test/auprc": 0.4, + "test/auroc": 0.7, + "wandb_run_id": "run-42", + }, + { + "experiment_type": "core", + "sprint": "1", + "paradigm": "mae", + "dataset": "miiv", + "task": "mortality_24h", + "seed": 123, + "protocol": "B", + "label_fraction": 1.0, + "model_size": "default", + "source_dataset": None, + "phase": "finetune", + "test/auprc": None, + "test/auroc": 0.8, + "wandb_run_id": "run-123", + }, + ] + ) + + aggregated = mod.build_aggregated_df(per_seed_df) + + assert aggregated.iloc[0]["n_seeds"] == 2 + assert aggregated.iloc[0]["test/auprc/n"] == 1 + assert aggregated.iloc[0]["test/auroc/n"] == 2 + + +def test_validate_warns_when_evaluation_row_missing_primary_metric(): + mod = importlib.import_module("scripts.export_results") + + per_seed_df = pd.DataFrame( + [ + { + "wandb_run_id": "run-1", + "paradigm": "mae", + "dataset": "miiv", + "task": "mortality_24h", + "seed": 42, + "phase": "finetune", + "test/auprc": None, + } + ] + ) + aggregated_df = pd.DataFrame() + + warnings = mod.validate(per_seed_df, aggregated_df) + + assert any("missing their primary test metric" in warning for warning in warnings) + + def test_parse_args_uses_revision_env_when_cli_omits_it(monkeypatch): mod = importlib.import_module("scripts.export_results") @@ -307,6 +420,20 @@ def test_parse_args_exposes_allow_incomplete_escape_hatch(monkeypatch): assert 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"], + ) + + args = mod.parse_args() + + assert args.allow_extraction_failures is True + + def test_main_exits_nonzero_when_no_runs_match(monkeypatch, tmp_path): mod = importlib.import_module("scripts.export_results") diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 9ccb0f4..79cb137 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -20,7 +20,6 @@ import torch import yaml from omegaconf import OmegaConf - from slices.data.labels import LabelBuilder, LabelBuilderFactory, LabelConfig from slices.data.labels.aki import AKILabelBuilder from slices.data.labels.los import LOSLabelBuilder @@ -1133,6 +1132,50 @@ def test_timestamp_precision_exact(self): 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) @@ -1620,6 +1663,133 @@ def test_apply_wandb_target_injects_project_and_entity(self): assert result[0].extra_overrides["logging.wandb_project"] == "slices-thesis" assert result[0].extra_overrides["logging.wandb_entity"] == "hannes-ill" + 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", + sprint="1", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir=str(pretrain_dir), + ) + finetune = Run( + id="finetune", + sprint="1", + 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) + + runner.recover_stale_running(state) + + 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_scheduler_exit_code_flags_pending_and_running_runs(self): + import scripts.internal.run_experiments as runner + + runs = [ + runner.Run( + id="pending-run", + sprint="1", + run_type="pretrain", + paradigm="mae", + dataset="miiv", + seed=42, + output_dir="outputs/pending", + ), + runner.Run( + id="running-run", + sprint="1", + 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_cmd_tag_filters_inherited_sources_by_revision(self, monkeypatch): + import scripts.internal.run_experiments as runner + + class FakeRun: + def __init__(self, tags): + self.tags = tags + self.name = "fake" + self.id = "fake" + self.updated = False + + def update(self): + self.updated = True + + matching = FakeRun(["sprint:1", "revision:thesis-v1"]) + stale = FakeRun(["sprint:1", "revision:old"]) + + class FakeApi: + default_entity = "entity" + + def runs(self, path, filters): + return [matching, stale] + + monkeypatch.setitem(sys.modules, "wandb", SimpleNamespace(Api=lambda: FakeApi())) + + args = SimpleNamespace( + sprint=["2"], + dry_run=False, + project="slices", + entity="entity", + revision="thesis-v1", + ) + + runner.cmd_tag(args) + + assert "sprint:2" in matching.tags + assert matching.updated is True + assert "sprint:2" not in stale.tags + assert stale.updated is False + def test_wandb_logger_includes_model_size_metadata(self, monkeypatch): import slices.training.utils as training_utils @@ -2656,6 +2826,9 @@ def test_sprint7p_expanded_to_five_seeds(self): assert len(builder.runs) == 100 assert sorted({run.seed for run in builder.runs}) == SEEDS_EXTENDED assert BASELINE_SPRINTS["7p"] == ["6", "10"] + assert "10" in BASELINE_SPRINTS["6"] + assert "10" in BASELINE_SPRINTS["7"] + assert "10" in BASELINE_SPRINTS["8"] assert {run.extra_overrides["+model_size"] for run in builder.runs} == { "medium", "large", diff --git a/tests/test_jepa_objective.py b/tests/test_jepa_objective.py index ff06be5..885a787 100644 --- a/tests/test_jepa_objective.py +++ b/tests/test_jepa_objective.py @@ -334,6 +334,20 @@ def capture_padding_mask(module, args, kwargs): 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 diff --git a/tests/test_task_builders.py b/tests/test_task_builders.py index 83059a3..f550d0a 100644 --- a/tests/test_task_builders.py +++ b/tests/test_task_builders.py @@ -207,9 +207,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) diff --git a/tests/test_transformer_encoder.py b/tests/test_transformer_encoder.py index d7fc29b..546df7c 100644 --- a/tests/test_transformer_encoder.py +++ b/tests/test_transformer_encoder.py @@ -931,6 +931,19 @@ def test_forward_obs_aware(self, encoder): # 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 From ef18cac8f385aeaaeac786a5273f95e79269dba5 Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 22 Apr 2026 12:18:01 -0400 Subject: [PATCH 095/121] Exclude derived RICU features from extraction Block future-derived and composite RICU concepts from model inputs, include the blocklist in the extraction fingerprint, and add extractor regression coverage. This prevents treatment-duration and derived-score leakage while forcing stale processed datasets to rebuild before thesis reruns. Behavioral implication: regenerated datasets now use 84 time-series features; existing processed parquet built with the old feature set is invalidated by the updated signature. --- src/slices/constants.py | 67 +++++++++++++++++++++++++-- src/slices/data/extractors/ricu.py | 1 + tests/test_ricu_extractor.py | 74 ++++++++++++++++++++++++++++++ 3 files changed, 137 insertions(+), 5 deletions(-) diff --git a/src/slices/constants.py b/src/slices/constants.py index 19249de..f2d2727 100644 --- a/src/slices/constants.py +++ b/src/slices/constants.py @@ -38,8 +38,65 @@ # ============================================================================= # Feature Blocklist # ============================================================================= -# RICU concepts that must be excluded from model input: -# - los_hosp / los_icu: leak downstream task labels (updated hourly by RICU, -# directly reveal the length-of-stay answer) -# - dur_var: internal ricu helper variable, not a clinical concept -FEATURE_BLOCKLIST: frozenset = frozenset({"los_hosp", "los_icu", "dur_var"}) +# 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/extractors/ricu.py b/src/slices/data/extractors/ricu.py index b57545f..91f511b 100644 --- a/src/slices/data/extractors/ricu.py +++ b/src/slices/data/extractors/ricu.py @@ -189,6 +189,7 @@ def _get_upstream_source_signature(self) -> dict: "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, } diff --git a/tests/test_ricu_extractor.py b/tests/test_ricu_extractor.py index 73065b0..5b9063f 100644 --- a/tests/test_ricu_extractor.py +++ b/tests/test_ricu_extractor.py @@ -373,6 +373,23 @@ def test_existing_extraction_is_invalidated_when_upstream_inputs_change( 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 @@ -651,6 +668,63 @@ def test_run_metadata_content(self, ricu_output_dir: Path, tmp_path: Path) -> No 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" config = ExtractorConfig( From e0a887e1d39a852cbcb8517e43a0bc6671c71677 Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 22 Apr 2026 13:09:35 -0400 Subject: [PATCH 096/121] Align downstream normalization with SSL cohorts Compute downstream normalization statistics from the full dataset/seed train split while preserving task-filtered indices for optimization and label statistics. Add AKI regression coverage, bump cache freshness, and record full-cohort normalizer provenance. Harden launch readiness by pinning the launch commit, validating 84-feature processed artifacts, clearing runtime caches before warmup, and refreshing stale AKI denominators in the experiment plan. Behavioral implication: AKI downstream and supervised runs now share SSL input scaling for each dataset/seed; existing hash-keyed normalization caches are invalidated. --- docs/internal/EXPERIMENT_PLAN.md | 21 ++-- scripts/internal/launch_thesis_tmux.sh | 113 +++++++++++++++++++++- scripts/internal/run_experiments.py | 37 ++++++- scripts/setup_and_extract.sh | 129 +++++++++++++++++++++---- src/slices/data/datamodule.py | 30 ++++-- src/slices/data/dataset.py | 52 ++++++---- src/slices/data/preparation.py | 4 +- src/slices/data/splits.py | 48 +++++++-- src/slices/data/tensor_cache.py | 7 +- src/slices/training/utils.py | 5 +- tests/test_dataset_datamodule.py | 53 +++++++++- 11 files changed, 424 insertions(+), 75 deletions(-) diff --git a/docs/internal/EXPERIMENT_PLAN.md b/docs/internal/EXPERIMENT_PLAN.md index 0ef2452..1bca3a0 100644 --- a/docs/internal/EXPERIMENT_PLAN.md +++ b/docs/internal/EXPERIMENT_PLAN.md @@ -47,9 +47,12 @@ with ascertainable creatinine-based KDIGO labels. A stay requires non-null creatinine in the 0-24h baseline window and at least one creatinine measurement in the 24-48h prediction window; stays without an ascertainable AKI label are excluded from AKI supervised training/evaluation but remain available for SSL -pretraining and non-AKI tasks. This leaves 56,403/74,829 MIIV stays, -88,239/132,900 eICU stays, and 144,642/207,729 combined stays. Most exclusions -are due to no 24-48h creatinine measurement. +pretraining and non-AKI tasks. In the current 84-feature processed artifacts, +metadata-derived counts are 56,403/74,829 MIIV stays, 88,239/132,686 eICU +stays, and 144,642/207,515 combined stays. Final tables should derive full +denominators from each processed `metadata.yaml` (`n_stays`) rather than +copying static denominators into prose. Most exclusions are due to no 24-48h +creatinine measurement. ### 1.3 Shared Encoder Architecture @@ -350,14 +353,18 @@ For concrete rerun outputs, use the exported result artifacts under `results/` and the thesis W&B project summaries. Final thesis reruns should use the tmux launcher so W&B scoping stays -consistent: +consistent and launch provenance is pinned to a reviewed commit. Push the +reviewed commit or check out that exact hash on the VM, then pass it explicitly: -`WANDB_ENTITY= bash scripts/internal/launch_thesis_tmux.sh` +`WANDB_ENTITY= LAUNCH_COMMIT= bash scripts/internal/launch_thesis_tmux.sh` + +The launcher validates the 84-feature processed artifacts before starting tmux +and clears runtime tensor/hash-normalizer caches before cache warmup. If launching a sprint manually, include the thesis W&B target on every training -command: +command and pass the same commit hash for W&B config provenance: -`--project slices-thesis --revision thesis-v1 --entity ` +`--project slices-thesis --revision thesis-v1 --entity --launch-commit ` ### Upcoming: Sprint 10 — Extra Seeds 789, 1011 (630 runs) diff --git a/scripts/internal/launch_thesis_tmux.sh b/scripts/internal/launch_thesis_tmux.sh index f159cc2..f3225d9 100755 --- a/scripts/internal/launch_thesis_tmux.sh +++ b/scripts/internal/launch_thesis_tmux.sh @@ -7,6 +7,10 @@ WANDB_PROJECT="${WANDB_PROJECT:-slices-thesis}" WANDB_ENTITY="${WANDB_ENTITY:-}" REVISION="${REVISION:-thesis-v1}" REASON="${REASON:-thesis-final canonical-sprint11-baselines}" +LAUNCH_COMMIT="${LAUNCH_COMMIT:-}" +ALLOW_DIRTY="${ALLOW_DIRTY:-0}" +EXPECTED_FEATURES="${EXPECTED_FEATURES:-84}" +PURGE_RUNTIME_CACHES="${PURGE_RUNTIME_CACHES:-1}" PARALLEL_MAIN="${PARALLEL_MAIN:-4}" PARALLEL_APPENDIX="${PARALLEL_APPENDIX:-4}" BATCH_SIZE_FAIRNESS="${BATCH_SIZE_FAIRNESS:-64}" @@ -34,6 +38,106 @@ if [[ -z "$WANDB_ENTITY" ]]; then exit 1 fi +if [[ -z "$LAUNCH_COMMIT" ]]; then + LAUNCH_COMMIT="$(git -C "$REPO_ROOT" rev-parse --verify HEAD)" +fi + +if ! git -C "$REPO_ROOT" cat-file -e "$LAUNCH_COMMIT^{commit}" 2>/dev/null; then + echo "LAUNCH_COMMIT is not available in this checkout: $LAUNCH_COMMIT" >&2 + exit 1 +fi + +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 + +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 + + 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" +echo "Validating processed artifacts..." +( + cd "$REPO_ROOT" + validate_processed_artifacts +) +echo "Processed artifacts are ready." +purge_runtime_caches + mkdir -p "$LOG_DIR" main_sprints=(1 1b 1c 2 3 4 5 6 7 8 10 11) @@ -71,7 +175,7 @@ export_args=(--project "$WANDB_PROJECT" --entity "$WANDB_ENTITY") fairness_args=(--project "$WANDB_PROJECT" --entity "$WANDB_ENTITY") tag_args=(--project "$WANDB_PROJECT" --entity "$WANDB_ENTITY") -run_args+=(--revision "$REVISION" --reason "$REASON") +run_args+=(--revision "$REVISION" --reason "$REASON" --launch-commit "$LAUNCH_COMMIT") export_args+=(--revision "$REVISION" --output-dir "$RESULTS_DIR") if [[ -n "$REVISION" ]]; then tag_args+=(--revision "$REVISION") @@ -124,10 +228,17 @@ echo "Start: \$(date -Iseconds)" echo "W&B project: $(printf "%q" "$WANDB_PROJECT")" echo "W&B entity: $(printf "%q" "$WANDB_ENTITY")" echo "Revision: $(printf "%q" "$REVISION")" +echo "Launch commit: $(printf "%q" "$LAUNCH_COMMIT")" echo "Main sprints: ${main_sprints[*]}" echo "Fairness sprints: ${fairness_sprints[*]}" echo "Appendix sprints: ${appendix_sprints[*]:-none}" +current_commit="\$(git rev-parse --verify HEAD)" +if [[ "\$current_commit" != "$(printf "%q" "$LAUNCH_COMMIT")" ]]; then + echo "Launch commit mismatch: current=\$current_commit expected=$(printf "%q" "$LAUNCH_COMMIT")" >&2 + exit 1 +fi + uv sync --dev ${warmup_line} ${main_line} diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index dd198c5..36f54b6 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -29,6 +29,7 @@ uv run python scripts/internal/run_experiments.py tag --sprint 2 3 5 \\ --project slices-thesis --entity """ + from __future__ import annotations import argparse @@ -975,7 +976,7 @@ def apply_revision(runs: list[Run], revision: str, reason: str | None = None) -> # Insert rev- after the sprint prefix (e.g. s1_ → s1_rev-v2_) prefix = f"s{r.sprint}_" if old_id.startswith(prefix): - new_id = f"{prefix}rev-{revision}_{old_id[len(prefix):]}" + new_id = f"{prefix}rev-{revision}_{old_id[len(prefix) :]}" else: new_id = f"rev-{revision}_{old_id}" old_to_new[old_id] = new_id @@ -1015,6 +1016,16 @@ def apply_wandb_target( 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 + + # --------------------------------------------------------------------------- # State Management # --------------------------------------------------------------------------- @@ -1218,7 +1229,7 @@ def handle_signal(signum, frame): if run.run_type == "pretrain": encoder_path = Path(run.output_dir) / "encoder.pt" if not encoder_path.exists(): - print(f" FAILED {rid}: exit 0 but encoder.pt " f"missing {elapsed}") + print(f" FAILED {rid}: exit 0 but encoder.pt missing {elapsed}") set_run_status(state, rid, "failed", finished_at=now, exit_code=exit_code) _propagate_failure(rid, runs, state) continue @@ -1492,7 +1503,7 @@ def cmd_run(args): extra = [all_by_id[d] for d in deps_needed if d in all_by_id] runs = extra + runs if extra and not args.dry_run: - print(f"Including {len(extra)} dependency run(s) from earlier " f"sprints") + print(f"Including {len(extra)} dependency run(s) from earlier sprints") if not runs: print(f"No runs found for sprint(s): {', '.join(sprints)}") @@ -1502,10 +1513,13 @@ def cmd_run(args): if args.revision: runs = apply_revision(runs, args.revision, args.reason) runs = apply_wandb_target(runs, args.project, args.entity) + runs = apply_launch_commit(runs, args.launch_commit) print(f"Sprint(s) {', '.join(sprints)}: {len(runs)} runs") if args.revision: print(f"Revision: {args.revision}" + (f" ({args.reason})" if args.reason else "")) + if args.launch_commit: + print(f"Launch commit: {args.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}") @@ -1564,6 +1578,8 @@ def _retry_command_suggestion(args, *, include_failed: bool, include_skipped: bo cmd.extend(["--project", str(args.project)]) if args.entity: cmd.extend(["--entity", str(args.entity)]) + if args.launch_commit: + cmd.extend(["--launch-commit", str(args.launch_commit)]) cmd.extend(["--parallel", str(args.parallel)]) return shlex.join(cmd) @@ -1624,6 +1640,7 @@ def cmd_retry(args): else: all_runs = apply_revision(all_runs, args.revision, args.reason) all_runs = apply_wandb_target(all_runs, args.project, args.entity) + all_runs = apply_launch_commit(all_runs, args.launch_commit) state = load_state() recover_stale_running(state) @@ -1679,6 +1696,8 @@ def cmd_retry(args): 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 args.launch_commit: + print(f"Launch commit: {args.launch_commit}") return run_scheduler(runs_to_retry, state, args.parallel, dry_run=False) @@ -1858,6 +1877,12 @@ def main(): 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 (e.g. v2)") 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, @@ -1888,6 +1913,12 @@ def main(): p_retry.add_argument("--sprint", nargs="+", default=None, help="Scope retry to sprint(s)") 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, diff --git a/scripts/setup_and_extract.sh b/scripts/setup_and_extract.sh index a3855ef..debafd6 100755 --- a/scripts/setup_and_extract.sh +++ b/scripts/setup_and_extract.sh @@ -10,6 +10,7 @@ # ./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 @@ -32,6 +33,8 @@ error() { echo -e "${RED}[ERROR]${NC} $1"; } # Parse arguments # --------------------------------------------------------------------------- SKIP_DEPS=false +FORCE_PROCESSED=false +EXPECTED_FEATURES="${EXPECTED_FEATURES:-84}" REQUESTED_DATASETS=() BASE_DATASETS=() @@ -50,9 +53,10 @@ append_base_dataset() { for arg in "$@"; do case "$arg" in --skip-deps) SKIP_DEPS=true ;; + --force-processed) FORCE_PROCESSED=true ;; miiv|eicu|combined) REQUESTED_DATASETS+=("$arg") ;; *) error "Unknown argument: $arg" - echo "Usage: $0 [--skip-deps] [miiv] [eicu] [combined]" + echo "Usage: $0 [--skip-deps] [--force-processed] [miiv] [eicu] [combined]" exit 1 ;; esac done @@ -85,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 # --------------------------------------------------------------------------- @@ -182,11 +251,20 @@ for ds in "${BASE_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" @@ -194,14 +272,10 @@ for ds in "${BASE_DATASETS[@]}"; do fi # --- Dataset preparation (splits + normalization) --- - if [ -f "$processed_dir/splits.yaml" ] && \ - [ -f "$processed_dir/normalization_stats.yaml" ]; then - info "Splits & normalization already exist: $processed_dir (skipping preparation)" - else - info "Running dataset preparation for $ds..." - uv run python scripts/preprocessing/prepare_dataset.py dataset="$ds" - info "Dataset preparation complete" - fi + 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 @@ -211,19 +285,36 @@ if [ "$BUILD_COMBINED" = true ]; then processed_dir="data/processed/combined" - if [ -f "$processed_dir/timeseries.parquet" ] && \ - [ -f "$processed_dir/static.parquet" ] && \ - [ -f "$processed_dir/labels.parquet" ] && \ - [ -f "$processed_dir/metadata.yaml" ] && \ + if processed_artifacts_exist "$processed_dir" && \ [ -f "$processed_dir/splits.yaml" ] && \ [ -f "$processed_dir/normalization_stats.yaml" ]; then - info "Combined dataset already exists and is prepared: $processed_dir (skipping)" + 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 "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 diff --git a/src/slices/data/datamodule.py b/src/slices/data/datamodule.py index f02d2ff..6504a2a 100644 --- a/src/slices/data/datamodule.py +++ b/src/slices/data/datamodule.py @@ -178,6 +178,8 @@ def __init__( # 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] = [] @@ -211,6 +213,7 @@ def setup(self, stage: Optional[str] = None) -> None: 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, @@ -221,10 +224,12 @@ def setup(self, stage: Optional[str] = None) -> None: ) # Save full train indices for normalization BEFORE subsampling. - # Normalization must always use the full training split to avoid - # noisy stats at low label fractions and pretrain/finetune mismatch. + # 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) - normalization_train_indices = list(self.full_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: @@ -232,16 +237,17 @@ def setup(self, stage: Optional[str] = None) -> None: self.train_indices, self.label_fraction, self.seed ) - # Create dataset with FULL training indices for normalization. - # self.train_indices (possibly subsampled) is used by train_dataloader - # to create a Subset for optimization. + # 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=normalization_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 @@ -283,11 +289,13 @@ def setup(self, stage: Optional[str] = None) -> None: # 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: @@ -420,6 +428,7 @@ def get_split_info(self) -> Dict[str, Any]: # 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, @@ -433,6 +442,11 @@ def get_split_info(self) -> Dict[str, Any]: "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, } diff --git a/src/slices/data/dataset.py b/src/slices/data/dataset.py index 0fead51..e8002c2 100644 --- a/src/slices/data/dataset.py +++ b/src/slices/data/dataset.py @@ -113,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: @@ -127,8 +128,13 @@ def __init__( 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) @@ -141,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() @@ -171,6 +182,8 @@ def __init__( 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: @@ -246,7 +259,7 @@ def _load_data(self) -> None: Uses a two-phase approach for efficiency: 1. Load raw tensors from cache (or extract from parquet on cache miss) - 2. Filter excluded stays, then normalize using run-specific train_indices + 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 @@ -303,17 +316,30 @@ def _load_data(self) -> None: self.n_features, ) - # Phase 2: Filter excluded stays, then normalize + # 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, + ) - # Apply stay exclusion filtering (task-specific, e.g. missing labels) + # 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) - timeseries_tensor = timeseries_tensor[idx_tensor] - masks_tensor = masks_tensor[idx_tensor] + 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 @@ -321,24 +347,12 @@ def _load_data(self) -> None: 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") - else: - self.stay_ids = all_stay_ids # Build stay_id -> index mapping logger.debug("Building stay_id index mapping") 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 existing normalization stats (for reproducibility) - cached_stats = load_normalization_stats(self.data_dir, self.train_indices, self.normalize) - - # Apply normalization and imputation - self._precompute_tensors( - precomputed_tensors=(timeseries_tensor, masks_tensor), - train_indices=self.train_indices, - cached_stats=cached_stats, - ) - # Pre-compute labels and static features self._precompute_labels_and_static() diff --git a/src/slices/data/preparation.py b/src/slices/data/preparation.py index 021f0cc..e6d36af 100644 --- a/src/slices/data/preparation.py +++ b/src/slices/data/preparation.py @@ -166,7 +166,7 @@ def prepare_processed_dataset( if not processed_dir.exists(): raise FileNotFoundError( - f"Processed directory not found: {processed_dir}\n" "Run extraction before preparation." + f"Processed directory not found: {processed_dir}\nRun extraction before preparation." ) print("\n1. Loading metadata...") @@ -225,6 +225,8 @@ def prepare_processed_dataset( ) 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() diff --git a/src/slices/data/splits.py b/src/slices/data/splits.py index 9f2d880..b3ad2fb 100644 --- a/src/slices/data/splits.py +++ b/src/slices/data/splits.py @@ -190,9 +190,7 @@ def filter_stays_with_missing_labels( # 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. " f"Available: {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: @@ -345,7 +343,15 @@ def compute_patient_level_splits( 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], + 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. @@ -374,7 +380,12 @@ def compute_patient_level_splits( Returns: Tuple of (train_indices, val_indices, test_indices, static_df, labels_df, - all_stay_ids, filtered_stay_ids, excluded_stay_ids). + 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...") @@ -401,6 +412,12 @@ def compute_patient_level_splits( 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( @@ -415,6 +432,14 @@ def compute_patient_level_splits( ) 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, @@ -424,14 +449,11 @@ def compute_patient_level_splits( all_stay_ids, stay_ids, excluded_stay_ids, + normalization_train_indices, ) logger.info("Computing full-cohort patient-level splits...") - # Get stay_id -> patient_id mapping - logger.debug("Building stay-to-patient mapping") - stay_to_patient = dict(zip(static_df["stay_id"].to_list(), static_df["patient_id"].to_list())) - # 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} @@ -501,6 +523,13 @@ def compute_patient_level_splits( 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) @@ -523,6 +552,7 @@ def compute_patient_level_splits( all_stay_ids, stay_ids, excluded_stay_ids, + normalization_train_indices, ) diff --git a/src/slices/data/tensor_cache.py b/src/slices/data/tensor_cache.py index 6dccb75..7b842f6 100644 --- a/src/slices/data/tensor_cache.py +++ b/src/slices/data/tensor_cache.py @@ -20,7 +20,8 @@ import yaml logger = logging.getLogger(__name__) -_CACHE_FINGERPRINT_VERSION = "2026-04-07" +_CACHE_FINGERPRINT_VERSION = "2026-04-22" +_NORMALIZATION_STATS_VERSION = "full-cohort-train-v2" def _fingerprint_payload(payload: Dict[str, Any]) -> str: @@ -112,7 +113,7 @@ def _compute_split_hash(train_indices: List[int], normalize: bool) -> str: Used to key normalization stats files so concurrent runs with different splits write to different files instead of overwriting each other. """ - content = f"{sorted(train_indices)}|{normalize}" + content = f"{_NORMALIZATION_STATS_VERSION}|{sorted(train_indices)}|{normalize}" return hashlib.md5(content.encode()).hexdigest()[:12] @@ -243,6 +244,8 @@ def save_normalization_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, diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 0830391..6ce4292 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -282,6 +282,8 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: if len(tag) > 64: tag = tag[:61] + "..." tags.append(tag) + if cfg.get("launch_commit") is not None: + tags.append(f"commit:{str(cfg.launch_commit)[:12]}") model_size = cfg.get("model_size") if model_size is not None: tags.append(f"model_size:{model_size}") @@ -433,8 +435,7 @@ def report_and_validate_train_label_support( print("\n Train label support:") print( - " - Dataset / task / seed / fraction: " - f"{dataset} / {task_name} / {seed} / {label_fraction}" + f" - Dataset / task / seed / fraction: {dataset} / {task_name} / {seed} / {label_fraction}" ) print( f" - Full train split: {full_pos} positive, {full_neg} negative, " diff --git a/tests/test_dataset_datamodule.py b/tests/test_dataset_datamodule.py index 24fa414..3f60c4e 100644 --- a/tests/test_dataset_datamodule.py +++ b/tests/test_dataset_datamodule.py @@ -8,7 +8,6 @@ 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 @@ -1116,6 +1115,50 @@ def test_cached_full_cohort_splits_are_filtered_within_patient_sets(self, mock_e 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 @@ -1134,6 +1177,7 @@ def test_ssl_train_disjoint_from_task_filtered_eval_splits_for_runner_matrix(sel ssl_stay_ids, _, _, + _, ) = compute_patient_level_splits( data_dir, task_name=None, @@ -1158,6 +1202,7 @@ def test_ssl_train_disjoint_from_task_filtered_eval_splits_for_runner_matrix(sel _, downstream_stay_ids, _, + _, ) = compute_patient_level_splits( data_dir, task_name=task_name, @@ -1610,9 +1655,9 @@ 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_tensor_cache_has_warnings_import(self): """tensor_cache.py should have warnings imported at module level. From c725bfea1a50a120565dc3b0971ce2d780319851 Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 22 Apr 2026 13:16:28 -0400 Subject: [PATCH 097/121] Harden export and fairness integrity Exclude HP-ablation provenance from inherited label-efficiency and transfer export rows, export LOS fairness metrics, validate checkpoint provenance, tighten fairness skip-existing completeness checks, and keep optional ICU mortality outside the primary thesis export unless explicitly requested. Behavioral implications: publication exports now fail on checkpoint provenance warnings unless --allow-incomplete is used; extension task exports require --include-extension-tasks. --- scripts/eval/evaluate_fairness.py | 49 +++++- scripts/export_results.py | 214 ++++++++++++++++++++++++-- src/slices/eval/fairness_evaluator.py | 15 +- tests/test_evaluate_fairness.py | 82 ++++++++-- tests/test_export_results.py | 163 ++++++++++++++++++++ tests/test_fairness_evaluator.py | 22 +++ 6 files changed, 514 insertions(+), 31 deletions(-) diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index 579edb2..2929d2a 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -62,6 +62,23 @@ CORE_SPRINTS = ["1", "2", "3", "4", "5", "6", "7", "8", "7p", "10", "11", "12", "13"] DEFAULT_PHASES = ["finetune", "supervised", "baseline"] DEFAULT_PROTECTED_ATTRIBUTES = ["gender", "age_group", "race"] +BINARY_FAIRNESS_REQUIRED_METRICS = [ + "n_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", +] # --------------------------------------------------------------------------- @@ -168,6 +185,26 @@ def _has_finite_summary_value(summary: dict[str, Any], key: str) -> bool: return True +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 has_fairness_metrics(run, protected_attributes: list[str]) -> bool: """Check whether all requested dataset-appropriate fairness keys exist. @@ -179,16 +216,12 @@ def has_fairness_metrics(run, protected_attributes: list[str]) -> bool: 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}/" - if not _has_finite_summary_value(sm, f"{prefix}n_valid_groups"): - return False - has_task_metric = _has_finite_summary_value( - sm, - f"{prefix}worst_group_auroc", - ) or _has_finite_summary_value(sm, f"{prefix}worst_group_mse") - if not has_task_metric: - return False + for metric_name in required_metrics: + if not _has_finite_summary_value(sm, f"{prefix}{metric_name}"): + return False return True except Exception: return False diff --git a/scripts/export_results.py b/scripts/export_results.py index 14b7d76..cde9d85 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -92,6 +92,12 @@ EXPECTED_FIXED_SEEDS = {42, 123, 456, 789, 1011} CAPACITY_PILOT_LABEL_FRACTIONS = {0.01, 0.1, 0.5} CAPACITY_PILOT_PARADIGMS = {"mae", "supervised"} +PRIMARY_THESIS_TASKS = { + "mortality_24h", + "mortality_hospital", + "aki_kdigo", + "los_remaining", +} # For core runs: the sprint tag is not part of the config identity. # lr/mask_ratio excluded because they're determined by protocol (redundant) or @@ -157,7 +163,9 @@ "test/r2", # Universal "test/loss", - # Fairness (populated by scripts/eval/evaluate_fairness.py) + # Fairness (populated by scripts/eval/evaluate_fairness.py). + # Numeric fairness/* summary keys are also extracted dynamically so per-group + # LOS regression summaries are not dropped by this canonical aggregate list. "fairness/gender/worst_group_auroc", "fairness/gender/worst_group_auprc", "fairness/gender/auroc_gap", @@ -165,6 +173,11 @@ "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/age_group/worst_group_auroc", "fairness/age_group/worst_group_auprc", "fairness/age_group/auroc_gap", @@ -172,6 +185,11 @@ "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/race/worst_group_auroc", "fairness/race/worst_group_auprc", "fairness/race/auroc_gap", @@ -179,6 +197,11 @@ "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", ] VAL_METRICS = [ @@ -188,6 +211,7 @@ ] ALL_METRICS = TEST_METRICS + VAL_METRICS +PERFORMANCE_TEST_METRICS = [metric for metric in TEST_METRICS if metric.startswith("test/")] # Canonical model variants used in the benchmark matrix. MODEL_VARIANTS = { @@ -444,6 +468,7 @@ def _is_label_efficiency_inherited_row(row: dict) -> bool: return ( math.isclose(label_fraction, 1.0, rel_tol=0.0, abs_tol=1e-9) + and not _has_hp_ablation_provenance(row) and row.get("phase") in {"finetune", "supervised"} and row.get("paradigm") in {"mae", "jepa", "contrastive", "supervised"} and row.get("model_size") == "default" @@ -460,6 +485,7 @@ def _is_transfer_inherited_row(row: dict) -> bool: return ( row.get("phase") == "finetune" + and not _has_hp_ablation_provenance(row) and row.get("paradigm") in {"mae", "jepa", "contrastive"} and row.get("dataset") in {"miiv", "eicu"} and row.get("protocol") == "B" @@ -469,6 +495,16 @@ def _is_transfer_inherited_row(row: dict) -> bool: ) +def _has_hp_ablation_provenance(row: dict) -> bool: + """Return whether a row carries upstream HP-ablation provenance.""" + subtype = row.get("experiment_subtype") + return ( + subtype in {"lr_ablation", "mask_ablation"} + or row.get("upstream_pretrain_lr") is not None + or row.get("upstream_pretrain_mask_ratio") is not None + ) + + def _is_hp_anchor_inherited_row(row: dict) -> bool: """Return True for default LR/mask rows inherited into HP ablations.""" try: @@ -589,6 +625,15 @@ def _infer_model_size(config: dict) -> str: 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], @@ -632,14 +677,15 @@ def extract_run( # Extract metrics from summary metrics = {} for key in metric_keys: - val = summary.get(key, None) - if val is not None and isinstance(val, (int, float)): - if math.isnan(val) or math.isinf(val): - metrics[key] = float("nan") - else: - metrics[key] = float(val) - else: - metrics[key] = float("nan") + 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)) + and not isinstance(value, bool) + ): + metrics[key] = _metric_value_or_nan(value) config_sprint = str(config.get("sprint", "")) sprint_str = _select_export_sprint(config_sprint, tags, requested_sprints) @@ -872,13 +918,46 @@ def build_per_seed_df( return df +def filter_thesis_tasks( + per_seed_df: pd.DataFrame, + include_extension_tasks: bool = False, +) -> pd.DataFrame: + """Apply the primary thesis task whitelist unless extension export is requested.""" + if include_extension_tasks or per_seed_df.empty or "task" not in per_seed_df.columns: + return per_seed_df + + mask = per_seed_df["task"].isin(PRIMARY_THESIS_TASKS) + excluded = per_seed_df.loc[~mask, "task"].dropna() + if not excluded.empty: + task_counts = excluded.value_counts().to_dict() + print( + " Excluding extension task rows from primary thesis export: " + + ", ".join(f"{task}={count}" for task, count in sorted(task_counts.items())), + file=sys.stderr, + ) + return per_seed_df.loc[mask].copy().reset_index(drop=True) + + +def _metric_columns(df: pd.DataFrame) -> list[str]: + """Return static and dynamically discovered numeric metric columns.""" + static_metrics = [c for c in ALL_METRICS if c in df.columns] + dynamic_fairness = [ + c + for c in sorted(df.columns) + if c.startswith("fairness/") + and c not in static_metrics + and pd.api.types.is_numeric_dtype(df[c]) + ] + return static_metrics + dynamic_fairness + + def _aggregate_group(df: pd.DataFrame, fingerprint_cols: list[str]) -> pd.DataFrame: """Aggregate a subset of runs by the given fingerprint columns. Returns a DataFrame with one row per unique fingerprint, containing mean/std/min/max of metrics and lists of run IDs/seeds/groups. """ - metric_cols = [c for c in ALL_METRICS if c in df.columns] + metric_cols = _metric_columns(df) # Sentinel for NaN in groupby work = df.copy() @@ -1423,7 +1502,7 @@ def validate( ) # Check for runs with no test metrics at all - test_cols = [c for c in TEST_METRICS if c in per_seed_df.columns] + test_cols = [c for c in PERFORMANCE_TEST_METRICS if c in per_seed_df.columns] if test_cols: all_nan = per_seed_df[test_cols].isna().all(axis=1) if all_nan.any(): @@ -1462,6 +1541,31 @@ def validate( + f" — missing {primary_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"{c}={row.get(c)}" + for c in [ + "wandb_run_id", + "paradigm", + "dataset", + "task", + "seed", + "phase", + ] + if c in row + ) + + f" — {reason}" + ) + # Summary by experiment type n_runs = len(per_seed_df) n_configs = len(aggregated_df) @@ -1477,6 +1581,78 @@ def validate( return warnings +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.""" + phase = row.get("phase") + if phase not in EVAL_PHASES: + return False + paradigm = str(row.get("paradigm", "")).lower() + if paradigm == "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 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 + + # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- @@ -1557,6 +1733,14 @@ def parse_args() -> argparse.Namespace: "publication exports fail closed by default." ), ) + parser.add_argument( + "--include-extension-tasks", + action="store_true", + help=( + "Include optional extension tasks such as ICU mortality. By default, " + f"primary thesis exports are restricted to {sorted(PRIMARY_THESIS_TASKS)}." + ), + ) args = parser.parse_args() if not args.revision: @@ -1600,8 +1784,16 @@ def main(): requested_sprints=args.sprint, allow_extraction_failures=args.allow_extraction_failures, ) + 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) diff --git a/src/slices/eval/fairness_evaluator.py b/src/slices/eval/fairness_evaluator.py index 3d9c816..65d742b 100644 --- a/src/slices/eval/fairness_evaluator.py +++ b/src/slices/eval/fairness_evaluator.py @@ -135,9 +135,10 @@ 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 def _encode_attribute( @@ -173,7 +174,7 @@ def _encode_attribute( 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, patient_ids @@ -430,6 +431,7 @@ def _evaluate_regression( 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 @@ -450,14 +452,21 @@ 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, diff --git a/tests/test_evaluate_fairness.py b/tests/test_evaluate_fairness.py index 87512d9..8fbed9d 100644 --- a/tests/test_evaluate_fairness.py +++ b/tests/test_evaluate_fairness.py @@ -7,6 +7,31 @@ import pytest +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}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 test_default_phases_include_baseline(): mod = importlib.import_module("scripts.eval.evaluate_fairness") @@ -82,12 +107,10 @@ def test_has_fairness_metrics_requires_requested_attribute_completeness(): mod = importlib.import_module("scripts.eval.evaluate_fairness") run = SimpleNamespace( - config={"dataset": "miiv"}, + config={"dataset": "miiv", "task": {"task_type": "binary"}}, summary_metrics={ - "fairness/gender/n_valid_groups": 2, - "fairness/gender/worst_group_auroc": 0.71, - "fairness/age_group/n_valid_groups": 4, - "fairness/age_group/worst_group_auroc": 0.69, + **_binary_fairness_summary("gender", 0.71), + **_binary_fairness_summary("age_group", 0.69), }, ) @@ -99,13 +122,54 @@ def test_has_fairness_metrics_ignores_race_for_eicu(): mod = importlib.import_module("scripts.eval.evaluate_fairness") run = SimpleNamespace( - config={"dataset": "eicu"}, + config={"dataset": "eicu", "task": {"task_type": "binary"}}, + summary_metrics={ + **_binary_fairness_summary("gender", 0.71), + **_binary_fairness_summary("age_group", 0.69), + }, + ) + + 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, - "fairness/age_group/n_valid_groups": 4, - "fairness/age_group/worst_group_auroc": 0.69, }, ) - assert mod.has_fairness_metrics(run, ["gender", "age_group", "race"]) is True + assert mod.has_fairness_metrics(run, ["gender"]) is False + + +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"), + }, + ) + 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"), + }, + ) + + assert mod.has_fairness_metrics(complete, ["gender", "age_group"]) is True + assert mod.has_fairness_metrics(partial, ["gender", "age_group"]) is False diff --git a/tests/test_export_results.py b/tests/test_export_results.py index aeff9e6..1147fb6 100644 --- a/tests/test_export_results.py +++ b/tests/test_export_results.py @@ -173,6 +173,40 @@ def test_build_per_seed_df_adds_unscoped_inherited_memberships_for_thesis_famili assert ("run-2", "8", "hp_ablation") in memberships +def test_build_per_seed_df_rejects_hp_ablation_from_label_and_transfer_inheritance(): + 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={ + "sprint": "10", + "dataset": "miiv", + "paradigm": "mae", + "seed": 42, + "output_dir": ( + "outputs/sprint10/" 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=["sprint:10", "sprint:6", "sprint:7", "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_ablation"} + assert set(df["sprint"]) == {"10"} + 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") @@ -394,6 +428,135 @@ def test_validate_warns_when_evaluation_row_missing_primary_metric(): 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={ + "sprint": "1", + "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=["sprint:1", "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") + + 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, + }, + ] + ) + + 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_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": "mortality", "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") diff --git a/tests/test_fairness_evaluator.py b/tests/test_fairness_evaluator.py index d312958..50b38d1 100644 --- a/tests/test_fairness_evaluator.py +++ b/tests/test_fairness_evaluator.py @@ -151,6 +151,10 @@ def test_singleton_prediction_dimension_is_squeezed(self): 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: @@ -185,6 +189,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.""" From be9d6efb0075d7395231a7c61d3b6bdd9577969f Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 22 Apr 2026 13:13:49 -0400 Subject: [PATCH 098/121] Harden baseline and protocol runner semantics Resolve XGBoost balanced weights to a logged numeric native scale_pos_weight, add explicit finetune Protocol A/B overlays, preserve requested sprint ordering, and emit documented W&B ablation tags. Clarify SMART as finite graded attention bias, restrict contrastive mask diagnostics to eligible timesteps, and pass task type into inline fairness so regression fairness remains task-safe. Behavioral implications: XGBoost provenance now records resolved weighting, manual finetune launches can select protocol=a or protocol=b, Sprint 10 can be requested before Sprint 7p, and masking-budget audits exclude invalid empty timesteps. --- configs/finetune.yaml | 12 +-- configs/protocol/a.yaml | 20 +++++ configs/protocol/b.yaml | 21 +++++ configs/xgboost.yaml | 4 +- docs/internal/EXPERIMENT_PLAN.md | 3 +- scripts/internal/run_experiments.py | 15 +++- scripts/training/xgboost_baseline.py | 58 ++++++++++--- src/slices/models/encoders/smart.py | 10 ++- src/slices/models/pretraining/contrastive.py | 89 ++++++++++++++------ src/slices/training/utils.py | 31 +++++-- tests/test_contrastive_objective.py | 8 +- tests/test_fixes.py | 81 ++++++++++++++++++ tests/test_smart_encoder.py | 28 ++++++ tests/test_training_utils.py | 73 ++++++++++++++++ 14 files changed, 392 insertions(+), 61 deletions(-) create mode 100644 configs/protocol/a.yaml create mode 100644 configs/protocol/b.yaml diff --git a/configs/finetune.yaml b/configs/finetune.yaml index 49e0c18..5656feb 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=a or protocol=b for thesis Protocol A/B-safe manual launches. # # Example usage: -# # Linear probing (frozen encoder, default) -# uv run python scripts/training/finetune.py dataset=miiv checkpoint=outputs/encoder.pt +# # Protocol A: strict 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=a # -# # Full finetuning (train encoder + head) -# uv run python scripts/training/finetune.py dataset=miiv checkpoint=outputs/encoder.pt training.freeze_encoder=false +# # Protocol B: 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=b # # # 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 @@ -24,6 +25,7 @@ defaults: - model: transformer - tasks: mortality_24h - _self_ + - optional protocol: null # Project info project_name: slices diff --git a/configs/protocol/a.yaml b/configs/protocol/a.yaml new file mode 100644 index 0000000..aebc935 --- /dev/null +++ b/configs/protocol/a.yaml @@ -0,0 +1,20 @@ +# @package _global_ +# Protocol A: strict linear probe for SSL representation quality. + +protocol: A + +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/protocol/b.yaml b/configs/protocol/b.yaml new file mode 100644 index 0000000..9ff4a70 --- /dev/null +++ b/configs/protocol/b.yaml @@ -0,0 +1,21 @@ +# @package _global_ +# Protocol B: full finetune for practical downstream utility. + +protocol: B + +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/xgboost.yaml b/configs/xgboost.yaml index 9d1e5c2..8ef7dd9 100644 --- a/configs/xgboost.yaml +++ b/configs/xgboost.yaml @@ -45,7 +45,9 @@ xgboost: colsample_bytree: 0.8 min_child_weight: 1 early_stopping_rounds: 20 - scale_pos_weight: balanced # auto-compute from label distribution for binary tasks + # "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 Sprint 11 runs many baselines in parallel # Data loading diff --git a/docs/internal/EXPERIMENT_PLAN.md b/docs/internal/EXPERIMENT_PLAN.md index 1bca3a0..28e6cb6 100644 --- a/docs/internal/EXPERIMENT_PLAN.md +++ b/docs/internal/EXPERIMENT_PLAN.md @@ -402,7 +402,7 @@ command and pass the same commit hash for W&B config provenance: **Matrix**: 2 baselines × 3 datasets × 4 tasks × 5 seeds = 120 full-label runs + 240 label-efficiency runs = **360 runs** -**XGBoost**: 8 summary statistics per input feature (mean, std, min, max, first, last, obs_count, obs_fraction) → 896 tabular features. Uses `ICUDataModule` for identical data splits. +**XGBoost**: 8 summary statistics per input feature (mean, std, min, max, first, last, obs_count, obs_fraction) → 896 tabular features. Uses `ICUDataModule` for identical data splits. For binary tasks, `scale_pos_weight=balanced` resolves to XGBoost's native `n_negative / n_positive` ratio and is logged numerically in the run config/summary; this is intentionally not the neural sqrt-balanced weighting scheme. **GRU-D**: d_model=64, 1 GRU layer, same downstream training protocol as supervised Transformer (lr=3e-4, wd=0.05, patience=10, label smoothing, class weights). Trains on GPU via `supervised.py --config-name gru_d`. @@ -478,6 +478,7 @@ parsing run names. | `paradigm:mae` / `paradigm:jepa` / `paradigm:contrastive` / `paradigm:supervised` / `paradigm:xgboost` / `paradigm:gru_d` | Paradigm | | `task:mortality_24h` / etc. | Downstream task | | `ablation:label-efficiency` / `ablation:transfer` | Ablation type | +| `protocol:A` / `protocol:B` | Downstream evaluation family. Supervised-from-scratch and classical baselines align with Protocol B for comparison; use `phase:*` to distinguish SSL finetunes from non-SSL baselines. | | `sprint:N` | Execution sprint | | `revision:` | Thesis rerun revision for fairness/export filtering | | `seed:42` / `seed:123` / `seed:456` / `seed:789` / `seed:1011` | Random seed | diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index 36f54b6..29aceda 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -960,6 +960,17 @@ def generate_all_runs() -> list[Run]: return builder.build_all() +def filter_runs_for_sprints(all_runs: list[Run], sprints: list[str]) -> list[Run]: + """Filter runs while preserving the caller's requested sprint order.""" + sprint_order: dict[str, int] = {} + for sprint in sprints: + sprint_order.setdefault(sprint, len(sprint_order)) + + indexed_runs = [(idx, run) for idx, run in enumerate(all_runs) if run.sprint in sprint_order] + indexed_runs.sort(key=lambda item: (sprint_order[item[1].sprint], item[0])) + return [run for _, run in indexed_runs] + + def apply_revision(runs: list[Run], revision: str, reason: str | None = None) -> list[Run]: """Post-process runs to inject revision into IDs, output dirs, and overrides. @@ -1489,7 +1500,7 @@ def cmd_run(args): # Filter to requested sprints sprints = [str(s) for s in args.sprint] - runs = [r for r in all_runs if r.sprint in sprints] + runs = filter_runs_for_sprints(all_runs, sprints) # Also include dependency runs from earlier sprints (pretrain reuse) run_ids = {r.id for r in runs} @@ -1715,7 +1726,7 @@ def cmd_warmup(args): # Filter to requested sprints sprints = [str(s) for s in args.sprint] - runs = [r for r in all_runs if r.sprint in sprints] + runs = filter_runs_for_sprints(all_runs, sprints) # Also include dependency runs run_ids = {r.id for r in runs} diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index 55796e5..cb4caf7 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -35,6 +35,29 @@ def _xgboost_eval_metric(task_type: str) -> str: return "mae" if task_type == "regression" else "aucpr" +def _add_wandb_tag(tags: list[str], tag: str) -> None: + """Append a W&B tag once, preserving caller order.""" + 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) @@ -61,16 +84,17 @@ 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("sprint") is not None: - tags.append(f"sprint:{cfg.sprint}") + _add_wandb_tag(tags, f"sprint:{cfg.sprint}") if cfg.get("revision") is not None: - tags.append(f"revision:{cfg.revision}") + _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] + "..." - tags.append(tag) + _add_wandb_tag(tags, tag) if cfg.get("label_fraction", 1.0) < 1.0: - tags.append(f"label_fraction:{cfg.label_fraction}") + _add_wandb_tag(tags, f"label_fraction:{cfg.label_fraction}") + _add_wandb_tag(tags, "ablation:label-efficiency") return tags or None @@ -171,22 +195,30 @@ def main(cfg: DictConfig) -> None: "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: - # Auto-compute scale_pos_weight if requested - scale_pos_weight = xgb_cfg.get("scale_pos_weight", None) - if scale_pos_weight == "balanced": - n_pos = y_train.sum() - n_neg = len(y_train) - n_pos - scale_pos_weight = n_neg / max(n_pos, 1) - print(f" scale_pos_weight (balanced): {scale_pos_weight:.2f}") + 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=scale_pos_weight, + scale_pos_weight=resolved_scale_pos_weight, eval_metric=_xgboost_eval_metric(task_type), ) @@ -220,6 +252,8 @@ def main(cfg: DictConfig) -> None: metrics["test/accuracy"] = accuracy_score(y_test, y_pred_class) metrics["test/brier_score"] = brier_score_loss(y_test, y_pred_proba) metrics["test/ece"] = _binary_ece(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(): diff --git a/src/slices/models/encoders/smart.py b/src/slices/models/encoders/smart.py index 15cb058..dfb2252 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) diff --git a/src/slices/models/pretraining/contrastive.py b/src/slices/models/pretraining/contrastive.py index 7d6ed03..cf6bf50 100644 --- a/src/slices/models/pretraining/contrastive.py +++ b/src/slices/models/pretraining/contrastive.py @@ -227,6 +227,7 @@ def forward( full_2, effective_mask_1, effective_mask_2, + valid_timestep_mask, ) else: # 4b. Instance mode: mean-pool -> project -> instance NT-Xent @@ -236,7 +237,13 @@ def forward( 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) + loss, metrics = self._nt_xent_loss( + z1, + z2, + effective_mask_1, + effective_mask_2, + valid_timestep_mask, + ) return loss, metrics @@ -259,6 +266,30 @@ 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, @@ -286,6 +317,7 @@ def _temporal_nt_xent_loss( 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. @@ -298,6 +330,7 @@ def _temporal_nt_xent_loss( 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) @@ -313,6 +346,11 @@ def _temporal_nt_xent_loss( 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(), @@ -320,13 +358,10 @@ def _temporal_nt_xent_loss( "contrastive_pos_similarity": torch.tensor(0.0), "contrastive_temperature": temperature, "contrastive_n_timesteps": T, - "contrastive_n_visible_view1": ssl_mask_1.sum().item() / B, - "contrastive_n_visible_view2": ssl_mask_2.sum().item() / B, - "contrastive_n_masked_view1": (~ssl_mask_1).sum().item() / B, - "contrastive_n_masked_view2": (~ssl_mask_2).sum().item() / B, "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 @@ -358,6 +393,11 @@ def _temporal_nt_xent_loss( 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(), @@ -366,13 +406,10 @@ def _temporal_nt_xent_loss( "contrastive_pos_similarity": pos_sim, "contrastive_temperature": temperature, "contrastive_n_timesteps": T, - "contrastive_n_visible_view1": ssl_mask_1.sum().item() / B, - "contrastive_n_visible_view2": ssl_mask_2.sum().item() / B, - "contrastive_n_masked_view1": (~ssl_mask_1).sum().item() / B, - "contrastive_n_masked_view2": (~ssl_mask_2).sum().item() / B, "contrastive_n_overlap_tokens": N, "contrastive_n_overlap_per_sample": N / B, } + metrics.update(count_metrics) return loss, metrics @@ -382,6 +419,7 @@ def _nt_xent_loss( z2: torch.Tensor, ssl_mask_1: torch.Tensor, ssl_mask_2: torch.Tensor, + eligible_mask: torch.Tensor, ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: """Compute NT-Xent (Normalized Temperature-scaled Cross Entropy) loss. @@ -390,6 +428,7 @@ def _nt_xent_loss( z2: (B, proj_dim) L2-normalized projections from view 2. 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) @@ -402,6 +441,12 @@ def _nt_xent_loss( 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(), @@ -412,23 +457,17 @@ def _nt_xent_loss( "contrastive_effective_rank": torch.tensor(0.0, device=z1.device), "contrastive_temperature": self.config.temperature, "contrastive_n_timesteps": T, - "contrastive_n_visible_view1": valid_mask_1.sum().item() - / max(n_valid_samples, 1), - "contrastive_n_visible_view2": valid_mask_2.sum().item() - / max(n_valid_samples, 1), - "contrastive_n_masked_view1": (~valid_mask_1).sum().item() - / max(n_valid_samples, 1), - "contrastive_n_masked_view2": (~valid_mask_2).sum().item() - / max(n_valid_samples, 1), "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 @@ -477,12 +516,13 @@ def _nt_xent_loss( p = p / p.sum().clamp(min=1e-12) eff_rank = (-p * p.clamp(min=1e-7).log()).sum().exp() - # Timestep statistics + # Timestep statistics over eligible non-empty timesteps only. T = ssl_mask_1.shape[1] - n_vis_1 = ssl_mask_1.sum().item() - n_vis_2 = ssl_mask_2.sum().item() - n_masked_1 = (~ssl_mask_1).sum().item() - n_masked_2 = (~ssl_mask_2).sum().item() + count_metrics = self._timestep_count_metrics( + ssl_mask_1, + ssl_mask_2, + eligible_mask, + ) metrics = { "contrastive_loss": loss.detach(), @@ -494,12 +534,9 @@ def _nt_xent_loss( "contrastive_effective_rank": eff_rank, "contrastive_temperature": temperature, "contrastive_n_timesteps": T, - "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_samples_used": B, "contrastive_n_samples_skipped": valid_samples.numel() - B, } + metrics.update(count_metrics) return loss, metrics diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 6ce4292..07c0630 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -264,6 +264,12 @@ def setup_finetune_callbacks(cfg: DictConfig, checkpoint_prefix: str = "finetune # ============================================================================= +def _add_wandb_tag(tags: list[str], tag: str) -> None: + """Append a W&B tag once, preserving caller order.""" + if tag not in tags: + tags.append(tag) + + def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: """Set up W&B experiment logger. @@ -274,35 +280,40 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: tags = list(cfg.logging.get("wandb_tags", [])) if cfg.get("sprint") is not None: - tags.append(f"sprint:{cfg.sprint}") + _add_wandb_tag(tags, f"sprint:{cfg.sprint}") if cfg.get("revision") is not None: - tags.append(f"revision:{cfg.revision}") + _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] + "..." - tags.append(tag) + _add_wandb_tag(tags, tag) if cfg.get("launch_commit") is not None: - tags.append(f"commit:{str(cfg.launch_commit)[:12]}") + _add_wandb_tag(tags, f"commit:{str(cfg.launch_commit)[:12]}") model_size = cfg.get("model_size") if model_size is not None: - tags.append(f"model_size:{model_size}") + _add_wandb_tag(tags, f"model_size:{model_size}") - # Add protocol tag for finetune runs (Protocol A = frozen, Protocol B = unfrozen) + # Add downstream protocol family tag. Supervised-from-scratch shares the + # Protocol B optimization budget; the phase tag distinguishes it from SSL. freeze_encoder = cfg.get("training", {}).get("freeze_encoder") if freeze_encoder is not None: protocol = "A" if freeze_encoder else "B" - tags.append(f"protocol:{protocol}") + _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: - tags.append(f"mask_ratio:{ssl_cfg.mask_ratio}") + _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: - tags.append(f"label_fraction:{label_fraction}") + _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, "ablation:transfer") tags = tags or None @@ -371,6 +382,7 @@ def run_fairness_evaluation( datamodule.test_dataloader(), device=model.device, ) + task_type = cfg.get("task", {}).get("task_type", "binary") evaluator = FairnessEvaluator( static_df=datamodule.dataset.static_df, @@ -378,6 +390,7 @@ 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) diff --git a/tests/test_contrastive_objective.py b/tests/test_contrastive_objective.py index 67bef93..eb0e704 100644 --- a/tests/test_contrastive_objective.py +++ b/tests/test_contrastive_objective.py @@ -188,7 +188,7 @@ def test_two_views_encoded(self, encoder, contrastive_config): assert metrics["contrastive_n_visible_view2"] > 0 def test_empty_timesteps_are_excluded_from_visible_counts(self, encoder, contrastive_config): - """Contrastive views should ignore hours with no observed variables.""" + """Contrastive diagnostics should ignore hours with no observed variables.""" obj = ContrastiveObjective(encoder, contrastive_config) B, T, D = 2, 8, 10 @@ -201,6 +201,12 @@ def test_empty_timesteps_are_excluded_from_visible_counts(self, encoder, contras 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.""" diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 79cb137..e5b3572 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -1515,6 +1515,7 @@ def test_xgboost_wandb_tags_include_revision_scope(self): assert "sprint:11" 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_build_per_seed_df_keeps_distinct_hp_ablation_configs(self): @@ -1663,6 +1664,54 @@ def test_apply_wandb_target_injects_project_and_entity(self): assert result[0].extra_overrides["logging.wandb_project"] == "slices-thesis" assert result[0].extra_overrides["logging.wandb_entity"] == "hannes-ill" + def test_cmd_run_respects_requested_sprint_order(self, monkeypatch): + import scripts.internal.run_experiments as runner + + sprint7p = runner.Run( + id="s7p_supervised_mortality_24h_miiv_seed42", + sprint="7p", + run_type="supervised", + paradigm="supervised", + dataset="miiv", + seed=42, + output_dir="outputs/sprint7p/supervised_mortality_24h_miiv_seed42", + task="mortality_24h", + ) + sprint10 = runner.Run( + id="s10_supervised_mortality_24h_miiv_seed789", + sprint="10", + run_type="supervised", + paradigm="supervised", + dataset="miiv", + seed=789, + output_dir="outputs/sprint10/supervised_mortality_24h_miiv_seed789", + task="mortality_24h", + ) + scheduled = {} + + monkeypatch.setattr(runner, "generate_all_runs", lambda: [sprint7p, sprint10]) + 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( + sprint=["10", "7p"], + revision=None, + reason=None, + project=None, + entity=None, + parallel=1, + dry_run=True, + ) + + runner.cmd_run(args) + + assert [run.sprint for run in scheduled["runs"]] == ["10", "7p"] + def test_pretrain_and_finetune_commands_resume_from_last_checkpoint(self, tmp_path): from scripts.internal.run_experiments import Run @@ -1814,6 +1863,7 @@ def __init__(self, **kwargs): "output_dir": "outputs/sprint7p/example", "sprint": "7p", "model_size": "medium", + "source_dataset": "eicu", "label_fraction": 0.1, "training": {"freeze_encoder": False}, "logging": { @@ -1833,6 +1883,8 @@ def __init__(self, **kwargs): assert captured["kwargs"]["name"] == "s7p_finetune_miiv_mae_mortality_24h_seed42_medium" assert captured["kwargs"]["group"] == "finetune_miiv_mae_mortality_24h_medium_frac01" assert "model_size:medium" in captured["kwargs"]["tags"] + assert "ablation:label-efficiency" in captured["kwargs"]["tags"] + assert "ablation:transfer" in captured["kwargs"]["tags"] def test_transfer_finetune_command_propagates_source_dataset(self): from scripts.internal.run_experiments import Run @@ -1925,6 +1977,28 @@ def test_protocol_b_finetune_command_keeps_task_mlp_head(self): assert "training.freeze_encoder=false" in cmd assert "task.head_type=linear" not in cmd + def test_manual_finetune_protocol_configs_encode_safe_defaults(self): + protocol_a = OmegaConf.load("configs/protocol/a.yaml") + protocol_b = OmegaConf.load("configs/protocol/b.yaml") + + assert protocol_a.protocol == "A" + assert protocol_a.training.freeze_encoder is True + assert protocol_a.training.max_epochs == 50 + assert protocol_a.training.early_stopping_patience == 10 + assert protocol_a.optimizer.lr == pytest.approx(1.0e-4) + assert protocol_a.task.head_type == "linear" + assert list(protocol_a.task.hidden_dims) == [] + assert protocol_a.task.dropout == 0.0 + + assert protocol_b.protocol == "B" + assert protocol_b.training.freeze_encoder is False + assert protocol_b.training.max_epochs == 100 + assert protocol_b.training.early_stopping_patience == 10 + assert protocol_b.optimizer.lr == pytest.approx(3.0e-4) + assert protocol_b.task.head_type == "mlp" + assert list(protocol_b.task.hidden_dims) == [64] + assert protocol_b.task.dropout == 0.3 + class TestExperimentRunnerRetry: """Tests for retry scoping and dependency preservation.""" @@ -2358,6 +2432,13 @@ def test_xgboost_config_caps_cpu_parallelism(self): 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.""" diff --git a/tests/test_smart_encoder.py b/tests/test_smart_encoder.py index bcb5abe..9de5e44 100644 --- a/tests/test_smart_encoder.py +++ b/tests/test_smart_encoder.py @@ -203,6 +203,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).""" diff --git a/tests/test_training_utils.py b/tests/test_training_utils.py index 1465f25..3d7dcb5 100644 --- a/tests/test_training_utils.py +++ b/tests/test_training_utils.py @@ -21,6 +21,7 @@ build_optimizer, build_scheduler, report_and_validate_train_label_support, + run_fairness_evaluation, save_encoder_checkpoint, setup_finetune_callbacks, ) @@ -353,6 +354,78 @@ def test_training_task_configs_cover_package_tasks(self): ) +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.""" From 3303d2663e3b17bfc56df27272b3951d75b29fae Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 22 Apr 2026 13:19:35 -0400 Subject: [PATCH 099/121] Migrate thesis runs to experiment classes Replace historical sprint metadata with explicit experiment_class and experiment_subtype metadata across configs, W&B tagging, the runner, export, fairness evaluation, launcher, and internal experiment docs. The final launch, export, fairness, retry, and warmup paths now use --experiment-class and no longer expose inherited sprint tagging. The generated final matrix remains 2305 runs including SMART and 2170 runs excluding SMART, with revision-scoped class metadata on every generated command. --- configs/finetune.yaml | 9 +- configs/gru_d.yaml | 7 +- configs/pretrain.yaml | 9 +- configs/supervised.yaml | 9 +- configs/xgboost.yaml | 9 +- docs/internal/EXPERIMENT_PLAN.md | 777 +++------- scripts/eval/evaluate_fairness.py | 44 +- scripts/export_results.py | 1341 ++++++---------- scripts/internal/launch_thesis_tmux.sh | 147 +- scripts/internal/run_experiments.py | 1937 ++++++++++-------------- scripts/training/xgboost_baseline.py | 25 +- src/slices/training/utils.py | 29 +- tests/test_config_schemas.py | 30 +- tests/test_evaluate_fairness.py | 10 +- tests/test_export_results.py | 414 ++--- tests/test_fixes.py | 392 +++-- 16 files changed, 1976 insertions(+), 3213 deletions(-) diff --git a/configs/finetune.yaml b/configs/finetune.yaml index 5656feb..0659b66 100644 --- a/configs/finetune.yaml +++ b/configs/finetune.yaml @@ -37,10 +37,9 @@ dataset: miiv # miiv | eicu | combined # Paradigm — auto-detected from checkpoint in finetune.py; override with paradigm=jepa etc. paradigm: mae -# Experiment group (stored as a sprint tag for historical W&B compatibility) -sprint: null - -# Revision tracking (for rerunning experiment groups with changed settings) +# Final rerun corpus metadata +experiment_class: null +experiment_subtype: null revision: null rerun_reason: null @@ -146,7 +145,7 @@ logging: use_wandb: true wandb_project: ${project_name} wandb_entity: null # Set to your W&B username/team - run_name: s${sprint}_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" diff --git a/configs/gru_d.yaml b/configs/gru_d.yaml index 2c265aa..685f2b2 100644 --- a/configs/gru_d.yaml +++ b/configs/gru_d.yaml @@ -25,8 +25,9 @@ dataset: miiv # Paradigm paradigm: gru_d -# Experiment group (stored as a sprint tag for historical W&B compatibility) -sprint: null +# Final rerun corpus metadata +experiment_class: null +experiment_subtype: null revision: null rerun_reason: null @@ -83,7 +84,7 @@ logging: use_wandb: true wandb_project: ${project_name} wandb_entity: null - run_name: s${sprint}_gru_d_${dataset}_${task.task_name}_seed${seed} + run_name: ${experiment_class}_gru_d_${dataset}_${task.task_name}_seed${seed} wandb_group: gru_d_${dataset}_${task.task_name} wandb_tags: - "phase:baseline" diff --git a/configs/pretrain.yaml b/configs/pretrain.yaml index d49ac99..2b4dd06 100644 --- a/configs/pretrain.yaml +++ b/configs/pretrain.yaml @@ -22,10 +22,9 @@ project_name: slices # Dataset identifier — drives data paths via ricu.yaml interpolation dataset: miiv # miiv | eicu | combined -# Experiment group (stored as a sprint tag for historical W&B compatibility) -sprint: null - -# Revision tracking (for rerunning experiment groups with changed settings) +# Final rerun corpus metadata +experiment_class: null +experiment_subtype: null revision: null rerun_reason: null @@ -97,7 +96,7 @@ logging: use_wandb: true wandb_project: ${project_name} wandb_entity: null # Set to your W&B username/team - run_name: s${sprint}_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" diff --git a/configs/supervised.yaml b/configs/supervised.yaml index 367cfb8..1cb40e6 100644 --- a/configs/supervised.yaml +++ b/configs/supervised.yaml @@ -33,10 +33,9 @@ dataset: miiv # miiv | eicu | combined # Paradigm — always "supervised" for this config paradigm: supervised -# Experiment group (stored as a sprint tag for historical W&B compatibility) -sprint: null - -# Revision tracking (for rerunning experiment groups with changed settings) +# Final rerun corpus metadata +experiment_class: null +experiment_subtype: null revision: null rerun_reason: null @@ -129,7 +128,7 @@ logging: use_wandb: true wandb_project: ${project_name} wandb_entity: null - run_name: s${sprint}_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" diff --git a/configs/xgboost.yaml b/configs/xgboost.yaml index 8ef7dd9..f7eda4b 100644 --- a/configs/xgboost.yaml +++ b/configs/xgboost.yaml @@ -22,8 +22,9 @@ dataset: miiv # Paradigm paradigm: xgboost -# Experiment group (stored as a sprint tag for historical W&B compatibility) -sprint: null +# Final rerun corpus metadata +experiment_class: null +experiment_subtype: null revision: null rerun_reason: null @@ -48,7 +49,7 @@ xgboost: # "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 Sprint 11 runs many baselines in parallel + n_jobs: 4 # cap per-process CPU use when many baselines run in parallel # Data loading training: @@ -66,7 +67,7 @@ logging: use_wandb: true wandb_project: ${project_name} wandb_entity: null - run_name: s${sprint}_xgboost_${dataset}_${task.task_name}_seed${seed} + run_name: ${experiment_class}_xgboost_${dataset}_${task.task_name}_seed${seed} wandb_group: xgboost_${dataset}_${task.task_name} wandb_tags: - "phase:baseline" diff --git a/docs/internal/EXPERIMENT_PLAN.md b/docs/internal/EXPERIMENT_PLAN.md index 28e6cb6..18ff862 100644 --- a/docs/internal/EXPERIMENT_PLAN.md +++ b/docs/internal/EXPERIMENT_PLAN.md @@ -1,559 +1,238 @@ -# SLICES Experiment Plan +# SLICES Final Experiment Plan -## Overview +This plan defines the final thesis rerun corpus for SLICES. The corpus is +organized by scientific experiment class, not by historical launch order. -This document defines the full experiment matrix for the SLICES benchmark. The goal is to answer: +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. -> 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? +## Benchmark Contract -Secondary questions addressed through ablations: +SLICES compares SSL paradigms for sparse, irregular ICU time series under +controlled conditions. -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? -4. Does increasing model capacity widen the SSL-supervised gap on low-label mortality prediction? -5. Does a temporal contrastive objective (TS2Vec) materially change the contrastive-family conclusion? +Controlled SSL objectives: ---- +- `mae` +- `jepa` +- `contrastive` -## 1. Experiment Matrix +Controlled SSL objectives share: -### 1.1 Independent Variables +- 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 -| Variable | Levels | Notes | -|----------|--------|-------| -| **Dataset** | MIMIC-IV, eICU, Combined | Combined = pooled pretraining corpus | -| **Paradigm** | MAE (reconstruction), JEPA (self-distillation), Contrastive, Supervised | Supervised = no pretraining baseline; final thesis reporting also includes GRU-D and XGBoost as contextual baselines | -| **Downstream Task** | mortality_24h, mortality_hospital, aki_kdigo, los_remaining | 3 classification + 1 regression | -| **Seed** | 42, 123, 456, 789, 1011 | Core matrix starts with 3 seeds; Sprint `10` extends the thesis corpus to 5 seeds, and Sprints `7p`, `11`, `12`, and `13` use the same 5-seed footing | +`ts2vec_extension` is a temporal-contrastive extension. It provides context for +the contrastive result but is not one of the three core thesis vertices. -**Formal thesis scope**: The controlled comparison remains the three-way MAE/JEPA/Contrastive benchmark plus supervised. The final thesis matrix also includes GRU-D and XGBoost as contextual baselines, canonicalized in Sprint `11` rather than duplicated across the early SSL/supervised sprints. Two additional thesis experiments are now tracked explicitly: Sprint `7p` is a focused capacity study (larger MAE + supervised encoders on MIIV mortality_24h), and Sprint `13` is a temporal-contrastive extension (TS2Vec) run under the same downstream A/B protocol. Sprint `12` (SMART) remains appendix-only because it changes the encoder family. - -### 1.2 Controlled Variables (Held Constant) - -| Variable | Value | Rationale | -|----------|-------|-----------| -| Observation window | 24 hours | Current benchmark input window | -| AKI prediction window | Hours 24–48 (forward-looking) | Prevents leakage from creatinine values used in label construction | -| Min stay | 24 hours | Ensures a full observation window before downstream labeling | -| Splits | 70/15/15 train/val/test | Patient-level, no leakage | -| Imputation | Normalize-then-zero-fill | Eliminates imputation as confound | -| Downstream head | Protocol A: linear; Protocol B/supervised: MLP, hidden_dims=[64] | Same head within each protocol across all paradigms | -| Precision | fp32 | Avoids bf16 numerical issues with small model | -| Weight decay | 0.05 | Same regularization across all downstream runs | - -**AKI labelable cohort**: `aki_kdigo` is evaluated on a task-specific subset -with ascertainable creatinine-based KDIGO labels. A stay requires non-null -creatinine in the 0-24h baseline window and at least one creatinine measurement -in the 24-48h prediction window; stays without an ascertainable AKI label are -excluded from AKI supervised training/evaluation but remain available for SSL -pretraining and non-AKI tasks. In the current 84-feature processed artifacts, -metadata-derived counts are 56,403/74,829 MIIV stays, 88,239/132,686 eICU -stays, and 144,642/207,515 combined stays. Final tables should derive full -denominators from each processed `metadata.yaml` (`n_stays`) rather than -copying static denominators into prose. Most exclusions are due to no 24-48h -creatinine measurement. - -### 1.3 Shared Encoder Architecture - -All three SSL paradigms share the same encoder architecture, removing tokenization as a confounding variable: - -| Paradigm | Tokenization | Encoder | Masking | Rationale | -|----------|-------------|---------|---------|-----------| -| MAE | Timestep-level (obs-aware) | Transformer (d=64, L=2, H=4, obs_aware=True) | Random timestep masking | Reconstruct masked timestep features | -| JEPA | Timestep-level (obs-aware) | Transformer (d=64, L=2, H=4, obs_aware=True) | Block masking (3 contiguous segments) | Predict in latent space; block masking prevents trivial interpolation | -| Contrastive | Timestep-level (obs-aware) | Transformer (d=64, L=2, H=4, obs_aware=True) | Two complementary masked views (instance mode, zero overlap) | Instance-level NT-Xent: mean-pool each view → align sequence embeddings | -| Supervised | Timestep-level (obs-aware) | Transformer (d=64, L=2, H=4, obs_aware=True) | N/A | Same encoder as SSL for fair comparison | - -**Obs-aware tokenization**: Each timestep token is produced by an MLP projection of `concat(values, obs_mask)` → `d_model`. This encodes both the observed values and the missingness pattern, avoiding the "mostly zeros" problem of naively feeding sparse timestep vectors through a linear projection. With ~22 observed features per timestep on average, the MLP maps ~44 non-zero inputs (22 values + 22 mask bits) to 64-dim tokens — no information bottleneck. - -**Design rationale**: An earlier observation-level design (1 token per observed value, using `ObservationTransformerEncoder`) produced ~1660 tokens per sample at 20% observation density, causing O(N²) attention to exceed L4 GPU memory. Timestep-level tokenization reduces this to 48 tokens (~35× reduction) while preserving all information through the obs-aware MLP. This also strengthens the controlled comparison by making all SSL paradigms share identical tokenization. - ---- - -## 2. Training Protocol - -### 2.1 Pretraining - -| Setting | MAE | JEPA | Contrastive | -|---------|-----|------|-------------| -| Epochs | 100 | 100 | 100 | -| Batch size | 256 | 256 | 256 | -| Learning rate | 1e-3 | 1e-3 | 1e-3 | -| Scheduler | Warmup cosine (10 warmup) | Warmup cosine (10 warmup) | Warmup cosine (10 warmup) | -| Mask ratio | 0.5 (random masking) | 0.5 (block masking) | 0.5 (complementary masks, zero overlap) | -| Gradient clipping | 1.0 | 1.0 | 1.0 | -| Early stopping | None (fixed schedule) | None (fixed schedule) | None (fixed schedule) | -| Checkpoint | Last epoch + best val loss | Last epoch + best val loss | Last epoch + best val loss | - -**Fixed-schedule training**: All paradigms train for the full epoch budget without early stopping, following the convention in self-distillation literature (I-JEPA, DINO, BYOL) where validation loss is unreliable as a stopping criterion — JEPA's val loss rises as representations get richer, not worse. Fixed schedules ensure identical gradient step budgets across paradigms. The cosine LR schedule decays to eta_min=1e-6, providing natural convergence. Both last-epoch and best-val-loss encoders are saved; the last-epoch encoder is used for all primary results, with best-val-loss as a robustness check in the appendix. - -### 2.2 Downstream Evaluation (Two Protocols) - -SSL encoders are evaluated under **two** protocols to answer different questions: - -#### Protocol A: Linear Probing (primary SSL comparison) - -Isolates representation quality — the only variable is the SSL pretraining objective. - -| 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 AUPRC (classification) or val MAE (regression) | -| Head | Single linear layer, dropout=0.0 | -| Gradient clipping | 1.0 | -| Label smoothing | 0.1 | -| Weight decay | 0.05 | -| Class weighting | sqrt(balanced) — square root of inverse frequency | - -**Rationale**: Linear probing answers "Which SSL objective produces the best representations?" by preventing the encoder from adapting to the downstream task and by removing nonlinear head capacity as a confound. - -#### Protocol B: Full Finetuning (practical utility) - -Measures how useful SSL pretraining is as weight initialization. - -| Setting | Value | -|---------|-------| -| Protocol | Full finetuning (freeze_encoder=false) | -| Epochs | 100 | -| Batch size | 64 | -| Learning rate | 3e-4 | -| Scheduler | Cosine decay (eta_min=1e-6) | -| Early stopping | Patience=10 on val AUPRC (classification) or val MAE (regression) | -| Head | MLP, hidden_dims=[64], dropout=0.3, ReLU | -| Gradient clipping | 1.0 | -| Label smoothing | 0.1 | -| Weight decay | 0.05 | -| Class weighting | sqrt(balanced) — square root of inverse frequency | - -**Rationale**: Full finetuning answers "Which SSL objective is most useful in practice as initialization?" Results may diverge from linear probing — an encoder with mediocre frozen representations can still provide excellent initialization. - -**Regularization rationale**: Early Sprint 1 runs showed severe overfitting across all paradigms (train AUROC ~0.98 by epoch 1–2, monotonically degrading val metrics). The regularization suite — reduced LR (1e-3 → 3e-4), increased dropout (0.1 → 0.3), higher weight decay (0.01 → 0.05), sqrt class weights, and label smoothing — was introduced to address this. These settings apply uniformly to all paradigms and the supervised baseline, preserving the controlled comparison. - -### 2.3 Supervised Baseline - -| Setting | Value | -|---------|-------| -| Epochs | 100 | -| Batch size | 64 | -| Learning rate | 3e-4 | -| Scheduler | Cosine decay (eta_min=1e-6) | -| Early stopping | Patience=10 on val AUPRC (classification) or val MAE (regression) | -| Encoder | Transformer (d=64, L=2, H=4, obs_aware=True), trained end-to-end | -| Head | MLP, hidden_dims=[64], dropout=0.3, ReLU | -| Gradient clipping | 1.0 | -| Label smoothing | 0.1 | -| Weight decay | 0.05 | -| Class weighting | sqrt(balanced) — square root of inverse frequency | - -**Same-architecture rationale**: The supervised baseline uses the identical encoder architecture and tokenization as the SSL paradigms to eliminate model capacity as a confounding variable. The only difference is the training procedure: supervised trains end-to-end on labeled data from random initialization, while SSL paradigms pretrain on unlabeled data and then evaluate via Protocol A (strict linear probe) and Protocol B (full finetune). Comparing supervised vs. Protocol B isolates the effect of SSL initialization; comparing SSL paradigms via Protocol A isolates representation quality. - ---- - -## 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. - -For `aki_kdigo`, prevalence and sample counts are reported on the labelable -cohort defined in Section 1.2, not on the full >=24h ICU stay cohort. - -### 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 5 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 | Transformer (obs_aware) | ssl=mae | -| P2 | MIMIC-IV | JEPA | Transformer (obs_aware) | ssl=jepa | -| P3 | MIMIC-IV | Contrastive | Transformer (obs_aware) | ssl=contrastive | -| P4 | eICU | MAE | Transformer (obs_aware) | ssl=mae dataset=eicu | -| P5 | eICU | JEPA | Transformer (obs_aware) | ssl=jepa dataset=eicu | -| P6 | eICU | Contrastive | Transformer (obs_aware) | ssl=contrastive dataset=eicu | -| P7 | Combined | MAE | Transformer (obs_aware) | ssl=mae dataset=combined | -| P8 | Combined | JEPA | Transformer (obs_aware) | ssl=jepa dataset=combined | -| P9 | Combined | Contrastive | Transformer (obs_aware) | ssl=contrastive dataset=combined | - -### 4.2 Phase 2: Downstream Evaluation (84 runs × 3 seeds = 252 runs) - -Each of the 9 pretrained encoders is evaluated on the **4-task primary thesis matrix** under **both** protocols: -- Protocol A (linear probe): 9 encoders × 4 tasks = 36 runs -- Protocol B (full finetune): 9 encoders × 4 tasks = 36 runs -- Supervised baseline: 3 datasets × 4 tasks = 12 runs - -The repo now also exposes an optional `mortality` ICU-mortality task, but it is not included in these run totals unless explicitly added as an extension. - -**Result table template** (one per dataset, two sub-tables per dataset): - -| | mortality_24h | mortality_hospital | aki_kdigo | los_remaining | -|---|---|---|---|---| -| **MAE** | AUPRC ± std | AUPRC ± std | AUPRC ± std | MAE ± std | -| **JEPA** | | | | | -| **Contrastive** | | | | | -| **Supervised** | (Protocol B only) | | | | - -### 4.3 Phase 3: Fairness Evaluation (on all downstream runs) - -No additional training. Compute fairness metrics by rerunning inference from the -same evaluation checkpoint provenance recorded for each downstream run. - ---- - -## 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**: Evaluate SSL under both downstream protocols plus the supervised Transformer baseline at low label fractions {1%, 5%, 10%, 25%, 50%}. The 100% endpoints are inherited from the main full-label runs rather than relaunched here. - -| Scope | Tasks | Label Fractions | Runs | -|-------|-------|----------------|------| -| Full sweep | mortality_24h (anchor task) | 1%, 5%, 10%, 25%, 50% | (3 SSL paradigms × 2 protocols + supervised) × 3 datasets × 5 fractions = 105 runs/seed | -| Trend check | Remaining 3 tasks | 10% | (3 SSL paradigms × 2 protocols + supervised) × 3 datasets × 3 tasks = 63 runs/seed | -| **Total** | | | **168 runs/seed × 3 seeds = 504 runs** | - -**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** (from main experiments, no extra runs): in-domain pretraining, supervised from scratch. - -### 5.3 Learning Rate Sensitivity - -**Question**: Are paradigm rankings robust to the shared learning rate? - -**Design**: Pretrain with LR ∈ {2e-4, 5e-4, 1e-3, 2e-3} on MIMIC-IV, finetune on mortality_24h. LR=1e-3 reused from main experiments. - -**Total**: 9 new pretraining + 9 new finetuning = 18 runs × 3 seeds = 54 runs - -### 5.4 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. 0.5 reused from main experiments. - -**Total**: 6 new pretraining + 6 new finetuning = 12 runs × 3 seeds = 36 runs - -### 5.5 Focused Capacity Study (Sprint 7p) - -**Question**: Are the modest SSL gains on low-label mortality prediction capacity-limited, or does supervised scale similarly when model size increases? - -**Design**: MIIV only, mortality_24h only, MAE + supervised only. Compare two larger encoder sizes (128d/4L and 256d/4L) against the inherited default-size baseline (64d/2L) from Sprints `6` and `10`, at label fractions {0.01, 0.1, 0.5}, across 5 seeds. - -**Total**: 2 sizes × [1 pretrain + 3 Protocol B finetunes + 3 Protocol A probes + 3 supervised] = 20 runs/seed × 5 seeds = **100 runs** - -### 5.6 Temporal Contrastive Extension (Sprint 13) - -**Question**: Does giving the contrastive family its natural temporal objective and augmentations overturn the conclusion from the instance-level contrastive baseline? - -**Design**: TS2Vec on all 3 datasets, the same 4-task primary thesis matrix, both Protocol A and Protocol B, 5 seeds. Same base encoder family and matched training budget as the controlled comparison. - -**Total**: 1 paradigm × 3 datasets × 5 seeds = 15 pretrains + 120 downstream probe/finetune runs = **135 runs** - ---- - -## 6. Run Summary - -| Phase | Description | Unit Runs | Planned Runs | Cumulative | -|-------|-------------|-----------|--------------|------------| -| Phase 1 | Pretraining | 9 | 27 | 27 | -| Phase 2 | Downstream (linear probe + full finetune + supervised) | 84 | 252 | 279 | -| Phase 3 | Fairness evaluation | 0 (eval only) | 0 | 279 | -| Ablation 5.1 | Label efficiency | 168 | 504 | 783 | -| Ablation 5.2 | Cross-dataset transfer | 24 | 72 | 855 | -| Ablation 5.3 | LR sensitivity | 18 | 54 | 909 | -| Ablation 5.4 | Mask ratio sensitivity | 12 | 36 | 945 | -| Sprint 7p | Focused capacity study (MAE + supervised, 5 seeds) | — | 100 | 1045 | -| Sprint 10 | Seeds 789, 1011 for Sprints 1–8 scope (SSL + supervised) | — | 630 | 1675 | -| Sprint 11 | Classical baselines (XGBoost + GRU-D), 5 seeds | — | 360 | 2035 | -| Sprint 13 | TS2Vec temporal contrastive extension, 5 seeds | — | 135 | 2170 | -| Sprint 12 | SMART external SSL reference, 5 seeds | — | 135 | 2305 | -| **Total** | | | **2305** | | - -**Planning note**: The formal thesis corpus through Sprint `13` totals **2170** runs. Sprint `12` adds **135** appendix-only SMART reference runs, bringing the full run budget to **2305**. - ---- - -## 7. Sprint Execution Order - -| Sprint | Description | Runs | Status | -|--------|-------------|------|--------| -| 1 | Pipeline validation (MIMIC, Protocol B + supervised) | 19 | COMPLETE | -| 1b | LR sensitivity (MIMIC, mortality_24h) | 18 | COMPLETE | -| 1c | Mask ratio sensitivity (MIMIC, mortality_24h) | 12 | COMPLETE | -| 2 | MIMIC Protocol A (completes MIMIC results table) | 12 | COMPLETE | -| 3 | eICU (both protocols + supervised) | 31 | COMPLETE | -| 4 | Combined (both protocols + supervised) | 31 | COMPLETE | -| 5 | Seeds 123, 456 for Sprints 1–4 | 186 | COMPLETE | -| 6 | Label efficiency ablation | 504 | COMPLETE | -| 7 | Cross-dataset transfer | 72 | COMPLETE | -| 8 | LR + mask ratio extra seeds | 60 | COMPLETE | -| 10 | Seeds 789, 1011 for Sprints 1–8 scope (SSL + supervised) | 630 | TODO | -| 7p | Focused capacity study (MIIV mortality_24h, MAE + supervised, 5 seeds) | 100 | TODO | -| 11 | Classical baselines (XGBoost + GRU-D), 5 seeds | 360 | TODO | -| 13 | TS2Vec temporal contrastive extension, both protocols, 5 seeds | 135 | TODO | -| 12 | SMART external SSL reference, 5 seeds | 135 | TODO | -| 9 | Fairness (eval only, runs last after all training) | 0 | TODO | - -For concrete rerun outputs, use the exported result artifacts under `results/` -and the thesis W&B project summaries. - -Final thesis reruns should use the tmux launcher so W&B scoping stays -consistent and launch provenance is pinned to a reviewed commit. Push the -reviewed commit or check out that exact hash on the VM, then pass it explicitly: - -`WANDB_ENTITY= LAUNCH_COMMIT= bash scripts/internal/launch_thesis_tmux.sh` - -The launcher validates the 84-feature processed artifacts before starting tmux -and clears runtime tensor/hash-normalizer caches before cache warmup. - -If launching a sprint manually, include the thesis W&B target on every training -command and pass the same commit hash for W&B config provenance: - -`--project slices-thesis --revision thesis-v1 --entity --launch-commit ` - -### Upcoming: Sprint 10 — Extra Seeds 789, 1011 (630 runs) - -**Motivation**: Go from 3 seeds to 5 seeds for all SSL and supervised experiments, improving statistical power (4 d.f. instead of 2 for variance estimates). - -**Scope**: Seeds 789 and 1011 for the full Sprints 1–8 scope. Classical baselines are not duplicated here; Sprint `11` remains their single canonical family. - -| Scope | Equivalent Sprint | Runs | -|-------|-------------------|------| -| Core experiments (pretrain + probe + finetune + supervised) | Sprints 1–4 via Sprint 5 | 186 | -| Label efficiency ablation (SSL + supervised) | Sprint 6 | 336 | -| Cross-dataset transfer | Sprint 7 | 48 | -| LR + mask ratio HP ablations | Sprint 8 | 60 | -| **Total** | | **630** | - -**Breakdown**: 48 pretrain (~30 min each, the bottleneck) + 510 finetune/probe (~3–5 min each) + 72 supervised (~3–5 min each). - -**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 10 --parallel 4 --project slices-thesis --revision thesis-v1 --entity ` (limited by GPU-bound pretraining) - -### Upcoming: Sprint 7p — Focused Capacity Study (100 runs) - -**Motivation**: The main benchmark suggests SSL gains are real but modest. Sprint `7p` tests whether that gap widens once model capacity increases, while keeping the question narrow enough to stay interpretable: MAE vs supervised, MIIV only, mortality_24h only, and low-label fractions where representation quality matters most. - -**Matrix**: 5 seeds × 2 model sizes × [1 pretrain + 3 Protocol A probes + 3 Protocol B finetunes + 3 supervised baselines] = **100 runs** - -**Comparison set**: Default-size baselines are inherited from Sprint `6` (seeds 42/123/456) and Sprint `10` (seeds 789/1011), so Sprint `7p` only launches the larger-capacity runs. - -**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 7p --parallel 4 --project slices-thesis --revision thesis-v1 --entity ` - -### Upcoming: Sprint 11 — Classical Baselines (360 runs) - -**Motivation**: Provide non-neural (XGBoost) and missingness-native neural (GRU-D) reference points. Without these, SSL AUROC numbers exist in a vacuum. XGBoost with engineered features is the perennially competitive clinical baseline (BAT 2025); GRU-D (Che et al. 2018) handles missing data natively via learnable decay. - -**Canonicalization note**: These baselines are part of the final thesis matrix, but they are launched only here. Earlier sprints intentionally contain only the controlled SSL + supervised benchmark so the classical baselines are reported once as a separate contextual comparison family. - -**Matrix**: 2 baselines × 3 datasets × 4 tasks × 5 seeds = 120 full-label runs + 240 label-efficiency runs = **360 runs** - -**XGBoost**: 8 summary statistics per input feature (mean, std, min, max, first, last, obs_count, obs_fraction) → 896 tabular features. Uses `ICUDataModule` for identical data splits. For binary tasks, `scale_pos_weight=balanced` resolves to XGBoost's native `n_negative / n_positive` ratio and is logged numerically in the run config/summary; this is intentionally not the neural sqrt-balanced weighting scheme. - -**GRU-D**: d_model=64, 1 GRU layer, same downstream training protocol as supervised Transformer (lr=3e-4, wd=0.05, patience=10, label smoothing, class weights). Trains on GPU via `supervised.py --config-name gru_d`. - -**Key comparisons enabled**: -1. XGBoost/GRU-D vs SSL Protocol B — absolute performance context -2. XGBoost/GRU-D vs supervised Transformer — is the Transformer architecture justified? -3. XGBoost cannot transfer across datasets — highlights SSL's cross-dataset transfer capability (Sprint 7) -4. XGBoost/GRU-D on label efficiency curves — direct comparison with SSL learning curves - -**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 11 --parallel 12 --project slices-thesis --revision thesis-v1 --entity ` (all runs independent, no pretraining dependency) - -### Upcoming: Sprint 13 — TS2Vec Temporal Contrastive Extension (135 runs) - -**Motivation**: Sprint `13` addresses the strongest foreseeable reviewer objection to the contrastive-family result: "instance-level contrastive was set up to fail." TS2Vec gives contrastive learning its natural temporal objective and augmentations while preserving the rest of the benchmark machinery. - -**Scope**: 3 datasets × 4 tasks × 2 protocols × 5 seeds, with matched downstream Protocol A/B evaluation and the same base encoder family. - -**Matrix**: 15 pretrains + 120 downstream probe/finetune runs = **135 runs** - -**Interpretation**: Sprint `13` is a formal thesis extension, not a replacement for the controlled MAE/JEPA/Contrastive triangle. It tests whether a better-instantiated contrastive family changes the conclusion. - -**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 13 --parallel 4 --project slices-thesis --revision thesis-v1 --entity ` - -### Upcoming: Sprint 12 — SMART External SSL Reference (135 runs) - -**Motivation**: SMART (NeurIPS 2024) is the most recent published SSL method for ICU time series. Including it provides an external SOTA reference point beyond the controlled three-way comparison. Addresses reviewer concern #7 ("no comparison to published clinical SSL methods"). - -**Important**: SMART uses a different encoder architecture (MART, d_model=32, per-variable query tokens, element-wise masking) and is therefore NOT part of the controlled comparison. Results go in the appendix as an external reference, not in the main results table. - -**Matrix**: 1 paradigm × 3 datasets × 4 tasks × 2 protocols × 5 seeds = 120 finetune + 15 pretrain = **135 runs** - -**Execution**: `uv run python scripts/internal/run_experiments.py run --sprint 12 --parallel 4 --project slices-thesis --revision thesis-v1 --entity ` - -### Upcoming: Sprint 9 — Fairness Analysis (0 extra runs) - -**Runs last**, after all training sprints are complete. Zero additional training — pure evaluation on existing test predictions. - -1. Compute fairness metrics on all downstream test predictions from Sprints 1–8, 7p, 10, 11, 12, and 13 -2. Run `scripts/eval/evaluate_fairness.py --project slices-thesis --revision thesis-v1 --entity ` for the thesis rerun corpus -3. Generate fairness tables and disparity plots -4. Protected attributes: sex (all datasets), age group (all), race/ethnicity (MIMIC-IV only) -5. The default standalone fairness sweep includes Sprint `11` classical baselines via `phase:baseline` - ---- - -## 8. W&B Tracking - -### Project Structure - -Legacy exploratory runs remain in W&B project `slices`. Final thesis reruns should log to `slices-thesis`, use the single revision tag `thesis-v1`, and set the W&B entity explicitly. - -### Run Naming Convention - -Run names are phase-specific and generated from the Hydra configs rather than a -single universal template. The current defaults use the `s${sprint}_...` -pattern from `configs/*.yaml`. - -Examples: -- `s10_pretrain_miiv_mae_seed42` -- `s10_finetune_eicu_jepa_mortality_24h_seed123` -- `s10_supervised_combined_aki_kdigo_seed456` -- `s11_xgboost_miiv_mortality_hospital_seed42` - -Downstream tooling should rely on tags and structured config fields rather than -parsing run names. - -### Tags - -| Tag | Purpose | -|-----|---------| -| `phase:pretrain` / `phase:finetune` / `phase:supervised` / `phase:baseline` | Training phase | -| `dataset:miiv` / `dataset:eicu` / `dataset:combined` | Dataset | -| `paradigm:mae` / `paradigm:jepa` / `paradigm:contrastive` / `paradigm:supervised` / `paradigm:xgboost` / `paradigm:gru_d` | Paradigm | -| `task:mortality_24h` / etc. | Downstream task | -| `ablation:label-efficiency` / `ablation:transfer` | Ablation type | -| `protocol:A` / `protocol:B` | Downstream evaluation family. Supervised-from-scratch and classical baselines align with Protocol B for comparison; use `phase:*` to distinguish SSL finetunes from non-SSL baselines. | -| `sprint:N` | Execution sprint | -| `revision:` | Thesis rerun revision for fairness/export filtering | -| `seed:42` / `seed:123` / `seed:456` / `seed:789` / `seed:1011` | Random seed | - -### Key Metrics to Log - -| Phase | Metrics | -|-------|---------| -| Pretraining | `train/loss`, `val/loss`, `train/gradient_steps`, `train/wall_clock_seconds` | -| Classification downstream | `val/auroc`, `val/auprc`, `test/auroc`, `test/auprc`, `test/brier_score`, `test/ece` | -| Regression downstream | `val/loss`, `test/mae`, `test/mse`, `test/r2` | -| Fairness | `fairness/gender/*`, `fairness/age_group/*`, `fairness/race/*` | - -### Baseline Inheritance Across Sprints - -Later sprints reuse runs from earlier sprints as comparison baselines. To enable filtering all relevant runs for a sprint in a single W&B query, baseline runs are tagged with later sprint tags using `run_experiments.py tag --sprint N --project slices-thesis --entity `. - -| Sprint | Inherits From | What's Inherited | -|--------|--------------|------------------| -| **1** | — | Self-contained (root sprint) | -| **1b** | 1 | Default-LR pretrain+finetune (mortality_24h) + supervised baseline | -| **1c** | 1 | Default-mask-ratio pretrain+finetune (mortality_24h) + supervised baseline | -| **2** | 1 | All Sprint 1 runs — completes the MIMIC results table | -| **3** | — | Self-contained (eICU has own pretrains + supervised) | -| **4** | — | Self-contained (Combined has own pretrains + supervised) | -| **5** | 1, 2, 3, 4 | All seed=42 runs (to aggregate with seeds 123, 456) | -| **6** | 1, 2, 3, 4, 5 | All fraction=1.0 finetune + supervised (learning curve right endpoints) | -| **7** | 1, 3, 5 | In-domain finetunes + supervised for MIMIC and eICU | -| **7p** | 6, 10 | Default-size label-efficiency baselines for seeds 42/123/456/789/1011 | -| **8** | 1, 1b, 1c, 5 | Seed=42 ablation runs + default-HP data points for seeds 123, 456 | -| **10** | — | Self-contained (new pretrains for seeds 789, 1011) | -| **11** | — | Self-contained (classical baselines, 5 seeds) | -| **13** | — | Self-contained (TS2Vec pretrains + both protocols, 5 seeds) | -| **12** | — | Self-contained (SMART pretrains + finetune, 5 seeds) | -| **9** | All | Evaluates test predictions from all training sprints (runs last) | - -**Usage**: `uv run python scripts/internal/run_experiments.py tag --sprint N --project slices-thesis --entity ` (idempotent). - ---- - -## 9. Expected Outputs - -### Tables for Paper - -1. **Table 1**: Main results — AUPRC (mean ± std) for each paradigm × dataset × task -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 - -- Per-seed results for transparency -- Training curves (loss vs steps) for all pretraining runs -- Hyperparameter sensitivity (if reviewer requests) -- Literature comparison table (see `docs/LIT_COMP.md`) - ---- - -## 10. Risk Mitigation - -| Risk | Mitigation | -|------|-----------| -| eICU lacks race/ethnicity data | Fairness analysis for race is MIMIC-only; document as limitation | -| Unequal training budgets across paradigms | Track gradient steps; normalize if >2x difference | -| Class imbalance affects metrics | Use AUPRC as primary metric; sqrt(balanced) class weighting; early stopping on val AUPRC (not val loss); report class ratios | -| Downstream overfitting | Addressed with regularization suite: dropout 0.3, LR 3e-4, weight decay 0.05, sqrt class weights, label smoothing 0.1 — applied uniformly to all paradigms | -| Compute budget exceeded | Sprint ordering ensures usable results at each checkpoint | -| Shared hyperparameters unfair to one paradigm | LR sensitivity (Sprint 1b) and mask ratio sensitivity (Sprint 1c) validate that rankings are robust to shared hyperparameter choices | -| Reviewer requests more seeds | Sprint `10` expands the main matrix to 5 seeds, and Sprints `7p`/`11`/`12`/`13` keep the same 4 d.f. footing for variance estimates | -| AKI label leakage | Fixed: forward-looking prediction window (hours 24–48) prevents model from seeing creatinine values used for KDIGO label construction | +`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` | `A` for linear probe, `B` for full finetune/full-training 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 remains scientifically equivalent to the prior +thesis matrix. + +| 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 | 840 | +| `cross_dataset_transfer` | MIMIC-IV to eICU and eICU to MIMIC-IV; SSL Protocol B; 5 seeds | 120 | +| `hp_robustness` | LR and mask-ratio robustness 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 | 360 | +| `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 | **2305** | + +The formal thesis corpus excluding `smart_external_reference` has 2170 launched +runs. + +## 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 \ + --entity +``` + +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 \ + --entity dummy \ + --dry-run +``` + +Expected dry-run count: 2305 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 +``` + +Revision filtering is mandatory to avoid mixing rerun corpora. + +## 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. + +## Validation Checklist + +Before launching final runs: + +- generated matrix has exactly 2305 launched runs including SMART +- generated matrix has exactly 2170 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 + +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/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index 2929d2a..8dd44a3 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -8,17 +8,17 @@ 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 --sprint/--paradigm/--dataset filters. +--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 sprint/dataset + # Scope to specific experiment class/dataset uv run python scripts/eval/evaluate_fairness.py \ --project slices-thesis --revision thesis-v1 --entity \ - --sprint 1 --dataset miiv + --experiment-class core_ssl_benchmark --dataset miiv # Preview which runs would be evaluated uv run python scripts/eval/evaluate_fairness.py \ @@ -59,7 +59,16 @@ # Constants # --------------------------------------------------------------------------- -CORE_SPRINTS = ["1", "2", "3", "4", "5", "6", "7", "8", "7p", "10", "11", "12", "13"] +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 = ["gender", "age_group", "race"] BINARY_FAIRNESS_REQUIRED_METRICS = [ @@ -107,7 +116,7 @@ def _retry(fn, max_retries=3, base_delay=5): def fetch_eval_runs( project: str, entity: Optional[str], - sprints: Optional[list[str]], + experiment_classes: Optional[list[str]], paradigms: Optional[list[str]], datasets: Optional[list[str]], phases: list[str], @@ -122,8 +131,8 @@ def fetch_eval_runs( filters: dict = {"state": "finished"} tag_filters: list[str] = [] - if sprints and len(sprints) == 1: - tag_filters.append(f"sprint:{sprints[0]}") + if experiment_classes and len(experiment_classes) == 1: + tag_filters.append(f"experiment_class:{experiment_classes[0]}") if paradigms and len(paradigms) == 1: tag_filters.append(f"paradigm:{paradigms[0]}") if datasets and len(datasets) == 1: @@ -140,7 +149,9 @@ def fetch_eval_runs( runs_iter = api.runs(path, filters=filters or {}, order="-created_at") # Client-side filtering for multi-value filters - sprint_set = set(sprints) if sprints and len(sprints) > 1 else None + experiment_class_set = ( + set(experiment_classes) if experiment_classes and len(experiment_classes) > 1 else None + ) paradigm_set = set(paradigms) if paradigms and len(paradigms) > 1 else None dataset_set = set(datasets) if datasets and len(datasets) > 1 else None phase_set = set(phases) if phases and len(phases) > 1 else None @@ -150,7 +161,10 @@ def fetch_eval_runs( for run in runs_iter: tags = set(run.tags) - if sprint_set and not any(f"sprint:{s}" in tags for s in sprint_set): + if experiment_class_set and not any( + f"experiment_class:{experiment_class}" in tags + for experiment_class in experiment_class_set + ): continue if paradigm_set and not any(f"paradigm:{p}" in tags for p in paradigm_set): continue @@ -609,9 +623,9 @@ def parse_args() -> argparse.Namespace: help="W&B entity (default: $WANDB_ENTITY)", ) parser.add_argument( - "--sprint", + "--experiment-class", nargs="+", - help="Filter to experiment group(s). Default: benchmark fairness groups", + 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)") @@ -696,16 +710,16 @@ def main() -> None: device = _resolve_device(args.device) log.info("Device: %s", device) - # Default to the benchmark fairness groups if not specified. - sprints = args.sprint or CORE_SPRINTS - log.info("Experiment groups: %s", sprints) + # 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, - sprints=sprints, + experiment_classes=experiment_classes, paradigms=args.paradigm, datasets=args.dataset, phases=args.phase, diff --git a/scripts/export_results.py b/scripts/export_results.py index cde9d85..4273a96 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -1,32 +1,23 @@ #!/usr/bin/env python3 -""" -Canonical results export pipeline for SLICES. - -Pulls all experiment runs from W&B, extracts config + test metrics, and produces -structured parquet files: +"""Canonical class-based results export pipeline for SLICES. - - results/per_seed_results.parquet — one row per W&B run/export membership - - results/aggregated_results.parquet — one row per unique config (~600 rows), - with mean/std/min/max across seeds - - results/statistical_tests.parquet — pairwise Wilcoxon + Bonferroni + - Cohen's d significance table - - results/ts2vec_vs_core_contrastive.parquet — direct TS2Vec vs core - contrastive comparisons +Pulls final rerun-corpus W&B runs, extracts config plus test metrics, and writes +publication-oriented parquet files: -Both files include wandb run IDs for traceability back to W&B. + - 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 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 \ - --sprint 1 --expected-seeds 42 123 456 - uv run python scripts/export_results.py \ - --project slices-thesis --revision thesis-v1 --entity \ - --paradigm mae jepa --dataset miiv - uv run python scripts/export_results.py \ - --project slices-thesis --revision thesis-v1 --entity \ - --output-dir results/sprint1 --sprint 1 + --experiment-class core_ssl_benchmark label_efficiency """ from __future__ import annotations @@ -52,58 +43,37 @@ # Constants # --------------------------------------------------------------------------- -# Historical sprint tag to benchmark experiment type mapping. -# Core groups share the same config and only add seeds; they merge across groups. -# Non-core groups have ablation-specific pretrain configs that are not visible -# in the finetune W&B config, so the sprint tag remains a grouping key. -EXPERIMENT_TYPES = { - "1": "core", - "2": "core", - "3": "core", - "4": "core", - "5": "core", - "10": "core", - "1b": "lr_ablation", - "1c": "mask_ablation", - "8": "hp_ablation", - "6": "label_efficiency", - "7": "transfer", - "7p": "capacity_pilot", - "11": "classical_baselines", - "12": "smart_reference", - "13": "temporal_contrastive", -} +EXPERIMENT_CLASSES = [ + "core_ssl_benchmark", + "label_efficiency", + "cross_dataset_transfer", + "hp_robustness", + "capacity_study", + "classical_baselines", + "ts2vec_extension", + "smart_external_reference", +] -CROSS_SPRINT_TYPES = {"core", "label_efficiency", "transfer"} -CONTEXTUAL_STAT_TYPES = { +CONTEXTUAL_STAT_CLASSES = { "classical_context_full", "classical_context_label_efficiency", } -FIXED_SEED_EXPERIMENT_TYPES = { - "core", - "label_efficiency", - "transfer", - "hp_ablation", - "capacity_pilot", - "classical_baselines", - "smart_reference", - "temporal_contrastive", -} + +FIXED_SEED_EXPERIMENT_CLASSES = set(EXPERIMENT_CLASSES) EXPECTED_FIXED_SEEDS = {42, 123, 456, 789, 1011} -CAPACITY_PILOT_LABEL_FRACTIONS = {0.01, 0.1, 0.5} -CAPACITY_PILOT_PARADIGMS = {"mae", "supervised"} -PRIMARY_THESIS_TASKS = { +CAPACITY_LABEL_FRACTIONS = {0.01, 0.1, 0.5} +CAPACITY_PARADIGMS = {"mae", "supervised"} +THESIS_TASKS = { "mortality_24h", "mortality_hospital", "aki_kdigo", "los_remaining", } -# For core runs: the sprint tag is not part of the config identity. -# lr/mask_ratio excluded because they're determined by protocol (redundant) or -# only in the pretrain config (NaN for finetune runs). -CORE_FINGERPRINT = [ +FINGERPRINT = [ + "experiment_class", "experiment_type", + "experiment_subtype", "paradigm", "dataset", "task", @@ -112,42 +82,17 @@ "model_size", "source_dataset", "phase", -] - -# HP ablations merge across groups, but must preserve the upstream pretrain -# config and subtype so LR and mask-ratio families cannot collapse together. -HP_ABLATION_FINGERPRINT = CORE_FINGERPRINT + [ - "experiment_subtype", "upstream_pretrain_lr", "upstream_pretrain_mask_ratio", ] -# For per-group experiment families, the sprint tag remains part of the identity. -ABLATION_FINGERPRINT = HP_ABLATION_FINGERPRINT + [ - "sprint", -] - FINGERPRINT_AUDIT_COLUMNS = [ - "experiment_type", - "experiment_subtype", - "paradigm", - "dataset", - "task", - "protocol", - "label_fraction", + *FINGERPRINT, "lr", "mask_ratio", - "model_size", - "source_dataset", - "phase", - "sprint", - "upstream_pretrain_lr", - "upstream_pretrain_mask_ratio", ] -# Metrics to extract from run summaries. TEST_METRICS = [ - # Binary classification "test/auroc", "test/auprc", "test/accuracy", @@ -157,15 +102,10 @@ "test/specificity", "test/brier_score", "test/ece", - # Regression (LOS) "test/mse", "test/mae", "test/r2", - # Universal "test/loss", - # Fairness (populated by scripts/eval/evaluate_fairness.py). - # Numeric fairness/* summary keys are also extracted dynamically so per-group - # LOS regression summaries are not dropped by this canonical aggregate list. "fairness/gender/worst_group_auroc", "fairness/gender/worst_group_auprc", "fairness/gender/auroc_gap", @@ -213,7 +153,6 @@ ALL_METRICS = TEST_METRICS + VAL_METRICS PERFORMANCE_TEST_METRICS = [metric for metric in TEST_METRICS if metric.startswith("test/")] -# Canonical model variants used in the benchmark matrix. MODEL_VARIANTS = { (64, 2): "default", (128, 4): "medium", @@ -234,8 +173,7 @@ "test/mse", } -# Phases that correspond to evaluation runs (not pretraining). -EVAL_PHASES = ["finetune", "supervised", "gru_d", "xgboost", "baseline"] +EVAL_PHASES = ["finetune", "supervised", "baseline"] # --------------------------------------------------------------------------- @@ -247,54 +185,39 @@ def fetch_all_runs( project: str, entity: str | None = None, state: str = "finished", - sprint: list[str] | None = None, + 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 filtering. - - Returns a list of wandb.Run objects. - """ + """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 - # Build server-side filters filters: dict = {} if state: filters["state"] = state - # Tag-based filters — use $all so runs must match ALL specified tags. tag_filters: list[str] = [] - - if sprint: - if len(sprint) == 1: - tag_filters.append(f"sprint:{sprint[0]}") - # Multiple sprints: fetch all, filter client-side (W&B $all requires ALL tags) - + if experiment_class and len(experiment_class) == 1: + tag_filters.append(f"experiment_class:{experiment_class[0]}") if paradigm and len(paradigm) == 1: tag_filters.append(f"paradigm:{paradigm[0]}") - if dataset and len(dataset) == 1: tag_filters.append(f"dataset:{dataset[0]}") - if phase and len(phase) == 1: tag_filters.append(f"phase:{phase[0]}") - if revision and len(revision) == 1: tag_filters.append(f"revision:{revision[0]}") - if tag_filters: filters["tags"] = {"$all": tag_filters} 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") - # Client-side filtering for multi-value filters - sprint_set = set(sprint) if sprint and len(sprint) > 1 else None + class_set = set(experiment_class) if experiment_class and len(experiment_class) > 1 else None paradigm_set = set(paradigm) if paradigm and len(paradigm) > 1 else None dataset_set = set(dataset) if dataset and len(dataset) > 1 else None phase_set = set(phase) if phase and len(phase) > 1 else None @@ -303,23 +226,16 @@ def fetch_all_runs( runs = [] for run in runs_iter: tags = set(run.tags) - - if sprint_set: - if not any(f"sprint:{s}" in tags for s in sprint_set): - continue - if paradigm_set: - if not any(f"paradigm:{p}" in tags for p in paradigm_set): - continue - if dataset_set: - if not any(f"dataset:{d}" in tags for d in dataset_set): - continue - if phase_set: - if not any(f"phase:{p}" in tags for p in phase_set): - continue - if revision_set: - if not any(f"revision:{r}" in tags for r in revision_set): - continue - + if class_set and not any(f"experiment_class:{value}" in tags for value in class_set): + continue + if paradigm_set and not any(f"paradigm:{value}" in tags for value in paradigm_set): + continue + if dataset_set and not any(f"dataset:{value}" in tags for value in dataset_set): + continue + if phase_set and not any(f"phase:{value}" in tags for value in phase_set): + continue + if revision_set and not any(f"revision:{value}" in tags for value in revision_set): + continue runs.append(run) print(f" Fetched {len(runs)} runs.", file=sys.stderr) @@ -332,44 +248,37 @@ def fetch_all_runs( 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] + 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): - """Retry a function with exponential backoff on timeout/connection errors.""" for attempt in range(max_retries + 1): try: return fn() - except Exception as e: - err_str = str(e).lower() + except Exception as exc: + err_str = str(exc).lower() is_retryable = any( - kw in err_str - for kw in ["timeout", "timed out", "connection", "429", "500", "502", "503"] + 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: {e!r}", + 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]: - """Load all data from a W&B run in one call (minimizes lazy-load API hits). - - Returns (config, summary_dict, run_id, run_url, run_name, tags, group, created_at). - """ - # Access all lazy-loaded properties together so retries cover them all config = dict(run.config) summary = dict(run.summary_metrics or {}) return ( @@ -384,10 +293,6 @@ def _load_run_data(run) -> tuple[dict, dict, str, str, str, list[str], str, str] ) -# Lookup tables for decoding run-name-encoded HP values. -# run_experiments.py encodes: str(v).replace(".", "") -# str(2e-4)="0.0002" -> "00002", str(5e-4)="0.0005" -> "00005", -# str(2e-3)="0.002" -> "0002", str(0.3)="0.3" -> "03", str(0.75)="0.75" -> "075" _LR_DECODE = { "00002": 2e-4, "00005": 5e-4, @@ -396,228 +301,68 @@ def _load_run_data(run) -> tuple[dict, dict, str, str, str, list[str], str, str] _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: - """Recover transfer provenance from explicit config or historical names.""" 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 _sprint_tags(tags: list[str]) -> list[str]: - """Return sprint values present in W&B tags, preserving tag order.""" - return [tag.split(":", 1)[1] for tag in tags if tag.startswith("sprint:")] - - -def _select_export_sprint( - config_sprint: str, - tags: list[str], - requested_sprints: list[str] | None = None, -) -> str: - """Choose the sprint family used for export classification. - - Historical inherited baselines keep their original `config.sprint`, but the - runner adds target sprint tags. When the export is scoped to target sprint - tags, classify the run in that requested target family. - """ - if not requested_sprints: - return config_sprint - - tagged_sprints = set(_sprint_tags(tags)) - requested_matches = [ - str(sprint) for sprint in requested_sprints if str(sprint) in tagged_sprints - ] - if not requested_matches: - return config_sprint - if config_sprint in requested_matches: - return config_sprint - return requested_matches[0] - - -def _is_capacity_pilot_inherited_row(row: dict) -> bool: - """Return True for default-size rows inherited into Sprint 7p comparisons.""" - try: - label_fraction = float(row.get("label_fraction", 1.0)) - except (TypeError, ValueError): - return False - - return ( - row.get("dataset") == "miiv" - and row.get("task") == "mortality_24h" - and row.get("paradigm") in CAPACITY_PILOT_PARADIGMS - and row.get("model_size") == "default" - and any( - math.isclose(label_fraction, allowed, rel_tol=0.0, abs_tol=1e-9) - for allowed in CAPACITY_PILOT_LABEL_FRACTIONS - ) - ) - - -def _is_label_efficiency_inherited_row(row: dict) -> bool: - """Return True for full-label rows inherited as 100% endpoints.""" - try: - label_fraction = float(row.get("label_fraction", 1.0)) - except (TypeError, ValueError): - return False - - return ( - math.isclose(label_fraction, 1.0, rel_tol=0.0, abs_tol=1e-9) - and not _has_hp_ablation_provenance(row) - and row.get("phase") in {"finetune", "supervised"} - and row.get("paradigm") in {"mae", "jepa", "contrastive", "supervised"} - and row.get("model_size") == "default" - and row.get("source_dataset") is None - ) - - -def _is_transfer_inherited_row(row: dict) -> bool: - """Return True for in-domain SSL rows inherited into transfer comparisons.""" - try: - label_fraction = float(row.get("label_fraction", 1.0)) - except (TypeError, ValueError): - return False - - return ( - row.get("phase") == "finetune" - and not _has_hp_ablation_provenance(row) - and row.get("paradigm") in {"mae", "jepa", "contrastive"} - and row.get("dataset") in {"miiv", "eicu"} - and row.get("protocol") == "B" - and math.isclose(label_fraction, 1.0, rel_tol=0.0, abs_tol=1e-9) - and row.get("model_size") == "default" - and row.get("source_dataset") is None - ) - - -def _has_hp_ablation_provenance(row: dict) -> bool: - """Return whether a row carries upstream HP-ablation provenance.""" - subtype = row.get("experiment_subtype") - return ( - subtype in {"lr_ablation", "mask_ablation"} - or row.get("upstream_pretrain_lr") is not None - or row.get("upstream_pretrain_mask_ratio") is not None - ) - - -def _is_hp_anchor_inherited_row(row: dict) -> bool: - """Return True for default LR/mask rows inherited into HP ablations.""" - try: - label_fraction = float(row.get("label_fraction", 1.0)) - except (TypeError, ValueError): - return False - - return ( - row.get("phase") == "finetune" - and row.get("paradigm") in {"mae", "jepa", "contrastive"} - and row.get("dataset") == "miiv" - and row.get("task") == "mortality_24h" - and row.get("protocol") == "B" - and math.isclose(label_fraction, 1.0, rel_tol=0.0, abs_tol=1e-9) - and row.get("model_size") == "default" - and row.get("source_dataset") is None - ) - - -def _is_inherited_export_row_for_target(row: dict, sprint: str) -> bool: - """Return whether a coarsely sprint-tagged source row belongs in target export.""" - if sprint == "6": - return _is_label_efficiency_inherited_row(row) - if sprint == "7": - return _is_transfer_inherited_row(row) - if sprint in {"1b", "1c", "8"}: - return _is_hp_anchor_inherited_row(row) - if sprint == "7p": - return _is_capacity_pilot_inherited_row(row) - return False - - -def _additional_inherited_export_sprints( - row: dict, - tags: list[str], - requested_sprints: list[str] | None = None, -) -> list[str]: - """Return inherited target sprints that need explicit export rows. - - Unscoped and multi-sprint thesis exports need both the source-family row and - target-family comparison row for inherited baselines. Sprint tags are - intentionally coarse, so target-specific eligibility keeps unrelated source - rows out of the target family. - """ - config_sprint = str(row.get("config_sprint", "")) - current_sprint = str(row.get("sprint", "")) - requested = {str(sprint) for sprint in requested_sprints} if requested_sprints else None - inherited: list[str] = [] - - for sprint in _sprint_tags(tags): - if requested is not None and sprint not in requested: - continue - if sprint in {config_sprint, current_sprint}: - continue - if _is_inherited_export_row_for_target(row, sprint): - inherited.append(sprint) - - return list(dict.fromkeys(inherited)) - - def _recover_pretrain_metadata( - run_name: str, config: dict + run_name: str, + config: dict, ) -> tuple[float | None, float | None, str | None]: - """Recover upstream pretrain metadata from config (new runs) or run name (historical). - - Returns (upstream_pretrain_lr, upstream_pretrain_mask_ratio, experiment_subtype). - """ - # New runs have explicit fields + """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", "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 - # Historical recovery: parse from run name / output_dir output_dir = config.get("output_dir", run_name) - lr_match = re.search(r"_lr([^_]+)", output_dir) if lr_match: - lr_str = lr_match.group(1) - up_lr = _LR_DECODE.get(lr_str) + up_lr = _LR_DECODE.get(lr_match.group(1)) if up_lr is not None: - subtype = "lr_ablation" + subtype = "lr_sensitivity" mr_match = re.search(r"_mask_ratio([^_]+)", output_dir) if mr_match: - mr_str = mr_match.group(1) - up_mr = _MR_DECODE.get(mr_str) + up_mr = _MR_DECODE.get(mr_match.group(1)) if up_mr is not None: - subtype = subtype or "mask_ablation" + subtype = subtype or "mask_ratio_sensitivity" return up_lr, up_mr, subtype def _infer_model_size(config: dict) -> str: - """Infer model-size label from the encoder config.""" + explicit = config.get("model_size") + if explicit: + return explicit d_model = _get_nested(config, "encoder.d_model") n_layers = _get_nested(config, "encoder.n_layers") - sprint = str(config.get("sprint", "")) - if d_model is None: return "default" - variant = MODEL_VARIANTS.get((d_model, n_layers)) if variant is not None: return variant - - if sprint == "7p": - if d_model == 128: - return "medium" - if d_model == 256: - return "large" - if d_model == 64: return "default" if n_layers is None: @@ -634,95 +379,55 @@ def _metric_value_or_nan(value) -> float: return float(value) -def extract_run( - run, - metric_keys: list[str], - requested_sprints: list[str] | None = None, -) -> dict: - """Extract config + metrics from a W&B run in a single retry-protected call.""" +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) ) - # Detect phase from tags - phase = None - for tag in tags: - if tag.startswith("phase:"): - phase = tag.split(":", 1)[1] - break + 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: - p = config.get("paradigm", "") - if p in ("supervised", "gru_d", "xgboost"): - phase = p + paradigm = config.get("paradigm", "") + if paradigm in {"supervised", "gru_d", "xgboost"}: + phase = "baseline" if paradigm in {"gru_d", "xgboost"} else "supervised" else: phase = "finetune" - # Derive protocol from freeze_encoder. XGBoost runs do not use this field, - # but they belong with the full-training / Protocol-B comparison family. freeze = _get_nested(config, "training.freeze_encoder") - if freeze is True: - protocol = "A" - elif freeze is False: - protocol = "B" - elif config.get("paradigm") == "xgboost" or phase == "baseline": - protocol = "B" - else: - protocol = None - - model_size = _infer_model_size(config) + protocol = config.get("protocol") or _tag_value(tags, "protocol") + if protocol is None: + if freeze is True: + protocol = "A" + elif freeze is False or config.get("paradigm") == "xgboost" or phase == "baseline": + protocol = "B" - # Detect source_dataset for transfer runs from config, output_dir, or run name. source_dataset = _recover_source_dataset(run_name, config) + up_lr, up_mr, experiment_subtype = _recover_pretrain_metadata(run_name, config) - # Extract metrics from summary 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)) - and not isinstance(value, bool) - ): + if key.startswith("fairness/") and isinstance(value, (int, float)): metrics[key] = _metric_value_or_nan(value) - config_sprint = str(config.get("sprint", "")) - sprint_str = _select_export_sprint(config_sprint, tags, requested_sprints) - experiment_type = EXPERIMENT_TYPES.get(sprint_str, "unknown") - - # Historical group 10 contains both core (label_fraction=1.0) and - # label-efficiency runs that merge with group 6 during aggregation. - label_frac = config.get("label_fraction", 1.0) - if experiment_type == "core" and label_frac is not None and float(label_frac) < 1.0: - experiment_type = "label_efficiency" - - # Recover upstream pretrain metadata (from config for new runs, run name for historical) - up_lr, up_mr, experiment_subtype = _recover_pretrain_metadata(run_name, config) - - if experiment_type in ("lr_ablation", "mask_ablation"): - experiment_subtype = experiment_subtype or experiment_type - experiment_type = "hp_ablation" - - # Reclassify HP ablation runs that would otherwise stay as "core" - if experiment_type == "core" and experiment_subtype in ("lr_ablation", "mask_ablation"): - experiment_type = "hp_ablation" - - # Historical group 10 also adds transfer seeds for group 7 scope. - if experiment_type == "core" and source_dataset is not None: - experiment_type = "transfer" - + paradigm = config.get("paradigm") or _get_nested(config, "ssl.name") row = { "wandb_run_id": run_id, "wandb_run_url": run_url, "wandb_run_name": run_name, "wandb_group": group, "created_at": created_at, - "experiment_type": experiment_type, - "sprint": sprint_str, - "config_sprint": config_sprint, - "sprint_tags": json.dumps(_sprint_tags(tags)), - "paradigm": config.get("paradigm", None), + "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), @@ -730,13 +435,12 @@ def extract_run( "label_fraction": config.get("label_fraction", 1.0), "lr": _get_nested(config, "optimizer.lr", default=None), "mask_ratio": _get_nested(config, "ssl.mask_ratio", default=None), - "model_size": model_size, + "model_size": _infer_model_size(config), "source_dataset": source_dataset, - "revision": config.get("revision", None), + "revision": config.get("revision", None) or _tag_value(tags, "revision"), "phase": phase, "upstream_pretrain_lr": up_lr, "upstream_pretrain_mask_ratio": up_mr, - "experiment_subtype": experiment_subtype, "_eval_checkpoint_source": summary.get("_eval_checkpoint_source", None), "_best_ckpt_path": summary.get("_best_ckpt_path", None), "_best_ckpt_load_ok": summary.get("_best_ckpt_load_ok", None), @@ -751,24 +455,7 @@ def extract_run( # --------------------------------------------------------------------------- -def _fingerprint_for_experiment_type(experiment_type: str) -> list[str]: - """Return the canonical fingerprint for a given experiment family.""" - if experiment_type in CROSS_SPRINT_TYPES or experiment_type in CONTEXTUAL_STAT_TYPES: - return CORE_FINGERPRINT - if experiment_type == "hp_ablation": - return HP_ABLATION_FINGERPRINT - return ABLATION_FINGERPRINT - - -def _allowed_varying_columns(experiment_type: str) -> set[str]: - """Columns allowed to vary within one canonical experiment family.""" - if experiment_type in CROSS_SPRINT_TYPES or experiment_type == "hp_ablation": - return {"sprint"} - return set() - - def _fillna_for_grouping(df: pd.DataFrame, columns: list[str]) -> pd.DataFrame: - """Fill NaNs so pandas groupby/drop_duplicates treat missing values deterministically.""" work = df.copy() for col in columns: if col in work.columns: @@ -777,42 +464,26 @@ def _fillna_for_grouping(df: pd.DataFrame, columns: list[str]) -> pd.DataFrame: def _assert_no_ambiguous_fingerprint_collisions(df: pd.DataFrame) -> None: - """Fail if the chosen fingerprint would collapse non-equivalent configurations.""" collisions = [] - - for experiment_type in sorted(df["experiment_type"].dropna().unique()): - subset = df[df["experiment_type"] == experiment_type] - if subset.empty: + 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 - - fingerprint = [*(_fingerprint_for_experiment_type(experiment_type)), "seed"] - fingerprint = [c for c in fingerprint if c in subset.columns] - if not fingerprint: - continue - - allowed_vary = _allowed_varying_columns(experiment_type) - identity_cols = [ - c for c in FINGERPRINT_AUDIT_COLUMNS if c in subset.columns and c not in allowed_vary - ] - - work = _fillna_for_grouping(subset, list(dict.fromkeys(fingerprint + identity_cols))) - grouped = work.groupby(fingerprint, dropna=False) - - for key, group in grouped: - if len(group) <= 1: - continue - - unique_identities = group[identity_cols].drop_duplicates() - if len(unique_identities) > 1: - if not isinstance(key, tuple): - key = (key,) - collisions.append((experiment_type, dict(zip(fingerprint, key)), len(group))) + 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 experiment_type, fingerprint, count in collisions[:5]: - desc = ", ".join(f"{k}={v}" for k, v in fingerprint.items()) - preview.append(f"{experiment_type}: {desc} ({count} rows)") + 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) ) @@ -820,43 +491,22 @@ def _assert_no_ambiguous_fingerprint_collisions(df: pd.DataFrame) -> None: def build_per_seed_df( runs: list, - requested_sprints: list[str] | None = None, allow_extraction_failures: bool = False, ) -> pd.DataFrame: - """Build the per-seed DataFrame from raw W&B runs. - - One row per run/export membership with config columns, metrics, and W&B IDs. - """ rows = [] failed = [] - for i, run in enumerate(runs): - if (i + 1) % 100 == 0: - print(f" Processing run {i + 1}/{len(runs)}...", file=sys.stderr) + for index, run in enumerate(runs): + if (index + 1) % 100 == 0: + print(f" Processing run {index + 1}/{len(runs)}...", file=sys.stderr) try: - row = extract_run(run, ALL_METRICS, requested_sprints=requested_sprints) - rows.append(row) - tags = [f"sprint:{s}" for s in json.loads(row.get("sprint_tags", "[]"))] - for inherited_sprint in _additional_inherited_export_sprints( - row, - tags, - requested_sprints=requested_sprints, - ): - rows.append( - extract_run( - run, - ALL_METRICS, - requested_sprints=[inherited_sprint], - ) - ) - except Exception as e: + 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}: {e!r}", file=sys.stderr) + 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, - ) + 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( @@ -868,32 +518,20 @@ def build_per_seed_df( if df.empty: return df - # Ensure correct dtypes 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") - df["label_fraction"] = df["label_fraction"].fillna(1.0) + 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") - # Deduplicate: keep only the most recent run per (fingerprint + seed). - # This handles reruns/revisions — the latest finished run is canonical. - # Different experiment families use different fingerprints. df = df.sort_values("created_at", ascending=False) before = len(df) _assert_no_ambiguous_fingerprint_collisions(df) - - parts = [] - for experiment_type in sorted(df["experiment_type"].dropna().unique()): - subset = df[df["experiment_type"] == experiment_type] - dedup_cols = [*(_fingerprint_for_experiment_type(experiment_type)), "seed"] - dedup_cols = [c for c in dedup_cols if c in subset.columns] - parts.append(subset.drop_duplicates(subset=dedup_cols, keep="first")) - df = pd.concat(parts, ignore_index=True) if parts else df.iloc[0:0] - + dedup_cols = [col for col in [*FINGERPRINT, "seed"] if col in df.columns] + df = df.drop_duplicates(subset=dedup_cols, keep="first") after = len(df) if before != after: print( @@ -902,19 +540,19 @@ def build_per_seed_df( file=sys.stderr, ) - # Sort for reproducibility - sort_cols = [c for c in ["sprint", "paradigm", "dataset", "task", "seed"] if c in df.columns] + 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) - # Row count validation by experiment_type for auditability - if "experiment_type" in df.columns: - counts = df["experiment_type"].value_counts().sort_index() - print("\n Row counts by experiment_type after dedup:", file=sys.stderr) - for etype, count in counts.items(): - print(f" {etype}: {count}", file=sys.stderr) - print(f" TOTAL: {len(df)}", file=sys.stderr) - + 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 @@ -922,69 +560,47 @@ def filter_thesis_tasks( per_seed_df: pd.DataFrame, include_extension_tasks: bool = False, ) -> pd.DataFrame: - """Apply the primary thesis task whitelist unless extension export is requested.""" + """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 - mask = per_seed_df["task"].isin(PRIMARY_THESIS_TASKS) - excluded = per_seed_df.loc[~mask, "task"].dropna() - if not excluded.empty: - task_counts = excluded.value_counts().to_dict() + 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( - " Excluding extension task rows from primary thesis export: " - + ", ".join(f"{task}={count}" for task, count in sorted(task_counts.items())), + " Dropped extension-task rows outside the thesis task set: " + f"{len(dropped)} ({counts})", file=sys.stderr, ) - return per_seed_df.loc[mask].copy().reset_index(drop=True) - - -def _metric_columns(df: pd.DataFrame) -> list[str]: - """Return static and dynamically discovered numeric metric columns.""" - static_metrics = [c for c in ALL_METRICS if c in df.columns] - dynamic_fairness = [ - c - for c in sorted(df.columns) - if c.startswith("fairness/") - and c not in static_metrics - and pd.api.types.is_numeric_dtype(df[c]) - ] - return static_metrics + dynamic_fairness + return per_seed_df.loc[task_mask].reset_index(drop=True) def _aggregate_group(df: pd.DataFrame, fingerprint_cols: list[str]) -> pd.DataFrame: - """Aggregate a subset of runs by the given fingerprint columns. - - Returns a DataFrame with one row per unique fingerprint, containing - mean/std/min/max of metrics and lists of run IDs/seeds/groups. - """ - metric_cols = _metric_columns(df) - - # Sentinel for NaN in groupby - work = df.copy() - for col in fingerprint_cols: - if col in work.columns: - work[col] = work[col].fillna("__none__") - + 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) - agg_rows = [] + rows = [] for fingerprint, group in grouped: if isinstance(fingerprint, str): fingerprint = (fingerprint,) row = dict(zip(fingerprint_cols, fingerprint)) - - # Restore None from sentinel for col in fingerprint_cols: if row.get(col) == "__none__": row[col] = None - # Run metadata 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["sprint_list"] = json.dumps(sorted(group["sprint"].dropna().unique().tolist())) + row["experiment_class_list"] = json.dumps( + sorted(group["experiment_class"].dropna().unique().tolist()) + ) - # Metric aggregation for metric in metric_cols: values = group[metric].dropna() row[f"{metric}/n"] = int(len(values)) @@ -998,65 +614,119 @@ def _aggregate_group(df: pd.DataFrame, fingerprint_cols: list[str]) -> pd.DataFr row[f"{metric}/std"] = float("nan") row[f"{metric}/min"] = float("nan") row[f"{metric}/max"] = float("nan") + rows.append(row) - agg_rows.append(row) - - return pd.DataFrame(agg_rows) + return pd.DataFrame(rows) def build_aggregated_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: - """Group by fingerprint, compute mean/std/min/max across seeds. + 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 + - Uses conditional grouping: - - Core / label-efficiency / transfer runs: group without the sprint tag - - HP ablations: group without the sprint tag, but include subtype + upstream pretrain config - - Other experiment families: group with the sprint tag - """ +# --------------------------------------------------------------------------- +# 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 = [] - cross_sprint = per_seed_df[per_seed_df["experiment_type"].isin(CROSS_SPRINT_TYPES)] - if len(cross_sprint) > 0: - print( - f" Aggregating {len(cross_sprint)} cross-sprint runs (merged experiment families)...", - file=sys.stderr, + 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") ) - parts.append(_aggregate_group(cross_sprint, CORE_FINGERPRINT)) + return pd.concat(parts, ignore_index=True) if parts else pd.DataFrame() - hp_ablation = per_seed_df[per_seed_df["experiment_type"] == "hp_ablation"] - if len(hp_ablation) > 0: - print( - " Aggregating " - f"{len(hp_ablation)} HP-ablation runs " - "(merged by subtype + upstream config)...", - file=sys.stderr, - ) - parts.append(_aggregate_group(hp_ablation, HP_ABLATION_FINGERPRINT)) - per_sprint = per_seed_df[ - ~per_seed_df["experiment_type"].isin(CROSS_SPRINT_TYPES | {"hp_ablation"}) - ] - if len(per_sprint) > 0: - print( - f" Aggregating {len(per_sprint)} per-sprint runs (ablations, etc.)...", file=sys.stderr +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") ) - parts.append(_aggregate_group(per_sprint, ABLATION_FINGERPRINT)) + return pd.concat(parts, ignore_index=True) if parts else pd.DataFrame() + - if not parts: +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) + protocol_b = per_seed_df["protocol"] == "B" + full = per_seed_df[ + protocol_b + & (label_fraction == 1.0) + & (per_seed_df["experiment_class"].isin(["core_ssl_benchmark", "classical_baselines"])) + ].copy() + low_label = per_seed_df[ + protocol_b + & (label_fraction < 1.0) + & (per_seed_df["experiment_class"].isin(["label_efficiency", "classical_baselines"])) + ].copy() - agg_df = pd.concat(parts, ignore_index=True) + 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() - # Sort - sort_cols = [ - c for c in ["experiment_type", "paradigm", "dataset", "task"] if c 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 +# --------------------------------------------------------------------------- +# Statistical Tables +# --------------------------------------------------------------------------- def _primary_metric_for_task(task_name: str | None) -> str | None: - """Return the benchmark-primary test metric for a downstream task.""" if task_name is None: return None if task_name in PRIMARY_TEST_METRIC_BY_TASK: @@ -1067,9 +737,7 @@ def _primary_metric_for_task(task_name: str | None) -> str | None: def _format_task_seed_pairs(keys: set[tuple[str, int]]) -> str: - """Serialize task/seed coverage gaps for audit columns.""" - formatted = [{"task": task, "seed": int(seed)} for task, seed in sorted(keys)] - return json.dumps(formatted) + return json.dumps([{"task": task, "seed": int(seed)} for task, seed in sorted(keys)]) def _metric_pairs_for_paradigm( @@ -1077,7 +745,6 @@ def _metric_pairs_for_paradigm( paradigm: str, value_col: str = "primary_metric_value", ) -> pd.DataFrame: - """Return finite task/seed metric rows for one paradigm.""" pairs = scope_group[scope_group["paradigm"] == paradigm][["task", "seed", value_col]].rename( columns={value_col: "value"} ) @@ -1094,21 +761,17 @@ def _paired_metric_frame( paradigm_b: str, value_col: str = "primary_metric_value", ) -> tuple[pd.DataFrame, dict[str, int | str]]: - """Build paired metric frame and task/seed coverage summary.""" 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), @@ -1128,22 +791,15 @@ def _paired_stat_row( paradigm_b: str, higher_is_better: bool, ) -> dict: - """Compute paired Wilcoxon/effect-size row for a comparison.""" 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) - median_improvement = float(pd.Series(improvement).median()) - better_paradigm = ( - paradigm_a if mean_improvement > 0 else paradigm_b if mean_improvement < 0 else None - ) - return { **scope, "paradigm_a": paradigm_a, @@ -1156,8 +812,10 @@ def _paired_stat_row( "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": median_improvement, - "better_paradigm": better_paradigm, + "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"], @@ -1170,89 +828,50 @@ def _paired_stat_row( def _add_contextual_classical_stat_rows(per_seed_df: pd.DataFrame) -> pd.DataFrame: - """Add synthetic statistic scopes for classical-vs-neural comparisons.""" - required = {"experiment_type", "protocol", "label_fraction", "paradigm"} + required = {"experiment_class", "protocol", "label_fraction", "paradigm"} if per_seed_df.empty or not required.issubset(per_seed_df.columns): return per_seed_df - work = per_seed_df.copy() - label_fraction = pd.to_numeric(work["label_fraction"], errors="coerce").fillna(1.0) - protocol_b = work["protocol"] == "B" - classical = work["experiment_type"] == "classical_baselines" - - parts = [work] - - full_context_mask = ( - protocol_b & (label_fraction == 1.0) & ((work["experiment_type"] == "core") | classical) - ) - label_efficiency_context_mask = ( - protocol_b - & (label_fraction < 1.0) - & ((work["experiment_type"] == "label_efficiency") | classical) - ) - - for mask, experiment_type in [ - (full_context_mask, "classical_context_full"), - (label_efficiency_context_mask, "classical_context_label_efficiency"), - ]: - context = work.loc[mask].copy() - if context.empty: - continue - paradigms = set(context["paradigm"].dropna()) - if not (paradigms & {"xgboost", "gru_d"}) or paradigms <= {"xgboost", "gru_d"}: - continue - context["experiment_type"] = experiment_type + 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: - """Build pairwise paradigm significance tables from per-seed results. - - Comparisons are scoped by the canonical experiment fingerprint with - `paradigm` and `task` removed, then pooled across matched `(task, seed)` - pairs for tasks that share the same primary metric. This matches the benchmark - plan's paired-across-seeds-and-tasks intent without mixing incompatible - scales such as AUPRC and MAE in the same test. - """ - if per_seed_df.empty or "experiment_type" not in per_seed_df.columns: + 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_type in sorted(per_seed_df["experiment_type"].dropna().unique()): - subset = per_seed_df[per_seed_df["experiment_type"] == experiment_type].copy() + 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.index, + for row, metric_name in zip( subset.to_dict("records"), subset["primary_metric_name"], strict=True, ) ] - fingerprint = _fingerprint_for_experiment_type(experiment_type) scope_cols = [ - c for c in fingerprint if c in subset.columns and c not in {"paradigm", "task", "phase"} + col + for col in FINGERPRINT + if col in subset.columns and col not in {"paradigm", "task", "phase"} ] if "primary_metric_name" not in scope_cols: scope_cols.append("primary_metric_name") - if not scope_cols: - continue - 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__": @@ -1268,30 +887,25 @@ def build_statistical_tests_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: 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_type in CONTEXTUAL_STAT_TYPES: + 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 - - row = _paired_stat_row( - scope=scope, - paired=paired, - coverage=coverage, - paradigm_a=paradigm_a, - paradigm_b=paradigm_b, - higher_is_better=higher_is_better, + 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, + ) ) - family_rows.append(row) - - if not family_rows: - continue corrected = bonferroni_correction([row["p_value"] for row in family_rows]) for row, corrected_p in zip(family_rows, corrected, strict=True): @@ -1303,15 +917,13 @@ def build_statistical_tests_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: if not rows: return pd.DataFrame() - stats_df = pd.DataFrame(rows) sort_cols = [ col for col in [ - "experiment_type", + "experiment_class", "dataset", "protocol", - "phase", "label_fraction", "primary_metric_name", "paradigm_a", @@ -1325,17 +937,19 @@ def build_statistical_tests_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: def build_ts2vec_vs_core_contrastive_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: - """Compare Sprint-13 TS2Vec directly against core contrastive runs.""" - required = {"experiment_type", "paradigm", "task", "seed"} + 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_type"] == "temporal_contrastive") + (per_seed_df["experiment_class"] == "ts2vec_extension") & (per_seed_df["paradigm"] == "ts2vec") ) - | ((per_seed_df["experiment_type"] == "core") & (per_seed_df["paradigm"] == "contrastive")) + | ( + (per_seed_df["experiment_class"] == "core_ssl_benchmark") + & (per_seed_df["paradigm"] == "contrastive") + ) ].copy() if subset.empty: return pd.DataFrame() @@ -1362,65 +976,44 @@ def build_ts2vec_vs_core_contrastive_df(per_seed_df: pd.DataFrame) -> pd.DataFra ] if col in subset.columns ] - work = _fillna_for_grouping(subset, scope_cols) 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 - - paradigms = set(scope_group["paradigm"].dropna()) - if not {"contrastive", "ts2vec"}.issubset(paradigms): + 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 - row = _paired_stat_row( - scope={ - **scope, - "comparison_type": "ts2vec_vs_core_contrastive", - "experiment_type_a": "temporal_contrastive", - "experiment_type_b": "core", - }, - paired=paired, - coverage=coverage, - paradigm_a="ts2vec", - paradigm_b="contrastive", - higher_is_better=higher_is_better, + 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, + ) ) - rows.append(row) 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) - - result = pd.DataFrame(rows) - sort_cols = [ - col - for col in [ - "dataset", - "protocol", - "label_fraction", - "model_size", - "primary_metric_name", - ] - if col in result.columns - ] - if sort_cols: - result = result.sort_values(sort_cols, na_position="last").reset_index(drop=True) - return result + return pd.DataFrame(rows) # --------------------------------------------------------------------------- @@ -1428,31 +1021,94 @@ def build_ts2vec_vs_core_contrastive_df(per_seed_df: pd.DataFrame) -> pd.DataFra # --------------------------------------------------------------------------- +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 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 validate( per_seed_df: pd.DataFrame, aggregated_df: pd.DataFrame, statistical_df: pd.DataFrame | None = None, expected_seeds: set[int] | None = None, ) -> list[str]: - """Validate results and return warning strings. - - Warns when fixed-seed experiment families do not have the exact expected - seed set, and when statistical comparisons have incomplete task/seed pairing. - """ warnings = [] expected_seeds = set(EXPECTED_FIXED_SEEDS if expected_seeds is None else expected_seeds) - # Check fixed-seed experiment families for exact expected seed set. - if "experiment_type" in aggregated_df.columns: + if "experiment_class" in aggregated_df.columns: fixed_seed = aggregated_df[ - aggregated_df["experiment_type"].isin(FIXED_SEED_EXPERIMENT_TYPES) + 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 " @@ -1460,12 +1116,12 @@ def validate( ) for row, seeds in bad_seed_rows: desc = ", ".join( - f"{c}={row[c]}" - for c in ["paradigm", "dataset", "task", "protocol"] - if c in row and row[c] is not None + 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_type={row['experiment_type']}, {desc} — " + f" experiment_class={row['experiment_class']}, {desc} - " f"seeds={sorted(seeds)}, missing={sorted(expected_seeds - seeds)}, " f"unexpected={sorted(seeds - expected_seeds)}" ) @@ -1482,64 +1138,31 @@ def validate( f"WARNING: {len(incomplete)}/{len(statistical_df)} statistical comparisons " "have incomplete paired task/seed coverage:" ) - for _, row in incomplete.head(10).iterrows(): - scope = ", ".join( - f"{c}={row[c]}" - for c in [ - "experiment_type", - "dataset", - "protocol", - "label_fraction", - "primary_metric_name", - "paradigm_a", - "paradigm_b", - ] - if c in row and row[c] is not None - ) - warnings.append( - f" {scope} — shared={row['n_shared_task_seed_pairs']}, " - f"union={row['n_union_task_seed_pairs']}" - ) - # Check for runs with no test metrics at all - test_cols = [c for c in PERFORMANCE_TEST_METRICS if c in per_seed_df.columns] + 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(): - n_empty = all_nan.sum() warnings.append( - f"WARNING: {n_empty} runs have no test metrics at all. " + f"WARNING: {all_nan.sum()} runs have no test metrics at all. " "These may be pretraining runs or crashed evaluations." ) - # Check primary metric presence for every evaluation row. if not per_seed_df.empty and "task" in per_seed_df.columns: missing_primary = [] for _, row in per_seed_df.iterrows(): - phase = row.get("phase") - if phase not in EVAL_PHASES: + 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." ) - for row, primary_metric in missing_primary[:10]: - warnings.append( - " " - + ", ".join( - f"{c}={row.get(c)}" - for c in ["wandb_run_id", "paradigm", "dataset", "task", "seed", "phase"] - if c in row - ) - + f" — missing {primary_metric}" - ) checkpoint_issues = _checkpoint_provenance_issues(per_seed_df) if checkpoint_issues: @@ -1552,8 +1175,8 @@ def validate( warnings.append( " " + ", ".join( - f"{c}={row.get(c)}" - for c in [ + f"{col}={row.get(col)}" + for col in [ "wandb_run_id", "paradigm", "dataset", @@ -1561,98 +1184,24 @@ def validate( "seed", "phase", ] - if c in row + if col in row ) - + f" — {reason}" + + f" - {reason}" ) - # Summary by experiment type - n_runs = len(per_seed_df) - n_configs = len(aggregated_df) print("\nValidation summary:", file=sys.stderr) - print(f" Total runs: {n_runs}", file=sys.stderr) - print(f" Unique configs: {n_configs}", file=sys.stderr) - if "experiment_type" in aggregated_df.columns: - for etype, group in aggregated_df.groupby("experiment_type"): + 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" {etype}: {len(group)} configs, seeds: {seed_dist}", file=sys.stderr) + print( + f" {experiment_class}: {len(group)} configs, seeds: {seed_dist}", file=sys.stderr + ) print(f" Warnings: {len(warnings)}", file=sys.stderr) - return warnings -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.""" - phase = row.get("phase") - if phase not in EVAL_PHASES: - return False - paradigm = str(row.get("paradigm", "")).lower() - if paradigm == "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 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 - - # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- @@ -1675,41 +1224,22 @@ def parse_args() -> argparse.Namespace: help="W&B entity name (default: WANDB_ENTITY env var)", ) parser.add_argument( - "--sprint", + "--experiment-class", nargs="+", - help="Filter to specific experiment group(s), e.g. --sprint 1 1b 2", - ) - parser.add_argument( - "--paradigm", - nargs="+", - help="Filter to specific paradigm(s), e.g. --paradigm mae jepa", - ) - parser.add_argument( - "--dataset", - nargs="+", - help="Filter to specific dataset(s), e.g. --dataset miiv eicu", + 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), e.g. --revision benchmark-v1", - ) - parser.add_argument( - "--state", - default="finished", - help="Run state filter (default: finished)", - ) - parser.add_argument( - "--output-dir", - default="results", - help="Output directory (default: results/)", - ) + parser.add_argument("--revision", nargs="+", help="Filter to specific revision tag(s)") + 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="+", @@ -1738,7 +1268,7 @@ def parse_args() -> argparse.Namespace: action="store_true", help=( "Include optional extension tasks such as ICU mortality. By default, " - f"primary thesis exports are restricted to {sorted(PRIMARY_THESIS_TASKS)}." + f"primary thesis exports are restricted to {sorted(THESIS_TASKS)}." ), ) args = parser.parse_args() @@ -1756,32 +1286,28 @@ def parse_args() -> argparse.Namespace: return args -def main(): +def main() -> None: args = parse_args() output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) - # Fetch runs runs = fetch_all_runs( project=args.project, entity=args.entity, state=args.state, - sprint=args.sprint, + 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) - # Build DataFrames print(f"\nBuilding per-seed DataFrame from {len(runs)} runs...", file=sys.stderr) per_seed_df = build_per_seed_df( runs, - requested_sprints=args.sprint, allow_extraction_failures=args.allow_extraction_failures, ) per_seed_df = filter_thesis_tasks( @@ -1798,6 +1324,11 @@ def main(): 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) @@ -1806,7 +1337,6 @@ def main(): ts2vec_contrastive_df = build_ts2vec_vs_core_contrastive_df(per_seed_df) print(f" Shape: {ts2vec_contrastive_df.shape}", file=sys.stderr) - # Validate 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() @@ -1817,46 +1347,35 @@ def main(): statistical_df=statistical_coverage_df, expected_seeds=set(args.expected_seeds), ) - for w in warnings: - print(w, file=sys.stderr) - - # Save - per_seed_path = output_dir / "per_seed_results.parquet" - aggregated_path = output_dir / "aggregated_results.parquet" - statistical_path = output_dir / "statistical_tests.parquet" - ts2vec_contrastive_path = output_dir / "ts2vec_vs_core_contrastive.parquet" - - per_seed_df.to_parquet(per_seed_path, index=False) - aggregated_df.to_parquet(aggregated_path, index=False) - statistical_df.to_parquet(statistical_path, index=False) - ts2vec_contrastive_df.to_parquet(ts2vec_contrastive_path, index=False) - - print("\nSaved:", file=sys.stderr) - print(f" {per_seed_path} ({len(per_seed_df)} rows)", file=sys.stderr) - print(f" {aggregated_path} ({len(aggregated_df)} rows)", file=sys.stderr) - print(f" {statistical_path} ({len(statistical_df)} rows)", file=sys.stderr) - print(f" {ts2vec_contrastive_path} ({len(ts2vec_contrastive_df)} rows)", file=sys.stderr) - - # Print quick summary to stdout for piping + for warning in warnings: + print(warning, file=sys.stderr) + + 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_type" in aggregated_df.columns: - for etype, group in aggregated_df.groupby("experiment_type"): + 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" {etype}: {len(group)} configs, seeds: {seed_dist}") - + print(f" {experiment_class}: {len(group)} configs, seeds: {seed_dist}") if warnings and not args.allow_incomplete: sys.exit(1) - if "paradigm" in aggregated_df.columns: - print(f"Paradigms: {sorted(aggregated_df['paradigm'].dropna().unique())}") - if "dataset" in aggregated_df.columns: - print(f"Datasets: {sorted(aggregated_df['dataset'].dropna().unique())}") - if "task" in aggregated_df.columns: - print(f"Tasks: {sorted(aggregated_df['task'].dropna().unique())}") if __name__ == "__main__": diff --git a/scripts/internal/launch_thesis_tmux.sh b/scripts/internal/launch_thesis_tmux.sh index f3225d9..e241b94 100755 --- a/scripts/internal/launch_thesis_tmux.sh +++ b/scripts/internal/launch_thesis_tmux.sh @@ -6,18 +6,20 @@ SESSION_NAME="${SESSION_NAME:-slices-thesis}" WANDB_PROJECT="${WANDB_PROJECT:-slices-thesis}" WANDB_ENTITY="${WANDB_ENTITY:-}" REVISION="${REVISION:-thesis-v1}" -REASON="${REASON:-thesis-final canonical-sprint11-baselines}" +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_SPRINT12="${INCLUDE_SPRINT12:-1}" -INCLUDE_SPRINT13="${INCLUDE_SPRINT13:-1}" -INCLUDE_SPRINT7P="${INCLUDE_SPRINT7P:-1}" +INCLUDE_SMART_REFERENCE="${INCLUDE_SMART_REFERENCE:-1}" +INCLUDE_TS2VEC_EXTENSION="${INCLUDE_TS2VEC_EXTENSION:-1}" +INCLUDE_CAPACITY_STUDY="${INCLUDE_CAPACITY_STUDY:-1}" RUN_EXPORT="${RUN_EXPORT:-1}" STATUS_INTERVAL="${STATUS_INTERVAL:-60}" RESULTS_DIR="${RESULTS_DIR:-results/${WANDB_PROJECT}_${REVISION}}" @@ -38,29 +40,33 @@ if [[ -z "$WANDB_ENTITY" ]]; then exit 1 fi -if [[ -z "$LAUNCH_COMMIT" ]]; then - LAUNCH_COMMIT="$(git -C "$REPO_ROOT" rev-parse --verify HEAD)" -fi - -if ! git -C "$REPO_ROOT" cat-file -e "$LAUNCH_COMMIT^{commit}" 2>/dev/null; then - echo "LAUNCH_COMMIT is not available in this checkout: $LAUNCH_COMMIT" >&2 - exit 1 -fi +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 -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 ! git -C "$REPO_ROOT" cat-file -e "$LAUNCH_COMMIT^{commit}" 2>/dev/null; then + echo "LAUNCH_COMMIT is not available in this checkout: $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 + 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() { @@ -97,9 +103,7 @@ for dataset in datasets: 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}" - ) + 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"): @@ -124,46 +128,55 @@ 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" -echo "Validating processed artifacts..." -( - cd "$REPO_ROOT" - validate_processed_artifacts -) -echo "Processed artifacts are ready." +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_sprints=(1 1b 1c 2 3 4 5 6 7 8 10 11) -fairness_sprints=(1 2 3 4 5 6 7 8 10 11) -tag_sprints=(1b 1c 2 5 6 7 8) -appendix_sprints=() +main_classes=( + core_ssl_benchmark + label_efficiency + cross_dataset_transfer + hp_robustness + classical_baselines +) +fairness_classes=("${main_classes[@]}") +appendix_classes=() -if [[ "$INCLUDE_SPRINT13" == "1" ]]; then - main_sprints+=(13) - fairness_sprints+=(13) +if [[ "$INCLUDE_TS2VEC_EXTENSION" == "1" ]]; then + main_classes+=(ts2vec_extension) + fairness_classes+=(ts2vec_extension) fi -if [[ "$INCLUDE_SPRINT7P" == "1" ]]; then - main_sprints+=(7p) - fairness_sprints+=(7p) - tag_sprints+=(7p) +if [[ "$INCLUDE_CAPACITY_STUDY" == "1" ]]; then + main_classes+=(capacity_study) + fairness_classes+=(capacity_study) fi -if [[ "$INCLUDE_SPRINT12" == "1" ]]; then - appendix_sprints+=(12) - fairness_sprints+=(12) +if [[ "$INCLUDE_SMART_REFERENCE" == "1" ]]; then + appendix_classes+=(smart_external_reference) + fairness_classes+=(smart_external_reference) fi -warmup_sprints=("${main_sprints[@]}") -if [[ "$INCLUDE_SPRINT12" == "1" ]]; then - warmup_sprints+=(12) +warmup_classes=("${main_classes[@]}") +if [[ "$INCLUDE_SMART_REFERENCE" == "1" ]]; then + warmup_classes+=(smart_external_reference) fi quote_cmd() { @@ -173,26 +186,22 @@ quote_cmd() { run_args=(--project "$WANDB_PROJECT" --entity "$WANDB_ENTITY") export_args=(--project "$WANDB_PROJECT" --entity "$WANDB_ENTITY") fairness_args=(--project "$WANDB_PROJECT" --entity "$WANDB_ENTITY") -tag_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 - tag_args+=(--revision "$REVISION") -fi if [[ -n "$REVISION" ]]; then fairness_args+=(--revision "$REVISION") fi -fairness_args+=(--sprint "${fairness_sprints[@]}" --batch-size "$BATCH_SIZE_FAIRNESS" --device "$DEVICE_FAIRNESS") +export_args+=(--experiment-class "${fairness_classes[@]}") +fairness_args+=(--experiment-class "${fairness_classes[@]}" --batch-size "$BATCH_SIZE_FAIRNESS" --device "$DEVICE_FAIRNESS") -warmup_cmd=(uv run python scripts/internal/run_experiments.py warmup --sprint "${warmup_sprints[@]}") -main_cmd=(uv run python scripts/internal/run_experiments.py run --sprint "${main_sprints[@]}" --parallel "$PARALLEL_MAIN" "${run_args[@]}") -tag_cmd=(uv run python scripts/internal/run_experiments.py tag --sprint "${tag_sprints[@]}" "${tag_args[@]}") +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_sprints[@]} > 0)); then - appendix_cmd=(uv run python scripts/internal/run_experiments.py run --sprint "${appendix_sprints[@]}" --parallel "$PARALLEL_APPENDIX" "${run_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 @@ -204,7 +213,6 @@ status_log="$LOG_DIR/thesis-status-${timestamp}.log" printf -v warmup_line "%q " "${warmup_cmd[@]}" printf -v main_line "%q " "${main_cmd[@]}" -printf -v tag_line "%q " "${tag_cmd[@]}" printf -v fairness_line "%q " "${fairness_cmd[@]}" printf -v export_line "%q " "${export_cmd[@]}" @@ -229,20 +237,21 @@ echo "W&B project: $(printf "%q" "$WANDB_PROJECT")" echo "W&B entity: $(printf "%q" "$WANDB_ENTITY")" echo "Revision: $(printf "%q" "$REVISION")" echo "Launch commit: $(printf "%q" "$LAUNCH_COMMIT")" -echo "Main sprints: ${main_sprints[*]}" -echo "Fairness sprints: ${fairness_sprints[*]}" -echo "Appendix sprints: ${appendix_sprints[*]:-none}" - -current_commit="\$(git rev-parse --verify HEAD)" -if [[ "\$current_commit" != "$(printf "%q" "$LAUNCH_COMMIT")" ]]; then - echo "Launch commit mismatch: current=\$current_commit expected=$(printf "%q" "$LAUNCH_COMMIT")" >&2 - exit 1 +echo "Main classes: ${main_classes[*]}" +echo "Fairness classes: ${fairness_classes[*]}" +echo "Appendix classes: ${appendix_classes[*]:-none}" + +if [[ "$(printf "%q" "$SKIP_LAUNCH_GIT_CHECK")" != "1" ]]; then + current_commit="\$(git rev-parse --verify HEAD)" + if [[ "\$current_commit" != "$(printf "%q" "$LAUNCH_COMMIT")" ]]; then + echo "Launch commit mismatch: current=\$current_commit expected=$(printf "%q" "$LAUNCH_COMMIT")" >&2 + exit 1 + fi fi uv sync --dev ${warmup_line} ${main_line} -${tag_line} ${appendix_block}${fairness_line} ${export_block}echo "Finished: \$(date -Iseconds)" EOF diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index 29aceda..917fc9d 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -1,35 +1,23 @@ #!/usr/bin/env python3 -""" -Parallel experiment runner for SLICES. +"""Class-based experiment runner for the final SLICES thesis rerun corpus. -Generates all experiment configurations across sprints, resolves dependencies, -and executes them in parallel with crash recovery and state persistence. +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 --sprint 1 - uv run python scripts/internal/run_experiments.py run --sprint 1 --parallel 4 \\ - --project slices-thesis --revision thesis-v1 --entity - uv run python scripts/internal/run_experiments.py run --sprint 1 2 3 --parallel 6 \\ - --dry-run --project slices-thesis --revision thesis-v1 --entity - uv run python scripts/internal/run_experiments.py run --sprint 1 \\ - --project slices-thesis --revision thesis-v1 --entity --reason "fix LR" - uv run python scripts/internal/run_experiments.py run --sprint 1 2 \\ - --project slices-thesis --revision thesis-v1 --entity - uv run python scripts/internal/run_experiments.py status - uv run python scripts/internal/run_experiments.py status --sprint 1 - uv run python scripts/internal/run_experiments.py retry --failed --parallel 4 \\ + 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 - uv run python scripts/internal/run_experiments.py retry --skipped --failed --parallel 4 \\ + uv run python scripts/internal/run_experiments.py run \ + --experiment-class core_ssl_benchmark --dry-run \ --project slices-thesis --revision thesis-v1 --entity - uv run python scripts/internal/run_experiments.py retry --failed \\ - --sprint 1 --revision thesis-v1 --parallel 4 \\ - --project slices-thesis --entity - uv run python scripts/internal/run_experiments.py tag --sprint 2 --dry-run \\ - --project slices-thesis --entity - uv run python scripts/internal/run_experiments.py tag --sprint 2 3 5 \\ + uv run python scripts/internal/run_experiments.py status --experiment-class core_ssl_benchmark + uv run python scripts/internal/run_experiments.py retry --failed \ + --experiment-class core_ssl_benchmark --revision thesis-v1 \ --project slices-thesis --entity """ - from __future__ import annotations import argparse @@ -49,20 +37,53 @@ # --------------------------------------------------------------------------- # 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": 840, + "cross_dataset_transfer": 120, + "hp_robustness": 150, + "capacity_study": 100, + "classical_baselines": 360, + "ts2vec_extension": 135, + "smart_external_reference": 135, +} + SSL_PARADIGMS = ["mae", "jepa", "contrastive"] DATASETS = ["miiv", "eicu", "combined"] TASKS = ["mortality_24h", "mortality_hospital", "aki_kdigo", "los_remaining"] -SEEDS = [42, 123, 456] -SEEDS_EXTENDED = [42, 123, 456, 789, 1011] # 5 seeds for cheap baselines +SEEDS_EXTENDED = [42, 123, 456, 789, 1011] LABEL_FRACTIONS_FULL = [0.01, 0.05, 0.1, 0.25, 0.5] LABEL_FRACTIONS_TREND = [0.1] -LABEL_FRACTIONS_PILOT = [0.01, 0.1, 0.5] +LABEL_FRACTIONS_CAPACITY = [0.01, 0.1, 0.5] -# Model capacity variants for Sprint 7 pilot (capacity ablation) MODEL_SIZES = { "medium": { "model": "transformer_medium", - "encoder": {"d_model": 128, "n_layers": 4, "n_heads": 8, "d_ff": 512}, "ssl_scale": { "mae": { "ssl.decoder_d_model": 128, @@ -70,21 +91,10 @@ "ssl.decoder_n_heads": 8, "ssl.decoder_d_ff": 512, }, - "jepa": { - "ssl.predictor_d_model": 64, - "ssl.predictor_n_layers": 2, - "ssl.predictor_n_heads": 4, - "ssl.predictor_d_ff": 256, - }, - "contrastive": { - "ssl.proj_hidden_dim": 512, - "ssl.proj_output_dim": 128, - }, }, }, "large": { "model": "transformer_large", - "encoder": {"d_model": 256, "n_layers": 4, "n_heads": 8, "d_ff": 1024}, "ssl_scale": { "mae": { "ssl.decoder_d_model": 256, @@ -92,46 +102,17 @@ "ssl.decoder_n_heads": 8, "ssl.decoder_d_ff": 1024, }, - "jepa": { - "ssl.predictor_d_model": 128, - "ssl.predictor_n_layers": 2, - "ssl.predictor_n_heads": 8, - "ssl.predictor_d_ff": 512, - }, - "contrastive": { - "ssl.proj_hidden_dim": 1024, - "ssl.proj_output_dim": 256, - }, }, }, } -LR_ABLATION = [2e-4, 5e-4, 2e-3] # 1e-3 reused from Phase 1 -MASK_RATIO_ABLATION = [0.3, 0.75] # 0.5 reused from Phase 1 -TRANSFER_PAIRS = [("miiv", "eicu"), ("eicu", "miiv")] -# Baseline inheritance: which earlier sprints' runs should be tagged as baselines -# for later sprints. See docs/internal/EXPERIMENT_PLAN.md § "Baseline Inheritance Across Sprints". -# Intentionally coarse-grained: ALL runs from source sprints are tagged, even when -# only a subset is relevant (e.g. 1b only needs mortality_24h from Sprint 1). Use -# secondary W&B tags (task:, protocol:, phase:, etc.) to filter down when querying. -BASELINE_SPRINTS: dict[str, list[str]] = { - "1": [], - "1b": ["1"], - "1c": ["1"], - "2": ["1"], - "3": [], - "4": [], - "5": ["1", "2", "3", "4"], - "6": ["1", "2", "3", "4", "5", "10"], - "7": ["1", "3", "5", "10"], - "7p": ["6", "10"], - "8": ["1", "1b", "1c", "5", "10"], -} +LR_ROBUSTNESS = [2e-4, 5e-4, 2e-3] +MASK_RATIO_ROBUSTNESS = [0.3, 0.75] +TRANSFER_PAIRS = [("miiv", "eicu"), ("eicu", "miiv")] STATE_FILE = Path("outputs/experiment_state.json") LOG_DIR = Path("logs/runner") -# Protocol defaults PROTO_A = {"freeze_encoder": True, "max_epochs": 50, "patience": 10, "lr": 1e-4} PROTO_B = {"freeze_encoder": False, "max_epochs": 100, "patience": 10, "lr": 3e-4} @@ -139,41 +120,84 @@ # --------------------------------------------------------------------------- # Data Model # --------------------------------------------------------------------------- + + @dataclass class Run: id: str - sprint: str + experiment_class: str run_type: str # "pretrain" | "finetune" | "supervised" | "gru_d" | "xgboost" - paradigm: str # "mae" | "jepa" | "contrastive" | "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) - # For transfer learning: dataset the encoder was pretrained on source_dataset: str | None = None - # Upstream pretrain metadata for provenance tracking in finetune runs upstream_pretrain_lr: float | None = None upstream_pretrain_mask_ratio: float | None = None - experiment_subtype: str | None = None # "lr_ablation" | "mask_ablation" | None + experiment_subtype: str | None = None + model_size: str | None = None - def build_command(self, runs_by_id: dict[str, Run]) -> list[str]: + 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: + if self.freeze_encoder is True: + return "A" + if self.freeze_encoder is False or self.run_type in {"supervised", "gru_d", "xgboost"}: + return "B" + return None + + 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() - elif self.run_type == "finetune": + if self.run_type == "finetune": return self._finetune_cmd(runs_by_id) - elif self.run_type == "supervised": + if self.run_type == "supervised": return self._supervised_cmd() - elif self.run_type == "gru_d": + if self.run_type == "gru_d": return self._gru_d_cmd() - elif self.run_type == "xgboost": + if self.run_type == "xgboost": return self._xgboost_cmd() - else: - raise ValueError(f"Unknown run_type: {self.run_type}") + 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 is not None: + overrides.append(f"+protocol={self.protocol}") + 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(): + cmd.append(f"ckpt_path={last_ckpt}") def _pretrain_cmd(self) -> list[str]: cmd = [ @@ -184,22 +208,17 @@ def _pretrain_cmd(self) -> list[str]: f"dataset={self.dataset}", f"ssl={self.paradigm}", f"seed={self.seed}", - f"sprint={self.sprint}", f"hydra.run.dir={self.output_dir}", ] - last_ckpt = Path(self.output_dir) / "checkpoints" / "last.ckpt" - if last_ckpt.exists(): - cmd.append(f"ckpt_path={last_ckpt}") - for k, v in self.extra_overrides.items(): - cmd.append(f"{k}={v}") + self._append_resume(cmd) + cmd.extend(self._metadata_overrides()) return cmd - def _finetune_cmd(self, runs_by_id: dict[str, Run]) -> list[str]: - # Find the pretrain dependency to get encoder path - pretrain_id = [d for d in self.depends_on if d in runs_by_id] - if not pretrain_id: - raise ValueError(f"Finetune run {self.id} has no pretrain dependency") - pretrain_dir = runs_by_id[pretrain_id[0]].output_dir + 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", @@ -210,42 +229,39 @@ def _finetune_cmd(self, runs_by_id: dict[str, Run]) -> list[str]: f"checkpoint={pretrain_dir}/encoder.pt", f"tasks={self.task}", f"seed={self.seed}", - f"sprint={self.sprint}", f"hydra.run.dir={self.output_dir}", ] - last_ckpt = Path(self.output_dir) / "checkpoints" / "last.ckpt" - if last_ckpt.exists(): - cmd.append(f"ckpt_path={last_ckpt}") + self._append_resume(cmd) if self.freeze_encoder is True: - cmd += [ - "training.freeze_encoder=true", - f"training.max_epochs={PROTO_A['max_epochs']}", - f"training.early_stopping_patience={PROTO_A['patience']}", - f"optimizer.lr={PROTO_A['lr']}", - "task.head_type=linear", - "task.hidden_dims=[]", - "task.dropout=0.0", - ] + cmd.extend( + [ + "training.freeze_encoder=true", + f"training.max_epochs={PROTO_A['max_epochs']}", + f"training.early_stopping_patience={PROTO_A['patience']}", + f"optimizer.lr={PROTO_A['lr']}", + "task.head_type=linear", + "task.hidden_dims=[]", + "task.dropout=0.0", + ] + ) elif self.freeze_encoder is False: - cmd += [ - "training.freeze_encoder=false", - f"training.max_epochs={PROTO_B['max_epochs']}", - f"training.early_stopping_patience={PROTO_B['patience']}", - f"optimizer.lr={PROTO_B['lr']}", - ] + cmd.extend( + [ + "training.freeze_encoder=false", + f"training.max_epochs={PROTO_B['max_epochs']}", + f"training.early_stopping_patience={PROTO_B['patience']}", + f"optimizer.lr={PROTO_B['lr']}", + ] + ) if self.label_fraction < 1.0: cmd.append(f"label_fraction={self.label_fraction}") - for k, v in self.extra_overrides.items(): - cmd.append(f"{k}={v}") if self.source_dataset is not None: cmd.append(f"+source_dataset={self.source_dataset}") - # Propagate upstream pretrain metadata so it lands in W&B config 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}") - if self.experiment_subtype is not None: - cmd.append(f"+experiment_subtype={self.experiment_subtype}") + cmd.extend(self._metadata_overrides()) return cmd def _supervised_cmd(self) -> list[str]: @@ -257,16 +273,12 @@ def _supervised_cmd(self) -> list[str]: f"dataset={self.dataset}", f"tasks={self.task}", f"seed={self.seed}", - f"sprint={self.sprint}", f"hydra.run.dir={self.output_dir}", ] - last_ckpt = Path(self.output_dir) / "checkpoints" / "last.ckpt" - if last_ckpt.exists(): - cmd.append(f"ckpt_path={last_ckpt}") + self._append_resume(cmd) if self.label_fraction < 1.0: cmd.append(f"label_fraction={self.label_fraction}") - for k, v in self.extra_overrides.items(): - cmd.append(f"{k}={v}") + cmd.extend(self._metadata_overrides()) return cmd def _gru_d_cmd(self) -> list[str]: @@ -280,16 +292,12 @@ def _gru_d_cmd(self) -> list[str]: f"dataset={self.dataset}", f"tasks={self.task}", f"seed={self.seed}", - f"sprint={self.sprint}", f"hydra.run.dir={self.output_dir}", ] - last_ckpt = Path(self.output_dir) / "checkpoints" / "last.ckpt" - if last_ckpt.exists(): - cmd.append(f"ckpt_path={last_ckpt}") + self._append_resume(cmd) if self.label_fraction < 1.0: cmd.append(f"label_fraction={self.label_fraction}") - for k, v in self.extra_overrides.items(): - cmd.append(f"{k}={v}") + cmd.extend(self._metadata_overrides()) return cmd def _xgboost_cmd(self) -> list[str]: @@ -301,77 +309,110 @@ def _xgboost_cmd(self) -> list[str]: f"dataset={self.dataset}", f"tasks={self.task}", f"seed={self.seed}", - f"sprint={self.sprint}", f"hydra.run.dir={self.output_dir}", ] if self.label_fraction < 1.0: cmd.append(f"label_fraction={self.label_fraction}") - for k, v in self.extra_overrides.items(): - cmd.append(f"{k}={v}") + 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: - """Canonical key for deduplicating pretrain runs.""" - key = f"pretrain_{paradigm}_{dataset}_seed{seed}" - if extra: - for k, v in sorted(extra.items()): - short_k = k.split(".")[-1] - short_v = str(v).replace(".", "") - key += f"_{short_k}{short_v}" - return key + return f"pretrain_{paradigm}_{dataset}_seed{seed}{_name_suffix(extra)}" -def _output_dir(sprint: str, name: str) -> str: - return f"outputs/sprint{sprint}/{name}" +def _output_dir(experiment_class: str, run_key: str) -> str: + return f"outputs/{experiment_class}/{run_key}" class MatrixBuilder: - """Generates all Run objects across sprints with pretrain deduplication.""" + """Generate the final thesis matrix by experiment class.""" - def __init__(self): + def __init__(self) -> None: self.runs: list[Run] = [] - self.pretrain_index: dict[str, Run] = {} # canonical_key -> 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, sprint: str, paradigm: str, dataset: str, seed: int, extra: dict | None = None + 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: - """Add a pretrain run, deduplicating by canonical key.""" extra = extra or {} - key = _pretrain_key(paradigm, dataset, seed, extra) - if key in self.pretrain_index: - return self.pretrain_index[key] + index_key = self._pretrain_index_key(paradigm, dataset, seed, extra) + if index_key in self.pretrain_index: + return self.pretrain_index[index_key] - dir_name = _pretrain_key(paradigm, dataset, seed, extra) + run_key = _pretrain_key(paradigm, dataset, seed, extra) run = Run( - id=f"s{sprint}_{dir_name}", - sprint=sprint, + 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(sprint, dir_name), + output_dir=_output_dir(experiment_class, run_key), extra_overrides=extra, + experiment_subtype=experiment_subtype, + model_size=model_size, ) - self.runs.append(run) - self.pretrain_index[key] = run + 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 + self, + paradigm: str, + dataset: str, + seed: int, + extra: dict | None = None, ) -> Run: - """Retrieve an existing pretrain run.""" - key = _pretrain_key(paradigm, dataset, seed, extra) - if key not in self.pretrain_index: - raise KeyError(f"Pretrain not found: {key}") - return self.pretrain_index[key] + 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, - sprint: str, + experiment_class: str, paradigm: str, dataset: str, seed: int, @@ -385,626 +426,432 @@ def _add_finetune( upstream_pretrain_lr: float | None = None, upstream_pretrain_mask_ratio: float | None = None, experiment_subtype: str | None = None, + model_size: str | None = None, ) -> Run: - """Add a finetune run. - - Args: - extra: Hydra overrides added to the finetune command AND name. - name_extra: Dict used only for directory/ID disambiguation - (e.g. pretrain ablation params). Not passed to command. - upstream_pretrain_lr: LR used in the upstream pretrain run (for provenance). - upstream_pretrain_mask_ratio: Mask ratio used in upstream pretrain run. - experiment_subtype: "lr_ablation", "mask_ablation", or None. - """ prefix = "probe" if freeze else "finetune" - name = f"{prefix}_{paradigm}_{task}_{dataset}_seed{seed}" + run_key = f"{prefix}_{paradigm}_{task}_{dataset}_seed{seed}" if source_dataset: - name += f"_from_{source_dataset}" + run_key += f"_from_{source_dataset}" if label_fraction < 1.0: - frac_str = str(label_fraction).replace(".", "") - name += f"_frac{frac_str}" - # Append both extra and name_extra to the directory name - for d in [name_extra, extra]: - if d: - for k, v in sorted(d.items()): - short_k = k.split(".")[-1] - short_v = str(v).replace(".", "") - name += f"_{short_k}{short_v}" - - run = Run( - id=f"s{sprint}_{name}", - sprint=sprint, - run_type="finetune", - paradigm=paradigm, - dataset=dataset, - seed=seed, - output_dir=_output_dir(sprint, name), - 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, + 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, + ) ) - self.runs.append(run) - return run def _add_supervised( - self, sprint: str, dataset: str, seed: int, task: str, label_fraction: float = 1.0 + 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: - name = f"supervised_{task}_{dataset}_seed{seed}" + run_key = f"supervised_{task}_{dataset}_seed{seed}" + if model_size: + run_key += f"_{model_size}" if label_fraction < 1.0: - frac_str = str(label_fraction).replace(".", "") - name += f"_frac{frac_str}" - run = Run( - id=f"s{sprint}_{name}", - sprint=sprint, - run_type="supervised", - paradigm="supervised", - dataset=dataset, - seed=seed, - output_dir=_output_dir(sprint, name), - task=task, - label_fraction=label_fraction, + 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, + ) ) - self.runs.append(run) - return run def _add_gru_d( - self, sprint: str, dataset: str, seed: int, task: str, label_fraction: float = 1.0 + self, + experiment_class: str, + dataset: str, + seed: int, + task: str, + label_fraction: float = 1.0, ) -> Run: - name = f"gru_d_{task}_{dataset}_seed{seed}" + run_key = f"gru_d_{task}_{dataset}_seed{seed}" if label_fraction < 1.0: - frac_str = str(label_fraction).replace(".", "") - name += f"_frac{frac_str}" - run = Run( - id=f"s{sprint}_{name}", - sprint=sprint, - run_type="gru_d", - paradigm="gru_d", - dataset=dataset, - seed=seed, - output_dir=_output_dir(sprint, name), - task=task, - label_fraction=label_fraction, + 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, + ) ) - self.runs.append(run) - return run def _add_xgboost( - self, sprint: str, dataset: str, seed: int, task: str, label_fraction: float = 1.0 + self, + experiment_class: str, + dataset: str, + seed: int, + task: str, + label_fraction: float = 1.0, ) -> Run: - name = f"xgboost_{task}_{dataset}_seed{seed}" + run_key = f"xgboost_{task}_{dataset}_seed{seed}" if label_fraction < 1.0: - frac_str = str(label_fraction).replace(".", "") - name += f"_frac{frac_str}" - run = Run( - id=f"s{sprint}_{name}", - sprint=sprint, - run_type="xgboost", - paradigm="xgboost", - dataset=dataset, - seed=seed, - output_dir=_output_dir(sprint, name), - task=task, - label_fraction=label_fraction, + 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, + ) ) - self.runs.append(run) - return run - @staticmethod - def _add_model_size_metadata(run: Run, size_name: str) -> Run: - """Add Hydra metadata used by W&B naming/tagging without changing run IDs.""" - run.extra_overrides = {**run.extra_overrides, "+model_size": size_name} - return run - - # --- Sprint builders --- - - def build_sprint1(self): - """MIMIC, all tasks, Protocol B + supervised, seed=42. - - Classical baselines (GRU-D, XGBoost) are canonicalized in Sprint 11 - so the final thesis matrix reports them once rather than duplicating - them across the early SSL/supervised sprints. - """ - ds, seed, sprint = "miiv", 42, "1" - for p in SSL_PARADIGMS: - pt = self._add_pretrain(sprint, p, ds, seed) - for task in TASKS: - self._add_finetune(sprint, p, ds, seed, task, False, pt) - for task in TASKS: - self._add_supervised(sprint, ds, seed, task) - - def build_sprint1b(self): - """LR sensitivity, MIMIC, mortality_24h, seed=42.""" - ds, seed, task, sprint = "miiv", 42, "mortality_24h", "1b" - for p in SSL_PARADIGMS: - for lr in LR_ABLATION: - extra = {"optimizer.lr": lr} - pt = self._add_pretrain(sprint, p, ds, seed, extra) - self._add_finetune( - sprint, - p, - ds, - seed, - task, - False, - pt, - name_extra=extra, - upstream_pretrain_lr=lr, - experiment_subtype="lr_ablation", - ) - - def build_sprint1c(self): - """Mask ratio sensitivity, MIMIC, mortality_24h, seed=42.""" - ds, seed, task, sprint = "miiv", 42, "mortality_24h", "1c" - for p in SSL_PARADIGMS: - for mr in MASK_RATIO_ABLATION: - extra = {"ssl.mask_ratio": mr} - pt = self._add_pretrain(sprint, p, ds, seed, extra) - self._add_finetune( - sprint, - p, - ds, - seed, - task, - False, - pt, - name_extra=extra, - upstream_pretrain_mask_ratio=mr, - experiment_subtype="mask_ablation", - ) - - def build_sprint2(self): - """MIMIC Protocol A, seed=42 — reuses Sprint 1 pretrains.""" - ds, seed, sprint = "miiv", 42, "2" - for p in SSL_PARADIGMS: - pt = self._get_pretrain(p, ds, seed) - for task in TASKS: - self._add_finetune(sprint, p, ds, seed, task, True, pt) - - def build_sprint3(self): - """eICU, both protocols + supervised, seed=42. - - Classical baselines are launched in Sprint 11 only. - """ - ds, seed, sprint = "eicu", 42, "3" - for p in SSL_PARADIGMS: - pt = self._add_pretrain(sprint, p, ds, seed) - for task in TASKS: - self._add_finetune(sprint, p, ds, seed, task, False, pt) - self._add_finetune(sprint, p, ds, seed, task, True, pt) - for task in TASKS: - self._add_supervised(sprint, ds, seed, task) - - def build_sprint4(self): - """Combined, both protocols + supervised, seed=42. - - Classical baselines are launched in Sprint 11 only. - """ - ds, seed, sprint = "combined", 42, "4" - for p in SSL_PARADIGMS: - pt = self._add_pretrain(sprint, p, ds, seed) - for task in TASKS: - self._add_finetune(sprint, p, ds, seed, task, False, pt) - self._add_finetune(sprint, p, ds, seed, task, True, pt) - for task in TASKS: - self._add_supervised(sprint, ds, seed, task) - - def build_sprint5(self): - """Seeds 123, 456 for datasets miiv, eicu, combined. - - Extends the SSL + supervised matrix to three seeds. Classical baselines - are canonicalized in Sprint 11. - """ - sprint = "5" - for seed in [123, 456]: - for ds in DATASETS: - for p in SSL_PARADIGMS: - pt = self._add_pretrain(sprint, p, ds, seed) + 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(sprint, p, ds, seed, task, False, pt) - self._add_finetune(sprint, p, ds, seed, task, True, pt) + 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(sprint, ds, seed, task) - - def build_sprint6(self): - """Label efficiency ablation — reuses Phase 1 encoders. - - Includes SSL Protocol A/B plus the supervised Transformer baseline. - Classical baselines are canonicalized separately in Sprint 11. - """ - sprint = "6" - for seed in SEEDS: - for ds in DATASETS: - for p in SSL_PARADIGMS: - pt = self._get_pretrain(p, ds, seed) - # mortality_24h gets full sweep - for frac in LABEL_FRACTIONS_FULL: - self._add_finetune(sprint, p, ds, seed, "mortality_24h", False, pt, frac) - self._add_finetune(sprint, p, ds, seed, "mortality_24h", True, pt, frac) - # Other tasks get trend check - for task in TASKS[1:]: - for frac in LABEL_FRACTIONS_TREND: - self._add_finetune(sprint, p, ds, seed, task, False, pt, frac) - self._add_finetune(sprint, p, ds, seed, task, True, pt, frac) - # Supervised label efficiency - for frac in LABEL_FRACTIONS_FULL: - self._add_supervised(sprint, ds, seed, "mortality_24h", frac) - for task in TASKS[1:]: - for frac in LABEL_FRACTIONS_TREND: - self._add_supervised(sprint, ds, seed, task, frac) - - def build_sprint7p(self): - """Focused model-capacity study for the thesis. - - MIIV only, mortality_24h only, MAE + supervised. - Two larger model sizes (medium=128d/4L, large=256d/4L) at 3 label fractions - across all five thesis seeds. Baseline comparisons come from Sprint 6 - (seeds 42/123/456) and Sprint 10 (seeds 789/1011), so the default-size - baseline is inherited rather than rerun here. - - Total: 5 seeds × [2 pretrain + 2×3 finetune + 2×3 probe + 2×3 supervised] = 100 runs. - """ - sprint = "7p" - ds, task = "miiv", "mortality_24h" + self._add_supervised(experiment_class, dataset, seed, task) + def build_label_efficiency(self) -> None: + experiment_class = "label_efficiency" for seed in SEEDS_EXTENDED: - for size_name, size_cfg in MODEL_SIZES.items(): - # Common model override for all runs at this size - model_extra = {"model": size_cfg["model"]} - - # --- MAE pretrain (new encoder size, full data) --- - pretrain_extra = { - **model_extra, - **size_cfg["ssl_scale"]["mae"], - } - pt = self._add_pretrain(sprint, "mae", ds, seed, extra=pretrain_extra) - self._add_model_size_metadata(pt, size_name) - - # --- MAE finetune (Protocol B) + probe (Protocol A) --- - finetune_extra = {**model_extra} - for frac in LABEL_FRACTIONS_PILOT: - self._add_model_size_metadata( + for dataset in DATASETS: + for paradigm in SSL_PARADIGMS: + pretrain = self._get_pretrain(paradigm, dataset, seed) + for frac in LABEL_FRACTIONS_FULL: self._add_finetune( - sprint, - "mae", - ds, + experiment_class, + paradigm, + dataset, seed, - task, + "mortality_24h", False, - pt, - label_fraction=frac, - extra=finetune_extra, - name_extra={"size": size_name}, - ), - size_name, - ) - self._add_model_size_metadata( + pretrain, + frac, + ) self._add_finetune( - sprint, - "mae", - ds, + experiment_class, + paradigm, + dataset, seed, - task, + "mortality_24h", True, - pt, - label_fraction=frac, - extra=finetune_extra, - name_extra={"size": size_name}, - ), - size_name, - ) - - # --- Supervised baseline (same larger encoder, from scratch) --- - sup_extra = {**model_extra} - for frac in LABEL_FRACTIONS_PILOT: - sup_name = f"supervised_{task}_{ds}_seed{seed}_{size_name}" - if frac < 1.0: - frac_str = str(frac).replace(".", "") - sup_name += f"_frac{frac_str}" - run = Run( - id=f"s{sprint}_{sup_name}", - sprint=sprint, - run_type="supervised", - paradigm="supervised", - dataset=ds, - seed=seed, - output_dir=_output_dir(sprint, sup_name), - task=task, - label_fraction=frac, - extra_overrides=sup_extra, - ) - self._add_model_size_metadata(run, size_name) - self.runs.append(run) - - def build_sprint7(self): - """Cross-dataset transfer — Protocol B only (full finetune).""" - sprint = "7" - for seed in SEEDS: - for source_ds, target_ds in TRANSFER_PAIRS: - for p in SSL_PARADIGMS: - pt = self._get_pretrain(p, source_ds, seed) - for task in TASKS: - self._add_finetune( - sprint, p, target_ds, seed, task, False, pt, source_dataset=source_ds + pretrain, + frac, ) - - def build_sprint8(self): - """LR + mask ablations, seeds 123, 456.""" - sprint = "8" - for seed in [123, 456]: - for p in SSL_PARADIGMS: - # LR ablation - for lr in LR_ABLATION: - extra = {"optimizer.lr": lr} - pt = self._add_pretrain(sprint, p, "miiv", seed, extra) - self._add_finetune( - sprint, - p, - "miiv", - seed, - "mortality_24h", - False, - pt, - name_extra=extra, - upstream_pretrain_lr=lr, - experiment_subtype="lr_ablation", - ) - # Mask ratio ablation - for mr in MASK_RATIO_ABLATION: - extra = {"ssl.mask_ratio": mr} - pt = self._add_pretrain(sprint, p, "miiv", seed, extra) - self._add_finetune( - sprint, - p, - "miiv", - seed, - "mortality_24h", - False, - pt, - name_extra=extra, - upstream_pretrain_mask_ratio=mr, - experiment_subtype="mask_ablation", - ) - - def build_sprint10(self): - """Seeds 789, 1011 for Sprints 1-8 scope (SSL + supervised).""" - sprint = "10" - new_seeds = [789, 1011] - - # --- Core experiments (Sprint 1-4 scope) --- - for seed in new_seeds: - for ds in DATASETS: - for p in SSL_PARADIGMS: - pt = self._add_pretrain(sprint, p, ds, seed) - for task in TASKS: - self._add_finetune(sprint, p, ds, seed, task, False, pt) - self._add_finetune(sprint, p, ds, seed, task, True, pt) - for task in TASKS: - self._add_supervised(sprint, ds, seed, task) - - # --- Label efficiency (Sprint 6 scope) --- - for seed in new_seeds: - for ds in DATASETS: - for p in SSL_PARADIGMS: - pt = self._get_pretrain(p, ds, seed) - for frac in LABEL_FRACTIONS_FULL: - self._add_finetune(sprint, p, ds, seed, "mortality_24h", False, pt, frac) - self._add_finetune(sprint, p, ds, seed, "mortality_24h", True, pt, frac) for task in TASKS[1:]: for frac in LABEL_FRACTIONS_TREND: - self._add_finetune(sprint, p, ds, seed, task, False, pt, frac) - self._add_finetune(sprint, p, ds, seed, task, True, pt, frac) + 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_FULL: - self._add_supervised(sprint, ds, seed, "mortality_24h", frac) + self._add_supervised(experiment_class, dataset, seed, "mortality_24h", frac) for task in TASKS[1:]: for frac in LABEL_FRACTIONS_TREND: - self._add_supervised(sprint, ds, seed, task, frac) + self._add_supervised(experiment_class, dataset, seed, task, frac) - # --- Cross-dataset transfer (Sprint 7 scope) --- - for seed in new_seeds: - for source_ds, target_ds in TRANSFER_PAIRS: - for p in SSL_PARADIGMS: - pt = self._get_pretrain(p, source_ds, seed) + 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( - sprint, - p, - target_ds, + experiment_class, + paradigm, + target_dataset, seed, task, False, - pt, - source_dataset=source_ds, + pretrain, + source_dataset=source_dataset, ) - # --- HP ablations (Sprint 8 scope) --- - for seed in new_seeds: - for p in SSL_PARADIGMS: - for lr in LR_ABLATION: + 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} - pt = self._add_pretrain(sprint, p, "miiv", seed, extra) + subtype = "lr_sensitivity" + pretrain = self._add_pretrain( + experiment_class, + paradigm, + dataset, + seed, + extra, + experiment_subtype=subtype, + ) self._add_finetune( - sprint, - p, - "miiv", + experiment_class, + paradigm, + dataset, seed, - "mortality_24h", + task, False, - pt, + pretrain, name_extra=extra, upstream_pretrain_lr=lr, - experiment_subtype="lr_ablation", + experiment_subtype=subtype, + ) + for mask_ratio in MASK_RATIO_ROBUSTNESS: + extra = {"ssl.mask_ratio": mask_ratio} + subtype = "mask_ratio_sensitivity" + pretrain = self._add_pretrain( + experiment_class, + paradigm, + dataset, + seed, + extra, + experiment_subtype=subtype, ) - for mr in MASK_RATIO_ABLATION: - extra = {"ssl.mask_ratio": mr} - pt = self._add_pretrain(sprint, p, "miiv", seed, extra) self._add_finetune( - sprint, - p, - "miiv", + experiment_class, + paradigm, + dataset, seed, - "mortality_24h", + task, False, - pt, + pretrain, name_extra=extra, - upstream_pretrain_mask_ratio=mr, - experiment_subtype="mask_ablation", + upstream_pretrain_mask_ratio=mask_ratio, + experiment_subtype=subtype, ) - def build_sprint11(self): - """Canonical classical baselines (XGBoost + GRU-D), 5 seeds. + 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, + ) - These baselines are intentionally centralized here rather than being - duplicated across earlier sprints. That keeps the final thesis matrix - aligned with the docs and export logic: one canonical experiment family - for contextual non-SSL baselines. - """ - sprint = "11" + def build_classical_baselines(self) -> None: + experiment_class = "classical_baselines" for seed in SEEDS_EXTENDED: - for ds in DATASETS: + for dataset in DATASETS: for task in TASKS: - self._add_xgboost(sprint, ds, seed, task) - self._add_gru_d(sprint, ds, seed, task) - # Baseline label-efficiency (mirrors Sprint 6 baseline scope) + self._add_xgboost(experiment_class, dataset, seed, task) + self._add_gru_d(experiment_class, dataset, seed, task) for frac in LABEL_FRACTIONS_FULL: - self._add_gru_d(sprint, ds, seed, "mortality_24h", frac) - self._add_xgboost(sprint, ds, seed, "mortality_24h", frac) - for task_name in TASKS[1:]: + 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_gru_d(sprint, ds, seed, task_name, frac) - self._add_xgboost(sprint, ds, seed, task_name, frac) - - def build_sprint12(self): - """SMART (NeurIPS 2024) as external SSL reference, 5 seeds. - - SMART uses its own encoder (MART architecture, d_model=32) and element-wise - masking — NOT part of the controlled comparison (different architecture). - Included as an external SSL SOTA reference point in the appendix. - """ - sprint = "12" - # model=smart selects the MART encoder; ssl=smart is set via paradigm arg + 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"} - for seed in SEEDS_EXTENDED: - for ds in DATASETS: - # Pretrain (ssl=smart is set by _pretrain_cmd via self.paradigm) - pt = self._add_pretrain(sprint, "smart", ds, seed, pretrain_extra) - # Both protocols + all tasks + for dataset in DATASETS: + pretrain = self._add_pretrain( + experiment_class, "smart", dataset, seed, pretrain_extra + ) for task in TASKS: self._add_finetune( - sprint, + experiment_class, "smart", - ds, + dataset, seed, task, False, - pt, + pretrain, extra=finetune_extra, ) self._add_finetune( - sprint, + experiment_class, "smart", - ds, + dataset, seed, task, True, - pt, + pretrain, extra=finetune_extra, ) - def build_sprint13(self): - """TS2Vec temporal contrastive variant, 5 seeds, both protocols. - - Addresses the "contrastive was set up to fail" vulnerability by giving - the contrastive paradigm its natural augmentations (noise + masking) - and a temporal contrastive loss. Same encoder, same training budget. - Evaluated under both Protocol A (probe) and Protocol B (full finetune) - so it fits the main thesis A/B framing. - """ - sprint = "13" - for seed in SEEDS_EXTENDED: - for ds in DATASETS: - pt = self._add_pretrain(sprint, "ts2vec", ds, seed) - for task in TASKS: - self._add_finetune(sprint, "ts2vec", ds, seed, task, False, pt) - self._add_finetune(sprint, "ts2vec", ds, seed, task, True, pt) - def build_all(self) -> list[Run]: - """Build full experiment matrix. Order matters for dedup.""" - self.build_sprint1() - self.build_sprint1b() - self.build_sprint1c() - self.build_sprint2() - self.build_sprint3() - self.build_sprint4() - self.build_sprint5() - self.build_sprint6() - self.build_sprint7p() - self.build_sprint7() - self.build_sprint8() - self.build_sprint10() - self.build_sprint11() - self.build_sprint12() - self.build_sprint13() + 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]: - builder = MatrixBuilder() - return builder.build_all() - - -def filter_runs_for_sprints(all_runs: list[Run], sprints: list[str]) -> list[Run]: - """Filter runs while preserving the caller's requested sprint order.""" - sprint_order: dict[str, int] = {} - for sprint in sprints: - sprint_order.setdefault(sprint, len(sprint_order)) - - indexed_runs = [(idx, run) for idx, run in enumerate(all_runs) if run.sprint in sprint_order] - indexed_runs.sort(key=lambda item: (sprint_order[item[1].sprint], item[0])) - return [run for _, run in indexed_runs] + 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]: - """Post-process runs to inject revision into IDs, output dirs, and overrides. - - Rewrites: - - Run IDs: s1_pretrain_... → s1_rev-v2_pretrain_... - - Output dirs: outputs/sprint1/... → outputs/sprint1_rev-v2/... - - depends_on: remap references within the revised set - - extra_overrides: inject revision= and rerun_reason= - """ + """Inject revision metadata into IDs, output dirs, dependencies, and overrides.""" old_to_new: dict[str, str] = {} - for r in runs: - old_id = r.id - # Insert rev- after the sprint prefix (e.g. s1_ → s1_rev-v2_) - prefix = f"s{r.sprint}_" - if old_id.startswith(prefix): - new_id = f"{prefix}rev-{revision}_{old_id[len(prefix) :]}" - else: - new_id = f"rev-{revision}_{old_id}" - old_to_new[old_id] = new_id - - r.id = new_id - # Rewrite output dir: sprint1/ → sprint1_rev-v2/ - r.output_dir = r.output_dir.replace( - f"outputs/sprint{r.sprint}/", - f"outputs/sprint{r.sprint}_rev-{revision}/", + 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 ) - r.extra_overrides["revision"] = revision + run.extra_overrides["revision"] = revision if reason: - r.extra_overrides["rerun_reason"] = reason + run.extra_overrides["rerun_reason"] = reason - # Remap depends_on references within the revised set - for r in runs: - r.depends_on = [old_to_new.get(d, d) for d in r.depends_on] + for run in runs: + run.depends_on = [old_to_new.get(dep_id, dep_id) for dep_id in run.depends_on] return runs @@ -1017,7 +864,6 @@ def apply_wandb_target( """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 @@ -1031,7 +877,6 @@ def apply_launch_commit(runs: list[Run], launch_commit: str | None = None) -> li """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 @@ -1040,15 +885,16 @@ def apply_launch_commit(runs: list[Run], launch_commit: str | None = None) -> li # --------------------------------------------------------------------------- # 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): +def save_state(state: dict) -> None: STATE_FILE.parent.mkdir(parents=True, exist_ok=True) - # Atomic write fd, tmp = tempfile.mkstemp(dir=STATE_FILE.parent, suffix=".tmp") try: with os.fdopen(fd, "w") as f: @@ -1066,12 +912,10 @@ 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): - if run_id not in state["runs"]: - state["runs"][run_id] = {} - entry = state["runs"][run_id] - entry["status"] = status - entry.update(kwargs) +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 is_pid_alive(pid: int) -> bool: @@ -1083,7 +927,6 @@ def is_pid_alive(pid: int) -> bool: def _pid_matches_command(pid: int, expected_command: str | None) -> bool: - """Best-effort guard against PID reuse before trusting a running state.""" if not expected_command: return True try: @@ -1098,23 +941,20 @@ def _pid_matches_command(pid: int, expected_command: str | None) -> bool: if result.returncode != 0: return False actual_command = result.stdout.strip() - if not actual_command: - return False - return expected_command == actual_command or expected_command in actual_command + return bool(actual_command) and ( + expected_command == actual_command or expected_command in actual_command + ) -def recover_stale_running(state: dict): +def recover_stale_running(state: dict) -> None: """Reset running entries whose PIDs are dead or no longer match the run.""" 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(pid) - or not _pid_matches_command( - int(pid), - info.get("command"), - ) + or not is_pid_alive(int(pid)) + or not _pid_matches_command(int(pid), info.get("command")) ): info["status"] = "pending" info.pop("pid", None) @@ -1122,9 +962,6 @@ def recover_stale_running(state: dict): print(f" Recovered stale run: {run_id}") -# Scheduler slot weights. Pretraining occupies the full default budget of 4 -# slots so `--parallel 4` does not mix heavyweight pretrains with downstream -# GPU jobs on a single device. RUN_SLOT_COSTS = { "pretrain": 4, "finetune": 1, @@ -1135,7 +972,6 @@ def recover_stale_running(state: dict): def _slot_cost(run: Run, slot_budget: int) -> int: - """Return the scheduler slot cost for a run under the active budget.""" return min(RUN_SLOT_COSTS.get(run.run_type, 1), slot_budget) @@ -1145,36 +981,31 @@ def _select_ready_runs( runs_by_id: dict[str, Run], slot_budget: int, ) -> list[Run]: - """Choose a launch batch that fits the available scheduler slots. - - Heavier runs are prioritized first so pretrains claim the budget before - lighter downstream jobs, avoiding accidental co-scheduling on one GPU. - """ - active_slots = sum(_slot_cost(runs_by_id[rid], slot_budget) for rid in active_run_ids) + 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 = {r.id: r for r in runs} - active: dict[str, subprocess.Popen] = {} # run_id -> Popen + runs_by_id = {run.id: run for run in runs} + active: dict[str, subprocess.Popen] = {} shutting_down = False def handle_signal(signum, frame): @@ -1183,22 +1014,19 @@ def handle_signal(signum, frame): return shutting_down = True print(f"\nReceived signal {signum}, shutting down gracefully...") - for rid, proc in active.items(): + for proc in active.values(): try: proc.terminate() except OSError: pass - # Wait up to 30s for children deadline = time.time() + 30 - for rid, proc in list(active.items()): - remaining = max(0, deadline - time.time()) + for proc in list(active.values()): try: - proc.wait(timeout=remaining) + proc.wait(timeout=max(0, deadline - time.time())) except subprocess.TimeoutExpired: proc.kill() - # Reset interrupted runs to pending - for rid in active: - set_run_status(state, rid, "pending") + 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) @@ -1213,65 +1041,55 @@ def handle_signal(signum, frame): return 0 print(f"Scheduler started (slot_budget={parallel})") - print(f"State file: {STATE_FILE}") - print() + print(f"State file: {STATE_FILE}\n") while not shutting_down: - # 1. Poll active processes finished = [] - for rid, proc in active.items(): - ret = proc.poll() - if ret is not None: - finished.append((rid, ret)) - - # 2. Handle completions - for rid, exit_code in finished: - proc = active.pop(rid) - fh = getattr(proc, "_log_fh", None) - if fh: - fh.close() - run = runs_by_id[rid] + 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, rid, now) + elapsed = _format_elapsed(state, run_id, now) if exit_code == 0: - # For pretrain, verify encoder.pt exists - if run.run_type == "pretrain": - encoder_path = Path(run.output_dir) / "encoder.pt" - if not encoder_path.exists(): - print(f" FAILED {rid}: exit 0 but encoder.pt missing {elapsed}") - set_run_status(state, rid, "failed", finished_at=now, exit_code=exit_code) - _propagate_failure(rid, runs, state) - continue - print(f" DONE {rid} {elapsed}") - set_run_status(state, rid, "completed", finished_at=now, 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 {rid} (exit {exit_code}) {elapsed}") - set_run_status(state, rid, "failed", finished_at=now, exit_code=exit_code) - _propagate_failure(rid, runs, state) + 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) - # 3. Find ready runs ready = [] - for r in runs: - status = get_run_status(state, r.id) - if status != "pending": + for run in runs: + if get_run_status(state, run.id) != "pending": continue - if r.id in active: + if run.id in active: continue - # Check all dependencies completed - deps_ok = all(get_run_status(state, d) == "completed" for d in r.depends_on) - if deps_ok: - ready.append(r) + if all(get_run_status(state, dep_id) == "completed" for dep_id in run.depends_on): + ready.append(run) - # 4. Launch runs that fit the remaining scheduler slots launch_batch = _select_ready_runs(ready, set(active), runs_by_id, parallel) - for r in launch_batch: - cmd = r.build_command(runs_by_id) - log_file = LOG_DIR / f"{r.id}.log" + 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 {r.id}") + print(f" START {run.id}") log_fh = open(log_file, "w") proc = subprocess.Popen( cmd, @@ -1279,12 +1097,11 @@ def handle_signal(signum, frame): stderr=subprocess.STDOUT, start_new_session=True, ) - # Store log handle for cleanup proc._log_fh = log_fh # type: ignore[attr-defined] - active[r.id] = proc + active[run.id] = proc set_run_status( state, - r.id, + run.id, "running", started_at=now, pid=proc.pid, @@ -1292,29 +1109,23 @@ def handle_signal(signum, frame): command=shlex.join(cmd), ) - # 5. Save state save_state(state) - # 6. Exit condition if not active and not ready: break - time.sleep(2) - # Close any remaining log handles for proc in active.values(): - fh = getattr(proc, "_log_fh", None) - if fh: - fh.close() + log_fh = getattr(proc, "_log_fh", None) + if log_fh: + log_fh.close() - # Print summary print() _print_summary(runs, state) return _scheduler_exit_code(runs, state) def _format_elapsed(state: dict, run_id: str, now_iso: str) -> str: - """Format elapsed time as (Xh Ym Zs).""" started = state["runs"].get(run_id, {}).get("started_at") if not started: return "" @@ -1324,98 +1135,85 @@ def _format_elapsed(state: dict, run_id: str, now_iso: str) -> str: secs = int((t1 - t0).total_seconds()) if secs < 60: return f"({secs}s)" - elif secs < 3600: + if secs < 3600: return f"({secs // 60}m{secs % 60:02d}s)" - else: - h, rem = divmod(secs, 3600) - m, s = divmod(rem, 60) - return f"({h}h{m:02d}m{s:02d}s)" - except (ValueError, TypeError): + 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): - """Skip runs that depend on a failed run (transitive).""" +def _propagate_failure(failed_id: str, runs: list[Run], state: dict) -> None: queue = [failed_id] while queue: - fid = queue.pop(0) - for r in runs: - if fid in r.depends_on and get_run_status(state, r.id) == "pending": - set_run_status(state, r.id, "skipped", reason=f"dependency {fid} failed") - print(f" SKIP {r.id} (dep {fid} failed)") - queue.append(r.id) + 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]): - """Print all runs and their commands without executing.""" +def _print_dry_run(runs: list[Run], runs_by_id: dict[str, Run]) -> None: print(f"DRY RUN: {len(runs)} runs\n") - for r in runs: - deps = ", ".join(r.depends_on) if r.depends_on else "(none)" - cmd = shlex.join(r.build_command(runs_by_id)) - print(f"[{r.id}]") - print(f" sprint={r.sprint} type={r.run_type} deps={deps}") - print(f" dir={r.output_dir}") - print(f" cmd: {cmd}") - print() + 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: - """Extract revision name from a run ID (e.g. 's1_rev-v2_pretrain_...' → 'v2').""" - m = re.search(r"_rev-([^_]+)_", run_id) - return m.group(1) if m else None -def _sprint_display(sprint: str, revision: str | None) -> str: - """Format sprint column, e.g. '1' or '1/v2'.""" - if revision: - return f"{sprint}/{revision}" - return sprint +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_sprint_from_id(run_id: str) -> str | None: - """Extract sprint from a run ID (e.g. 's1_rev-v2_pretrain_...' → '1').""" - m = re.match(r"s([^_]+)_", run_id) - return m.group(1) if m 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(sprint_filter: list[str] | None = None): +def print_status(experiment_class_filter: list[str] | None = None) -> None: all_runs = generate_all_runs() state = load_state() - generated_ids = {r.id for r in all_runs} + generated_ids = {run.id for run in all_runs} - # Group generated runs by (sprint, revision=None) groups: dict[tuple[str, str | None], list[str]] = {} - for r in all_runs: - key = (r.sprint, None) - groups.setdefault(key, []).append(r.id) + for run in all_runs: + groups.setdefault((run.experiment_class, None), []).append(run.id) - # Add revised runs from state (not in generated set) by parsing their IDs for run_id in state.get("runs", {}): if run_id not in generated_ids and "_rev-" in run_id: - sprint = _extract_sprint_from_id(run_id) - rev = _extract_revision_from_id(run_id) - if sprint and rev: - key = (sprint, rev) - groups.setdefault(key, []).append(run_id) + experiment_class = _extract_experiment_class_from_id(run_id) + revision = _extract_revision_from_id(run_id) + if experiment_class and revision: + groups.setdefault((experiment_class, revision), []).append(run_id) - group_keys = sorted(groups.keys(), key=lambda k: (_sprint_sort_key(k[0]), k[1] or "")) - if sprint_filter: - group_keys = [k for k in group_keys if k[0] in sprint_filter] + group_keys = sorted(groups) + if experiment_class_filter: + allowed = set(experiment_class_filter) + group_keys = [key for key in group_keys if key[0] in allowed] print( - f"{'Sprint':>8} | {'Total':>5} | {'Done':>4} | {'Run':>3} | " + f"{'Experiment Class':>28} | {'Total':>5} | {'Done':>4} | {'Run':>3} | " f"{'Fail':>4} | {'Pend':>4} | {'Skip':>4}" ) - print("-" * 54) + print("-" * 76) totals = {"total": 0, "completed": 0, "running": 0, "failed": 0, "pending": 0, "skipped": 0} - - for sprint, rev in group_keys: - run_ids = groups[(sprint, rev)] - label = _sprint_display(sprint, rev) + 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, @@ -1424,49 +1222,42 @@ def print_status(sprint_filter: list[str] | None = None): "pending": 0, "skipped": 0, } - for rid in run_ids: - status = get_run_status(state, rid) - if status in counts: - counts[status] += 1 - else: - counts["pending"] += 1 # unknown → pending + 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:>8} | {counts['total']:>5} | {counts['completed']:>4} | " + f"{label:>28} | {counts['total']:>5} | {counts['completed']:>4} | " f"{counts['running']:>3} | {counts['failed']:>4} | " f"{counts['pending']:>4} | {counts['skipped']:>4}" ) - for k in totals: - totals[k] += counts[k] + for key in totals: + totals[key] += counts[key] - print("-" * 54) + print("-" * 76) print( - f"{'TOTAL':>8} | {totals['total']:>5} | {totals['completed']:>4} | " + 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): - """Print completion summary for a scheduler run.""" +def _print_summary(runs: list[Run], state: dict) -> None: counts = {"completed": 0, "failed": 0, "skipped": 0, "pending": 0} - for r in runs: - s = get_run_status(state, r.id) - counts[s] = counts.get(s, 0) + 1 + 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, " - f"{counts['pending']} pending" + f"{counts['failed']} failed, {counts['skipped']} skipped, {counts['pending']} pending" ) def _scheduler_exit_code(runs: list[Run], state: dict) -> int: - """Return nonzero when the requested run set did not complete cleanly.""" - failed = [r.id for r in runs if get_run_status(state, r.id) == "failed"] - skipped = [r.id for r in runs if get_run_status(state, r.id) == "skipped"] - running = [r.id for r in runs if get_run_status(state, r.id) == "running"] - pending = [r.id for r in runs if get_run_status(state, r.id) == "pending"] - + 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: " @@ -1474,113 +1265,42 @@ def _scheduler_exit_code(runs: list[Run], state: dict) -> int: f"{len(running)} running, {len(pending)} pending" ) return 1 - return 0 -def _sprint_sort_key(s: str) -> tuple[int, str]: - """Sort sprints: 1, 1b, 2, 3, ...""" - num = "" - suffix = "" - for c in s: - if c.isdigit(): - num += c - else: - suffix += c - return (int(num) if num else 0, suffix) - - -# --------------------------------------------------------------------------- -# CLI -# --------------------------------------------------------------------------- -def cmd_run(args): - all_runs = generate_all_runs() - state = load_state() - recover_stale_running(state) - - # Filter to requested sprints - sprints = [str(s) for s in args.sprint] - runs = filter_runs_for_sprints(all_runs, sprints) - - # Also include dependency runs from earlier sprints (pretrain reuse) - run_ids = {r.id for r in runs} - deps_needed = set() - for r in runs: - for d in r.depends_on: - if d not in run_ids: - deps_needed.add(d) - if deps_needed: - all_by_id = {r.id: r for r in all_runs} - extra = [all_by_id[d] for d in deps_needed if d in all_by_id] - runs = extra + runs - if extra and not args.dry_run: - print(f"Including {len(extra)} dependency run(s) from earlier sprints") - - if not runs: - print(f"No runs found for sprint(s): {', '.join(sprints)}") - return 0 - - # Apply revision transform if requested - if args.revision: - runs = apply_revision(runs, args.revision, args.reason) - runs = apply_wandb_target(runs, args.project, args.entity) - runs = apply_launch_commit(runs, args.launch_commit) - - print(f"Sprint(s) {', '.join(sprints)}: {len(runs)} runs") - if args.revision: - print(f"Revision: {args.revision}" + (f" ({args.reason})" if args.reason else "")) - if args.launch_commit: - print(f"Launch commit: {args.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): - sprint_filter = [str(s) for s in args.sprint] if args.sprint else None - print_status(sprint_filter) - - def _collect_dependency_closure( runs: list[Run], all_by_id: dict[str, Run], ) -> tuple[list[Run], dict[str, set[str]], set[str]]: - """Return transitive dependencies for runs, with provenance for reporting.""" deps: list[Run] = [] required_by: dict[str, set[str]] = {} missing: set[str] = set() seen: set[str] = set() - queue: list[tuple[str, str]] = [(dep_id, run.id) for run in runs for dep_id in run.depends_on] - + 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 _retry_command_suggestion(args, *, include_failed: bool, include_skipped: bool) -> str: - """Build a retry command that preserves the user's current retry scope.""" cmd = ["uv", "run", "python", "scripts/internal/run_experiments.py", "retry"] if include_failed: cmd.append("--failed") if include_skipped: cmd.append("--skipped") - if args.sprint: - cmd.append("--sprint") - cmd.extend(str(s) for s in args.sprint) + 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: @@ -1589,8 +1309,9 @@ def _retry_command_suggestion(args, *, include_failed: bool, include_skipped: bo cmd.extend(["--project", str(args.project)]) if args.entity: cmd.extend(["--entity", str(args.entity)]) - if args.launch_commit: - cmd.extend(["--launch-commit", str(args.launch_commit)]) + 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) @@ -1602,7 +1323,6 @@ def _fail_for_blocked_retry_dependencies( required_by: dict[str, set[str]], args, ) -> None: - """Print a dependency report for retry requests that cannot make progress.""" print("Cannot retry the selected runs because required dependencies are blocked.") if blocked: print("\nBlocked dependencies:") @@ -1614,7 +1334,6 @@ def _fail_for_blocked_retry_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 @@ -1627,53 +1346,94 @@ def _fail_for_blocked_retry_dependencies( raise SystemExit(1) +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) + + def cmd_retry(args): all_runs = generate_all_runs() - sprint_filter = {str(s) for s in args.sprint} if args.sprint else None + class_filter = set(args.experiment_class) if args.experiment_class else None - # Apply revision if specified (must happen before ID matching against state) if args.revision: - if sprint_filter: - revised = [r for r in all_runs if r.sprint in sprint_filter] - revised_ids = {r.id for r in revised} - all_by_id = {r.id: r for r in all_runs} - deps_needed = { - dep_id - for r in revised - for dep_id in r.depends_on - if dep_id not in revised_ids and dep_id in all_by_id - } - revised = [all_by_id[dep_id] for dep_id in deps_needed] + revised - revised_ids = {r.id for r in revised} - rest = [r for r in all_runs if r.id not in revised_ids] - revised = apply_revision(revised, args.revision, args.reason) - all_runs = rest + revised - else: - all_runs = apply_revision(all_runs, args.revision, args.reason) - all_runs = apply_wandb_target(all_runs, args.project, args.entity) - all_runs = apply_launch_commit(all_runs, args.launch_commit) + revised = _filter_requested_runs(all_runs, args.experiment_class) + revised_ids = {run.id for run in revised} + rest = [run for run in all_runs if run.id not in revised_ids] + revised = apply_revision(revised, args.revision, args.reason) + all_runs = rest + revised + 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 = [r for r in all_runs if sprint_filter is None or r.sprint in sprint_filter] + candidate_runs = [ + run for run in all_runs if class_filter is None or run.experiment_class in class_filter + ] runs_to_retry = [] - for r in candidate_runs: - status = get_run_status(state, r.id) + for run in candidate_runs: + status = get_run_status(state, run.id) if args.failed and status == "failed": - runs_to_retry.append(r) + runs_to_retry.append(run) elif args.skipped and status == "skipped": - runs_to_retry.append(r) + runs_to_retry.append(run) if not runs_to_retry: print("No runs to retry.") return 0 - all_by_id = {r.id: r for r in all_runs} + all_by_id = {run.id: run for run in all_runs} dependencies, required_by, missing = _collect_dependency_closure(runs_to_retry, all_by_id) - - # Failed/skipped dependencies cannot satisfy downstream runs. Treat --failed - # and --skipped as the explicit request to reset those dependency states too. blocked = [ dep for dep in dependencies @@ -1684,77 +1444,39 @@ def cmd_retry(args): 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, - ) + _fail_for_blocked_retry_dependencies(blocked, unresolved_missing, state, required_by, args) - runs_by_retry_id: dict[str, Run] = {} + runs_by_id: dict[str, Run] = {} for run in dependencies + runs_to_retry: - runs_by_retry_id.setdefault(run.id, run) - - for run in runs_by_retry_id.values(): + runs_by_id.setdefault(run.id, run) + for run in runs_by_id.values(): if get_run_status(state, run.id) in {"failed", "skipped"}: set_run_status(state, run.id, "pending") - runs_to_retry = list(runs_by_retry_id.values()) - save_state(state) - print(f"Retrying {len(runs_to_retry)} runs") + 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 args.launch_commit: - print(f"Launch commit: {args.launch_commit}") - return run_scheduler(runs_to_retry, state, args.parallel, dry_run=False) + 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): - """Pre-build raw tensor caches for requested sprints. - - Creates one raw tensor cache per dataset (not per task/seed/fraction). - The cache stores unnormalized tensors; normalization is applied at runtime. - This means only ~0.8GB total cache instead of 140GB+. - """ if args.revision: print(f"Note: --revision={args.revision} ignored (warmup is revision-independent)") - all_runs = generate_all_runs() + runs = _filter_requested_runs(generate_all_runs(), args.experiment_class) + datasets = dict.fromkeys(run.dataset for run in runs) - # Filter to requested sprints - sprints = [str(s) for s in args.sprint] - runs = filter_runs_for_sprints(all_runs, sprints) - - # Also include dependency runs - run_ids = {r.id for r in runs} - deps_needed = set() - for r in runs: - for d in r.depends_on: - if d not in run_ids: - deps_needed.add(d) - if deps_needed: - all_by_id = {r.id: r for r in all_runs} - runs = [all_by_id[d] for d in deps_needed if d in all_by_id] + runs - - # Collect unique datasets — raw tensor cache is per-dataset only - datasets: dict[str, None] = {} # ordered set via dict - for r in runs: - datasets[r.dataset] = None - - print(f"Warming up tensor caches for sprint(s) {', '.join(sprints)}") + 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 i, dataset in enumerate(datasets, 1): + for index, dataset in enumerate(datasets, 1): processed_dir = f"data/processed/{dataset}" - print(f"[{i}/{len(datasets)}] {dataset}") - - # Run in a subprocess so that memory is fully returned to the OS. - # Uses task_name=None (unsupervised) to load ALL stays without - # any task-specific filtering — this populates the raw cache. + print(f"[{index}/{len(datasets)}] {dataset}") result = subprocess.run( [ "uv", @@ -1762,12 +1484,8 @@ def cmd_warmup(args): "python", "-c", "from slices.data.dataset import ICUDataset; " - f"ds = ICUDataset(" - f" data_dir={processed_dir!r}," - f" task_name=None," - f" normalize=False," - f"); " - f"print(len(ds))", + f"ds = ICUDataset(data_dir={processed_dir!r}, task_name=None, normalize=False); " + "print(len(ds))", ], capture_output=True, text=True, @@ -1783,102 +1501,16 @@ def cmd_warmup(args): 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/") - print("You can now run experiments in parallel without OOM.") return 0 -def cmd_tag(args): - """Add sprint tags to inherited baseline runs in W&B. - - For each requested sprint, finds runs from inherited baseline sprints - (per BASELINE_SPRINTS mapping) and adds the target sprint tag via the - W&B API. Idempotent — already-tagged runs are skipped. - - Usage: - uv run python scripts/internal/run_experiments.py tag --sprint 2 \\ - --project slices-thesis --entity - uv run python scripts/internal/run_experiments.py tag --sprint 2 3 --dry-run \\ - --project slices-thesis --entity - uv run python scripts/internal/run_experiments.py tag --sprint 2 \\ - --project slices-thesis --entity - """ - try: - import wandb # noqa: F811 - except ImportError: - print("Error: wandb is required for tagging. Install with: pip install wandb") - sys.exit(1) - - api = wandb.Api() - - project = args.project - entity = args.entity - # Resolve entity: CLI flag > WANDB_ENTITY env var > wandb default entity - if entity is None: - entity = os.environ.get("WANDB_ENTITY") or api.default_entity - if entity is None: - print("Error: Could not determine W&B entity. Set --entity or WANDB_ENTITY env var.") - sys.exit(1) - - sprints = [str(s) for s in args.sprint] - dry_run = args.dry_run - revision_tag = f"revision:{args.revision}" if args.revision else None - - for target_sprint in sprints: - source_sprints = BASELINE_SPRINTS.get(target_sprint, []) - if not source_sprints: - print(f"Sprint {target_sprint}: no baseline sprints to inherit from — skipping.") - continue - - target_tag = f"sprint:{target_sprint}" - sources = ", ".join(source_sprints) - print(f"\nSprint {target_sprint}: inheriting baselines from sprint(s) {sources}") - if revision_tag: - print(f" Filtering source runs to {revision_tag}") - - # Query runs from each source sprint - tagged = 0 - skipped = 0 - skipped_revision = 0 - for source_sprint in source_sprints: - source_tag = f"sprint:{source_sprint}" - runs = api.runs( - f"{entity}/{project}", - filters={"tags": {"$in": [source_tag]}}, - ) - - for run in runs: - if revision_tag and revision_tag not in run.tags: - skipped_revision += 1 - continue - if target_tag in run.tags: - skipped += 1 - continue - - if dry_run: - print(f" [dry-run] Would tag: {run.name} ({run.id}) with {target_tag}") - tagged += 1 - else: - run.tags = run.tags + [target_tag] - run.update() - tagged += 1 - - action = "would tag" if dry_run else "tagged" - print(f" {action} {tagged} runs, skipped {skipped} (already tagged)") - if skipped_revision: - print(f" skipped {skipped_revision} source runs outside {revision_tag}") - - print("\nDone.") - - -def main(): - parser = argparse.ArgumentParser(description="SLICES parallel experiment runner") +def main() -> None: + parser = argparse.ArgumentParser(description="SLICES class-based experiment runner") sub = parser.add_subparsers(dest="command", required=True) - # run - p_run = sub.add_parser("run", help="Run experiments for given sprints") - p_run.add_argument("--sprint", nargs="+", required=True, help="Sprint(s) to run (e.g. 1 1b 2)") + 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, @@ -1886,7 +1518,7 @@ def main(): 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 (e.g. v2)") + 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", @@ -1907,11 +1539,9 @@ def main(): help="W&B entity override for launched runs (default: WANDB_ENTITY env var)", ) - # status p_status = sub.add_parser("status", help="Show experiment status") - p_status.add_argument("--sprint", nargs="*", default=None, help="Filter by sprint(s)") + p_status.add_argument("--experiment-class", nargs="*", choices=EXPERIMENT_CLASSES, default=None) - # retry 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 skipped runs") @@ -1921,7 +1551,7 @@ def main(): default=4, help="Max scheduler slots (default: 4). Pretrains consume 4 slots, other runs 1.", ) - p_retry.add_argument("--sprint", nargs="+", default=None, help="Scope retry to sprint(s)") + 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( @@ -1943,55 +1573,23 @@ def main(): help="W&B entity override for relaunched runs (default: WANDB_ENTITY env var)", ) - # warmup p_warmup = sub.add_parser( "warmup", help="Pre-build tensor caches to avoid OOM during parallel runs" ) p_warmup.add_argument( - "--sprint", nargs="+", required=True, help="Sprint(s) to warmup (e.g. 1 1b 2)" + "--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") - # tag - p_tag = sub.add_parser("tag", help="Tag inherited baseline runs in W&B with target sprint tags") - p_tag.add_argument( - "--sprint", - nargs="+", - required=True, - help="Sprint(s) whose baselines to tag (e.g. 2 or 2 3 5)", - ) - p_tag.add_argument("--dry-run", action="store_true", help="Show what would be tagged") - p_tag.add_argument( - "--revision", - type=str, - default=os.environ.get("REVISION") or os.environ.get("WANDB_REVISION"), - help="Only inherit source runs with this revision tag (default: REVISION/WANDB_REVISION)", - ) - p_tag.add_argument( - "--project", - type=str, - default=os.environ.get("WANDB_PROJECT", "slices"), - help="W&B project name (default: WANDB_PROJECT env var or 'slices')", - ) - p_tag.add_argument( - "--entity", - type=str, - default=None, - help="W&B entity (default: WANDB_ENTITY env var or wandb default)", - ) - args = parser.parse_args() - # Validate --reason requires --revision if getattr(args, "reason", None) and not getattr(args, "revision", None): parser.error("--reason requires --revision") - - # Validate retry --revision requires --sprint - if args.command == "retry" and args.revision and not args.sprint: - parser.error("--revision requires --sprint to scope which sprints to revise") + if args.command == "retry" and args.revision and not args.experiment_class: + parser.error("--revision requires --experiment-class to scope which classes to revise") if args.command == "run": exit_code = cmd_run(args) @@ -2001,11 +1599,8 @@ def main(): exit_code = cmd_retry(args) elif args.command == "warmup": exit_code = cmd_warmup(args) - elif args.command == "tag": - exit_code = cmd_tag(args) else: exit_code = 0 - sys.exit(exit_code or 0) diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index cb4caf7..28630b6 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -35,8 +35,10 @@ def _xgboost_eval_metric(task_type: str) -> str: return "mae" if task_type == "regression" else "aucpr" -def _add_wandb_tag(tags: list[str], tag: str) -> None: - """Append a W&B tag once, preserving caller order.""" +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) @@ -83,8 +85,10 @@ def _binary_ece(y_true, y_pred_proba, n_bins: int = 15) -> float: 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("sprint") is not None: - _add_wandb_tag(tags, 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: @@ -92,6 +96,19 @@ def _build_wandb_tags(cfg: DictConfig) -> list[str] | None: if len(tag) > 64: tag = tag[:61] + "..." _add_wandb_tag(tags, tag) + 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:{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") diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 07c0630..29bf374 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -264,8 +264,10 @@ def setup_finetune_callbacks(cfg: DictConfig, checkpoint_prefix: str = "finetune # ============================================================================= -def _add_wandb_tag(tags: list[str], tag: str) -> None: - """Append a W&B tag once, preserving caller order.""" +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) @@ -279,8 +281,10 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: return None tags = list(cfg.logging.get("wandb_tags", [])) - if cfg.get("sprint") is not None: - _add_wandb_tag(tags, 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: @@ -290,6 +294,19 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: _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}") @@ -297,7 +314,9 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: # Add downstream protocol family tag. Supervised-from-scratch shares the # Protocol B optimization budget; the phase tag distinguishes it from SSL. freeze_encoder = cfg.get("training", {}).get("freeze_encoder") - if freeze_encoder is not None: + if cfg.get("protocol") is not None: + _add_wandb_tag(tags, f"protocol:{cfg.protocol}") + elif freeze_encoder is not None: protocol = "A" if freeze_encoder else "B" _add_wandb_tag(tags, f"protocol:{protocol}") diff --git a/tests/test_config_schemas.py b/tests/test_config_schemas.py index bebf302..c225977 100644 --- a/tests/test_config_schemas.py +++ b/tests/test_config_schemas.py @@ -9,7 +9,6 @@ import pytest import yaml from pydantic import ValidationError - from slices.data.config_schemas import DataConfig from slices.training.config_schemas import ( OptimizerConfig, @@ -247,3 +246,32 @@ def test_training_yaml_has_allow_best_ckpt_fallback(self, config_name, expected) 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_evaluate_fairness.py b/tests/test_evaluate_fairness.py index 8fbed9d..7488f6a 100644 --- a/tests/test_evaluate_fairness.py +++ b/tests/test_evaluate_fairness.py @@ -38,16 +38,18 @@ def test_default_phases_include_baseline(): assert "baseline" in mod.DEFAULT_PHASES -def test_default_core_sprints_include_sprint11(): +def test_default_experiment_classes_include_classical_baselines(): mod = importlib.import_module("scripts.eval.evaluate_fairness") - assert "11" in mod.CORE_SPRINTS + assert "classical_baselines" in mod.DEFAULT_EXPERIMENT_CLASSES -def test_default_core_sprints_include_ablation_and_transfer_sprints(): +def test_default_experiment_classes_include_downstream_families(): mod = importlib.import_module("scripts.eval.evaluate_fairness") - assert {"6", "7", "8"}.issubset(set(mod.CORE_SPRINTS)) + assert {"label_efficiency", "cross_dataset_transfer", "hp_robustness"}.issubset( + set(mod.DEFAULT_EXPERIMENT_CLASSES) + ) def test_parse_args_exposes_allow_incomplete_escape_hatch(monkeypatch): diff --git a/tests/test_export_results.py b/tests/test_export_results.py index 1147fb6..2fc480c 100644 --- a/tests/test_export_results.py +++ b/tests/test_export_results.py @@ -1,4 +1,4 @@ -"""Tests for the results export pipeline.""" +"""Tests for the class-based results export pipeline.""" import importlib import sys @@ -19,35 +19,38 @@ def __init__(self, config, tags, summary=None, name="dummy"): 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": "B", + "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 test_build_statistical_tests_df_produces_pairwise_significance_rows(): mod = importlib.import_module("scripts.export_results") rows = [] - for paradigm, phase, offset in [ - ("mae", "finetune", 0.10), - ("supervised", "supervised", 0.0), - ]: + 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)]: - rows.append( - { - "experiment_type": "core", - "sprint": "1", - "paradigm": paradigm, - "dataset": "miiv", - "task": task, - "seed": seed, - "protocol": "A", - "label_fraction": 1.0, - "model_size": "default", - "source_dataset": None, - "phase": phase, - "test/auprc": base + offset + (seed / 100000.0), - } - ) - - per_seed_df = pd.DataFrame(rows) - stats_df = mod.build_statistical_tests_df(per_seed_df) + 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)) assert len(stats_df) == 1 row = stats_df.iloc[0] @@ -55,144 +58,117 @@ def test_build_statistical_tests_df_produces_pairwise_significance_rows(): assert row["paradigm_a"] == "mae" assert row["paradigm_b"] == "supervised" assert row["n_pairs"] == 6 - assert row["n_tasks"] == 2 assert row["better_paradigm"] == "mae" - assert row["p_value_bonferroni"] >= row["p_value"] assert row["n_shared_task_seed_pairs"] == 6 - assert row["n_union_task_seed_pairs"] == 6 -def test_extract_run_uses_requested_inherited_sprint_tag_for_family(): +def test_extract_run_requires_experiment_class(): mod = importlib.import_module("scripts.export_results") run = DummyRun( config={ - "sprint": "6", "dataset": "miiv", "paradigm": "mae", "seed": 42, - "encoder": {"d_model": 64, "n_layers": 2}, "training": {"freeze_encoder": False}, "task": {"task_name": "mortality_24h"}, - "label_fraction": 0.1, }, - tags=["sprint:6", "sprint:7p", "phase:finetune"], - name="s6_inherited_capacity_baseline", + tags=["phase:finetune"], ) - default_row = mod.extract_run(run, []) - requested_row = mod.extract_run(run, [], requested_sprints=["7p"]) + with pytest.raises(RuntimeError, match="no experiment_class"): + mod.extract_run(run, []) - assert default_row["experiment_type"] == "label_efficiency" - assert requested_row["experiment_type"] == "capacity_pilot" - assert requested_row["sprint"] == "7p" - assert requested_row["config_sprint"] == "6" - -def test_build_per_seed_df_adds_revision_wide_capacity_membership_for_inherited_row(): +def test_extract_run_uses_class_metadata_from_config_or_tags(): mod = importlib.import_module("scripts.export_results") - run = DummyRun( + config_run = DummyRun( config={ - "sprint": "6", + "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": "mortality_24h"}, - "label_fraction": 0.1, + "encoder": {"d_model": 64, "n_layers": 2}, }, - tags=["sprint:6", "sprint:7p", "phase:finetune"], - name="s6_inherited_capacity_baseline", + tags=["phase:finetune", "revision:thesis-v1"], ) - - df = mod.build_per_seed_df([run]) - - assert set(df["experiment_type"]) == {"label_efficiency", "capacity_pilot"} - capacity = df[df["experiment_type"] == "capacity_pilot"].iloc[0] - assert capacity["sprint"] == "7p" - assert capacity["config_sprint"] == "6" - assert capacity["wandb_run_id"] == "dummy-id" - - scoped_df = mod.build_per_seed_df([run], requested_sprints=["6", "7p"]) - assert set(scoped_df["experiment_type"]) == {"label_efficiency", "capacity_pilot"} - - -def test_build_per_seed_df_does_not_add_capacity_membership_for_irrelevant_tagged_rows(): - mod = importlib.import_module("scripts.export_results") - - run = DummyRun( + tag_run = DummyRun( config={ - "sprint": "6", - "dataset": "eicu", - "paradigm": "jepa", + "dataset": "miiv", + "paradigm": "xgboost", "seed": 42, - "encoder": {"d_model": 64, "n_layers": 2}, - "training": {"freeze_encoder": False}, - "task": {"task_name": "mortality_hospital"}, - "label_fraction": 0.05, + "task": {"task_name": "mortality_24h"}, }, - tags=["sprint:6", "sprint:7p", "phase:finetune"], - name="s6_coarsely_tagged_but_not_capacity_scope", + tags=["experiment_class:classical_baselines", "phase:baseline"], ) - df = mod.build_per_seed_df([run]) - - assert df["experiment_type"].tolist() == ["label_efficiency"] + 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"] == "B" -def test_build_per_seed_df_adds_unscoped_inherited_memberships_for_thesis_families(): +def test_fetch_all_runs_single_class_adds_server_side_filter(monkeypatch): mod = importlib.import_module("scripts.export_results") + captured = {} + + 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 [] + + monkeypatch.setattr(mod.wandb, "Api", DummyApi) + + mod.fetch_all_runs( + project="proj", + entity="entity", + experiment_class=["core_ssl_benchmark"], + paradigm=["mae"], + dataset=["miiv"], + phase=["finetune"], + revision=["thesis-v1"], + ) - base_config = { - "sprint": "1", - "dataset": "miiv", - "paradigm": "mae", - "seed": 42, - "encoder": {"d_model": 64, "n_layers": 2}, - "training": {"freeze_encoder": False}, - "task": {"task_name": "mortality_24h"}, - "label_fraction": 1.0, - } - runs = [ - DummyRun(base_config, ["sprint:1", "sprint:6", "phase:finetune"], name="core_for_s6"), - DummyRun(base_config, ["sprint:1", "sprint:7", "phase:finetune"], name="core_for_s7"), - DummyRun(base_config, ["sprint:1", "sprint:8", "phase:finetune"], name="core_for_s8"), + assert captured["path"] == "entity/proj" + assert captured["filters"]["tags"]["$all"] == [ + "experiment_class:core_ssl_benchmark", + "paradigm:mae", + "dataset:miiv", + "phase:finetune", + "revision:thesis-v1", ] - for i, run in enumerate(runs): - run.id = f"run-{i}" - run.url = f"https://wandb.test/run-{i}" - run.created_at = f"2026-04-21T00:00:0{i}" - df = mod.build_per_seed_df(runs) - memberships = set(zip(df["wandb_run_id"], df["sprint"], df["experiment_type"], strict=True)) - assert ("run-0", "6", "label_efficiency") in memberships - assert ("run-1", "7", "transfer") in memberships - assert ("run-2", "8", "hp_ablation") in memberships - - -def test_build_per_seed_df_rejects_hp_ablation_from_label_and_transfer_inheritance(): +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={ - "sprint": "10", + "experiment_class": "hp_robustness", + "experiment_subtype": "lr_sensitivity", "dataset": "miiv", "paradigm": "mae", "seed": 42, "output_dir": ( - "outputs/sprint10/" f"finetune_mae_mortality_24h_miiv_seed42_lr{lr_code}" + "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=["sprint:10", "sprint:6", "sprint:7", "phase:finetune"], + tags=["experiment_class:hp_robustness", "phase:finetune"], name=f"finetune_mae_mortality_24h_miiv_seed42_lr{lr_code}", ) run.id = run_id @@ -202,126 +178,72 @@ def test_build_per_seed_df_rejects_hp_ablation_from_label_and_transfer_inheritan df = mod.build_per_seed_df(runs) assert len(df) == 2 - assert set(df["experiment_type"]) == {"hp_ablation"} - assert set(df["sprint"]) == {"10"} + 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") - bad_run = object() - with pytest.raises(RuntimeError, match="failed extraction"): - mod.build_per_seed_df([bad_run]) + mod.build_per_seed_df([object()]) + + assert mod.build_per_seed_df([object()], allow_extraction_failures=True).empty + + +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}, + ] - df = mod.build_per_seed_df([bad_run], allow_extraction_failures=True) - assert df.empty + 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_type, paradigm, phase, offset in [ - ("core", "mae", "finetune", 0.10), - ("core", "supervised", "supervised", 0.0), - ("classical_baselines", "xgboost", "baseline", 0.03), - ("classical_baselines", "gru_d", "baseline", 0.05), + 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, base in [("mortality_24h", 0.30), ("aki_kdigo", 0.25)]: - rows.append( - { - "experiment_type": experiment_type, - "sprint": "11" if experiment_type == "classical_baselines" else "1", - "paradigm": paradigm, - "dataset": "miiv", - "task": task, - "seed": seed, - "protocol": "B", - "label_fraction": 1.0, - "model_size": "default", - "source_dataset": None, - "phase": phase, - "test/auprc": base + offset + (seed / 100000.0), - } - ) + 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_type"] == "classical_context_full"] + 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 ("supervised", "xgboost") in pairs assert ("mae", "xgboost") in pairs assert ("gru_d", "supervised") in pairs - assert ("gru_d", "mae") in pairs assert ("gru_d", "xgboost") not in pairs - assert ("mae", "supervised") not in pairs -def test_build_statistical_tests_df_reports_incomplete_pair_coverage(): +def test_build_ts2vec_vs_core_contrastive_df_compares_across_classes(): mod = importlib.import_module("scripts.export_results") rows = [] - for paradigm, seeds in [("mae", [42, 123]), ("supervised", [42])]: - for seed in seeds: - rows.append( - { - "experiment_type": "core", - "sprint": "1", - "paradigm": paradigm, - "dataset": "miiv", - "task": "mortality_24h", - "seed": seed, - "protocol": "B", - "label_fraction": 1.0, - "model_size": "default", - "source_dataset": None, - "phase": "finetune", - "test/auprc": 0.3 + (0.1 if paradigm == "mae" else 0.0), - } - ) - - stats_df = mod.build_statistical_tests_df(pd.DataFrame(rows)) - row = stats_df.iloc[0] - - assert row["n_pairs"] == 1 - assert row["n_task_seed_pairs_a"] == 2 - assert row["n_task_seed_pairs_b"] == 1 - assert row["n_shared_task_seed_pairs"] == 1 - assert row["n_union_task_seed_pairs"] == 2 - assert '"seed": 123' in row["missing_task_seed_pairs_b"] - - -def test_build_ts2vec_vs_core_contrastive_df_compares_across_experiment_types(): - mod = importlib.import_module("scripts.export_results") - - rows = [] - for experiment_type, paradigm, offset in [ - ("temporal_contrastive", "ts2vec", 0.04), - ("core", "contrastive", 0.0), + for experiment_class, paradigm, offset in [ + ("ts2vec_extension", "ts2vec", 0.04), + ("core_ssl_benchmark", "contrastive", 0.0), ]: for seed in [42, 123]: - for task, base in [("mortality_24h", 0.30), ("aki_kdigo", 0.25)]: - rows.append( - { - "experiment_type": experiment_type, - "sprint": "13" if paradigm == "ts2vec" else "1", - "paradigm": paradigm, - "dataset": "miiv", - "task": task, - "seed": seed, - "protocol": "B", - "label_fraction": 1.0, - "model_size": "default", - "source_dataset": None, - "phase": "finetune", - "test/auprc": base + offset + (seed / 100000.0), - } - ) + 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)) @@ -331,7 +253,6 @@ def test_build_ts2vec_vs_core_contrastive_df_compares_across_experiment_types(): assert row["paradigm_a"] == "ts2vec" assert row["paradigm_b"] == "contrastive" assert row["n_pairs"] == 4 - assert row["better_paradigm"] == "ts2vec" def test_validate_flags_wrong_fixed_seed_set_with_correct_count(): @@ -340,7 +261,7 @@ def test_validate_flags_wrong_fixed_seed_set_with_correct_count(): aggregated_df = pd.DataFrame( [ { - "experiment_type": "core", + "experiment_class": "core_ssl_benchmark", "paradigm": "mae", "dataset": "miiv", "task": "mortality_24h", @@ -350,9 +271,7 @@ def test_validate_flags_wrong_fixed_seed_set_with_correct_count(): } ] ) - per_seed_df = pd.DataFrame([{"test/auprc": 0.5}]) - - warnings = mod.validate(per_seed_df, aggregated_df) + 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) @@ -361,44 +280,17 @@ def test_validate_flags_wrong_fixed_seed_set_with_correct_count(): def test_build_aggregated_df_records_per_metric_non_null_counts(): mod = importlib.import_module("scripts.export_results") - per_seed_df = pd.DataFrame( - [ - { - "experiment_type": "core", - "sprint": "1", - "paradigm": "mae", - "dataset": "miiv", - "task": "mortality_24h", - "seed": 42, - "protocol": "B", - "label_fraction": 1.0, - "model_size": "default", - "source_dataset": None, - "phase": "finetune", - "test/auprc": 0.4, - "test/auroc": 0.7, - "wandb_run_id": "run-42", - }, - { - "experiment_type": "core", - "sprint": "1", - "paradigm": "mae", - "dataset": "miiv", - "task": "mortality_24h", - "seed": 123, - "protocol": "B", - "label_fraction": 1.0, - "model_size": "default", - "source_dataset": None, - "phase": "finetune", - "test/auprc": None, - "test/auroc": 0.8, - "wandb_run_id": "run-123", - }, - ] - ) + 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(per_seed_df) + aggregated = mod.build_aggregated_df(pd.DataFrame(rows)) assert aggregated.iloc[0]["n_seeds"] == 2 assert aggregated.iloc[0]["test/auprc/n"] == 1 @@ -408,22 +300,22 @@ def test_build_aggregated_df_records_per_metric_non_null_counts(): def test_validate_warns_when_evaluation_row_missing_primary_metric(): mod = importlib.import_module("scripts.export_results") - per_seed_df = pd.DataFrame( - [ - { - "wandb_run_id": "run-1", - "paradigm": "mae", - "dataset": "miiv", - "task": "mortality_24h", - "seed": 42, - "phase": "finetune", - "test/auprc": None, - } - ] + 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(), ) - aggregated_df = pd.DataFrame() - - warnings = mod.validate(per_seed_df, aggregated_df) assert any("missing their primary test metric" in warning for warning in warnings) @@ -433,7 +325,7 @@ def test_extract_run_exports_los_regression_fairness_keys(): run = DummyRun( config={ - "sprint": "1", + "experiment_class": "core_ssl_benchmark", "dataset": "miiv", "paradigm": "mae", "seed": 42, @@ -441,7 +333,7 @@ def test_extract_run_exports_los_regression_fairness_keys(): "training": {"freeze_encoder": False}, "task": {"task_name": "los_remaining", "task_type": "regression"}, }, - tags=["sprint:1", "phase:finetune"], + tags=["experiment_class:core_ssl_benchmark", "phase:finetune"], summary={ "test/mae": 1.4, "fairness/gender/worst_group_mse": 2.5, @@ -578,9 +470,7 @@ def test_parse_args_exposes_allow_incomplete_escape_hatch(monkeypatch): ["export_results.py", "--revision", "thesis-v1", "--allow-incomplete"], ) - args = mod.parse_args() - - assert args.allow_incomplete is True + assert mod.parse_args().allow_incomplete is True def test_parse_args_exposes_extraction_failure_escape_hatch(monkeypatch): @@ -592,9 +482,7 @@ def test_parse_args_exposes_extraction_failure_escape_hatch(monkeypatch): ["export_results.py", "--revision", "thesis-v1", "--allow-extraction-failures"], ) - args = mod.parse_args() - - assert args.allow_extraction_failures is True + assert mod.parse_args().allow_extraction_failures is True def test_main_exits_nonzero_when_no_runs_match(monkeypatch, tmp_path): @@ -620,9 +508,7 @@ def test_parse_args_requires_revision_without_cli_or_env(monkeypatch): monkeypatch.delenv("WANDB_REVISION", raising=False) monkeypatch.setattr(sys, "argv", ["export_results.py"]) - try: + with pytest.raises(SystemExit) as excinfo: mod.parse_args() - except SystemExit as exc: - assert exc.code == 2 - else: - raise AssertionError("parse_args() should require an explicit revision scope") + + assert excinfo.value.code == 2 diff --git a/tests/test_fixes.py b/tests/test_fixes.py index e5b3572..bbf3927 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -729,13 +729,21 @@ def test_parse_args_requires_revision_when_no_scope_is_available(self, monkeypat assert excinfo.value.code == 2 - def test_default_core_sprints_include_all_downstream_training_sprints(self): - """The standalone thesis fairness corpus should include every downstream sprint.""" + 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") - assert set(["1", "2", "3", "4", "5", "6", "7", "8", "7p", "10", "11", "12", "13"]).issubset( - set(mod.CORE_SPRINTS) - ) + 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_adds_server_side_filter(self, monkeypatch): """A single revision should be sent as a W&B tag filter.""" @@ -757,7 +765,7 @@ def runs(self, path, filters, order): mod.fetch_eval_runs( project="proj", entity="entity", - sprints=["1"], + experiment_classes=["core_ssl_benchmark"], paradigms=["mae"], datasets=["miiv"], phases=["finetune"], @@ -766,7 +774,7 @@ def runs(self, path, filters, order): assert captured["path"] == "entity/proj" assert captured["filters"]["tags"]["$all"] == [ - "sprint:1", + "experiment_class:core_ssl_benchmark", "paradigm:mae", "dataset:miiv", "phase:finetune", @@ -796,7 +804,7 @@ def runs(self, path, filters, order): filtered = mod.fetch_eval_runs( project="proj", entity=None, - sprints=None, + experiment_classes=None, paradigms=None, datasets=None, phases=["finetune", "supervised"], @@ -910,6 +918,9 @@ def test_tmux_launcher_passes_revision_to_fairness(self, tmp_path): 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"], @@ -929,8 +940,8 @@ def test_tmux_launcher_passes_revision_to_fairness(self, tmp_path): assert "--project slices-thesis" in runner_text assert "--entity test-entity" in runner_text - def test_tmux_launcher_includes_sprint11_in_fairness(self, tmp_path): - """The fairness sweep should include the canonical baseline sprint by default.""" + 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) @@ -959,6 +970,9 @@ def test_tmux_launcher_includes_sprint11_in_fairness(self, tmp_path): 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"], @@ -973,7 +987,9 @@ def test_tmux_launcher_includes_sprint11_in_fairness(self, tmp_path): assert len(runner_scripts) == 1 runner_text = runner_scripts[0].read_text() - assert "--sprint 1 2 3 4 5 6 7 8 10 11" in runner_text + assert "--experiment-class" in runner_text + assert "classical_baselines" in runner_text + assert "core_ssl_benchmark" in runner_text def test_tmux_status_pane_uses_uv_run_python(self, tmp_path): """The status pane should use uv-managed Python.""" @@ -1005,6 +1021,9 @@ def test_tmux_status_pane_uses_uv_run_python(self, tmp_path): 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"], @@ -1338,8 +1357,8 @@ def test_legacy_schema_migration(self): # ============================================================================ -class TestExporterHistoricalRecovery: - """Tests for _recover_pretrain_metadata run-name parsing.""" +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.""" @@ -1365,64 +1384,64 @@ def test_mr_decode_matches_run_experiments_encoding(self): assert _MR_DECODE["075"] == 0.75 def test_recover_lr_from_output_dir(self): - """Historical LR ablation finetunes recovered from output_dir pattern.""" + """LR sensitivity finetunes recover upstream LR from output_dir pattern.""" from scripts.export_results import _recover_pretrain_metadata - config = {"output_dir": "outputs/sprint10/finetune_mae_mortality_24h_miiv_seed789_lr00002"} + 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_ablation" + assert subtype == "lr_sensitivity" def test_recover_mask_ratio_from_output_dir(self): - """Historical mask_ratio ablation finetunes recovered from output_dir.""" + """Mask-ratio sensitivity finetunes recover upstream mask ratio from output_dir.""" from scripts.export_results import _recover_pretrain_metadata config = { - "output_dir": "outputs/sprint1c/finetune_jepa_mortality_24h_miiv_seed42_mask_ratio03" + "output_dir": ( + "outputs/hp_robustness/" "finetune_jepa_mortality_24h_miiv_seed42_mask_ratio03" + ) } up_lr, up_mr, subtype = _recover_pretrain_metadata("some_run", config) assert up_mr == 0.3 - assert subtype == "mask_ablation" + 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_ablation"} + 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_ablation" + assert subtype == "lr_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/sprint1/finetune_mae_mortality_24h_miiv_seed42"} + 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 Sprint 7p medium/large labels.""" + """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({"sprint": "7p", "encoder": {"d_model": 128, "n_layers": 4}}) == ( - "medium" - ) - assert _infer_model_size({"sprint": "7p", "encoder": {"d_model": 256, "n_layers": 4}}) == ( - "large" - ) + assert _infer_model_size({"encoder": {"d_model": 128, "n_layers": 4}}) == "medium" + assert _infer_model_size({"encoder": {"d_model": 256, "n_layers": 4}}) == "large" - def test_extract_run_reclassifies_transfer_and_sprint13(self): - """Sprint-10 transfer seeds and Sprint 13 should export with stable experiment types.""" + 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={ - "sprint": "10", + "experiment_class": "cross_dataset_transfer", "dataset": "eicu", "paradigm": "mae", "source_dataset": "miiv", @@ -1434,12 +1453,12 @@ def test_extract_run_reclassifies_transfer_and_sprint13(self): tags=["phase:finetune"], name="finetune_mae_mortality_24h_eicu_from_miiv_seed789", ) - assert extract_run(transfer, [])["experiment_type"] == "transfer" + assert extract_run(transfer, [])["experiment_class"] == "cross_dataset_transfer" ts2vec = DummyWandbRun( run_id="ts2vec", config={ - "sprint": "13", + "experiment_class": "ts2vec_extension", "dataset": "miiv", "paradigm": "ts2vec", "seed": 42, @@ -1450,20 +1469,23 @@ def test_extract_run_reclassifies_transfer_and_sprint13(self): tags=["phase:finetune"], name="finetune_ts2vec_mortality_24h_miiv_seed42", ) - assert extract_run(ts2vec, [])["experiment_type"] == "temporal_contrastive" + assert extract_run(ts2vec, [])["experiment_class"] == "ts2vec_extension" def test_extract_run_recovers_transfer_from_output_dir(self): - """Historical transfer runs should recover source_dataset from output_dir.""" + """Transfer runs should recover source_dataset from output_dir.""" from scripts.export_results import extract_run transfer = DummyWandbRun( run_id="transfer_output_dir", config={ - "sprint": "10", + "experiment_class": "cross_dataset_transfer", "dataset": "eicu", "paradigm": "mae", "seed": 789, - "output_dir": "outputs/sprint10/finetune_mae_mortality_24h_eicu_seed789_from_miiv", + "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"}, @@ -1474,7 +1496,7 @@ def test_extract_run_recovers_transfer_from_output_dir(self): row = extract_run(transfer, []) assert row["source_dataset"] == "miiv" - assert row["experiment_type"] == "transfer" + assert row["experiment_class"] == "cross_dataset_transfer" def test_extract_run_assigns_protocol_b_to_xgboost_baselines(self): """XGBoost baselines must align with the Protocol-B comparison family.""" @@ -1483,7 +1505,7 @@ def test_extract_run_assigns_protocol_b_to_xgboost_baselines(self): run = DummyWandbRun( run_id="xgb", config={ - "sprint": "11", + "experiment_class": "classical_baselines", "dataset": "miiv", "paradigm": "xgboost", "seed": 42, @@ -1502,7 +1524,7 @@ def test_xgboost_wandb_tags_include_revision_scope(self): cfg = OmegaConf.create( { "logging": {"wandb_tags": ["phase:baseline"]}, - "sprint": "11", + "experiment_class": "classical_baselines", "revision": "thesis-v1", "rerun_reason": "rerun canonical thesis baseline sweep with fixed tags", "label_fraction": 0.1, @@ -1512,7 +1534,7 @@ def test_xgboost_wandb_tags_include_revision_scope(self): tags = _build_wandb_tags(cfg) assert "phase:baseline" in tags - assert "sprint:11" 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 @@ -1526,11 +1548,13 @@ def test_build_per_seed_df_keeps_distinct_hp_ablation_configs(self): DummyWandbRun( run_id="hp_lr_2e4", config={ - "sprint": "10", + "experiment_class": "hp_robustness", "dataset": "miiv", "paradigm": "mae", "seed": 42, - "output_dir": "outputs/sprint10/finetune_mae_mortality_24h_miiv_seed42_lr00002", + "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"}, @@ -1542,11 +1566,13 @@ def test_build_per_seed_df_keeps_distinct_hp_ablation_configs(self): DummyWandbRun( run_id="hp_lr_5e4", config={ - "sprint": "10", + "experiment_class": "hp_robustness", "dataset": "miiv", "paradigm": "mae", "seed": 42, - "output_dir": "outputs/sprint10/finetune_mae_mortality_24h_miiv_seed42_lr00005", + "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"}, @@ -1560,20 +1586,20 @@ def test_build_per_seed_df_keeps_distinct_hp_ablation_configs(self): df = build_per_seed_df(runs) assert len(df) == 2 - assert set(df["experiment_type"]) == {"hp_ablation"} + 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_ablation_across_sprints(self): - """Sprint 10 historical seeds should merge with Sprint 1b/1c HP-ablation families.""" + 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_type": "hp_ablation", - "experiment_subtype": "lr_ablation", - "sprint": sprint, + "experiment_class": "hp_robustness", + "experiment_type": "hp_robustness", + "experiment_subtype": "lr_sensitivity", "paradigm": "mae", "dataset": "miiv", "task": "mortality_24h", @@ -1586,19 +1612,17 @@ def test_build_aggregated_df_merges_hp_ablation_across_sprints(self): "upstream_pretrain_lr": 2e-4, "upstream_pretrain_mask_ratio": None, } - for i, (sprint, seed) in enumerate( - [("1b", 42), ("1b", 123), ("10", 456), ("10", 789), ("10", 1011)] - ) + 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_type"] == "hp_ablation" - assert agg.iloc[0]["experiment_subtype"] == "lr_ablation" + 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]["sprint_list"])) == {"1b", "10"} + 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.""" @@ -1608,7 +1632,7 @@ def test_build_per_seed_df_raises_on_ambiguous_core_collision(self): DummyWandbRun( run_id="core_lr_1e4", config={ - "sprint": "1", + "experiment_class": "core_ssl_benchmark", "dataset": "miiv", "paradigm": "mae", "seed": 42, @@ -1624,7 +1648,7 @@ def test_build_per_seed_df_raises_on_ambiguous_core_collision(self): DummyWandbRun( run_id="core_lr_3e4", config={ - "sprint": "2", + "experiment_class": "core_ssl_benchmark", "dataset": "miiv", "paradigm": "mae", "seed": 42, @@ -1650,13 +1674,13 @@ def test_apply_wandb_target_injects_project_and_entity(self): from scripts.internal.run_experiments import Run, apply_wandb_target run = Run( - id="s1_supervised_mortality_24h_miiv_seed42", - sprint="1", + 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/sprint1/supervised_mortality_24h_miiv_seed42", + output_dir="outputs/core_ssl_benchmark/supervised_mortality_24h_miiv_seed42", ) result = apply_wandb_target([run], project="slices-thesis", entity="hannes-ill") @@ -1664,32 +1688,32 @@ def test_apply_wandb_target_injects_project_and_entity(self): assert result[0].extra_overrides["logging.wandb_project"] == "slices-thesis" assert result[0].extra_overrides["logging.wandb_entity"] == "hannes-ill" - def test_cmd_run_respects_requested_sprint_order(self, monkeypatch): + def test_cmd_run_respects_requested_class_order(self, monkeypatch): import scripts.internal.run_experiments as runner - sprint7p = runner.Run( - id="s7p_supervised_mortality_24h_miiv_seed42", - sprint="7p", + 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/sprint7p/supervised_mortality_24h_miiv_seed42", + output_dir="outputs/capacity_study/supervised_mortality_24h_miiv_seed42", task="mortality_24h", ) - sprint10 = runner.Run( - id="s10_supervised_mortality_24h_miiv_seed789", - sprint="10", + 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/sprint10/supervised_mortality_24h_miiv_seed789", + output_dir="outputs/classical_baselines/supervised_mortality_24h_miiv_seed789", task="mortality_24h", ) scheduled = {} - monkeypatch.setattr(runner, "generate_all_runs", lambda: [sprint7p, sprint10]) + 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( @@ -1699,7 +1723,7 @@ def test_cmd_run_respects_requested_sprint_order(self, monkeypatch): ) args = SimpleNamespace( - sprint=["10", "7p"], + experiment_class=["classical_baselines", "capacity_study"], revision=None, reason=None, project=None, @@ -1710,7 +1734,10 @@ def test_cmd_run_respects_requested_sprint_order(self, monkeypatch): runner.cmd_run(args) - assert [run.sprint for run in scheduled["runs"]] == ["10", "7p"] + 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 @@ -1724,7 +1751,7 @@ def test_pretrain_and_finetune_commands_resume_from_last_checkpoint(self, tmp_pa pretrain = Run( id="pretrain", - sprint="1", + experiment_class="core_ssl_benchmark", run_type="pretrain", paradigm="mae", dataset="miiv", @@ -1733,7 +1760,7 @@ def test_pretrain_and_finetune_commands_resume_from_last_checkpoint(self, tmp_pa ) finetune = Run( id="finetune", - sprint="1", + experiment_class="core_ssl_benchmark", run_type="finetune", paradigm="mae", dataset="miiv", @@ -1773,7 +1800,7 @@ def test_scheduler_exit_code_flags_pending_and_running_runs(self): runs = [ runner.Run( id="pending-run", - sprint="1", + experiment_class="core_ssl_benchmark", run_type="pretrain", paradigm="mae", dataset="miiv", @@ -1782,7 +1809,7 @@ def test_scheduler_exit_code_flags_pending_and_running_runs(self): ), runner.Run( id="running-run", - sprint="1", + experiment_class="core_ssl_benchmark", run_type="pretrain", paradigm="mae", dataset="miiv", @@ -1800,44 +1827,10 @@ def test_scheduler_exit_code_flags_pending_and_running_runs(self): assert runner._scheduler_exit_code(runs, state) == 1 - def test_cmd_tag_filters_inherited_sources_by_revision(self, monkeypatch): + def test_runner_no_longer_exposes_inherited_tag_command(self): import scripts.internal.run_experiments as runner - class FakeRun: - def __init__(self, tags): - self.tags = tags - self.name = "fake" - self.id = "fake" - self.updated = False - - def update(self): - self.updated = True - - matching = FakeRun(["sprint:1", "revision:thesis-v1"]) - stale = FakeRun(["sprint:1", "revision:old"]) - - class FakeApi: - default_entity = "entity" - - def runs(self, path, filters): - return [matching, stale] - - monkeypatch.setitem(sys.modules, "wandb", SimpleNamespace(Api=lambda: FakeApi())) - - args = SimpleNamespace( - sprint=["2"], - dry_run=False, - project="slices", - entity="entity", - revision="thesis-v1", - ) - - runner.cmd_tag(args) - - assert "sprint:2" in matching.tags - assert matching.updated is True - assert "sprint:2" not in stale.tags - assert stale.updated is False + assert not hasattr(runner, "cmd_tag") def test_wandb_logger_includes_model_size_metadata(self, monkeypatch): import slices.training.utils as training_utils @@ -1860,8 +1853,8 @@ def __init__(self, **kwargs): cfg = OmegaConf.create( { - "output_dir": "outputs/sprint7p/example", - "sprint": "7p", + "output_dir": "outputs/capacity_study/example", + "experiment_class": "capacity_study", "model_size": "medium", "source_dataset": "eicu", "label_fraction": 0.1, @@ -1870,7 +1863,7 @@ def __init__(self, **kwargs): "use_wandb": True, "wandb_project": "slices", "wandb_entity": None, - "run_name": "s7p_finetune_miiv_mae_mortality_24h_seed42", + "run_name": "capacity_study_finetune_miiv_mae_mortality_24h_seed42", "wandb_group": "finetune_miiv_mae_mortality_24h", "wandb_tags": ["phase:finetune"], }, @@ -1880,8 +1873,11 @@ def __init__(self, **kwargs): logger = training_utils.setup_wandb_logger(cfg) assert isinstance(logger, DummyWandbLogger) - assert captured["kwargs"]["name"] == "s7p_finetune_miiv_mae_mortality_24h_seed42_medium" + assert captured["kwargs"]["name"] == ( + "capacity_study_finetune_miiv_mae_mortality_24h_seed42_medium" + ) assert captured["kwargs"]["group"] == "finetune_miiv_mae_mortality_24h_medium_frac01" + assert "experiment_class:capacity_study" in captured["kwargs"]["tags"] assert "model_size:medium" in captured["kwargs"]["tags"] assert "ablation:label-efficiency" in captured["kwargs"]["tags"] assert "ablation:transfer" in captured["kwargs"]["tags"] @@ -1891,21 +1887,21 @@ def test_transfer_finetune_command_propagates_source_dataset(self): pretrain = Run( id="pretrain_mae_miiv_seed42", - sprint="7", + experiment_class="cross_dataset_transfer", run_type="pretrain", paradigm="mae", dataset="miiv", seed=42, - output_dir="outputs/sprint7/pretrain_mae_miiv_seed42", + output_dir="outputs/cross_dataset_transfer/pretrain_mae_miiv_seed42", ) finetune = Run( id="finetune_mae_eicu_seed42", - sprint="7", + experiment_class="cross_dataset_transfer", run_type="finetune", paradigm="mae", dataset="eicu", seed=42, - output_dir="outputs/sprint7/finetune_mae_mortality_24h_eicu_seed42_from_miiv", + output_dir="outputs/cross_dataset_transfer/finetune_mae_mortality_24h_eicu_seed42_from_miiv", depends_on=[pretrain.id], task="mortality_24h", freeze_encoder=False, @@ -1920,21 +1916,21 @@ def test_protocol_a_finetune_command_uses_strict_linear_probe(self): pretrain = Run( id="pretrain_mae_miiv_seed42", - sprint="2", + experiment_class="core_ssl_benchmark", run_type="pretrain", paradigm="mae", dataset="miiv", seed=42, - output_dir="outputs/sprint1/pretrain_mae_miiv_seed42", + output_dir="outputs/core_ssl_benchmark/pretrain_mae_miiv_seed42", ) probe = Run( id="probe_mae_miiv_seed42", - sprint="2", + experiment_class="core_ssl_benchmark", run_type="finetune", paradigm="mae", dataset="miiv", seed=42, - output_dir="outputs/sprint2/probe_mae_mortality_24h_miiv_seed42", + output_dir="outputs/core_ssl_benchmark/probe_mae_mortality_24h_miiv_seed42", depends_on=[pretrain.id], task="mortality_24h", freeze_encoder=True, @@ -1952,21 +1948,21 @@ def test_protocol_b_finetune_command_keeps_task_mlp_head(self): pretrain = Run( id="pretrain_mae_miiv_seed42", - sprint="1", + experiment_class="core_ssl_benchmark", run_type="pretrain", paradigm="mae", dataset="miiv", seed=42, - output_dir="outputs/sprint1/pretrain_mae_miiv_seed42", + output_dir="outputs/core_ssl_benchmark/pretrain_mae_miiv_seed42", ) finetune = Run( id="finetune_mae_miiv_seed42", - sprint="1", + experiment_class="core_ssl_benchmark", run_type="finetune", paradigm="mae", dataset="miiv", seed=42, - output_dir="outputs/sprint1/finetune_mae_mortality_24h_miiv_seed42", + output_dir="outputs/core_ssl_benchmark/finetune_mae_mortality_24h_miiv_seed42", depends_on=[pretrain.id], task="mortality_24h", freeze_encoder=False, @@ -2003,37 +1999,37 @@ def test_manual_finetune_protocol_configs_encode_safe_defaults(self): class TestExperimentRunnerRetry: """Tests for retry scoping and dependency preservation.""" - def test_cmd_retry_respects_sprint_filter_but_keeps_dependencies(self, monkeypatch): + 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", - sprint="0", + experiment_class="core_ssl_benchmark", run_type="pretrain", paradigm="mae", dataset="miiv", seed=42, - output_dir="outputs/sprint0/pretrain_mae_miiv_seed42", + output_dir="outputs/core_ssl_benchmark/pretrain_mae_miiv_seed42", ) target = runner.Run( id="target_finetune", - sprint="1", + experiment_class="label_efficiency", run_type="finetune", paradigm="mae", dataset="miiv", seed=42, - output_dir="outputs/sprint1/finetune_mae_miiv_seed42", + output_dir="outputs/label_efficiency/finetune_mae_miiv_seed42", depends_on=[dependency.id], task="mortality_24h", ) other = runner.Run( id="other_failed", - sprint="2", + experiment_class="classical_baselines", run_type="supervised", paradigm="supervised", dataset="miiv", seed=42, - output_dir="outputs/sprint2/supervised_miiv_seed42", + output_dir="outputs/classical_baselines/supervised_miiv_seed42", task="mortality_24h", ) @@ -2058,7 +2054,7 @@ def test_cmd_retry_respects_sprint_filter_but_keeps_dependencies(self, monkeypat ) args = SimpleNamespace( - sprint=["1"], + experiment_class=["label_efficiency"], failed=True, skipped=False, revision=None, @@ -2080,21 +2076,21 @@ def test_cmd_retry_skipped_reports_failed_dependencies(self, monkeypatch, capsys dependency = runner.Run( id="dep_pretrain", - sprint="0", + experiment_class="core_ssl_benchmark", run_type="pretrain", paradigm="mae", dataset="miiv", seed=42, - output_dir="outputs/sprint0/pretrain_mae_miiv_seed42", + output_dir="outputs/core_ssl_benchmark/pretrain_mae_miiv_seed42", ) target = runner.Run( id="target_finetune", - sprint="1", + experiment_class="label_efficiency", run_type="finetune", paradigm="mae", dataset="miiv", seed=42, - output_dir="outputs/sprint1/finetune_mae_miiv_seed42", + output_dir="outputs/label_efficiency/finetune_mae_miiv_seed42", depends_on=[dependency.id], task="mortality_24h", ) @@ -2117,7 +2113,7 @@ def test_cmd_retry_skipped_reports_failed_dependencies(self, monkeypatch, capsys ) args = SimpleNamespace( - sprint=["1"], + experiment_class=["label_efficiency"], failed=False, skipped=True, revision=None, @@ -2143,21 +2139,21 @@ def test_cmd_retry_skipped_resets_failed_dependencies_when_requested(self, monke dependency = runner.Run( id="dep_pretrain", - sprint="0", + experiment_class="core_ssl_benchmark", run_type="pretrain", paradigm="mae", dataset="miiv", seed=42, - output_dir="outputs/sprint0/pretrain_mae_miiv_seed42", + output_dir="outputs/core_ssl_benchmark/pretrain_mae_miiv_seed42", ) target = runner.Run( id="target_finetune", - sprint="1", + experiment_class="label_efficiency", run_type="finetune", paradigm="mae", dataset="miiv", seed=42, - output_dir="outputs/sprint1/finetune_mae_miiv_seed42", + output_dir="outputs/label_efficiency/finetune_mae_miiv_seed42", depends_on=[dependency.id], task="mortality_24h", ) @@ -2182,7 +2178,7 @@ def test_cmd_retry_skipped_resets_failed_dependencies_when_requested(self, monke ) args = SimpleNamespace( - sprint=["1"], + experiment_class=["label_efficiency"], failed=True, skipped=True, revision=None, @@ -2204,28 +2200,30 @@ def test_cmd_retry_revises_dependencies_for_revision_scoped_retry(self, monkeypa import scripts.internal.run_experiments as runner dependency = runner.Run( - id="s1_pretrain_mae_miiv_seed42", - sprint="1", + 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/sprint1/pretrain_mae_miiv_seed42", + output_dir="outputs/core_ssl_benchmark/pretrain_mae_miiv_seed42", ) target = runner.Run( - id="s2_finetune_mae_mortality_24h_miiv_seed42", - sprint="2", + 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/sprint2/finetune_mae_mortality_24h_miiv_seed42", + output_dir="outputs/label_efficiency/finetune_mae_mortality_24h_miiv_seed42", depends_on=[dependency.id], task="mortality_24h", ) - revised_dep_id = "s1_rev-thesis-v1_pretrain_mae_miiv_seed42" - revised_target_id = "s2_rev-thesis-v1_finetune_mae_mortality_24h_miiv_seed42" + 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": { @@ -2246,7 +2244,7 @@ def test_cmd_retry_revises_dependencies_for_revision_scoped_retry(self, monkeypa ) args = SimpleNamespace( - sprint=["2"], + experiment_class=["label_efficiency"], failed=True, skipped=False, revision="thesis-v1", @@ -2270,21 +2268,21 @@ def test_select_ready_runs_prioritizes_pretrains_with_slot_budget(self): pretrain = runner.Run( id="pretrain", - sprint="10", + experiment_class="core_ssl_benchmark", run_type="pretrain", paradigm="mae", dataset="miiv", seed=42, - output_dir="outputs/sprint10/pretrain_mae_miiv_seed42", + output_dir="outputs/core_ssl_benchmark/pretrain_mae_miiv_seed42", ) finetune = runner.Run( id="finetune", - sprint="10", + experiment_class="core_ssl_benchmark", run_type="finetune", paradigm="mae", dataset="miiv", seed=42, - output_dir="outputs/sprint10/finetune_mae_miiv_seed42", + output_dir="outputs/core_ssl_benchmark/finetune_mae_miiv_seed42", task="mortality_24h", ) @@ -2302,21 +2300,21 @@ def test_select_ready_runs_does_not_mix_pretrain_with_active_gpu_work(self): pretrain = runner.Run( id="pretrain", - sprint="10", + experiment_class="core_ssl_benchmark", run_type="pretrain", paradigm="mae", dataset="miiv", seed=42, - output_dir="outputs/sprint10/pretrain_mae_miiv_seed42", + output_dir="outputs/core_ssl_benchmark/pretrain_mae_miiv_seed42", ) finetune = runner.Run( id="finetune", - sprint="10", + experiment_class="core_ssl_benchmark", run_type="finetune", paradigm="mae", dataset="miiv", seed=42, - output_dir="outputs/sprint10/finetune_mae_miiv_seed42", + output_dir="outputs/core_ssl_benchmark/finetune_mae_miiv_seed42", task="mortality_24h", ) @@ -2334,12 +2332,12 @@ def test_run_scheduler_returns_nonzero_for_failed_run(self, tmp_path, monkeypatc failed_run = runner.Run( id="failed", - sprint="1", + experiment_class="core_ssl_benchmark", run_type="supervised", paradigm="supervised", dataset="miiv", seed=42, - output_dir="outputs/sprint1/supervised_miiv_seed42", + output_dir="outputs/core_ssl_benchmark/supervised_miiv_seed42", task="mortality_24h", ) state = {"version": 1, "runs": {failed_run.id: {"status": "failed"}}} @@ -2356,12 +2354,12 @@ def test_scheduler_exit_code_nonzero_for_dependency_skips(self): skipped_run = runner.Run( id="skipped", - sprint="1", + experiment_class="core_ssl_benchmark", run_type="finetune", paradigm="mae", dataset="miiv", seed=42, - output_dir="outputs/sprint1/finetune_mae_miiv_seed42", + output_dir="outputs/core_ssl_benchmark/finetune_mae_miiv_seed42", task="mortality_24h", ) state = {"version": 1, "runs": {skipped_run.id: {"status": "skipped"}}} @@ -2385,8 +2383,8 @@ def test_build_statistical_tests_df_includes_xgboost_comparisons(self): for task, base in [("mortality_24h", 0.30), ("aki_kdigo", 0.25)]: rows.append( { + "experiment_class": "classical_baselines", "experiment_type": "classical_baselines", - "sprint": "11", "paradigm": paradigm, "dataset": "miiv", "task": task, @@ -2875,51 +2873,49 @@ def fake_finetune_module(config, checkpoint_path=None, pretrain_checkpoint_path= class TestExperimentRunnerMatrix: - """Tests for thesis-final sprint matrix coverage.""" + """Tests for thesis-final experiment-class matrix coverage.""" - def test_classical_baselines_are_canonicalized_to_sprint11(self): + def test_matrix_counts_match_final_experiment_classes(self): from collections import Counter - from scripts.internal.run_experiments import generate_all_runs + from scripts.internal.run_experiments import ( + EXPECTED_CLASS_COUNTS, + generate_all_runs, + scientific_fingerprint, + ) runs = generate_all_runs() - by_sprint = Counter(run.sprint for run in runs) + by_class = Counter(run.experiment_class for run in runs) assert len(runs) == 2305 - assert by_sprint["1"] == 19 - assert by_sprint["3"] == 31 - assert by_sprint["4"] == 31 - assert by_sprint["5"] == 186 - assert by_sprint["6"] == 504 - assert by_sprint["11"] == 360 - - non_s11_classical = [ - run for run in runs if run.run_type in ("gru_d", "xgboost") and run.sprint != "11" + assert sum(run.experiment_class != "smart_external_reference" for run in runs) == 2170 + 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) + + non_classical_baselines = [ + run + for run in runs + if run.run_type in ("gru_d", "xgboost") + and run.experiment_class != "classical_baselines" ] - assert non_s11_classical == [] + assert non_classical_baselines == [] - def test_sprint7p_expanded_to_five_seeds(self): - from scripts.internal.run_experiments import BASELINE_SPRINTS, SEEDS_EXTENDED, MatrixBuilder + def test_capacity_study_expanded_to_five_seeds(self): + from scripts.internal.run_experiments import SEEDS_EXTENDED, MatrixBuilder builder = MatrixBuilder() - builder.build_sprint7p() + builder.build_capacity_study() assert len(builder.runs) == 100 assert sorted({run.seed for run in builder.runs}) == SEEDS_EXTENDED - assert BASELINE_SPRINTS["7p"] == ["6", "10"] - assert "10" in BASELINE_SPRINTS["6"] - assert "10" in BASELINE_SPRINTS["7"] - assert "10" in BASELINE_SPRINTS["8"] - assert {run.extra_overrides["+model_size"] for run in builder.runs} == { - "medium", - "large", - } + assert {run.model_size for run in builder.runs} == {"medium", "large"} - def test_sprint13_includes_both_protocols(self): + def test_ts2vec_extension_includes_both_protocols(self): from scripts.internal.run_experiments import MatrixBuilder builder = MatrixBuilder() - builder.build_sprint13() + builder.build_ts2vec_extension() assert len(builder.runs) == 135 pretrains = [run for run in builder.runs if run.run_type == "pretrain"] From f8c8a754aeb3bf9a8c9e8b86963702d8a643dd79 Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 22 Apr 2026 15:14:30 -0400 Subject: [PATCH 100/121] Harden thesis launch and export validation Fix generated downstream protocol overrides so SSL finetune and probe jobs compose while baseline configs keep protocol metadata without loading protocol groups. Add publication export safeguards for expected matrix coverage, duplicate fingerprints, W&B tag filters, and validation-before-write behavior. Adjust contrastive mask sweeps and AKI missingness documentation to preserve benchmark interpretation. Implications: generated finetune/probe commands now launch with protocol=A/B, XGBoost exports report xgboost.learning_rate, and publication exports fail closed on absent or duplicate matrix cells unless explicitly allowed. --- configs/gru_d.yaml | 1 + configs/supervised.yaml | 1 + configs/tasks/aki_kdigo.yaml | 5 +- configs/xgboost.yaml | 1 + scripts/eval/evaluate_fairness.py | 2 +- scripts/export_results.py | 165 ++++++++++++++++++++++++++- scripts/internal/run_experiments.py | 80 +++++++++++-- scripts/training/xgboost_baseline.py | 11 +- src/slices/data/tasks/aki_kdigo.yaml | 5 +- src/slices/training/utils.py | 3 + tests/test_export_results.py | 124 ++++++++++++++++++-- tests/test_fixes.py | 98 +++++++++++++--- 12 files changed, 452 insertions(+), 44 deletions(-) diff --git a/configs/gru_d.yaml b/configs/gru_d.yaml index 685f2b2..179eb6f 100644 --- a/configs/gru_d.yaml +++ b/configs/gru_d.yaml @@ -24,6 +24,7 @@ dataset: miiv # Paradigm paradigm: gru_d +protocol: B # Final rerun corpus metadata experiment_class: null diff --git a/configs/supervised.yaml b/configs/supervised.yaml index 1cb40e6..5d26067 100644 --- a/configs/supervised.yaml +++ b/configs/supervised.yaml @@ -32,6 +32,7 @@ dataset: miiv # miiv | eicu | combined # Paradigm — always "supervised" for this config paradigm: supervised +protocol: B # Final rerun corpus metadata experiment_class: null diff --git a/configs/tasks/aki_kdigo.yaml b/configs/tasks/aki_kdigo.yaml index 409957b..e953bbc 100644 --- a/configs/tasks/aki_kdigo.yaml +++ b/configs/tasks/aki_kdigo.yaml @@ -17,7 +17,10 @@ label_params: relative_rise_threshold: 1.5 relative_window_hours: 168 quality_checks: - max_missing_percentage: 28.0 + # 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 diff --git a/configs/xgboost.yaml b/configs/xgboost.yaml index f7eda4b..9578034 100644 --- a/configs/xgboost.yaml +++ b/configs/xgboost.yaml @@ -21,6 +21,7 @@ dataset: miiv # Paradigm paradigm: xgboost +protocol: B # Final rerun corpus metadata experiment_class: null diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index 8dd44a3..a054b42 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -143,7 +143,7 @@ def fetch_eval_runs( tag_filters.append(f"revision:{revisions[0]}") if tag_filters: - filters["tags"] = {"$all": tag_filters} + filters["$and"] = [{"tags": tag} for tag in tag_filters] 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") diff --git a/scripts/export_results.py b/scripts/export_results.py index 4273a96..8565861 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -211,7 +211,7 @@ def fetch_all_runs( if revision and len(revision) == 1: tag_filters.append(f"revision:{revision[0]}") if tag_filters: - filters["tags"] = {"$all": tag_filters} + filters["$and"] = [{"tags": tag} for tag in tag_filters] print(f"Fetching runs from {path}...", file=sys.stderr) print(f" Server-side filters: {json.dumps(filters, default=str)}", file=sys.stderr) @@ -343,7 +343,7 @@ def _recover_pretrain_metadata( if up_lr is not None: subtype = "lr_sensitivity" - mr_match = re.search(r"_mask_ratio([^_]+)", output_dir) + 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: @@ -418,6 +418,9 @@ def extract_run(run, metric_keys: list[str]) -> dict: metrics[key] = _metric_value_or_nan(value) 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, @@ -433,7 +436,7 @@ def extract_run(run, metric_keys: list[str]) -> dict: "seed": config.get("seed", None), "protocol": protocol, "label_fraction": config.get("label_fraction", 1.0), - "lr": _get_nested(config, "optimizer.lr", default=None), + "lr": lr, "mask_ratio": _get_nested(config, "ssl.mask_ratio", default=None), "model_size": _infer_model_size(config), "source_dataset": source_dataset, @@ -492,6 +495,7 @@ def _assert_no_ambiguous_fingerprint_collisions(df: pd.DataFrame) -> None: def build_per_seed_df( runs: list, allow_extraction_failures: bool = False, + allow_duplicate_fingerprints: bool = False, ) -> pd.DataFrame: rows = [] failed = [] @@ -531,6 +535,18 @@ def build_per_seed_df( 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: @@ -1091,15 +1107,135 @@ def _checkpoint_provenance_issues(per_seed_df: pd.DataFrame) -> list[tuple[pd.Se return issues +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) @@ -1263,6 +1399,14 @@ def parse_args() -> argparse.Namespace: "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", @@ -1309,6 +1453,7 @@ def main() -> None: 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, @@ -1337,6 +1482,13 @@ def main() -> None: 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() @@ -1346,10 +1498,15 @@ def main() -> None: 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, @@ -1374,8 +1531,6 @@ def main() -> None: 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 warnings and not args.allow_incomplete: - sys.exit(1) if __name__ == "__main__": diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index 917fc9d..bbc7a17 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -186,8 +186,8 @@ def _metadata_overrides(self) -> list[str]: ] if self.experiment_subtype is not None: overrides.append(f"experiment_subtype={self.experiment_subtype}") - if self.protocol is not None: - overrides.append(f"+protocol={self.protocol}") + if self.protocol is not None and self.run_type == "finetune": + overrides.append(f"protocol={self.protocol}") if self.model_size is not None: overrides.append(f"+model_size={self.model_size}") for k, v in self.extra_overrides.items(): @@ -664,6 +664,8 @@ def build_hp_robustness(self) -> None: ) for mask_ratio in MASK_RATIO_ROBUSTNESS: extra = {"ssl.mask_ratio": mask_ratio} + if paradigm == "contrastive": + extra["ssl.complementary_masks"] = False subtype = "mask_ratio_sensitivity" pretrain = self._add_pretrain( experiment_class, @@ -918,6 +920,59 @@ def set_run_status(state: dict, run_id: str, status: str, **kwargs) -> None: state["runs"][run_id].update(kwargs) +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 reset_state_for_launch_commit_mismatch(runs: list[Run], state: dict) -> None: + """Do not trust completed/skipped state from a different reviewed commit.""" + reset = 0 + for run in runs: + expected_commit = _run_launch_commit(run) + if expected_commit is None: + continue + + info = state["runs"].get(run.id) + if not info: + continue + + status = info.get("status", "pending") + if status == "pending": + continue + + actual_commit = _state_launch_commit(info) + if actual_commit == expected_commit: + continue + + if status == "running": + print( + f" WARNING {run.id}: state is running from launch_commit=" + f"{actual_commit or 'unknown'}, expected {expected_commit}; leaving it untouched." + ) + continue + + state["runs"][run.id] = { + "status": "pending", + "reset_reason": ( + f"launch_commit changed from {actual_commit or 'unknown'} " f"to {expected_commit}" + ), + } + reset += 1 + + if reset: + print(f"Reset {reset} stale run state entries for launch_commit mismatch.") + + def is_pid_alive(pid: int) -> bool: try: os.kill(pid, 0) @@ -1040,6 +1095,8 @@ def handle_signal(signum, frame): _print_dry_run(runs, runs_by_id) return 0 + reset_state_for_launch_commit_mismatch(runs, state) + print(f"Scheduler started (slot_budget={parallel})") print(f"State file: {STATE_FILE}\n") @@ -1099,15 +1156,16 @@ def handle_signal(signum, frame): ) proc._log_fh = log_fh # type: ignore[attr-defined] active[run.id] = proc - set_run_status( - state, - run.id, - "running", - started_at=now, - pid=proc.pid, - log_file=str(log_file), - command=shlex.join(cmd), - ) + status_updates = { + "started_at": now, + "pid": proc.pid, + "log_file": str(log_file), + "command": shlex.join(cmd), + } + launch_commit = _run_launch_commit(run) + if launch_commit is not None: + status_updates["launch_commit"] = launch_commit + set_run_status(state, run.id, "running", **status_updates) save_state(state) diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index 28630b6..0341a43 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -325,11 +325,18 @@ def main(cfg: DictConfig) -> None: if cfg.logging.get("use_wandb", False): import wandb + 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}" + run = wandb.init( project=cfg.logging.wandb_project, entity=cfg.logging.get("wandb_entity", None), - name=cfg.logging.get("run_name", f"xgboost_{cfg.dataset}_{task_name}"), - group=cfg.logging.get("wandb_group", f"xgboost_{cfg.dataset}_{task_name}"), + name=run_name, + group=group_name, tags=_build_wandb_tags(cfg), config=OmegaConf.to_container(cfg, resolve=True), ) diff --git a/src/slices/data/tasks/aki_kdigo.yaml b/src/slices/data/tasks/aki_kdigo.yaml index ac06651..57ae24d 100644 --- a/src/slices/data/tasks/aki_kdigo.yaml +++ b/src/slices/data/tasks/aki_kdigo.yaml @@ -13,7 +13,10 @@ label_params: relative_rise_threshold: 1.5 relative_window_hours: 168 # 7 days per KDIGO 2012 specification quality_checks: - max_missing_percentage: 28.0 # Warn if 24h baseline pushes exclusions above prior range + # 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/training/utils.py b/src/slices/training/utils.py index 29bf374..9e1be8b 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -342,6 +342,9 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: run_name = run_name.replace("_finetune_", "_probe_", 1) 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. diff --git a/tests/test_export_results.py b/tests/test_export_results.py index 2fc480c..f7e4f1d 100644 --- a/tests/test_export_results.py +++ b/tests/test_export_results.py @@ -138,12 +138,12 @@ def runs(self, path, filters, order): ) assert captured["path"] == "entity/proj" - assert captured["filters"]["tags"]["$all"] == [ - "experiment_class:core_ssl_benchmark", - "paradigm:mae", - "dataset:miiv", - "phase:finetune", - "revision:thesis-v1", + assert captured["filters"]["$and"] == [ + {"tags": "experiment_class:core_ssl_benchmark"}, + {"tags": "paradigm:mae"}, + {"tags": "dataset:miiv"}, + {"tags": "phase:finetune"}, + {"tags": "revision:thesis-v1"}, ] @@ -160,8 +160,7 @@ def test_build_per_seed_df_keeps_hp_robustness_rows_out_of_derived_views(): "paradigm": "mae", "seed": 42, "output_dir": ( - "outputs/hp_robustness/" - f"finetune_mae_mortality_24h_miiv_seed42_lr{lr_code}" + "outputs/hp_robustness/" f"finetune_mae_mortality_24h_miiv_seed42_lr{lr_code}" ), "encoder": {"d_model": 64, "n_layers": 2}, "training": {"freeze_encoder": False}, @@ -193,6 +192,53 @@ def test_build_per_seed_df_fails_closed_on_extraction_errors(): 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": "B", + "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") @@ -277,6 +323,27 @@ def test_validate_flags_wrong_fixed_seed_set_with_correct_count(): 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"] == "B")].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") @@ -485,6 +552,18 @@ def test_parse_args_exposes_extraction_failure_escape_hatch(monkeypatch): 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_main_exits_nonzero_when_no_runs_match(monkeypatch, tmp_path): mod = importlib.import_module("scripts.export_results") @@ -501,6 +580,35 @@ def test_main_exits_nonzero_when_no_runs_match(monkeypatch, tmp_path): 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") diff --git a/tests/test_fixes.py b/tests/test_fixes.py index bbf3927..7222e30 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -773,12 +773,12 @@ def runs(self, path, filters, order): ) assert captured["path"] == "entity/proj" - assert captured["filters"]["tags"]["$all"] == [ - "experiment_class:core_ssl_benchmark", - "paradigm:mae", - "dataset:miiv", - "phase:finetune", - "revision:thesis-v1", + assert captured["filters"]["$and"] == [ + {"tags": "experiment_class:core_ssl_benchmark"}, + {"tags": "paradigm:mae"}, + {"tags": "dataset:miiv"}, + {"tags": "phase:finetune"}, + {"tags": "revision:thesis-v1"}, ] def test_fetch_eval_runs_multi_revision_filters_client_side(self, monkeypatch): @@ -1398,14 +1398,15 @@ 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 - config = { - "output_dir": ( - "outputs/hp_robustness/" "finetune_jepa_mortality_24h_miiv_seed42_mask_ratio03" - ) - } - up_lr, up_mr, subtype = _recover_pretrain_metadata("some_run", config) - assert up_mr == 0.3 - assert subtype == "mask_ratio_sensitivity" + 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.""" @@ -1688,6 +1689,36 @@ def test_apply_wandb_target_injects_project_and_entity(self): 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_cmd_run_respects_requested_class_order(self, monkeypatch): import scripts.internal.run_experiments as runner @@ -1874,7 +1905,7 @@ def __init__(self, **kwargs): assert isinstance(logger, DummyWandbLogger) assert captured["kwargs"]["name"] == ( - "capacity_study_finetune_miiv_mae_mortality_24h_seed42_medium" + "capacity_study_finetune_miiv_mae_mortality_24h_seed42_medium_frac01" ) assert captured["kwargs"]["group"] == "finetune_miiv_mae_mortality_24h_medium_frac01" assert "experiment_class:capacity_study" in captured["kwargs"]["tags"] @@ -1942,6 +1973,8 @@ def test_protocol_a_finetune_command_uses_strict_linear_probe(self): assert "task.head_type=linear" in cmd assert "task.hidden_dims=[]" in cmd assert "task.dropout=0.0" in cmd + assert "protocol=A" in cmd + assert "+protocol=A" not in cmd def test_protocol_b_finetune_command_keeps_task_mlp_head(self): from scripts.internal.run_experiments import Run @@ -1972,6 +2005,27 @@ def test_protocol_b_finetune_command_keeps_task_mlp_head(self): assert "training.freeze_encoder=false" in cmd assert "task.head_type=linear" not in cmd + assert "protocol=B" 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=B" not in cmd + assert "protocol=B" not in cmd def test_manual_finetune_protocol_configs_encode_safe_defaults(self): protocol_a = OmegaConf.load("configs/protocol/a.yaml") @@ -2901,6 +2955,20 @@ def test_matrix_counts_match_final_experiment_classes(self): ] 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 == "mask_ratio_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 From 7f1c076a2f83cff4800b73f12fefa002791c23f5 Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 22 Apr 2026 15:58:44 -0400 Subject: [PATCH 101/121] Harden experiment rerun validation Quarantine stale run outputs when launch commits change, require complete fairness reports before writing W&B metrics, and skip undefined single-class bootstrap samples for AUROC/AUPRC metrics. This prevents cross-commit checkpoint reuse and makes fairness/statistical analysis fail closed; no migration is required. --- scripts/eval/evaluate_fairness.py | 40 ++++++++++++++++- scripts/internal/run_experiments.py | 49 +++++++++++++++++++-- src/slices/eval/statistical.py | 48 +++++++++++++++++---- tests/test_evaluate_fairness.py | 66 +++++++++++++++++++++++++++++ tests/test_fixes.py | 41 ++++++++++++++++++ tests/test_statistical.py | 31 ++++++++++++++ 6 files changed, 262 insertions(+), 13 deletions(-) diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index a054b42..7f85792 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -47,6 +47,7 @@ import re import sys import time +from numbers import Real from pathlib import Path from typing import Any, Optional @@ -194,7 +195,7 @@ def _has_finite_summary_value(summary: dict[str, Any], key: str) -> bool: value = summary.get(key) if value is None: return False - if isinstance(value, float) and math.isnan(value): + if isinstance(value, Real) and not isinstance(value, bool) and not math.isfinite(float(value)): return False return True @@ -241,6 +242,31 @@ def has_fairness_metrics(run, protected_attributes: list[str]) -> bool: 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_finite_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], @@ -845,6 +871,18 @@ def _sort_key(r): 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) run_path = ( diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index bbc7a17..1a3a42a 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -934,9 +934,39 @@ def _state_launch_commit(info: dict) -> str | None: 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_commit: str | None, + expected_commit: 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_commit) + expected = _safe_path_component(expected_commit) + 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_commit_mismatch(runs: list[Run], state: dict) -> None: - """Do not trust completed/skipped state from a different reviewed commit.""" + """Do not trust state or artifacts from a different reviewed commit.""" reset = 0 + quarantined = 0 for run in runs: expected_commit = _run_launch_commit(run) if expected_commit is None: @@ -947,12 +977,11 @@ def reset_state_for_launch_commit_mismatch(runs: list[Run], state: dict) -> None continue status = info.get("status", "pending") - if status == "pending": - continue - actual_commit = _state_launch_commit(info) if actual_commit == expected_commit: continue + if status == "pending" and actual_commit is None and not Path(run.output_dir).exists(): + continue if status == "running": print( @@ -961,16 +990,28 @@ def reset_state_for_launch_commit_mismatch(runs: list[Run], state: dict) -> None ) continue + quarantined_output_dir = _quarantine_stale_output_dir( + run, + actual_commit, + expected_commit, + ) + if quarantined_output_dir: + quarantined += 1 + state["runs"][run.id] = { "status": "pending", "reset_reason": ( f"launch_commit changed from {actual_commit or 'unknown'} " f"to {expected_commit}" ), } + 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_commit mismatch.") + if quarantined: + print(f"Quarantined {quarantined} stale output directories before relaunch.") def is_pid_alive(pid: int) -> bool: diff --git a/src/slices/eval/statistical.py b/src/slices/eval/statistical.py index 98c3ac3..1be84e8 100644 --- a/src/slices/eval/statistical.py +++ b/src/slices/eval/statistical.py @@ -45,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): @@ -53,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: @@ -115,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 @@ -130,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 @@ -149,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, @@ -305,6 +320,23 @@ def _compute_metric( 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], diff --git a/tests/test_evaluate_fairness.py b/tests/test_evaluate_fairness.py index 7488f6a..1c7969c 100644 --- a/tests/test_evaluate_fairness.py +++ b/tests/test_evaluate_fairness.py @@ -175,3 +175,69 @@ def test_has_fairness_metrics_requires_regression_metric_family(): 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, + "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, + "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_rejects_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, + "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 "gender: missing worst_group_auroc" in missing diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 7222e30..7c432f2 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -1719,6 +1719,47 @@ def test_launch_commit_mismatch_resets_stale_completed_state(self): 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_cmd_run_respects_requested_class_order(self, monkeypatch): import scripts.internal.run_experiments as runner diff --git a/tests/test_statistical.py b/tests/test_statistical.py index 5fe995d..47f39fb 100644 --- a/tests/test_statistical.py +++ b/tests/test_statistical.py @@ -17,6 +17,7 @@ paired_bootstrap_test, paired_wilcoxon_signed_rank, ) +from torchmetrics import AUROC # ── Helpers ────────────────────────────────────────────────────────────────── @@ -95,6 +96,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 ──────────────────────────────────────────────────── @@ -215,6 +228,24 @@ def test_reproducibility_with_seed(self): 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): From 6f0ceaf9f56b0f2eb52a37fac01fa6840bfcd621 Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 22 Apr 2026 16:12:37 -0400 Subject: [PATCH 102/121] Fix launch identity and checkpoint safeguards Bind runner state and checkpoint resume decisions to revision, W&B target, and launch commit so stale state cannot silently skip or resume a different corpus. Reject ambiguous finetune checkpoint inputs and emit per-task statistical comparisons alongside the omnibus tests. Behavioral implications: scoped reruns quarantine mismatched output directories and only resume checkpoints with matching launch markers; exported statistical tests now include comparison_scope rows for omnibus and per-task claims. --- configs/finetune.yaml | 3 +- scripts/export_results.py | 103 +++++----- scripts/internal/run_experiments.py | 250 ++++++++++++++++++++++--- scripts/training/finetune.py | 7 +- src/slices/training/finetune_module.py | 5 + tests/test_export_results.py | 9 +- tests/test_finetune_module.py | 9 + tests/test_fixes.py | 137 ++++++++++++++ 8 files changed, 441 insertions(+), 82 deletions(-) diff --git a/configs/finetune.yaml b/configs/finetune.yaml index 0659b66..1d39c91 100644 --- a/configs/finetune.yaml +++ b/configs/finetune.yaml @@ -46,8 +46,7 @@ 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 diff --git a/scripts/export_results.py b/scripts/export_results.py index 8565861..9c77c4f 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -876,60 +876,69 @@ def build_statistical_tests_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: ) ] - scope_cols = [ + 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 scope_cols: - scope_cols.append("primary_metric_name") - - 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 - - paradigms = sorted( - paradigm - for paradigm in scope_group["paradigm"].dropna().unique().tolist() - if paradigm != "__none__" - ) - if len(paradigms) < 2: - continue + if "primary_metric_name" not in base_scope_cols: + base_scope_cols.append("primary_metric_name") - 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: + 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 - 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, - ) + 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 - 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) + 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() @@ -938,7 +947,9 @@ def build_statistical_tests_df(per_seed_df: pd.DataFrame) -> pd.DataFrame: col for col in [ "experiment_class", + "comparison_scope", "dataset", + "task", "protocol", "label_fraction", "primary_metric_name", diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index 1a3a42a..72a6f1a 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -112,6 +112,8 @@ 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") PROTO_A = {"freeze_encoder": True, "max_epochs": 50, "patience": 10, "lr": 1e-4} PROTO_B = {"freeze_encoder": False, "max_epochs": 100, "patience": 10, "lr": 3e-4} @@ -196,7 +198,7 @@ def _metadata_overrides(self) -> list[str]: def _append_resume(self, cmd: list[str]) -> None: last_ckpt = Path(self.output_dir) / "checkpoints" / "last.ckpt" - if last_ckpt.exists(): + if last_ckpt.exists() and _can_resume_from_last_checkpoint(self): cmd.append(f"ckpt_path={last_ckpt}") def _pretrain_cmd(self) -> list[str]: @@ -910,6 +912,164 @@ def save_state(state: dict) -> None: 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") @@ -941,8 +1101,8 @@ def _safe_path_component(value: str | None) -> str: def _quarantine_stale_output_dir( run: Run, - actual_commit: str | None, - expected_commit: str, + actual_identity: str | None, + expected_identity: str, ) -> str | None: output_dir = Path(run.output_dir) if not output_dir.exists(): @@ -950,8 +1110,8 @@ def _quarantine_stale_output_dir( parent = output_dir.parent timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") - actual = _safe_path_component(actual_commit) - expected = _safe_path_component(expected_commit) + 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 @@ -963,57 +1123,83 @@ def _quarantine_stale_output_dir( return str(candidate) -def reset_state_for_launch_commit_mismatch(runs: list[Run], state: dict) -> None: - """Do not trust state or artifacts from a different reviewed commit.""" +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_commit = _run_launch_commit(run) - if expected_commit is None: - continue - + expected_identity = _run_launch_identity(run) info = state["runs"].get(run.id) - if not info: - continue + 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)}" + ) - status = info.get("status", "pending") - actual_commit = _state_launch_commit(info) - if actual_commit == expected_commit: - continue - if status == "pending" and actual_commit is None and not Path(run.output_dir).exists(): + 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 is running from launch_commit=" - f"{actual_commit or 'unknown'}, expected {expected_commit}; leaving it untouched." + 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, - actual_commit, - expected_commit, + _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", - "reset_reason": ( - f"launch_commit changed from {actual_commit or 'unknown'} " f"to {expected_commit}" - ), + "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_commit mismatch.") + 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) @@ -1136,7 +1322,7 @@ def handle_signal(signum, frame): _print_dry_run(runs, runs_by_id) return 0 - reset_state_for_launch_commit_mismatch(runs, state) + reset_state_for_launch_identity_mismatch(runs, state) print(f"Scheduler started (slot_budget={parallel})") print(f"State file: {STATE_FILE}\n") @@ -1188,6 +1374,7 @@ def handle_signal(signum, frame): 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, @@ -1202,10 +1389,11 @@ def handle_signal(signum, frame): "pid": proc.pid, "log_file": str(log_file), "command": shlex.join(cmd), + "launch_identity": _run_launch_identity(run), } - launch_commit = _run_launch_commit(run) - if launch_commit is not None: - status_updates["launch_commit"] = launch_commit + 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) @@ -1298,7 +1486,7 @@ def print_status(experiment_class_filter: list[str] | None = None) -> None: if experiment_class and revision: groups.setdefault((experiment_class, revision), []).append(run_id) - group_keys = sorted(groups) + 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] diff --git a/scripts/training/finetune.py b/scripts/training/finetune.py index fede201..0a8ad8a 100644 --- a/scripts/training/finetune.py +++ b/scripts/training/finetune.py @@ -90,7 +90,12 @@ 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 " diff --git a/src/slices/training/finetune_module.py b/src/slices/training/finetune_module.py index 96f308a..66e4f84 100644 --- a/src/slices/training/finetune_module.py +++ b/src/slices/training/finetune_module.py @@ -74,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) diff --git a/tests/test_export_results.py b/tests/test_export_results.py index f7e4f1d..fa91b8c 100644 --- a/tests/test_export_results.py +++ b/tests/test_export_results.py @@ -52,8 +52,9 @@ def test_build_statistical_tests_df_produces_pairwise_significance_rows(): stats_df = mod.build_statistical_tests_df(pd.DataFrame(rows)) - assert len(stats_df) == 1 - row = stats_df.iloc[0] + 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" @@ -61,6 +62,10 @@ def test_build_statistical_tests_df_produces_pairwise_significance_rows(): 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") diff --git a/tests/test_finetune_module.py b/tests/test_finetune_module.py index b7b0f0e..7c32f9c 100644 --- a/tests/test_finetune_module.py +++ b/tests/test_finetune_module.py @@ -442,6 +442,15 @@ def test_load_encoder_checkpoint(self, sample_config, 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. diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 7c432f2..b5854e4 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -1760,6 +1760,129 @@ def test_launch_commit_mismatch_quarantines_stale_resume_artifacts(self, tmp_pat ) 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_cmd_run_respects_requested_class_order(self, monkeypatch): import scripts.internal.run_experiments as runner @@ -2501,6 +2624,8 @@ def test_build_statistical_tests_df_includes_xgboost_comparisons(self): 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: @@ -2841,6 +2966,18 @@ def test_dataloader(self): 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 = {} From 00bbcca9758141fa1ce9af43d5ccf3732ba37298 Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 22 Apr 2026 16:54:21 -0400 Subject: [PATCH 103/121] Require scoped thesis experiment launches Fail closed when final experiment runs or retries omit revision and W&B project metadata, and make status reporting revision-aware for tmux monitoring. Tighten data and analysis safeguards by validating benchmark metadata, disambiguating W&B run names, logging label-support summaries, filtering fairness/export tasks to the fixed thesis set, and treating written undefined fairness metrics as complete. Behavioral implication: final launches must pass --revision and --project or set WANDB_PROJECT; exports and fairness now default to slices-thesis and exclude non-thesis task rows unless explicitly requested. --- scripts/eval/evaluate_fairness.py | 47 ++++++--- scripts/export_results.py | 7 +- scripts/internal/launch_thesis_tmux.sh | 10 +- scripts/internal/run_experiments.py | 37 +++++-- scripts/training/finetune.py | 5 +- scripts/training/supervised.py | 5 +- scripts/training/xgboost_baseline.py | 6 +- src/slices/training/utils.py | 129 +++++++++++++++++++++++-- tests/test_evaluate_fairness.py | 28 +++++- tests/test_export_results.py | 4 +- tests/test_fixes.py | 87 ++++++++++++++++- 11 files changed, 323 insertions(+), 42 deletions(-) diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index 7f85792..71d7d54 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -42,12 +42,10 @@ import argparse import json import logging -import math import os import re import sys import time -from numbers import Real from pathlib import Path from typing import Any, Optional @@ -72,6 +70,12 @@ ] DEFAULT_PHASES = ["finetune", "supervised", "baseline"] DEFAULT_PROTECTED_ATTRIBUTES = ["gender", "age_group", "race"] +THESIS_TASKS = { + "mortality_24h", + "mortality_hospital", + "aki_kdigo", + "los_remaining", +} BINARY_FAIRNESS_REQUIRED_METRICS = [ "n_valid_groups", "worst_group_auroc", @@ -191,13 +195,9 @@ def _expected_fairness_attributes(run, protected_attributes: list[str]) -> list[ return attrs -def _has_finite_summary_value(summary: dict[str, Any], key: str) -> bool: - value = summary.get(key) - if value is None: - return False - if isinstance(value, Real) and not isinstance(value, bool) and not math.isfinite(float(value)): - return False - return True +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: @@ -235,7 +235,7 @@ def has_fairness_metrics(run, protected_attributes: list[str]) -> bool: for attr in expected_attrs: prefix = f"fairness/{attr}/" for metric_name in required_metrics: - if not _has_finite_summary_value(sm, f"{prefix}{metric_name}"): + if not _has_summary_value(sm, f"{prefix}{metric_name}"): return False return True except Exception: @@ -261,7 +261,7 @@ def missing_fairness_report_requirements( prefix = f"fairness/{attr}/" for metric_name in required_metrics: key = f"{prefix}{metric_name}" - if not _has_finite_summary_value(flat, key): + if not _has_summary_value(flat, key): missing.append(f"{attr}: missing {metric_name}") return missing @@ -502,6 +502,15 @@ def _get_nested(config: dict, dotted_key: str, default=None): 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, @@ -640,8 +649,8 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument( "--project", - default=os.environ.get("WANDB_PROJECT", "slices"), - help="W&B project name (default: $WANDB_PROJECT or 'slices')", + default=os.environ.get("WANDB_PROJECT", "slices-thesis"), + help="W&B project name (default: $WANDB_PROJECT or 'slices-thesis')", ) parser.add_argument( "--entity", @@ -691,6 +700,11 @@ def parse_args() -> argparse.Namespace: "--min-subgroup-size", type=int, default=50, 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", @@ -753,6 +767,13 @@ def main() -> None: ) ) + 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) diff --git a/scripts/export_results.py b/scripts/export_results.py index 9c77c4f..6319f7b 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -162,7 +162,6 @@ PRIMARY_TEST_METRIC_BY_TASK = { "mortality_24h": "test/auprc", "mortality_hospital": "test/auprc", - "mortality": "test/auprc", "aki_kdigo": "test/auprc", "los_remaining": "test/mae", } @@ -1362,8 +1361,8 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument( "--project", - default=os.environ.get("WANDB_PROJECT", "slices"), - help="W&B project name (default: WANDB_PROJECT env var or 'slices')", + 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", @@ -1422,7 +1421,7 @@ def parse_args() -> argparse.Namespace: "--include-extension-tasks", action="store_true", help=( - "Include optional extension tasks such as ICU mortality. By default, " + "Include task rows outside the fixed thesis task set. By default, " f"primary thesis exports are restricted to {sorted(THESIS_TASKS)}." ), ) diff --git a/scripts/internal/launch_thesis_tmux.sh b/scripts/internal/launch_thesis_tmux.sh index e241b94..22d0750 100755 --- a/scripts/internal/launch_thesis_tmux.sh +++ b/scripts/internal/launch_thesis_tmux.sh @@ -258,11 +258,11 @@ EOF chmod +x "$runner_script" -status_cmd=( - bash - -lc - "cd $(printf "%q" "$REPO_ROOT"); while true; do clear; date; echo; uv run python scripts/internal/run_experiments.py status; sleep $(printf "%q" "$STATUS_INTERVAL"); done" -) +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")" diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index 72a6f1a..831b7f5 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -13,7 +13,8 @@ uv run python scripts/internal/run_experiments.py run \ --experiment-class core_ssl_benchmark --dry-run \ --project slices-thesis --revision thesis-v1 --entity - uv run python scripts/internal/run_experiments.py status --experiment-class core_ssl_benchmark + 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 \ --experiment-class core_ssl_benchmark --revision thesis-v1 \ --project slices-thesis --entity @@ -1470,19 +1471,26 @@ def _extract_experiment_class_from_id(run_id: str) -> str | None: return None -def print_status(experiment_class_filter: list[str] | None = None) -> 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() 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, None), []).append(run.id) + 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) @@ -1684,7 +1692,7 @@ def cmd_run(args): def cmd_status(args): - print_status(args.experiment_class) + print_status(args.experiment_class, args.revision) def cmd_retry(args): @@ -1828,6 +1836,7 @@ def main() -> None: 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") @@ -1875,8 +1884,24 @@ def main() -> None: if getattr(args, "reason", None) and not getattr(args, "revision", None): parser.error("--reason requires --revision") - if args.command == "retry" and args.revision and not args.experiment_class: - parser.error("--revision requires --experiment-class to scope which classes to revise") + 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 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.failed and not args.skipped: + parser.error("retry requires --failed and/or --skipped") if args.command == "run": exit_code = cmd_run(args) diff --git a/scripts/training/finetune.py b/scripts/training/finetune.py index 0a8ad8a..49b987c 100644 --- a/scripts/training/finetune.py +++ b/scripts/training/finetune.py @@ -40,6 +40,7 @@ run_fairness_evaluation, setup_finetune_callbacks, setup_wandb_logger, + train_label_support_summary, validate_data_prerequisites, ) @@ -189,7 +190,7 @@ def main(cfg: DictConfig) -> None: f"({(1 - stats.get('prevalence', 0))*100:.1f}%)" ) - report_and_validate_train_label_support( + train_support_stats = report_and_validate_train_label_support( datamodule=datamodule, task_name=task_name, task_type=task_type, @@ -262,6 +263,8 @@ def main(cfg: DictConfig) -> None: callbacks = setup_finetune_callbacks(cfg, checkpoint_prefix="finetune") logger = setup_wandb_logger(cfg) + if logger: + logger.experiment.summary.update(train_label_support_summary(train_support_stats)) trainer = pl.Trainer( max_epochs=cfg.training.max_epochs, diff --git a/scripts/training/supervised.py b/scripts/training/supervised.py index e0142c2..0aa47f0 100644 --- a/scripts/training/supervised.py +++ b/scripts/training/supervised.py @@ -40,6 +40,7 @@ save_encoder_checkpoint, setup_finetune_callbacks, setup_wandb_logger, + train_label_support_summary, validate_data_prerequisites, ) @@ -229,7 +230,7 @@ def main(cfg: DictConfig) -> None: f"({(1 - stats.get('prevalence', 0))*100:.1f}%)" ) - report_and_validate_train_label_support( + train_support_stats = report_and_validate_train_label_support( datamodule=datamodule, task_name=task_name, task_type=task_type, @@ -300,6 +301,8 @@ def main(cfg: DictConfig) -> None: callbacks = setup_finetune_callbacks(cfg, checkpoint_prefix="supervised") logger = setup_wandb_logger(cfg) + if logger: + logger.experiment.summary.update(train_label_support_summary(train_support_stats)) trainer = pl.Trainer( max_epochs=cfg.training.max_epochs, diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index 0341a43..e5354b4 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -96,6 +96,8 @@ def _build_wandb_tags(cfg: DictConfig) -> list[str] | None: 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: @@ -134,6 +136,7 @@ def main(cfg: DictConfig) -> None: from slices.training.utils import ( report_and_validate_train_label_support, + train_label_support_summary, validate_data_prerequisites, ) @@ -163,7 +166,7 @@ def main(cfg: DictConfig) -> None: print(f" Val: {len(datamodule.val_indices)} stays") print(f" Test: {len(datamodule.test_indices)} stays") - report_and_validate_train_label_support( + train_support_stats = report_and_validate_train_label_support( datamodule=datamodule, task_name=task_name, task_type=task_type, @@ -340,6 +343,7 @@ def main(cfg: DictConfig) -> None: tags=_build_wandb_tags(cfg), config=OmegaConf.to_container(cfg, resolve=True), ) + run.summary.update(train_label_support_summary(train_support_stats)) run.summary.update(metrics) if fairness_report: run.summary.update(flatten_fairness_report(fairness_report)) diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 9e1be8b..ef4bbeb 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -20,6 +20,12 @@ 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, +) from slices.data.labels import LabelBuilder, LabelBuilderFactory, LabelConfig # ============================================================================= @@ -272,6 +278,62 @@ def _add_wandb_tag(tags: list[str], tag: str | None) -> None: 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 == "mask_ratio_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. @@ -332,7 +394,12 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: _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 @@ -340,6 +407,7 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: 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: @@ -352,6 +420,7 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: 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: @@ -557,6 +626,50 @@ class weights do not apply. Other task types fail closed instead of 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, @@ -567,7 +680,8 @@ def validate_data_prerequisites( 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. + 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 @@ -593,6 +707,13 @@ def validate_data_prerequisites( 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( @@ -649,17 +770,13 @@ def get_training_tasks_path() -> Path: # Validate label manifest if task configs are provided or can be resolved. if resolved_task_configs: - metadata_path = path / "metadata.yaml" - if not metadata_path.exists(): + 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}" ) - with open(metadata_path) as f: - metadata = yaml.safe_load(f) - label_manifest = metadata.get("label_manifest") if label_manifest is None: raise RuntimeError( diff --git a/tests/test_evaluate_fairness.py b/tests/test_evaluate_fairness.py index 1c7969c..0930e41 100644 --- a/tests/test_evaluate_fairness.py +++ b/tests/test_evaluate_fairness.py @@ -55,6 +55,7 @@ def test_default_experiment_classes_include_downstream_families(): 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", @@ -64,6 +65,7 @@ def test_parse_args_exposes_allow_incomplete_escape_hatch(monkeypatch): 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): @@ -148,6 +150,19 @@ def test_has_fairness_metrics_rejects_partial_binary_summaries(): 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") + 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_requires_regression_metric_family(): mod = importlib.import_module("scripts.eval.evaluate_fairness") @@ -221,7 +236,7 @@ def test_missing_fairness_report_requirements_ignores_eicu_race(): assert missing == [] -def test_missing_fairness_report_requirements_rejects_nan_required_metric(): +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"}}) @@ -240,4 +255,13 @@ def test_missing_fairness_report_requirements_rejects_nan_required_metric(): missing = mod.missing_fairness_report_requirements(run, report, ["gender"]) - assert "gender: missing worst_group_auroc" in missing + 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] diff --git a/tests/test_export_results.py b/tests/test_export_results.py index fa91b8c..5a094df 100644 --- a/tests/test_export_results.py +++ b/tests/test_export_results.py @@ -496,7 +496,7 @@ def test_filter_thesis_tasks_excludes_optional_mortality_by_default(): df = pd.DataFrame( [ {"task": "mortality_24h", "wandb_run_id": "main"}, - {"task": "mortality", "wandb_run_id": "extension"}, + {"task": "sepsis", "wandb_run_id": "extension"}, ] ) @@ -526,11 +526,13 @@ def test_parse_args_uses_revision_env_when_cli_omits_it(monkeypatch): 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): diff --git a/tests/test_fixes.py b/tests/test_fixes.py index b5854e4..25c1563 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -225,6 +225,30 @@ def test_validate_data_prerequisites_matching_passes(self, tmp_path): # 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 @@ -1883,6 +1907,62 @@ def test_revised_status_handles_base_and_revision_groups(self, monkeypatch, caps 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_cmd_run_respects_requested_class_order(self, monkeypatch): import scripts.internal.run_experiments as runner @@ -2069,11 +2149,14 @@ def __init__(self, **kwargs): assert isinstance(logger, DummyWandbLogger) assert captured["kwargs"]["name"] == ( - "capacity_study_finetune_miiv_mae_mortality_24h_seed42_medium_frac01" + "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 captured["kwargs"]["group"] == "finetune_miiv_mae_mortality_24h_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"] From fba0e7b82fc68027d48bfcce70247a909a8e8188 Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 22 Apr 2026 16:58:55 -0400 Subject: [PATCH 104/121] Clarify controlled SSL masking semantics Update README, experiment plan, and SSL config comments to distinguish the shared mask budget and interface from objective-specific mask geometry. Document JEPA block masking and contrastive two-view masking so the benchmark contract matches the current experiment design. No behavioral, benchmark, or migration impact. --- README.md | 8 +++++--- configs/ssl/contrastive.yaml | 3 ++- configs/ssl/jepa.yaml | 7 +++++-- docs/internal/EXPERIMENT_PLAN.md | 6 ++++++ 4 files changed, 18 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index eb7c043..1b509f0 100644 --- a/README.md +++ b/README.md @@ -29,7 +29,7 @@ and model-capacity studies. | 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 | @@ -42,7 +42,9 @@ Fair comparison of SSL objectives for clinical time series is currently impossib ## SSL Paradigms The controlled SSL objectives share the same timestep-level obs-aware Transformer -encoder and differ only in the SSL objective and masking logic: +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 | |---|---|---|---| @@ -268,7 +270,7 @@ uv run 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. - **Patient-level splits**: No data leakage between train/val/test. - **Config-driven ablations**: Change one YAML default to switch paradigm, encoder, or task. diff --git a/configs/ssl/contrastive.yaml b/configs/ssl/contrastive.yaml index b7b96ba..e8ef1ff 100644 --- a/configs/ssl/contrastive.yaml +++ b/configs/ssl/contrastive.yaml @@ -13,7 +13,8 @@ name: contrastive # Mode: "temporal" (per-timestep overlap pairs) or "instance" (mean-pool, SimCLR-style) mode: instance -# Masking parameters (same as MAE/JEPA for fair comparison) +# 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 diff --git a/configs/ssl/jepa.yaml b/configs/ssl/jepa.yaml index 53ca21b..3c0031a 100644 --- a/configs/ssl/jepa.yaml +++ b/configs/ssl/jepa.yaml @@ -1,8 +1,9 @@ # @package ssl # JEPA (Joint-Embedding Predictive Architecture) SSL Configuration # -# Timestep-level tokenization: same masking interface as MAE, but predicts -# latent timestep representations (d_model vectors) instead of raw features. +# 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) @@ -12,6 +13,8 @@ name: jepa # 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) diff --git a/docs/internal/EXPERIMENT_PLAN.md b/docs/internal/EXPERIMENT_PLAN.md index 18ff862..e3af9d4 100644 --- a/docs/internal/EXPERIMENT_PLAN.md +++ b/docs/internal/EXPERIMENT_PLAN.md @@ -27,6 +27,12 @@ Controlled SSL objectives share: - 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. From 32b406e0b03bb349a679e81839f37d908d7ef0ba Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 22 Apr 2026 18:13:36 -0400 Subject: [PATCH 105/121] Harden thesis export and masking checks Fix JEPA block masking to honor the requested valid-timestep budget, stabilize SMART external-reference model size identity, and fail closed on missing fairness summaries or accidental multi-revision exports. Add subgroup comparability accounting for one-class fairness groups, document the export semantics, and add a root experiment-plan pointer. Behavioral implications: JEPA high-mask robustness now matches configured masking pressure, SMART appendix rows export as model_size=default, publication export now requires fairness metrics, and multi-revision exports require an explicit exploratory flag. --- EXPERIMENTS_PLAN.md | 6 + docs/EVAL_FAIRNESS.md | 11 +- scripts/eval/evaluate_fairness.py | 1 + scripts/export_results.py | 180 +++++++++++++++++++++++ scripts/internal/run_experiments.py | 10 +- src/slices/eval/fairness_evaluator.py | 42 ++++-- src/slices/models/pretraining/masking.py | 139 +++++++++-------- tests/test_evaluate_fairness.py | 4 + tests/test_export_results.py | 122 +++++++++++++++ tests/test_fairness_evaluator.py | 29 ++++ tests/test_fixes.py | 15 ++ tests/test_jepa_objective.py | 52 +++++++ 12 files changed, 531 insertions(+), 80 deletions(-) create mode 100644 EXPERIMENTS_PLAN.md 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/docs/EVAL_FAIRNESS.md b/docs/EVAL_FAIRNESS.md index cab93cd..0cabbbd 100644 --- a/docs/EVAL_FAIRNESS.md +++ b/docs/EVAL_FAIRNESS.md @@ -64,9 +64,16 @@ For binary tasks it logs: - `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` @@ -97,7 +104,9 @@ For regression tasks it logs: - `results/statistical_tests.parquet` Those exports include fairness columns when the runs already have `fairness/*` -summary metrics. +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. ## Design Decisions diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index 71d7d54..9067160 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -78,6 +78,7 @@ } BINARY_FAIRNESS_REQUIRED_METRICS = [ "n_valid_groups", + "n_metric_valid_groups", "worst_group_auroc", "worst_group_auprc", "auroc_gap", diff --git a/scripts/export_results.py b/scripts/export_results.py index 6319f7b..c1006b3 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -118,6 +118,7 @@ "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", @@ -130,6 +131,7 @@ "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", @@ -142,6 +144,7 @@ "fairness/race/mse_gap", "fairness/race/mae_gap", "fairness/race/n_valid_groups", + "fairness/race/n_metric_valid_groups", ] VAL_METRICS = [ @@ -173,6 +176,26 @@ } EVAL_PHASES = ["finetune", "supervised", "baseline"] +FAIRNESS_ATTRIBUTES = ["gender", "age_group", "race"] +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", +] # --------------------------------------------------------------------------- @@ -355,6 +378,17 @@ 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: @@ -415,6 +449,9 @@ def extract_run(run, metric_keys: list[str]) -> dict: 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) @@ -447,6 +484,7 @@ def extract_run(run, metric_keys: list[str]) -> dict: "_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), } row.update(metrics) return row @@ -1117,6 +1155,105 @@ def _checkpoint_provenance_issues(per_seed_df: pd.DataFrame) -> list[tuple[pd.Se 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 _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 _expected_row_from_run(run) -> dict: return { "experiment_class": run.experiment_class, @@ -1335,6 +1472,36 @@ def validate( + 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}" + ) + + warnings.extend(_mixed_revision_group_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) @@ -1384,6 +1551,14 @@ def parse_args() -> argparse.Namespace: 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( @@ -1436,6 +1611,11 @@ def parse_args() -> argparse.Namespace: "--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 diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index 831b7f5..1d98c74 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -776,10 +776,16 @@ 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 + experiment_class, + "smart", + dataset, + seed, + pretrain_extra, + model_size=model_size, ) for task in TASKS: self._add_finetune( @@ -791,6 +797,7 @@ def build_smart_external_reference(self) -> None: False, pretrain, extra=finetune_extra, + model_size=model_size, ) self._add_finetune( experiment_class, @@ -801,6 +808,7 @@ def build_smart_external_reference(self) -> None: True, pretrain, extra=finetune_extra, + model_size=model_size, ) def build_all(self) -> list[Run]: diff --git a/src/slices/eval/fairness_evaluator.py b/src/slices/eval/fairness_evaluator.py index 65d742b..c2d99b0 100644 --- a/src/slices/eval/fairness_evaluator.py +++ b/src/slices/eval/fairness_evaluator.py @@ -361,35 +361,49 @@ def _evaluate_binary( 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: @@ -412,6 +426,8 @@ 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, } diff --git a/src/slices/models/pretraining/masking.py b/src/slices/models/pretraining/masking.py index 86817d0..bc860fe 100644 --- a/src/slices/models/pretraining/masking.py +++ b/src/slices/models/pretraining/masking.py @@ -171,6 +171,45 @@ def create_complementary_timestep_masks( 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, @@ -181,14 +220,10 @@ def create_block_timestep_mask( ) -> torch.Tensor: """Create contiguous block mask at timestep level. - Masks n_blocks contiguous segments that together cover approximately - mask_ratio of the sequence. Forces the model to predict from distant - context rather than interpolate from adjacent visible timesteps. - - Strategy: divides the sequence into n_blocks equal zones, then places one - randomly-sized masked block within each zone. Block sizes are drawn from a - Dirichlet-like split of the total masked budget. Fully vectorized — no - Python loops over batch elements. + 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. @@ -202,69 +237,43 @@ def create_block_timestep_mask( Returns: ssl_mask: (B, T) bool mask, True = visible, False = masked. """ - n_masked_total = max(int(n_timesteps * mask_ratio), 1) - n_masked_total = min(n_masked_total, n_timesteps - 1) - - # Split total masked budget into n_blocks random lengths per sample. - # Use Dirichlet-like splitting: draw n_blocks uniform values, normalize, - # then scale to sum to n_masked_total. Add 1 to each to ensure min length 1. - raw = torch.rand(batch_size, n_blocks) # CPU for speed - raw = raw / raw.sum(dim=1, keepdim=True) # normalize to sum=1 - # Reserve 1 per block, distribute the rest proportionally - extra = n_masked_total - n_blocks - if extra > 0: - block_lengths = 1 + (raw * extra).int() - # Fix rounding: adjust last block to hit exact total - block_lengths[:, -1] = n_masked_total - block_lengths[:, :-1].sum(dim=1) - else: - # Edge case: fewer masked timesteps than blocks — give 1 to first blocks - block_lengths = torch.zeros(batch_size, n_blocks, dtype=torch.int) - block_lengths[:, :n_masked_total] = 1 - - # Clamp to valid range - block_lengths = block_lengths.clamp(min=0, max=n_timesteps) - - # Divide sequence into n_blocks equal zones. Each block is placed randomly - # within its zone, avoiding cross-zone overlap by construction. - zone_size = n_timesteps // n_blocks - zone_starts = torch.arange(n_blocks) * zone_size # (n_blocks,) - - # Random offset within each zone (vectorized over batch) - # Max offset = zone_size - block_length (so block fits within zone) - max_offsets = zone_size - block_lengths # (B, n_blocks) - max_offsets = max_offsets.clamp(min=0) - offsets = (torch.rand(batch_size, n_blocks) * (max_offsets.float() + 1)).int() - offsets = offsets.clamp(max=max_offsets) - - # Compute absolute start positions: zone_start + offset - starts = zone_starts.unsqueeze(0) + offsets # (B, n_blocks) - - # Build mask: create time indices and compare with block ranges - # (B, n_blocks, T) — True where timestep t falls in block k - t = torch.arange(n_timesteps).unsqueeze(0).unsqueeze(0) # (1, 1, T) - starts_3d = starts.unsqueeze(2) # (B, n_blocks, 1) - ends_3d = (starts + block_lengths).unsqueeze(2) # (B, n_blocks, 1) - in_block = (t >= starts_3d) & (t < ends_3d) # (B, n_blocks, T) - - # Union across blocks -> masked positions - masked = in_block.any(dim=1) # (B, T) - ssl_mask = (~masked).to(device=device) - 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 = ssl_mask | (~valid_timestep_mask) + ssl_mask = torch.ones(batch_size, n_timesteps, dtype=torch.bool, device=device) + requested_blocks = max(int(n_blocks), 1) - n_visible = (ssl_mask & valid_timestep_mask).sum(dim=1) - has_valid = valid_timestep_mask.any(dim=1) - needs_fix = (n_visible == 0) & has_valid - if needs_fix.any(): - for b in range(batch_size): - if needs_fix[b]: - first_valid = valid_timestep_mask[b].nonzero(as_tuple=True)[0][0] - ssl_mask[b, first_valid] = True + 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 diff --git a/tests/test_evaluate_fairness.py b/tests/test_evaluate_fairness.py index 0930e41..4871d58 100644 --- a/tests/test_evaluate_fairness.py +++ b/tests/test_evaluate_fairness.py @@ -11,6 +11,7 @@ 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, @@ -199,6 +200,7 @@ def test_missing_fairness_report_requirements_flags_missing_requested_attribute( 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, @@ -221,6 +223,7 @@ def test_missing_fairness_report_requirements_ignores_eicu_race(): 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, @@ -243,6 +246,7 @@ def test_missing_fairness_report_requirements_accepts_written_nan_required_metri 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, diff --git a/tests/test_export_results.py b/tests/test_export_results.py index 5a094df..fbee38d 100644 --- a/tests/test_export_results.py +++ b/tests/test_export_results.py @@ -1,6 +1,7 @@ """Tests for the class-based results export pipeline.""" import importlib +import json import sys import pandas as pd @@ -432,6 +433,11 @@ def test_extract_run_exports_los_regression_fairness_keys(): 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( [ @@ -475,6 +481,7 @@ def test_validate_warns_on_missing_or_failed_checkpoint_provenance(): "seed": 42, "phase": "baseline", "test/auprc": 0.4, + mod.FAIRNESS_SUMMARY_KEY_COLUMN: json.dumps(fairness_keys), }, ] ) @@ -490,6 +497,74 @@ def test_validate_warns_on_missing_or_failed_checkpoint_provenance(): 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_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") @@ -571,6 +646,53 @@ def test_parse_args_exposes_duplicate_fingerprint_escape_hatch(monkeypatch): 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_main_exits_nonzero_when_no_runs_match(monkeypatch, tmp_path): mod = importlib.import_module("scripts.export_results") diff --git a/tests/test_fairness_evaluator.py b/tests/test_fairness_evaluator.py index 50b38d1..37237a0 100644 --- a/tests/test_fairness_evaluator.py +++ b/tests/test_fairness_evaluator.py @@ -122,6 +122,34 @@ 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.""" @@ -355,6 +383,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): diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 25c1563..79a305d 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -1458,6 +1458,12 @@ def test_infer_model_size_for_capacity_variants(self): 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.""" @@ -3253,3 +3259,12 @@ def test_ts2vec_extension_includes_both_protocols(self): 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_jepa_objective.py b/tests/test_jepa_objective.py index 885a787..d4399e4 100644 --- a/tests/test_jepa_objective.py +++ b/tests/test_jepa_objective.py @@ -13,6 +13,58 @@ 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]) + + # ============================================================================= # Predictor tests # ============================================================================= From 02d6c9dff9f1485502349ddb70587f971a54ae69 Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 22 Apr 2026 21:49:29 -0400 Subject: [PATCH 106/121] Prepare benchmark release exports Relabel the contrastive robustness slice, expand default binary metrics, add aggregate confidence intervals, and fix LOS label and FLOPs reporting helpers. Benchmark implications: final exports now include threshold metrics and seed-level CIs, contrastive HP rows are no longer described as pure mask-ratio sweeps, and compute/label helper outputs are aligned with the 84-feature thesis corpus. --- README.md | 21 +++++ configs/eval/default.yaml | 6 +- docs/internal/EXPERIMENT_PLAN.md | 25 +++++- scripts/analyze_labels.py | 133 ++++++++++++++++++++++++++-- scripts/export_results.py | 32 ++++++- scripts/internal/run_experiments.py | 9 +- scripts/measure_flops.py | 37 ++++++-- src/slices/eval/README.md | 10 +-- src/slices/eval/metrics.py | 14 ++- src/slices/training/README.md | 5 ++ src/slices/training/utils.py | 2 +- tests/test_analyze_labels.py | 17 ++++ tests/test_export_results.py | 3 + tests/test_fixes.py | 16 +++- tests/test_metrics.py | 17 +++- 15 files changed, 311 insertions(+), 36 deletions(-) diff --git a/README.md b/README.md index 1b509f0..148ca8a 100644 --- a/README.md +++ b/README.md @@ -155,6 +155,26 @@ uv run python scripts/training/supervised.py dataset=miiv 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 the class-based launcher 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 ``` @@ -271,6 +291,7 @@ uv run mypy src/ - **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. - **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. diff --git a/configs/eval/default.yaml b/configs/eval/default.yaml index a042712..6a481cc 100644 --- a/configs/eval/default.yaml +++ b/configs/eval/default.yaml @@ -21,7 +21,11 @@ # 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 diff --git a/docs/internal/EXPERIMENT_PLAN.md b/docs/internal/EXPERIMENT_PLAN.md index e3af9d4..a1438b1 100644 --- a/docs/internal/EXPERIMENT_PLAN.md +++ b/docs/internal/EXPERIMENT_PLAN.md @@ -101,7 +101,7 @@ thesis matrix. | `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 | 840 | | `cross_dataset_transfer` | MIMIC-IV to eICU and eICU to MIMIC-IV; SSL Protocol B; 5 seeds | 120 | -| `hp_robustness` | LR and mask-ratio robustness on MIMIC-IV `mortality_24h`; 5 seeds | 150 | +| `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 | 360 | | `ts2vec_extension` | TS2Vec temporal contrastive extension; 3 datasets; 4 tasks; 5 seeds; Protocol A/B | 135 | @@ -143,9 +143,15 @@ uv run python scripts/internal/run_experiments.py run \ 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 @@ -161,6 +167,7 @@ uv run python scripts/internal/run_experiments.py run \ smart_external_reference \ --project slices-thesis \ --revision thesis-v1 \ + --launch-commit \ --entity dummy \ --dry-run ``` @@ -207,6 +214,21 @@ uv run python scripts/export_results.py \ 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`, or `--allow-duplicate-fingerprints` 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 @@ -225,6 +247,7 @@ Before launching final runs: - 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: diff --git a/scripts/analyze_labels.py b/scripts/analyze_labels.py index 8223476..d9bc373 100644 --- a/scripts/analyze_labels.py +++ b/scripts/analyze_labels.py @@ -4,6 +4,7 @@ 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 @@ -69,16 +70,19 @@ def load_metadata(processed_dir: Path) -> Dict[str, Any]: 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.""" +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: - config = yaml.safe_load(f) or {} + return yaml.safe_load(f) or {} - return config.get("quality_checks", {}) 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]: @@ -192,9 +196,44 @@ 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. @@ -230,10 +269,31 @@ def analyze_task( } 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) @@ -287,6 +347,7 @@ def analyze_task( 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, @@ -337,7 +398,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] = { @@ -347,8 +411,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 = {} @@ -363,6 +440,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, } @@ -412,13 +490,24 @@ 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:") for class_val, class_stats in stats["class_distribution"].items(): @@ -446,6 +535,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}") @@ -630,10 +741,12 @@ def main(): print("=" * 60) for task_name in task_columns: + task_config = load_task_config(task_name) stats = analyze_task( labels_df, task_name, - quality_checks=load_task_quality_checks(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) @@ -648,10 +761,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/export_results.py b/scripts/export_results.py index c1006b3..16252c5 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -12,6 +12,9 @@ - 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 @@ -33,6 +36,7 @@ import pandas as pd import wandb +from scipy import stats as scipy_stats from slices.eval.statistical import ( bonferroni_correction, cohens_d, @@ -353,7 +357,11 @@ def _recover_pretrain_metadata( 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", "mask_ablation"}: + 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 @@ -629,6 +637,23 @@ def filter_thesis_tasks( 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 @@ -662,11 +687,16 @@ def _aggregate_group(df: pd.DataFrame, fingerprint_cols: list[str]) -> pd.DataFr 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) diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index 1d98c74..85bd8e9 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -109,6 +109,7 @@ 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") @@ -667,9 +668,15 @@ def build_hp_robustness(self) -> None: ) 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 = "mask_ratio_sensitivity" + subtype = VIEW_MASK_EXPERIMENT_SUBTYPE pretrain = self._add_pretrain( experiment_class, paradigm, diff --git a/scripts/measure_flops.py b/scripts/measure_flops.py index 5a700f8..e638c24 100644 --- a/scripts/measure_flops.py +++ b/scripts/measure_flops.py @@ -7,7 +7,7 @@ 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 35 # override feature dim + 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 """ @@ -26,8 +26,10 @@ # Paradigm configs — mirror the YAML defaults used in actual runs # --------------------------------------------------------------------------- +DEFAULT_N_FEATURES = 84 + ENCODER_CONFIG = { - "d_input": 35, + "d_input": DEFAULT_N_FEATURES, "d_model": 64, "max_seq_length": SEQ_LENGTH_HOURS, "n_layers": 2, @@ -134,7 +136,7 @@ def format_params(n: int) -> str: return str(n) -def query_wandb_runs(entity: str, project: str) -> Dict[str, dict]: +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}. @@ -144,8 +146,11 @@ def query_wandb_runs(entity: str, project: str) -> Dict[str, dict]: api = wandb.Api(timeout=300) path = f"{entity}/{project}" if entity else project - # Server-side filter: only pretrain runs (tagged phase:pretrain) - filters = {"tags": {"$all": ["phase:pretrain"]}, "state": "finished"} + # 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] = {} @@ -190,7 +195,10 @@ def query_wandb_runs(entity: str, project: str) -> Dict[str, dict]: def main(): parser = argparse.ArgumentParser(description="Measure FLOPs per SSL paradigm") parser.add_argument( - "--n-features", type=int, default=35, help="Number of input features (d_input)" + "--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)" @@ -206,9 +214,15 @@ def main(): parser.add_argument( "--wandb-project", type=str, - default=os.environ.get("WANDB_PROJECT", "SLICES"), + 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 @@ -239,7 +253,14 @@ def main(): wandb_data = {} if not args.no_wandb: try: - wandb_data = query_wandb_runs(args.wandb_entity, args.wandb_project) + 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: diff --git a/src/slices/eval/README.md b/src/slices/eval/README.md index 7ee5505..fc8c19f 100644 --- a/src/slices/eval/README.md +++ b/src/slices/eval/README.md @@ -20,11 +20,7 @@ inference helpers, and SSL imputation checks. ```python from slices.eval import MetricConfig, build_metrics -config = MetricConfig( - task_type="binary", - metrics=["auroc", "auprc", "brier_score", "ece"], - threshold=0.5, -) +config = MetricConfig(task_type="binary", threshold=0.5) val_metrics = build_metrics(config, prefix="val") ``` @@ -36,9 +32,9 @@ Supported task types: - `multilabel`: `auroc`, `auprc`, `accuracy`, `f1` - `regression`: `mse`, `mae`, `rmse`, `r2` -Default metric sets are intentionally small: +Default metric sets match the public benchmark export surface: -- 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/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/training/README.md b/src/slices/training/README.md index b1cfde6..25eed9a 100644 --- a/src/slices/training/README.md +++ b/src/slices/training/README.md @@ -55,6 +55,11 @@ Checkpoint inputs: - full Lightning pretrain checkpoint via `pretrain_checkpoint` - no checkpoint for supervised-from-scratch runs +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. + What it handles: - linear probing via `training.freeze_encoder=true` diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index ef4bbeb..52aab91 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -309,7 +309,7 @@ def append_wandb_identity_suffixes(name: str | None, cfg: DictConfig) -> str | N 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 == "mask_ratio_sensitivity": + 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: diff --git a/tests/test_analyze_labels.py b/tests/test_analyze_labels.py index 7e7e0d9..a7bcf03 100644 --- a/tests/test_analyze_labels.py +++ b/tests/test_analyze_labels.py @@ -41,3 +41,20 @@ def test_analyze_task_has_no_quality_warning_below_threshold(): 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_export_results.py b/tests/test_export_results.py index fbee38d..3f227c0 100644 --- a/tests/test_export_results.py +++ b/tests/test_export_results.py @@ -368,6 +368,9 @@ def test_build_aggregated_df_records_per_metric_non_null_counts(): 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(): diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 79a305d..f719094 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -1441,6 +1441,20 @@ def test_new_runs_use_config_directly(self): 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 @@ -3227,7 +3241,7 @@ def test_matrix_counts_match_final_experiment_classes(self): for run in runs if run.experiment_class == "hp_robustness" and run.paradigm == "contrastive" - and run.experiment_subtype == "mask_ratio_sensitivity" + and run.experiment_subtype == "view_mask_sensitivity" and run.run_type == "pretrain" ] assert contrastive_mask_sweep diff --git a/tests/test_metrics.py b/tests/test_metrics.py index 468d43a..fddbdd6 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -341,11 +341,20 @@ 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_all_task_types_covered(self): """Test all task types have entries.""" From 9f2f79f58e8355cf2a9038ece09088cebdbe475f Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 22 Apr 2026 22:31:53 -0400 Subject: [PATCH 107/121] Fix mortality label handling for unknown eICU outcomes Preserve nullable hospital mortality outcomes, treat death-coded discharge locations as death evidence, and use discharge timing when it can place deaths inside mortality windows. Bumps the mortality label semantic version to invalidate stale processed labels; regenerated worktree artifacts update eICU, miiv, and combined manifests to 2.2.0. --- scripts/preprocessing/extract_with_ricu.R | 143 +++++++++++++---- src/slices/data/labels/mortality.py | 187 ++++++++++++++++++---- tests/test_task_builders.py | 139 ++++++++++++++++ 3 files changed, 401 insertions(+), 68 deletions(-) diff --git a/scripts/preprocessing/extract_with_ricu.R b/scripts/preprocessing/extract_with_ricu.R index 2bbda39..360cad5 100644 --- a/scripts/preprocessing/extract_with_ricu.R +++ b/scripts/preprocessing/extract_with_ricu.R @@ -775,54 +775,125 @@ extract_mortality_eicu <- function(dataset, dict) { names(pat)[1] } - # eICU uses status strings and offsets rather than timestamps. - # We derive date_of_death from hospitaldischargeoffset for expired patients, - # using the same synthetic epoch as the stays extraction so timelines are - # consistent (intime = epoch, outtime = epoch + unitdischargeoffset*60). - 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)) + } } - # Derive date_of_death from hospitaldischargeoffset for expired patients. - # eICU does not store death timestamps directly, but hospitaldischargeoffset - # gives minutes from unit admission to hospital discharge. For patients who - # expired (hospitaldischargestatus == "Expired"), this is the approximate - # time of death. 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. + 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") - has_offset <- "hospitaldischargeoffset" %in% names(pat) - expired <- !is.na(hosp_flag) & hosp_flag == 1L + 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") - if (has_offset) { - dod[expired] <- epoch + pat$hospitaldischargeoffset[expired] * 60 - } - - # Also derive dischtime from hospitaldischargeoffset for all patients + 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") - if (has_offset) { - dischtime <- epoch + pat$hospitaldischargeoffset * 60 - } + 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 = dod, hospital_expire_flag = hosp_flag, dischtime = dischtime, - discharge_location = if ("unitdischargelocation" %in% names(pat)) { - pat$unitdischargelocation - } else { - NA_character_ - }, + discharge_location = discharge_location, # Precision-aware schema: eICU timestamps are derived from offsets - # (minute-level precision) — these are true timestamps, not date-only. + # (minute-level precision) when an offset is available. death_time = dod, death_date = as.Date(rep(NA, nrow(pat))), - death_time_precision = ifelse(expired, "timestamp", NA_character_), - death_source = ifelse(expired, "derived", NA_character_), + death_time_precision = ifelse( + has_death_offset, + "timestamp", + ifelse(death_evidence, "unknown", NA_character_) + ), + death_source = death_source, stringsAsFactors = FALSE ) df @@ -1035,8 +1106,10 @@ main <- function(opts) { } else if (!is.null(deprecated_seq_length_hours) && raw_export_horizon_hours != deprecated_seq_length_hours) { stop( - "Conflicting values provided for --raw_export_horizon_hours and " - "--seq_length_hours. Use only --raw_export_horizon_hours." + paste0( + "Conflicting values provided for --raw_export_horizon_hours and ", + "--seq_length_hours. Use only --raw_export_horizon_hours." + ) ) } diff --git a/src/slices/data/labels/mortality.py b/src/slices/data/labels/mortality.py index 30371cb..d2f96aa 100644 --- a/src/slices/data/labels/mortality.py +++ b/src/slices/data/labels/mortality.py @@ -34,10 +34,11 @@ class MortalityLabelBuilder(LabelBuilder): - "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: falls back to hospital_expire_flag + - "unknown" or missing: uses explicit death evidence and preserves + unknown outcomes as null """ - SEMANTIC_VERSION = "2.1.2" + SEMANTIC_VERSION = "2.2.0" def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: """Build mortality labels from stay and mortality data. @@ -73,7 +74,8 @@ 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 exist (backward compat with legacy data) + # 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 @@ -106,7 +108,13 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: labels = merged.select( [ "stay_id", - pl.col("hospital_expire_flag").cast(pl.Int32).alias("label"), + pl.when(self._death_evidence_expr()) + .then(1) + .when(self._known_survivor_expr()) + .then(0) + .otherwise(None) + .cast(pl.Int32) + .alias("label"), ] ) @@ -134,6 +142,7 @@ def build_labels(self, raw_data: Dict[str, pl.DataFrame]) -> pl.DataFrame: 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. @@ -149,6 +158,14 @@ def _ensure_precision_columns(merged: pl.DataFrame) -> pl.DataFrame: merged = merged.with_columns( pl.col("death_time").cast(pl.Datetime("us")).alias("death_time") ) + else: + 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 @@ -198,6 +215,40 @@ def _ensure_precision_columns(merged: pl.DataFrame) -> pl.DataFrame: pl.lit("legacy").alias("death_source"), ) + @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. @@ -220,15 +271,66 @@ def _normalize_datetime_columns(merged: pl.DataFrame) -> pl.DataFrame: 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 _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)) + ) + 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._effective_death_time_expr() + effective_precision = self._effective_death_precision_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 = pl.col("death_time").is_not_null() & (pl.col("death_time") < obs_end) + 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 @@ -241,9 +343,9 @@ def _build_hospital_mortality_with_obs( ) died_during_obs = ( - pl.when(pl.col("death_time_precision") == "timestamp") + pl.when(effective_precision == "timestamp") .then(ts_died_during_obs) - .when(pl.col("death_time_precision") == "date") + .when(effective_precision == "date") .then(date_died_during_obs | date_obs_boundary) .otherwise(pl.lit(False)) ) @@ -253,9 +355,11 @@ def _build_hospital_mortality_with_obs( "stay_id", pl.when(died_during_obs) .then(None) - .when(pl.col("hospital_expire_flag") == 1) + .when(self._death_evidence_expr()) .then(1) - .otherwise(0) + .when(self._known_survivor_expr()) + .then(0) + .otherwise(None) .cast(pl.Int32) .alias("label"), ] @@ -263,31 +367,43 @@ def _build_hospital_mortality_with_obs( def _build_icu_mortality(self, merged: pl.DataFrame) -> pl.DataFrame: """ICU mortality (died during or at ICU discharge), no observation window.""" - ts_died = pl.col("death_time").is_not_null() & (pl.col("death_time") <= pl.col("outtime")) + 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(pl.col("death_time_precision") == "timestamp") + pl.when(effective_precision == "timestamp") .then(ts_died) - .when(pl.col("death_time_precision") == "date") + .when(effective_precision == "date") .then(date_died) - .otherwise(pl.lit(False)) + .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).otherwise(0).cast(pl.Int32).alias("label"), + 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 = pl.col("death_time").is_not_null() & (pl.col("death_time") <= boundary) + 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() & ( ( @@ -309,9 +425,9 @@ def _build_legacy_windowed(self, merged: pl.DataFrame, window_hours: int) -> pl. ) died = ( - pl.when(pl.col("death_time_precision") == "timestamp") + pl.when(effective_precision == "timestamp") .then(ts_died) - .when(pl.col("death_time_precision") == "date") + .when(effective_precision == "date") .then( pl.when(date_died_definite) .then(pl.lit(True)) @@ -319,13 +435,15 @@ def _build_legacy_windowed(self, merged: pl.DataFrame, window_hours: int) -> pl. .then(pl.lit(None)) # null = ambiguous .otherwise(pl.lit(False)) ) - .otherwise(pl.lit(False)) + .otherwise(pl.lit(None)) ) return merged.select( [ "stay_id", - pl.when(died.is_null()) + pl.when(died.is_null() & self._known_survivor_expr()) + .then(0) + .when(died.is_null()) .then(None) .when(died) .then(1) @@ -352,10 +470,12 @@ def _build_windowed_mortality_labels( - "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: null for windowed tasks + - "unknown"/missing: true survivors remain 0, unknown outcomes are null """ prediction_start_hours = obs_hours + gap_hours until_icu_discharge = prediction_hours == -1 + effective_death_time = self._effective_death_time_expr() + effective_precision = self._effective_death_precision_expr() obs_end = pl.col("intime") + pl.duration(hours=obs_hours) pred_start = pl.col("intime") + pl.duration(hours=prediction_start_hours) @@ -367,20 +487,20 @@ def _build_windowed_mortality_labels( pred_end = pl.col("intime") + pl.duration(hours=prediction_end_hours) # --- Timestamp precision: exact comparisons --- - ts_died_during_obs = pl.col("death_time").is_not_null() & (pl.col("death_time") < obs_end) + ts_died_during_obs = effective_death_time.is_not_null() & (effective_death_time < obs_end) ts_died_during_gap = ( ( - pl.col("death_time").is_not_null() - & (pl.col("death_time") > obs_end) - & (pl.col("death_time") < pred_start) + 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 = ( - pl.col("death_time").is_not_null() - & (pl.col("death_time") >= pred_start) - & (pl.col("death_time") <= pred_end) + effective_death_time.is_not_null() + & (effective_death_time >= pred_start) + & (effective_death_time <= pred_end) ) # --- Date precision: interval-based conservative logic --- @@ -453,15 +573,16 @@ def _build_windowed_mortality_labels( # Combine by precision label_expr = ( - pl.when(pl.col("death_time_precision") == "timestamp") + pl.when(effective_precision == "timestamp") .then(ts_label) - .when(pl.col("death_time_precision") == "date") + .when(effective_precision == "date") .then(date_label) .otherwise( - # Unknown precision or no death info → fall back - pl.when(pl.col("hospital_expire_flag") == 1) - .then(None) # can't resolve timing → exclude - .otherwise(0) # not flagged as dead → survived + # Unknown death timing or unknown outcome → exclude. Explicit + # survivors remain negatives. + pl.when(self._known_survivor_expr()) + .then(0) + .otherwise(None) ) ) diff --git a/tests/test_task_builders.py b/tests/test_task_builders.py index f550d0a..ea5a768 100644 --- a/tests/test_task_builders.py +++ b/tests/test_task_builders.py @@ -546,6 +546,145 @@ 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, + ) -> pl.DataFrame: + return pl.DataFrame( + { + "stay_id": [1], + "death_time": [None], + "death_date": [None], + "death_time_precision": [None], + "death_source": [None], + "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_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). From 9f8d714a4ee0e504018e360e28f6ced5d9b7d04c Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 22 Apr 2026 22:26:55 -0400 Subject: [PATCH 108/121] Tighten final launch provenance Require launch commits and clean tracked worktrees for direct final run and retry commands while keeping dry-run audits convenient. Validate exported revision launch-commit homogeneity, surface skipped retry dependents, use locked tmux dependency sync, keep fairness/export jobs active for shutdown, and initialize XGBoost W&B runs before expensive setup. Final launch behavior now fails closed for missing or ambiguous provenance outside dry runs; retry --failed reports skipped downstream work and final exports warn on missing or mixed launch commits. --- scripts/export_results.py | 51 +++++++++ scripts/internal/auto_shutdown.sh | 4 +- scripts/internal/launch_thesis_tmux.sh | 2 +- scripts/internal/run_experiments.py | 153 ++++++++++++++++++++++++- scripts/training/xgboost_baseline.py | 57 +++++---- tests/test_export_results.py | 76 ++++++++++++ tests/test_fixes.py | 140 ++++++++++++++++++++++ 7 files changed, 453 insertions(+), 30 deletions(-) diff --git a/scripts/export_results.py b/scripts/export_results.py index 16252c5..a47c453 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -485,6 +485,7 @@ def extract_run(run, metric_keys: list[str]) -> dict: "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, @@ -1284,6 +1285,55 @@ def _mixed_revision_group_warnings(per_seed_df: pd.DataFrame) -> list[str]: 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, @@ -1531,6 +1581,7 @@ def validate( ) 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) diff --git a/scripts/internal/auto_shutdown.sh b/scripts/internal/auto_shutdown.sh index c30c24e..4b301af 100755 --- a/scripts/internal/auto_shutdown.sh +++ b/scripts/internal/auto_shutdown.sh @@ -8,8 +8,8 @@ set -euo pipefail IDLE_THRESHOLD_MIN=59 STAMP_FILE="/tmp/slices_last_training_activity" -# Check if any training process or Claude Code session is running -if pgrep -f "scripts/(training/(pretrain|finetune|supervised)|internal/run_experiments|preprocessing/prepare_dataset)\.py" > /dev/null 2>&1 \ +# Check if any final-run process or Claude Code 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; then # Activity detected — update timestamp and exit date +%s > "$STAMP_FILE" diff --git a/scripts/internal/launch_thesis_tmux.sh b/scripts/internal/launch_thesis_tmux.sh index 22d0750..2e83e3b 100755 --- a/scripts/internal/launch_thesis_tmux.sh +++ b/scripts/internal/launch_thesis_tmux.sh @@ -249,7 +249,7 @@ if [[ "$(printf "%q" "$SKIP_LAUNCH_GIT_CHECK")" != "1" ]]; then fi fi -uv sync --dev +uv sync --dev --locked ${warmup_line} ${main_line} ${appendix_block}${fairness_line} diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index 85bd8e9..159a3b4 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -15,9 +15,9 @@ --project slices-thesis --revision thesis-v1 --entity 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 \ + uv run python scripts/internal/run_experiments.py retry --failed --skipped \ --experiment-class core_ssl_benchmark --revision thesis-v1 \ - --project slices-thesis --entity + --project slices-thesis --entity --launch-commit """ from __future__ import annotations @@ -902,6 +902,81 @@ def apply_launch_commit(runs: list[Run], launch_commit: str | None = None) -> li 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 _validate_clean_final_launch_state(launch_commit: str) -> str | None: + """Return an error string if a final launch would have ambiguous provenance.""" + try: + resolved_launch_commit = _git_stdout( + ["rev-parse", "--verify", f"{launch_commit}^{{commit}}"] + ) + current_commit = _git_stdout(["rev-parse", "--verify", "HEAD"]) + if current_commit != resolved_launch_commit: + return ( + "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 ( + "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 f"could not validate final launch git provenance: {exc}" + return None + + +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, + ) + return + + if not launch_commit: + parser.error( + "final run/retry requires --launch-commit or SLICES_LAUNCH_COMMIT " "for W&B provenance" + ) + + error = _validate_clean_final_launch_state(str(launch_commit)) + if error: + parser.error(error) + + # --------------------------------------------------------------------------- # State Management # --------------------------------------------------------------------------- @@ -1602,6 +1677,33 @@ def _collect_dependency_closure( 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: @@ -1656,6 +1758,34 @@ def _fail_for_blocked_retry_dependencies( 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 _filter_requested_runs( all_runs: list[Run], experiment_classes: list[str], @@ -1756,10 +1886,17 @@ def cmd_retry(args): 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: + for run in dependencies + runs_to_retry + accounting_only: runs_by_id.setdefault(run.id, run) - for run in runs_by_id.values(): + for run in dependencies + runs_to_retry: if get_run_status(state, run.id) in {"failed", "skipped"}: set_run_status(state, run.id, "pending") @@ -1855,7 +1992,11 @@ def main() -> None: 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 skipped 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, @@ -1918,6 +2059,8 @@ def main() -> None: 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": diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index e5354b4..5341533 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -117,6 +117,35 @@ def _build_wandb_tags(cfg: DictConfig) -> list[str] | None: 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) + run = wandb.init( + project=cfg.logging.wandb_project, + entity=cfg.logging.get("wandb_entity", None), + name=run_name, + group=group_name, + tags=_build_wandb_tags(cfg), + config=OmegaConf.to_container(cfg, resolve=True), + ) + return wandb, run + + @hydra.main(version_base=None, config_path="../../configs", config_name="xgboost") def main(cfg: DictConfig) -> None: """Train XGBoost baseline.""" @@ -142,6 +171,7 @@ def main(cfg: DictConfig) -> None: task_name = cfg.task.get("task_name", "mortality_24h") task_type = cfg.task.get("task_type", "binary") + wandb_module, wandb_run = _init_wandb_run(cfg, task_name) # Validate data prerequisites including label freshness validate_data_prerequisites( @@ -325,29 +355,12 @@ def main(cfg: DictConfig) -> None: # ========================================================================= # 7. Log to W&B # ========================================================================= - if cfg.logging.get("use_wandb", False): - import wandb - - 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}" - - run = wandb.init( - project=cfg.logging.wandb_project, - entity=cfg.logging.get("wandb_entity", None), - name=run_name, - group=group_name, - tags=_build_wandb_tags(cfg), - config=OmegaConf.to_container(cfg, resolve=True), - ) - run.summary.update(train_label_support_summary(train_support_stats)) - run.summary.update(metrics) + if wandb_run is not None: + wandb_run.summary.update(train_label_support_summary(train_support_stats)) + wandb_run.summary.update(metrics) if fairness_report: - run.summary.update(flatten_fairness_report(fairness_report)) - wandb.finish() + wandb_run.summary.update(flatten_fairness_report(fairness_report)) + wandb_module.finish() print("\n" + "=" * 80) print("XGBoost Baseline Complete!") diff --git a/tests/test_export_results.py b/tests/test_export_results.py index 3f227c0..27353d9 100644 --- a/tests/test_export_results.py +++ b/tests/test_export_results.py @@ -117,6 +117,36 @@ def test_extract_run_uses_class_metadata_from_config_or_tags(): assert tag_row["protocol"] == "B" +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_single_class_adds_server_side_filter(monkeypatch): mod = importlib.import_module("scripts.export_results") captured = {} @@ -696,6 +726,52 @@ def test_validate_warns_when_group_mixes_revisions(): 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") diff --git a/tests/test_fixes.py b/tests/test_fixes.py index f719094..1d99d7a 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -7,6 +7,7 @@ import importlib import os +import re import shutil import subprocess import sys @@ -963,6 +964,21 @@ def test_tmux_launcher_passes_revision_to_fairness(self, tmp_path): assert "--revision thesis-v2" in runner_text assert "--project slices-thesis" in runner_text assert "--entity test-entity" in runner_text + assert "uv sync --dev --locked" 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_includes_classical_baselines_in_fairness(self, tmp_path): """The fairness sweep should include the canonical baseline class by default.""" @@ -1585,6 +1601,14 @@ def test_xgboost_wandb_tags_include_revision_scope(self): 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 @@ -1983,6 +2007,57 @@ def test_main_retry_requires_revision_scope_and_project(self, monkeypatch): 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", + "--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_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", + ], + ) + + with pytest.raises(SystemExit) as excinfo: + runner.main() + + assert excinfo.value.code == 2 + def test_cmd_run_respects_requested_class_order(self, monkeypatch): import scripts.internal.run_experiments as runner @@ -2517,6 +2592,71 @@ def test_cmd_retry_skipped_resets_failed_dependencies_when_requested(self, monke 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 From 89e757fddcc4889a2a04d89f20c851b1b6a95b96 Mon Sep 17 00:00:00 2001 From: hill Date: Wed, 22 Apr 2026 22:28:27 -0400 Subject: [PATCH 109/121] Align publication fairness and XGBoost metrics Log the full binary publication metric set for XGBoost baselines and add tests that keep the documented metric surface in sync. Stamp post-run fairness summaries with artifact, checkpoint, script, and subgroup metadata, and make export validation fail closed when summaries are missing or stale. Update fairness documentation with the class/revision-based corpus and final publication command sequence. --- docs/EVAL_FAIRNESS.md | 80 +++++++++++- docs/internal/EXPERIMENT_PLAN.md | 11 +- scripts/eval/evaluate_fairness.py | 187 +++++++++++++++++++++++++-- scripts/export_results.py | 176 ++++++++++++++++++++++++- scripts/training/xgboost_baseline.py | 35 ++++- src/slices/eval/fairness_metadata.py | 80 ++++++++++++ tests/test_evaluate_fairness.py | 98 ++++++++++++++ tests/test_export_results.py | 87 +++++++++++++ tests/test_metrics.py | 20 +++ 9 files changed, 748 insertions(+), 26 deletions(-) create mode 100644 src/slices/eval/fairness_metadata.py diff --git a/docs/EVAL_FAIRNESS.md b/docs/EVAL_FAIRNESS.md index 0cabbbd..b61ecd4 100644 --- a/docs/EVAL_FAIRNESS.md +++ b/docs/EVAL_FAIRNESS.md @@ -12,11 +12,17 @@ to W&B summaries and the parquet exports produced by `scripts/export_results.py` - The standalone script is revision-scoped. Pass `--revision ` so reruns do not get mixed together. - Default benchmark fairness corpus: - - Internal experiment groups `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`, `7p`, - `10`, `11`, `12`, `13` + - 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 @@ -80,6 +86,9 @@ For regression tasks it logs: - `per_group_mae` - `per_group_r2` - `worst_group_mse` +- `worst_group_mae` +- `mse_gap` +- `mae_gap` - `group_sizes` - `group_sample_sizes` @@ -87,13 +96,15 @@ For regression tasks it logs: `scripts/eval/evaluate_fairness.py`: -1. Fetches finished W&B runs matching the requested sprint/paradigm/dataset - filters. +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 @@ -108,6 +119,12 @@ 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 @@ -121,13 +138,16 @@ race for MIMIC/combined rows. eICU race is not required. ## How To Refresh Fairness Results -1. Run the fairness sweep for a specific revision: +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 + --entity \ + --force ``` 2. Re-export the result tables: @@ -139,6 +159,54 @@ uv run python scripts/export_results.py \ --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: diff --git a/docs/internal/EXPERIMENT_PLAN.md b/docs/internal/EXPERIMENT_PLAN.md index a1438b1..d0eb5b6 100644 --- a/docs/internal/EXPERIMENT_PLAN.md +++ b/docs/internal/EXPERIMENT_PLAN.md @@ -195,10 +195,13 @@ Run fairness after training: uv run python scripts/eval/evaluate_fairness.py \ --project slices-thesis \ --revision thesis-v1 \ - --entity + --entity \ + --force ``` -Revision filtering is mandatory to avoid mixing rerun corpora. +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 @@ -218,8 +221,8 @@ 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`, or `--allow-duplicate-fingerprints` for final -tables. +`--allow-extraction-failures`, `--allow-duplicate-fingerprints`, or +`--allow-multiple-revisions` for final tables. ## Checkpoint Policy diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index 9067160..76085d8 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -51,6 +51,24 @@ import torch from slices.eval.fairness_evaluator import flatten_fairness_report +from slices.eval.fairness_metadata import ( + FAIRNESS_ARTIFACT_PATH_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, + normalize_protected_attributes, +) log = logging.getLogger("evaluate_fairness") @@ -69,7 +87,7 @@ "smart_external_reference", ] DEFAULT_PHASES = ["finetune", "supervised", "baseline"] -DEFAULT_PROTECTED_ATTRIBUTES = ["gender", "age_group", "race"] +DEFAULT_PROTECTED_ATTRIBUTES = FAIRNESS_DEFAULT_PROTECTED_ATTRIBUTES THESIS_TASKS = { "mortality_24h", "mortality_hospital", @@ -221,11 +239,144 @@ def _required_fairness_metrics_for_run(run) -> list[str]: return BINARY_FAIRNESS_REQUIRED_METRICS -def has_fairness_metrics(run, protected_attributes: list[str]) -> bool: +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_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}" + ) + + 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. + ``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 {} @@ -238,7 +389,7 @@ def has_fairness_metrics(run, protected_attributes: list[str]) -> bool: for metric_name in required_metrics: if not _has_summary_value(sm, f"{prefix}{metric_name}"): return False - return True + return not fairness_summary_metadata_issues(run, protected_attributes, min_subgroup_size) except Exception: return False @@ -296,7 +447,7 @@ def _do_update(): except (TypeError, ValueError): existing_summary = dict(getattr(run.summary, "_json_dict", {}) or {}) for key in existing_summary: - if key.startswith("fairness/") and key not in fairness_flat: + if key.startswith(FAIRNESS_CLEAR_PREFIXES) and key not in fairness_flat: run.summary[key] = None run.summary.update(fairness_flat) run.summary.save() @@ -698,7 +849,10 @@ def parse_args() -> argparse.Namespace: help=f"Attributes to evaluate (default: {DEFAULT_PROTECTED_ATTRIBUTES})", ) parser.add_argument( - "--min-subgroup-size", type=int, default=50, help="Min patients per subgroup" + "--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( @@ -782,10 +936,18 @@ def main() -> None: # 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)] + 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 existing fairness metrics.", 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 @@ -907,6 +1069,15 @@ def _sort_key(r): # 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 diff --git a/scripts/export_results.py b/scripts/export_results.py index a47c453..7c5b5d0 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -37,6 +37,23 @@ import pandas as pd import wandb from scipy import stats as scipy_stats +from slices.eval.fairness_metadata import ( + FAIRNESS_ARTIFACT_PATH_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, @@ -180,7 +197,7 @@ } EVAL_PHASES = ["finetune", "supervised", "baseline"] -FAIRNESS_ATTRIBUTES = ["gender", "age_group", "race"] +FAIRNESS_ATTRIBUTES = FAIRNESS_DEFAULT_PROTECTED_ATTRIBUTES FAIRNESS_SUMMARY_KEY_COLUMN = "_fairness_summary_keys" BINARY_FAIRNESS_REQUIRED_METRICS = [ "n_valid_groups", @@ -489,12 +506,15 @@ def extract_run(run, metric_keys: list[str]) -> dict: "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), "_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 @@ -1255,6 +1275,131 @@ def _fairness_completeness_issues(per_seed_df: pd.DataFrame) -> list[tuple[pd.Se 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}" + ) + + 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: @@ -1580,6 +1725,35 @@ def validate( + 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)) diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index 5341533..7a6f1b5 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -20,9 +20,13 @@ 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 slices.data.datamodule import ICUDataModule @@ -82,6 +86,29 @@ def _binary_ece(y_true, y_pred_proba, n_bins: int = 15) -> float: 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", [])) @@ -295,13 +322,7 @@ def main(cfg: DictConfig) -> None: metrics["test/r2"] = r2_score(y_test, y_pred) else: y_pred_proba = model.predict_proba(X_test)[:, 1] - y_pred_class = (y_pred_proba >= 0.5).astype(int) - - metrics["test/auroc"] = roc_auc_score(y_test, y_pred_proba) - metrics["test/auprc"] = average_precision_score(y_test, y_pred_proba) - metrics["test/accuracy"] = accuracy_score(y_test, y_pred_class) - metrics["test/brier_score"] = brier_score_loss(y_test, y_pred_proba) - metrics["test/ece"] = _binary_ece(y_test, y_pred_proba) + 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 diff --git a/src/slices/eval/fairness_metadata.py b/src/slices/eval/fairness_metadata.py new file mode 100644 index 0000000..c421331 --- /dev/null +++ b/src/slices/eval/fairness_metadata.py @@ -0,0 +1,80 @@ +"""Shared metadata contract for post-run fairness summaries.""" + +from __future__ import annotations + +import json +from typing import Any + +FAIRNESS_SUMMARY_SCHEMA_VERSION = "2026-04-23.v1" +FAIRNESS_SCRIPT_VERSION = "evaluate_fairness.py:2026-04-23.v1" +FAIRNESS_DEFAULT_PROTECTED_ATTRIBUTES = ["gender", "age_group", "race"] +FAIRNESS_DEFAULT_MIN_SUBGROUP_SIZE = 50 + +FAIRNESS_SCHEMA_VERSION_KEY = "_fairness_summary_schema_version" +FAIRNESS_SCRIPT_VERSION_KEY = "_fairness_script_version" +FAIRNESS_ARTIFACT_PATH_KEY = "_fairness_artifact_path" +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_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("/") diff --git a/tests/test_evaluate_fairness.py b/tests/test_evaluate_fairness.py index 4871d58..d318449 100644 --- a/tests/test_evaluate_fairness.py +++ b/tests/test_evaluate_fairness.py @@ -5,6 +5,18 @@ from types import SimpleNamespace import pytest +from slices.eval.fairness_metadata import ( + FAIRNESS_ARTIFACT_PATH_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]: @@ -33,6 +45,28 @@ def _regression_fairness_summary(attr: str) -> dict[str, float]: } +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", +) -> dict[str, object]: + metadata = { + "_eval_checkpoint_source": ( + checkpoint_source if checkpoint_source in {"best", "final"} else None + ), + FAIRNESS_SCHEMA_VERSION_KEY: FAIRNESS_SUMMARY_SCHEMA_VERSION, + FAIRNESS_SCRIPT_VERSION_KEY: FAIRNESS_SCRIPT_VERSION, + FAIRNESS_ARTIFACT_PATH_KEY: 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: 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") @@ -108,6 +142,36 @@ def test_resolve_evaluation_artifact_supports_xgboost(tmp_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" + 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_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_has_fairness_metrics_requires_requested_attribute_completeness(): mod = importlib.import_module("scripts.eval.evaluate_fairness") @@ -116,6 +180,7 @@ def test_has_fairness_metrics_requires_requested_attribute_completeness(): summary_metrics={ **_binary_fairness_summary("gender", 0.71), **_binary_fairness_summary("age_group", 0.69), + **_fresh_fairness_metadata(["gender", "age_group"]), }, ) @@ -131,6 +196,7 @@ def test_has_fairness_metrics_ignores_race_for_eicu(): summary_metrics={ **_binary_fairness_summary("gender", 0.71), **_binary_fairness_summary("age_group", 0.69), + **_fresh_fairness_metadata(["gender", "age_group", "race"]), }, ) @@ -156,6 +222,7 @@ def test_has_fairness_metrics_accepts_written_nan_summary_values(): 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, @@ -164,6 +231,35 @@ def test_has_fairness_metrics_accepts_written_nan_summary_values(): 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_requires_regression_metric_family(): mod = importlib.import_module("scripts.eval.evaluate_fairness") @@ -175,6 +271,7 @@ def test_has_fairness_metrics_requires_regression_metric_family(): summary_metrics={ **_regression_fairness_summary("gender"), **_regression_fairness_summary("age_group"), + **_fresh_fairness_metadata(["gender", "age_group"]), }, ) partial = SimpleNamespace( @@ -186,6 +283,7 @@ def test_has_fairness_metrics_requires_regression_metric_family(): "fairness/gender/n_valid_groups": 2, "fairness/gender/worst_group_mse": 2.0, **_regression_fairness_summary("age_group"), + **_fresh_fairness_metadata(["gender", "age_group"]), }, ) diff --git a/tests/test_export_results.py b/tests/test_export_results.py index 27353d9..8f1ab1f 100644 --- a/tests/test_export_results.py +++ b/tests/test_export_results.py @@ -40,6 +40,39 @@ def _row(experiment_class, paradigm, seed, task="mortality_24h", offset=0.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, +): + protected_attributes = protected_attributes or mod.FAIRNESS_ATTRIBUTES + return { + 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_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") @@ -514,7 +547,14 @@ def test_validate_warns_on_missing_or_failed_checkpoint_provenance(): "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", + ), }, ] ) @@ -572,6 +612,53 @@ def test_validate_accepts_written_nan_fairness_summary_metrics(): 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_does_not_require_eicu_race_fairness(): mod = importlib.import_module("scripts.export_results") diff --git a/tests/test_metrics.py b/tests/test_metrics.py index fddbdd6..e87d69c 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -8,6 +8,7 @@ - Error handling """ +import numpy as np import pytest import torch from slices.eval.metrics import ( @@ -356,6 +357,25 @@ def test_binary_defaults_match_public_export_surface(self): "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.""" task_types = ["binary", "multiclass", "multilabel", "regression"] From eee115509c5e043b485cf49f20e46eb5454aa943 Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 23 Apr 2026 00:06:39 -0400 Subject: [PATCH 110/121] Fix final experiment matrix safeguards Update the thesis run matrix to remove unreliable 1% mortality_24h cells while adding the full hospital-mortality label-efficiency curve. Correct hospital mortality labeling, make result and fairness filtering robust to incomplete W&B tags, force fairness refreshes, and cover SMART imputation diagnostics and unknown gender handling with tests. Behavioral implications: the final matrix is now 2590 launched runs including SMART and 2455 excluding SMART. Processed data was regenerated locally after the mortality label semantic change, but ignored data artifacts are not part of this commit. --- docs/internal/EXPERIMENT_PLAN.md | 27 ++++++--- scripts/eval/evaluate_fairness.py | 61 +++++++++---------- scripts/export_results.py | 60 ++++++++++--------- scripts/internal/launch_thesis_tmux.sh | 2 +- scripts/internal/run_experiments.py | 27 ++++++--- src/slices/data/labels/mortality.py | 81 ++++++++++++++++++++++++-- src/slices/eval/fairness_evaluator.py | 31 +++++++++- src/slices/eval/imputation.py | 30 +++++++++- tests/test_export_results.py | 47 +++++++++++---- tests/test_fairness_evaluator.py | 20 +++++++ tests/test_fixes.py | 66 ++++++++++++++++----- tests/test_imputation_eval.py | 43 ++++++++++++++ tests/test_task_builders.py | 37 ++++++++++-- 13 files changed, 419 insertions(+), 113 deletions(-) diff --git a/docs/internal/EXPERIMENT_PLAN.md b/docs/internal/EXPERIMENT_PLAN.md index d0eb5b6..deb437b 100644 --- a/docs/internal/EXPERIMENT_PLAN.md +++ b/docs/internal/EXPERIMENT_PLAN.md @@ -93,24 +93,33 @@ project-history knowledge. ## Experiment Classes -The final launched run budget remains scientifically equivalent to the prior -thesis matrix. +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 | 840 | +| `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 | 360 | +| `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 | **2305** | +| **Total** | Includes appendix SMART reference | **2590** | -The formal thesis corpus excluding `smart_external_reference` has 2170 launched +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 @@ -172,7 +181,7 @@ uv run python scripts/internal/run_experiments.py run \ --dry-run ``` -Expected dry-run count: 2305 launched runs. +Expected dry-run count: 2590 launched runs. ## Fairness @@ -237,8 +246,8 @@ post-hoc fairness evaluation. Before launching final runs: -- generated matrix has exactly 2305 launched runs including SMART -- generated matrix has exactly 2170 launched runs excluding SMART +- 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=` diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index 76085d8..7a2ad51 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -137,6 +137,28 @@ def _retry(fn, max_retries=3, base_delay=5): 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], @@ -153,50 +175,21 @@ def fetch_eval_runs( path = f"{entity}/{project}" if entity else project filters: dict = {"state": "finished"} - tag_filters: list[str] = [] - - if experiment_classes and len(experiment_classes) == 1: - tag_filters.append(f"experiment_class:{experiment_classes[0]}") - if paradigms and len(paradigms) == 1: - tag_filters.append(f"paradigm:{paradigms[0]}") - if datasets and len(datasets) == 1: - tag_filters.append(f"dataset:{datasets[0]}") - if phases and len(phases) == 1: - tag_filters.append(f"phase:{phases[0]}") - if revisions and len(revisions) == 1: - tag_filters.append(f"revision:{revisions[0]}") - - if tag_filters: - filters["$and"] = [{"tags": tag} for tag in tag_filters] 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") - # Client-side filtering for multi-value filters - experiment_class_set = ( - set(experiment_classes) if experiment_classes and len(experiment_classes) > 1 else None - ) - paradigm_set = set(paradigms) if paradigms and len(paradigms) > 1 else None - dataset_set = set(datasets) if datasets and len(datasets) > 1 else None - phase_set = set(phases) if phases and len(phases) > 1 else None - revision_set = set(revisions) if revisions and len(revisions) > 1 else None - runs = [] for run in runs_iter: - tags = set(run.tags) - - if experiment_class_set and not any( - f"experiment_class:{experiment_class}" in tags - for experiment_class in experiment_class_set - ): + if not _run_matches_filter(run, "experiment_class", "experiment_class", experiment_classes): continue - if paradigm_set and not any(f"paradigm:{p}" in tags for p in paradigm_set): + if not _run_matches_filter(run, "paradigm", "paradigm", paradigms): continue - if dataset_set and not any(f"dataset:{d}" in tags for d in dataset_set): + if not _run_matches_filter(run, "dataset", "dataset", datasets): continue - if phase_set and not any(f"phase:{p}" in tags for p in phase_set): + if not _run_matches_filter(run, "phase", "phase", phases): continue - if revision_set and not any(f"revision:{r}" in tags for r in revision_set): + if not _run_matches_filter(run, "revision", "revision", revisions): continue runs.append(run) diff --git a/scripts/export_results.py b/scripts/export_results.py index 7c5b5d0..d7d91ba 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -82,7 +82,7 @@ FIXED_SEED_EXPERIMENT_CLASSES = set(EXPERIMENT_CLASSES) EXPECTED_FIXED_SEEDS = {42, 123, 456, 789, 1011} -CAPACITY_LABEL_FRACTIONS = {0.01, 0.1, 0.5} +CAPACITY_LABEL_FRACTIONS = {0.05, 0.1, 0.5} CAPACITY_PARADIGMS = {"mae", "supervised"} THESIS_TASKS = { "mortality_24h", @@ -224,6 +224,33 @@ # --------------------------------------------------------------------------- +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, @@ -242,42 +269,21 @@ def fetch_all_runs( if state: filters["state"] = state - tag_filters: list[str] = [] - if experiment_class and len(experiment_class) == 1: - tag_filters.append(f"experiment_class:{experiment_class[0]}") - if paradigm and len(paradigm) == 1: - tag_filters.append(f"paradigm:{paradigm[0]}") - if dataset and len(dataset) == 1: - tag_filters.append(f"dataset:{dataset[0]}") - if phase and len(phase) == 1: - tag_filters.append(f"phase:{phase[0]}") - if revision and len(revision) == 1: - tag_filters.append(f"revision:{revision[0]}") - if tag_filters: - filters["$and"] = [{"tags": tag} for tag in tag_filters] - 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") - class_set = set(experiment_class) if experiment_class and len(experiment_class) > 1 else None - paradigm_set = set(paradigm) if paradigm and len(paradigm) > 1 else None - dataset_set = set(dataset) if dataset and len(dataset) > 1 else None - phase_set = set(phase) if phase and len(phase) > 1 else None - revision_set = set(revision) if revision and len(revision) > 1 else None - runs = [] for run in runs_iter: - tags = set(run.tags) - if class_set and not any(f"experiment_class:{value}" in tags for value in class_set): + if not _run_matches_filter(run, "experiment_class", "experiment_class", experiment_class): continue - if paradigm_set and not any(f"paradigm:{value}" in tags for value in paradigm_set): + if not _run_matches_filter(run, "paradigm", "paradigm", paradigm): continue - if dataset_set and not any(f"dataset:{value}" in tags for value in dataset_set): + if not _run_matches_filter(run, "dataset", "dataset", dataset): continue - if phase_set and not any(f"phase:{value}" in tags for value in phase_set): + if not _run_matches_filter(run, "phase", "phase", phase): continue - if revision_set and not any(f"revision:{value}" in tags for value in revision_set): + if not _run_matches_filter(run, "revision", "revision", revision): continue runs.append(run) diff --git a/scripts/internal/launch_thesis_tmux.sh b/scripts/internal/launch_thesis_tmux.sh index 2e83e3b..7b7ceaf 100755 --- a/scripts/internal/launch_thesis_tmux.sh +++ b/scripts/internal/launch_thesis_tmux.sh @@ -193,7 +193,7 @@ 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") +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[@]}") diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index 159a3b4..9083b8c 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -65,11 +65,11 @@ EXPECTED_CLASS_COUNTS = { "core_ssl_benchmark": 465, - "label_efficiency": 840, + "label_efficiency": 1155, "cross_dataset_transfer": 120, "hp_robustness": 150, "capacity_study": 100, - "classical_baselines": 360, + "classical_baselines": 330, "ts2vec_extension": 135, "smart_external_reference": 135, } @@ -79,8 +79,9 @@ TASKS = ["mortality_24h", "mortality_hospital", "aki_kdigo", "los_remaining"] 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.01, 0.1, 0.5] +LABEL_FRACTIONS_CAPACITY = [0.05, 0.1, 0.5] MODEL_SIZES = { "medium": { @@ -570,7 +571,7 @@ def build_label_efficiency(self) -> None: for dataset in DATASETS: for paradigm in SSL_PARADIGMS: pretrain = self._get_pretrain(paradigm, dataset, seed) - for frac in LABEL_FRACTIONS_FULL: + for frac in LABEL_FRACTIONS_MORTALITY_24H: self._add_finetune( experiment_class, paradigm, @@ -592,7 +593,12 @@ def build_label_efficiency(self) -> None: frac, ) for task in TASKS[1:]: - for frac in LABEL_FRACTIONS_TREND: + fractions = ( + LABEL_FRACTIONS_FULL + if task == "mortality_hospital" + else LABEL_FRACTIONS_TREND + ) + for frac in fractions: self._add_finetune( experiment_class, paradigm, @@ -613,10 +619,15 @@ def build_label_efficiency(self) -> None: pretrain, frac, ) - for frac in LABEL_FRACTIONS_FULL: + for frac in LABEL_FRACTIONS_MORTALITY_24H: self._add_supervised(experiment_class, dataset, seed, "mortality_24h", frac) for task in TASKS[1:]: - for frac in LABEL_FRACTIONS_TREND: + 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: @@ -758,7 +769,7 @@ def build_classical_baselines(self) -> None: 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_FULL: + 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:]: diff --git a/src/slices/data/labels/mortality.py b/src/slices/data/labels/mortality.py index d2f96aa..1712f90 100644 --- a/src/slices/data/labels/mortality.py +++ b/src/slices/data/labels/mortality.py @@ -38,7 +38,7 @@ class MortalityLabelBuilder(LabelBuilder): unknown outcomes as null """ - SEMANTIC_VERSION = "2.2.0" + 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. @@ -285,6 +285,48 @@ def _death_evidence_expr() -> pl.Expr: 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.""" @@ -320,13 +362,42 @@ def _effective_death_precision_expr() -> pl.Expr: .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._effective_death_time_expr() - effective_precision = self._effective_death_precision_expr() + 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. @@ -355,9 +426,9 @@ def _build_hospital_mortality_with_obs( "stay_id", pl.when(died_during_obs) .then(None) - .when(self._death_evidence_expr()) + .when(hospital_death_evidence) .then(1) - .when(self._known_survivor_expr()) + .when(self._known_hospital_survivor_expr()) .then(0) .otherwise(None) .cast(pl.Int32) diff --git a/src/slices/eval/fairness_evaluator.py b/src/slices/eval/fairness_evaluator.py index c2d99b0..84ca485 100644 --- a/src/slices/eval/fairness_evaluator.py +++ b/src/slices/eval/fairness_evaluator.py @@ -141,6 +141,35 @@ def _bin_age(self, ages: torch.Tensor) -> torch.Tensor: 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], @@ -184,7 +213,7 @@ def _encode_attribute( unique_vals = set() raw_vals = [] for row in rows: - val = row.get("gender") + val = self._canonicalize_gender(row.get("gender")) raw_vals.append(val) if val is not None: unique_vals.add(val) diff --git a/src/slices/eval/imputation.py b/src/slices/eval/imputation.py index fe9b2f5..2534873 100644 --- a/src/slices/eval/imputation.py +++ b/src/slices/eval/imputation.py @@ -23,7 +23,7 @@ import torch.nn as nn from torch.utils.data import DataLoader -from slices.models.encoders import ObservationTransformerEncoder +from slices.models.encoders import ObservationTransformerEncoder, SMARTEncoder logger = logging.getLogger(__name__) @@ -88,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. @@ -195,6 +219,8 @@ def from_mae_checkpoint( # 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() @@ -248,6 +274,8 @@ def from_encoder_checkpoint( seq_length = encoder_config.get("max_seq_length", 168) 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. " diff --git a/tests/test_export_results.py b/tests/test_export_results.py index 8f1ab1f..8e06d89 100644 --- a/tests/test_export_results.py +++ b/tests/test_export_results.py @@ -180,9 +180,41 @@ def test_extract_run_exports_launch_commit_from_config_or_tags(): assert mod.extract_run(tag_run, [])["launch_commit"] == "tag-commit" -def test_fetch_all_runs_single_class_adds_server_side_filter(monkeypatch): +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): @@ -192,11 +224,11 @@ def runs(self, path, filters, order): captured["path"] = path captured["filters"] = filters captured["order"] = order - return [] + return [keep_config, keep_tags, drop, drop_stale_tag] monkeypatch.setattr(mod.wandb, "Api", DummyApi) - mod.fetch_all_runs( + runs = mod.fetch_all_runs( project="proj", entity="entity", experiment_class=["core_ssl_benchmark"], @@ -207,13 +239,8 @@ def runs(self, path, filters, order): ) assert captured["path"] == "entity/proj" - assert captured["filters"]["$and"] == [ - {"tags": "experiment_class:core_ssl_benchmark"}, - {"tags": "paradigm:mae"}, - {"tags": "dataset:miiv"}, - {"tags": "phase:finetune"}, - {"tags": "revision:thesis-v1"}, - ] + 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(): diff --git a/tests/test_fairness_evaluator.py b/tests/test_fairness_evaluator.py index 37237a0..b974530 100644 --- a/tests/test_fairness_evaluator.py +++ b/tests/test_fairness_evaluator.py @@ -267,6 +267,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.""" diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 1d99d7a..700d251 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -770,10 +770,38 @@ def test_default_experiment_classes_include_all_downstream_training_classes(self } assert expected.issubset(set(mod.DEFAULT_EXPERIMENT_CLASSES)) - def test_fetch_eval_runs_single_revision_adds_server_side_filter(self, monkeypatch): - """A single revision should be sent as a W&B tag filter.""" + 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): @@ -783,11 +811,11 @@ def runs(self, path, filters, order): captured["path"] = path captured["filters"] = filters captured["order"] = order - return [] + return [keep_config, keep_tags, drop, drop_stale_tag] monkeypatch.setitem(sys.modules, "wandb", SimpleNamespace(Api=DummyApi)) - mod.fetch_eval_runs( + runs = mod.fetch_eval_runs( project="proj", entity="entity", experiment_classes=["core_ssl_benchmark"], @@ -798,13 +826,8 @@ def runs(self, path, filters, order): ) assert captured["path"] == "entity/proj" - assert captured["filters"]["$and"] == [ - {"tags": "experiment_class:core_ssl_benchmark"}, - {"tags": "paradigm:mae"}, - {"tags": "dataset:miiv"}, - {"tags": "phase:finetune"}, - {"tags": "revision:thesis-v1"}, - ] + 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.""" @@ -964,6 +987,7 @@ def test_tmux_launcher_passes_revision_to_fairness(self, tmp_path): 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 def test_auto_shutdown_counts_fairness_and_export_processes(self): @@ -3362,11 +3386,27 @@ def test_matrix_counts_match_final_experiment_classes(self): runs = generate_all_runs() by_class = Counter(run.experiment_class for run in runs) - assert len(runs) == 2305 - assert sum(run.experiment_class != "smart_external_reference" for run in runs) == 2170 + 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 diff --git a/tests/test_imputation_eval.py b/tests/test_imputation_eval.py index db962db..b58b6fc 100644 --- a/tests/test_imputation_eval.py +++ b/tests/test_imputation_eval.py @@ -393,6 +393,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.""" diff --git a/tests/test_task_builders.py b/tests/test_task_builders.py index ea5a768..834aa86 100644 --- a/tests/test_task_builders.py +++ b/tests/test_task_builders.py @@ -3,7 +3,7 @@ Comprehensive tests for mortality tasks and task builder infrastructure. """ -from datetime import datetime +from datetime import date, datetime import polars as pl import pytest @@ -566,14 +566,16 @@ 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": [None], - "death_time_precision": [None], - "death_source": [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], @@ -615,6 +617,33 @@ def test_null_hospital_flag_without_death_evidence_is_null_for_windowed_mortalit 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( From 2cc5fc3fbc5ca513c3391e5425854e432ef9f718 Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 23 Apr 2026 08:37:11 -0400 Subject: [PATCH 111/121] Harden benchmark launch and artifact provenance Emit lowercase Hydra protocol selectors for generated finetune runs, revalidate completed pretrain encoder artifacts, and require explicit W&B entity ownership for direct final launches. Add thesis task defaults and fail-closed combined rebuild checks, checkpoint/model SHA256 provenance for fairness/export linkage, and durable imputation result output. Behavioral implications: final reruns now fail closed on missing pretrain encoders, stale fairness artifacts, missing artifact digests, missing direct-launch W&B entity, and source datasets missing required thesis tasks. --- scripts/eval/evaluate_fairness.py | 24 ++- scripts/eval/evaluate_imputation.py | 15 ++ scripts/export_results.py | 34 +++- scripts/internal/run_experiments.py | 66 +++++++- .../preprocessing/create_combined_dataset.py | 42 ++++- scripts/preprocessing/extract_ricu.py | 2 +- scripts/training/finetune.py | 28 ++++ scripts/training/supervised.py | 26 ++++ scripts/training/xgboost_baseline.py | 11 ++ src/slices/constants.py | 8 + src/slices/data/config_schemas.py | 3 +- src/slices/data/extractors/base.py | 4 +- src/slices/eval/fairness_metadata.py | 19 ++- tests/test_base_extractor.py | 1 + tests/test_config_schemas.py | 6 + tests/test_evaluate_fairness.py | 27 ++++ tests/test_export_results.py | 29 ++++ tests/test_extractor_config.py | 6 +- tests/test_fixes.py | 146 ++++++++++++++++-- tests/test_imputation_eval.py | 7 +- 20 files changed, 467 insertions(+), 37 deletions(-) diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index 7a2ad51..315aa2b 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -50,9 +50,12 @@ from typing import Any, Optional import torch +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, @@ -67,6 +70,7 @@ canonical_artifact_id, decode_protected_attributes, encode_protected_attributes, + file_sha256, normalize_protected_attributes, ) @@ -88,12 +92,7 @@ ] DEFAULT_PHASES = ["finetune", "supervised", "baseline"] DEFAULT_PROTECTED_ATTRIBUTES = FAIRNESS_DEFAULT_PROTECTED_ATTRIBUTES -THESIS_TASKS = { - "mortality_24h", - "mortality_hospital", - "aki_kdigo", - "los_remaining", -} +THESIS_TASKS = set(BENCHMARK_THESIS_TASKS) BINARY_FAIRNESS_REQUIRED_METRICS = [ "n_valid_groups", "n_metric_valid_groups", @@ -289,6 +288,7 @@ def build_fairness_summary_metadata( 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), @@ -356,6 +356,18 @@ def fairness_summary_metadata_issues( 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 diff --git a/scripts/eval/evaluate_imputation.py b/scripts/eval/evaluate_imputation.py index 451e89b..17fed9f 100644 --- a/scripts/eval/evaluate_imputation.py +++ b/scripts/eval/evaluate_imputation.py @@ -20,6 +20,9 @@ masking.strategies=[random] """ +import json +from pathlib import Path + import hydra import lightning.pytorch as pl import torch @@ -203,10 +206,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 index d7d91ba..e6bb016 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -37,8 +37,12 @@ import pandas as pd import wandb from scipy import stats as scipy_stats +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, @@ -84,12 +88,7 @@ EXPECTED_FIXED_SEEDS = {42, 123, 456, 789, 1011} CAPACITY_LABEL_FRACTIONS = {0.05, 0.1, 0.5} CAPACITY_PARADIGMS = {"mae", "supervised"} -THESIS_TASKS = { - "mortality_24h", - "mortality_hospital", - "aki_kdigo", - "los_remaining", -} +THESIS_TASKS = set(BENCHMARK_THESIS_TASKS) FINGERPRINT = [ "experiment_class", @@ -514,6 +513,8 @@ def extract_run(run, metric_keys: list[str]) -> dict: "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), @@ -1202,6 +1203,10 @@ def _checkpoint_provenance_issues(per_seed_df: pd.DataFrame) -> list[tuple[pd.Se 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")) @@ -1400,6 +1405,23 @@ def _fairness_metadata_staleness_issues( 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)) diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index 9083b8c..db24f0c 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -35,6 +35,8 @@ from datetime import datetime, timezone from pathlib import Path +from slices.constants import THESIS_TASKS + # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- @@ -76,7 +78,7 @@ SSL_PARADIGMS = ["mae", "jepa", "contrastive"] DATASETS = ["miiv", "eicu", "combined"] -TASKS = ["mortality_24h", "mortality_hospital", "aki_kdigo", "los_remaining"] +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] @@ -170,6 +172,11 @@ def protocol(self) -> str | None: return "B" return None + @property + def protocol_selector(self) -> str | None: + protocol = self.protocol + return protocol.lower() if protocol is not None else None + def build_command(self, runs_by_id: dict[str, "Run"]) -> list[str]: """Build the subprocess command for this run.""" if self.run_type == "pretrain": @@ -191,8 +198,8 @@ def _metadata_overrides(self) -> list[str]: ] if self.experiment_subtype is not None: overrides.append(f"experiment_subtype={self.experiment_subtype}") - if self.protocol is not None and self.run_type == "finetune": - overrides.append(f"protocol={self.protocol}") + 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(): @@ -1182,6 +1189,52 @@ def set_run_status(state: dict, run_id: str, status: str, **kwargs) -> None: 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 @@ -1425,6 +1478,7 @@ def handle_signal(signum, frame): 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") @@ -1467,7 +1521,7 @@ def handle_signal(signum, frame): continue if run.id in active: continue - if all(get_run_status(state, dep_id) == "completed" for dep_id in run.depends_on): + if _dependencies_ready(run, runs, runs_by_id, state): ready.append(run) launch_batch = _select_ready_runs(ready, set(active), runs_by_id, parallel) @@ -2058,6 +2112,8 @@ def main() -> None: 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") @@ -2067,6 +2123,8 @@ def main() -> None: 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") diff --git a/scripts/preprocessing/create_combined_dataset.py b/scripts/preprocessing/create_combined_dataset.py index 8fd3f4c..e3b7582 100644 --- a/scripts/preprocessing/create_combined_dataset.py +++ b/scripts/preprocessing/create_combined_dataset.py @@ -23,7 +23,7 @@ import polars as pl import yaml -from slices.constants import MIN_STAY_HOURS, SEQ_LENGTH_HOURS +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 @@ -174,6 +174,39 @@ def validate_frame_schema(name: str, df_a: pl.DataFrame, df_b: pl.DataFrame) -> ) +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.""" cols_a = set(labels_a.columns) @@ -319,6 +352,13 @@ 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 diff --git a/scripts/preprocessing/extract_ricu.py b/scripts/preprocessing/extract_ricu.py index 771fe20..11d1d31 100644 --- a/scripts/preprocessing/extract_ricu.py +++ b/scripts/preprocessing/extract_ricu.py @@ -51,7 +51,7 @@ def main(cfg: DictConfig) -> None: 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/training/finetune.py b/scripts/training/finetune.py index 49b987c..0344fbf 100644 --- a/scripts/training/finetune.py +++ b/scripts/training/finetune.py @@ -28,11 +28,18 @@ 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 ( report_and_validate_train_label_support, @@ -45,6 +52,26 @@ ) +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( @@ -380,6 +407,7 @@ def main(cfg: DictConfig) -> None: 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 "", diff --git a/scripts/training/supervised.py b/scripts/training/supervised.py index 0aa47f0..9356f2a 100644 --- a/scripts/training/supervised.py +++ b/scripts/training/supervised.py @@ -31,6 +31,11 @@ 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 ( @@ -45,6 +50,26 @@ ) +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 = [] @@ -480,6 +505,7 @@ def main(cfg: DictConfig) -> None: 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 "", diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index 7a6f1b5..40af205 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -30,6 +30,11 @@ roc_auc_score, ) 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 from xgboost import XGBClassifier, XGBRegressor @@ -379,6 +384,12 @@ def main(cfg: DictConfig) -> None: 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() diff --git a/src/slices/constants.py b/src/slices/constants.py index f2d2727..f1ec123 100644 --- a/src/slices/constants.py +++ b/src/slices/constants.py @@ -17,6 +17,14 @@ # (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", +) + # ============================================================================= # Extraction # ============================================================================= diff --git a/src/slices/data/config_schemas.py b/src/slices/data/config_schemas.py index 0696af2..83f44c4 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, ) @@ -191,7 +192,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 diff --git a/src/slices/data/extractors/base.py b/src/slices/data/extractors/base.py index cb9546c..02781bb 100644 --- a/src/slices/data/extractors/base.py +++ b/src/slices/data/extractors/base.py @@ -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 diff --git a/src/slices/eval/fairness_metadata.py b/src/slices/eval/fairness_metadata.py index c421331..eb45c11 100644 --- a/src/slices/eval/fairness_metadata.py +++ b/src/slices/eval/fairness_metadata.py @@ -2,17 +2,22 @@ from __future__ import annotations +import hashlib import json +from pathlib import Path from typing import Any -FAIRNESS_SUMMARY_SCHEMA_VERSION = "2026-04-23.v1" -FAIRNESS_SCRIPT_VERSION = "evaluate_fairness.py:2026-04-23.v1" +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" @@ -22,6 +27,7 @@ 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, @@ -78,3 +84,12 @@ def canonical_artifact_id(path_value: Any) -> str: 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/tests/test_base_extractor.py b/tests/test_base_extractor.py index c1c505b..069b6bc 100644 --- a/tests/test_base_extractor.py +++ b/tests/test_base_extractor.py @@ -1000,6 +1000,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 c225977..71f2a1b 100644 --- a/tests/test_config_schemas.py +++ b/tests/test_config_schemas.py @@ -9,6 +9,7 @@ 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, @@ -109,6 +110,11 @@ def test_label_definition_fields_are_accepted(self): 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( diff --git a/tests/test_evaluate_fairness.py b/tests/test_evaluate_fairness.py index d318449..b76098f 100644 --- a/tests/test_evaluate_fairness.py +++ b/tests/test_evaluate_fairness.py @@ -6,7 +6,9 @@ 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, @@ -51,14 +53,17 @@ def _fresh_fairness_metadata( 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), @@ -146,6 +151,8 @@ 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={ @@ -163,6 +170,7 @@ def test_build_fairness_summary_metadata_stamps_artifact_and_settings(tmp_path): ) 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 @@ -260,6 +268,25 @@ def test_has_fairness_metrics_rejects_stale_metadata_settings(): ) +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") diff --git a/tests/test_export_results.py b/tests/test_export_results.py index 8e06d89..172fd15 100644 --- a/tests/test_export_results.py +++ b/tests/test_export_results.py @@ -47,12 +47,15 @@ def _fresh_fairness_metadata( 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), @@ -686,6 +689,32 @@ def test_validate_accepts_fresh_fairness_summary_metadata(): 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") diff --git a/tests/test_extractor_config.py b/tests/test_extractor_config.py index 64bac23..ddf873a 100644 --- a/tests/test_extractor_config.py +++ b/tests/test_extractor_config.py @@ -2,6 +2,7 @@ import pytest from pydantic import ValidationError +from slices.constants import THESIS_TASKS from slices.data.extractors.base import ExtractorConfig @@ -175,11 +176,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_fixes.py b/tests/test_fixes.py index 700d251..58f4a93 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -21,6 +21,7 @@ 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 @@ -517,6 +518,39 @@ def test_patient_ids_are_namespaced_even_without_stay_collision(self): 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.""" @@ -553,10 +587,20 @@ def _write_processed_dataset( { "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", @@ -565,13 +609,8 @@ def _write_processed_dataset( "seq_length_hours": 2, "min_stay_hours": 2, "label_horizon_hours": 24, - "task_names": ["mortality_24h"], - "label_manifest": { - "mortality_24h": { - "builder_version": "1.0.0", - "config_hash": "abc123", - } - }, + "task_names": list(THESIS_TASKS), + "label_manifest": label_manifest, } with open(processed_dir / "metadata.yaml", "w") as f: yaml.dump(metadata, f) @@ -2047,6 +2086,8 @@ def test_main_run_dry_run_warns_without_launch_commit(self, monkeypatch, capsys) "dryrun-audit", "--project", "slices-thesis", + "--entity", + "hannes-ill", "--dry-run", ], ) @@ -2074,6 +2115,34 @@ def test_main_run_requires_launch_commit_for_final_launch(self, monkeypatch): "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_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", ], ) @@ -2339,7 +2408,8 @@ def test_protocol_a_finetune_command_uses_strict_linear_probe(self): assert "task.head_type=linear" in cmd assert "task.hidden_dims=[]" in cmd assert "task.dropout=0.0" in cmd - assert "protocol=A" in cmd + assert "protocol=a" in cmd + assert "protocol=A" not in cmd assert "+protocol=A" not in cmd def test_protocol_b_finetune_command_keeps_task_mlp_head(self): @@ -2371,7 +2441,8 @@ def test_protocol_b_finetune_command_keeps_task_mlp_head(self): assert "training.freeze_encoder=false" in cmd assert "task.head_type=linear" not in cmd - assert "protocol=B" in cmd + assert "protocol=b" 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): @@ -2393,6 +2464,17 @@ def test_xgboost_command_does_not_load_protocol_config_group(self): assert "+protocol=B" not in cmd assert "protocol=B" not in cmd + def test_generated_finetune_protocol_selectors_are_lowercase(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 sum("protocol=a" in command for command in commands) == 825 + assert sum("protocol=b" in command for command in commands) == 1020 + def test_manual_finetune_protocol_configs_encode_safe_defaults(self): protocol_a = OmegaConf.load("configs/protocol/a.yaml") protocol_b = OmegaConf.load("configs/protocol/b.yaml") @@ -2834,6 +2916,52 @@ def test_run_scheduler_returns_nonzero_for_failed_run(self, tmp_path, monkeypatc 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 diff --git a/tests/test_imputation_eval.py b/tests/test_imputation_eval.py index b58b6fc..dc397a3 100644 --- a/tests/test_imputation_eval.py +++ b/tests/test_imputation_eval.py @@ -457,7 +457,7 @@ def test_custom_decoder(self): class TestEvaluateImputationScript: """Regression tests for the standalone imputation evaluation script.""" - def test_pretrain_checkpoint_still_trains_probe_decoder(self, monkeypatch): + 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 = {} @@ -555,7 +555,7 @@ def evaluate(self, dataloader, mask_strategy, mask_ratio, feature_stds): "logging": { "use_wandb": False, }, - "output_dir": "outputs/test-imputation", + "output_dir": str(tmp_path / "test-imputation"), } ) @@ -569,3 +569,6 @@ def evaluate(self, dataloader, mask_strategy, mask_ratio, feature_stds): 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() From ffb2679476c652a671412b45fc3bea1b7b187c50 Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 23 Apr 2026 09:40:28 -0400 Subject: [PATCH 112/121] Harden thesis benchmark guardrails Add fairness matrix coverage and artifact digest validation before post-hoc fairness inference. Enforce online W&B mode for final thesis launches, align timestep masking budgets across SSL objectives, and harden pooling/checkpoint edge cases. Behavioral implications: final launches now reject offline W&B mode, fairness evaluation fails on incomplete scoped matrices unless explicitly allowed, random timestep SSL masks use an exact per-sample budget, and CLS pooling checkpoint overrides are rejected instead of silently rebuilding incompatible modules. --- configs/finetune.yaml | 4 +- configs/supervised.yaml | 2 +- scripts/eval/evaluate_fairness.py | 185 ++++++++++++++++++++++ scripts/internal/launch_thesis_tmux.sh | 8 + scripts/internal/run_experiments.py | 11 ++ src/slices/models/common.py | 10 +- src/slices/models/encoders/observation.py | 6 +- src/slices/models/pretraining/masking.py | 24 ++- src/slices/training/checkpoint_loading.py | 58 ++++--- tests/test_evaluate_fairness.py | 97 ++++++++++++ tests/test_fixes.py | 85 ++++++++++ tests/test_jepa_objective.py | 37 +++++ tests/test_model_common.py | 35 ++++ 13 files changed, 516 insertions(+), 46 deletions(-) diff --git a/configs/finetune.yaml b/configs/finetune.yaml index 1d39c91..8d2b487 100644 --- a/configs/finetune.yaml +++ b/configs/finetune.yaml @@ -59,11 +59,11 @@ ckpt_path: null 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) diff --git a/configs/supervised.yaml b/configs/supervised.yaml index 5d26067..f77cdb8 100644 --- a/configs/supervised.yaml +++ b/configs/supervised.yaml @@ -54,7 +54,7 @@ checkpoint_dir: ${output_dir}/checkpoints # 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) diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index 315aa2b..056e8d2 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -111,6 +111,22 @@ "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", +] # --------------------------------------------------------------------------- @@ -197,6 +213,137 @@ def fetch_eval_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 = config.get("protocol") or _tag_value(tags, "protocol") + if protocol is None: + if freeze is True: + protocol = "A" + elif freeze is False or phase == "baseline": + protocol = "B" + + 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() @@ -642,6 +789,28 @@ def resolve_evaluation_artifact( 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 # --------------------------------------------------------------------------- @@ -938,6 +1107,21 @@ def main() -> None: 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) @@ -1012,6 +1196,7 @@ def _sort_key(r): 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) diff --git a/scripts/internal/launch_thesis_tmux.sh b/scripts/internal/launch_thesis_tmux.sh index 7b7ceaf..e1e3833 100755 --- a/scripts/internal/launch_thesis_tmux.sh +++ b/scripts/internal/launch_thesis_tmux.sh @@ -40,6 +40,14 @@ if [[ -z "$WANDB_ENTITY" ]]; then 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 diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index db24f0c..8843f01 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -969,6 +969,13 @@ def _validate_clean_final_launch_state(launch_commit: str) -> str | None: return None +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 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"}: @@ -990,6 +997,10 @@ def validate_direct_final_launch_policy(args, parser: argparse.ArgumentParser) - "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) + error = _validate_clean_final_launch_state(str(launch_commit)) if error: parser.error(error) diff --git a/src/slices/models/common.py b/src/slices/models/common.py index dd86806..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] diff --git a/src/slices/models/encoders/observation.py b/src/slices/models/encoders/observation.py index 94145fb..48cc25c 100644 --- a/src/slices/models/encoders/observation.py +++ b/src/slices/models/encoders/observation.py @@ -210,7 +210,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/pretraining/masking.py b/src/slices/models/pretraining/masking.py index bc860fe..b273597 100644 --- a/src/slices/models/pretraining/masking.py +++ b/src/slices/models/pretraining/masking.py @@ -102,27 +102,21 @@ def create_timestep_mask( Returns: ssl_mask: (B, T) bool mask, True = visible, False = masked. """ - rand_vals = torch.rand(batch_size, n_timesteps, device=device) - ssl_mask = rand_vals >= mask_ratio # True = visible - 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) - # Empty timesteps are not eligible SSL tokens; mark them visible so they - # are excluded from masked-token accounting and downstream losses. - ssl_mask = ssl_mask | (~valid_timestep_mask) + 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 - # Ensure at least 1 visible eligible timestep per sample - n_visible = (ssl_mask & valid_timestep_mask).sum(dim=1) # (B,) - has_valid = valid_timestep_mask.any(dim=1) - needs_fix = (n_visible == 0) & has_valid - if needs_fix.any(): - for b in range(batch_size): - if needs_fix[b]: - first_valid = valid_timestep_mask[b].nonzero(as_tuple=True)[0][0] - ssl_mask[b, first_valid] = True + masked_ordinals = torch.randperm(n_valid, device=device)[:n_masked] + ssl_mask[b, valid_idx[masked_ordinals]] = False return ssl_mask diff --git a/src/slices/training/checkpoint_loading.py b/src/slices/training/checkpoint_loading.py index 79ac6b5..1807889 100644 --- a/src/slices/training/checkpoint_loading.py +++ b/src/slices/training/checkpoint_loading.py @@ -115,6 +115,34 @@ def wrap_encoder_with_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, @@ -156,19 +184,7 @@ def load_encoder_weights( 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 = 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, - ) + 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)", @@ -281,18 +297,10 @@ def load_from_pretrain_checkpoint( # 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 = 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, - ) + 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)", diff --git a/tests/test_evaluate_fairness.py b/tests/test_evaluate_fairness.py index b76098f..1cd4b3a 100644 --- a/tests/test_evaluate_fairness.py +++ b/tests/test_evaluate_fairness.py @@ -4,6 +4,7 @@ import sys from types import SimpleNamespace +import pandas as pd import pytest from slices.eval.fairness_metadata import ( EVAL_ARTIFACT_SHA256_KEY, @@ -394,3 +395,99 @@ def test_filter_thesis_task_runs_drops_non_thesis_tasks(): 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": "B", + "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": "B", + "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": "B", + "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_fixes.py b/tests/test_fixes.py index 58f4a93..4379bce 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -974,6 +974,23 @@ def test_missing_checkpoint_provenance_raises(self, 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] @@ -1043,6 +1060,46 @@ def test_auto_shutdown_counts_fairness_and_export_processes(self): "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] @@ -2151,6 +2208,34 @@ def test_main_run_requires_wandb_entity_for_final_launch(self, monkeypatch): 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_cmd_run_respects_requested_class_order(self, monkeypatch): import scripts.internal.run_experiments as runner diff --git a/tests/test_jepa_objective.py b/tests/test_jepa_objective.py index d4399e4..1ba83c6 100644 --- a/tests/test_jepa_objective.py +++ b/tests/test_jepa_objective.py @@ -65,6 +65,43 @@ def test_block_mask_budgets_over_valid_timesteps_only(self): 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 # ============================================================================= diff --git a/tests/test_model_common.py b/tests/test_model_common.py index 1b575fb..233556e 100644 --- a/tests/test_model_common.py +++ b/tests/test_model_common.py @@ -15,6 +15,8 @@ from slices.models.encoders import ( LinearConfig, LinearEncoder, + ObservationTransformerConfig, + ObservationTransformerEncoder, TransformerConfig, TransformerEncoder, ) @@ -56,6 +58,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 +75,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 +101,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") From d40e12d473b1d957b8057cf176a24b4f1e8c2b28 Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 23 Apr 2026 10:17:04 -0400 Subject: [PATCH 113/121] Rename downstream protocols for public benchmark Replace opaque A/B downstream protocol selectors with linear_probe and full_finetune across configs, runner commands, logging, export, fairness coverage, and regression tests. Legacy A/B protocol values are canonicalized during export and fairness matrix reconstruction so older W&B runs remain comparable. --- configs/finetune.yaml | 10 +-- configs/gru_d.yaml | 2 +- .../protocol/{b.yaml => full_finetune.yaml} | 4 +- .../protocol/{a.yaml => linear_probe.yaml} | 4 +- configs/supervised.yaml | 2 +- configs/xgboost.yaml | 2 +- scripts/eval/evaluate_fairness.py | 18 +++-- scripts/export_results.py | 26 ++++--- scripts/internal/run_experiments.py | 39 ++++++---- scripts/training/xgboost_baseline.py | 3 +- src/slices/constants.py | 42 +++++++++++ src/slices/training/utils.py | 9 ++- tests/test_evaluate_fairness.py | 6 +- tests/test_export_results.py | 12 ++-- tests/test_fixes.py | 72 ++++++++++--------- 15 files changed, 163 insertions(+), 88 deletions(-) rename configs/protocol/{b.yaml => full_finetune.yaml} (71%) rename configs/protocol/{a.yaml => linear_probe.yaml} (69%) diff --git a/configs/finetune.yaml b/configs/finetune.yaml index 8d2b487..2588f34 100644 --- a/configs/finetune.yaml +++ b/configs/finetune.yaml @@ -2,14 +2,14 @@ # # This config orchestrates finetuning a pretrained encoder on downstream tasks. # By default, the encoder is frozen and only the task head is trained. Use -# protocol=a or protocol=b for thesis Protocol A/B-safe manual launches. +# protocol=linear_probe or protocol=full_finetune for benchmark-safe manual launches. # # Example usage: -# # Protocol A: strict 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=a +# # 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 # -# # Protocol B: 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=b +# # 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 diff --git a/configs/gru_d.yaml b/configs/gru_d.yaml index 179eb6f..384d265 100644 --- a/configs/gru_d.yaml +++ b/configs/gru_d.yaml @@ -24,7 +24,7 @@ dataset: miiv # Paradigm paradigm: gru_d -protocol: B +protocol: full_finetune # Final rerun corpus metadata experiment_class: null diff --git a/configs/protocol/b.yaml b/configs/protocol/full_finetune.yaml similarity index 71% rename from configs/protocol/b.yaml rename to configs/protocol/full_finetune.yaml index 9ff4a70..f68be34 100644 --- a/configs/protocol/b.yaml +++ b/configs/protocol/full_finetune.yaml @@ -1,7 +1,7 @@ # @package _global_ -# Protocol B: full finetune for practical downstream utility. +# Full finetune: train encoder and downstream head for practical utility. -protocol: B +protocol: full_finetune training: freeze_encoder: false diff --git a/configs/protocol/a.yaml b/configs/protocol/linear_probe.yaml similarity index 69% rename from configs/protocol/a.yaml rename to configs/protocol/linear_probe.yaml index aebc935..fd7150c 100644 --- a/configs/protocol/a.yaml +++ b/configs/protocol/linear_probe.yaml @@ -1,7 +1,7 @@ # @package _global_ -# Protocol A: strict linear probe for SSL representation quality. +# Linear probe: frozen encoder, linear head, representation-quality evaluation. -protocol: A +protocol: linear_probe training: freeze_encoder: true diff --git a/configs/supervised.yaml b/configs/supervised.yaml index f77cdb8..0c8ae9a 100644 --- a/configs/supervised.yaml +++ b/configs/supervised.yaml @@ -32,7 +32,7 @@ dataset: miiv # miiv | eicu | combined # Paradigm — always "supervised" for this config paradigm: supervised -protocol: B +protocol: full_finetune # Final rerun corpus metadata experiment_class: null diff --git a/configs/xgboost.yaml b/configs/xgboost.yaml index 9578034..a0a6009 100644 --- a/configs/xgboost.yaml +++ b/configs/xgboost.yaml @@ -21,7 +21,7 @@ dataset: miiv # Paradigm paradigm: xgboost -protocol: B +protocol: full_finetune # Final rerun corpus metadata experiment_class: null diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index 056e8d2..c62c93a 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -50,7 +50,14 @@ from typing import Any, Optional import torch -from slices.constants import THESIS_TASKS as BENCHMARK_THESIS_TASKS +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, @@ -263,12 +270,11 @@ def _run_matrix_row(run) -> dict[str, Any]: phase = "finetune" freeze = _get_nested(config, "training.freeze_encoder") - protocol = config.get("protocol") or _tag_value(tags, "protocol") + protocol = canonical_downstream_protocol(config.get("protocol") or _tag_value(tags, "protocol")) if protocol is None: - if freeze is True: - protocol = "A" - elif freeze is False or phase == "baseline": - protocol = "B" + 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 { diff --git a/scripts/export_results.py b/scripts/export_results.py index e6bb016..c15aa62 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -37,7 +37,14 @@ import pandas as pd import wandb from scipy import stats as scipy_stats -from slices.constants import THESIS_TASKS as BENCHMARK_THESIS_TASKS +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, @@ -463,12 +470,11 @@ def extract_run(run, metric_keys: list[str]) -> dict: phase = "finetune" freeze = _get_nested(config, "training.freeze_encoder") - protocol = config.get("protocol") or _tag_value(tags, "protocol") + protocol = canonical_downstream_protocol(config.get("protocol") or _tag_value(tags, "protocol")) if protocol is None: - if freeze is True: - protocol = "A" - elif freeze is False or config.get("paradigm") == "xgboost" or phase == "baseline": - protocol = "B" + 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) @@ -806,14 +812,16 @@ 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) - protocol_b = per_seed_df["protocol"] == "B" + full_finetune = ( + per_seed_df["protocol"].map(canonical_downstream_protocol) == FULL_FINETUNE_PROTOCOL + ) full = per_seed_df[ - protocol_b + full_finetune & (label_fraction == 1.0) & (per_seed_df["experiment_class"].isin(["core_ssl_benchmark", "classical_baselines"])) ].copy() low_label = per_seed_df[ - protocol_b + full_finetune & (label_fraction < 1.0) & (per_seed_df["experiment_class"].isin(["label_efficiency", "classical_baselines"])) ].copy() diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index 8843f01..568cc73 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -35,7 +35,11 @@ from datetime import datetime, timezone from pathlib import Path -from slices.constants import THESIS_TASKS +from slices.constants import ( + FULL_FINETUNE_PROTOCOL, + THESIS_TASKS, + downstream_protocol_from_freeze, +) # --------------------------------------------------------------------------- # Constants @@ -120,8 +124,13 @@ LAUNCH_IDENTITY_FILE = ".runner_launch_identity.json" LAUNCH_IDENTITY_KEYS = ("revision", "wandb_project", "wandb_entity", "launch_commit") -PROTO_A = {"freeze_encoder": True, "max_epochs": 50, "patience": 10, "lr": 1e-4} -PROTO_B = {"freeze_encoder": False, "max_epochs": 100, "patience": 10, "lr": 3e-4} +LINEAR_PROBE_SETTINGS = {"freeze_encoder": True, "max_epochs": 50, "patience": 10, "lr": 1e-4} +FULL_FINETUNE_SETTINGS = { + "freeze_encoder": False, + "max_epochs": 100, + "patience": 10, + "lr": 3e-4, +} # --------------------------------------------------------------------------- @@ -166,16 +175,16 @@ def phase(self) -> str: @property def protocol(self) -> str | None: - if self.freeze_encoder is True: - return "A" - if self.freeze_encoder is False or self.run_type in {"supervised", "gru_d", "xgboost"}: - return "B" + 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: - protocol = self.protocol - return protocol.lower() if protocol is not None else None + return self.protocol def build_command(self, runs_by_id: dict[str, "Run"]) -> list[str]: """Build the subprocess command for this run.""" @@ -248,9 +257,9 @@ def _finetune_cmd(self, runs_by_id: dict[str, "Run"]) -> list[str]: cmd.extend( [ "training.freeze_encoder=true", - f"training.max_epochs={PROTO_A['max_epochs']}", - f"training.early_stopping_patience={PROTO_A['patience']}", - f"optimizer.lr={PROTO_A['lr']}", + f"training.max_epochs={LINEAR_PROBE_SETTINGS['max_epochs']}", + f"training.early_stopping_patience={LINEAR_PROBE_SETTINGS['patience']}", + f"optimizer.lr={LINEAR_PROBE_SETTINGS['lr']}", "task.head_type=linear", "task.hidden_dims=[]", "task.dropout=0.0", @@ -260,9 +269,9 @@ def _finetune_cmd(self, runs_by_id: dict[str, "Run"]) -> list[str]: cmd.extend( [ "training.freeze_encoder=false", - f"training.max_epochs={PROTO_B['max_epochs']}", - f"training.early_stopping_patience={PROTO_B['patience']}", - f"optimizer.lr={PROTO_B['lr']}", + f"training.max_epochs={FULL_FINETUNE_SETTINGS['max_epochs']}", + f"training.early_stopping_patience={FULL_FINETUNE_SETTINGS['patience']}", + f"optimizer.lr={FULL_FINETUNE_SETTINGS['lr']}", ] ) if self.label_fraction < 1.0: diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index 40af205..f652469 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -29,6 +29,7 @@ recall_score, roc_auc_score, ) +from slices.constants import canonical_downstream_protocol from slices.data.datamodule import ICUDataModule from slices.eval.fairness_metadata import ( EVAL_ARTIFACT_PATH_KEY, @@ -142,7 +143,7 @@ def _build_wandb_tags(cfg: DictConfig) -> list[str] | None: 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:{cfg.protocol}") + _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") diff --git a/src/slices/constants.py b/src/slices/constants.py index f1ec123..33d4b55 100644 --- a/src/slices/constants.py +++ b/src/slices/constants.py @@ -25,6 +25,48 @@ "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 # ============================================================================= diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 52aab91..6c3acb1 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -25,6 +25,8 @@ LABEL_HORIZON_HOURS, MIN_STAY_HOURS, SEQ_LENGTH_HOURS, + canonical_downstream_protocol, + downstream_protocol_from_freeze, ) from slices.data.labels import LabelBuilder, LabelBuilderFactory, LabelConfig @@ -374,12 +376,13 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: _add_wandb_tag(tags, f"model_size:{model_size}") # Add downstream protocol family tag. Supervised-from-scratch shares the - # Protocol B optimization budget; the phase tag distinguishes it from SSL. + # 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: - _add_wandb_tag(tags, f"protocol:{cfg.protocol}") + protocol = canonical_downstream_protocol(cfg.protocol) + _add_wandb_tag(tags, f"protocol:{protocol}") elif freeze_encoder is not None: - protocol = "A" if freeze_encoder else "B" + 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) diff --git a/tests/test_evaluate_fairness.py b/tests/test_evaluate_fairness.py index 1cd4b3a..eeb8ab3 100644 --- a/tests/test_evaluate_fairness.py +++ b/tests/test_evaluate_fairness.py @@ -407,7 +407,7 @@ def test_fairness_matrix_coverage_accepts_complete_scoped_corpus(monkeypatch): "paradigm": "mae", "dataset": "miiv", "task": "mortality_24h", - "protocol": "B", + "protocol": "full_finetune", "label_fraction": 1.0, "model_size": "default", "source_dataset": None, @@ -423,7 +423,7 @@ def test_fairness_matrix_coverage_accepts_complete_scoped_corpus(monkeypatch): "paradigm": "mae", "dataset": "miiv", "phase": "finetune", - "protocol": "B", + "protocol": "full_finetune", "seed": 42, "task": {"task_name": "mortality_24h"}, }, @@ -451,7 +451,7 @@ def test_fairness_matrix_coverage_flags_missing_scoped_rows(monkeypatch): "paradigm": "mae", "dataset": "miiv", "task": "mortality_24h", - "protocol": "B", + "protocol": "full_finetune", "label_fraction": 1.0, "model_size": "default", "source_dataset": None, diff --git a/tests/test_export_results.py b/tests/test_export_results.py index 172fd15..7998cdc 100644 --- a/tests/test_export_results.py +++ b/tests/test_export_results.py @@ -29,7 +29,7 @@ def _row(experiment_class, paradigm, seed, task="mortality_24h", offset=0.0): "dataset": "miiv", "task": task, "seed": seed, - "protocol": "B", + "protocol": "full_finetune", "label_fraction": 1.0, "model_size": "default", "source_dataset": None, @@ -150,7 +150,7 @@ def test_extract_run_uses_class_metadata_from_config_or_tags(): 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"] == "B" + assert tag_row["protocol"] == "full_finetune" def test_extract_run_exports_launch_commit_from_config_or_tags(): @@ -327,7 +327,7 @@ def test_extract_run_uses_xgboost_learning_rate_for_lr_metadata(): "dataset": "miiv", "paradigm": "xgboost", "seed": 42, - "protocol": "B", + "protocol": "full_finetune", "optimizer": {"lr": 3e-4}, "xgboost": {"learning_rate": 0.1}, "task": {"task_name": "mortality_24h"}, @@ -410,7 +410,7 @@ def test_validate_flags_wrong_fixed_seed_set_with_correct_count(): "paradigm": "mae", "dataset": "miiv", "task": "mortality_24h", - "protocol": "B", + "protocol": "full_finetune", "n_seeds": 5, "seed_list": "[1, 2, 3, 4, 5]", } @@ -431,7 +431,9 @@ def test_validate_flags_absent_expected_matrix_rows(): dataset=["miiv"], phase=["finetune"], ) - present = expected[(expected["task"] == "mortality_24h") & (expected["protocol"] == "B")].copy() + 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) diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 4379bce..115eef3 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -1679,8 +1679,8 @@ def test_extract_run_recovers_transfer_from_output_dir(self): assert row["source_dataset"] == "miiv" assert row["experiment_class"] == "cross_dataset_transfer" - def test_extract_run_assigns_protocol_b_to_xgboost_baselines(self): - """XGBoost baselines must align with the Protocol-B comparison family.""" + 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( @@ -1696,7 +1696,7 @@ def test_extract_run_assigns_protocol_b_to_xgboost_baselines(self): name="s11_xgboost_miiv_mortality_24h_seed42", ) - assert extract_run(run, [])["protocol"] == "B" + 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.""" @@ -1793,7 +1793,7 @@ def test_build_aggregated_df_merges_hp_robustness_across_seeds(self): "dataset": "miiv", "task": "mortality_24h", "seed": seed, - "protocol": "B", + "protocol": "full_finetune", "label_fraction": 1.0, "model_size": "default", "source_dataset": None, @@ -2462,7 +2462,7 @@ def test_transfer_finetune_command_propagates_source_dataset(self): cmd = finetune.build_command({pretrain.id: pretrain, finetune.id: finetune}) assert "+source_dataset=miiv" in cmd - def test_protocol_a_finetune_command_uses_strict_linear_probe(self): + def test_linear_probe_finetune_command_uses_strict_linear_probe(self): from scripts.internal.run_experiments import Run pretrain = Run( @@ -2493,11 +2493,12 @@ def test_protocol_a_finetune_command_uses_strict_linear_probe(self): assert "task.head_type=linear" in cmd assert "task.hidden_dims=[]" in cmd assert "task.dropout=0.0" in cmd - assert "protocol=a" in cmd + assert "protocol=linear_probe" in cmd + assert "protocol=a" not in cmd assert "protocol=A" not in cmd assert "+protocol=A" not in cmd - def test_protocol_b_finetune_command_keeps_task_mlp_head(self): + def test_full_finetune_command_keeps_task_mlp_head(self): from scripts.internal.run_experiments import Run pretrain = Run( @@ -2526,7 +2527,8 @@ def test_protocol_b_finetune_command_keeps_task_mlp_head(self): assert "training.freeze_encoder=false" in cmd assert "task.head_type=linear" not in cmd - assert "protocol=b" in cmd + assert "protocol=full_finetune" in cmd + assert "protocol=b" not in cmd assert "protocol=B" not in cmd assert "+protocol=B" not in cmd @@ -2546,10 +2548,11 @@ def test_xgboost_command_does_not_load_protocol_config_group(self): 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_lowercase(self): + def test_generated_finetune_protocol_selectors_are_public_names(self): import scripts.internal.run_experiments as runner runs = runner.generate_all_runs() @@ -2557,30 +2560,31 @@ def test_generated_finetune_protocol_selectors_are_lowercase(self): 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 sum("protocol=a" in command for command in commands) == 825 - assert sum("protocol=b" in command for command in commands) == 1020 + 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): - protocol_a = OmegaConf.load("configs/protocol/a.yaml") - protocol_b = OmegaConf.load("configs/protocol/b.yaml") - - assert protocol_a.protocol == "A" - assert protocol_a.training.freeze_encoder is True - assert protocol_a.training.max_epochs == 50 - assert protocol_a.training.early_stopping_patience == 10 - assert protocol_a.optimizer.lr == pytest.approx(1.0e-4) - assert protocol_a.task.head_type == "linear" - assert list(protocol_a.task.hidden_dims) == [] - assert protocol_a.task.dropout == 0.0 - - assert protocol_b.protocol == "B" - assert protocol_b.training.freeze_encoder is False - assert protocol_b.training.max_epochs == 100 - assert protocol_b.training.early_stopping_patience == 10 - assert protocol_b.optimizer.lr == pytest.approx(3.0e-4) - assert protocol_b.task.head_type == "mlp" - assert list(protocol_b.task.hidden_dims) == [64] - assert protocol_b.task.dropout == 0.3 + 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: @@ -3073,9 +3077,9 @@ def test_build_statistical_tests_df_includes_xgboost_comparisons(self): rows = [] for paradigm, phase, protocol, offset in [ - ("xgboost", "baseline", "B", 0.08), - ("gru_d", "baseline", "B", 0.03), - ("supervised", "supervised", "B", 0.0), + ("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)]: From fe99fb7e0be8453e4e83ecd2b8afadf8663d9872 Mon Sep 17 00:00:00 2001 From: Hannes Ill Date: Thu, 23 Apr 2026 15:48:51 +0000 Subject: [PATCH 114/121] Track codex sessions in auto shutdown --- scripts/internal/auto_shutdown.sh | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/scripts/internal/auto_shutdown.sh b/scripts/internal/auto_shutdown.sh index 4b301af..3e86d11 100755 --- a/scripts/internal/auto_shutdown.sh +++ b/scripts/internal/auto_shutdown.sh @@ -8,10 +8,11 @@ set -euo pipefail IDLE_THRESHOLD_MIN=59 STAMP_FILE="/tmp/slices_last_training_activity" -# Check if any final-run process or Claude Code session is running +# 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; then - # Activity detected — update timestamp and exit + || 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 From ffd6dc52c2ed224da5290cfb5d1a6634b5ba6ab6 Mon Sep 17 00:00:00 2001 From: Hannes Ill Date: Thu, 23 Apr 2026 16:54:50 +0000 Subject: [PATCH 115/121] Handle invalid W&B entity errors cleanly --- scripts/training/finetune.py | 7 +++++- scripts/training/pretrain.py | 7 +++++- scripts/training/supervised.py | 7 +++++- scripts/training/xgboost_baseline.py | 36 +++++++++++++++++++++------- src/slices/training/utils.py | 21 ++++++++++++++-- tests/test_fixes.py | 35 +++++++++++++++++++++++++++ 6 files changed, 99 insertions(+), 14 deletions(-) diff --git a/scripts/training/finetune.py b/scripts/training/finetune.py index 0344fbf..64eb421 100644 --- a/scripts/training/finetune.py +++ b/scripts/training/finetune.py @@ -49,6 +49,7 @@ setup_wandb_logger, train_label_support_summary, validate_data_prerequisites, + WandbEntityNotFoundError, ) @@ -289,7 +290,11 @@ 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}") + return if logger: logger.experiment.summary.update(train_label_support_summary(train_support_stats)) diff --git a/scripts/training/pretrain.py b/scripts/training/pretrain.py index babb3e3..05c7d93 100644 --- a/scripts/training/pretrain.py +++ b/scripts/training/pretrain.py @@ -26,6 +26,7 @@ setup_pretrain_callbacks, setup_wandb_logger, validate_data_prerequisites, + WandbEntityNotFoundError, ) # SSL -> compatible model encoder mappings @@ -151,7 +152,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}") + return trainer = pl.Trainer( max_epochs=cfg.training.max_epochs, diff --git a/scripts/training/supervised.py b/scripts/training/supervised.py index 9356f2a..1233b79 100644 --- a/scripts/training/supervised.py +++ b/scripts/training/supervised.py @@ -47,6 +47,7 @@ setup_wandb_logger, train_label_support_summary, validate_data_prerequisites, + WandbEntityNotFoundError, ) @@ -325,7 +326,11 @@ 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}") + return if logger: logger.experiment.summary.update(train_label_support_summary(train_support_stats)) diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index f652469..ba38e3e 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -168,14 +168,27 @@ def _init_wandb_run(cfg: DictConfig, task_name: str): import wandb run_name, group_name = _wandb_run_name_and_group(cfg, task_name) - run = wandb.init( - project=cfg.logging.wandb_project, - entity=cfg.logging.get("wandb_entity", None), - name=run_name, - group=group_name, - tags=_build_wandb_tags(cfg), - config=OmegaConf.to_container(cfg, resolve=True), - ) + 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 @@ -197,6 +210,7 @@ def main(cfg: DictConfig) -> None: print("=" * 80) from slices.training.utils import ( + WandbEntityNotFoundError, report_and_validate_train_label_support, train_label_support_summary, validate_data_prerequisites, @@ -204,7 +218,11 @@ def main(cfg: DictConfig) -> None: task_name = cfg.task.get("task_name", "mortality_24h") task_type = cfg.task.get("task_type", "binary") - wandb_module, wandb_run = _init_wandb_run(cfg, task_name) + try: + wandb_module, wandb_run = _init_wandb_run(cfg, task_name) + except WandbEntityNotFoundError as exc: + print(f"\nError: {exc}") + return # Validate data prerequisites including label freshness validate_data_prerequisites( diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 6c3acb1..218fec4 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -12,6 +12,7 @@ import torch import torch.nn as nn import yaml +import wandb from lightning.pytorch.callbacks import ( EarlyStopping, LearningRateMonitor, @@ -30,6 +31,10 @@ ) from slices.data.labels import LabelBuilder, LabelBuilderFactory, LabelConfig + +class WandbEntityNotFoundError(ValueError): + """Raised when an explicit W&B entity does not exist.""" + # ============================================================================= # Optimizer / Scheduler # ============================================================================= @@ -430,9 +435,11 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: 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), + entity=wandb_entity, name=run_name, group=group, tags=tags, @@ -440,7 +447,17 @@ def setup_wandb_logger(cfg: DictConfig) -> Optional[WandbLogger]: 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 diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 115eef3..8f6c785 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -2433,6 +2433,41 @@ def __init__(self, **kwargs): 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_transfer_finetune_command_propagates_source_dataset(self): from scripts.internal.run_experiments import Run From d47fe182a25ae4e1a6ffa187a0b641c3bc641568 Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 23 Apr 2026 15:40:36 -0400 Subject: [PATCH 116/121] Harden final experiment launch failures Make W&B entity setup failures exit nonzero so the experiment scheduler cannot mark failed launches as completed. Propagate the explicit launch git-check skip into generated tmux runners and teach the Python runner to honor it for unchecked provenance runs. Behavioral implication: final thesis launches now fail closed on missing W&B entities, while SKIP_LAUNCH_GIT_CHECK=1 remains an explicit audited escape hatch. --- scripts/internal/launch_thesis_tmux.sh | 1 + scripts/internal/run_experiments.py | 17 ++++++ scripts/training/finetune.py | 4 +- scripts/training/pretrain.py | 4 +- scripts/training/supervised.py | 4 +- scripts/training/xgboost_baseline.py | 2 +- tests/test_fixes.py | 82 ++++++++++++++++++++++++-- 7 files changed, 103 insertions(+), 11 deletions(-) diff --git a/scripts/internal/launch_thesis_tmux.sh b/scripts/internal/launch_thesis_tmux.sh index e1e3833..fa410a8 100755 --- a/scripts/internal/launch_thesis_tmux.sh +++ b/scripts/internal/launch_thesis_tmux.sh @@ -248,6 +248,7 @@ echo "Launch commit: $(printf "%q" "$LAUNCH_COMMIT")" echo "Main classes: ${main_classes[*]}" echo "Fairness classes: ${fairness_classes[*]}" echo "Appendix classes: ${appendix_classes[*]:-none}" +export SLICES_SKIP_LAUNCH_GIT_CHECK=$(printf "%q" "$SKIP_LAUNCH_GIT_CHECK") if [[ "$(printf "%q" "$SKIP_LAUNCH_GIT_CHECK")" != "1" ]]; then current_commit="\$(git rev-parse --verify HEAD)" diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index 568cc73..53f2f44 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -985,6 +985,15 @@ def _validate_wandb_online_mode() -> str | None: 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"}: @@ -1010,6 +1019,14 @@ def validate_direct_final_launch_policy(args, parser: argparse.ArgumentParser) - 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 + error = _validate_clean_final_launch_state(str(launch_commit)) if error: parser.error(error) diff --git a/scripts/training/finetune.py b/scripts/training/finetune.py index 64eb421..4415191 100644 --- a/scripts/training/finetune.py +++ b/scripts/training/finetune.py @@ -42,6 +42,7 @@ ) from slices.training import FineTuneModule from slices.training.utils import ( + WandbEntityNotFoundError, report_and_validate_train_label_support, resolve_balanced_class_weights, run_fairness_evaluation, @@ -49,7 +50,6 @@ setup_wandb_logger, train_label_support_summary, validate_data_prerequisites, - WandbEntityNotFoundError, ) @@ -294,7 +294,7 @@ def main(cfg: DictConfig) -> None: logger = setup_wandb_logger(cfg) except WandbEntityNotFoundError as exc: print(f"\nError: {exc}") - return + raise SystemExit(1) from exc if logger: logger.experiment.summary.update(train_label_support_summary(train_support_stats)) diff --git a/scripts/training/pretrain.py b/scripts/training/pretrain.py index 05c7d93..caabbf5 100644 --- a/scripts/training/pretrain.py +++ b/scripts/training/pretrain.py @@ -23,10 +23,10 @@ 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, - WandbEntityNotFoundError, ) # SSL -> compatible model encoder mappings @@ -156,7 +156,7 @@ def main(cfg: DictConfig) -> None: logger = setup_wandb_logger(cfg) except WandbEntityNotFoundError as exc: print(f"\nError: {exc}") - return + raise SystemExit(1) from exc trainer = pl.Trainer( max_epochs=cfg.training.max_epochs, diff --git a/scripts/training/supervised.py b/scripts/training/supervised.py index 1233b79..7d8bea6 100644 --- a/scripts/training/supervised.py +++ b/scripts/training/supervised.py @@ -39,6 +39,7 @@ 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, @@ -47,7 +48,6 @@ setup_wandb_logger, train_label_support_summary, validate_data_prerequisites, - WandbEntityNotFoundError, ) @@ -330,7 +330,7 @@ def main(cfg: DictConfig) -> None: logger = setup_wandb_logger(cfg) except WandbEntityNotFoundError as exc: print(f"\nError: {exc}") - return + raise SystemExit(1) from exc if logger: logger.experiment.summary.update(train_label_support_summary(train_support_stats)) diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index ba38e3e..0b7147e 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -222,7 +222,7 @@ def main(cfg: DictConfig) -> None: wandb_module, wandb_run = _init_wandb_run(cfg, task_name) except WandbEntityNotFoundError as exc: print(f"\nError: {exc}") - return + raise SystemExit(1) from exc # Validate data prerequisites including label freshness validate_data_prerequisites( diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 8f6c785..0fe2924 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -1045,6 +1045,7 @@ def test_tmux_launcher_passes_revision_to_fairness(self, tmp_path): 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.""" @@ -2236,6 +2237,41 @@ def test_main_run_rejects_offline_wandb_mode_for_final_launch(self, monkeypatch) 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 @@ -2434,9 +2470,8 @@ def __init__(self, **kwargs): 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 + import wandb class DummyWandbLogger: def __init__(self, **kwargs): @@ -2445,7 +2480,9 @@ def __init__(self, **kwargs): @property def experiment(self): if self._entity == "missing-team": - raise wandb.errors.CommError("entity missing-team not found during upsertBucket") + raise wandb.errors.CommError( + "entity missing-team not found during upsertBucket" + ) raise AssertionError("unexpected entity") monkeypatch.setattr(training_utils, "WandbLogger", DummyWandbLogger) @@ -2465,9 +2502,46 @@ def experiment(self): } ) - with pytest.raises(training_utils.WandbEntityNotFoundError, match="W&B entity 'missing-team' could not be found"): + 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 From 600d6a6f9a8cf78e7a7e9538aa2bf11f345cdd78 Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 23 Apr 2026 17:22:09 -0400 Subject: [PATCH 117/121] Fix revision retry dependent recovery Apply revision metadata to the full generated retry graph before class filtering, then include skipped downstream dependents caused by selected retry roots. This prevents class-scoped recovery from leaving label-efficiency or transfer runs skipped after an upstream SSL pretrain failure. Adds regression coverage for cross-class skipped dependents in revision-scoped retry. No experiment matrix definitions or benchmark semantics change. --- scripts/internal/run_experiments.py | 34 ++++++++++++-- tests/test_fixes.py | 73 +++++++++++++++++++++++++++++ 2 files changed, 102 insertions(+), 5 deletions(-) diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index 53f2f44..721510c 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -1888,6 +1888,20 @@ def _warn_for_skipped_dependents_after_failed_retry( 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], @@ -1947,11 +1961,7 @@ def cmd_retry(args): class_filter = set(args.experiment_class) if args.experiment_class else None if args.revision: - revised = _filter_requested_runs(all_runs, args.experiment_class) - revised_ids = {run.id for run in revised} - rest = [run for run in all_runs if run.id not in revised_ids] - revised = apply_revision(revised, args.revision, args.reason) - all_runs = rest + revised + 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) @@ -1974,6 +1984,20 @@ def cmd_retry(args): 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 = [ diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 0fe2924..e93bfbf 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -3028,6 +3028,79 @@ def test_cmd_retry_revises_dependencies_for_revision_scoped_retry(self, monkeypa 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 From 926acf472a657976fd7b4ac377361c908b26ec68 Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 23 Apr 2026 17:58:09 -0400 Subject: [PATCH 118/121] Harden final experiment launch guardrails Resolve symbolic launch commits to immutable hashes before recording run metadata, fail closed when the tmux thesis launcher selects a partial final corpus, and stop exposing los_days through model-facing static batch payloads while preserving static_df for analysis. This changes launch behavior for final runs: partial corpus tmux launches now require VALIDATE_FINAL_CORPUS=0, and direct launches record the resolved commit hash. Verified with uv run pytest -q and a 2590-run dry run. --- scripts/internal/launch_thesis_tmux.sh | 57 ++++++++++++++++- scripts/internal/run_experiments.py | 32 ++++++--- src/slices/data/dataset.py | 7 +- tests/test_dataset_datamodule.py | 3 +- tests/test_fixes.py | 89 ++++++++++++++++++++++++++ 5 files changed, 173 insertions(+), 15 deletions(-) diff --git a/scripts/internal/launch_thesis_tmux.sh b/scripts/internal/launch_thesis_tmux.sh index fa410a8..72a88b2 100755 --- a/scripts/internal/launch_thesis_tmux.sh +++ b/scripts/internal/launch_thesis_tmux.sh @@ -20,6 +20,8 @@ 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}}" @@ -55,10 +57,11 @@ else LAUNCH_COMMIT="$(git -C "$REPO_ROOT" rev-parse --verify HEAD)" fi - if ! git -C "$REPO_ROOT" cat-file -e "$LAUNCH_COMMIT^{commit}" 2>/dev/null; then + 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 @@ -187,6 +190,58 @@ 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 " "$@" } diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index 721510c..077661c 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -956,26 +956,37 @@ def _git_quiet(args: list[str]) -> bool: raise RuntimeError(message) -def _validate_clean_final_launch_state(launch_commit: str) -> str | None: - """Return an error string if a final launch would have ambiguous provenance.""" +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 = _git_stdout( - ["rev-parse", "--verify", f"{launch_commit}^{{commit}}"] - ) + 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})" + 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." + "Commit or stash tracked changes first; dry runs may be used for local audits.", ) except RuntimeError as exc: - return f"could not validate final launch git provenance: {exc}" - return None + 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: @@ -1027,9 +1038,10 @@ def validate_direct_final_launch_policy(args, parser: argparse.ArgumentParser) - ) return - error = _validate_clean_final_launch_state(str(launch_commit)) + resolved_launch_commit, error = _validated_final_launch_commit(str(launch_commit)) if error: parser.error(error) + args.launch_commit = resolved_launch_commit # --------------------------------------------------------------------------- diff --git a/src/slices/data/dataset.py b/src/slices/data/dataset.py index e8002c2..8d246e3 100644 --- a/src/slices/data/dataset.py +++ b/src/slices/data/dataset.py @@ -552,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, {}) @@ -561,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") diff --git a/tests/test_dataset_datamodule.py b/tests/test_dataset_datamodule.py index 3f60c4e..2e86e4b 100644 --- a/tests/test_dataset_datamodule.py +++ b/tests/test_dataset_datamodule.py @@ -402,7 +402,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). diff --git a/tests/test_fixes.py b/tests/test_fixes.py index e93bfbf..2c6f3eb 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -1152,6 +1152,48 @@ def test_tmux_launcher_includes_classical_baselines_in_fairness(self, tmp_path): 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] @@ -2183,6 +2225,53 @@ def test_main_run_requires_launch_commit_for_final_launch(self, monkeypatch): 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 From cde200b796203c8ba55988dd3a89a17d4fe9d776 Mon Sep 17 00:00:00 2001 From: hill Date: Thu, 23 Apr 2026 19:18:04 -0400 Subject: [PATCH 119/121] Tighten SSL and training configuration contracts Centralize the controlled SSL tokenization contract, move intrinsic missingness handling behind an encoder capability, and let the experiment runner rely on protocol Hydra configs instead of duplicating overrides. Also tighten data and pretraining config schemas, load FLOPs defaults from YAML, validate imputation data prerequisites, align the ruff pre-commit hook, sort imports, and make setup dependency sync lockfile-frozen. No script names or CLI arguments changed; setup now fails if the lockfile is stale. Verification: uv run ruff check src scripts tests; uv run mypy src; uv run pytest -q. --- .pre-commit-config.yaml | 2 +- README.md | 18 +++-- configs/eval/default.yaml | 2 +- scripts/analyze_labels.py | 1 + scripts/debug/explore_missingness.py | 1 + scripts/debug/export_snapshots.py | 1 + scripts/debug/inspect_embeddings.py | 2 + scripts/eval/evaluate_fairness.py | 2 + scripts/eval/evaluate_imputation.py | 6 +- scripts/export_results.py | 1 + scripts/internal/run_experiments.py | 30 -------- scripts/measure_flops.py | 75 ++++++------------- .../preprocessing/create_combined_dataset.py | 1 + scripts/preprocessing/extract_ricu.py | 1 + scripts/preprocessing/prepare_dataset.py | 1 + scripts/sanity_checks/sc_mae_pretraining.py | 3 +- .../sanity_checks/sc_supervised_learning.py | 3 +- scripts/setup_and_extract.sh | 2 +- scripts/training/finetune.py | 1 + scripts/training/pretrain.py | 1 + scripts/training/supervised.py | 1 + scripts/training/xgboost_baseline.py | 3 +- src/slices/data/config_schemas.py | 16 +++- src/slices/models/encoders/README.md | 6 +- src/slices/models/encoders/__init__.py | 3 +- src/slices/models/encoders/base.py | 38 +++++++++- src/slices/models/encoders/gru_d.py | 4 + src/slices/models/encoders/observation.py | 4 + src/slices/models/encoders/smart.py | 4 + src/slices/models/encoders/transformer.py | 4 + src/slices/models/pretraining/base.py | 45 ++++++++++- src/slices/models/pretraining/contrastive.py | 19 +---- src/slices/models/pretraining/jepa.py | 18 +---- src/slices/models/pretraining/mae.py | 18 +---- src/slices/models/pretraining/ts2vec.py | 14 +--- src/slices/training/checkpoint_loading.py | 21 +----- src/slices/training/config_schemas.py | 16 ++++ src/slices/training/pretrain_module.py | 5 ++ src/slices/training/utils.py | 3 +- tests/conftest.py | 1 + tests/test_aki_builder.py | 1 + tests/test_base_extractor.py | 1 + tests/test_config_schemas.py | 34 +++++++++ tests/test_contrastive_objective.py | 1 + tests/test_dataset_datamodule.py | 1 + tests/test_embeddings.py | 1 + tests/test_encoder_wrapper.py | 1 + tests/test_evaluate_fairness.py | 1 + tests/test_extractor_config.py | 1 + tests/test_extractor_integration.py | 1 + tests/test_factories.py | 1 + tests/test_fairness.py | 1 + tests/test_fairness_evaluator.py | 1 + tests/test_finetune_module.py | 3 + tests/test_fixes.py | 16 ++-- tests/test_gru_d_encoder.py | 1 + tests/test_imputation_eval.py | 5 +- tests/test_inference.py | 1 + tests/test_jepa_objective.py | 1 + tests/test_los_builder.py | 1 + tests/test_mae_objective.py | 1 + tests/test_metrics.py | 3 +- tests/test_model_common.py | 1 + tests/test_package.py | 4 +- tests/test_pretrain_module.py | 9 +++ tests/test_regression_metrics.py | 5 +- tests/test_ricu_extractor.py | 1 + tests/test_sampling.py | 1 + tests/test_sliding_window.py | 1 + tests/test_smart_encoder.py | 1 + tests/test_smart_objective.py | 1 + tests/test_snapshots.py | 1 + tests/test_statistical.py | 3 +- tests/test_task_builders.py | 1 + tests/test_training_utils.py | 1 + tests/test_transformer_encoder.py | 1 + tests/test_transforms.py | 1 + tests/test_ts2vec_objective.py | 1 + 78 files changed, 309 insertions(+), 199 deletions(-) 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/README.md b/README.md index 148ca8a..f8571b5 100644 --- a/README.md +++ b/README.md @@ -163,11 +163,12 @@ contrastive SSL under the shared RICU pipeline, canonical obs-aware 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 the class-based launcher 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. +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, @@ -248,9 +249,11 @@ uv run python scripts/training/pretrain.py dataset=miiv ssl=jepa training.overfi | Group | Options | Purpose | |---|---|---| | `ssl/` | `mae`, `jepa`, `contrastive`, `ts2vec`, `smart` | SSL objective hyperparameters | -| `model/` | `transformer`, `observation_transformer`, `smart` | Encoder architecture | +| `model/` | `transformer`, `transformer_medium`, `transformer_large`, `observation_transformer`, `smart`, `gru_d`, `linear` | Encoder architecture | | `data/` | `ricu` | Dataset and paths | | `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 @@ -302,7 +305,8 @@ The framework uses factory patterns throughout, making it straightforward to add - `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 task definitions +- `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/eval/default.yaml b/configs/eval/default.yaml index 6a481cc..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) diff --git a/scripts/analyze_labels.py b/scripts/analyze_labels.py index d9bc373..133a463 100644 --- a/scripts/analyze_labels.py +++ b/scripts/analyze_labels.py @@ -31,6 +31,7 @@ import numpy as np import polars as pl import yaml + from slices.data.utils import get_package_data_dir diff --git a/scripts/debug/explore_missingness.py b/scripts/debug/explore_missingness.py index a4d1b8c..55f652d 100644 --- a/scripts/debug/explore_missingness.py +++ b/scripts/debug/explore_missingness.py @@ -20,6 +20,7 @@ import numpy as np import polars as pl import yaml + from slices.constants import SEQ_LENGTH_HOURS diff --git a/scripts/debug/export_snapshots.py b/scripts/debug/export_snapshots.py index a96fb09..ec398d7 100644 --- a/scripts/debug/export_snapshots.py +++ b/scripts/debug/export_snapshots.py @@ -12,6 +12,7 @@ import hydra import polars as pl from omegaconf import DictConfig + from slices.debug.snapshots import create_snapshots_from_processed diff --git a/scripts/debug/inspect_embeddings.py b/scripts/debug/inspect_embeddings.py index 44bccc2..ac455ba 100644 --- a/scripts/debug/inspect_embeddings.py +++ b/scripts/debug/inspect_embeddings.py @@ -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 diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index c62c93a..c76e73c 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -50,6 +50,7 @@ from typing import Any, Optional import torch + from slices.constants import ( FULL_FINETUNE_PROTOCOL, canonical_downstream_protocol, @@ -878,6 +879,7 @@ def build_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) diff --git a/scripts/eval/evaluate_imputation.py b/scripts/eval/evaluate_imputation.py index 17fed9f..799a09e 100644 --- a/scripts/eval/evaluate_imputation.py +++ b/scripts/eval/evaluate_imputation.py @@ -27,9 +27,11 @@ import lightning.pytorch as pl import torch from omegaconf import DictConfig, OmegaConf +from torch.utils.data import DataLoader, Subset + from slices.data.datamodule import ICUDataModule, icu_collate_fn from slices.eval.imputation import ImputationEvaluator -from torch.utils.data import DataLoader, Subset +from slices.training.utils import validate_data_prerequisites @hydra.main( @@ -50,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) diff --git a/scripts/export_results.py b/scripts/export_results.py index c15aa62..a3452a4 100644 --- a/scripts/export_results.py +++ b/scripts/export_results.py @@ -37,6 +37,7 @@ import pandas as pd import wandb from scipy import stats as scipy_stats + from slices.constants import ( FULL_FINETUNE_PROTOCOL, canonical_downstream_protocol, diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index 077661c..47ba7bb 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -124,15 +124,6 @@ LAUNCH_IDENTITY_FILE = ".runner_launch_identity.json" LAUNCH_IDENTITY_KEYS = ("revision", "wandb_project", "wandb_entity", "launch_commit") -LINEAR_PROBE_SETTINGS = {"freeze_encoder": True, "max_epochs": 50, "patience": 10, "lr": 1e-4} -FULL_FINETUNE_SETTINGS = { - "freeze_encoder": False, - "max_epochs": 100, - "patience": 10, - "lr": 3e-4, -} - - # --------------------------------------------------------------------------- # Data Model # --------------------------------------------------------------------------- @@ -253,27 +244,6 @@ def _finetune_cmd(self, runs_by_id: dict[str, "Run"]) -> list[str]: f"hydra.run.dir={self.output_dir}", ] self._append_resume(cmd) - if self.freeze_encoder is True: - cmd.extend( - [ - "training.freeze_encoder=true", - f"training.max_epochs={LINEAR_PROBE_SETTINGS['max_epochs']}", - f"training.early_stopping_patience={LINEAR_PROBE_SETTINGS['patience']}", - f"optimizer.lr={LINEAR_PROBE_SETTINGS['lr']}", - "task.head_type=linear", - "task.hidden_dims=[]", - "task.dropout=0.0", - ] - ) - elif self.freeze_encoder is False: - cmd.extend( - [ - "training.freeze_encoder=false", - f"training.max_epochs={FULL_FINETUNE_SETTINGS['max_epochs']}", - f"training.early_stopping_patience={FULL_FINETUNE_SETTINGS['patience']}", - f"optimizer.lr={FULL_FINETUNE_SETTINGS['lr']}", - ] - ) if self.label_fraction < 1.0: cmd.append(f"label_fraction={self.label_fraction}") if self.source_dataset is not None: diff --git a/scripts/measure_flops.py b/scripts/measure_flops.py index e638c24..b0c77cf 100644 --- a/scripts/measure_flops.py +++ b/scripts/measure_flops.py @@ -13,77 +13,46 @@ 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 -from torch.utils.flop_counter import FlopCounterMode # --------------------------------------------------------------------------- -# Paradigm configs — mirror the YAML defaults used in actual runs +# 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, - "d_model": 64, "max_seq_length": SEQ_LENGTH_HOURS, - "n_layers": 2, - "n_heads": 4, - "d_ff": 256, - "dropout": 0.1, - "activation": "gelu", - "prenorm": True, - "layer_norm_eps": 1e-5, - "use_positional_encoding": True, "pooling": "none", # Required for SSL pretraining - "obs_aware": True, # All SSL paradigms use obs-aware tokenization + "obs_aware": True, # All controlled SSL paradigms use obs-aware tokenization } -# SSL configs — exact values from configs/ssl/*.yaml SSL_CONFIGS: Dict[str, dict] = { - "mae": { - "mask_ratio": 0.5, - "decoder_d_model": 64, - "decoder_n_layers": 2, - "decoder_n_heads": 4, - "decoder_d_ff": 256, - "decoder_dropout": 0.1, - }, - "jepa": { - "mask_ratio": 0.5, - "mask_strategy": "block", - "mask_n_blocks": 3, - "predictor_d_model": 32, - "predictor_n_layers": 2, - "predictor_n_heads": 4, - "predictor_d_ff": 128, - "predictor_dropout": 0.1, - "momentum_base": 0.999, - "momentum_final": 1.0, - "loss_type": "mse", - }, - "contrastive": { - "mode": "instance", - "mask_ratio": 0.5, - "proj_hidden_dim": 256, - "proj_output_dim": 64, - "temperature": 0.07, - "complementary_masks": True, - }, - "ts2vec": { - "mask_ratio": 0.5, - "noise_scale": 0.01, - "crop_ratio": 1.0, - "proj_hidden_dim": 256, - "proj_output_dim": 64, - "temperature": 0.05, - "n_hierarchical_scales": 4, - }, + name: load_yaml_config(CONFIG_DIR / "ssl" / f"{name}.yaml") for name in PARADIGMS } @@ -243,8 +212,6 @@ def main(): if not obs_mask[b, t].any(): obs_mask[b, t, 0] = True - paradigms = ["mae", "jepa", "contrastive", "ts2vec"] - 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() @@ -268,7 +235,7 @@ def main(): # Measure each paradigm rows = [] - for name in paradigms: + for name in PARADIGMS: objective = build_objective(name, encoder_config) trainable, total = count_parameters(objective) diff --git a/scripts/preprocessing/create_combined_dataset.py b/scripts/preprocessing/create_combined_dataset.py index e3b7582..db910a4 100644 --- a/scripts/preprocessing/create_combined_dataset.py +++ b/scripts/preprocessing/create_combined_dataset.py @@ -23,6 +23,7 @@ import polars as pl import yaml + from slices.constants import MIN_STAY_HOURS, SEQ_LENGTH_HOURS, THESIS_TASKS from slices.data.preparation import prepare_processed_dataset diff --git a/scripts/preprocessing/extract_ricu.py b/scripts/preprocessing/extract_ricu.py index 11d1d31..c3fd8fe 100644 --- a/scripts/preprocessing/extract_ricu.py +++ b/scripts/preprocessing/extract_ricu.py @@ -21,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 diff --git a/scripts/preprocessing/prepare_dataset.py b/scripts/preprocessing/prepare_dataset.py index c787d7a..ab61276 100644 --- a/scripts/preprocessing/prepare_dataset.py +++ b/scripts/preprocessing/prepare_dataset.py @@ -4,6 +4,7 @@ import hydra from omegaconf import DictConfig + from slices.data.preparation import prepare_processed_dataset diff --git a/scripts/sanity_checks/sc_mae_pretraining.py b/scripts/sanity_checks/sc_mae_pretraining.py index 5ea9068..21bdbd4 100644 --- a/scripts/sanity_checks/sc_mae_pretraining.py +++ b/scripts/sanity_checks/sc_mae_pretraining.py @@ -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 diff --git a/scripts/sanity_checks/sc_supervised_learning.py b/scripts/sanity_checks/sc_supervised_learning.py index 99cac8c..2f121a9 100644 --- a/scripts/sanity_checks/sc_supervised_learning.py +++ b/scripts/sanity_checks/sc_supervised_learning.py @@ -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 diff --git a/scripts/setup_and_extract.sh b/scripts/setup_and_extract.sh index debafd6..852d52f 100755 --- a/scripts/setup_and_extract.sh +++ b/scripts/setup_and_extract.sh @@ -225,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 diff --git a/scripts/training/finetune.py b/scripts/training/finetune.py index 4415191..e2d191e 100644 --- a/scripts/training/finetune.py +++ b/scripts/training/finetune.py @@ -34,6 +34,7 @@ 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, diff --git a/scripts/training/pretrain.py b/scripts/training/pretrain.py index caabbf5..075b946 100644 --- a/scripts/training/pretrain.py +++ b/scripts/training/pretrain.py @@ -19,6 +19,7 @@ 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 diff --git a/scripts/training/supervised.py b/scripts/training/supervised.py index 7d8bea6..a8cd6ce 100644 --- a/scripts/training/supervised.py +++ b/scripts/training/supervised.py @@ -30,6 +30,7 @@ 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, diff --git a/scripts/training/xgboost_baseline.py b/scripts/training/xgboost_baseline.py index 0b7147e..7b6c1e8 100644 --- a/scripts/training/xgboost_baseline.py +++ b/scripts/training/xgboost_baseline.py @@ -29,6 +29,8 @@ 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 ( @@ -37,7 +39,6 @@ file_sha256, ) from slices.eval.inference import extract_tabular_features -from xgboost import XGBClassifier, XGBRegressor def _xgboost_eval_metric(task_type: str) -> str: diff --git a/src/slices/data/config_schemas.py b/src/slices/data/config_schemas.py index 83f44c4..91d1c1a 100644 --- a/src/slices/data/config_schemas.py +++ b/src/slices/data/config_schemas.py @@ -181,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) @@ -211,7 +212,12 @@ class DataConfig(BaseModel): # 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": @@ -236,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/models/encoders/README.md b/src/slices/models/encoders/README.md index 721fe05..b250ec9 100644 --- a/src/slices/models/encoders/README.md +++ b/src/slices/models/encoders/README.md @@ -131,8 +131,4 @@ Relevant test coverage lives in: - `tests/test_smart_encoder.py` - `tests/test_factories.py` -Example script: - -```bash -uv run python examples/transformer_encoder_example.py -``` +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 956d1df..a39de3f 100644 --- a/src/slices/models/encoders/__init__.py +++ b/src/slices/models/encoders/__init__.py @@ -1,6 +1,6 @@ """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 @@ -21,6 +21,7 @@ "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/gru_d.py b/src/slices/models/encoders/gru_d.py index 5a7e225..2228a77 100644 --- a/src/slices/models/encoders/gru_d.py +++ b/src/slices/models/encoders/gru_d.py @@ -66,6 +66,10 @@ def __init__(self, config: GRUDConfig) -> None: 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) diff --git a/src/slices/models/encoders/observation.py b/src/slices/models/encoders/observation.py index 48cc25c..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( diff --git a/src/slices/models/encoders/smart.py b/src/slices/models/encoders/smart.py index dfb2252..0854015 100644 --- a/src/slices/models/encoders/smart.py +++ b/src/slices/models/encoders/smart.py @@ -433,6 +433,10 @@ def __init__(self, config: SMARTEncoderConfig) -> None: ] ) + 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.""" if self.config.d_model % self.config.n_heads != 0: diff --git a/src/slices/models/encoders/transformer.py b/src/slices/models/encoders/transformer.py index 4520d4d..be5ac00 100644 --- a/src/slices/models/encoders/transformer.py +++ b/src/slices/models/encoders/transformer.py @@ -234,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: 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 cf6bf50..38144ce 100644 --- a/src/slices/models/pretraining/contrastive.py +++ b/src/slices/models/pretraining/contrastive.py @@ -37,7 +37,7 @@ import torch.nn as nn import torch.nn.functional as F -from .base import BaseSSLObjective, SSLConfig +from .base import BaseSSLObjective, SSLConfig, require_ssl_tokenizing_encoder from .masking import ( create_complementary_timestep_masks, create_timestep_mask, @@ -128,22 +128,7 @@ def __init__(self, encoder: nn.Module, config: ContrastiveConfig) -> None: super().__init__(encoder, config) self.config: ContrastiveConfig = config - # Validate encoder has obs-aware tokenization - if not getattr(getattr(encoder, "config", None), "obs_aware", False): - raise ValueError( - "Contrastive requires an encoder with obs_aware=True " - "(e.g., TransformerEncoder with obs_aware=True). 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() diff --git a/src/slices/models/pretraining/jepa.py b/src/slices/models/pretraining/jepa.py index 135a033..81489bb 100644 --- a/src/slices/models/pretraining/jepa.py +++ b/src/slices/models/pretraining/jepa.py @@ -25,7 +25,7 @@ from slices.models.common import build_sinusoidal_pe -from .base import BaseSSLObjective, SSLConfig +from .base import BaseSSLObjective, SSLConfig, require_ssl_tokenizing_encoder from .masking import ( create_block_timestep_mask, create_timestep_mask, @@ -184,21 +184,7 @@ def __init__(self, encoder: nn.Module, config: JEPAConfig) -> None: super().__init__(encoder, config) self.config: JEPAConfig = config - # Validate encoder has obs-aware tokenization - if not getattr(getattr(encoder, "config", None), "obs_aware", False): - raise ValueError( - "JEPA requires an encoder with obs_aware=True " - "(e.g., TransformerEncoder with obs_aware=True). 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() self.d_encoder = d_encoder diff --git a/src/slices/models/pretraining/mae.py b/src/slices/models/pretraining/mae.py index 75750cc..aec636f 100644 --- a/src/slices/models/pretraining/mae.py +++ b/src/slices/models/pretraining/mae.py @@ -19,7 +19,7 @@ from slices.models.common import build_sinusoidal_pe -from .base import BaseSSLObjective, SSLConfig +from .base import BaseSSLObjective, SSLConfig, require_ssl_tokenizing_encoder from .masking import create_timestep_mask, extract_visible_timesteps, scatter_visible_timesteps @@ -177,21 +177,7 @@ def __init__(self, encoder: nn.Module, config: MAEConfig) -> None: super().__init__(encoder, config) self.config: MAEConfig = config - # Validate encoder has obs-aware tokenization - if not getattr(getattr(encoder, "config", None), "obs_aware", False): - raise ValueError( - "MAE requires an encoder with obs_aware=True " - "(e.g., TransformerEncoder with obs_aware=True). 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 diff --git a/src/slices/models/pretraining/ts2vec.py b/src/slices/models/pretraining/ts2vec.py index ea661bd..74844f2 100644 --- a/src/slices/models/pretraining/ts2vec.py +++ b/src/slices/models/pretraining/ts2vec.py @@ -36,7 +36,7 @@ import torch.nn as nn import torch.nn.functional as F -from .base import BaseSSLObjective, SSLConfig +from .base import BaseSSLObjective, SSLConfig, require_ssl_tokenizing_encoder from .masking import create_timestep_mask, extract_visible_timesteps, scatter_visible_timesteps @@ -110,17 +110,7 @@ def __init__(self, encoder: nn.Module, config: TS2VecConfig) -> None: super().__init__(encoder, config) self.config: TS2VecConfig = config - # Validate encoder - if not getattr(getattr(encoder, "config", None), "obs_aware", False): - raise ValueError( - "TS2Vec requires an encoder with obs_aware=True. " f"Got: {type(encoder).__name__}" - ) - encoder_pooling = getattr(encoder.config, "pooling", "none") - if encoder_pooling != "none": - raise ValueError( - "TS2Vec requires encoder with pooling='none' for per-token " - f"representations, but got pooling='{encoder_pooling}'" - ) + require_ssl_tokenizing_encoder(encoder, "TS2Vec") d_encoder = encoder.get_output_dim() self.missing_token = None diff --git a/src/slices/training/checkpoint_loading.py b/src/slices/training/checkpoint_loading.py index 1807889..6893248 100644 --- a/src/slices/training/checkpoint_loading.py +++ b/src/slices/training/checkpoint_loading.py @@ -14,14 +14,8 @@ import torch.nn as nn from omegaconf import DictConfig -from slices.models.encoders import ( - EncoderWithMissingToken, - GRUDEncoder, - ObservationTransformerEncoder, - SMARTEncoder, - TransformerEncoder, - build_encoder, -) +from slices.models.encoders.factory import build_encoder +from slices.models.encoders.wrapper import EncoderWithMissingToken logger = logging.getLogger(__name__) @@ -77,8 +71,8 @@ def wrap_encoder_with_missing_token( Returns: The encoder, possibly wrapped with EncoderWithMissingToken. """ - # Skip for encoders that handle missingness intrinsically - if isinstance(encoder, (ObservationTransformerEncoder, SMARTEncoder, GRUDEncoder)): + 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)", @@ -86,13 +80,6 @@ def wrap_encoder_with_missing_token( ) return encoder - if isinstance(encoder, TransformerEncoder) and getattr(encoder.config, "obs_aware", False): - logger.info( - "Skipping EncoderWithMissingToken wrapper for obs_aware " - "TransformerEncoder (forward() handles mask via obs_proj)" - ) - return encoder - d_input = encoder.config.d_input if missing_token is not None: diff --git a/src/slices/training/config_schemas.py b/src/slices/training/config_schemas.py index b937d38..4be294d 100644 --- a/src/slices/training/config_schemas.py +++ b/src/slices/training/config_schemas.py @@ -102,6 +102,22 @@ class TrainingConfig(BaseModel): 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 + + class OptimizerConfig(BaseModel): """Validated optimizer configuration.""" diff --git a/src/slices/training/pretrain_module.py b/src/slices/training/pretrain_module.py index c1705e9..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) diff --git a/src/slices/training/utils.py b/src/slices/training/utils.py index 218fec4..9ada5f1 100644 --- a/src/slices/training/utils.py +++ b/src/slices/training/utils.py @@ -11,8 +11,8 @@ import torch import torch.nn as nn -import yaml import wandb +import yaml from lightning.pytorch.callbacks import ( EarlyStopping, LearningRateMonitor, @@ -35,6 +35,7 @@ class WandbEntityNotFoundError(ValueError): """Raised when an explicit W&B entity does not exist.""" + # ============================================================================= # Optimizer / Scheduler # ============================================================================= diff --git a/tests/conftest.py b/tests/conftest.py index ebe7f70..97a9b77 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -9,6 +9,7 @@ import polars as pl import pytest import torch + from slices.constants import SEQ_LENGTH_HOURS diff --git a/tests/test_aki_builder.py b/tests/test_aki_builder.py index ec8f50a..4be7013 100644 --- a/tests/test_aki_builder.py +++ b/tests/test_aki_builder.py @@ -2,6 +2,7 @@ import polars as pl import pytest + from slices.data.labels import AKILabelBuilder, LabelBuilderFactory, LabelConfig diff --git a/tests/test_base_extractor.py b/tests/test_base_extractor.py index 069b6bc..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 diff --git a/tests/test_config_schemas.py b/tests/test_config_schemas.py index 71f2a1b..5976559 100644 --- a/tests/test_config_schemas.py +++ b/tests/test_config_schemas.py @@ -9,10 +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, @@ -133,6 +135,24 @@ def test_split_ratios_must_sum_to_one(self): 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.""" @@ -181,6 +201,20 @@ def test_allow_best_ckpt_fallback_field_exists(self): 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.""" diff --git a/tests/test_contrastive_objective.py b/tests/test_contrastive_objective.py index eb0e704..8f8687e 100644 --- a/tests/test_contrastive_objective.py +++ b/tests/test_contrastive_objective.py @@ -2,6 +2,7 @@ import pytest import torch + from slices.models.encoders import ( TransformerConfig, TransformerEncoder, diff --git a/tests/test_dataset_datamodule.py b/tests/test_dataset_datamodule.py index 2e86e4b..53b69a6 100644 --- a/tests/test_dataset_datamodule.py +++ b/tests/test_dataset_datamodule.py @@ -8,6 +8,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.splits import compute_patient_level_splits, load_cached_splits 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 index eeb8ab3..3d62283 100644 --- a/tests/test_evaluate_fairness.py +++ b/tests/test_evaluate_fairness.py @@ -6,6 +6,7 @@ import pandas as pd import pytest + from slices.eval.fairness_metadata import ( EVAL_ARTIFACT_SHA256_KEY, FAIRNESS_ARTIFACT_PATH_KEY, diff --git a/tests/test_extractor_config.py b/tests/test_extractor_config.py index ddf873a..70ce8d8 100644 --- a/tests/test_extractor_config.py +++ b/tests/test_extractor_config.py @@ -2,6 +2,7 @@ import pytest from pydantic import ValidationError + from slices.constants import THESIS_TASKS from slices.data.extractors.base import ExtractorConfig 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 41f53a8..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, 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 b974530..fba4770 100644 --- a/tests/test_fairness_evaluator.py +++ b/tests/test_fairness_evaluator.py @@ -13,6 +13,7 @@ import polars as pl import pytest import torch + from slices.eval.fairness_evaluator import FairnessEvaluator, flatten_fairness_report diff --git a/tests/test_finetune_module.py b/tests/test_finetune_module.py index 7c32f9c..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, @@ -690,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).""" @@ -733,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: diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 2c6f3eb..6657437 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -21,6 +21,7 @@ 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 @@ -2559,9 +2560,10 @@ def __init__(self, **kwargs): assert "ablation:transfer" in captured["kwargs"]["tags"] def test_setup_wandb_logger_raises_clear_error_for_missing_entity(self, monkeypatch): - import slices.training.utils as training_utils import wandb + import slices.training.utils as training_utils + class DummyWandbLogger: def __init__(self, **kwargs): self._entity = kwargs["entity"] @@ -2687,11 +2689,11 @@ def test_linear_probe_finetune_command_uses_strict_linear_probe(self): cmd = probe.build_command({pretrain.id: pretrain, probe.id: probe}) - assert "training.freeze_encoder=true" in cmd - assert "task.head_type=linear" in cmd - assert "task.hidden_dims=[]" in cmd - assert "task.dropout=0.0" in cmd 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 @@ -2723,9 +2725,9 @@ def test_full_finetune_command_keeps_task_mlp_head(self): cmd = finetune.build_command({pretrain.id: pretrain, finetune.id: finetune}) - assert "training.freeze_encoder=false" in cmd - assert "task.head_type=linear" not in cmd 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 diff --git a/tests/test_gru_d_encoder.py b/tests/test_gru_d_encoder.py index 47d9512..473ecd3 100644 --- a/tests/test_gru_d_encoder.py +++ b/tests/test_gru_d_encoder.py @@ -3,6 +3,7 @@ import math import torch + from slices.models.encoders.gru_d import GRUDConfig, GRUDEncoder diff --git a/tests/test_imputation_eval.py b/tests/test_imputation_eval.py index dc397a3..bb4b33b 100644 --- a/tests/test_imputation_eval.py +++ b/tests/test_imputation_eval.py @@ -15,9 +15,10 @@ import torch import torch.nn as nn from omegaconf import OmegaConf -from slices.eval.imputation import ImputationEvaluator from torch.utils.data import DataLoader +from slices.eval.imputation import ImputationEvaluator + class SimpleEncoder(nn.Module): """Minimal encoder for testing.""" @@ -525,6 +526,7 @@ def evaluate(self, dataloader, mask_strategy, mask_ratio, feature_stds): 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) @@ -536,6 +538,7 @@ def evaluate(self, dataloader, mask_strategy, mask_ratio, feature_stds): cfg = OmegaConf.create( { + "dataset": "miiv", "seed": 42, "data": { "processed_dir": "data/processed/miiv", diff --git a/tests/test_inference.py b/tests/test_inference.py index 5cb3b3e..ca2e76d 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -1,6 +1,7 @@ """Tests for shared inference utilities.""" import torch + from slices.eval.inference import run_inference diff --git a/tests/test_jepa_objective.py b/tests/test_jepa_objective.py index 1ba83c6..37ef8d7 100644 --- a/tests/test_jepa_objective.py +++ b/tests/test_jepa_objective.py @@ -2,6 +2,7 @@ import pytest import torch + from slices.models.encoders import ( TransformerConfig, TransformerEncoder, 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 6356fdd..5234782 100644 --- a/tests/test_mae_objective.py +++ b/tests/test_mae_objective.py @@ -2,6 +2,7 @@ import pytest import torch + from slices.data.transforms import ( MaskingStrategy, apply_mask, diff --git a/tests/test_metrics.py b/tests/test_metrics.py index e87d69c..335fdd2 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -11,6 +11,8 @@ import numpy as np import pytest import torch +from torchmetrics import MetricCollection + from slices.eval.metrics import ( AVAILABLE_METRICS, DEFAULT_METRICS, @@ -19,7 +21,6 @@ build_metrics, get_default_metrics, ) -from torchmetrics import MetricCollection class TestMetricConfig: diff --git a/tests/test_model_common.py b/tests/test_model_common.py index 233556e..0436dd7 100644 --- a/tests/test_model_common.py +++ b/tests/test_model_common.py @@ -11,6 +11,7 @@ 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, 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 a08faf7..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 @@ -519,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 5b9063f..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 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 9de5e44..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, diff --git a/tests/test_smart_objective.py b/tests/test_smart_objective.py index 9cfd6fb..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 ( diff --git a/tests/test_snapshots.py b/tests/test_snapshots.py index 8520e2a..a1f9c78 100644 --- a/tests/test_snapshots.py +++ b/tests/test_snapshots.py @@ -1,6 +1,7 @@ from pathlib import Path import polars as pl + from slices.debug.snapshots import ( LegacyPipelineStage, capture_dense_snapshot, diff --git a/tests/test_statistical.py b/tests/test_statistical.py index 47f39fb..f27983e 100644 --- a/tests/test_statistical.py +++ b/tests/test_statistical.py @@ -10,6 +10,8 @@ import pandas as pd import pytest import torch +from torchmetrics import AUROC + from slices.eval.statistical import ( bonferroni_correction, bootstrap_ci, @@ -17,7 +19,6 @@ paired_bootstrap_test, paired_wilcoxon_signed_rank, ) -from torchmetrics import AUROC # ── Helpers ────────────────────────────────────────────────────────────────── diff --git a/tests/test_task_builders.py b/tests/test_task_builders.py index 834aa86..814511d 100644 --- a/tests/test_task_builders.py +++ b/tests/test_task_builders.py @@ -7,6 +7,7 @@ import polars as pl import pytest + from slices.data.labels import LabelBuilderFactory, LabelConfig from slices.data.labels.mortality import MortalityLabelBuilder diff --git a/tests/test_training_utils.py b/tests/test_training_utils.py index 3d7dcb5..c60814c 100644 --- a/tests/test_training_utils.py +++ b/tests/test_training_utils.py @@ -15,6 +15,7 @@ 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 ( diff --git a/tests/test_transformer_encoder.py b/tests/test_transformer_encoder.py index 546df7c..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 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 index 91c394a..32e9324 100644 --- a/tests/test_ts2vec_objective.py +++ b/tests/test_ts2vec_objective.py @@ -2,6 +2,7 @@ import pytest import torch + from slices.models.encoders import TransformerConfig, TransformerEncoder from slices.models.pretraining import TS2VecConfig, TS2VecObjective From 24cbc24681b41a074075b343e9e7846423899b3c Mon Sep 17 00:00:00 2001 From: hill Date: Fri, 24 Apr 2026 15:28:50 -0400 Subject: [PATCH 120/121] Harden final experiment guardrails Replace existing W&B fairness summary keys during recomputation so stale subgroup metrics cannot leak into publication exports. Validate provided launch commits during dry runs, recover stale running entries from status, and update the experiment plan protocol metadata. Behavioral implications: fairness recomputation now clears old fairness summary keys by default, and status may update runner state when it detects dead or mismatched running processes. --- docs/internal/EXPERIMENT_PLAN.md | 2 +- scripts/eval/evaluate_fairness.py | 10 ++- scripts/internal/run_experiments.py | 18 ++++- tests/test_evaluate_fairness.py | 44 +++++++++++ tests/test_fixes.py | 114 +++++++++++++++++++++++++++- 5 files changed, 180 insertions(+), 8 deletions(-) diff --git a/docs/internal/EXPERIMENT_PLAN.md b/docs/internal/EXPERIMENT_PLAN.md index deb437b..655ad09 100644 --- a/docs/internal/EXPERIMENT_PLAN.md +++ b/docs/internal/EXPERIMENT_PLAN.md @@ -70,7 +70,7 @@ Every final run must carry class-based metadata: | `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` | `A` for linear probe, `B` for full finetune/full-training where applicable | +| `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 | diff --git a/scripts/eval/evaluate_fairness.py b/scripts/eval/evaluate_fairness.py index c76e73c..7a57de3 100644 --- a/scripts/eval/evaluate_fairness.py +++ b/scripts/eval/evaluate_fairness.py @@ -582,9 +582,14 @@ def write_fairness_to_wandb( run_path: str, fairness_flat: dict[str, Any], dry_run: bool = False, - clear_existing: bool = False, + clear_existing: bool = True, ) -> None: - """Write fairness metrics to W&B run summary.""" + """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: @@ -1285,7 +1290,6 @@ def _sort_key(r): run_path, fairness_flat, args.dry_run, - clear_existing=args.force, ) results["processed"] += 1 diff --git a/scripts/internal/run_experiments.py b/scripts/internal/run_experiments.py index 47ba7bb..0988c91 100644 --- a/scripts/internal/run_experiments.py +++ b/scripts/internal/run_experiments.py @@ -9,10 +9,12 @@ --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 + --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 + --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 \ @@ -989,6 +991,11 @@ def validate_direct_final_launch_policy(args, parser: argparse.ArgumentParser) - "--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: @@ -1402,8 +1409,9 @@ def _pid_matches_command(pid: int, expected_command: str | None) -> bool: ) -def recover_stale_running(state: dict) -> None: +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") @@ -1416,6 +1424,8 @@ def recover_stale_running(state: dict) -> None: info.pop("pid", None) info.pop("command", None) print(f" Recovered stale run: {run_id}") + recovered += 1 + return recovered RUN_SLOT_COSTS = { @@ -1653,6 +1663,8 @@ def print_status( 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]] = {} diff --git a/tests/test_evaluate_fairness.py b/tests/test_evaluate_fairness.py index 3d62283..d08e7ea 100644 --- a/tests/test_evaluate_fairness.py +++ b/tests/test_evaluate_fairness.py @@ -182,6 +182,50 @@ def test_build_fairness_summary_metadata_stamps_artifact_and_settings(tmp_path): 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") diff --git a/tests/test_fixes.py b/tests/test_fixes.py index 6657437..7da56f2 100644 --- a/tests/test_fixes.py +++ b/tests/test_fixes.py @@ -2200,6 +2200,86 @@ def test_main_run_dry_run_warns_without_launch_commit(self, monkeypatch, capsys) 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 @@ -2462,12 +2542,44 @@ def test_recover_stale_running_resets_pid_reuse(self, monkeypatch): monkeypatch.setattr(runner, "is_pid_alive", lambda pid: True) monkeypatch.setattr(runner, "_pid_matches_command", lambda pid, command: False) - runner.recover_stale_running(state) + 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 From 4d640250b7d3d3670f10a1795167600f13f2969a Mon Sep 17 00:00:00 2001 From: hill Date: Fri, 15 May 2026 20:41:03 -0700 Subject: [PATCH 121/121] Clarify SSL objective scope Update the pretraining README to distinguish the core controlled SSL objectives from TS2Vec as a temporal-contrastive extension. This keeps the documentation aligned with the thesis benchmark contract without changing runtime behavior. --- src/slices/models/pretraining/README.md | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/slices/models/pretraining/README.md b/src/slices/models/pretraining/README.md index 95fe90d..e3ec45b 100644 --- a/src/slices/models/pretraining/README.md +++ b/src/slices/models/pretraining/README.md @@ -35,15 +35,17 @@ These constraints are enforced in the objective constructors and by ## Benchmark Design -The controlled benchmark objectives share the same timestep-level obs-aware -transformer encoder and differ only in the SSL objective and masking logic. +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. Highlights: - 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 stronger temporal objective +- 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 ## Example @@ -97,6 +99,8 @@ pretraining/ - 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