combatlearn 1.2.0__tar.gz → 1.3.0__tar.gz

{combatlearn-1.2.0 → combatlearn-1.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: combatlearn
-Version: 1.2.0
+Version: 1.3.0
 Summary: Batch-effect harmonization for machine learning frameworks.
 Author-email: Ettore Rocchi <ettoreroc@gmail.com>
 License: MIT
@@ -16,6 +16,7 @@ Requires-Dist: pandas>=1.3
 Requires-Dist: numpy>=1.21
 Requires-Dist: scikit-learn>=1.2
 Requires-Dist: matplotlib>=3.4
+Requires-Dist: seaborn>=0.12
 Requires-Dist: plotly>=5.0
 Requires-Dist: nbformat>=4.2
 Requires-Dist: umap-learn>=0.5
@@ -26,7 +27,7 @@ Requires-Dist: pre-commit>=3.0; extra == "dev"
 Provides-Extra: docs
 Requires-Dist: sphinx>=7.0.0; extra == "docs"
 Requires-Dist: pydata-sphinx-theme>=0.15.0; extra == "docs"
-Requires-Dist: myst-parser>=3.0.0; extra == "docs"
+Requires-Dist: myst-nb>=1.1.0; extra == "docs"
 Requires-Dist: sphinx-copybutton>=0.5.0; extra == "docs"
 Requires-Dist: linkify-it-py>=2.0.0; extra == "docs"
 Dynamic: license-file
@@ -76,7 +77,7 @@ from combatlearn import ComBat
 df = pd.read_csv("data.csv", index_col=0)
 X, y = df.drop(columns="y"), df["y"]
 
-batch = pd.read_csv("batch.csv", index_col=0, squeeze=True)
+batch = pd.read_csv("batch.csv", index_col=0).squeeze("columns")
 diag = pd.read_csv("diagnosis.csv", index_col=0) # categorical
 age = pd.read_csv("age.csv", index_col=0) # continuous
 

{combatlearn-1.2.0 → combatlearn-1.3.0}/README.md

@@ -43,7 +43,7 @@ from combatlearn import ComBat
 df = pd.read_csv("data.csv", index_col=0)
 X, y = df.drop(columns="y"), df["y"]
 
-batch = pd.read_csv("batch.csv", index_col=0, squeeze=True)
+batch = pd.read_csv("batch.csv", index_col=0).squeeze("columns")
 diag = pd.read_csv("diagnosis.csv", index_col=0) # categorical
 age = pd.read_csv("age.csv", index_col=0) # continuous
 

{combatlearn-1.2.0 → combatlearn-1.3.0}/combatlearn/__init__.py

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

{combatlearn-1.2.0 → combatlearn-1.3.0}/combatlearn/core.py

@@ -117,7 +117,26 @@ class ComBatModel:
         discrete_covariates: ArrayLike | None = None,
         continuous_covariates: ArrayLike | None = None,
     ) -> ComBatModel:
-        """Fit the ComBat model."""
+        """Fit the ComBat model.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Input data to fit.
+        y : None
+            Ignored. Present for API compatibility.
+        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).
+        continuous_covariates : array-like or None, default=None
+            Continuous covariates to protect (Fortin/Chen only).
+
+        Returns
+        -------
+        self : ComBatModel
+            Fitted model.
+        """
         method = self.method.lower()
         if method not in {"johnson", "fortin", "chen"}:
             raise ValueError("method must be 'johnson', 'fortin', or 'chen'.")
@@ -459,7 +478,29 @@ class ComBatModel:
         discrete_covariates: ArrayLike | None = None,
         continuous_covariates: ArrayLike | None = None,
     ) -> pd.DataFrame:
-        """Transform the data using fitted ComBat parameters."""
+        """Transform the data using fitted ComBat parameters.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Input data to transform.
+        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).
+        continuous_covariates : array-like or None, default=None
+            Continuous covariates (Fortin/Chen only).
+
+        Returns
+        -------
+        X_adjusted : pd.DataFrame
+            Batch-corrected data.
+
+        Raises
+        ------
+        ValueError
+            If the model is not fitted or if unseen batch levels are present.
+        """
         if not hasattr(self, "_gamma_star"):
             raise ValueError(
                 "This ComBatModel instance is not fitted yet. Call 'fit' before 'transform'."

{combatlearn-1.2.0 → combatlearn-1.3.0}/combatlearn/metrics.py

@@ -111,6 +111,7 @@ def _kbet_score(
     batch_labels: np.ndarray,
     k0: int,
     alpha: float = 0.05,
+    nn_algorithm: str = "auto",
 ) -> tuple[float, float]:
     """
     Compute kBET (k-nearest neighbor Batch Effect Test) acceptance rate.
@@ -130,6 +131,9 @@ def _kbet_score(
         Neighborhood size.
     alpha : float
         Significance level for chi-squared test.
+    nn_algorithm : {'auto', 'ball_tree', 'kd_tree', 'brute'}, default='auto'
+        Algorithm used for nearest neighbor computation. Passed to
+        ``sklearn.neighbors.NearestNeighbors``.
 
     Returns
     -------
@@ -148,7 +152,7 @@ def _kbet_score(
     global_freq = batch_counts / n_samples
     k0 = min(k0, n_samples - 1)
 
-    nn = NearestNeighbors(n_neighbors=k0 + 1, algorithm="auto")
+    nn = NearestNeighbors(n_neighbors=k0 + 1, algorithm=nn_algorithm)
     nn.fit(X)
     _, indices = nn.kneighbors(X)
 
@@ -236,6 +240,7 @@ def _lisi_score(
     X: np.ndarray,
     batch_labels: np.ndarray,
     perplexity: int = 30,
+    nn_algorithm: str = "auto",
 ) -> float:
     """
     Compute mean Local Inverse Simpson's Index (LISI).
@@ -253,6 +258,9 @@ def _lisi_score(
         Batch labels for each sample.
     perplexity : int
         Perplexity for Gaussian kernel.
+    nn_algorithm : {'auto', 'ball_tree', 'kd_tree', 'brute'}, default='auto'
+        Algorithm used for nearest neighbor computation. Passed to
+        ``sklearn.neighbors.NearestNeighbors``.
 
     Returns
     -------
@@ -269,7 +277,7 @@ def _lisi_score(
 
     k = min(3 * perplexity, n_samples - 1)
 
-    nn = NearestNeighbors(n_neighbors=k + 1, algorithm="auto")
+    nn = NearestNeighbors(n_neighbors=k + 1, algorithm=nn_algorithm)
     nn.fit(X)
     distances, indices = nn.kneighbors(X)
 
@@ -354,6 +362,7 @@ def _knn_preservation(
     X_after: np.ndarray,
     k_values: list[int],
     n_jobs: int = 1,
+    nn_algorithm: str = "auto",
 ) -> dict[int, float]:
     """
     Compute fraction of k-nearest neighbors preserved after correction.
@@ -371,6 +380,9 @@ def _knn_preservation(
         Values of k for k-NN.
     n_jobs : int
         Number of parallel jobs.
+    nn_algorithm : {'auto', 'ball_tree', 'kd_tree', 'brute'}, default='auto'
+        Algorithm used for nearest neighbor computation. Passed to
+        ``sklearn.neighbors.NearestNeighbors``.
 
     Returns
     -------
@@ -381,11 +393,11 @@ def _knn_preservation(
     max_k = max(k_values)
     max_k = min(max_k, X_before.shape[0] - 1)
 
-    nn_before = NearestNeighbors(n_neighbors=max_k + 1, algorithm="auto", n_jobs=n_jobs)
+    nn_before = NearestNeighbors(n_neighbors=max_k + 1, algorithm=nn_algorithm, n_jobs=n_jobs)
     nn_before.fit(X_before)
     _, indices_before = nn_before.kneighbors(X_before)
 
-    nn_after = NearestNeighbors(n_neighbors=max_k + 1, algorithm="auto", n_jobs=n_jobs)
+    nn_after = NearestNeighbors(n_neighbors=max_k + 1, algorithm=nn_algorithm, n_jobs=n_jobs)
     nn_after.fit(X_after)
     _, indices_after = nn_after.kneighbors(X_after)
 
@@ -557,6 +569,7 @@ class ComBatMetricsMixin:
         kbet_k0: int | None = None,
         lisi_perplexity: int = 30,
         n_jobs: int = 1,
+        nn_algorithm: str = "auto",
     ) -> dict[str, Any]:
         """
         Compute batch effect metrics before and after ComBat correction.
@@ -579,6 +592,9 @@ class ComBatMetricsMixin:
             Perplexity for LISI computation.
         n_jobs : int, default=1
             Number of parallel jobs for neighbor computations.
+        nn_algorithm : {'auto', 'ball_tree', 'kd_tree', 'brute'}, default='auto'
+            Algorithm used for nearest neighbor computation. Passed to
+            ``sklearn.neighbors.NearestNeighbors``.
 
         Returns
         -------
@@ -596,7 +612,11 @@ class ComBatMetricsMixin:
         ValueError
             If the model is not fitted or if pca_components is invalid.
         """
-        if not hasattr(self._model, "_gamma_star"):
+        _valid_nn = {"auto", "ball_tree", "kd_tree", "brute"}
+        if nn_algorithm not in _valid_nn:
+            raise ValueError(f"nn_algorithm must be one of {_valid_nn}, got '{nn_algorithm}'")
+
+        if not hasattr(self, "_model") or not hasattr(self._model, "_gamma_star"):
             raise ValueError(
                 "This ComBat instance is not fitted yet. Call 'fit' before 'compute_batch_metrics'."
             )
@@ -646,16 +666,22 @@ class ComBatMetricsMixin:
         db_before = _davies_bouldin_batch(X_before_pca, batch_labels)
         db_after = _davies_bouldin_batch(X_after_pca, batch_labels)
 
-        kbet_before, _ = _kbet_score(X_before_pca, batch_labels, kbet_k0)
-        kbet_after, _ = _kbet_score(X_after_pca, batch_labels, kbet_k0)
+        kbet_before, _ = _kbet_score(X_before_pca, batch_labels, kbet_k0, nn_algorithm=nn_algorithm)
+        kbet_after, _ = _kbet_score(X_after_pca, batch_labels, kbet_k0, nn_algorithm=nn_algorithm)
 
-        lisi_before = _lisi_score(X_before_pca, batch_labels, lisi_perplexity)
-        lisi_after = _lisi_score(X_after_pca, batch_labels, lisi_perplexity)
+        lisi_before = _lisi_score(
+            X_before_pca, batch_labels, lisi_perplexity, nn_algorithm=nn_algorithm
+        )
+        lisi_after = _lisi_score(
+            X_after_pca, batch_labels, lisi_perplexity, nn_algorithm=nn_algorithm
+        )
 
         var_ratio_before = _variance_ratio(X_before_pca, batch_labels)
         var_ratio_after = _variance_ratio(X_after_pca, batch_labels)
 
-        knn_results = _knn_preservation(X_before_pca, X_after_pca, k_neighbors, n_jobs)
+        knn_results = _knn_preservation(
+            X_before_pca, X_after_pca, k_neighbors, n_jobs, nn_algorithm=nn_algorithm
+        )
         dist_corr = _pairwise_distance_correlation(X_before_pca, X_after_pca)
 
         centroid_before = _mean_centroid_distance(X_before_pca, batch_labels)
@@ -743,7 +769,7 @@ class ComBatMetricsMixin:
         ValueError
             If the model is not fitted or if mode is invalid.
         """
-        if not hasattr(self._model, "_gamma_star"):
+        if not hasattr(self, "_model") or not hasattr(self._model, "_gamma_star"):
             raise ValueError(
                 "This ComBat instance is not fitted yet. "
                 "Call 'fit' before 'feature_batch_importance'."

combatlearn-1.3.0/combatlearn/sklearn_api.py

@@ -0,0 +1,277 @@
+"""Scikit-learn compatible ComBat wrapper."""
+
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+from sklearn.base import BaseEstimator, TransformerMixin
+from sklearn.utils.validation import check_array, check_is_fitted
+
+from .core import ArrayLike, ComBatModel
+from .metrics import ComBatMetricsMixin
+from .visualization import ComBatVisualizationMixin
+
+
+class ComBat(ComBatMetricsMixin, ComBatVisualizationMixin, BaseEstimator, TransformerMixin):
+    """Pipeline-friendly wrapper around `ComBatModel`.
+
+    Stores batch (and optional covariates) passed at construction and
+    appropriately uses them for separate `fit` and `transform`.
+
+    Parameters
+    ----------
+    batch : array-like of shape (n_samples,)
+        Batch labels for each sample.
+    discrete_covariates : array-like, optional
+        Categorical covariates to protect (Fortin/Chen only).
+    continuous_covariates : array-like, optional
+        Continuous covariates to protect (Fortin/Chen only).
+    method : {'johnson', 'fortin', 'chen'}, default='johnson'
+        ComBat variant to use.
+    parametric : bool, default=True
+        Use parametric empirical Bayes.
+    mean_only : bool, default=False
+        Adjust only the mean (ignore variance).
+    reference_batch : str, optional
+        Batch level to leave unchanged.
+    eps : float, default=1e-8
+        Numerical jitter for stability.
+    covbat_cov_thresh : float or int, default=0.9
+        CovBat variance threshold for PCs.
+    compute_metrics : bool, default=False
+        If True, ``fit_transform`` caches batch metrics in ``metrics_``.
+    """
+
+    def __init__(
+        self,
+        batch: ArrayLike,
+        *,
+        discrete_covariates: ArrayLike | None = None,
+        continuous_covariates: ArrayLike | None = None,
+        method: str = "johnson",
+        parametric: bool = True,
+        mean_only: bool = False,
+        reference_batch: str | None = None,
+        eps: float = 1e-8,
+        covbat_cov_thresh: float | int = 0.9,
+        compute_metrics: bool = False,
+    ) -> None:
+        self.batch = batch
+        self.discrete_covariates = discrete_covariates
+        self.continuous_covariates = continuous_covariates
+        self.method = method
+        self.parametric = parametric
+        self.mean_only = mean_only
+        self.reference_batch = reference_batch
+        self.eps = eps
+        self.covbat_cov_thresh = covbat_cov_thresh
+        self.compute_metrics = compute_metrics
+
+    def _validate_inputs(self, X: ArrayLike, *, fitting: bool = False) -> None:
+        """Validate X and, during fitting, batch/covariates for NaN/Inf."""
+        check_array(X, ensure_all_finite=True, dtype="numeric")
+
+        if fitting:
+            batch_ser = (
+                pd.Series(self.batch) if not isinstance(self.batch, pd.Series) else self.batch
+            )
+            if batch_ser.isna().any():
+                raise ValueError("batch contains NaN values.")
+
+            if self.discrete_covariates is not None:
+                disc_df = (
+                    pd.DataFrame(self.discrete_covariates)
+                    if not isinstance(self.discrete_covariates, (pd.Series, pd.DataFrame))
+                    else self.discrete_covariates
+                )
+                if (
+                    disc_df.isna().any().any()
+                    if isinstance(disc_df, pd.DataFrame)
+                    else disc_df.isna().any()
+                ):
+                    raise ValueError("discrete_covariates contains NaN values.")
+
+            if self.continuous_covariates is not None:
+                cont_arr = self.continuous_covariates
+                if isinstance(cont_arr, (pd.Series, pd.DataFrame)):
+                    cont_arr = cont_arr.values
+                check_array(
+                    np.atleast_2d(cont_arr) if np.asarray(cont_arr).ndim == 1 else cont_arr,
+                    ensure_all_finite=True,
+                    dtype="numeric",
+                )
+
+    def fit(self, X: ArrayLike, y: ArrayLike | None = None) -> ComBat:
+        """Fit the ComBat model.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Input data to fit.
+        y : None
+            Ignored. Present for API compatibility.
+
+        Returns
+        -------
+        self : ComBat
+            Fitted estimator.
+        """
+        self._validate_inputs(X, fitting=True)
+        idx = X.index if isinstance(X, pd.DataFrame) else pd.RangeIndex(len(X))
+
+        if isinstance(X, pd.DataFrame):
+            self.feature_names_in_ = np.asarray(X.columns, dtype=object)
+        else:
+            self.feature_names_in_ = np.asarray(
+                [f"x{i}" for i in range(np.asarray(X).shape[1])], dtype=object
+            )
+
+        self._model = ComBatModel(
+            method=self.method,
+            parametric=self.parametric,
+            mean_only=self.mean_only,
+            reference_batch=self.reference_batch,
+            eps=self.eps,
+            covbat_cov_thresh=self.covbat_cov_thresh,
+        )
+
+        batch_vec = self._subset(self.batch, idx)
+        disc = self._subset(self.discrete_covariates, idx)
+        cont = self._subset(self.continuous_covariates, idx)
+        self._model.fit(
+            X,
+            batch=batch_vec,
+            discrete_covariates=disc,
+            continuous_covariates=cont,
+        )
+        self._fitted_batch = batch_vec
+        return self
+
+    def transform(self, X: ArrayLike) -> pd.DataFrame:
+        """Transform the data using fitted ComBat parameters.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Input data to transform.
+
+        Returns
+        -------
+        X_transformed : pd.DataFrame
+            Batch-corrected data.
+        """
+        self._validate_inputs(X)
+        idx = X.index if isinstance(X, pd.DataFrame) else pd.RangeIndex(len(X))
+        batch_vec = self._subset(self.batch, idx)
+        disc = self._subset(self.discrete_covariates, idx)
+        cont = self._subset(self.continuous_covariates, idx)
+        return self._model.transform(
+            X,
+            batch=batch_vec,
+            discrete_covariates=disc,
+            continuous_covariates=cont,
+        )
+
+    @staticmethod
+    def _subset(obj: ArrayLike | None, idx: pd.Index) -> pd.DataFrame | pd.Series | None:
+        """Subset array-like object by index."""
+        if obj is None:
+            return None
+        if isinstance(obj, (pd.Series, pd.DataFrame)):
+            return obj.loc[idx]
+        else:
+            if isinstance(obj, np.ndarray) and obj.ndim == 1:
+                return pd.Series(obj, index=idx)
+            else:
+                return pd.DataFrame(obj, index=idx)
+
+    def fit_transform(self, X: ArrayLike, y: ArrayLike | None = None) -> pd.DataFrame:
+        """
+        Fit and transform the data, optionally computing metrics.
+
+        If ``compute_metrics=True`` was set at construction, batch effect
+        metrics are computed and cached in the ``metrics_`` property.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Input data to fit and transform.
+        y : None
+            Ignored. Present for API compatibility.
+
+        Returns
+        -------
+        X_transformed : pd.DataFrame
+            Batch-corrected data.
+        """
+        self.fit(X, y)
+        X_transformed = self.transform(X)
+
+        if self.compute_metrics:
+            self._metrics_cache = self.compute_batch_metrics(X)
+
+        return X_transformed
+
+    def get_feature_names_out(self, input_features: ArrayLike | None = None) -> np.ndarray:
+        """Get output feature names for transform.
+
+        Parameters
+        ----------
+        input_features : array-like of str or None, default=None
+            Ignored. Present for API compatibility.
+
+        Returns
+        -------
+        feature_names_out : ndarray of str objects
+            Feature names.
+
+        Raises
+        ------
+        sklearn.exceptions.NotFittedError
+            If the estimator is not fitted.
+        """
+        check_is_fitted(self, "feature_names_in_")
+        return self.feature_names_in_
+
+    def summary(self) -> str:
+        """Return a human-readable diagnostic report after fitting.
+
+        Returns
+        -------
+        str
+            Multi-line summary string.
+
+        Raises
+        ------
+        ValueError
+            If the model is not fitted.
+        """
+        if not hasattr(self, "_model") or not hasattr(self._model, "_gamma_star"):
+            raise ValueError("This ComBat instance is not fitted yet. Call 'fit' before 'summary'.")
+
+        lines = []
+        lines.append("ComBat Summary")
+        lines.append("=" * 40)
+        lines.append(f"Method:          {self.method}")
+        lines.append(f"Parametric:      {self.parametric}")
+        lines.append(f"Mean only:       {self.mean_only}")
+        lines.append(f"Reference batch: {self.reference_batch or 'None'}")
+
+        batch_levels = self._model._batch_levels
+        n_per_batch = self._model._n_per_batch
+        lines.append(f"Number of batches: {len(batch_levels)}")
+        lines.append("Samples per batch:")
+        for lvl in batch_levels:
+            lines.append(f"  {lvl}: {n_per_batch[str(lvl)]}")
+
+        n_features = len(self._model._grand_mean)
+        lines.append(f"Number of features: {n_features}")
+
+        lines.append("")
+        lines.append("Top 5 features by batch effect (combined):")
+        importance = self.feature_batch_importance()
+        top5 = importance.head(5)
+        for feat, row in top5.iterrows():
+            lines.append(f"  {feat}: {row['combined']:.4f}")
+
+        return "\n".join(lines)

{combatlearn-1.2.0 → combatlearn-1.3.0}/combatlearn/visualization.py

@@ -10,6 +10,7 @@ import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import plotly.graph_objects as go
+import seaborn as sns
 import umap
 from plotly.subplots import make_subplots
 from sklearn.decomposition import PCA
@@ -59,11 +60,30 @@ class ComBatVisualizationMixin:
             - `'static'`: matplotlib plots (can be saved as images)
             - `'interactive'`: plotly plots (explorable, requires plotly)
 
+        figsize : tuple of int, default=(12, 5)
+            Figure size in inches (width, height). Only used for static plots.
+
+        alpha : float, default=0.7
+            Marker transparency. Only used for static plots.
+
+        point_size : int, default=50
+            Marker size. Only used for static plots.
+
+        cmap : str, default='Set1'
+            Matplotlib colormap name for batch colors.
+
+        title : str or None, default=None
+            Custom figure title. If None, a default title is generated.
+
+        show_legend : bool, default=True
+            Whether to display the batch legend.
+
         return_embeddings : bool, default=False
             If `True`, return embeddings along with the plot.
 
         **reduction_kwargs : dict
-            Additional parameters for reduction methods.
+            Additional keyword arguments passed to the reduction method
+            (e.g., ``perplexity`` for t-SNE, ``n_neighbors`` for UMAP).
 
         Returns
         -------
@@ -75,7 +95,7 @@ class ComBatVisualizationMixin:
             - `'original'`: embedding of original data
             - `'transformed'`: embedding of ComBat-transformed data
         """
-        if not hasattr(self._model, "_gamma_star"):
+        if not hasattr(self, "_model") or not hasattr(self._model, "_gamma_star"):
             raise ValueError(
                 "This ComBat instance is not fitted yet. Call 'fit' before 'plot_transformation'."
             )
@@ -439,7 +459,7 @@ class ComBatVisualizationMixin:
         ValueError
             If the model is not fitted, or if kind/mode is invalid.
         """
-        if not hasattr(self._model, "_gamma_star"):
+        if not hasattr(self, "_model") or not hasattr(self._model, "_gamma_star"):
             raise ValueError(
                 "This ComBat instance is not fitted yet. "
                 "Call 'fit' before 'plot_feature_importance'."
@@ -531,3 +551,93 @@ class ComBatVisualizationMixin:
             print(f"Top {top_n} features explain {cumulative_pct:.1f}% of total {effect_label}")
 
         return fig
+
+    def plot_batch_effect_heatmap(
+        self,
+        top_n: int = 50,
+        figsize: tuple[int, int] = (12, 8),
+    ) -> Any:
+        """Plot a heatmap of batch effect parameters across features and batches.
+
+        Displays the estimated batch-specific location shifts (gamma) and,
+        unless ``mean_only=True``, log-scale shifts (log delta) for the
+        ``top_n`` most affected features.
+
+        Parameters
+        ----------
+        top_n : int, default=50
+            Number of top features (by combined batch effect) to display.
+        figsize : tuple of int, default=(12, 8)
+            Figure size in inches.
+
+        Returns
+        -------
+        matplotlib.figure.Figure
+            Figure containing the heatmap(s).
+
+        Raises
+        ------
+        ValueError
+            If the model is not fitted.
+        ImportError
+            If seaborn is not installed.
+        """
+        if not hasattr(self, "_model") or not hasattr(self._model, "_gamma_star"):
+            raise ValueError(
+                "This ComBat instance is not fitted yet. "
+                "Call 'fit' before 'plot_batch_effect_heatmap'."
+            )
+
+        feature_names = self._model._grand_mean.index
+        batch_levels = self._model._batch_levels
+        gamma_star = self._model._gamma_star
+        delta_star = self._model._delta_star
+
+        importance = self.feature_batch_importance()
+        top_features = importance.head(top_n).index
+        feat_idx = [feature_names.get_loc(f) for f in top_features]
+
+        gamma_df = pd.DataFrame(
+            gamma_star[:, feat_idx],
+            index=[str(b) for b in batch_levels],
+            columns=top_features,
+        )
+
+        n_plots = 1 if self.mean_only else 2
+        fig, axes = plt.subplots(1, n_plots, figsize=figsize)
+        if n_plots == 1:
+            axes = [axes]
+
+        sns.heatmap(
+            gamma_df,
+            cmap="RdBu_r",
+            center=0,
+            ax=axes[0],
+            xticklabels=True,
+            yticklabels=True,
+        )
+        axes[0].set_title("Location shifts (gamma)")
+        axes[0].set_xlabel("Feature")
+        axes[0].set_ylabel("Batch")
+
+        if not self.mean_only:
+            log_delta_df = pd.DataFrame(
+                np.log(delta_star[:, feat_idx]),
+                index=[str(b) for b in batch_levels],
+                columns=top_features,
+            )
+            sns.heatmap(
+                log_delta_df,
+                cmap="RdBu_r",
+                center=0,
+                ax=axes[1],
+                xticklabels=True,
+                yticklabels=True,
+            )
+            axes[1].set_title("Scale shifts (log delta)")
+            axes[1].set_xlabel("Feature")
+            axes[1].set_ylabel("Batch")
+
+        fig.suptitle(f"Batch Effect Heatmap (top {top_n} features)", fontsize=14, fontweight="bold")
+        plt.tight_layout()
+        return fig

{combatlearn-1.2.0 → combatlearn-1.3.0}/combatlearn.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: combatlearn
-Version: 1.2.0
+Version: 1.3.0
 Summary: Batch-effect harmonization for machine learning frameworks.
 Author-email: Ettore Rocchi <ettoreroc@gmail.com>
 License: MIT
@@ -16,6 +16,7 @@ Requires-Dist: pandas>=1.3
 Requires-Dist: numpy>=1.21
 Requires-Dist: scikit-learn>=1.2
 Requires-Dist: matplotlib>=3.4
+Requires-Dist: seaborn>=0.12
 Requires-Dist: plotly>=5.0
 Requires-Dist: nbformat>=4.2
 Requires-Dist: umap-learn>=0.5
@@ -26,7 +27,7 @@ Requires-Dist: pre-commit>=3.0; extra == "dev"
 Provides-Extra: docs
 Requires-Dist: sphinx>=7.0.0; extra == "docs"
 Requires-Dist: pydata-sphinx-theme>=0.15.0; extra == "docs"
-Requires-Dist: myst-parser>=3.0.0; extra == "docs"
+Requires-Dist: myst-nb>=1.1.0; extra == "docs"
 Requires-Dist: sphinx-copybutton>=0.5.0; extra == "docs"
 Requires-Dist: linkify-it-py>=2.0.0; extra == "docs"
 Dynamic: license-file
@@ -76,7 +77,7 @@ from combatlearn import ComBat
 df = pd.read_csv("data.csv", index_col=0)
 X, y = df.drop(columns="y"), df["y"]
 
-batch = pd.read_csv("batch.csv", index_col=0, squeeze=True)
+batch = pd.read_csv("batch.csv", index_col=0).squeeze("columns")
 diag = pd.read_csv("diagnosis.csv", index_col=0) # categorical
 age = pd.read_csv("age.csv", index_col=0) # continuous
 

{combatlearn-1.2.0 → combatlearn-1.3.0}/combatlearn.egg-info/requires.txt

@@ -2,6 +2,7 @@ pandas>=1.3
 numpy>=1.21
 scikit-learn>=1.2
 matplotlib>=3.4
+seaborn>=0.12
 plotly>=5.0
 nbformat>=4.2
 umap-learn>=0.5
@@ -14,6 +15,6 @@ pre-commit>=3.0
 [docs]
 sphinx>=7.0.0
 pydata-sphinx-theme>=0.15.0
-myst-parser>=3.0.0
+myst-nb>=1.1.0
 sphinx-copybutton>=0.5.0
 linkify-it-py>=2.0.0

{combatlearn-1.2.0 → combatlearn-1.3.0}/docs_requirements.txt

@@ -1,5 +1,5 @@
 sphinx>=7.0.0
 pydata-sphinx-theme>=0.15.0
-myst-parser>=3.0.0
+myst-nb>=1.1.0
 sphinx-copybutton>=0.5.0
 linkify-it-py>=2.0.0

{combatlearn-1.2.0 → combatlearn-1.3.0}/pyproject.toml

@@ -44,4 +44,3 @@ ignore = ["E501", "B008"]
 
 [tool.ruff.lint.per-file-ignores]
 "tests/*" = ["S101", "PLR2004"]
-"examples/*" = ["F401"]

{combatlearn-1.2.0 → combatlearn-1.3.0}/requirements.txt

@@ -2,6 +2,7 @@ pandas>=1.3
 numpy>=1.21
 scikit-learn>=1.2
 matplotlib>=3.4
+seaborn>=0.12
 plotly>=5.0
 nbformat>=4.2
 umap-learn>=0.5

{combatlearn-1.2.0 → combatlearn-1.3.0}/tests/test_combat.py

@@ -2,6 +2,7 @@ import numpy as np
 import pandas as pd
 import pytest
 from sklearn.base import clone
+from sklearn.exceptions import NotFittedError
 from sklearn.pipeline import Pipeline
 from sklearn.preprocessing import StandardScaler
 from utils import simulate_covariate_data, simulate_data
@@ -653,3 +654,198 @@ def test_plot_feature_importance_distribution_shows_cumulative(capsys):
     captured = capsys.readouterr()
     assert "features explain" in captured.out
     assert "%" in captured.out
+
+
+def test_fit_nan_in_X_raises():
+    """Fitting with NaN in X must raise ValueError."""
+    X, batch = simulate_data()
+    X.iloc[0, 0] = np.nan
+    with pytest.raises(ValueError):
+        ComBat(batch=batch).fit(X)
+
+
+def test_fit_inf_in_X_raises():
+    """Fitting with Inf in X must raise ValueError."""
+    X, batch = simulate_data()
+    X.iloc[0, 0] = np.inf
+    with pytest.raises(ValueError):
+        ComBat(batch=batch).fit(X)
+
+
+def test_transform_nan_in_X_raises():
+    """Transforming with NaN in X must raise ValueError."""
+    X, batch = simulate_data()
+    combat = ComBat(batch=batch).fit(X)
+    X_bad = X.copy()
+    X_bad.iloc[0, 0] = np.nan
+    with pytest.raises(ValueError):
+        combat.transform(X_bad)
+
+
+def test_fit_nan_in_batch_raises():
+    """Fitting with NaN in batch must raise ValueError."""
+    X, batch = simulate_data()
+    batch.iloc[0] = np.nan
+    with pytest.raises(ValueError, match="batch contains NaN"):
+        ComBat(batch=batch).fit(X)
+
+
+def test_fit_nan_in_discrete_covariates_raises():
+    """Fitting with NaN in discrete covariates must raise ValueError."""
+    X, batch, disc, cont = simulate_covariate_data()
+    disc.iloc[0, 0] = np.nan
+    with pytest.raises(ValueError, match="discrete_covariates contains NaN"):
+        ComBat(
+            batch=batch, discrete_covariates=disc, continuous_covariates=cont, method="fortin"
+        ).fit(X)
+
+
+def test_fit_nan_in_continuous_covariates_raises():
+    """Fitting with NaN in continuous covariates must raise ValueError."""
+    X, batch, disc, cont = simulate_covariate_data()
+    cont.iloc[0, 0] = np.nan
+    with pytest.raises(ValueError):
+        ComBat(
+            batch=batch, discrete_covariates=disc, continuous_covariates=cont, method="fortin"
+        ).fit(X)
+
+
+def test_metrics_nn_algorithm_ball_tree():
+    """compute_batch_metrics with nn_algorithm='ball_tree' must work."""
+    X, batch = simulate_data(n_samples=100, n_features=20)
+    combat = ComBat(batch=batch).fit(X)
+    metrics = combat.compute_batch_metrics(X, k_neighbors=[5], nn_algorithm="ball_tree")
+    assert "batch_effect" in metrics
+
+
+def test_metrics_nn_algorithm_kd_tree():
+    """compute_batch_metrics with nn_algorithm='kd_tree' must work."""
+    X, batch = simulate_data(n_samples=100, n_features=20)
+    combat = ComBat(batch=batch).fit(X)
+    metrics = combat.compute_batch_metrics(X, k_neighbors=[5], nn_algorithm="kd_tree")
+    assert "batch_effect" in metrics
+
+
+def test_metrics_nn_algorithm_brute():
+    """compute_batch_metrics with nn_algorithm='brute' must work."""
+    X, batch = simulate_data(n_samples=100, n_features=20)
+    combat = ComBat(batch=batch).fit(X)
+    metrics = combat.compute_batch_metrics(X, k_neighbors=[5], nn_algorithm="brute")
+    assert "batch_effect" in metrics
+
+
+def test_metrics_nn_algorithm_invalid_raises():
+    """compute_batch_metrics with invalid nn_algorithm must raise ValueError."""
+    X, batch = simulate_data(n_samples=100, n_features=20)
+    combat = ComBat(batch=batch).fit(X)
+    with pytest.raises(ValueError, match="nn_algorithm"):
+        combat.compute_batch_metrics(X, nn_algorithm="invalid")
+
+
+def test_get_feature_names_out_dataframe():
+    """get_feature_names_out returns column names from DataFrame input."""
+    X, batch = simulate_data(n_samples=100, n_features=10)
+    X.columns = [f"feat_{i}" for i in range(10)]
+    combat = ComBat(batch=batch).fit(X)
+    names = combat.get_feature_names_out()
+    assert list(names) == [f"feat_{i}" for i in range(10)]
+
+
+def test_get_feature_names_out_ndarray():
+    """get_feature_names_out returns default names for numpy input."""
+    X, batch = simulate_data(n_samples=100, n_features=10)
+    combat = ComBat(batch=batch).fit(X.values)
+    names = combat.get_feature_names_out()
+    assert list(names) == [f"x{i}" for i in range(10)]
+
+
+def test_get_feature_names_out_not_fitted_raises():
+    """get_feature_names_out before fit raises NotFittedError."""
+    _X, batch = simulate_data()
+    combat = ComBat(batch=batch)
+    with pytest.raises(NotFittedError):
+        combat.get_feature_names_out()
+
+
+def test_set_output_pandas_in_pipeline():
+    """Pipeline with set_output(transform='pandas') works with ComBat."""
+    X, batch = simulate_data(n_samples=100, n_features=10)
+    X.columns = [f"feat_{i}" for i in range(10)]
+    pipe = Pipeline([("combat", ComBat(batch=batch))])
+    pipe.set_output(transform="pandas")
+    result = pipe.fit_transform(X)
+    assert isinstance(result, pd.DataFrame)
+    assert list(result.columns) == [f"feat_{i}" for i in range(10)]
+
+
+def test_summary_contains_method():
+    """summary() output contains the method name."""
+    X, batch = simulate_data()
+    combat = ComBat(batch=batch, method="johnson").fit(X)
+    s = combat.summary()
+    assert "johnson" in s
+
+
+def test_summary_contains_batch_info():
+    """summary() output contains batch count and sample counts."""
+    X, batch = simulate_data()
+    combat = ComBat(batch=batch).fit(X)
+    s = combat.summary()
+    assert "Number of batches: 3" in s
+    assert "A:" in s
+    assert "B:" in s
+    assert "C:" in s
+
+
+def test_summary_not_fitted_raises():
+    """summary() before fit raises ValueError."""
+    _X, batch = simulate_data()
+    combat = ComBat(batch=batch)
+    with pytest.raises(ValueError, match="not fitted"):
+        combat.summary()
+
+
+def test_heatmap_returns_figure():
+    """plot_batch_effect_heatmap returns a matplotlib Figure."""
+    import matplotlib.figure
+
+    X, batch = simulate_data(n_samples=100, n_features=25)
+    combat = ComBat(batch=batch).fit(X)
+    fig = combat.plot_batch_effect_heatmap(top_n=10)
+    assert isinstance(fig, matplotlib.figure.Figure)
+
+
+def test_heatmap_top_n():
+    """plot_batch_effect_heatmap respects top_n parameter."""
+    X, batch = simulate_data(n_samples=100, n_features=25)
+    combat = ComBat(batch=batch).fit(X)
+    fig = combat.plot_batch_effect_heatmap(top_n=10)
+    # With mean_only=False, 2 subplots
+    assert len(fig.axes) >= 2
+
+
+def test_heatmap_mean_only():
+    """plot_batch_effect_heatmap with mean_only=True shows only 1 heatmap."""
+    X, batch = simulate_data(n_samples=100, n_features=25)
+    combat = ComBat(batch=batch, mean_only=True).fit(X)
+    fig = combat.plot_batch_effect_heatmap(top_n=10)
+    # mean_only → 1 subplot (gamma only), but colorbar adds an extra axis
+    # so check the first axis title
+    assert "gamma" in fig.axes[0].get_title().lower()
+
+
+def test_heatmap_not_fitted_raises():
+    """plot_batch_effect_heatmap before fit raises ValueError."""
+    _X, batch = simulate_data()
+    combat = ComBat(batch=batch)
+    with pytest.raises(ValueError, match="not fitted"):
+        combat.plot_batch_effect_heatmap()
+
+
+def test_set_params_updates_method():
+    """set_params must propagate to the model created at fit time."""
+    X, batch = simulate_data()
+    combat = ComBat(batch=batch, method="johnson")
+    combat.set_params(method="fortin")
+    combat.fit(X)
+    assert combat._model.method == "fortin"

combatlearn-1.2.0/combatlearn/sklearn_api.py

@@ -1,143 +0,0 @@
-"""Scikit-learn compatible ComBat wrapper."""
-
-from __future__ import annotations
-
-import numpy as np
-import pandas as pd
-from sklearn.base import BaseEstimator, TransformerMixin
-
-from .core import ArrayLike, ComBatModel
-from .metrics import ComBatMetricsMixin
-from .visualization import ComBatVisualizationMixin
-
-
-class ComBat(ComBatMetricsMixin, ComBatVisualizationMixin, BaseEstimator, TransformerMixin):
-    """Pipeline-friendly wrapper around `ComBatModel`.
-
-    Stores batch (and optional covariates) passed at construction and
-    appropriately uses them for separate `fit` and `transform`.
-
-    Parameters
-    ----------
-    batch : array-like of shape (n_samples,)
-        Batch labels for each sample.
-    discrete_covariates : array-like, optional
-        Categorical covariates to protect (Fortin/Chen only).
-    continuous_covariates : array-like, optional
-        Continuous covariates to protect (Fortin/Chen only).
-    method : {'johnson', 'fortin', 'chen'}, default='johnson'
-        ComBat variant to use.
-    parametric : bool, default=True
-        Use parametric empirical Bayes.
-    mean_only : bool, default=False
-        Adjust only the mean (ignore variance).
-    reference_batch : str, optional
-        Batch level to leave unchanged.
-    eps : float, default=1e-8
-        Numerical jitter for stability.
-    covbat_cov_thresh : float or int, default=0.9
-        CovBat variance threshold for PCs.
-    compute_metrics : bool, default=False
-        If True, ``fit_transform`` caches batch metrics in ``metrics_``.
-    """
-
-    def __init__(
-        self,
-        batch: ArrayLike,
-        *,
-        discrete_covariates: ArrayLike | None = None,
-        continuous_covariates: ArrayLike | None = None,
-        method: str = "johnson",
-        parametric: bool = True,
-        mean_only: bool = False,
-        reference_batch: str | None = None,
-        eps: float = 1e-8,
-        covbat_cov_thresh: float | int = 0.9,
-        compute_metrics: bool = False,
-    ) -> None:
-        self.batch = batch
-        self.discrete_covariates = discrete_covariates
-        self.continuous_covariates = continuous_covariates
-        self.method = method
-        self.parametric = parametric
-        self.mean_only = mean_only
-        self.reference_batch = reference_batch
-        self.eps = eps
-        self.covbat_cov_thresh = covbat_cov_thresh
-        self.compute_metrics = compute_metrics
-        self._model = ComBatModel(
-            method=method,
-            parametric=parametric,
-            mean_only=mean_only,
-            reference_batch=reference_batch,
-            eps=eps,
-            covbat_cov_thresh=covbat_cov_thresh,
-        )
-
-    def fit(self, X: ArrayLike, y: ArrayLike | None = None) -> ComBat:
-        """Fit the ComBat model."""
-        idx = X.index if isinstance(X, pd.DataFrame) else pd.RangeIndex(len(X))
-        batch_vec = self._subset(self.batch, idx)
-        disc = self._subset(self.discrete_covariates, idx)
-        cont = self._subset(self.continuous_covariates, idx)
-        self._model.fit(
-            X,
-            batch=batch_vec,
-            discrete_covariates=disc,
-            continuous_covariates=cont,
-        )
-        self._fitted_batch = batch_vec
-        return self
-
-    def transform(self, X: ArrayLike) -> pd.DataFrame:
-        """Transform the data using fitted ComBat parameters."""
-        idx = X.index if isinstance(X, pd.DataFrame) else pd.RangeIndex(len(X))
-        batch_vec = self._subset(self.batch, idx)
-        disc = self._subset(self.discrete_covariates, idx)
-        cont = self._subset(self.continuous_covariates, idx)
-        return self._model.transform(
-            X,
-            batch=batch_vec,
-            discrete_covariates=disc,
-            continuous_covariates=cont,
-        )
-
-    @staticmethod
-    def _subset(obj: ArrayLike | None, idx: pd.Index) -> pd.DataFrame | pd.Series | None:
-        """Subset array-like object by index."""
-        if obj is None:
-            return None
-        if isinstance(obj, (pd.Series, pd.DataFrame)):
-            return obj.loc[idx]
-        else:
-            if isinstance(obj, np.ndarray) and obj.ndim == 1:
-                return pd.Series(obj, index=idx)
-            else:
-                return pd.DataFrame(obj, index=idx)
-
-    def fit_transform(self, X: ArrayLike, y: ArrayLike | None = None) -> pd.DataFrame:
-        """
-        Fit and transform the data, optionally computing metrics.
-
-        If ``compute_metrics=True`` was set at construction, batch effect
-        metrics are computed and cached in the ``metrics_`` property.
-
-        Parameters
-        ----------
-        X : array-like of shape (n_samples, n_features)
-            Input data to fit and transform.
-        y : None
-            Ignored. Present for API compatibility.
-
-        Returns
-        -------
-        X_transformed : pd.DataFrame
-            Batch-corrected data.
-        """
-        self.fit(X, y)
-        X_transformed = self.transform(X)
-
-        if self.compute_metrics:
-            self._metrics_cache = self.compute_batch_metrics(X)
-
-        return X_transformed