combatlearn 2.1.0__tar.gz → 2.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (29) hide show
  1. {combatlearn-2.1.0 → combatlearn-2.2.0}/PKG-INFO +13 -5
  2. {combatlearn-2.1.0 → combatlearn-2.2.0}/README.md +13 -5
  3. {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn/__init__.py +1 -1
  4. {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn/core.py +291 -19
  5. {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn/sklearn_api.py +38 -7
  6. {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn/visualization.py +105 -49
  7. {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn.egg-info/PKG-INFO +13 -5
  8. {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn.egg-info/SOURCES.txt +2 -0
  9. {combatlearn-2.1.0 → combatlearn-2.2.0}/tests/test_combat.py +46 -0
  10. combatlearn-2.2.0/tests/test_eb_degenerate.py +111 -0
  11. combatlearn-2.2.0/tests/test_gam.py +298 -0
  12. {combatlearn-2.1.0 → combatlearn-2.2.0}/tests/test_longitudinal.py +12 -0
  13. {combatlearn-2.1.0 → combatlearn-2.2.0}/tests/test_visualization.py +35 -0
  14. {combatlearn-2.1.0 → combatlearn-2.2.0}/LICENSE +0 -0
  15. {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn/_mixed.py +0 -0
  16. {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn/_utils.py +0 -0
  17. {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn/inspection.py +0 -0
  18. {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn/metrics.py +0 -0
  19. {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn/py.typed +0 -0
  20. {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn.egg-info/dependency_links.txt +0 -0
  21. {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn.egg-info/requires.txt +0 -0
  22. {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn.egg-info/top_level.txt +0 -0
  23. {combatlearn-2.1.0 → combatlearn-2.2.0}/pyproject.toml +0 -0
  24. {combatlearn-2.1.0 → combatlearn-2.2.0}/requirements-dev.txt +0 -0
  25. {combatlearn-2.1.0 → combatlearn-2.2.0}/requirements-docs.txt +0 -0
  26. {combatlearn-2.1.0 → combatlearn-2.2.0}/requirements.txt +0 -0
  27. {combatlearn-2.1.0 → combatlearn-2.2.0}/setup.cfg +0 -0
  28. {combatlearn-2.1.0 → combatlearn-2.2.0}/tests/test_inspection.py +0 -0
  29. {combatlearn-2.1.0 → combatlearn-2.2.0}/tests/test_metrics.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: combatlearn
3
- Version: 2.1.0
3
+ Version: 2.2.0
4
4
  Summary: Batch-effect harmonization for machine learning frameworks.
5
5
  Author-email: Ettore Rocchi <ettoreroc@gmail.com>
6
6
  License: MIT
@@ -50,11 +50,13 @@ Dynamic: license-file
50
50
 
51
51
  **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**.
52
52
 
53
- **Four methods**:
53
+ **Six methods**:
54
54
  - `method="johnson"` - classic ComBat (Johnson _et al._, 2007)
55
55
  - `method="fortin"` - neuroComBat (Fortin _et al._, 2018)
56
56
  - `method="chen"` - CovBat (Chen _et al._, 2022)
57
57
  - `method="longitudinal"` - Longitudinal ComBat for repeated measures (Beer _et al._, 2020)
58
+ - `method="gam"` - ComBat-GAM, nonlinear (spline) covariate effects (Pomponio _et al._, 2020)
59
+ - `method="covbat_gam"` - CovBat with the same nonlinear covariate modeling
58
60
 
59
61
  ## Installation
60
62
 
@@ -128,8 +130,8 @@ The following section provides a detailed explanation of all parameters availabl
128
130
  | Parameter | Type | Default | Description |
129
131
  | --- | --- | --- | --- |
130
132
  | `batch` | array-like or pd.Series | **required** | Vector indicating batch assignment for each sample. This is used to estimate and remove batch effects. |
131
- | `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. |
132
- | `continuous_covariates` | array-like, pd.Series or pd.DataFrame | `None` | Optional continuous covariates (e.g., age). Only used in `"fortin"`, `"chen"`, and `"longitudinal"` methods. |
133
+ | `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. |
134
+ | `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. |
133
135
  | `subject_id` | array-like or pd.Series | `None` | Subject/individual labels for the random intercept. **Required** for `method="longitudinal"`, ignored otherwise. |
134
136
  | `time_covariate` | array-like or pd.Series | `None` | Optional continuous time variable for repeated measures. Only used in `"longitudinal"`. |
135
137
 
@@ -137,11 +139,15 @@ The following section provides a detailed explanation of all parameters availabl
137
139
 
138
140
  | Parameter | Type | Default | Description |
139
141
  | --- | --- | --- | --- |
140
- | `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> |
142
+ | `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"`). |
141
143
  | `parametric` | bool | `True` | Whether to use the **parametric empirical Bayes** formulation. If `False`, a non-parametric iterative scheme is used. |
142
144
  | `mean_only` | bool | `False` | If `True`, only the **mean** is corrected, while variances are left unchanged. Useful for preserving variance structure in the data. |
143
145
  | `reference_batch` | str or `None` | `None` | If specified, acts as a reference batch - other batches will be corrected to match this one. |
144
146
  | `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. |
147
+ | `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. |
148
+ | `spline_df` | int | `10` | For the GAM methods: B-spline degrees of freedom (basis functions) per smooth term. |
149
+ | `spline_degree` | int | `3` | For the GAM methods: B-spline degree (3 = cubic). |
150
+ | `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. |
145
151
  | `eps` | float | `1e-8` | Small jitter value added to variances to prevent divide-by-zero errors during standardization. |
146
152
 
147
153
 
@@ -197,3 +203,5 @@ Please consider citing the original papers:
197
203
  - [**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)
198
204
 
199
205
  - [**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)
206
+
207
+ - [**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)
@@ -12,11 +12,13 @@
12
12
 
13
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**.
14
14
 
15
- **Four methods**:
15
+ **Six methods**:
16
16
  - `method="johnson"` - classic ComBat (Johnson _et al._, 2007)
17
17
  - `method="fortin"` - neuroComBat (Fortin _et al._, 2018)
18
18
  - `method="chen"` - CovBat (Chen _et al._, 2022)
19
19
  - `method="longitudinal"` - Longitudinal ComBat for repeated measures (Beer _et al._, 2020)
20
+ - `method="gam"` - ComBat-GAM, nonlinear (spline) covariate effects (Pomponio _et al._, 2020)
21
+ - `method="covbat_gam"` - CovBat with the same nonlinear covariate modeling
20
22
 
21
23
  ## Installation
22
24
 
@@ -90,8 +92,8 @@ The following section provides a detailed explanation of all parameters availabl
90
92
  | Parameter | Type | Default | Description |
91
93
  | --- | --- | --- | --- |
92
94
  | `batch` | array-like or pd.Series | **required** | Vector indicating batch assignment for each sample. This is used to estimate and remove batch effects. |
93
- | `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. |
94
- | `continuous_covariates` | array-like, pd.Series or pd.DataFrame | `None` | Optional continuous covariates (e.g., age). Only used in `"fortin"`, `"chen"`, and `"longitudinal"` methods. |
95
+ | `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. |
96
+ | `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. |
95
97
  | `subject_id` | array-like or pd.Series | `None` | Subject/individual labels for the random intercept. **Required** for `method="longitudinal"`, ignored otherwise. |
96
98
  | `time_covariate` | array-like or pd.Series | `None` | Optional continuous time variable for repeated measures. Only used in `"longitudinal"`. |
97
99
 
@@ -99,11 +101,15 @@ The following section provides a detailed explanation of all parameters availabl
99
101
 
100
102
  | Parameter | Type | Default | Description |
101
103
  | --- | --- | --- | --- |
102
- | `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> |
104
+ | `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"`). |
103
105
  | `parametric` | bool | `True` | Whether to use the **parametric empirical Bayes** formulation. If `False`, a non-parametric iterative scheme is used. |
104
106
  | `mean_only` | bool | `False` | If `True`, only the **mean** is corrected, while variances are left unchanged. Useful for preserving variance structure in the data. |
105
107
  | `reference_batch` | str or `None` | `None` | If specified, acts as a reference batch - other batches will be corrected to match this one. |
106
108
  | `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. |
109
+ | `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. |
110
+ | `spline_df` | int | `10` | For the GAM methods: B-spline degrees of freedom (basis functions) per smooth term. |
111
+ | `spline_degree` | int | `3` | For the GAM methods: B-spline degree (3 = cubic). |
112
+ | `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. |
107
113
  | `eps` | float | `1e-8` | Small jitter value added to variances to prevent divide-by-zero errors during standardization. |
108
114
 
109
115
 
@@ -158,4 +164,6 @@ Please consider citing the original papers:
158
164
 
159
165
  - [**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)
160
166
 
161
- - [**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)
167
+ - [**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)
168
+
169
+ - [**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)
@@ -1,5 +1,5 @@
1
1
  from .sklearn_api import ComBat
2
2
 
3
3
  __all__ = ["ComBat"]
4
- __version__ = "2.1.0"
4
+ __version__ = "2.2.0"
5
5
  __author__ = "Ettore Rocchi"
@@ -1,10 +1,12 @@
1
1
  """ComBat algorithm core.
2
2
 
3
- `ComBatModel` implements four variants of the ComBat algorithm:
3
+ `ComBatModel` implements several variants of the ComBat algorithm:
4
4
  * Johnson et al. (2007) vanilla ComBat (method="johnson")
5
5
  * Fortin et al. (2018) extension with covariates (method="fortin")
6
6
  * Chen et al. (2022) CovBat (method="chen")
7
7
  * Beer et al. (2020) Longitudinal ComBat (method="longitudinal")
8
+ * Pomponio et al. (2020) ComBat-GAM, spline covariates (method="gam")
9
+ * CovBat with GAM covariates (method="covbat_gam")
8
10
  """
9
11
 
10
12
  from __future__ import annotations
@@ -16,6 +18,7 @@ import numpy as np
16
18
  import numpy.linalg as la
17
19
  import numpy.typing as npt
18
20
  import pandas as pd
21
+ from scipy.interpolate import BSpline
19
22
  from sklearn.decomposition import PCA
20
23
 
21
24
  from ._mixed import fit_random_intercept
@@ -23,18 +26,54 @@ from ._mixed import fit_random_intercept
23
26
  ArrayLike: TypeAlias = pd.DataFrame | pd.Series | npt.NDArray[Any]
24
27
  FloatArray: TypeAlias = npt.NDArray[np.float64]
25
28
 
29
+ _METHOD_ALIASES: dict[str, str] = {
30
+ "classiccombat": "johnson",
31
+ "neurocombat": "fortin",
32
+ "covbat": "chen",
33
+ "longcombat": "longitudinal",
34
+ "combatgam": "gam",
35
+ "covbatgam": "covbat_gam",
36
+ }
37
+
38
+ _GAM_METHODS: frozenset[str] = frozenset({"gam", "covbat_gam"})
39
+
40
+
41
+ def _resolve_method(method: str) -> str:
42
+ """Resolve a method name or literature alias to its canonical (author) name.
43
+
44
+ Author keys pass through unchanged; method-name aliases map to them, e.g.
45
+ ``"classic_combat"`` -> ``"johnson"``, ``"neurocombat"`` -> ``"fortin"``,
46
+ ``"covbat"`` -> ``"chen"``, ``"combat_gam"`` -> ``"gam"``. Case- and
47
+ separator-insensitive (hyphens, underscores, and spaces are ignored for
48
+ alias matching), so ``"covbat_gam"``, ``"covbat-gam"`` and ``"CovBatGAM"``
49
+ all resolve to the canonical ``"covbat_gam"``.
50
+ """
51
+ raw = method.lower()
52
+ stripped = raw.replace("-", "").replace("_", "").replace(" ", "")
53
+ return _METHOD_ALIASES.get(stripped, raw)
54
+
26
55
 
27
56
  class ComBatModel:
28
57
  """ComBat algorithm.
29
58
 
30
59
  Parameters
31
60
  ----------
32
- method : {'johnson', 'fortin', 'chen', 'longitudinal'}, default='johnson'
61
+ method : {'johnson', 'fortin', 'chen', 'longitudinal', 'gam', \
62
+ 'covbat_gam'}, default='johnson'
33
63
  * 'johnson' - classic ComBat.
34
64
  * 'fortin' - covariate-aware ComBat.
35
65
  * 'chen' - CovBat, PCA-based ComBat.
36
66
  * 'longitudinal' - Longitudinal ComBat; requires a per-sample
37
67
  ``subject_id`` for the random intercept.
68
+ * 'gam' - ComBat-GAM (Pomponio et al. 2020); like 'fortin' but the
69
+ continuous covariates in ``smooth_terms`` are modeled with B-spline
70
+ bases instead of linearly.
71
+ * 'covbat_gam' - CovBat with the same spline covariate modeling.
72
+
73
+ Literature aliases are also accepted (case- and separator-insensitive):
74
+ ``'classic_combat'`` (johnson), ``'neurocombat'`` (fortin),
75
+ ``'covbat'`` (chen), ``'longcombat'`` (longitudinal),
76
+ ``'combat_gam'`` (gam).
38
77
  parametric : bool, default=True
39
78
  Use the parametric empirical Bayes variant.
40
79
  mean_only : bool, default=False
@@ -45,6 +84,22 @@ class ComBatModel:
45
84
  covbat_cov_thresh : float or int, default=0.9
46
85
  CovBat: cumulative variance threshold (0, 1] to retain PCs, or
47
86
  integer >= 1 specifying the number of components directly.
87
+ smooth_terms : list of str or int, optional
88
+ Continuous-covariate columns to model nonlinearly with B-splines
89
+ (``'gam'``/``'covbat_gam'`` only). Column names, or integer positions
90
+ into ``continuous_covariates`` when it is an unnamed array. Default
91
+ (``None``) smooths every continuous covariate.
92
+ spline_df : int, default=10
93
+ Degrees of freedom (number of B-spline basis functions) per smooth
94
+ term, before the identifiability column is dropped. Must be
95
+ ``>= spline_degree + 1``.
96
+ spline_degree : int, default=3
97
+ B-spline degree (3 = cubic). Must be ``>= 1``.
98
+ smooth_term_bounds : tuple of (float, float) or dict, optional
99
+ Boundary knots for the spline basis. A single ``(lo, hi)`` tuple is
100
+ applied to every smooth term; a ``{term: (lo, hi)}`` dict sets them per
101
+ term. Bounds should contain the full data range (training and any
102
+ held-out data). Default uses each term's training min/max.
48
103
  eps : float, default=1e-8
49
104
  Numerical jitter to avoid division-by-zero.
50
105
  """
@@ -52,12 +107,18 @@ class ComBatModel:
52
107
  def __init__(
53
108
  self,
54
109
  *,
55
- method: Literal["johnson", "fortin", "chen", "longitudinal"] = "johnson",
110
+ method: Literal[
111
+ "johnson", "fortin", "chen", "longitudinal", "gam", "covbat_gam"
112
+ ] = "johnson",
56
113
  parametric: bool = True,
57
114
  mean_only: bool = False,
58
115
  reference_batch: str | None = None,
59
116
  eps: float = 1e-8,
60
117
  covbat_cov_thresh: float | int = 0.9,
118
+ smooth_terms: list[str | int] | None = None,
119
+ spline_df: int = 10,
120
+ spline_degree: int = 3,
121
+ smooth_term_bounds: tuple[float, float] | dict[Any, tuple[float, float]] | None = None,
61
122
  ) -> None:
62
123
  self.method: str = method
63
124
  self.parametric: bool = parametric
@@ -65,6 +126,12 @@ class ComBatModel:
65
126
  self.reference_batch: str | None = reference_batch
66
127
  self.eps: float = float(eps)
67
128
  self.covbat_cov_thresh: float | int = covbat_cov_thresh
129
+ self.smooth_terms: list[str | int] | None = smooth_terms
130
+ self.spline_df: int = int(spline_df)
131
+ self.spline_degree: int = int(spline_degree)
132
+ self.smooth_term_bounds: tuple[float, float] | dict[Any, tuple[float, float]] | None = (
133
+ smooth_term_bounds
134
+ )
68
135
 
69
136
  self._batch_levels: pd.Index
70
137
  self._grand_mean: pd.Series
@@ -84,6 +151,22 @@ class ComBatModel:
84
151
  self._pc_delta_star: FloatArray
85
152
  self._re_subject_levels: pd.Index
86
153
  self._re_blup: FloatArray
154
+ self._smooth_cols: list[Any]
155
+ self._smooth_knots: dict[Any, FloatArray]
156
+ self._smooth_basis_columns: dict[Any, list[str]]
157
+
158
+ # Validate spline parameters (used by the gam / covbat_gam methods)
159
+ if self.spline_degree < 1:
160
+ raise ValueError(
161
+ f"spline_degree={self.spline_degree!r} is invalid. Expected an int >= 1 "
162
+ f"(3 = cubic, the ComBat-GAM default)."
163
+ )
164
+ if self.spline_df < self.spline_degree + 1:
165
+ raise ValueError(
166
+ f"spline_df={self.spline_df!r} is too small for spline_degree="
167
+ f"{self.spline_degree!r}. Expected spline_df >= spline_degree + 1 "
168
+ f"(>= {self.spline_degree + 1}); the default is 10."
169
+ )
87
170
 
88
171
  # Validate covbat_cov_thresh
89
172
  if isinstance(self.covbat_cov_thresh, float):
@@ -174,14 +257,22 @@ class ComBatModel:
174
257
  self : ComBatModel
175
258
  Fitted model.
176
259
  """
177
- method = self.method.lower()
178
- if method not in {"johnson", "fortin", "chen", "longitudinal"}:
260
+ method = _resolve_method(self.method)
261
+ if method not in {"johnson", "fortin", "chen", "longitudinal", "gam", "covbat_gam"}:
179
262
  raise ValueError(
180
- f"method={self.method!r} is not recognized. "
181
- f"Expected one of 'johnson', 'fortin', 'chen', or 'longitudinal'."
263
+ f"method={self.method!r} is not recognized. Expected an author name "
264
+ f"('johnson', 'fortin', 'chen', 'longitudinal', 'gam', 'covbat_gam') or a "
265
+ f"method-name alias ('classic_combat', 'neurocombat', 'covbat', "
266
+ f"'longcombat', 'combat_gam')."
182
267
  )
183
268
  if not isinstance(X, pd.DataFrame):
184
269
  X = pd.DataFrame(X)
270
+ if X.shape[1] < 2:
271
+ warnings.warn(
272
+ "Only one feature: empirical-Bayes shrinkage is inactive and ComBat "
273
+ "reduces to per-batch standardization.",
274
+ stacklevel=2,
275
+ )
185
276
  idx = X.index
186
277
  batch = self._as_series(batch, idx, "batch")
187
278
 
@@ -206,9 +297,9 @@ class ComBatModel:
206
297
  if disc is not None or cont is not None:
207
298
  warnings.warn("Covariates are ignored when using method='johnson'.", stacklevel=2)
208
299
  self._fit_johnson(X, batch)
209
- elif method == "fortin":
300
+ elif method in {"fortin", "gam"}:
210
301
  self._fit_fortin(X, batch, disc, cont)
211
- elif method == "chen":
302
+ elif method in {"chen", "covbat_gam"}:
212
303
  self._fit_chen(X, batch, disc, cont)
213
304
  elif method == "longitudinal":
214
305
  self._fit_longitudinal(X, batch, disc, cont, subject_id, time_covariate)
@@ -286,6 +377,140 @@ class ComBatModel:
286
377
  self._n_per_batch = n_per_batch
287
378
  self._design_cond: float | None = None
288
379
 
380
+ def _is_gam(self) -> bool:
381
+ """Whether the resolved method models smooth terms with B-splines."""
382
+ return _resolve_method(self.method) in _GAM_METHODS
383
+
384
+ def _resolve_smooth_columns(self, cont: pd.DataFrame | None) -> list[Any]:
385
+ """Resolve ``smooth_terms`` to concrete column labels in ``cont``.
386
+
387
+ Raises if no continuous covariates are available (a GAM method needs at
388
+ least one smooth term) or if a requested term cannot be located.
389
+ """
390
+ if cont is None or cont.shape[1] == 0:
391
+ raise ValueError(
392
+ f"method={self.method!r} (ComBat-GAM) requires at least one continuous "
393
+ f"covariate to model with splines, but `continuous_covariates` is empty. "
394
+ f"Pass continuous covariates, or use method='fortin'/'chen' for a linear fit."
395
+ )
396
+ cols = list(cont.columns)
397
+ if self.smooth_terms is None:
398
+ return cols
399
+ if len(self.smooth_terms) == 0:
400
+ raise ValueError(
401
+ "smooth_terms is an empty list, so no covariate would be modeled with "
402
+ "splines. Pass at least one term, set smooth_terms=None to smooth all "
403
+ "continuous covariates, or use method='fortin'/'chen'."
404
+ )
405
+ resolved: list[Any] = []
406
+ for t in self.smooth_terms:
407
+ if t in cont.columns:
408
+ resolved.append(t)
409
+ elif isinstance(t, int | np.integer) and 0 <= int(t) < len(cols):
410
+ resolved.append(cols[int(t)])
411
+ else:
412
+ raise ValueError(
413
+ f"smooth_term {t!r} not found in continuous_covariates "
414
+ f"(columns: {cols}). Use a column name or a valid integer position."
415
+ )
416
+ return resolved
417
+
418
+ def _term_bounds(self, col: Any, x: FloatArray) -> tuple[float, float]:
419
+ """Boundary knots for a smooth term, from ``smooth_term_bounds`` or data."""
420
+ b = self.smooth_term_bounds
421
+ if b is None:
422
+ return float(np.min(x)), float(np.max(x))
423
+ if isinstance(b, dict):
424
+ if col in b:
425
+ lo, hi = b[col]
426
+ return float(lo), float(hi)
427
+ return float(np.min(x)), float(np.max(x))
428
+ lo, hi = b
429
+ return float(lo), float(hi)
430
+
431
+ def _spline_knots(self, x: FloatArray, lo: float, hi: float) -> FloatArray:
432
+ """Knot vector: ``degree+1`` boundary knots plus interior sample quantiles.
433
+
434
+ Matches ``statsmodels.gam.smooth_basis.BSplines`` (and hence neuroHarmonize):
435
+ ``spline_df - spline_degree - 1`` interior knots at equally spaced quantiles.
436
+ """
437
+ degree = self.spline_degree
438
+ n_interior = self.spline_df - degree - 1
439
+ if n_interior > 0:
440
+ probs = np.linspace(0.0, 1.0, n_interior + 2)[1:-1]
441
+ interior = np.quantile(x, probs)
442
+ else:
443
+ interior = np.empty(0, dtype=np.float64)
444
+ knots = np.concatenate([np.full(degree + 1, lo), interior, np.full(degree + 1, hi)])
445
+ return np.asarray(knots, dtype=np.float64)
446
+
447
+ def _spline_basis(self, x: FloatArray, knots: FloatArray) -> FloatArray:
448
+ """B-spline design matrix with the first (identifiability) column dropped.
449
+
450
+ Out-of-range values are clamped to the boundary knots (constant
451
+ extrapolation), so held-out data never falls outside the basis support.
452
+ """
453
+ xc = np.clip(np.asarray(x, dtype=np.float64), knots[0], knots[-1])
454
+ basis = BSpline.design_matrix(xc, knots, self.spline_degree).toarray()
455
+ return np.asarray(basis[:, 1:], dtype=np.float64)
456
+
457
+ def _apply_smooth_terms(self, cont: pd.DataFrame | None, *, fit: bool) -> pd.DataFrame:
458
+ """Replace each smooth continuous column with its B-spline basis columns.
459
+
460
+ Non-smooth continuous columns are kept (linear). On ``fit`` the knots are
461
+ computed from the training data and stored; on transform the stored knots
462
+ are reused so the basis is reproducible and leakage-free.
463
+ """
464
+ if cont is None or cont.shape[1] == 0:
465
+ raise ValueError(
466
+ f"method={self.method!r} (ComBat-GAM) requires at least one continuous "
467
+ f"covariate to model with splines, but `continuous_covariates` is empty. "
468
+ f"Pass continuous covariates, or use method='fortin'/'chen' for a linear fit."
469
+ )
470
+ if fit:
471
+ self._smooth_cols = self._resolve_smooth_columns(cont)
472
+ self._smooth_knots = {}
473
+ self._smooth_basis_columns = {}
474
+ smooth_cols = set(self._smooth_cols)
475
+
476
+ blocks: list[pd.DataFrame] = []
477
+ for col in cont.columns:
478
+ if col not in smooth_cols:
479
+ blocks.append(cont[[col]].astype(float))
480
+ continue
481
+ x = cont[col].to_numpy(dtype=np.float64)
482
+ if fit:
483
+ if np.unique(x).size < self.spline_df:
484
+ raise ValueError(
485
+ f"Smooth term {col!r} has only {int(np.unique(x).size)} distinct "
486
+ f"value(s), too few for spline_df={self.spline_df} (need >= "
487
+ f"{self.spline_df}). Lower spline_df, drop it from smooth_terms, "
488
+ f"or model it linearly with method='fortin'/'chen'."
489
+ )
490
+ lo, hi = self._term_bounds(col, x)
491
+ if lo >= hi:
492
+ raise ValueError(
493
+ f"smooth_term_bounds for {col!r} must satisfy lo < hi, got ({lo}, {hi})."
494
+ )
495
+ if x.min() < lo or x.max() > hi:
496
+ raise ValueError(
497
+ f"smooth_term_bounds ({lo}, {hi}) for {col!r} do not contain the "
498
+ f"training range [{x.min():.4g}, {x.max():.4g}]. Bounds must cover the "
499
+ f"full data range (set them wider to include held-out data); held-out "
500
+ f"values beyond the bounds are clamped at transform."
501
+ )
502
+ knots = self._spline_knots(x, lo, hi)
503
+ basis = self._spline_basis(x, knots)
504
+ names = [f"{col}__spline{j}" for j in range(basis.shape[1])]
505
+ self._smooth_knots[col] = knots
506
+ self._smooth_basis_columns[col] = names
507
+ else:
508
+ knots = self._smooth_knots[col]
509
+ basis = self._spline_basis(x, knots)
510
+ names = self._smooth_basis_columns[col]
511
+ blocks.append(pd.DataFrame(basis, index=cont.index, columns=names))
512
+ return pd.concat(blocks, axis=1)
513
+
289
514
  def _fit_fortin(
290
515
  self,
291
516
  X: pd.DataFrame,
@@ -293,7 +518,15 @@ class ComBatModel:
293
518
  disc: pd.DataFrame | None,
294
519
  cont: pd.DataFrame | None,
295
520
  ) -> None:
296
- """Fortin et al. (2018) neuroComBat."""
521
+ """Fortin et al. (2018) neuroComBat (and Pomponio et al. 2020 ComBat-GAM).
522
+
523
+ For the ``gam``/``covbat_gam`` methods the smooth continuous covariates are
524
+ replaced by B-spline bases before the (otherwise identical) least-squares ->
525
+ standardize -> empirical-Bayes pipeline.
526
+ """
527
+ if self._is_gam():
528
+ cont = self._apply_smooth_terms(cont, fit=True)
529
+
297
530
  self._batch_levels = batch.cat.categories
298
531
  n_batch = len(self._batch_levels)
299
532
  n_samples = len(X)
@@ -410,6 +643,15 @@ class ComBatModel:
410
643
  self._covbat_pca = pca
411
644
  self._covbat_n_pc = n_pc
412
645
 
646
+ if n_pc < 2:
647
+ warnings.warn(
648
+ f"CovBat retained a single principal component (n_pc={n_pc}); covariance "
649
+ f"harmonization is not meaningful with one component, and the correction "
650
+ f"reduces to the mean/variance adjustment already done by Fortin. Consider "
651
+ f"method='fortin', or raise covbat_cov_thresh to retain more components.",
652
+ stacklevel=2,
653
+ )
654
+
413
655
  scores = pca.transform(X_meanvar_adj)[:, :n_pc]
414
656
  scores_df = pd.DataFrame(scores, index=X.index, columns=[f"PC{i + 1}" for i in range(n_pc)])
415
657
  self._batch_levels_pc = self._batch_levels
@@ -618,14 +860,34 @@ class ComBatModel:
618
860
  Both the parametric and non-parametric branches use an iterative
619
861
  scheme that alternates between updating gamma (location) and delta
620
862
  (scale) posteriors until convergence.
863
+
864
+ With fewer than two units (features, or PCs for CovBat) the prior
865
+ variance is non-identifiable, so the posterior reduces to the per-unit
866
+ MLE (no shrinkage) - the exact boundary limit of the estimator.
621
867
  """
622
868
  B, _p = gamma_hat.shape
623
869
  n_vec = (
624
870
  np.array(list(n_per_batch.values())) if isinstance(n_per_batch, dict) else n_per_batch
625
871
  )
872
+ convergence_info: list[dict[str, object]] = []
873
+
874
+ if _p < 2:
875
+ # < 2 units: no cross-unit prior, so reduce to the per-unit MLE.
876
+ self._convergence_info = [
877
+ {
878
+ "batch": batch_names[i] if batch_names else str(i),
879
+ "converged": True,
880
+ "iterations": 0,
881
+ "final_gamma_change": 0.0,
882
+ "final_delta_change": 0.0,
883
+ }
884
+ for i in range(B)
885
+ ]
886
+ delta_star = np.ones_like(delta_hat) if self.mean_only else delta_hat.copy()
887
+ return gamma_hat.copy(), delta_star
888
+
626
889
  gamma_bar = gamma_hat.mean(axis=0)
627
890
  t2 = np.maximum(gamma_hat.var(axis=0, ddof=1), self.eps)
628
- convergence_info: list[dict[str, object]] = []
629
891
 
630
892
  def postmean(
631
893
  g_hat: FloatArray,
@@ -702,14 +964,19 @@ class ComBatModel:
702
964
 
703
965
  else:
704
966
 
967
+ def _moment_var(delta: FloatArray) -> float:
968
+ # Floor non-finite / tiny prior variance (defense-in-depth).
969
+ s2 = float(delta.var(ddof=1))
970
+ return self.eps if (not np.isfinite(s2) or s2 < self.eps) else s2
971
+
705
972
  def aprior(delta: FloatArray) -> float:
706
973
  m = float(delta.mean())
707
- s2 = float(max(delta.var(ddof=1), self.eps))
974
+ s2 = _moment_var(delta)
708
975
  return (2 * s2 + m**2) / s2
709
976
 
710
977
  def bprior(delta: FloatArray) -> float:
711
978
  m = float(delta.mean())
712
- s2 = float(max(delta.var(ddof=1), self.eps))
979
+ s2 = _moment_var(delta)
713
980
  return (m * s2 + m**3) / s2
714
981
 
715
982
  gamma_star = np.empty_like(gamma_hat)
@@ -836,19 +1103,21 @@ class ComBatModel:
836
1103
  disc = self._to_df(discrete_covariates, idx, "discrete_covariates")
837
1104
  cont = self._to_df(continuous_covariates, idx, "continuous_covariates")
838
1105
 
839
- method = self.method.lower()
1106
+ method = _resolve_method(self.method)
840
1107
  if method == "johnson":
841
1108
  return self._transform_johnson(X, batch)
842
- elif method == "fortin":
1109
+ elif method in {"fortin", "gam"}:
843
1110
  return self._transform_fortin(X, batch, disc, cont)
844
- elif method == "chen":
1111
+ elif method in {"chen", "covbat_gam"}:
845
1112
  return self._transform_chen(X, batch, disc, cont)
846
1113
  elif method == "longitudinal":
847
1114
  return self._transform_longitudinal(X, batch, disc, cont, subject_id, time_covariate)
848
1115
  else:
849
1116
  raise ValueError(
850
- f"Unknown method={method!r}. Expected one of 'johnson', 'fortin', "
851
- f"'chen', or 'longitudinal'."
1117
+ f"Unknown method={self.method!r}. Expected an author name "
1118
+ f"('johnson', 'fortin', 'chen', 'longitudinal', 'gam', 'covbat_gam') or a "
1119
+ f"method-name alias ('classic_combat', 'neurocombat', 'covbat', "
1120
+ f"'longcombat', 'combat_gam')."
852
1121
  )
853
1122
 
854
1123
  def _transform_johnson(self, X: pd.DataFrame, batch: pd.Series) -> pd.DataFrame:
@@ -880,7 +1149,10 @@ class ComBatModel:
880
1149
  disc: pd.DataFrame | None,
881
1150
  cont: pd.DataFrame | None,
882
1151
  ) -> pd.DataFrame:
883
- """Fortin transform implementation."""
1152
+ """Fortin transform implementation (also serves the gam method)."""
1153
+ if self._is_gam():
1154
+ cont = self._apply_smooth_terms(cont, fit=False)
1155
+
884
1156
  batch_dummies = pd.get_dummies(batch, drop_first=False).astype(float)
885
1157
  batch_dummies = batch_dummies.reindex(columns=self._batch_levels, fill_value=0.0)
886
1158
  if self.reference_batch is not None: