diff --git a/README.md b/README.md index 57ac61f..34bf474 100644 --- a/README.md +++ b/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 diff --git a/configs/dataset/addneuromed.yaml b/configs/dataset/addneuromed.yaml index 547a88a..383272f 100644 --- a/configs/dataset/addneuromed.yaml +++ b/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 diff --git a/configs/dataset/motrpac.yaml b/configs/dataset/motrpac.yaml index 2ba020c..bc25b12 100644 --- a/configs/dataset/motrpac.yaml +++ b/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 diff --git a/configs/dataset/parkinsons.yaml b/configs/dataset/parkinsons.yaml index f61ac18..7f8bc94 100644 --- a/configs/dataset/parkinsons.yaml +++ b/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 diff --git a/configs/dataset/smoking.yaml b/configs/dataset/smoking.yaml index b136c2a..75e6b23 100644 --- a/configs/dataset/smoking.yaml +++ b/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 diff --git a/configs/hf/default.yaml b/configs/hf/default.yaml index 67049c9..6e19efc 100644 --- a/configs/hf/default.yaml +++ b/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 diff --git a/ogbench/baseline.py b/ogbench/baseline.py index 7547bc6..18518bf 100644 --- a/ogbench/baseline.py +++ b/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 diff --git a/ogbench/data/corrections/__init__.py b/ogbench/data/corrections/__init__.py new file mode 100644 index 0000000..8ac5ad3 --- /dev/null +++ b/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', +] diff --git a/ogbench/data/corrections/center.py b/ogbench/data/corrections/center.py new file mode 100644 index 0000000..f58bbe0 --- /dev/null +++ b/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) diff --git a/ogbench/data/corrections/combat.py b/ogbench/data/corrections/combat.py new file mode 100644 index 0000000..8926d7f --- /dev/null +++ b/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) diff --git a/ogbench/data/corrections/covariates.py b/ogbench/data/corrections/covariates.py new file mode 100644 index 0000000..b525d9e --- /dev/null +++ b/ogbench/data/corrections/covariates.py @@ -0,0 +1,102 @@ +"""Train-only linear covariate adjustment (MoTrPAC-style).""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +from sklearn.linear_model import LinearRegression + + +class CovariateAdjuster: + """Fit ``feature ~ covariates`` on train and apply the same coefficients. + + Categorical covariates in ``{sex, race}`` are one-hot encoded with + ``drop_first=True``. Encoding columns and train covariate means are frozen + at ``fit`` time so val/test cannot change the model. + """ + + def __init__(self, covariate_names: list[str] | None = None) -> None: + self.covariate_names = list(covariate_names or ['age', 'sex', 'bmi', 'race']) + self.feature_names_: list[str] | None = None + self.encoded_columns_: list[str] | None = None + self.coef_: dict[str, np.ndarray] = {} + self.intercept_: dict[str, float] = {} + self.cov_mean_: np.ndarray | None = None + + def fit(self, data: pd.DataFrame, covariates: pd.DataFrame) -> CovariateAdjuster: + """Fit one linear model per feature on training rows.""" + data = data.reset_index(drop=True) + covariates = covariates.reset_index(drop=True) + if len(data) != len(covariates): + raise ValueError('data and covariates must have the same number of rows') + + x_cov = self._encode_covariates(covariates, fit=True) + valid_cov_mask = x_cov.notna().all(axis=1) + if int(valid_cov_mask.sum()) == 0: + raise ValueError('No training samples with complete covariates') + + self.cov_mean_ = x_cov.loc[valid_cov_mask].mean(axis=0).to_numpy(dtype=float) + self.feature_names_ = [str(c) for c in data.columns] + self.coef_.clear() + self.intercept_.clear() + + x_cov_values = x_cov.to_numpy(dtype=float) + for col in self.feature_names_: + y = data[col].to_numpy(dtype=float) + valid = valid_cov_mask.to_numpy() & ~np.isnan(y) + if int(valid.sum()) < 10: + continue + model = LinearRegression() + model.fit(x_cov_values[valid], y[valid]) + self.coef_[col] = model.coef_.astype(float) + self.intercept_[col] = float(model.intercept_) + return self + + def transform(self, data: pd.DataFrame, covariates: pd.DataFrame) -> pd.DataFrame: + """Apply train-fitted adjustment; leave incomplete-covariate rows unchanged.""" + if self.encoded_columns_ is None or self.cov_mean_ is None: + raise RuntimeError('CovariateAdjuster must be fit before transform') + data = data.reset_index(drop=True) + covariates = covariates.reset_index(drop=True) + if len(data) != len(covariates): + raise ValueError('data and covariates must have the same number of rows') + + x_cov = self._encode_covariates(covariates, fit=False) + valid_cov_mask = x_cov.notna().all(axis=1).to_numpy() + x_valid = x_cov.to_numpy(dtype=float)[valid_cov_mask] + adjusted = data.copy() + + for col, coef in self.coef_.items(): + if col not in adjusted.columns: + continue + values = np.array(adjusted[col].to_numpy(dtype=float), copy=True) + predicted = x_valid @ coef + self.intercept_[col] + predicted_mean = float(np.dot(self.cov_mean_, coef) + self.intercept_[col]) + values[valid_cov_mask] = values[valid_cov_mask] - (predicted - predicted_mean) + adjusted[col] = values + return adjusted + + def _encode_covariates(self, covariates: pd.DataFrame, *, fit: bool) -> pd.DataFrame: + missing = [c for c in self.covariate_names if c not in covariates.columns] + if missing: + raise ValueError(f'Missing covariate columns: {missing}') + cov_df = covariates[self.covariate_names].copy() + categorical_cols = [c for c in self.covariate_names if c in {'sex', 'race'}] + continuous_cols = [c for c in self.covariate_names if c not in {'sex', 'race'}] + + if categorical_cols: + encoded = pd.get_dummies(cov_df[categorical_cols], drop_first=True, dtype=float) + else: + encoded = pd.DataFrame(index=cov_df.index) + + if fit: + self.encoded_columns_ = list(encoded.columns) + else: + if self.encoded_columns_ is None: + raise RuntimeError('CovariateAdjuster must be fit before transform') + encoded = encoded.reindex(columns=self.encoded_columns_, fill_value=0.0) + + if continuous_cols: + continuous = cov_df[continuous_cols].astype(float) + return pd.concat([continuous, encoded], axis=1) + return encoded diff --git a/ogbench/data/corrections/promoter.py b/ogbench/data/corrections/promoter.py new file mode 100644 index 0000000..4cd45c5 --- /dev/null +++ b/ogbench/data/corrections/promoter.py @@ -0,0 +1,92 @@ +"""Train-only promoter probe selection (smoking / Illumina 450k).""" + +from __future__ import annotations + +import numpy as np +import pandas as pd + + +def select_min_promoter_per_gene( + beta: pd.DataFrame, mapping: pd.DataFrame, never_mask: np.ndarray +) -> tuple[pd.DataFrame, pd.Series]: + """For each gene pick the candidate probe with the minimum mean beta in never-smokers. + + ``never_mask`` must align with ``beta`` rows. Means are computed only on the + True entries of that mask (typically training never-smokers). + + Returns + ------- + gene_data: + DataFrame (samples x genes) of selected probe betas, columns renamed to genes. + gene_to_probe: + Series mapping gene -> chosen probe_id. + """ + if 'probe_id' not in mapping.columns or 'gene' not in mapping.columns: + raise ValueError('probe mapping must have probe_id and gene columns') + never_mask = np.asarray(never_mask, dtype=bool) + if len(never_mask) != len(beta): + raise ValueError('never_mask length must match number of samples') + if never_mask.sum() == 0: + raise ValueError('No never-smoker samples found (label == 0)') + + available = mapping[mapping['probe_id'].isin(beta.columns)].copy() + if available.empty: + raise ValueError('No manifest probes overlap with the beta matrix columns') + + candidate_probes = available['probe_id'].unique().tolist() + never_means = beta.loc[never_mask, candidate_probes].mean(axis=0) + available['never_mean'] = available['probe_id'].map(never_means) + + available = available.dropna(subset=['never_mean']) + if available.empty: + raise ValueError('All candidate probes have NaN mean across never-smoker samples') + + idx_min = available.groupby('gene')['never_mean'].idxmin() + chosen = available.loc[idx_min, ['gene', 'probe_id']] + gene_to_probe = pd.Series( + chosen['probe_id'].values, index=chosen['gene'].values, name='probe_id' + ) + + gene_data = beta.loc[:, gene_to_probe.values].copy() + gene_data.columns = gene_to_probe.index.astype(str) + return gene_data, gene_to_probe + + +class PromoterMinBetaSelector: + """Collapse promoter probes to genes using train never-smoker means only. + + Never-smokers are samples with label ``0`` (GEO smoking status never). + The chosen probe per gene is frozen and applied to val/test. + """ + + def __init__(self) -> None: + self.gene_to_probe_: pd.Series | None = None + + def fit( + self, + data: pd.DataFrame, + labels: np.ndarray, + mapping: pd.DataFrame, + ) -> PromoterMinBetaSelector: + """Select one promoter probe per gene from training never-smokers.""" + labels = np.asarray(labels).reshape(-1) + if len(labels) != len(data): + raise ValueError('labels length must match number of samples') + never_mask = labels == 0 + gene_data, gene_to_probe = select_min_promoter_per_gene(data, mapping, never_mask) + all_nan_cols = gene_data.columns[gene_data.isna().all(axis=0)] + if len(all_nan_cols) > 0: + gene_to_probe = gene_to_probe.drop(index=all_nan_cols) + self.gene_to_probe_ = gene_to_probe + return self + + def transform(self, data: pd.DataFrame) -> pd.DataFrame: + """Apply the frozen probe-to-gene map.""" + if self.gene_to_probe_ is None: + raise RuntimeError('PromoterMinBetaSelector must be fit before transform') + missing = [p for p in self.gene_to_probe_.values if p not in data.columns] + if missing: + raise ValueError(f'Selected probes missing from data: {missing[:5]}') + gene_data = data.loc[:, self.gene_to_probe_.values].copy() + gene_data.columns = self.gene_to_probe_.index.astype(str) + return gene_data diff --git a/ogbench/data/datasets/hf_omics.py b/ogbench/data/datasets/hf_omics.py index df96a06..13e434d 100644 --- a/ogbench/data/datasets/hf_omics.py +++ b/ogbench/data/datasets/hf_omics.py @@ -91,7 +91,7 @@ def __init__( node_sample_ratio: float | str = 1.0, train_val_test_split: list[float] | None = None, hf_repo_id: str = 'geometric-intelligence/ogbench', - revision: str = '4da96d838e81dc3f3da3c559925eea9bd356111e', + revision: str = '056dfdc4f434fd35355ffbe5f7b63910d785a97a', string_data_dir: str | None = None, species: int = 9606, split_type: str = 'fixed', @@ -118,7 +118,7 @@ def __init__( split_type: ``fixed`` (default) or ``k-fold`` k: Number of CV folds when ``split_type='k-fold'`` fold: Test-fold index when ``split_type='k-fold'`` (usually ``data_seed``) - corrections: Optional cache-path tags for train-only corrections + corrections: Optional train-only corrections (``covariate_adjust``, ``combat``) grouping: Optional group key for k-fold (``batch``) using a sidecar file **kwargs: Additional keyword arguments """ @@ -258,6 +258,83 @@ def _load_batch_labels(self) -> np.ndarray | None: return None return meta[batch_col].to_numpy() + def _apply_corrections( + self, + train_data: pd.DataFrame, + val_data: pd.DataFrame, + test_data: pd.DataFrame, + *, + train_ids: np.ndarray, + valid_ids: np.ndarray, + test_ids: np.ndarray, + train_targets: np.ndarray, + covariates_df: pd.DataFrame | None, + batches: np.ndarray | None, + probe_map_df: pd.DataFrame | None = None, + names: list[str] | None = None, + ) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: + """Fit configured corrections on train and transform val/test.""" + from ogbench.data.corrections import ( + CombatCorrector, + CovariateAdjuster, + MedianCenterer, + PromoterMinBetaSelector, + ) + + for name in names if names is not None else self.corrections: + if name == 'covariate_adjust': + if covariates_df is None: + raise FileNotFoundError( + f'corrections includes covariate_adjust but ' + f'{self.data_name}_covariates.parquet was not found' + ) + adjuster = CovariateAdjuster() + adjuster.fit(train_data, covariates_df.iloc[train_ids]) + train_data = adjuster.transform(train_data, covariates_df.iloc[train_ids]) + val_data = adjuster.transform(val_data, covariates_df.iloc[valid_ids]) + test_data = adjuster.transform(test_data, covariates_df.iloc[test_ids]) + logger.info('Applied train-only covariate adjustment') + elif name == 'combat': + if batches is None: + raise FileNotFoundError( + f'corrections includes combat but no batch labels were found for ' + f'{self.data_name}' + ) + corrector = CombatCorrector() + corrector.fit(train_data, batches[train_ids]) + train_data = corrector.transform(train_data, batches[train_ids]) + val_data = corrector.transform(val_data, batches[valid_ids]) + test_data = corrector.transform(test_data, batches[test_ids]) + logger.info('Applied train-only ComBat correction') + elif name == 'promoter_min_beta': + if probe_map_df is None: + raise FileNotFoundError( + f'corrections includes promoter_min_beta but ' + f'{self.data_name}_probe_map.parquet was not found' + ) + selector = PromoterMinBetaSelector() + selector.fit(train_data, train_targets, probe_map_df) + train_data = selector.transform(train_data) + val_data = selector.transform(val_data) + test_data = selector.transform(test_data) + logger.info( + 'Applied train-only promoter min-beta probe selection ' + f'({train_data.shape[1]} genes)' + ) + elif name == 'median_center': + centerer = MedianCenterer() + centerer.fit(train_data) + train_data = centerer.transform(train_data) + val_data = centerer.transform(val_data) + test_data = centerer.transform(test_data) + logger.info('Applied train-only median centering') + else: + raise ValueError( + f"Unknown correction {name!r}; use 'covariate_adjust', 'combat', " + f"'promoter_min_beta', or 'median_center'" + ) + return train_data, val_data, test_data + def download(self) -> None: r"""Download the dataset from HuggingFace and saves it to the raw directory.""" logger.info(f'Downloading raw data for {self.data_name} from HuggingFace...') @@ -282,6 +359,8 @@ def download(self) -> None: logger.info(f'Downloaded {len(targets)} samples with {raw_data.shape[1]} features') + covariates_df = self._download_optional_parquet(f'{self.data_name}_covariates.parquet') + probe_map_df = self._download_optional_parquet(f'{self.data_name}_probe_map.parquet') batches = self._load_batch_labels() groups = None if self.grouping == 'batch': @@ -337,6 +416,23 @@ def download(self) -> None: + f'): Train={len(train_targets)}, Val={len(val_targets)}, Test={len(test_targets)}' ) + pre_impute = [name for name in self.corrections if name != 'median_center'] + post_impute = [name for name in self.corrections if name == 'median_center'] + if pre_impute: + train_data, val_data, test_data = self._apply_corrections( + train_data, + val_data, + test_data, + train_ids=train_ids, + valid_ids=valid_ids, + test_ids=test_ids, + train_targets=train_targets, + covariates_df=covariates_df, + batches=batches, + probe_map_df=probe_map_df, + names=pre_impute, + ) + # Impute missing values - FIT on training data only, TRANSFORM on all splits nan_count = train_data.isna().sum().sum() if nan_count > 0 or raw_data.isna().sum().sum() > 0: @@ -367,6 +463,21 @@ def download(self) -> None: f'Val NaN={val_data.isna().sum().sum()}, Test NaN={test_data.isna().sum().sum()}' ) + if post_impute: + train_data, val_data, test_data = self._apply_corrections( + train_data, + val_data, + test_data, + train_ids=train_ids, + valid_ids=valid_ids, + test_ids=test_ids, + train_targets=train_targets, + covariates_df=covariates_df, + batches=batches, + probe_map_df=probe_map_df, + names=post_impute, + ) + # Calculate number of nodes to select based on TRAINING data only n_training_samples = len(train_targets) if self.node_sample_ratio == 'full': diff --git a/scripts/plot_adjacency_threshold_analysis.py b/scripts/plot_adjacency_threshold_analysis.py index 3b8d16c..558efd4 100755 --- a/scripts/plot_adjacency_threshold_analysis.py +++ b/scripts/plot_adjacency_threshold_analysis.py @@ -44,7 +44,7 @@ def load_dataset_config(dataset_name: str) -> dict[str, Any]: return { 'data_name': dataset_name, - 'revision': hf_cfg.get('revision', '83299150394717f0646b1bd44d6a55392ab789db'), + 'revision': hf_cfg.get('revision', '056dfdc4f434fd35355ffbe5f7b63910d785a97a'), 'train_val_test_split': cfg['loader']['parameters']['train_val_test_split'], } diff --git a/scripts/processors/addneuromed.py b/scripts/processors/addneuromed.py index 755fec5..6d37954 100644 --- a/scripts/processors/addneuromed.py +++ b/scripts/processors/addneuromed.py @@ -3,142 +3,67 @@ import gzip import os +from io import StringIO import numpy as np import pandas as pd -import requests -from combat.pycombat import pycombat from tqdm import tqdm from scripts.utils import create_dataset_metadata, download_file, upload_to_huggingface -def download_platform_file(url: str) -> str: - """Download platform file from URL and return content.""" - response = requests.get(url, timeout=30) - response.raise_for_status() - return response.text - - -def parse_gpl6947(content: str) -> pd.DataFrame: - """Parse GPL6947 platform file content.""" - lines = content.strip().split('\n') - - # Find the header line (starts with ID) - header_idx = None - for i, line in enumerate(lines): - if line.startswith('ID\t'): - header_idx = i - break - +def _load_illumina_probes(url: str, dest_path: str) -> pd.DataFrame: + """Download an Illumina probe table (bgx/txt) and return probe / Entrez IDs.""" + if not os.path.exists(dest_path): + print(f'Downloading {os.path.basename(dest_path)}...') + download_file(url, dest_path) + with gzip.open(dest_path, 'rt', encoding='utf-8', errors='replace') as handle: + lines = handle.readlines() + header_idx = next((i for i, line in enumerate(lines) if line.startswith('[Probes]')), None) if header_idx is None: - raise ValueError('Could not find header line in GPL6947 file') - - # Extract header and data - header_line = lines[header_idx] - data_lines = lines[header_idx + 1 :] - - # Create DataFrame - data = [] - for line in data_lines: - if line.strip(): # Skip empty lines - data.append(line.split('\t')) - - df = pd.DataFrame(data, columns=header_line.split('\t')) - - # Select relevant columns - relevant_cols = ['ID', 'ILMN_Gene', 'RefSeq_ID', 'Entrez_Gene_ID'] - available_cols = [col for col in relevant_cols if col in df.columns] - - return df[available_cols] - - -def parse_gpl10558(content: str) -> pd.DataFrame: - """Parse GPL10558 platform file content.""" - lines = content.strip().split('\n') - - # Find the header line by looking for ILMN_Gene column - header_idx = None - for i, line in enumerate(lines): - if 'ILMN_Gene' in line and '\t' in line: - header_idx = i - break - - if header_idx is None: - raise ValueError('Could not find header line with ILMN_Gene in GPL10558 file') - - # Extract header and data - header_line = lines[header_idx] - data_lines = lines[header_idx + 1 :] - - # Create DataFrame - data = [] - for line in data_lines: - if line.strip(): # Skip empty lines - data.append(line.split('\t')) - - df = pd.DataFrame(data, columns=header_line.split('\t')) - - # Select relevant columns - GPL10558 has different column names - relevant_cols = ['ID', 'ILMN_Gene', 'RefSeq_ID', 'Entrez_Gene_ID'] - available_cols = [col for col in relevant_cols if col in df.columns] - - return df[available_cols] - - -def create_gene_probe_mapping() -> dict[str, set[str]]: - """Create mapping of genes to sets of probes from both platform files.""" - - # URLs for the platform files - gpl6947_url = 'https://www.ncbi.nlm.nih.gov/geo/query/acc.cgi?mode=raw&is_datatable=true&acc=GPL6947&id=13512&db=GeoDb_blob107' - gpl10558_url = 'https://www.ncbi.nlm.nih.gov/geo/query/acc.cgi?mode=raw&is_datatable=true&acc=GPL10558&id=50081&db=GeoDb_blob135' - - print('Downloading GPL6947 platform file...') - gpl6947_content = download_platform_file(gpl6947_url) - gpl6947_df = parse_gpl6947(gpl6947_content) + raise ValueError(f'Could not find [Probes] section in {dest_path}') + df = pd.read_csv( + StringIO(''.join(lines[header_idx + 1 :])), + sep='\t', + dtype=str, + low_memory=False, + ) + if 'Probe_Id' not in df.columns or 'Entrez_Gene_ID' not in df.columns: + raise ValueError(f'Missing Probe_Id or Entrez_Gene_ID in {dest_path}') + df = df.rename(columns={'Probe_Id': 'ID', 'Entrez_Gene_ID': 'gene_id'}) + df['gene_id'] = df['gene_id'].replace('', np.nan) + return df[['ID', 'gene_id']].dropna(subset=['ID']) + + +def create_gene_probe_mapping(output_dir: str) -> dict[str, set[str]]: + """Create mapping of Entrez genes to probes from both Illumina platforms.""" + gpl6947_df = _load_illumina_probes( + 'https://ftp.ncbi.nlm.nih.gov/geo/platforms/GPL6nnn/GPL6947/suppl/' + 'GPL6947_HumanHT-12_V3_0_R1_11283641_A.bgx.gz', + os.path.join(output_dir, 'GPL6947_HumanHT-12_V3.bgx.gz'), + ) + gpl10558_df = _load_illumina_probes( + 'https://ftp.ncbi.nlm.nih.gov/geo/platforms/GPL10nnn/GPL10558/suppl/' + 'GPL10558_HumanHT-12_V4_0_R2_15002873_B.txt.gz', + os.path.join(output_dir, 'GPL10558_HumanHT-12_V4.txt.gz'), + ) print(f'GPL6947: {len(gpl6947_df)} probes') - - print('Downloading GPL10558 platform file...') - gpl10558_content = download_platform_file(gpl10558_url) - gpl10558_df = parse_gpl10558(gpl10558_content) print(f'GPL10558: {len(gpl10558_df)} probes') - gpl6947_gene_col = 'Entrez_Gene_ID' if 'Entrez_Gene_ID' in gpl6947_df.columns else 'ILMN_Gene' - gpl10558_gene_col = ( - 'Entrez_Gene_ID' if 'Entrez_Gene_ID' in gpl10558_df.columns else 'ILMN_Gene' - ) - - gpl6947_genes = set(gpl6947_df[gpl6947_gene_col].dropna().unique()) - gpl10558_genes = set(gpl10558_df[gpl10558_gene_col].dropna().unique()) + gpl6947_genes = set(gpl6947_df['gene_id'].dropna().unique()) + gpl10558_genes = set(gpl10558_df['gene_id'].dropna().unique()) print(f'GPL6947 unique genes: {len(gpl6947_genes)}') print(f'GPL10558 unique genes: {len(gpl10558_genes)}') - # Find intersection of genes common_genes = gpl6947_genes.intersection(gpl10558_genes) print(f'Common genes: {len(common_genes)}') - # Create mapping: gene symbol -> set of probe_ids gene_probe_mapping: dict[str, set[str]] = {} - - # Add probes from GPL6947 - for _, row in gpl6947_df.iterrows(): - symbol = row[gpl6947_gene_col] - probe_id = row['ID'] - if pd.notna(symbol) and symbol in common_genes: - if symbol not in gene_probe_mapping: - gene_probe_mapping[symbol] = set() - gene_probe_mapping[symbol].add(probe_id) - - # Add probes from GPL10558 - for _, row in gpl10558_df.iterrows(): - symbol = row[gpl10558_gene_col] - probe_id = row['ID'] - if pd.notna(symbol) and symbol in common_genes: - if symbol not in gene_probe_mapping: - gene_probe_mapping[symbol] = set() - gene_probe_mapping[symbol].add(probe_id) - + for frame in (gpl6947_df, gpl10558_df): + for probe_id, gene_id in zip(frame['ID'], frame['gene_id'], strict=True): + if pd.notna(gene_id) and gene_id in common_genes: + gene_probe_mapping.setdefault(gene_id, set()).add(probe_id) return gene_probe_mapping @@ -146,17 +71,14 @@ def process_addneuromed(output_dir: str = 'temp_data') -> None: """Download and process AddNeuroMed dataset.""" os.makedirs(output_dir, exist_ok=True) - # Download URLs urls = { 'GPL10558': 'https://ftp.ncbi.nlm.nih.gov/geo/series/GSE63nnn/GSE63063/matrix/GSE63063-GPL10558_series_matrix.txt.gz', 'GPL6947': 'https://ftp.ncbi.nlm.nih.gov/geo/series/GSE63nnn/GSE63063/matrix/GSE63063-GPL6947_series_matrix.txt.gz', } - # Create gene-probe mapping print('Creating gene-probe mapping...') - gene_probe_mapping = create_gene_probe_mapping() + gene_probe_mapping = create_gene_probe_mapping(output_dir) - raw_data: pd.DataFrame = pd.DataFrame() frames: list[pd.DataFrame] = [] statuses: list[str] = [] batches: list[str] = [] @@ -172,11 +94,9 @@ def process_addneuromed(output_dir: str = 'temp_data') -> None: print(f'Error downloading {dataset}: {str(e)}') raise - # Read microarray data with gzip.open(gz_path, 'rt') as f: data = pd.read_csv(f, sep='\t', comment='!', index_col='ID_REF').transpose() - # Extract statuses from the data dataset_statuses: list[str] = [] with gzip.open(gz_path, 'rt') as f: for line in f: @@ -186,45 +106,33 @@ def process_addneuromed(output_dir: str = 'temp_data') -> None: ) break - # Track batch labels for this dataset dataset_batches = [dataset] * len(data.index) batches.extend(dataset_batches) statuses.extend(dataset_statuses) - # Convert probe-level data to gene-level data print(f'Converting {dataset} from probe-level to gene-level data...') - - # Build gene expression data more efficiently + values = data.to_numpy(dtype=np.float64, copy=False) + col_index = {probe: i for i, probe in enumerate(data.columns)} gene_columns = {} for gene_symbol, probe_ids in tqdm(gene_probe_mapping.items(), desc='Processing genes'): - # Find probes that exist in this dataset - available_probes = [probe for probe in probe_ids if probe in data.columns] - - if available_probes: - # Average expression across all probes for this gene - gene_columns[gene_symbol] = data[available_probes].mean(axis=1) - # Create DataFrame from all gene columns at once + idxs = [col_index[probe] for probe in probe_ids if probe in col_index] + if idxs: + gene_columns[gene_symbol] = values[:, idxs].mean(axis=1) gene_data = pd.DataFrame(gene_columns, index=data.index) frames.append(gene_data) - # Verify no common patients between datasets common_patients = set(frames[0].index).intersection(set(frames[1].index)) assert len(common_patients) == 0, 'Common patients found between the two datasets' - # Find and use common genes common_genes = list(set(frames[0].columns).intersection(set(frames[1].columns))) frames[0] = frames[0][common_genes] frames[1] = frames[1][common_genes] - # Combine datasets raw_data = pd.concat(frames, axis=0) targets = np.array(statuses) batch_labels = np.array(batches) - # Define classes to remove classes_to_remove = {'CTL to AD', 'MCI to CTL', 'OTHER', 'borderline MCI'} - - # Filter out unwanted classes mask = ( (~pd.isna(raw_data)).any(axis=1).values & ~pd.isna(targets) @@ -237,62 +145,36 @@ def process_addneuromed(output_dir: str = 'temp_data') -> None: gene_map = pd.DataFrame( { 'node_id': list(raw_data.columns), - 'string_id': list(raw_data.columns), # Entrez IDs map directly to STRING + 'string_id': list(raw_data.columns), } ) gene_map['node_id'] = gene_map['node_id'].astype(str) gene_map['string_id'] = gene_map['string_id'].astype(str) print(f'Mapping: {len(gene_map)} genes with Entrez IDs for STRING') - - # Apply ComBat batch correction - print('Applying ComBat batch correction...') - - # Create design matrix for status to preserve biological signal - unique_statuses = np.unique(targets) - design_matrix = pd.DataFrame(index=raw_data.index) - for status in unique_statuses: - design_matrix[f'status_{status}'] = (targets == status).astype(int) - - # Apply ComBat correction - # pycombat expects genes x samples, so we transpose raw_data - raw_data_transposed = raw_data.T # Transpose to genes x samples - - # Convert batch labels to numeric indices for pycombat - unique_batches = np.unique(batch_labels) - batch_to_idx = {batch: i for i, batch in enumerate(unique_batches)} - batch_indices = np.array([batch_to_idx[batch] for batch in batch_labels]) - - # Apply ComBat correction without covariates for now - # This will correct batch effects while preserving overall biological signal - corrected_data = pycombat( - raw_data_transposed, # Pass DataFrame, not numpy array - batch_indices, + print( + 'Skipping full-data ComBat; platform batch labels are saved for train-only ' + 'correction in HFOmics. Unique batches:', + np.unique(batch_labels), ) - # Convert back to DataFrame (transpose back to samples x genes) - raw_data = corrected_data.T # Transpose back to samples x features - - # Raise if raw data or targets have nan values assert not raw_data.isna().any().any(), 'Raw data has nan values' assert not (targets == '').any(), 'Targets have empty strings' - # Convert string classes to integers - unique_classes = np.unique(targets) + unique_classes = [str(name) for name in np.unique(targets)] class_to_int = {class_name: i for i, class_name in enumerate(unique_classes)} - targets_int = np.array([class_to_int[class_name] for class_name in targets]) + targets_int = np.array([class_to_int[str(class_name)] for class_name in targets]) - # Save as parquet data_file = os.path.join(output_dir, 'addneuromed_data.parquet') targets_file = os.path.join(output_dir, 'addneuromed_targets.parquet') map_file = os.path.join(output_dir, 'addneuromed_map.parquet') + batches_file = os.path.join(output_dir, 'addneuromed_batches.parquet') - # Reset index to make it a proper DataFrame raw_data = raw_data.reset_index(drop=True) raw_data.to_parquet(data_file) pd.DataFrame({'target': targets_int}).to_parquet(targets_file) gene_map.reset_index(drop=True).to_parquet(map_file, index=False) + pd.DataFrame({'batch': batch_labels}).to_parquet(batches_file, index=False) - # Create metadata target_stats = { 'class_mapping': class_to_int, 'num_classes': len(unique_classes), @@ -305,10 +187,19 @@ def process_addneuromed(output_dir: str = 'temp_data') -> None: num_samples=len(targets_int), num_features=raw_data.shape[1], target_stats=target_stats, + preprocessing_notes=( + 'Platforms GPL10558 and GPL6947 are concatenated without ComBat at ingest. ' + 'ogbench fits ComBat on the training split only and applies frozen batch ' + 'location/scale to val/test. Labels are not used as ComBat covariates.' + ), ) - # Upload to HuggingFace - data_files = {'data': data_file, 'targets': targets_file, 'map': map_file} + data_files = { + 'data': data_file, + 'targets': targets_file, + 'map': map_file, + 'batches': batches_file, + } upload_to_huggingface('addneuromed', data_files, metadata) diff --git a/scripts/processors/motrpac.py b/scripts/processors/motrpac.py index 1f32147..e6f309d 100644 --- a/scripts/processors/motrpac.py +++ b/scripts/processors/motrpac.py @@ -4,116 +4,10 @@ import numpy as np import pandas as pd -from sklearn.linear_model import LinearRegression from scripts.utils import create_dataset_metadata, download_file, upload_to_huggingface -def adjust_for_covariates( - data: pd.DataFrame, - covariates: pd.DataFrame, - covariate_names: list[str], -) -> pd.DataFrame: - """Adjust protein data for specified covariates using linear regression. - - For each protein, fits: protein ~ covariates, then adjusts to remove covariate - effects while centering at mean covariate values. This avoids using target - labels, preventing data leakage. - Categorical variables (sex, race) are automatically one-hot encoded. - - Args: - data: DataFrame with protein columns (already log-transformed) - covariates: DataFrame with covariate columns - covariate_names: List of covariate column names to adjust for - - Returns: - Adjusted data with same shape as input - """ - adjusted_data = data.copy() - - # Build covariate matrix with one-hot encoding for categoricals - cov_df = covariates[covariate_names].copy() - - # Identify categorical columns (sex, race) - categorical_cols = [col for col in covariate_names if col in {'sex', 'race'}] - continuous_cols = [col for col in covariate_names if col not in categorical_cols] - - # One-hot encode categoricals and combine with continuous - if categorical_cols: - cov_encoded = pd.get_dummies(cov_df[categorical_cols], drop_first=True, dtype=float) - else: - cov_encoded = pd.DataFrame(index=cov_df.index) - - if continuous_cols: - cov_continuous = cov_df[continuous_cols].astype(float) - X_cov = pd.concat([cov_continuous, cov_encoded], axis=1) - else: - X_cov = cov_encoded - - # Find rows with complete covariate data - valid_cov_mask = X_cov.notna().all(axis=1) - - if valid_cov_mask.sum() == 0: - print('Warning: No samples with complete covariate data. Returning original data.') - return adjusted_data - - X_cov_valid = X_cov.loc[valid_cov_mask].values - X_cov_mean = X_cov_valid.mean(axis=0) - - print(f'Adjusting for covariates: {covariate_names}') - print(f' Encoded covariate columns: {list(X_cov.columns)}') - print(f' Samples with complete covariate data: {valid_cov_mask.sum()} / {len(data)}') - - # Track adjustment statistics - adjustment_stats: dict[str, dict[str, float]] = {} - - # For each protein, fit protein ~ covariates and adjust - for protein_col in data.columns: - protein_values = data[protein_col].values - - # Find valid rows: complete covariates AND non-NaN protein - valid_mask = valid_cov_mask.values & ~np.isnan(protein_values) - - if valid_mask.sum() < 10: - # Not enough data to fit regression, skip adjustment - continue - - # Get valid data for this protein - X_valid = X_cov.loc[valid_mask].values - y_valid = protein_values[valid_mask] - - # Fit: protein ~ covariates - model = LinearRegression() - model.fit(X_valid, y_valid) - - # Calculate expected value at mean covariates - expected_at_mean = model.predict(X_cov_mean.reshape(1, -1))[0] - - # For all samples with valid covariates, calculate adjustment - # adjusted = original - (predicted - expected_at_mean) - predicted_all = model.predict(X_cov_valid) - adjustment = predicted_all - expected_at_mean - - # Apply adjustment only to samples with valid covariates - adjusted_values = protein_values.copy() - adjusted_values[valid_cov_mask.values] = protein_values[valid_cov_mask.values] - adjustment - adjusted_data[protein_col] = adjusted_values - - # Track mean absolute adjustment for this protein - adjustment_stats[protein_col] = { - 'mean_abs_adjustment': float(np.abs(adjustment).mean()), - 'valid_samples': int(valid_mask.sum()), - } - - # Print summary statistics - mean_adjustments = [s['mean_abs_adjustment'] for s in adjustment_stats.values()] - print(f' Proteins adjusted: {len(adjustment_stats)} / {len(data.columns)}') - print(f' Mean absolute adjustment across proteins: {np.mean(mean_adjustments):.4f}') - print(f' Max absolute adjustment: {np.max(mean_adjustments):.4f}') - - return adjusted_data - - def process_motrpac(output_dir: str = 'temp_data') -> None: """Download and process MotrPac dataset.""" os.makedirs(output_dir, exist_ok=True) @@ -227,16 +121,10 @@ def process_motrpac(output_dir: str = 'temp_data') -> None: raw_data = np.log2(raw_data) raw_data = pd.DataFrame(raw_data, columns=raw_data.columns).reset_index(drop=True) - # 6.5) Adjust for covariates (age, sex, bmi, race) - print('\nCovariate selection:') - print(f' Available covariates in cov dataframe: {cov.columns.tolist()}') + print('\nCovariate columns available for train-only adjustment at graph-build time:') + print(f' {cov.columns.tolist()}') print(f' Non-null counts: {cov.notna().sum().to_dict()}') - - covariates_to_adjust = ['age', 'sex', 'bmi', 'race'] - - print(f' Selected for adjustment: {covariates_to_adjust}') - - raw_data = adjust_for_covariates(raw_data, cov, covariates_to_adjust) + print(' Skipping full-data adjustment; HFOmics applies covariate_adjust on train only.') # 7) Emit classification target only out = os.path.join(output_dir, 'motrpac') @@ -267,10 +155,9 @@ def process_motrpac(output_dir: str = 'temp_data') -> None: num_features=raw_data.shape[1], target_stats=target_stats, preprocessing_notes=( - 'Data is log2-transformed and adjusted for covariates (age, sex, bmi, race) ' - 'using linear regression. For each protein, fits protein ~ covariates, then adjusts ' - 'to remove covariate effects centered at mean covariate values. This approach avoids ' - 'using target labels, preventing data leakage.' + 'Data is log2-transformed. Covariate adjustment (age, sex, bmi, race) is NOT ' + 'applied at ingest; ogbench fits protein ~ covariates on the training split only ' + 'and applies the frozen coefficients to val/test.' ), ) @@ -279,6 +166,7 @@ def process_motrpac(output_dir: str = 'temp_data') -> None: 'data': os.path.join(out, 'motrpac_data.parquet'), 'targets': os.path.join(out, 'motrpac_targets.parquet'), 'map': os.path.join(out, 'motrpac_map.parquet'), + 'covariates': os.path.join(out, 'motrpac_covariates.parquet'), } upload_to_huggingface('motrpac', data_files, metadata) diff --git a/scripts/processors/parkinsons.py b/scripts/processors/parkinsons.py index 9ce8166..7236875 100644 --- a/scripts/processors/parkinsons.py +++ b/scripts/processors/parkinsons.py @@ -6,6 +6,7 @@ import numpy as np import pandas as pd +from ogbench.data.utils.split_utils import print_group_split_inventory from scripts.utils import create_dataset_metadata, download_file, upload_to_huggingface @@ -68,6 +69,34 @@ def _map_probes_to_genes( return df, symbol_to_entrez +def _parse_geo_sample_characteristics(metadata_lines: list[list[str]]) -> pd.DataFrame: + """Parse GEO ``!Sample_characteristics_ch1`` lines into a sample-by-field table.""" + n_samples = len(metadata_lines[0]) + records: list[dict[str, str]] = [{} for _ in range(n_samples)] + for line in metadata_lines: + if len(line) != n_samples: + raise ValueError('GEO characteristic lines have inconsistent sample counts') + for i, cell in enumerate(line): + cell = str(cell).strip().strip('"') + if ':' not in cell: + continue + key, value = cell.split(':', 1) + records[i][key.strip().lower()] = value.strip() + return pd.DataFrame.from_records(records) + + +def _infer_batch_field(sample_meta: pd.DataFrame) -> str | None: + preferred = ('batch', 'hybridization', 'hyb_batch', 'scan date', 'scan_date', 'chip') + columns = list(sample_meta.columns) + for key in preferred: + if key in columns: + return key + for col in columns: + if 'batch' in str(col): + return col + return None + + def process_parkinsons(output_dir: str = 'temp_data') -> None: """Download and process Parkinsons dataset.""" os.makedirs(output_dir, exist_ok=True) @@ -99,22 +128,10 @@ def process_parkinsons(output_dir: str = 'temp_data') -> None: if not metadata_lines: raise ValueError('No !Sample_characteristics_ch1 lines found.') - # Transpose so each item corresponds to a sample - sample_metadata = list(zip(*metadata_lines, strict=True)) - - # Extract 'moca score' for each sample - moca_scores = [] - for fields in sample_metadata: - moca = None - for field in fields: - if 'moca score:' in field.lower(): - try: - moca = field.split(':')[1].strip().strip('"') - except IndexError: - pass - moca_scores.append(moca) - - moca_scores = pd.to_numeric(moca_scores, errors='coerce') + sample_meta = _parse_geo_sample_characteristics(metadata_lines) + if 'moca score' not in sample_meta.columns: + raise ValueError('No moca score field found in GEO sample characteristics.') + moca_scores = pd.to_numeric(sample_meta['moca score'], errors='coerce') # Load gene expression data (after metadata ends) with gzip.open(gz_path, 'rt') as f: @@ -130,9 +147,10 @@ def process_parkinsons(output_dir: str = 'temp_data') -> None: len(moca_scores) == expression_df.shape[0] ), f'Mismatched samples: {len(moca_scores)} scores vs {expression_df.shape[0]} samples' - valid_mask = ~np.isnan(moca_scores) - raw_data = expression_df.loc[valid_mask] - targets = moca_scores[valid_mask] + valid_mask = ~np.isnan(moca_scores.to_numpy()) + raw_data = expression_df.iloc[valid_mask] + targets = moca_scores.to_numpy()[valid_mask] + sample_meta = sample_meta.iloc[valid_mask].reset_index(drop=True) assert not raw_data.isna().any().any(), 'Raw data contains NaNs' assert not np.isnan(targets).any(), 'Targets contain NaNs' @@ -171,16 +189,36 @@ def moca_to_class(moca_score): for class_id, count in zip(unique_classes, counts, strict=True): print(f' {class_names[class_id]}: {count} samples ({count/len(targets_class)*100:.1f}%)') + batch_field = _infer_batch_field(sample_meta) + print(f'GEO characteristic fields: {list(sample_meta.columns)}') + if batch_field is None: + print( + 'No batch-like GEO field found. Keep sample-stratified splits; ' + 'do not set grouping=batch.' + ) + else: + if batch_field != 'batch': + sample_meta['batch'] = sample_meta[batch_field] + print(f'Using batch field {batch_field!r}') + print_group_split_inventory( + sample_meta['batch'].to_numpy(), + targets_class, + k=5, + group_name='batch', + ) + # Save as parquet data_file = os.path.join(output_dir, 'parkinsons_data.parquet') targets_file = os.path.join(output_dir, 'parkinsons_targets.parquet') map_file = os.path.join(output_dir, 'parkinsons_map.parquet') + meta_file = os.path.join(output_dir, 'parkinsons_sample_meta.parquet') # Reset index to make it a proper DataFrame raw_data = raw_data.reset_index(drop=True) raw_data.to_parquet(data_file) pd.DataFrame({'target': targets_class}).to_parquet(targets_file) gene_map.reset_index(drop=True).to_parquet(map_file, index=False) + sample_meta.reset_index(drop=True).to_parquet(meta_file, index=False) # Create metadata target_stats = { @@ -201,10 +239,20 @@ def moca_to_class(moca_score): num_samples=len(targets), num_features=raw_data.shape[1], target_stats=target_stats, + preprocessing_notes=( + 'GEO sample characteristics are stored in parkinsons_sample_meta.parquet. ' + 'Unmixed-batch k-fold is enabled only when the processor inventory reports ' + 'FEASIBLE (at least k batches with both classes in every 3/1/1 split).' + ), ) # Upload to HuggingFace - data_files = {'data': data_file, 'targets': targets_file, 'map': map_file} + data_files = { + 'data': data_file, + 'targets': targets_file, + 'map': map_file, + 'sample_meta': meta_file, + } upload_to_huggingface('parkinsons', data_files, metadata) diff --git a/scripts/processors/smoking.py b/scripts/processors/smoking.py index c9c11bf..540db18 100644 --- a/scripts/processors/smoking.py +++ b/scripts/processors/smoking.py @@ -128,43 +128,6 @@ def _parse_series_matrix(gz_path: str) -> tuple[list[str], np.ndarray, pd.DataFr return sample_ids, targets, beta -def _select_min_promoter_per_gene( - beta: pd.DataFrame, mapping: pd.DataFrame, never_mask: np.ndarray -) -> tuple[pd.DataFrame, pd.Series]: - """For each gene pick the candidate probe with the minimum mean beta in never-smokers. - - Returns: - gene_data: DataFrame (samples x genes) of selected probe betas, columns renamed to genes. - gene_to_probe: Series mapping gene -> chosen probe_id. - """ - available = mapping[mapping['probe_id'].isin(beta.columns)].copy() - if available.empty: - raise ValueError('No manifest probes overlap with the beta matrix columns') - - candidate_probes = available['probe_id'].unique().tolist() - never_means = beta.loc[never_mask, candidate_probes].mean(axis=0) - available['never_mean'] = available['probe_id'].map(never_means) - - available = available.dropna(subset=['never_mean']) - if available.empty: - raise ValueError('All candidate probes have NaN mean across never-smoker samples') - - idx_min = available.groupby('gene')['never_mean'].idxmin() - chosen = available.loc[idx_min, ['gene', 'probe_id']] - gene_to_probe = pd.Series( - chosen['probe_id'].values, index=chosen['gene'].values, name='probe_id' - ) - - gene_data = beta.loc[:, gene_to_probe.values].copy() - gene_data.columns = gene_to_probe.index.astype(str) - - print( - f'Selected {gene_data.shape[1]} gene-level features from ' - f'{len(candidate_probes)} candidate promoter probes' - ) - return gene_data, gene_to_probe - - def process_smoking(output_dir: str = 'temp_data') -> None: """Download and process the smoking (GSE50660) methylation dataset.""" os.makedirs(output_dir, exist_ok=True) @@ -191,32 +154,20 @@ def process_smoking(output_dir: str = 'temp_data') -> None: print('Building probe-to-gene promoter mapping...') mapping = _build_probe_gene_mapping(manifest_path) - print('Selecting per-gene minimum-beta promoter in never-smokers...') - never_mask = targets == 0 - if never_mask.sum() == 0: - raise ValueError('No never-smoker samples found (smoking == 0)') - gene_data, gene_to_probe = _select_min_promoter_per_gene(beta, mapping, never_mask) - - all_nan_cols = gene_data.columns[gene_data.isna().all(axis=0)] - if len(all_nan_cols) > 0: - print(f'Dropping {len(all_nan_cols)} all-NaN gene columns') - gene_data = gene_data.drop(columns=all_nan_cols) - gene_to_probe = gene_to_probe.drop(index=all_nan_cols) - - if gene_data.isna().any().any(): - nan_cells = int(gene_data.isna().sum().sum()) - print(f'Imputing {nan_cells} remaining NaN cells with column means') - gene_data = gene_data.fillna(gene_data.mean(axis=0)) - - print('Median-centering per gene across samples...') - gene_data = gene_data - gene_data.median(axis=0) - - assert not gene_data.isna().any().any(), 'Gene data has NaN values after processing' - assert not np.isnan(targets).any(), 'Targets have NaN values' - assert gene_data.shape[0] == len(targets), 'Sample count mismatch between data and targets' + available = mapping[mapping['probe_id'].isin(beta.columns)].copy() + if available.empty: + raise ValueError('No manifest probes overlap with the beta matrix columns') + candidate_probes = available['probe_id'].unique().tolist() + probe_data = beta.loc[:, candidate_probes].copy() + print( + f'Keeping {probe_data.shape[1]} TSS1500/TSS200 promoter probes covering ' + f'{available["gene"].nunique()} genes; probe pick, impute, and median-center ' + 'run train-only in HFOmics.' + ) # Collapse the original 3-class GEO encoding (0=never, 1=former, 2=current) into a binary # never (0) vs ever-smoker (1 = former + current) target for downstream modeling. + # Class 0 remains never-smoker so PromoterMinBetaSelector can use train labels. targets_binary = (targets > 0).astype(np.int64) class_names = ['never', 'ever'] class_mapping = {'never': 0, 'ever': 1} @@ -225,25 +176,31 @@ def process_smoking(output_dir: str = 'temp_data') -> None: 'former': int((targets == 1).sum()), 'current': int((targets == 2).sum()), } + if original_class_counts['never'] == 0: + raise ValueError('No never-smoker samples found (smoking == 0)') + + assert probe_data.shape[0] == len( + targets_binary + ), 'Sample count mismatch between data and targets' - # Build gene map: feature columns are gene symbols, which STRING resolves - # directly via its alias lookup (same approach as brca/addneuromed). gene_map = pd.DataFrame( { - 'node_id': list(gene_data.columns), - 'string_id': list(gene_data.columns), + 'node_id': available['gene'].astype(str).unique(), + 'string_id': available['gene'].astype(str).unique(), } ) - gene_map['node_id'] = gene_map['node_id'].astype(str) - gene_map['string_id'] = gene_map['string_id'].astype(str) data_file = os.path.join(output_dir, 'smoking_data.parquet') targets_file = os.path.join(output_dir, 'smoking_targets.parquet') map_file = os.path.join(output_dir, 'smoking_map.parquet') + probe_map_file = os.path.join(output_dir, 'smoking_probe_map.parquet') - gene_data.reset_index(drop=True).to_parquet(data_file) + probe_data.reset_index(drop=True).to_parquet(data_file) pd.DataFrame({'target': targets_binary}).to_parquet(targets_file) gene_map.reset_index(drop=True).to_parquet(map_file, index=False) + available[['probe_id', 'gene']].drop_duplicates().reset_index(drop=True).to_parquet( + probe_map_file, index=False + ) target_stats: dict = { 'class_mapping': class_mapping, @@ -253,29 +210,36 @@ def process_smoking(output_dir: str = 'temp_data') -> None: name: int((targets_binary == idx).sum()) for name, idx in class_mapping.items() }, 'original_geo_class_counts': original_class_counts, + 'n_promoter_probes': int(probe_data.shape[1]), + 'n_mapped_genes': int(available['gene'].nunique()), } metadata = create_dataset_metadata( dataset_name='smoking', download_urls=urls, num_samples=len(targets_binary), - num_features=gene_data.shape[1], + num_features=probe_data.shape[1], target_stats=target_stats, preprocessing_notes=( - 'GSE50660 Illumina 450k beta values mapped to genes using the HumanMethylation450 ' - 'v1.2 manifest, restricted to probes annotated as TSS1500 or TSS200 promoter ' - 'regions. For each gene, the candidate promoter probe with the minimum mean beta ' - 'across never-smoker samples (original GEO smoking == 0) is kept as the gene-level ' - 'feature. Values are then median-centered per gene across samples. The original ' - '3-class GEO smoking status (0=never, 1=former, 2=current) is collapsed into a ' - 'binary target: 0=never, 1=ever (former or current).' + 'GSE50660 Illumina 450k beta values restricted to HumanMethylation450 v1.2 ' + 'probes annotated as TSS1500 or TSS200. Hub matrices are uncorrected promoter ' + 'probes. ogbench picks, for each gene, the candidate probe with the minimum mean ' + 'beta on training never-smokers (label 0), imputes remaining NaNs with training ' + 'column means, and median-centers each gene on the training split only. The ' + 'original 3-class GEO smoking status (0=never, 1=former, 2=current) is collapsed ' + 'into a binary target: 0=never, 1=ever (former or current).' ), ) - data_files = {'data': data_file, 'targets': targets_file, 'map': map_file} + data_files = { + 'data': data_file, + 'targets': targets_file, + 'map': map_file, + 'probe_map': probe_map_file, + } upload_to_huggingface('smoking', data_files, metadata) print('Successfully processed and uploaded smoking dataset') print(f' Samples: {len(targets_binary)}') - print(f' Features (genes): {gene_data.shape[1]}') + print(f' Features (promoter probes): {probe_data.shape[1]}') print(f' Target stats: {target_stats}') diff --git a/tests/data/test_corrections.py b/tests/data/test_corrections.py new file mode 100644 index 0000000..1358225 --- /dev/null +++ b/tests/data/test_corrections.py @@ -0,0 +1,125 @@ +"""Tests for train-only sample corrections.""" + +import numpy as np +import pandas as pd +import pytest + +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, select_min_promoter_per_gene + + +def test_covariate_adjuster_fit_does_not_use_val_rows(): + rng = np.random.default_rng(0) + n, p = 80, 5 + age = rng.normal(50, 10, size=n) + sex = np.where(rng.random(n) > 0.5, 'F', 'M') + data = pd.DataFrame({f'p{i}': 0.2 * age + rng.normal(0, 0.1, size=n) for i in range(p)}) + cov = pd.DataFrame({'age': age, 'sex': sex, 'bmi': rng.normal(25, 3, n), 'race': 'A'}) + train, val = slice(0, 60), slice(60, 80) + adj = CovariateAdjuster() + adj.fit(data.iloc[train], cov.iloc[train]) + coef_before = {k: v.copy() for k, v in adj.coef_.items()} + adj.transform(data.iloc[val], cov.iloc[val]) + for key, value in coef_before.items(): + np.testing.assert_array_equal(adj.coef_[key], value) + + +def test_covariate_adjuster_reduces_age_correlation_on_train(): + rng = np.random.default_rng(1) + n = 100 + age = rng.normal(50, 10, size=n) + y = 0.5 * age + rng.normal(0, 0.05, size=n) + data = pd.DataFrame({'prot': y}) + cov = pd.DataFrame({'age': age, 'sex': 'F', 'bmi': 24.0, 'race': 'A'}) + adj = CovariateAdjuster(covariate_names=['age']) + adj.fit(data, cov) + out = adj.transform(data, cov) + corr_before = abs(np.corrcoef(data['prot'], age)[0, 1]) + corr_after = abs(np.corrcoef(out['prot'], age)[0, 1]) + assert corr_after < corr_before * 0.2 + + +def test_combat_train_only_estimates_frozen(): + rng = np.random.default_rng(2) + n_genes = 12 + b0 = rng.normal(0, 1, size=(40, n_genes)) + 3.0 + b1 = rng.normal(0, 1, size=(40, n_genes)) - 2.0 + x = np.vstack([b0, b1]) + batches = np.array(['a'] * 40 + ['b'] * 40) + data = pd.DataFrame(x, columns=[f'g{i}' for i in range(n_genes)]) + combat = CombatCorrector() + combat.fit(data.iloc[:60], batches[:60]) + gamma_before = {k: v.copy() for k, v in combat.gamma_.items()} + combat.transform(data.iloc[60:], batches[60:]) + for key, value in gamma_before.items(): + np.testing.assert_array_equal(combat.gamma_[key], value) + + +def test_combat_unseen_batch_raises(): + rng = np.random.default_rng(3) + data = pd.DataFrame(rng.normal(size=(20, 4)), columns=list('abcd')) + batches = np.array(['a'] * 10 + ['b'] * 10) + combat = CombatCorrector() + combat.fit(data.iloc[:10], batches[:10]) + with pytest.raises(ValueError, match='Unseen ComBat batches'): + combat.transform(data.iloc[10:], batches[10:]) + + +def test_combat_reduces_batch_mean_gap(): + rng = np.random.default_rng(4) + n_genes = 8 + a = rng.normal(0, 1, size=(30, n_genes)) + 5.0 + b = rng.normal(0, 1, size=(30, n_genes)) + data = pd.DataFrame(np.vstack([a, b]), columns=[f'g{i}' for i in range(n_genes)]) + batches = np.array(['a'] * 30 + ['b'] * 30) + gap_before = abs(data.iloc[:30].mean().to_numpy() - data.iloc[30:].mean().to_numpy()).mean() + combat = CombatCorrector().fit(data, batches) + out = combat.transform(data, batches) + gap_after = abs(out.iloc[:30].mean().to_numpy() - out.iloc[30:].mean().to_numpy()).mean() + assert gap_after < gap_before * 0.25 + + +def test_promoter_selector_ignores_val_never_smokers(): + """Val never-smokers would prefer probe B; train-only fit must keep probe A.""" + probes = ['cgA', 'cgB'] + mapping = pd.DataFrame({'probe_id': probes, 'gene': ['GENE1', 'GENE1']}) + # 4 train never-smokers: A is lower; 4 val never-smokers: B is lower + train_never = pd.DataFrame({'cgA': [0.1, 0.1, 0.1, 0.1], 'cgB': [0.9, 0.9, 0.9, 0.9]}) + val_never = pd.DataFrame({'cgA': [0.9, 0.9, 0.9, 0.9], 'cgB': [0.1, 0.1, 0.1, 0.1]}) + ever = pd.DataFrame({'cgA': [0.5, 0.5], 'cgB': [0.5, 0.5]}) + train = pd.concat([train_never, ever], ignore_index=True) + val = pd.concat([val_never, ever], ignore_index=True) + y_train = np.array([0, 0, 0, 0, 1, 1]) + selector = PromoterMinBetaSelector().fit(train, y_train, mapping) + assert list(selector.gene_to_probe_.values) == ['cgA'] + out_val = selector.transform(val) + np.testing.assert_allclose(out_val['GENE1'].to_numpy()[:4], [0.9, 0.9, 0.9, 0.9]) + + +def test_select_min_promoter_matches_full_never_mask(): + mapping = pd.DataFrame({'probe_id': ['p1', 'p2', 'p3'], 'gene': ['G', 'G', 'H']}) + beta = pd.DataFrame( + { + 'p1': [0.2, 0.3, 0.8], + 'p2': [0.4, 0.5, 0.1], + 'p3': [0.6, 0.7, 0.9], + } + ) + never_mask = np.array([True, True, False]) + gene_data, gene_to_probe = select_min_promoter_per_gene(beta, mapping, never_mask) + assert gene_to_probe['G'] == 'p1' + assert gene_to_probe['H'] == 'p3' + np.testing.assert_allclose(gene_data['G'].to_numpy(), beta['p1'].to_numpy()) + + +def test_median_centerer_uses_train_only(): + train = pd.DataFrame({'g': [1.0, 3.0, 5.0]}) + val = pd.DataFrame({'g': [10.0, 20.0]}) + centerer = MedianCenterer().fit(train) + np.testing.assert_allclose(centerer.median_['g'], 3.0) + out_val = centerer.transform(val) + np.testing.assert_allclose(out_val['g'].to_numpy(), [7.0, 17.0]) + out_train = centerer.transform(train) + np.testing.assert_allclose(out_train['g'].median(), 0.0) diff --git a/tests/data/test_dataset_config_consistency.py b/tests/data/test_dataset_config_consistency.py index f299f6f..5adc3aa 100644 --- a/tests/data/test_dataset_config_consistency.py +++ b/tests/data/test_dataset_config_consistency.py @@ -74,6 +74,8 @@ def test_all_configs_identical_except_dataset_specific_keys(self, all_configs): ('loader', 'parameters', 'data_name'), ('loader', 'parameters', 'adjacency_threshold'), ('loader', 'parameters', 'species'), + ('loader', 'parameters', 'corrections'), + ('split_params', 'grouping'), ('parameters', 'num_classes'), ('parameters', 'num_samples'), ('parameters', 'full_num_nodes'), diff --git a/tests/data/test_parkinsons_geo_meta.py b/tests/data/test_parkinsons_geo_meta.py new file mode 100644 index 0000000..05de3d6 --- /dev/null +++ b/tests/data/test_parkinsons_geo_meta.py @@ -0,0 +1,39 @@ +"""Tests for Parkinson GEO characteristic parsing and batch inventory.""" + +from pathlib import Path + +import pandas as pd +import yaml + +from ogbench.data.utils.split_utils import group_kfold_is_feasible +from scripts.processors.parkinsons import ( + _infer_batch_field, + _parse_geo_sample_characteristics, +) + + +def test_parse_geo_sample_characteristics(): + lines = [ + ['moca score: 22', 'moca score: 18'], + ['batch: hyb1', 'batch: hyb2'], + ] + meta = _parse_geo_sample_characteristics(lines) + assert list(meta.columns) == ['moca score', 'batch'] + assert meta.loc[0, 'moca score'] == '22' + assert meta.loc[1, 'batch'] == 'hyb2' + assert _infer_batch_field(meta) == 'batch' + + +def test_two_batches_cannot_support_five_fold_groups(): + groups = pd.Series(['GPL1'] * 20 + ['GPL2'] * 20).to_numpy() + labels = [0, 1] * 20 + ok, reason = group_kfold_is_feasible(groups, labels, k=5) + assert not ok + assert 'need at least k=5 groups' in reason + + +def test_parkinsons_config_requests_batch_grouping(): + cfg_path = Path(__file__).resolve().parents[2] / 'configs' / 'dataset' / 'parkinsons.yaml' + with open(cfg_path) as f: + cfg = yaml.safe_load(f) + assert cfg['split_params']['grouping'] == 'batch'