combatlearn 2.0.0__tar.gz → 2.1.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: combatlearn
-Version: 2.0.0
+Version: 2.1.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,11 @@ 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**:
+**Four 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)
 
 ## Installation
 
@@ -75,6 +77,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 +128,20 @@ 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"`, 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. |
+| `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></ul> |
 | `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. |
 | `eps` | float | `1e-8` | Small jitter value added to variances to prevent divide-by-zero errors during standardization. |
 
 
@@ -190,3 +195,5 @@ 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)

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

@@ -12,10 +12,11 @@
 
 **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**:
+**Four 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)
 
 ## Installation
 
@@ -38,6 +39,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 +90,20 @@ 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"`, 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. |
+| `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></ul> |
 | `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. |
 | `eps` | float | `1e-8` | Small jitter value added to variances to prevent divide-by-zero errors during standardization. |
 
 
@@ -152,4 +156,6 @@ 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)

{combatlearn-2.0.0 → combatlearn-2.1.0}/combatlearn/__init__.py

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

combatlearn-2.1.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

{combatlearn-2.0.0 → combatlearn-2.1.0}/combatlearn/core.py

@@ -1,9 +1,10 @@
 """ComBat algorithm core.
 
-`ComBatModel` implements three variants of the ComBat algorithm:
+`ComBatModel` implements four variants of the ComBat algorithm:
     * Johnson et al. (2007) vanilla ComBat (method="johnson")
     * Fortin et al. (2018) extension with covariates (method="fortin")
     * Chen et al. (2022) CovBat (method="chen")
+    * Beer et al. (2020) Longitudinal ComBat (method="longitudinal")
 """
 
 from __future__ import annotations
@@ -17,6 +18,8 @@ import numpy.typing as npt
 import pandas as pd
 from sklearn.decomposition import PCA
 
+from ._mixed import fit_random_intercept
+
 ArrayLike: TypeAlias = pd.DataFrame | pd.Series | npt.NDArray[Any]
 FloatArray: TypeAlias = npt.NDArray[np.float64]
 
@@ -26,10 +29,12 @@ class ComBatModel:
 
     Parameters
     ----------
-    method : {'johnson', 'fortin', 'chen'}, default='johnson'
+    method : {'johnson', 'fortin', 'chen', 'longitudinal'}, default='johnson'
         * 'johnson' - classic ComBat.
         * 'fortin' - covariate-aware ComBat.
         * 'chen' - CovBat, PCA-based ComBat.
+        * 'longitudinal' - Longitudinal ComBat; requires a per-sample
+          ``subject_id`` for the random intercept.
     parametric : bool, default=True
         Use the parametric empirical Bayes variant.
     mean_only : bool, default=False
@@ -47,7 +52,7 @@ class ComBatModel:
     def __init__(
         self,
         *,
-        method: Literal["johnson", "fortin", "chen"] = "johnson",
+        method: Literal["johnson", "fortin", "chen", "longitudinal"] = "johnson",
         parametric: bool = True,
         mean_only: bool = False,
         reference_batch: str | None = None,
@@ -77,6 +82,8 @@ class ComBatModel:
         self._batch_levels_pc: pd.Index
         self._pc_gamma_star: FloatArray
         self._pc_delta_star: FloatArray
+        self._re_subject_levels: pd.Index
+        self._re_blup: FloatArray
 
         # Validate covbat_cov_thresh
         if isinstance(self.covbat_cov_thresh, float):
@@ -138,6 +145,8 @@ class ComBatModel:
         batch: ArrayLike,
         discrete_covariates: ArrayLike | None = None,
         continuous_covariates: ArrayLike | None = None,
+        subject_id: ArrayLike | None = None,
+        time_covariate: ArrayLike | None = None,
     ) -> ComBatModel:
         """Fit the ComBat model.
 
@@ -150,9 +159,15 @@ class ComBatModel:
         batch : array-like of shape (n_samples,)
             Batch labels for each sample.
         discrete_covariates : array-like or None, default=None
-            Categorical covariates to protect (Fortin/Chen only).
+            Categorical covariates to protect (Fortin/Chen/Longitudinal only).
         continuous_covariates : array-like or None, default=None
-            Continuous covariates to protect (Fortin/Chen only).
+            Continuous covariates to protect (Fortin/Chen/Longitudinal only).
+        subject_id : array-like or None, default=None
+            Subject/individual labels for the random intercept. Required for
+            ``method='longitudinal'``, ignored otherwise.
+        time_covariate : array-like or None, default=None
+            Continuous time variable for repeated measures (Longitudinal only).
+            Added to the fixed-effects design alongside ``continuous_covariates``.
 
         Returns
         -------
@@ -160,10 +175,10 @@ class ComBatModel:
             Fitted model.
         """
         method = self.method.lower()
-        if method not in {"johnson", "fortin", "chen"}:
+        if method not in {"johnson", "fortin", "chen", "longitudinal"}:
             raise ValueError(
                 f"method={self.method!r} is not recognized. "
-                f"Expected one of 'johnson', 'fortin', or 'chen'."
+                f"Expected one of 'johnson', 'fortin', 'chen', or 'longitudinal'."
             )
         if not isinstance(X, pd.DataFrame):
             X = pd.DataFrame(X)
@@ -180,6 +195,13 @@ class ComBatModel:
                 f"Check for typos or trailing whitespace in batch labels."
             )
 
+        if method != "longitudinal" and (subject_id is not None or time_covariate is not None):
+            warnings.warn(
+                "`subject_id` and `time_covariate` are only used when "
+                "method='longitudinal'; they are ignored here.",
+                stacklevel=2,
+            )
+
         if method == "johnson":
             if disc is not None or cont is not None:
                 warnings.warn("Covariates are ignored when using method='johnson'.", stacklevel=2)
@@ -188,6 +210,8 @@ class ComBatModel:
             self._fit_fortin(X, batch, disc, cont)
         elif method == "chen":
             self._fit_chen(X, batch, disc, cont)
+        elif method == "longitudinal":
+            self._fit_longitudinal(X, batch, disc, cont, subject_id, time_covariate)
         return self
 
     def _fit_johnson(self, X: pd.DataFrame, batch: pd.Series) -> None:
@@ -431,6 +455,153 @@ class ComBatModel:
         self._pc_gamma_star = gamma_star
         self._pc_delta_star = delta_star
 
+    def _build_longitudinal_design(
+        self,
+        X: pd.DataFrame,
+        batch: pd.Series,
+        disc: pd.DataFrame | None,
+        cont: pd.DataFrame | None,
+        time: pd.DataFrame | None,
+    ) -> tuple[FloatArray, list[str]]:
+        """Assemble the fixed-effects design [batch dummies | covariates | time].
+
+        Returns the design matrix and the list of non-batch column names.
+        """
+        batch_dummies = pd.get_dummies(batch, drop_first=False).astype(float)
+        batch_dummies = batch_dummies.reindex(columns=self._batch_levels, fill_value=0.0)
+
+        nonbatch_parts: list[pd.DataFrame] = []
+        if disc is not None:
+            nonbatch_parts.append(
+                pd.get_dummies(disc.astype("category"), drop_first=True).astype(float)
+            )
+        if cont is not None:
+            nonbatch_parts.append(cont.astype(float))
+        if time is not None:
+            nonbatch_parts.append(time.astype(float))
+
+        if nonbatch_parts:
+            nonbatch_df = pd.concat(nonbatch_parts, axis=1)
+        else:
+            nonbatch_df = pd.DataFrame(index=batch_dummies.index)
+
+        design_df = pd.concat([batch_dummies, nonbatch_df], axis=1)
+        n_batch = len(self._batch_levels)
+        nonbatch_columns = design_df.columns[n_batch:].tolist()
+        return design_df.values.astype(np.float64), nonbatch_columns
+
+    def _fit_longitudinal(
+        self,
+        X: pd.DataFrame,
+        batch: pd.Series,
+        disc: pd.DataFrame | None,
+        cont: pd.DataFrame | None,
+        subject_id: ArrayLike | None,
+        time_covariate: ArrayLike | None,
+    ) -> None:
+        """Longitudinal ComBat (*Beer et al.* 2020).
+
+        Identical to Fortin's standardize -> per-batch gamma/delta -> empirical
+        Bayes -> remove pipeline, except the fixed-effects mean model is fit with
+        a per-subject random intercept (REML), and the random-intercept BLUP is
+        included in the standardization mean.
+        """
+        if subject_id is None:
+            raise ValueError(
+                "method='longitudinal' requires `subject_id` (one label per sample "
+                "identifying the subject/individual used for the random intercept)."
+            )
+
+        self._batch_levels = batch.cat.categories
+        n_batch = len(self._batch_levels)
+        n_samples = len(X)
+
+        counts = batch.value_counts()
+        small = counts[counts < 2]
+        if len(small):
+            lvl = small.index[0]
+            raise ValueError(
+                f"Batch '{lvl}' has only {int(small.iloc[0])} sample(s). ComBat requires "
+                f">= 2 samples per batch. Consider merging small batches or removing them."
+            )
+
+        time = self._to_df(time_covariate, X.index, "time_covariate")
+        design, self._nonbatch_columns = self._build_longitudinal_design(X, batch, disc, cont, time)
+        self._design_cond = float(la.cond(design)) if design.shape[1] else 1.0
+        self._n_batch = n_batch
+        self._p_design = design.shape[1]
+
+        subj = self._as_series(subject_id, X.index, "subject_id")
+        self._re_subject_levels = subj.cat.categories
+        group_idx = subj.cat.codes.to_numpy().astype(np.intp)
+        n_groups = len(self._re_subject_levels)
+        n_k = np.bincount(group_idx, minlength=n_groups).astype(np.float64)
+
+        beta, blup, sigma2, _ = fit_random_intercept(
+            design, X.values.astype(np.float64), group_idx, n_groups, n_k, eps=self.eps
+        )
+        self._re_blup = blup
+
+        beta_batch = beta[:n_batch]
+        self._beta_hat_nonbatch = beta[n_batch:]
+
+        n_per_batch_arr = np.asarray(
+            batch.value_counts().reindex(self._batch_levels).astype(int).values, dtype=np.float64
+        )
+        self._n_per_batch = dict(zip(self._batch_levels, n_per_batch_arr.astype(int), strict=True))
+
+        grand_mean = (n_per_batch_arr / n_samples) @ beta_batch
+        self._grand_mean = pd.Series(grand_mean, index=X.columns)
+        var_pooled = sigma2 + self.eps
+        self._pooled_var = pd.Series(var_pooled, index=X.columns)
+
+        blup_per_row = blup[group_idx]
+        stand_mean = grand_mean + design[:, n_batch:] @ self._beta_hat_nonbatch + blup_per_row
+        Xs = (X.values - stand_mean) / np.sqrt(var_pooled)
+
+        gamma_hat = np.vstack(
+            [Xs[(batch == lvl).to_numpy()].mean(axis=0) for lvl in self._batch_levels]
+        )
+        delta_hat = np.vstack(
+            [
+                Xs[(batch == lvl).to_numpy()].var(axis=0, ddof=1) + self.eps
+                for lvl in self._batch_levels
+            ]
+        )
+
+        b_names = [str(lvl) for lvl in self._batch_levels]
+        if self.mean_only:
+            gamma_star = self._shrink_gamma(
+                gamma_hat,
+                delta_hat,
+                n_per_batch_arr,
+                parametric=self.parametric,
+                batch_names=b_names,
+            )
+            delta_star = np.ones_like(delta_hat)
+        else:
+            gamma_star, delta_star = self._shrink_gamma_delta(
+                gamma_hat,
+                delta_hat,
+                n_per_batch_arr,
+                parametric=self.parametric,
+                batch_names=b_names,
+            )
+
+        if self.reference_batch is not None:
+            ref_idx = list(self._batch_levels).index(self.reference_batch)
+            gamma_ref = gamma_star[ref_idx]
+            delta_ref = delta_star[ref_idx]
+            gamma_star = gamma_star - gamma_ref
+            if not self.mean_only:
+                delta_star = delta_star / delta_ref
+            self._reference_batch_idx = ref_idx
+        else:
+            self._reference_batch_idx = None
+
+        self._gamma_star = gamma_star
+        self._delta_star = delta_star
+
     def _shrink_gamma_delta(
         self,
         gamma_hat: FloatArray,
@@ -620,6 +791,8 @@ class ComBatModel:
         batch: ArrayLike,
         discrete_covariates: ArrayLike | None = None,
         continuous_covariates: ArrayLike | None = None,
+        subject_id: ArrayLike | None = None,
+        time_covariate: ArrayLike | None = None,
     ) -> pd.DataFrame:
         """Transform the data using fitted ComBat parameters.
 
@@ -630,9 +803,14 @@ class ComBatModel:
         batch : array-like of shape (n_samples,)
             Batch labels for each sample.
         discrete_covariates : array-like or None, default=None
-            Categorical covariates (Fortin/Chen only).
+            Categorical covariates (Fortin/Chen/Longitudinal only).
         continuous_covariates : array-like or None, default=None
-            Continuous covariates (Fortin/Chen only).
+            Continuous covariates (Fortin/Chen/Longitudinal only).
+        subject_id : array-like or None, default=None
+            Subject labels for the random intercept (Longitudinal only). Unseen
+            subjects fall back to a zero random intercept.
+        time_covariate : array-like or None, default=None
+            Continuous time variable (Longitudinal only).
 
         Returns
         -------
@@ -665,9 +843,12 @@ class ComBatModel:
             return self._transform_fortin(X, batch, disc, cont)
         elif method == "chen":
             return self._transform_chen(X, batch, disc, cont)
+        elif method == "longitudinal":
+            return self._transform_longitudinal(X, batch, disc, cont, subject_id, time_covariate)
         else:
             raise ValueError(
-                f"Unknown method={method!r}. Expected one of 'johnson', 'fortin', or 'chen'."
+                f"Unknown method={method!r}. Expected one of 'johnson', 'fortin', "
+                f"'chen', or 'longitudinal'."
             )
 
     def _transform_johnson(self, X: pd.DataFrame, batch: pd.Series) -> pd.DataFrame:
@@ -773,3 +954,55 @@ class ComBatModel:
 
         X_recon = self._covbat_pca.inverse_transform(scores_adj)
         return pd.DataFrame(X_recon, index=X.index, columns=X.columns)
+
+    def _transform_longitudinal(
+        self,
+        X: pd.DataFrame,
+        batch: pd.Series,
+        disc: pd.DataFrame | None,
+        cont: pd.DataFrame | None,
+        subject_id: ArrayLike | None,
+        time_covariate: ArrayLike | None,
+    ) -> pd.DataFrame:
+        """Longitudinal transform implementation."""
+        if subject_id is None:
+            raise ValueError("method='longitudinal' requires `subject_id` at transform time.")
+
+        n_features = len(self._grand_mean)
+        time = self._to_df(time_covariate, X.index, "time_covariate")
+        design, nonbatch_columns = self._build_longitudinal_design(X, batch, disc, cont, time)
+        # Realign non-batch columns to those seen at fit (handles unseen covariate levels).
+        nonbatch_design = pd.DataFrame(
+            design[:, self._n_batch :],
+            index=X.index,
+            columns=nonbatch_columns,
+        ).reindex(columns=self._nonbatch_columns, fill_value=0.0)
+
+        subj = self._as_series(subject_id, X.index, "subject_id")
+        level_pos = {lvl: i for i, lvl in enumerate(self._re_subject_levels)}
+        codes = np.array([level_pos.get(s, -1) for s in subj.to_numpy()], dtype=np.intp)
+        blup_per_row = np.zeros((len(X.index), n_features), dtype=np.float64)
+        seen = codes >= 0
+        if seen.any():
+            blup_per_row[seen] = self._re_blup[codes[seen]]
+
+        grand_vals = np.asarray(self._grand_mean.values, dtype=np.float64)
+        pooled_vals = np.asarray(self._pooled_var.values, dtype=np.float64)
+        stand_mean = grand_vals + nonbatch_design.values @ self._beta_hat_nonbatch + blup_per_row
+        Xs = (X.values - stand_mean) / np.sqrt(pooled_vals)
+
+        for i, lvl in enumerate(self._batch_levels):
+            idx = (batch == lvl).to_numpy()
+            if not idx.any():
+                continue
+            if self.reference_batch is not None and lvl == self.reference_batch:
+                continue
+            g = self._gamma_star[i]
+            d = self._delta_star[i]
+            if self.mean_only:
+                Xs[idx] = Xs[idx] - g
+            else:
+                Xs[idx] = (Xs[idx] - g) / np.sqrt(d)
+
+        X_adj = Xs * np.sqrt(pooled_vals) + stand_mean
+        return pd.DataFrame(X_adj, index=X.index, columns=X.columns, dtype=float)

{combatlearn-2.0.0 → combatlearn-2.1.0}/combatlearn/sklearn_api.py

@@ -27,10 +27,15 @@ class ComBat(BaseEstimator, TransformerMixin):  # type: ignore[misc]
     batch : array-like of shape (n_samples,)
         Batch labels for each sample.
     discrete_covariates : array-like, optional
-        Categorical covariates to protect (Fortin/Chen only).
+        Categorical covariates to protect (Fortin/Chen/Longitudinal only).
     continuous_covariates : array-like, optional
-        Continuous covariates to protect (Fortin/Chen only).
-    method : {'johnson', 'fortin', 'chen'}, default='johnson'
+        Continuous covariates to protect (Fortin/Chen/Longitudinal only).
+    subject_id : array-like, optional
+        Subject/individual labels for the random intercept. Required for
+        ``method='longitudinal'``, ignored otherwise.
+    time_covariate : array-like, optional
+        Continuous time variable for repeated measures (Longitudinal only).
+    method : {'johnson', 'fortin', 'chen', 'longitudinal'}, default='johnson'
         ComBat variant to use.
     parametric : bool, default=True
         Use parametric empirical Bayes.
@@ -53,6 +58,8 @@ class ComBat(BaseEstimator, TransformerMixin):  # type: ignore[misc]
         *,
         discrete_covariates: ArrayLike | None = None,
         continuous_covariates: ArrayLike | None = None,
+        subject_id: ArrayLike | None = None,
+        time_covariate: ArrayLike | None = None,
         method: str = "johnson",
         parametric: bool = True,
         mean_only: bool = False,
@@ -63,6 +70,8 @@ class ComBat(BaseEstimator, TransformerMixin):  # type: ignore[misc]
         self.batch = batch
         self.discrete_covariates = discrete_covariates
         self.continuous_covariates = continuous_covariates
+        self.subject_id = subject_id
+        self.time_covariate = time_covariate
         self.method = method
         self.parametric = parametric
         self.mean_only = mean_only
@@ -85,6 +94,19 @@ class ComBat(BaseEstimator, TransformerMixin):  # type: ignore[misc]
                     f"All batch labels must be non-null. Check your data for missing entries."
                 )
 
+            if self.subject_id is not None:
+                subj_ser = (
+                    pd.Series(self.subject_id)
+                    if not isinstance(self.subject_id, pd.Series)
+                    else self.subject_id
+                )
+                nan_count = int(subj_ser.isna().sum())
+                if nan_count:
+                    raise ValueError(
+                        f"subject_id contains {nan_count} NaN value(s). "
+                        f"All subject labels must be non-null."
+                    )
+
             if self.discrete_covariates is not None:
                 disc_df = (
                     pd.DataFrame(self.discrete_covariates)
@@ -118,6 +140,7 @@ class ComBat(BaseEstimator, TransformerMixin):  # type: ignore[misc]
         batch_ser: pd.Series,
         disc: pd.DataFrame | pd.Series | None,
         cont: pd.DataFrame | pd.Series | None,
+        time: pd.DataFrame | pd.Series | None = None,
     ) -> None:
         """Issue warnings for data quality issues that may affect results."""
         # Near-zero variance features
@@ -148,8 +171,10 @@ class ComBat(BaseEstimator, TransformerMixin):  # type: ignore[misc]
                 stacklevel=3,
             )
 
-        # Covariate collinearity with batch (fortin/chen only)
-        if self.method.lower() in {"fortin", "chen"} and (disc is not None or cont is not None):
+        # Covariate collinearity with batch (fortin/chen/longitudinal only)
+        if self.method.lower() in {"fortin", "chen", "longitudinal"} and (
+            disc is not None or cont is not None or time is not None
+        ):
             batch_dummies = pd.get_dummies(batch_ser, drop_first=True).astype(float)
             cov_parts: list[pd.DataFrame] = []
             if disc is not None:
@@ -160,6 +185,9 @@ class ComBat(BaseEstimator, TransformerMixin):  # type: ignore[misc]
             if cont is not None:
                 c = cont if isinstance(cont, pd.DataFrame) else cont.to_frame()
                 cov_parts.append(c.astype(float))
+            if time is not None:
+                t = time if isinstance(time, pd.DataFrame) else time.to_frame()
+                cov_parts.append(t.astype(float))
             cov_df = pd.concat(cov_parts, axis=1)
             design = pd.concat([batch_dummies, cov_df], axis=1)
             rank = la.matrix_rank(design.values)
@@ -227,14 +255,18 @@ class ComBat(BaseEstimator, TransformerMixin):  # type: ignore[misc]
         batch_vec = _subset(self.batch, idx)
         disc = _subset(self.discrete_covariates, idx)
         cont = _subset(self.continuous_covariates, idx)
+        subj = _subset(self.subject_id, idx)
+        time = _subset(self.time_covariate, idx)
 
-        self._check_data_quality(X_df, batch_vec, disc, cont)  # type: ignore[arg-type]
+        self._check_data_quality(X_df, batch_vec, disc, cont, time)  # type: ignore[arg-type]
 
         self._model.fit(
             X,
             batch=batch_vec,  # type: ignore[arg-type]
             discrete_covariates=disc,
             continuous_covariates=cont,
+            subject_id=subj,
+            time_covariate=time,
         )
         self._fitted_batch = batch_vec
 
@@ -261,11 +293,15 @@ class ComBat(BaseEstimator, TransformerMixin):  # type: ignore[misc]
         batch_vec = _subset(self.batch, idx)
         disc = _subset(self.discrete_covariates, idx)
         cont = _subset(self.continuous_covariates, idx)
+        subj = _subset(self.subject_id, idx)
+        time = _subset(self.time_covariate, idx)
         X_transformed = self._model.transform(
             X,
             batch=batch_vec,  # type: ignore[arg-type]
             discrete_covariates=disc,
             continuous_covariates=cont,
+            subject_id=subj,
+            time_covariate=time,
         )
 
         batch_arr = np.asarray(batch_vec)

{combatlearn-2.0.0 → combatlearn-2.1.0}/combatlearn.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: combatlearn
-Version: 2.0.0
+Version: 2.1.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,11 @@ 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**:
+**Four 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)
 
 ## Installation
 
@@ -75,6 +77,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 +128,20 @@ 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"`, 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. |
+| `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></ul> |
 | `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. |
 | `eps` | float | `1e-8` | Small jitter value added to variances to prevent divide-by-zero errors during standardization. |
 
 
@@ -190,3 +195,5 @@ 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)

{combatlearn-2.0.0 → combatlearn-2.1.0}/combatlearn.egg-info/SOURCES.txt

@@ -5,6 +5,7 @@ requirements-dev.txt
 requirements-docs.txt
 requirements.txt
 combatlearn/__init__.py
+combatlearn/_mixed.py
 combatlearn/_utils.py
 combatlearn/core.py
 combatlearn/inspection.py
@@ -19,5 +20,6 @@ combatlearn.egg-info/requires.txt
 combatlearn.egg-info/top_level.txt
 tests/test_combat.py
 tests/test_inspection.py
+tests/test_longitudinal.py
 tests/test_metrics.py
 tests/test_visualization.py

{combatlearn-2.0.0 → combatlearn-2.1.0}/combatlearn.egg-info/requires.txt

@@ -14,6 +14,7 @@ ruff>=0.8
 pre-commit>=3.0
 mypy>=1.0
 pandas-stubs
+statsmodels>=0.14
 
 [docs]
 sphinx>=7.0.0

{combatlearn-2.0.0 → combatlearn-2.1.0}/requirements-dev.txt

@@ -4,3 +4,4 @@ ruff>=0.8
 pre-commit>=3.0
 mypy>=1.0
 pandas-stubs
+statsmodels>=0.14

combatlearn-2.1.0/tests/test_longitudinal.py

@@ -0,0 +1,204 @@
+import numpy as np
+import pandas as pd
+import pytest
+from sklearn.base import clone
+from sklearn.pipeline import Pipeline
+from sklearn.preprocessing import StandardScaler
+from utils import simulate_longitudinal_data
+
+from combatlearn import ComBat
+from combatlearn._mixed import fit_random_intercept
+from combatlearn.core import ComBatModel
+
+
+def _fit_kwargs():
+    X, batch, subject, time = simulate_longitudinal_data()
+    return X, {
+        "batch": batch,
+        "subject_id": subject,
+        "time_covariate": time,
+        "method": "longitudinal",
+    }
+
+
+def test_longitudinal_shape_index_columns_preserved():
+    """Output keeps shape, index, and column names."""
+    X, kwargs = _fit_kwargs()
+    X.index = pd.Index([f"obs_{i}" for i in range(len(X))])
+    kwargs["batch"].index = X.index
+    kwargs["subject_id"].index = X.index
+    kwargs["time_covariate"].index = X.index
+    X.columns = [f"feat_{i}" for i in range(X.shape[1])]
+
+    X_corr = ComBat(**kwargs).fit_transform(X)
+    assert X_corr.shape == X.shape
+    assert X_corr.index.equals(X.index)
+    assert list(X_corr.columns) == list(X.columns)
+
+
+def test_longitudinal_no_nan_or_inf():
+    """Correction introduces no NaN/Inf."""
+    X, kwargs = _fit_kwargs()
+    X_corr = ComBat(**kwargs).fit_transform(X)
+    assert not np.isnan(X_corr.values).any()
+    assert not np.isinf(X_corr.values).any()
+
+
+@pytest.mark.parametrize("mean_only", [True, False])
+def test_longitudinal_mean_only(mean_only):
+    """mean_only mode works."""
+    X, kwargs = _fit_kwargs()
+    X_corr = ComBat(mean_only=mean_only, **kwargs).fit_transform(X)
+    assert X_corr.shape == X.shape
+    assert not np.isnan(X_corr.values).any()
+
+
+@pytest.mark.parametrize("parametric", [True, False])
+def test_longitudinal_parametric_flag(parametric):
+    """Both EB variants run."""
+    X, kwargs = _fit_kwargs()
+    X_corr = ComBat(parametric=parametric, **kwargs).fit_transform(X)
+    assert X_corr.shape == X.shape
+    assert not np.isnan(X_corr.values).any()
+
+
+def test_longitudinal_reduces_batch_variance():
+    """Harmonization should reduce the fraction of variance explained by batch."""
+    X, kwargs = _fit_kwargs()
+    cb = ComBat(**kwargs).fit(X)
+    cb.transform(X)
+    assert cb._batch_var_after_ < cb._batch_var_before_
+
+
+def test_longitudinal_requires_subject_id():
+    """method='longitudinal' without subject_id raises."""
+    X, batch, _subject, time = simulate_longitudinal_data()
+    with pytest.raises(ValueError, match=r"requires `subject_id`"):
+        ComBat(batch=batch, time_covariate=time, method="longitudinal").fit(X)
+
+
+def test_subject_id_ignored_warning_for_other_methods():
+    """Passing subject_id with a non-longitudinal method warns that it is ignored."""
+    X, batch, subject, _time = simulate_longitudinal_data()
+    with pytest.warns(UserWarning, match="only used when"):
+        ComBat(batch=batch, subject_id=subject, method="johnson").fit(X)
+
+
+def test_longitudinal_reference_batch_unchanged():
+    """Reference-batch samples come out numerically identical."""
+    X, batch, subject, time = simulate_longitudinal_data()
+    ref = batch.iloc[0]
+    cb = ComBat(
+        batch=batch,
+        subject_id=subject,
+        time_covariate=time,
+        method="longitudinal",
+        reference_batch=ref,
+    ).fit(X)
+    X_corr = cb.transform(X)
+    mask = (batch == ref).to_numpy()
+    np.testing.assert_allclose(X_corr.values[mask], X.values[mask], rtol=0, atol=1e-8)
+
+
+def test_longitudinal_clone_and_pipeline():
+    """Wrapper clones and runs inside a Pipeline reproducibly."""
+    X, kwargs = _fit_kwargs()
+    pipe = Pipeline([("combat", ComBat(**kwargs)), ("scaler", StandardScaler())])
+    X_corr = pipe.fit_transform(X)
+    X_corr2 = clone(pipe).fit_transform(X)
+    np.testing.assert_allclose(X_corr, X_corr2, rtol=1e-6, atol=1e-6)
+
+
+def test_longitudinal_unseen_subject_falls_back_to_zero_re():
+    """Transforming subjects unseen at fit (zero random intercept) must not crash."""
+    X, batch, subject, time = simulate_longitudinal_data()
+    model = ComBatModel(method="longitudinal").fit(
+        X, batch=batch, subject_id=subject, time_covariate=time
+    )
+    subject_new = subject.map(lambda s: f"new_{s}")
+    X_corr = model.transform(X, batch=batch, subject_id=subject_new, time_covariate=time)
+    assert X_corr.shape == X.shape
+    assert not np.isnan(X_corr.values).any()
+
+
+def test_longitudinal_no_covariates():
+    """Longitudinal with only a subject random intercept (no covariates) works."""
+    X, batch, subject, _time = simulate_longitudinal_data()
+    X_corr = ComBat(batch=batch, subject_id=subject, method="longitudinal").fit_transform(X)
+    assert X_corr.shape == X.shape
+    assert not np.isnan(X_corr.values).any()
+
+
+def test_longitudinal_collinear_discrete_covariate_warns():
+    """A discrete covariate collinear with batch should warn for the longitudinal method too."""
+    X, batch, subject, _time = simulate_longitudinal_data()
+    disc_collinear = batch.to_frame("cov")
+    with pytest.warns(UserWarning, match="rank-deficient"):
+        ComBat(
+            batch=batch,
+            subject_id=subject,
+            discrete_covariates=disc_collinear,
+            method="longitudinal",
+        ).fit(X)
+
+
+def test_longitudinal_collinear_time_covariate_warns():
+    """A time covariate collinear with batch should also trigger the rank-deficiency warning."""
+    X, batch, subject, _time = simulate_longitudinal_data()
+    time_collinear = batch.map({"A": 0.0, "B": 1.0, "C": 2.0}).rename("visit")
+    with pytest.warns(UserWarning, match="rank-deficient"):
+        ComBat(
+            batch=batch,
+            subject_id=subject,
+            time_covariate=time_collinear,
+            method="longitudinal",
+        ).fit(X)
+
+
+def test_random_intercept_recovers_fixed_effects():
+    """The REML fitter recovers known fixed effects on simulated data."""
+    rng = np.random.default_rng(0)
+    n_subjects, n_times = 60, 4
+    n = n_subjects * n_times
+    group_idx = np.repeat(np.arange(n_subjects), n_times).astype(np.intp)
+    n_k = np.bincount(group_idx).astype(np.float64)
+
+    x1 = rng.standard_normal(n)
+    beta_true = np.array([2.0, 1.5])
+    u = rng.normal(0.0, 1.0, n_subjects)
+    eps = rng.normal(0.0, 1.0, n)
+    y = beta_true[0] + beta_true[1] * x1 + u[group_idx] + eps
+
+    X = np.column_stack([np.ones(n), x1])
+    beta, _blup, sigma2, lam_opt = fit_random_intercept(X, y[:, None], group_idx, n_subjects, n_k)
+
+    np.testing.assert_allclose(beta[:, 0], beta_true, rtol=0, atol=0.25)
+    assert 0.5 < sigma2[0] < 2.0
+    assert lam_opt[0] > 0.0  # a real random intercept is detected
+
+
+def test_random_intercept_matches_statsmodels():
+    """Cross-check REML estimates against statsmodels.MixedLM."""
+    sm = pytest.importorskip("statsmodels.api")
+
+    rng = np.random.default_rng(1)
+    n_subjects, n_times = 60, 5
+    n = n_subjects * n_times
+    group_idx = np.repeat(np.arange(n_subjects), n_times).astype(np.intp)
+    n_k = np.bincount(group_idx).astype(np.float64)
+
+    x1 = rng.standard_normal(n)
+    u = rng.normal(0.0, 1.2, n_subjects)
+    eps = rng.normal(0.0, 0.9, n)
+    y = 1.0 + 0.8 * x1 + u[group_idx] + eps
+
+    X = np.column_stack([np.ones(n), x1])
+    beta, _blup, sigma2, lam_opt = fit_random_intercept(X, y[:, None], group_idx, n_subjects, n_k)
+    tau2 = lam_opt[0] * sigma2[0]
+
+    res = sm.MixedLM(endog=y, exog=X, groups=group_idx).fit(reml=True, method="lbfgs")
+    sm_tau2 = float(np.asarray(res.cov_re)[0, 0])
+
+    np.testing.assert_allclose(beta[:, 0], res.fe_params, rtol=0.02, atol=0.02)
+    np.testing.assert_allclose(sigma2[0], res.scale, rtol=0.05)
+    np.testing.assert_allclose(tau2, sm_tau2, rtol=0.12, atol=0.05)