combatlearn 2.0.0__tar.gz → 2.2.0__tar.gz

{combatlearn-2.0.0 → combatlearn-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: combatlearn
-Version: 2.0.0
+Version: 2.2.0
 Summary: Batch-effect harmonization for machine learning frameworks.
 Author-email: Ettore Rocchi <ettoreroc@gmail.com>
 License: MIT
@@ -27,6 +27,7 @@ Requires-Dist: ruff>=0.8; extra == "dev"
 Requires-Dist: pre-commit>=3.0; extra == "dev"
 Requires-Dist: mypy>=1.0; extra == "dev"
 Requires-Dist: pandas-stubs; extra == "dev"
+Requires-Dist: statsmodels>=0.14; extra == "dev"
 Provides-Extra: docs
 Requires-Dist: sphinx>=7.0.0; extra == "docs"
 Requires-Dist: pydata-sphinx-theme>=0.15.0; extra == "docs"
@@ -49,10 +50,13 @@ 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**.
 
-**Three methods**:
+**Six methods**:
 - `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
 
 ## Installation
 
@@ -75,6 +79,7 @@ import pandas as pd
 from sklearn.pipeline import Pipeline
 from sklearn.preprocessing import StandardScaler
 from sklearn.linear_model import LogisticRegression
+from sklearn.model_selection import GridSearchCV
 from combatlearn import ComBat
 
 df = pd.read_csv("data.csv", index_col=0)
@@ -125,18 +130,24 @@ 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"` and `"chen"` methods. |
-| `continuous_covariates` | array-like, pd.Series or pd.DataFrame | `None` | Optional continuous covariates (e.g., age). Only used in `"fortin"` and `"chen"` 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"`. |
 
 ### Algorithm Options
 
 | 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></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_)</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 $]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. |
+| `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. |
 
 
@@ -190,3 +201,7 @@ Please consider citing the original papers:
 - [**neuroCombat**](https://github.com/Jfortin1/neuroCombat) - Fortin JP et al. _Neuroimage_. 2018. doi: [10.1016/j.neuroimage.2017.11.024](https://doi.org/10.1016/j.neuroimage.2017.11.024)
 
 - [**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.0.0 → combatlearn-2.2.0}/README.md

@@ -12,10 +12,13 @@
 
 **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**.
 
-**Three methods**:
+**Six methods**:
 - `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
 
 ## Installation
 
@@ -38,6 +41,7 @@ import pandas as pd
 from sklearn.pipeline import Pipeline
 from sklearn.preprocessing import StandardScaler
 from sklearn.linear_model import LogisticRegression
+from sklearn.model_selection import GridSearchCV
 from combatlearn import ComBat
 
 df = pd.read_csv("data.csv", index_col=0)
@@ -88,18 +92,24 @@ 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"` and `"chen"` methods. |
-| `continuous_covariates` | array-like, pd.Series or pd.DataFrame | `None` | Optional continuous covariates (e.g., age). Only used in `"fortin"` and `"chen"` 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"`. |
 
 ### Algorithm Options
 
 | 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></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_)</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 $]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. |
+| `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. |
 
 
@@ -152,4 +162,8 @@ Please consider citing the original papers:
 
 - [**neuroCombat**](https://github.com/Jfortin1/neuroCombat) - Fortin JP et al. _Neuroimage_. 2018. doi: [10.1016/j.neuroimage.2017.11.024](https://doi.org/10.1016/j.neuroimage.2017.11.024)
 
-- [**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)
+- [**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.0.0 → combatlearn-2.2.0}/combatlearn/__init__.py

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

combatlearn-2.2.0/combatlearn/_mixed.py

@@ -0,0 +1,165 @@
+"""Random-intercept linear mixed model fitting for Longitudinal ComBat.
+
+Fits, for every feature (column of ``Y``) simultaneously, the model
+
+    y = X @ beta + Z @ u + eps,    u ~ N(0, tau^2 I),   eps ~ N(0, sigma^2 I)
+
+where ``Z`` is the subject (random-intercept) indicator matrix. Estimation is
+by REML, profiling out ``beta`` and ``sigma^2`` so that only the scalar
+variance ratio ``lambda = tau^2 / sigma^2`` is optimized per feature.
+
+The single random-intercept grouping makes the marginal covariance
+``V = sigma^2 (I + lambda Z Z^T)`` block diagonal (one block per subject), so
+``V^{-1}`` is available in closed form via Sherman-Morrison and every operation
+is O(n) per feature. The variance ratio is found by a vectorized grid search
+over ``lambda`` followed by a parabolic refinement, after which the fixed
+effects, residual variance, and random-intercept BLUPs are computed in a single
+per-feature pass.
+
+This keeps Longitudinal ComBat dependency-free (numpy only); correctness is
+cross-checked against ``statsmodels.MixedLM`` in the test suite.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import numpy.typing as npt
+
+FloatArray = npt.NDArray[np.float64]
+IntArray = npt.NDArray[np.intp]
+
+
+def _group_sums(M: FloatArray, group_idx: IntArray, n_groups: int) -> FloatArray:
+    """Sum the rows of ``M`` (n, c) within each group, returning (n_groups, c)."""
+    out = np.zeros((n_groups, M.shape[1]), dtype=np.float64)
+    np.add.at(out, group_idx, M)
+    return out
+
+
+def _w_apply(M: FloatArray, factor: FloatArray, group_idx: IntArray, n_groups: int) -> FloatArray:
+    """Apply ``W = (I + lambda Z Z^T)^{-1}`` to the columns of ``M``.
+
+    ``factor[k] = lambda / (1 + lambda * n_k)`` is the Sherman-Morrison
+    per-group correction (``n_k`` = number of observations in group ``k``).
+    """
+    gs = _group_sums(M, group_idx, n_groups)
+    return M - factor[group_idx][:, None] * gs[group_idx]
+
+
+def _reml_objective(
+    lam: float,
+    X: FloatArray,
+    Y: FloatArray,
+    group_idx: IntArray,
+    n_groups: int,
+    n_k: FloatArray,
+    n: int,
+    p: int,
+    ridge: float,
+) -> FloatArray:
+    """Profiled REML objective (to minimize) evaluated for all features at ``lam``."""
+    factor = lam / (1.0 + lam * n_k)
+    WX = _w_apply(X, factor, group_idx, n_groups)
+    A = X.T @ WX
+    A.flat[:: p + 1] += ridge
+    beta = np.linalg.solve(A, X.T @ _w_apply(Y, factor, group_idx, n_groups))
+    R = Y - X @ beta
+    WR = _w_apply(R, factor, group_idx, n_groups)
+    rss = np.maximum(np.einsum("ng,ng->g", R, WR), 1e-12)
+    logdet_v = float(np.sum(np.log1p(lam * n_k)))
+    _, logdet_a = np.linalg.slogdet(A)
+    result = logdet_v + (n - p) * np.log(rss) + float(logdet_a)
+    return np.asarray(result, dtype=np.float64)
+
+
+def fit_random_intercept(
+    X: FloatArray,
+    Y: FloatArray,
+    group_idx: IntArray,
+    n_groups: int,
+    n_k: FloatArray,
+    *,
+    eps: float = 1e-8,
+    n_grid: int = 80,
+) -> tuple[FloatArray, FloatArray, FloatArray, FloatArray]:
+    """Fit a random-intercept LME per feature by REML.
+
+    Parameters
+    ----------
+    X : ndarray of shape (n, p)
+        Fixed-effects design matrix (shared across features).
+    Y : ndarray of shape (n, n_features)
+        Response matrix; one column per feature.
+    group_idx : ndarray of shape (n,)
+        Integer group (subject) code in ``[0, n_groups)`` for each observation.
+    n_groups : int
+        Number of distinct subjects.
+    n_k : ndarray of shape (n_groups,)
+        Number of observations per subject.
+    eps : float, default=1e-8
+        Numerical jitter; also scales the diagonal ridge on ``X^T W X``.
+    n_grid : int, default=80
+        Number of log-spaced ``lambda`` grid points for the variance-ratio search.
+
+    Returns
+    -------
+    beta : ndarray of shape (p, n_features)
+        Estimated fixed-effects coefficients.
+    blup : ndarray of shape (n_groups, n_features)
+        Best linear unbiased predictors of the subject random intercepts.
+    sigma2 : ndarray of shape (n_features,)
+        Estimated residual variance per feature.
+    lam_opt : ndarray of shape (n_features,)
+        Estimated variance ratio ``tau^2 / sigma^2`` per feature.
+    """
+    X = np.asarray(X, dtype=np.float64)
+    Y = np.asarray(Y, dtype=np.float64)
+    group_idx = np.asarray(group_idx, dtype=np.intp)
+    n_k = np.asarray(n_k, dtype=np.float64)
+    n, p = X.shape
+    n_features = Y.shape[1]
+    ridge = eps * float(np.trace(X.T @ X)) / max(p, 1)
+
+    # Vectorized grid search for the variance ratio (per feature).
+    log_grid = np.linspace(-4.0, 4.0, n_grid)
+    lam_grid = np.concatenate([[0.0], 10.0**log_grid])
+    fvals = np.empty((lam_grid.size, n_features), dtype=np.float64)
+    for i, lam in enumerate(lam_grid):
+        fvals[i] = _reml_objective(lam, X, Y, group_idx, n_groups, n_k, n, p, ridge)
+    best = np.argmin(fvals, axis=0)
+    lam_opt = lam_grid[best].astype(np.float64)
+
+    # Vectorized parabolic refinement (log10 space) for interior positive optima.
+    interior = (best >= 2) & (best <= n_grid - 1)
+    if np.any(interior):
+        bi = best[interior]
+        f_lo = fvals[bi - 1, interior]
+        f0 = fvals[bi, interior]
+        f_hi = fvals[bi + 1, interior]
+        denom = f_lo - 2.0 * f0 + f_hi
+        delta = np.where(np.abs(denom) > 1e-12, 0.5 * (f_lo - f_hi) / denom, 0.0)
+        delta = np.clip(delta, -1.0, 1.0)
+        step = log_grid[1] - log_grid[0]
+        lam_opt[interior] = 10.0 ** (log_grid[bi - 1] + delta * step)
+
+    # Final per-feature fixed effects, residual variance, and random-intercept BLUPs.
+    beta = np.empty((p, n_features), dtype=np.float64)
+    sigma2 = np.empty(n_features, dtype=np.float64)
+    blup = np.zeros((n_groups, n_features), dtype=np.float64)
+    denom_sigma = max(n - p, 1)
+    for g in range(n_features):
+        lam = float(lam_opt[g])
+        factor = lam / (1.0 + lam * n_k)
+        A = X.T @ _w_apply(X, factor, group_idx, n_groups)
+        A.flat[:: p + 1] += ridge
+        y = Y[:, g : g + 1]
+        b = np.linalg.solve(A, X.T @ _w_apply(y, factor, group_idx, n_groups))[:, 0]
+        beta[:, g] = b
+        r = Y[:, g] - X @ b
+        Wr = _w_apply(r[:, None], factor, group_idx, n_groups)[:, 0]
+        sigma2[g] = float(r @ Wr) / denom_sigma
+        s_k = np.zeros(n_groups, dtype=np.float64)
+        np.add.at(s_k, group_idx, r)
+        blup[:, g] = lam * s_k / (1.0 + lam * n_k)
+
+    return beta, blup, sigma2, lam_opt