Fit omics batch, covariate, and smoking corrections on train only

README.md

@@ -65,7 +65,9 @@ Omics datasets use `dataset.split_params.split_type` (default **`fixed`**):
 - **`fixed`** — shuffle with seed 42, then cut 70 / 15 / 15. Graph caches keep the historical path.
 - **`k-fold`** — stratified 3/1/1 rotation over `k` folds (`k` defaults to 5 → about 60 / 20 / 20). `data_seed` is the **test fold**; validation is the next fold. Each sample is test once and validation once across folds `0 .. k-1`.
 
-Imputation, gene selection, adjacency, and feature normalization are always fit on **training samples only**, then applied to val/test. Optional `dataset.split_params.grouping` (for example `batch`) keeps whole groups inside a single k-fold fold; it is ignored for `fixed` splits.
+Imputation, gene selection, adjacency, and feature normalization are always fit on **training samples only**, then applied to val/test.
+
+MoTrPAC covariate adjustment, AddNeuroMed ComBat, and smoking promoter-probe pick plus median-centering are also train-only. Hub matrices are uncorrected; sidecars are `motrpac_covariates.parquet`, `addneuromed_batches.parquet`, and `smoking_probe_map.parquet`. Defaults are `corrections=[covariate_adjust]` (MoTrPAC), `corrections=[combat]` (AddNeuroMed), and `corrections=[promoter_min_beta, median_center]` (smoking). Smoking Hub columns are TSS1500/TSS200 probes; after the split, each gene keeps the candidate probe with the lowest mean beta on **training** never-smokers, remaining NaNs are imputed with training column means, and each gene is median-centered on train. Parkinson GEO characteristics are in `parkinsons_sample_meta.parquet`, including the hybridization-date `batch` field. Parkinson defaults to `dataset.split_params.grouping=batch`, so `StratifiedGroupKFold` keeps every batch inside a single fold and no batch is split across train / val / test. Batch sizes are very uneven (70 down to 1), so grouped folds are not equal sized: for `k=5` the test fold ranges from 95 to 120 samples. `grouping` applies to `k-fold` only and is ignored for `fixed`. Grouped caches are stored separately (`..._group_batch`).
 
 ```bash
 python -m ogbench dataset=brca model=gcn

configs/dataset/addneuromed.yaml

@@ -23,12 +23,13 @@ loader:
     k: ${dataset.split_params.k}
     fold: ${dataset.split_params.data_seed}
     grouping: ${dataset.split_params.grouping}
-    corrections: []
+    corrections:
+      - combat
 parameters:
   num_features: 1
   num_classes: 3
   num_samples: 711
-  full_num_nodes: 17198
+  full_num_nodes: 17197
   num_nodes: ${calculate_num_nodes:${dataset.parameters.num_samples},${dataset.loader.parameters.train_val_test_split},${dataset.loader.parameters.node_sample_ratio},${dataset.parameters.full_num_nodes},${dataset.split_params.split_type},${dataset.split_params.k}}
   task: classification
   loss_type: cross_entropy

configs/dataset/motrpac.yaml

@@ -23,7 +23,8 @@ loader:
     k: ${dataset.split_params.k}
     fold: ${dataset.split_params.data_seed}
     grouping: ${dataset.split_params.grouping}
-    corrections: []
+    corrections:
+      - covariate_adjust
 parameters:
   num_features: 1
   num_classes: 2

configs/dataset/parkinsons.yaml

@@ -51,7 +51,7 @@ split_params:
   data_seed: 0 # k-fold test-fold index; validation is the next fold
   split_type: fixed # or k-fold
   k: 5
-  grouping: null
+  grouping: batch # hybridization-date batches stay unmixed across k-fold splits
 dataloader_params:
   batch_size: 16
   num_workers: 0

configs/dataset/smoking.yaml

@@ -19,12 +19,14 @@ loader:
     k: ${dataset.split_params.k}
     fold: ${dataset.split_params.data_seed}
     grouping: ${dataset.split_params.grouping}
-    corrections: []
+    corrections:
+      - promoter_min_beta
+      - median_center
 parameters:
   num_features: 1
   num_classes: 2
   num_samples: 464
-  full_num_nodes: 20763
+  full_num_nodes: 139125
   num_nodes: ${calculate_num_nodes:${dataset.parameters.num_samples},${dataset.loader.parameters.train_val_test_split},${dataset.loader.parameters.node_sample_ratio},${dataset.parameters.full_num_nodes},${dataset.split_params.split_type},${dataset.split_params.k}}
   task: classification
   loss_type: cross_entropy

configs/hf/default.yaml

@@ -1,4 +1,6 @@
 # HuggingFace configuration
 # Single source of truth for the dataset revision hash.
 # Override per-dataset in dataset configs if needed.
-revision: 4da96d838e81dc3f3da3c559925eea9bd356111e
+# Uncorrected matrices plus sidecars (covariates, batches, Parkinson sample_meta,
+# smoking promoter probe map).
+revision: 056dfdc4f434fd35355ffbe5f7b63910d785a97a

ogbench/baseline.py

@@ -34,6 +34,12 @@
 from sklearn.model_selection import GridSearchCV
 from sklearn.pipeline import Pipeline
 
+from ogbench.data.corrections import (
+    CombatCorrector,
+    CovariateAdjuster,
+    MedianCenterer,
+    PromoterMinBetaSelector,
+)
 from ogbench.data.utils.split_utils import (
     build_omics_cache_relative_name,
     compute_omics_split_indices,
@@ -182,7 +188,7 @@ def load_metadata(data_name: str, cfg: DictConfig) -> dict[str, Any] | None:
         logger.info('Downloading metadata from HuggingFace...')
         hf_repo_id = 'geometric-intelligence/ogbench'
         revision = cfg.dataset.loader.parameters.get(
-            'revision', '4da96d838e81dc3f3da3c559925eea9bd356111e'
+            'revision', '056dfdc4f434fd35355ffbe5f7b63910d785a97a'
         )
 
         metadata_file = hf_hub_download(  # nosec
@@ -218,6 +224,13 @@ def _resolve_grouping(cfg: DictConfig) -> str | None:
     return str(grouping)
 
 
+def _resolve_corrections(cfg: DictConfig) -> list[str]:
+    raw = cfg.dataset.loader.parameters.get('corrections', [])
+    if raw is None:
+        return []
+    return [str(item) for item in list(raw)]
+
+
 def _load_optional_sidecar(
     cfg: DictConfig, data_name: str, suffix: str, local_dir: str
 ) -> pd.DataFrame | None:
@@ -229,7 +242,7 @@ def _load_optional_sidecar(
             repo_id='geometric-intelligence/ogbench',
             repo_type='dataset',
             revision=cfg.dataset.loader.parameters.get(
-                'revision', '4da96d838e81dc3f3da3c559925eea9bd356111e'
+                'revision', '056dfdc4f434fd35355ffbe5f7b63910d785a97a'
             ),
             filename=f'{data_name}_{suffix}.parquet',
         )
@@ -358,7 +371,7 @@ def load_and_prepare_data(cfg: DictConfig) -> DatasetContainer:
 
         hf_repo_id = 'geometric-intelligence/ogbench'
         revision = cfg.dataset.loader.parameters.get(
-            'revision', '4da96d838e81dc3f3da3c559925eea9bd356111e'
+            'revision', '056dfdc4f434fd35355ffbe5f7b63910d785a97a'
         )
 
         data_file = hf_hub_download(  # nosec
@@ -425,6 +438,45 @@ def load_and_prepare_data(cfg: DictConfig) -> DatasetContainer:
     y_val = targets[split_arrays['valid']]
     y_test = targets[split_arrays['test']]
 
+    for name in _resolve_corrections(cfg):
+        if name == 'median_center':
+            continue
+        if name == 'covariate_adjust':
+            cov = _load_optional_sidecar(cfg, data_name, 'covariates', 'temp_data')
+            if cov is None:
+                raise FileNotFoundError(f'{data_name}_covariates.parquet is required')
+            adjuster = CovariateAdjuster()
+            adjuster.fit(train_data, cov.iloc[split_arrays['train']].reset_index(drop=True))
+            train_data = adjuster.transform(
+                train_data, cov.iloc[split_arrays['train']].reset_index(drop=True)
+            )
+            val_data = adjuster.transform(
+                val_data, cov.iloc[split_arrays['valid']].reset_index(drop=True)
+            )
+            test_data = adjuster.transform(
+                test_data, cov.iloc[split_arrays['test']].reset_index(drop=True)
+            )
+        elif name == 'combat':
+            batches = _batch_labels_from_sidecars(cfg, data_name, len(targets))
+            if batches is None:
+                raise FileNotFoundError(f'batch labels are required for combat on {data_name}')
+            corrector = CombatCorrector()
+            corrector.fit(train_data, batches[split_arrays['train']])
+            train_data = corrector.transform(train_data, batches[split_arrays['train']])
+            val_data = corrector.transform(val_data, batches[split_arrays['valid']])
+            test_data = corrector.transform(test_data, batches[split_arrays['test']])
+        elif name == 'promoter_min_beta':
+            probe_map = _load_optional_sidecar(cfg, data_name, 'probe_map', 'temp_data')
+            if probe_map is None:
+                raise FileNotFoundError(f'{data_name}_probe_map.parquet is required')
+            selector = PromoterMinBetaSelector()
+            selector.fit(train_data, y_train, probe_map)
+            train_data = selector.transform(train_data)
+            val_data = selector.transform(val_data)
+            test_data = selector.transform(test_data)
+        else:
+            raise ValueError(f'Unknown correction {name!r}')
+
     X_train = train_data.values
     X_val = val_data.values
     X_test = test_data.values
@@ -455,6 +507,12 @@ def load_and_prepare_data(cfg: DictConfig) -> DatasetContainer:
     test_data = pd.DataFrame(
         imputer.transform(test_data), columns=test_data.columns, index=test_data.index
     )
+    if 'median_center' in _resolve_corrections(cfg):
+        centerer = MedianCenterer()
+        centerer.fit(train_data)
+        train_data = centerer.transform(train_data)
+        val_data = centerer.transform(val_data)
+        test_data = centerer.transform(test_data)
 
     X_train_imputed = train_data.values
     X_val_imputed = val_data.values

ogbench/data/corrections/__init__.py

@@ -0,0 +1,13 @@
+"""Train-only sample corrections (covariate adjustment, ComBat, smoking)."""
+
+from ogbench.data.corrections.center import MedianCenterer
+from ogbench.data.corrections.combat import CombatCorrector
+from ogbench.data.corrections.covariates import CovariateAdjuster
+from ogbench.data.corrections.promoter import PromoterMinBetaSelector
+
+__all__ = [
+    'CombatCorrector',
+    'CovariateAdjuster',
+    'MedianCenterer',
+    'PromoterMinBetaSelector',
+]

ogbench/data/corrections/center.py

@@ -0,0 +1,30 @@
+"""Train-only per-feature median centering."""
+
+from __future__ import annotations
+
+import pandas as pd
+
+
+class MedianCenterer:
+    """Subtract the training-split column median from every split."""
+
+    def __init__(self) -> None:
+        self.median_: pd.Series | None = None
+
+    def fit(self, data: pd.DataFrame) -> MedianCenterer:
+        """Store per-column medians from training samples."""
+        self.median_ = data.median(axis=0)
+        if self.median_.isna().any():
+            n_bad = int(self.median_.isna().sum())
+            raise ValueError(f'{n_bad} features have undefined training medians')
+        return self
+
+    def transform(self, data: pd.DataFrame) -> pd.DataFrame:
+        """Apply frozen training medians."""
+        if self.median_ is None:
+            raise RuntimeError('MedianCenterer must be fit before transform')
+        missing = [c for c in self.median_.index if c not in data.columns]
+        if missing:
+            raise ValueError(f'Features missing from data: {missing[:5]}')
+        centered = data.loc[:, self.median_.index] - self.median_
+        return pd.DataFrame(centered, columns=self.median_.index, index=data.index)

ogbench/data/corrections/combat.py

@@ -0,0 +1,86 @@
+"""Train-only parametric ComBat (location/scale) correction."""
+
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+
+
+class CombatCorrector:
+    """Parametric ComBat with estimates frozen on the training split.
+
+    For each gene, training data are standardized with the train grand mean and
+    pooled variance. Per-batch location (``gamma``) and scale (``delta``) are
+    then estimated on those standardized train values and applied to val/test
+    samples whose batch was seen in train.
+
+    Unseen batches raise ``ValueError``. Labels are never used.
+    """
+
+    def __init__(self, min_batch_samples: int = 2) -> None:
+        self.min_batch_samples = int(min_batch_samples)
+        self.batch_levels_: np.ndarray | None = None
+        self.grand_mean_: np.ndarray | None = None
+        self.var_pooled_: np.ndarray | None = None
+        self.gamma_: dict[object, np.ndarray] = {}
+        self.delta_: dict[object, np.ndarray] = {}
+
+    def fit(self, data: pd.DataFrame, batches: np.ndarray) -> CombatCorrector:
+        """Estimate ComBat parameters on training samples only."""
+        x = np.asarray(data, dtype=float)
+        batches = np.asarray(batches)
+        if x.ndim != 2:
+            raise ValueError('data must be 2-D (samples x features)')
+        if len(batches) != x.shape[0]:
+            raise ValueError('batches length must match number of samples')
+        if np.isnan(x).any():
+            raise ValueError('ComBat requires finite values; impute after correction or drop NaNs')
+
+        self.batch_levels_ = np.unique(batches)
+        n_samples, n_features = x.shape
+        self.grand_mean_ = x.mean(axis=0)
+        var_pooled = np.zeros(n_features, dtype=float)
+        for batch in self.batch_levels_:
+            mask = batches == batch
+            n_b = int(mask.sum())
+            if n_b < self.min_batch_samples:
+                raise ValueError(
+                    f'batch {batch!r} has {n_b} train samples; need >= {self.min_batch_samples}'
+                )
+            var_pooled += x[mask].var(axis=0, ddof=1) * (n_b - 1)
+        self.var_pooled_ = np.clip(var_pooled / max(n_samples - 1, 1), 1e-8, None)
+        scale = np.sqrt(self.var_pooled_)
+        standardized = (x - self.grand_mean_) / scale
+
+        self.gamma_.clear()
+        self.delta_.clear()
+        for batch in self.batch_levels_:
+            mask = batches == batch
+            batch_z = standardized[mask]
+            self.gamma_[batch] = batch_z.mean(axis=0)
+            delta = batch_z.var(axis=0, ddof=1)
+            self.delta_[batch] = np.clip(delta, 1e-8, None)
+        return self
+
+    def transform(self, data: pd.DataFrame, batches: np.ndarray) -> pd.DataFrame:
+        """Apply frozen train ComBat parameters."""
+        if self.grand_mean_ is None or self.var_pooled_ is None:
+            raise RuntimeError('CombatCorrector must be fit before transform')
+        x = np.asarray(data, dtype=float)
+        batches = np.asarray(batches)
+        if len(batches) != x.shape[0]:
+            raise ValueError('batches length must match number of samples')
+        unseen = sorted(set(np.unique(batches)) - set(self.batch_levels_))
+        if unseen:
+            raise ValueError(f'Unseen ComBat batches in transform: {unseen}')
+
+        scale = np.sqrt(self.var_pooled_)
+        standardized = (x - self.grand_mean_) / scale
+        adjusted = np.empty_like(standardized)
+        for batch in np.unique(batches):
+            mask = batches == batch
+            gamma = self.gamma_[batch]
+            delta = self.delta_[batch]
+            adjusted[mask] = (standardized[mask] - gamma) / np.sqrt(delta)
+        restored = adjusted * scale + self.grand_mean_
+        return pd.DataFrame(restored, columns=data.columns, index=data.index)