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.
- {combatlearn-2.1.0 → combatlearn-2.2.0}/PKG-INFO +13 -5
- {combatlearn-2.1.0 → combatlearn-2.2.0}/README.md +13 -5
- {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn/__init__.py +1 -1
- {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn/core.py +291 -19
- {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn/sklearn_api.py +38 -7
- {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn/visualization.py +105 -49
- {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn.egg-info/PKG-INFO +13 -5
- {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn.egg-info/SOURCES.txt +2 -0
- {combatlearn-2.1.0 → combatlearn-2.2.0}/tests/test_combat.py +46 -0
- combatlearn-2.2.0/tests/test_eb_degenerate.py +111 -0
- combatlearn-2.2.0/tests/test_gam.py +298 -0
- {combatlearn-2.1.0 → combatlearn-2.2.0}/tests/test_longitudinal.py +12 -0
- {combatlearn-2.1.0 → combatlearn-2.2.0}/tests/test_visualization.py +35 -0
- {combatlearn-2.1.0 → combatlearn-2.2.0}/LICENSE +0 -0
- {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn/_mixed.py +0 -0
- {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn/_utils.py +0 -0
- {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn/inspection.py +0 -0
- {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn/metrics.py +0 -0
- {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn/py.typed +0 -0
- {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn.egg-info/dependency_links.txt +0 -0
- {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn.egg-info/requires.txt +0 -0
- {combatlearn-2.1.0 → combatlearn-2.2.0}/combatlearn.egg-info/top_level.txt +0 -0
- {combatlearn-2.1.0 → combatlearn-2.2.0}/pyproject.toml +0 -0
- {combatlearn-2.1.0 → combatlearn-2.2.0}/requirements-dev.txt +0 -0
- {combatlearn-2.1.0 → combatlearn-2.2.0}/requirements-docs.txt +0 -0
- {combatlearn-2.1.0 → combatlearn-2.2.0}/requirements.txt +0 -0
- {combatlearn-2.1.0 → combatlearn-2.2.0}/setup.cfg +0 -0
- {combatlearn-2.1.0 → combatlearn-2.2.0}/tests/test_inspection.py +0 -0
- {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.
|
|
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
|
-
**
|
|
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 `"
|
|
132
|
-
| `continuous_covariates` | array-like, pd.Series or pd.DataFrame | `None` | Optional continuous covariates (e.g., age). Only used in `"fortin"`, `"chen"`, and `"
|
|
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
|
-
**
|
|
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 `"
|
|
94
|
-
| `continuous_covariates` | array-like, pd.Series or pd.DataFrame | `None` | Optional continuous covariates (e.g., age). Only used in `"fortin"`, `"chen"`, and `"
|
|
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,10 +1,12 @@
|
|
|
1
1
|
"""ComBat algorithm core.
|
|
2
2
|
|
|
3
|
-
`ComBatModel` implements
|
|
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'
|
|
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[
|
|
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
|
|
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"
|
|
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
|
|
300
|
+
elif method in {"fortin", "gam"}:
|
|
210
301
|
self._fit_fortin(X, batch, disc, cont)
|
|
211
|
-
elif method
|
|
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 =
|
|
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 =
|
|
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
|
|
1106
|
+
method = _resolve_method(self.method)
|
|
840
1107
|
if method == "johnson":
|
|
841
1108
|
return self._transform_johnson(X, batch)
|
|
842
|
-
elif method
|
|
1109
|
+
elif method in {"fortin", "gam"}:
|
|
843
1110
|
return self._transform_fortin(X, batch, disc, cont)
|
|
844
|
-
elif method
|
|
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
|
|
851
|
-
f"'
|
|
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:
|