combatlearn 2.1.0__tar.gz → 2.3.0__tar.gz

{combatlearn-2.1.0 → combatlearn-2.3.0}/PKG-INFO

@@ -1,14 +1,18 @@
 Metadata-Version: 2.4
 Name: combatlearn
-Version: 2.1.0
+Version: 2.3.0
 Summary: Batch-effect harmonization for machine learning frameworks.
 Author-email: Ettore Rocchi <ettoreroc@gmail.com>
-License: MIT
+License-Expression: MIT
 Keywords: machine-learning,harmonization,combat,preprocessing
-Classifier: Development Status :: 3 - Alpha
+Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Science/Research
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
@@ -23,6 +27,7 @@ Requires-Dist: umap-learn>=0.5
 Provides-Extra: dev
 Requires-Dist: pytest>=7; extra == "dev"
 Requires-Dist: pytest-cov>=4.0; extra == "dev"
+Requires-Dist: hypothesis>=6.0; extra == "dev"
 Requires-Dist: ruff>=0.8; extra == "dev"
 Requires-Dist: pre-commit>=3.0; extra == "dev"
 Requires-Dist: mypy>=1.0; extra == "dev"
@@ -50,11 +55,16 @@ Dynamic: license-file
 
 **combatlearn** makes the popular _ComBat_ (and _CovBat_) batch-effect correction algorithm available for use into machine learning frameworks. It lets you harmonise high-dimensional data inside a scikit-learn `Pipeline`, so that cross-validation and grid-search automatically take batch structure into account, **without data leakage**.
 
-**Four methods**:
+**Inductive `ComBat` methods** (fit on train, apply to held-out data, cross-validation-safe):
 - `method="johnson"` - classic ComBat (Johnson _et al._, 2007)
 - `method="fortin"` - neuroComBat (Fortin _et al._, 2018)
 - `method="chen"` - CovBat (Chen _et al._, 2022)
-- `method="longitudinal"` - Longitudinal ComBat for repeated measures (Beer _et al._, 2020)
+- `method="gam"` - ComBat-GAM, nonlinear (spline) covariate effects (Pomponio _et al._, 2020)
+- `method="covbat_gam"` - CovBat with the same nonlinear covariate modeling
+
+**Multiple batch variables** - `NestedComBat` harmonizes over several batch variables at once (e.g. site, scanner, protocol) with optional order optimization and Gaussian-mixture grouping (Nested / OPNested / GMM ComBat, Horng _et al._, 2022).
+
+**Whole-cohort (transductive)** - `combatlearn.transductive.TransductiveComBat(method="longitudinal")` for repeated-measures / longitudinal designs (Beer _et al._, 2020), used as a one-shot `fit_transform`. (`ComBat(method="longitudinal")` still works but is deprecated in favor of this and will be removed in v3.0.0.)
 
 ## Installation
 
@@ -119,6 +129,28 @@ print(f"Best CV AUROC: {grid.best_score_:.3f}")
 
 For a full example of how to use **combatlearn** see the [notebook demo](https://github.com/EttoreRocchi/combatlearn/blob/main/docs/source/demo/combatlearn_demo.ipynb)
 
+## Multiple batch variables (`NestedComBat`)
+
+When more than one technical variable needs harmonizing, `NestedComBat` applies ComBat to each in sequence. It picks the harmonization order that minimizes the residual batch effect, and can optionally add a latent Gaussian-mixture grouping. It is inductive and fits inside a `Pipeline` just like `ComBat`.
+
+```python
+import pandas as pd
+from combatlearn import NestedComBat
+
+batch = pd.DataFrame({"site": site, "scanner": scanner})  # one column per batch variable
+
+nested = NestedComBat(
+    batch=batch,
+    continuous_covariates=age,
+    discrete_covariates=diag,
+    method="fortin",      # per-step engine: "fortin"/"chen"/"gam"/"covbat_gam"
+    optimize_order=True,  # choose the order (OPNested)
+    gmm=None,             # or "batch" (+GMM) / "covariate" (-GMM)
+)
+X_corrected = nested.fit_transform(X)
+print("Harmonization order:", nested.order_)
+```
+
 ## `ComBat` parameters
 
 The following section provides a detailed explanation of all parameters available in the scikit-learn-compatible `ComBat` class. For complete API documentation, see the [API Reference](https://combatlearn.readthedocs.io/en/latest/api/).
@@ -128,8 +160,8 @@ The following section provides a detailed explanation of all parameters availabl
 | Parameter | Type | Default | Description |
 | --- | ---  | --- | --- |
 | `batch` | array-like or pd.Series | **required** | Vector indicating batch assignment for each sample. This is used to estimate and remove batch effects. |
-| `discrete_covariates` | array-like, pd.Series, or pd.DataFrame | `None` | Optional categorical covariates (e.g., sex, site). Only used in `"fortin"`, `"chen"`, and `"longitudinal"` methods. |
-| `continuous_covariates` | array-like, pd.Series or pd.DataFrame | `None` | Optional continuous covariates (e.g., age). Only used in `"fortin"`, `"chen"`, and `"longitudinal"` methods. |
+| `discrete_covariates` | array-like, pd.Series, or pd.DataFrame | `None` | Optional categorical covariates (e.g., sex, site). Only used in `"fortin"`, `"chen"`, `"longitudinal"`, `"gam"`, and `"covbat_gam"` methods. |
+| `continuous_covariates` | array-like, pd.Series or pd.DataFrame | `None` | Optional continuous covariates (e.g., age). Only used in `"fortin"`, `"chen"`, `"longitudinal"`, `"gam"`, and `"covbat_gam"` methods. **Required** for the GAM methods. |
 | `subject_id` | array-like or pd.Series | `None` | Subject/individual labels for the random intercept. **Required** for `method="longitudinal"`, ignored otherwise. |
 | `time_covariate` | array-like or pd.Series | `None` | Optional continuous time variable for repeated measures. Only used in `"longitudinal"`. |
 
@@ -137,11 +169,15 @@ The following section provides a detailed explanation of all parameters availabl
 
 | Parameter | Type | Default | Description |
 | --- | --- | --- | --- |
-| `method` | str | `"johnson"` | ComBat method to use: <ul><li>`"johnson"` - Classical ComBat (_Johnson et al. 2007_)</li><li>`"fortin"` - ComBat with covariates (_Fortin et al. 2018_)</li><li>`"chen"` - CovBat, PCA-based correction (_Chen et al. 2022_)</li><li>`"longitudinal"` - Longitudinal ComBat with a per-subject random intercept (_Beer et al. 2020_)</li></ul> |
+| `method` | str | `"johnson"` | ComBat method to use: <ul><li>`"johnson"` - Classical ComBat (_Johnson et al. 2007_)</li><li>`"fortin"` - ComBat with covariates (_Fortin et al. 2018_)</li><li>`"chen"` - CovBat, PCA-based correction (_Chen et al. 2022_)</li><li>`"longitudinal"` - Longitudinal ComBat with a per-subject random intercept (_Beer et al. 2020_); **deprecated**, use `combatlearn.transductive.TransductiveComBat`</li><li>`"gam"` - ComBat-GAM, nonlinear (spline) covariates (_Pomponio et al. 2020_)</li><li>`"covbat_gam"` - CovBat with the same nonlinear covariate modeling</li></ul> Case- and separator-insensitive literature aliases are also accepted (e.g. `"covbat"`, `"neurocombat"`, `"combat_gam"`). |
 | `parametric` | bool | `True` | Whether to use the **parametric empirical Bayes** formulation. If `False`, a non-parametric iterative scheme is used. |
 | `mean_only` | bool | `False` | If `True`, only the **mean** is corrected, while variances are left unchanged. Useful for preserving variance structure in the data. |
 | `reference_batch` | str or `None` | `None` | If specified, acts as a reference batch - other batches will be corrected to match this one. |
 | `covbat_cov_thresh` | float, int | `0.9` | For `"chen"` method only: Cumulative variance threshold in `(0, 1]` to retain PCs in PCA space (e.g., 0.9 = retain 90% explained variance). If an integer is provided, it represents the number of principal components to use. |
+| `smooth_terms` | list of str or int, or `None` | `None` | For `"gam"`/`"covbat_gam"` only: continuous covariates to model nonlinearly with B-splines, by column name or integer position. `None` smooths all continuous covariates. |
+| `spline_df` | int | `10` | For the GAM methods: B-spline degrees of freedom (basis functions) per smooth term. |
+| `spline_degree` | int | `3` | For the GAM methods: B-spline degree (3 = cubic). |
+| `smooth_term_bounds` | tuple or dict | `None` | For the GAM methods: boundary knots, as a single `(lo, hi)` for all terms or a `{term: (lo, hi)}` dict. Default uses each term's training min/max; widen to cover held-out data. |
 | `eps` | float | `1e-8` | Small jitter value added to variances to prevent divide-by-zero errors during standardization. |
 
 
@@ -197,3 +233,5 @@ Please consider citing the original papers:
 - [**CovBat**](https://github.com/andy1764/CovBat_Harmonization) - Chen AA et al. _Hum Brain Mapp_. 2022. doi: [10.1002/hbm.25688](https://doi.org/10.1002/hbm.25688)
 
 - [**Longitudinal ComBat**](https://github.com/jcbeer/longCombat) - Beer JC et al. _Neuroimage_. 2020. doi: [10.1016/j.neuroimage.2020.117129](https://doi.org/10.1016/j.neuroimage.2020.117129)
+
+- [**ComBat-GAM**](https://github.com/rpomponio/neuroHarmonize) - Pomponio R et al. _Neuroimage_. 2020. doi: [10.1016/j.neuroimage.2019.116450](https://doi.org/10.1016/j.neuroimage.2019.116450)

{combatlearn-2.1.0 → combatlearn-2.3.0}/README.md

@@ -12,11 +12,16 @@
 
 **combatlearn** makes the popular _ComBat_ (and _CovBat_) batch-effect correction algorithm available for use into machine learning frameworks. It lets you harmonise high-dimensional data inside a scikit-learn `Pipeline`, so that cross-validation and grid-search automatically take batch structure into account, **without data leakage**.
 
-**Four methods**:
+**Inductive `ComBat` methods** (fit on train, apply to held-out data, cross-validation-safe):
 - `method="johnson"` - classic ComBat (Johnson _et al._, 2007)
 - `method="fortin"` - neuroComBat (Fortin _et al._, 2018)
 - `method="chen"` - CovBat (Chen _et al._, 2022)
-- `method="longitudinal"` - Longitudinal ComBat for repeated measures (Beer _et al._, 2020)
+- `method="gam"` - ComBat-GAM, nonlinear (spline) covariate effects (Pomponio _et al._, 2020)
+- `method="covbat_gam"` - CovBat with the same nonlinear covariate modeling
+
+**Multiple batch variables** - `NestedComBat` harmonizes over several batch variables at once (e.g. site, scanner, protocol) with optional order optimization and Gaussian-mixture grouping (Nested / OPNested / GMM ComBat, Horng _et al._, 2022).
+
+**Whole-cohort (transductive)** - `combatlearn.transductive.TransductiveComBat(method="longitudinal")` for repeated-measures / longitudinal designs (Beer _et al._, 2020), used as a one-shot `fit_transform`. (`ComBat(method="longitudinal")` still works but is deprecated in favor of this and will be removed in v3.0.0.)
 
 ## Installation
 
@@ -81,6 +86,28 @@ print(f"Best CV AUROC: {grid.best_score_:.3f}")
 
 For a full example of how to use **combatlearn** see the [notebook demo](https://github.com/EttoreRocchi/combatlearn/blob/main/docs/source/demo/combatlearn_demo.ipynb)
 
+## Multiple batch variables (`NestedComBat`)
+
+When more than one technical variable needs harmonizing, `NestedComBat` applies ComBat to each in sequence. It picks the harmonization order that minimizes the residual batch effect, and can optionally add a latent Gaussian-mixture grouping. It is inductive and fits inside a `Pipeline` just like `ComBat`.
+
+```python
+import pandas as pd
+from combatlearn import NestedComBat
+
+batch = pd.DataFrame({"site": site, "scanner": scanner})  # one column per batch variable
+
+nested = NestedComBat(
+    batch=batch,
+    continuous_covariates=age,
+    discrete_covariates=diag,
+    method="fortin",      # per-step engine: "fortin"/"chen"/"gam"/"covbat_gam"
+    optimize_order=True,  # choose the order (OPNested)
+    gmm=None,             # or "batch" (+GMM) / "covariate" (-GMM)
+)
+X_corrected = nested.fit_transform(X)
+print("Harmonization order:", nested.order_)
+```
+
 ## `ComBat` parameters
 
 The following section provides a detailed explanation of all parameters available in the scikit-learn-compatible `ComBat` class. For complete API documentation, see the [API Reference](https://combatlearn.readthedocs.io/en/latest/api/).
@@ -90,8 +117,8 @@ The following section provides a detailed explanation of all parameters availabl
 | Parameter | Type | Default | Description |
 | --- | ---  | --- | --- |
 | `batch` | array-like or pd.Series | **required** | Vector indicating batch assignment for each sample. This is used to estimate and remove batch effects. |
-| `discrete_covariates` | array-like, pd.Series, or pd.DataFrame | `None` | Optional categorical covariates (e.g., sex, site). Only used in `"fortin"`, `"chen"`, and `"longitudinal"` methods. |
-| `continuous_covariates` | array-like, pd.Series or pd.DataFrame | `None` | Optional continuous covariates (e.g., age). Only used in `"fortin"`, `"chen"`, and `"longitudinal"` methods. |
+| `discrete_covariates` | array-like, pd.Series, or pd.DataFrame | `None` | Optional categorical covariates (e.g., sex, site). Only used in `"fortin"`, `"chen"`, `"longitudinal"`, `"gam"`, and `"covbat_gam"` methods. |
+| `continuous_covariates` | array-like, pd.Series or pd.DataFrame | `None` | Optional continuous covariates (e.g., age). Only used in `"fortin"`, `"chen"`, `"longitudinal"`, `"gam"`, and `"covbat_gam"` methods. **Required** for the GAM methods. |
 | `subject_id` | array-like or pd.Series | `None` | Subject/individual labels for the random intercept. **Required** for `method="longitudinal"`, ignored otherwise. |
 | `time_covariate` | array-like or pd.Series | `None` | Optional continuous time variable for repeated measures. Only used in `"longitudinal"`. |
 
@@ -99,11 +126,15 @@ The following section provides a detailed explanation of all parameters availabl
 
 | Parameter | Type | Default | Description |
 | --- | --- | --- | --- |
-| `method` | str | `"johnson"` | ComBat method to use: <ul><li>`"johnson"` - Classical ComBat (_Johnson et al. 2007_)</li><li>`"fortin"` - ComBat with covariates (_Fortin et al. 2018_)</li><li>`"chen"` - CovBat, PCA-based correction (_Chen et al. 2022_)</li><li>`"longitudinal"` - Longitudinal ComBat with a per-subject random intercept (_Beer et al. 2020_)</li></ul> |
+| `method` | str | `"johnson"` | ComBat method to use: <ul><li>`"johnson"` - Classical ComBat (_Johnson et al. 2007_)</li><li>`"fortin"` - ComBat with covariates (_Fortin et al. 2018_)</li><li>`"chen"` - CovBat, PCA-based correction (_Chen et al. 2022_)</li><li>`"longitudinal"` - Longitudinal ComBat with a per-subject random intercept (_Beer et al. 2020_); **deprecated**, use `combatlearn.transductive.TransductiveComBat`</li><li>`"gam"` - ComBat-GAM, nonlinear (spline) covariates (_Pomponio et al. 2020_)</li><li>`"covbat_gam"` - CovBat with the same nonlinear covariate modeling</li></ul> Case- and separator-insensitive literature aliases are also accepted (e.g. `"covbat"`, `"neurocombat"`, `"combat_gam"`). |
 | `parametric` | bool | `True` | Whether to use the **parametric empirical Bayes** formulation. If `False`, a non-parametric iterative scheme is used. |
 | `mean_only` | bool | `False` | If `True`, only the **mean** is corrected, while variances are left unchanged. Useful for preserving variance structure in the data. |
 | `reference_batch` | str or `None` | `None` | If specified, acts as a reference batch - other batches will be corrected to match this one. |
 | `covbat_cov_thresh` | float, int | `0.9` | For `"chen"` method only: Cumulative variance threshold in `(0, 1]` to retain PCs in PCA space (e.g., 0.9 = retain 90% explained variance). If an integer is provided, it represents the number of principal components to use. |
+| `smooth_terms` | list of str or int, or `None` | `None` | For `"gam"`/`"covbat_gam"` only: continuous covariates to model nonlinearly with B-splines, by column name or integer position. `None` smooths all continuous covariates. |
+| `spline_df` | int | `10` | For the GAM methods: B-spline degrees of freedom (basis functions) per smooth term. |
+| `spline_degree` | int | `3` | For the GAM methods: B-spline degree (3 = cubic). |
+| `smooth_term_bounds` | tuple or dict | `None` | For the GAM methods: boundary knots, as a single `(lo, hi)` for all terms or a `{term: (lo, hi)}` dict. Default uses each term's training min/max; widen to cover held-out data. |
 | `eps` | float | `1e-8` | Small jitter value added to variances to prevent divide-by-zero errors during standardization. |
 
 
@@ -158,4 +189,6 @@ Please consider citing the original papers:
 
 - [**CovBat**](https://github.com/andy1764/CovBat_Harmonization) - Chen AA et al. _Hum Brain Mapp_. 2022. doi: [10.1002/hbm.25688](https://doi.org/10.1002/hbm.25688)
 
-- [**Longitudinal ComBat**](https://github.com/jcbeer/longCombat) - Beer JC et al. _Neuroimage_. 2020. doi: [10.1016/j.neuroimage.2020.117129](https://doi.org/10.1016/j.neuroimage.2020.117129)
+- [**Longitudinal ComBat**](https://github.com/jcbeer/longCombat) - Beer JC et al. _Neuroimage_. 2020. doi: [10.1016/j.neuroimage.2020.117129](https://doi.org/10.1016/j.neuroimage.2020.117129)
+
+- [**ComBat-GAM**](https://github.com/rpomponio/neuroHarmonize) - Pomponio R et al. _Neuroimage_. 2020. doi: [10.1016/j.neuroimage.2019.116450](https://doi.org/10.1016/j.neuroimage.2019.116450)

combatlearn-2.3.0/combatlearn/__init__.py

@@ -0,0 +1,5 @@
+from .sklearn_api import ComBat, NestedComBat
+
+__all__ = ["ComBat", "NestedComBat"]
+__version__ = "2.3.0"
+__author__ = "Ettore Rocchi"

combatlearn-2.3.0/combatlearn/_nested.py

@@ -0,0 +1,326 @@
+"""Orchestration helpers for NestedComBat (*Horng et al.* 2022).
+
+Nested ComBat harmonizes several batch variables (e.g. site, scanner, protocol)
+by applying single-batch ComBat to each one in sequence, every step delegating
+to a :class:`~combatlearn.core.ComBatModel`. This module holds the pieces that
+surround that sequence and add no new empirical-Bayes math:
+
+* :func:`fit_nested_sequence` - fit one ``ComBatModel`` per batch variable in a
+  given order, chaining the transforms.
+* :func:`anderson_darling_batch_count` - the residual-batch-effect objective the
+  OPNested order search minimizes (Anderson-Darling k-sample feature count).
+* :func:`search_order` - exhaustive order search over all ``k!`` orderings, with
+  a factorial guardrail that falls back to greedy forward selection.
+* :class:`GMMGrouping` - the per-feature Gaussian-mixture grouping fed back into
+  the sequence as an extra batch variable (``+GMM``) or covariate (``-GMM``).
+
+Everything here is learned on the training data and frozen for transform, so the
+resulting estimator stays inductive / cross-validation-safe.
+"""
+
+from __future__ import annotations
+
+import warnings
+from collections.abc import Callable, Mapping, Sequence
+from itertools import permutations
+from math import factorial
+
+import numpy as np
+import pandas as pd
+from scipy.stats import anderson_ksamp
+from sklearn.mixture import GaussianMixture
+
+from .core import ComBatModel
+
+GMM_VAR_NAME = "GMM"
+
+ModelStep = tuple[str, ComBatModel]
+
+
+def variance_explained_by_batch(X: pd.DataFrame, batch: pd.Series) -> float:
+    """Fraction of total variance explained by ``batch`` (summed across features)."""
+    values = X.to_numpy(dtype=np.float64)
+    labels = np.asarray(batch)
+    grand = values.mean(axis=0)
+    ss_total = float(((values - grand) ** 2).sum())
+    if ss_total == 0.0:
+        return 0.0
+    ss_between = 0.0
+    for lvl in pd.unique(labels):
+        mask = labels == lvl
+        batch_mean = values[mask].mean(axis=0)
+        ss_between += float(mask.sum()) * float(((batch_mean - grand) ** 2).sum())
+    return ss_between / ss_total
+
+
+def anderson_darling_batch_count(
+    X: pd.DataFrame,
+    batch_vars: Mapping[str, pd.Series],
+    *,
+    alpha: float = 0.05,
+) -> int:
+    """Number of features with a significant residual batch effect.
+
+    For every batch variable and every feature the Anderson-Darling k-sample test
+    (:func:`scipy.stats.anderson_ksamp`) compares the feature distribution across
+    that variable's levels; a feature is counted once per batch variable whose
+    test is significant at ``alpha``. The counts are summed over all batch
+    variables, so a lower value means less residual batch structure. This is the
+    objective minimized by the OPNested order search.
+
+    Parameters
+    ----------
+    X : pd.DataFrame of shape (n_samples, n_features)
+        Harmonized data to score.
+    batch_vars : mapping of str to pd.Series
+        Batch variables, each a per-sample label vector.
+    alpha : float, default=0.05
+        Significance threshold for the Anderson-Darling test.
+
+    Returns
+    -------
+    int
+        Total number of significant (feature, batch variable) pairs.
+    """
+    col_values = X.to_numpy(dtype=np.float64)
+    n_features = col_values.shape[1]
+    total = 0
+    with warnings.catch_warnings():
+        # anderson_ksamp warns when the p-value is capped to [0.001, 0.25]; that
+        # capping never changes the alpha=0.05 decision, so the warning is considered noise.
+        warnings.simplefilter("ignore")
+        for batch in batch_vars.values():
+            labels = np.asarray(batch)
+            masks = [labels == lvl for lvl in pd.unique(labels)]
+            masks = [m for m in masks if m.any()]
+            if len(masks) < 2:
+                continue
+            for j in range(n_features):
+                col = col_values[:, j]
+                samples = [col[m] for m in masks]
+                try:
+                    result = anderson_ksamp(samples)
+                except ValueError:
+                    # Degenerate feature (e.g. constant within the pooled sample):
+                    # no detectable batch effect, so it does not count.
+                    continue
+                if result.significance_level < alpha:
+                    total += 1
+    return total
+
+
+def fit_nested_sequence(
+    order: Sequence[str],
+    X: pd.DataFrame,
+    batch_vars: Mapping[str, pd.Series],
+    disc: pd.DataFrame | None,
+    cont: pd.DataFrame | None,
+    make_model: Callable[[str], ComBatModel],
+) -> tuple[list[ModelStep], pd.DataFrame]:
+    """Fit one ``ComBatModel`` per batch variable in ``order``, chaining transforms.
+
+    The covariates ``disc`` / ``cont`` are preserved across every step, so nesting
+    composes with the covariate-aware, CovBat, and GAM engines.
+
+    Returns the list of ``(batch_variable, fitted_model)`` steps and the fully
+    harmonized training data.
+    """
+    models: list[ModelStep] = []
+    X_current = X
+    for name in order:
+        model = make_model(name)
+        batch = batch_vars[name]
+        model.fit(X_current, batch=batch, discrete_covariates=disc, continuous_covariates=cont)
+        X_current = model.transform(
+            X_current, batch=batch, discrete_covariates=disc, continuous_covariates=cont
+        )
+        models.append((name, model))
+    return models, X_current
+
+
+def transform_nested_sequence(
+    models: Sequence[ModelStep],
+    X: pd.DataFrame,
+    batch_vars: Mapping[str, pd.Series],
+    disc: pd.DataFrame | None,
+    cont: pd.DataFrame | None,
+) -> pd.DataFrame:
+    """Replay a fitted nested sequence on new data, in the stored order."""
+    X_current = X
+    for name, model in models:
+        batch = batch_vars[name]
+        X_current = model.transform(
+            X_current, batch=batch, discrete_covariates=disc, continuous_covariates=cont
+        )
+    return X_current
+
+
+def _search_exhaustive(
+    names: Sequence[str],
+    X: pd.DataFrame,
+    batch_vars: Mapping[str, pd.Series],
+    disc: pd.DataFrame | None,
+    cont: pd.DataFrame | None,
+    make_model: Callable[[str], ComBatModel],
+    alpha: float,
+) -> tuple[list[str], list[ModelStep]]:
+    """Try every ordering, keeping the first that minimizes the AD feature count."""
+    best_score: int | None = None
+    best_order: list[str] = list(names)
+    best_models: list[ModelStep] = []
+    for order in permutations(names):
+        models, X_harmonized = fit_nested_sequence(order, X, batch_vars, disc, cont, make_model)
+        score = anderson_darling_batch_count(X_harmonized, batch_vars, alpha=alpha)
+        if best_score is None or score < best_score:
+            best_score = score
+            best_order = list(order)
+            best_models = models
+    return best_order, best_models
+
+
+def _search_greedy(
+    names: Sequence[str],
+    X: pd.DataFrame,
+    batch_vars: Mapping[str, pd.Series],
+    disc: pd.DataFrame | None,
+    cont: pd.DataFrame | None,
+    make_model: Callable[[str], ComBatModel],
+    alpha: float,
+) -> tuple[list[str], list[ModelStep]]:
+    """Greedy forward selection: at each position pick the step that most reduces
+    the AD feature count, harmonize with it, then continue with the rest."""
+    remaining = list(names)
+    order: list[str] = []
+    models: list[ModelStep] = []
+    X_current = X
+    while remaining:
+        best_score: int | None = None
+        best_name = remaining[0]
+        best_model: ComBatModel | None = None
+        best_X = X_current
+        for name in remaining:
+            model = make_model(name)
+            batch = batch_vars[name]
+            model.fit(X_current, batch=batch, discrete_covariates=disc, continuous_covariates=cont)
+            X_try = model.transform(
+                X_current, batch=batch, discrete_covariates=disc, continuous_covariates=cont
+            )
+            score = anderson_darling_batch_count(X_try, batch_vars, alpha=alpha)
+            if best_score is None or score < best_score:
+                best_score = score
+                best_name = name
+                best_model = model
+                best_X = X_try
+        assert best_model is not None
+        order.append(best_name)
+        models.append((best_name, best_model))
+        X_current = best_X
+        remaining.remove(best_name)
+    return order, models
+
+
+def search_order(
+    names: Sequence[str],
+    X: pd.DataFrame,
+    batch_vars: Mapping[str, pd.Series],
+    disc: pd.DataFrame | None,
+    cont: pd.DataFrame | None,
+    make_model: Callable[[str], ComBatModel],
+    *,
+    max_exhaustive_vars: int,
+    alpha: float = 0.05,
+) -> tuple[list[str], list[ModelStep], bool]:
+    """Select a harmonization order that minimizes the residual batch effect.
+
+    All ``k!`` orderings are tried exhaustively when ``k <= max_exhaustive_vars``;
+    above the cap the factorial cost is prohibitive, so a warning is issued and
+    the search falls back to greedy forward selection. A large exhaustive search
+    (``k >= 4``) also warns up front with the number of ComBat fits it will run.
+
+    Returns the chosen order, the fitted sequence for that order, and whether the
+    greedy fallback was used.
+    """
+    k = len(names)
+    if k <= max_exhaustive_vars:
+        if k >= 4:
+            n_orderings = factorial(k)
+            warnings.warn(
+                f"OPNested exhaustive order search over {k} batch variables will fit "
+                f"{n_orderings} orderings x {k} steps = {n_orderings * k} ComBat models, each "
+                f"scored with a per-feature Anderson-Darling test; on high-dimensional data "
+                f"this can be slow. Lower max_exhaustive_vars to use the greedy search instead, "
+                f"or set optimize_order=False to skip the search.",
+                UserWarning,
+                stacklevel=2,
+            )
+        order, models = _search_exhaustive(names, X, batch_vars, disc, cont, make_model, alpha)
+        return order, models, False
+    warnings.warn(
+        f"{len(names)} batch variables exceed max_exhaustive_vars="
+        f"{max_exhaustive_vars}; an exhaustive search over {len(names)}! orderings is "
+        f"impractical, so falling back to greedy forward selection. Raise "
+        f"max_exhaustive_vars to force the exhaustive search.",
+        UserWarning,
+        stacklevel=2,
+    )
+    order, models = _search_greedy(names, X, batch_vars, disc, cont, make_model, alpha)
+    return order, models, True
+
+
+class GMMGrouping:
+    """A per-feature Gaussian-mixture grouping (*Horng et al.* 2022, GMM ComBat).
+
+    A 2-component Gaussian mixture is fit to each feature; features whose split
+    leaves either cluster below ``min_cluster_frac`` of the samples are discarded
+    (an unbalanced split rarely reflects a real latent group), and the balanced
+    feature with the lowest AIC is retained. Its two-component assignment is the
+    latent grouping fed back into NestedComBat as an extra batch variable
+    (``+GMM``) or a protected covariate (``-GMM``).
+
+    The mixture is fit on the training data and stored, so held-out samples are
+    assigned by :meth:`assign` (``predict`` on the same fitted mixture), keeping
+    the grouping leakage-free.
+    """
+
+    def __init__(self, feature: object, gmm: GaussianMixture) -> None:
+        self.feature = feature
+        self.gmm = gmm
+
+    @classmethod
+    def fit(
+        cls,
+        X: pd.DataFrame,
+        *,
+        min_cluster_frac: float = 0.25,
+        random_state: int | None = 0,
+    ) -> GMMGrouping | None:
+        """Fit the grouping, or return ``None`` if no feature yields a balanced split."""
+        n_samples = len(X)
+        best_feature: object = None
+        best_aic = np.inf
+        best_gmm: GaussianMixture | None = None
+        for col in X.columns:
+            x = X[col].to_numpy(dtype=np.float64).reshape(-1, 1)
+            gmm = GaussianMixture(n_components=2, random_state=random_state)
+            try:
+                gmm.fit(x)
+                labels = gmm.predict(x)
+            except ValueError:
+                continue
+            smaller = min(int((labels == 0).sum()), int((labels == 1).sum()))
+            if smaller <= min_cluster_frac * n_samples:
+                continue
+            aic = float(gmm.aic(x))
+            if aic < best_aic:
+                best_aic = aic
+                best_feature = col
+                best_gmm = gmm
+        if best_gmm is None:
+            return None
+        return cls(best_feature, best_gmm)
+
+    def assign(self, X: pd.DataFrame) -> pd.Series:
+        """Assign each sample to its mixture component (``'gmm0'`` / ``'gmm1'``)."""
+        x = X[self.feature].to_numpy(dtype=np.float64).reshape(-1, 1)
+        labels = self.gmm.predict(x)
+        return pd.Series([f"gmm{int(c)}" for c in labels], index=X.index, name=GMM_VAR_NAME)

combatlearn-2.3.0/combatlearn/_utils.py

@@ -0,0 +1,66 @@
+"""Internal utilities shared across combatlearn modules."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import numpy as np
+import pandas as pd
+
+from .core import ArrayLike
+
+
+def _check_positional_alignment(
+    X: ArrayLike, auxiliaries: Iterable[tuple[str, ArrayLike | None]]
+) -> None:
+    """Require length-matched auxiliaries when ``X`` carries no index labels.
+
+    ``batch`` and the covariates are supplied at construction and label-aligned to
+    ``X`` via its pandas index. When ``X`` is a plain array it has no labels, so the
+    only possible alignment is positional, which needs equal lengths. This most
+    often breaks inside cross-validation: the CV splitter subsets ``X`` positionally
+    while the construction-time ``batch``/covariates stay full-length, and the
+    silent length mismatch would otherwise surface as a confusing pandas reindexing
+    error. A DataFrame ``X`` label-aligns and is skipped.
+
+    Parameters
+    ----------
+    X : array-like
+        The data passed to ``fit``/``transform``.
+    auxiliaries : iterable of (name, object)
+        The construction-time vectors to check; ``None`` entries are ignored.
+
+    Raises
+    ------
+    ValueError
+        If ``X`` is not a DataFrame and any auxiliary has a different length.
+    """
+    if isinstance(X, pd.DataFrame):
+        return
+    n = len(X)
+    for name, obj in auxiliaries:
+        if obj is None:
+            continue
+        if len(obj) != n:
+            raise ValueError(
+                f"`{name}` has length {len(obj)} but X has {n} rows. When X is not a "
+                f"pandas DataFrame it has no index labels, so combatlearn can only align "
+                f"`{name}` to X positionally, which requires equal lengths. This typically "
+                f"happens inside cross-validation, where the splitter subsets X but the "
+                f"batch/covariates passed at construction stay full-length. Pass X and "
+                f"`{name}` as pandas objects sharing the same index so folds stay "
+                f"label-aligned, or pass a same-length `{name}`."
+            )
+
+
+def _subset(obj: ArrayLike | None, idx: pd.Index) -> pd.DataFrame | pd.Series | None:
+    """Subset array-like object by index."""
+    if obj is None:
+        return None
+    if isinstance(obj, pd.Series | pd.DataFrame):
+        return obj.loc[idx]
+    else:
+        if isinstance(obj, np.ndarray) and obj.ndim == 1:
+            return pd.Series(obj, index=idx)
+        else:
+            return pd.DataFrame(obj, index=idx)