combatlearn 0.2.2__py3-none-any.whl → 1.1.0__py3-none-any.whl

combatlearn/__init__.py

@@ -1,4 +1,5 @@
-from .combat import ComBatModel, ComBat
+from .combat import ComBat
 
-__all__ = ["ComBatModel", "ComBat"]
-__version__ = "0.2.2"
+__all__ = ["ComBat"]
+__version__ = "1.1.0"
+__author__ = "Ettore Rocchi"

combatlearn/combat.py

@@ -14,32 +14,544 @@ import numpy as np
 import numpy.linalg as la
 import pandas as pd
 from sklearn.base import BaseEstimator, TransformerMixin
-from sklearn.utils.validation import check_is_fitted
 from sklearn.decomposition import PCA
 from sklearn.manifold import TSNE
+from sklearn.neighbors import NearestNeighbors
+from sklearn.metrics import silhouette_score, davies_bouldin_score
+from scipy.stats import levene, spearmanr, chi2
+from scipy.spatial.distance import pdist
+import matplotlib
 import matplotlib.pyplot as plt
 import matplotlib.colors as mcolors
-from typing import Literal, Optional, Union, Dict, Tuple, Any, cast
+from typing import Literal, Optional, Union, Dict, Tuple, Any, List
 import numpy.typing as npt
 import warnings
+import umap
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
 
-try:
-    import umap
-    UMAP_AVAILABLE = True
-except ImportError:
-    UMAP_AVAILABLE = False
+ArrayLike = Union[pd.DataFrame, pd.Series, npt.NDArray[Any]]
+FloatArray = npt.NDArray[np.float64]
 
-try:
-    import plotly.graph_objects as go
-    from plotly.subplots import make_subplots
-    PLOTLY_AVAILABLE = True
-except ImportError:
-    PLOTLY_AVAILABLE = False
+def _compute_pca_embedding(
+    X_before: np.ndarray,
+    X_after: np.ndarray,
+    n_components: int,
+) -> Tuple[np.ndarray, np.ndarray, PCA]:
+    """
+    Compute PCA embeddings for both datasets.
 
-__author__ = "Ettore Rocchi"
+    Fits PCA on X_before and applies to both datasets.
 
-ArrayLike = Union[pd.DataFrame, pd.Series, npt.NDArray[Any]]
-FloatArray = npt.NDArray[np.float64]
+    Parameters
+    ----------
+    X_before : np.ndarray
+        Original data before correction.
+    X_after : np.ndarray
+        Corrected data.
+    n_components : int
+        Number of PCA components.
+
+    Returns
+    -------
+    X_before_pca : np.ndarray
+        PCA-transformed original data.
+    X_after_pca : np.ndarray
+        PCA-transformed corrected data.
+    pca : PCA
+        Fitted PCA model.
+    """
+    n_components = min(n_components, X_before.shape[1], X_before.shape[0] - 1)
+    pca = PCA(n_components=n_components, random_state=42)
+    X_before_pca = pca.fit_transform(X_before)
+    X_after_pca = pca.transform(X_after)
+    return X_before_pca, X_after_pca, pca
+
+
+def _silhouette_batch(X: np.ndarray, batch_labels: np.ndarray) -> float:
+    """
+    Compute silhouette coefficient using batch as cluster labels.
+
+    Lower values after correction indicate better batch mixing.
+    Range: [-1, 1], where -1 = batch mixing, 1 = batch separation.
+
+    Parameters
+    ----------
+    X : np.ndarray
+        Data matrix.
+    batch_labels : np.ndarray
+        Batch labels for each sample.
+
+    Returns
+    -------
+    float
+        Silhouette coefficient.
+    """
+    unique_batches = np.unique(batch_labels)
+    if len(unique_batches) < 2:
+        return 0.0
+    try:
+        return silhouette_score(X, batch_labels, metric='euclidean')
+    except Exception:
+        return 0.0
+
+
+def _davies_bouldin_batch(X: np.ndarray, batch_labels: np.ndarray) -> float:
+    """
+    Compute Davies-Bouldin index using batch labels.
+
+    Lower values indicate better batch mixing.
+    Range: [0, inf), 0 = perfect batch overlap.
+
+    Parameters
+    ----------
+    X : np.ndarray
+        Data matrix.
+    batch_labels : np.ndarray
+        Batch labels for each sample.
+
+    Returns
+    -------
+    float
+        Davies-Bouldin index.
+    """
+    unique_batches = np.unique(batch_labels)
+    if len(unique_batches) < 2:
+        return 0.0
+    try:
+        return davies_bouldin_score(X, batch_labels)
+    except Exception:
+        return 0.0
+
+
+def _kbet_score(
+    X: np.ndarray,
+    batch_labels: np.ndarray,
+    k0: int,
+    alpha: float = 0.05,
+) -> Tuple[float, float]:
+    """
+    Compute kBET (k-nearest neighbor Batch Effect Test) acceptance rate.
+
+    Tests if local batch proportions match global batch proportions.
+    Higher acceptance rate = better batch mixing.
+
+    Reference: Buttner et al. (2019) Nature Methods
+
+    Parameters
+    ----------
+    X : np.ndarray
+        Data matrix.
+    batch_labels : np.ndarray
+        Batch labels for each sample.
+    k0 : int
+        Neighborhood size.
+    alpha : float
+        Significance level for chi-squared test.
+
+    Returns
+    -------
+    acceptance_rate : float
+        Fraction of samples where H0 (uniform mixing) is accepted.
+    mean_stat : float
+        Mean chi-squared statistic across samples.
+    """
+    n_samples = X.shape[0]
+    unique_batches, batch_counts = np.unique(batch_labels, return_counts=True)
+    n_batches = len(unique_batches)
+
+    if n_batches < 2:
+        return 1.0, 0.0
+
+    global_freq = batch_counts / n_samples
+    k0 = min(k0, n_samples - 1)
+
+    nn = NearestNeighbors(n_neighbors=k0 + 1, algorithm='auto')
+    nn.fit(X)
+    _, indices = nn.kneighbors(X)
+
+    chi2_stats = []
+    p_values = []
+    batch_to_idx = {b: i for i, b in enumerate(unique_batches)}
+
+    for i in range(n_samples):
+        neighbors = indices[i, 1:k0+1]
+        neighbor_batches = batch_labels[neighbors]
+
+        observed = np.zeros(n_batches)
+        for nb in neighbor_batches:
+            observed[batch_to_idx[nb]] += 1
+
+        expected = global_freq * k0
+
+        mask = expected > 0
+        if mask.sum() < 2:
+            continue
+
+        stat = np.sum((observed[mask] - expected[mask])**2 / expected[mask])
+        df = max(1, mask.sum() - 1)
+        p_val = 1 - chi2.cdf(stat, df)
+
+        chi2_stats.append(stat)
+        p_values.append(p_val)
+
+    if len(p_values) == 0:
+        return 1.0, 0.0
+
+    acceptance_rate = np.mean(np.array(p_values) > alpha)
+    mean_stat = np.mean(chi2_stats)
+
+    return acceptance_rate, mean_stat
+
+
+def _find_sigma(distances: np.ndarray, target_perplexity: float, tol: float = 1e-5) -> float:
+    """
+    Binary search for sigma to achieve target perplexity.
+
+    Used in LISI computation.
+
+    Parameters
+    ----------
+    distances : np.ndarray
+        Distances to neighbors.
+    target_perplexity : float
+        Target perplexity value.
+    tol : float
+        Tolerance for convergence.
+
+    Returns
+    -------
+    float
+        Sigma value.
+    """
+    target_H = np.log2(target_perplexity + 1e-10)
+
+    sigma_min, sigma_max = 1e-10, 1e10
+    sigma = 1.0
+
+    for _ in range(50):
+        P = np.exp(-distances**2 / (2 * sigma**2 + 1e-10))
+        P_sum = P.sum()
+        if P_sum < 1e-10:
+            sigma = (sigma + sigma_max) / 2
+            continue
+        P = P / P_sum
+        P = np.clip(P, 1e-10, 1.0)
+        H = -np.sum(P * np.log2(P))
+
+        if abs(H - target_H) < tol:
+            break
+        elif H < target_H:
+            sigma_min = sigma
+        else:
+            sigma_max = sigma
+        sigma = (sigma_min + sigma_max) / 2
+
+    return sigma
+
+
+def _lisi_score(
+    X: np.ndarray,
+    batch_labels: np.ndarray,
+    perplexity: int = 30,
+) -> float:
+    """
+    Compute mean Local Inverse Simpson's Index (LISI).
+
+    Range: [1, n_batches], where n_batches = perfect mixing.
+    Higher = better batch mixing.
+
+    Reference: Korsunsky et al. (2019) Nature Methods (Harmony paper)
+
+    Parameters
+    ----------
+    X : np.ndarray
+        Data matrix.
+    batch_labels : np.ndarray
+        Batch labels for each sample.
+    perplexity : int
+        Perplexity for Gaussian kernel.
+
+    Returns
+    -------
+    float
+        Mean LISI score.
+    """
+    n_samples = X.shape[0]
+    unique_batches = np.unique(batch_labels)
+    n_batches = len(unique_batches)
+    batch_to_idx = {b: i for i, b in enumerate(unique_batches)}
+
+    if n_batches < 2:
+        return 1.0
+
+    k = min(3 * perplexity, n_samples - 1)
+
+    nn = NearestNeighbors(n_neighbors=k + 1, algorithm='auto')
+    nn.fit(X)
+    distances, indices = nn.kneighbors(X)
+
+    distances = distances[:, 1:]
+    indices = indices[:, 1:]
+
+    lisi_values = []
+
+    for i in range(n_samples):
+        sigma = _find_sigma(distances[i], perplexity)
+
+        P = np.exp(-distances[i]**2 / (2 * sigma**2 + 1e-10))
+        P_sum = P.sum()
+        if P_sum < 1e-10:
+            lisi_values.append(1.0)
+            continue
+        P = P / P_sum
+
+        neighbor_batches = batch_labels[indices[i]]
+        batch_probs = np.zeros(n_batches)
+        for j, nb in enumerate(neighbor_batches):
+            batch_probs[batch_to_idx[nb]] += P[j]
+
+        simpson = np.sum(batch_probs**2)
+        if simpson < 1e-10:
+            lisi = n_batches
+        else:
+            lisi = 1.0 / simpson
+        lisi_values.append(lisi)
+
+    return np.mean(lisi_values)
+
+
+def _variance_ratio(X: np.ndarray, batch_labels: np.ndarray) -> float:
+    """
+    Compute between-batch to within-batch variance ratio.
+
+    Similar to F-statistic in one-way ANOVA.
+    Lower ratio after correction = better batch effect removal.
+
+    Parameters
+    ----------
+    X : np.ndarray
+        Data matrix.
+    batch_labels : np.ndarray
+        Batch labels for each sample.
+
+    Returns
+    -------
+    float
+        Variance ratio (between/within).
+    """
+    unique_batches = np.unique(batch_labels)
+    n_batches = len(unique_batches)
+    n_samples = X.shape[0]
+
+    if n_batches < 2:
+        return 0.0
+
+    grand_mean = np.mean(X, axis=0)
+
+    between_var = 0.0
+    within_var = 0.0
+
+    for batch in unique_batches:
+        mask = batch_labels == batch
+        n_b = np.sum(mask)
+        X_batch = X[mask]
+        batch_mean = np.mean(X_batch, axis=0)
+
+        between_var += n_b * np.sum((batch_mean - grand_mean)**2)
+        within_var += np.sum((X_batch - batch_mean)**2)
+
+    between_var /= (n_batches - 1)
+    within_var /= (n_samples - n_batches)
+
+    if within_var < 1e-10:
+        return 0.0
+
+    return between_var / within_var
+
+
+def _knn_preservation(
+    X_before: np.ndarray,
+    X_after: np.ndarray,
+    k_values: List[int],
+    n_jobs: int = 1,
+) -> Dict[int, float]:
+    """
+    Compute fraction of k-nearest neighbors preserved after correction.
+
+    Range: [0, 1], where 1 = perfect preservation.
+    Higher = better biological structure preservation.
+
+    Parameters
+    ----------
+    X_before : np.ndarray
+        Original data.
+    X_after : np.ndarray
+        Corrected data.
+    k_values : list of int
+        Values of k for k-NN.
+    n_jobs : int
+        Number of parallel jobs.
+
+    Returns
+    -------
+    dict
+        Mapping from k to preservation fraction.
+    """
+    results = {}
+    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.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.fit(X_after)
+    _, indices_after = nn_after.kneighbors(X_after)
+
+    for k in k_values:
+        if k > max_k:
+            results[k] = 0.0
+            continue
+
+        overlaps = []
+        for i in range(X_before.shape[0]):
+            neighbors_before = set(indices_before[i, 1:k+1])
+            neighbors_after = set(indices_after[i, 1:k+1])
+            overlap = len(neighbors_before & neighbors_after) / k
+            overlaps.append(overlap)
+
+        results[k] = np.mean(overlaps)
+
+    return results
+
+
+def _pairwise_distance_correlation(
+    X_before: np.ndarray,
+    X_after: np.ndarray,
+    subsample: int = 1000,
+    random_state: int = 42,
+) -> float:
+    """
+    Compute Spearman correlation of pairwise distances.
+
+    Range: [-1, 1], where 1 = perfect rank preservation.
+    Higher = better relative relationship preservation.
+
+    Parameters
+    ----------
+    X_before : np.ndarray
+        Original data.
+    X_after : np.ndarray
+        Corrected data.
+    subsample : int
+        Maximum samples to use (for efficiency).
+    random_state : int
+        Random seed for subsampling.
+
+    Returns
+    -------
+    float
+        Spearman correlation coefficient.
+    """
+    n_samples = X_before.shape[0]
+
+    if n_samples > subsample:
+        rng = np.random.default_rng(random_state)
+        idx = rng.choice(n_samples, subsample, replace=False)
+        X_before = X_before[idx]
+        X_after = X_after[idx]
+
+    dist_before = pdist(X_before, metric='euclidean')
+    dist_after = pdist(X_after, metric='euclidean')
+
+    if len(dist_before) == 0:
+        return 1.0
+
+    corr, _ = spearmanr(dist_before, dist_after)
+
+    if np.isnan(corr):
+        return 1.0
+
+    return corr
+
+
+def _mean_centroid_distance(X: np.ndarray, batch_labels: np.ndarray) -> float:
+    """
+    Compute mean pairwise Euclidean distance between batch centroids.
+
+    Lower after correction = better batch alignment.
+
+    Parameters
+    ----------
+    X : np.ndarray
+        Data matrix.
+    batch_labels : np.ndarray
+        Batch labels for each sample.
+
+    Returns
+    -------
+    float
+        Mean pairwise distance between centroids.
+    """
+    unique_batches = np.unique(batch_labels)
+    n_batches = len(unique_batches)
+
+    if n_batches < 2:
+        return 0.0
+
+    centroids = []
+    for batch in unique_batches:
+        mask = batch_labels == batch
+        centroid = np.mean(X[mask], axis=0)
+        centroids.append(centroid)
+
+    centroids = np.array(centroids)
+    distances = pdist(centroids, metric='euclidean')
+
+    return np.mean(distances)
+
+
+def _levene_median_statistic(X: np.ndarray, batch_labels: np.ndarray) -> float:
+    """
+    Compute median Levene test statistic across features.
+
+    Lower statistic = more homogeneous variances across batches.
+
+    Parameters
+    ----------
+    X : np.ndarray
+        Data matrix.
+    batch_labels : np.ndarray
+        Batch labels for each sample.
+
+    Returns
+    -------
+    float
+        Median Levene test statistic.
+    """
+    unique_batches = np.unique(batch_labels)
+    if len(unique_batches) < 2:
+        return 0.0
+
+    levene_stats = []
+    for j in range(X.shape[1]):
+        groups = [X[batch_labels == b, j] for b in unique_batches]
+        groups = [g for g in groups if len(g) > 0]
+        if len(groups) < 2:
+            continue
+        try:
+            stat, _ = levene(*groups, center='median')
+            if not np.isnan(stat):
+                levene_stats.append(stat)
+        except Exception:
+            continue
+
+    if len(levene_stats) == 0:
+        return 0.0
+
+    return np.median(levene_stats)
 
 
 class ComBatModel:
@@ -58,8 +570,9 @@ class ComBatModel:
         ignoring the variance (`delta_star`).
     reference_batch : str, optional
         If specified, the batch level to use as reference.
-    covbat_cov_thresh : float, default=0.9
-        CovBat: cumulative explained variance threshold for PCA.
+    covbat_cov_thresh : float or int, default=0.9
+        CovBat: cumulative variance threshold (0, 1] to retain PCs, or
+        integer >= 1 specifying the number of components directly.
     eps : float, default=1e-8
         Numerical jitter to avoid division-by-zero.
     """
@@ -67,19 +580,19 @@ class ComBatModel:
     def __init__(
         self,
         *,
-        method: Literal["johnson", "fortin", "chen"] = "johnson", 
+        method: Literal["johnson", "fortin", "chen"] = "johnson",
         parametric: bool = True,
         mean_only: bool = False,
         reference_batch: Optional[str] = None,
         eps: float = 1e-8,
-        covbat_cov_thresh: float = 0.9,
+        covbat_cov_thresh: Union[float, int] = 0.9,
     ) -> None:
         self.method: str = method
         self.parametric: bool = parametric
         self.mean_only: bool = bool(mean_only)
         self.reference_batch: Optional[str] = reference_batch
         self.eps: float = float(eps)
-        self.covbat_cov_thresh: float = float(covbat_cov_thresh)
+        self.covbat_cov_thresh: Union[float, int] = covbat_cov_thresh
 
         self._batch_levels: pd.Index
         self._grand_mean: pd.Series
@@ -96,9 +609,16 @@ class ComBatModel:
         self._batch_levels_pc: pd.Index
         self._pc_gamma_star: FloatArray
         self._pc_delta_star: FloatArray
-        
-        if not (0.0 < self.covbat_cov_thresh <= 1.0):
-            raise ValueError("covbat_cov_thresh must be in (0, 1].")
+
+        # Validate covbat_cov_thresh
+        if isinstance(self.covbat_cov_thresh, float):
+            if not (0.0 < self.covbat_cov_thresh <= 1.0):
+                raise ValueError("covbat_cov_thresh must be in (0, 1] when float.")
+        elif isinstance(self.covbat_cov_thresh, int):
+            if self.covbat_cov_thresh < 1:
+                raise ValueError("covbat_cov_thresh must be >= 1 when int.")
+        else:
+            raise TypeError("covbat_cov_thresh must be float or int.")
 
     @staticmethod
     def _as_series(
@@ -336,8 +856,14 @@ class ComBatModel:
         X_meanvar_adj = self._transform_fortin(X, batch, disc, cont)
         X_centered = X_meanvar_adj - X_meanvar_adj.mean(axis=0)
         pca = PCA(svd_solver="full", whiten=False).fit(X_centered)
-        cumulative = np.cumsum(pca.explained_variance_ratio_)
-        n_pc = int(np.searchsorted(cumulative, self.covbat_cov_thresh) + 1)
+
+        # Determine number of components based on threshold type
+        if isinstance(self.covbat_cov_thresh, int):
+            n_pc = min(self.covbat_cov_thresh, len(pca.explained_variance_ratio_))
+        else:
+            cumulative = np.cumsum(pca.explained_variance_ratio_)
+            n_pc = int(np.searchsorted(cumulative, self.covbat_cov_thresh) + 1)
+
         self._covbat_pca = pca
         self._covbat_n_pc = n_pc
 
@@ -488,7 +1014,8 @@ class ComBatModel:
         continuous_covariates: Optional[ArrayLike] = None,
     ) -> pd.DataFrame:
         """Transform the data using fitted ComBat parameters."""
-        check_is_fitted(self, ["_gamma_star"])
+        if not hasattr(self, "_gamma_star"):
+            raise ValueError("This ComBatModel instance is not fitted yet. Call 'fit' before 'transform'.")
         if not isinstance(X, pd.DataFrame):
             X = pd.DataFrame(X)
         idx = X.index
@@ -600,7 +1127,7 @@ class ComBatModel:
         """Chen transform implementation."""
         X_meanvar_adj = self._transform_fortin(X, batch, disc, cont)
         X_centered = X_meanvar_adj - self._covbat_pca.mean_
-        scores = self._covbat_pca.transform(X_centered.values)
+        scores = self._covbat_pca.transform(X_centered)
         n_pc = self._covbat_n_pc
         scores_adj = scores.copy()
 
@@ -639,7 +1166,8 @@ class ComBat(BaseEstimator, TransformerMixin):
         mean_only: bool = False,
         reference_batch: Optional[str] = None,
         eps: float = 1e-8,
-        covbat_cov_thresh: float = 0.9,
+        covbat_cov_thresh: Union[float, int] = 0.9,
+        compute_metrics: bool = False,
     ) -> None:
         self.batch = batch
         self.discrete_covariates = discrete_covariates
@@ -650,6 +1178,7 @@ class ComBat(BaseEstimator, TransformerMixin):
         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,
@@ -707,6 +1236,221 @@ class ComBat(BaseEstimator, TransformerMixin):
             else:
                 return pd.DataFrame(obj, index=idx)
 
+    @property
+    def metrics_(self) -> Optional[Dict[str, Any]]:
+        """Return cached metrics from last fit_transform with compute_metrics=True.
+
+        Returns
+        -------
+        dict or None
+            Cached metrics dictionary, or None if no metrics have been computed.
+        """
+        return getattr(self, '_metrics_cache', None)
+
+    def compute_batch_metrics(
+        self,
+        X: ArrayLike,
+        batch: Optional[ArrayLike] = None,
+        *,
+        pca_components: Optional[int] = None,
+        k_neighbors: List[int] = [5, 10, 50],
+        kbet_k0: Optional[int] = None,
+        lisi_perplexity: int = 30,
+        n_jobs: int = 1,
+    ) -> Dict[str, Any]:
+        """
+        Compute batch effect metrics before and after ComBat correction.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Input data to evaluate.
+        batch : array-like of shape (n_samples,), optional
+            Batch labels. If None, uses the batch stored at construction.
+        pca_components : int, optional
+            Number of PCA components for dimensionality reduction before
+            computing metrics. If None (default), metrics are computed in
+            the original feature space. Must be less than min(n_samples, n_features).
+        k_neighbors : list of int, default=[5, 10, 50]
+            Values of k for k-NN preservation metric.
+        kbet_k0 : int, optional
+            Neighborhood size for kBET. Default is 10% of samples.
+        lisi_perplexity : int, default=30
+            Perplexity for LISI computation.
+        n_jobs : int, default=1
+            Number of parallel jobs for neighbor computations.
+
+        Returns
+        -------
+        metrics : dict
+            Dictionary with structure:
+            {
+                'batch_effect': {
+                    'silhouette': {'before': float, 'after': float},
+                    'davies_bouldin': {...},
+                    'kbet': {...},
+                    'lisi': {..., 'max_value': n_batches},
+                    'variance_ratio': {...},
+                },
+                'preservation': {
+                    'knn': {k: fraction for k in k_neighbors},
+                    'distance_correlation': float,
+                },
+                'alignment': {
+                    'centroid_distance': {...},
+                    'levene_statistic': {...},
+                },
+            }
+
+        Raises
+        ------
+        ValueError
+            If the model is not fitted or if pca_components is invalid.
+        """
+        if not hasattr(self._model, "_gamma_star"):
+            raise ValueError(
+                "This ComBat instance is not fitted yet. "
+                "Call 'fit' before 'compute_batch_metrics'."
+            )
+
+        if not isinstance(X, pd.DataFrame):
+            X = pd.DataFrame(X)
+
+        idx = X.index
+
+        if batch is None:
+            batch_vec = self._subset(self.batch, idx)
+        else:
+            if isinstance(batch, (pd.Series, pd.DataFrame)):
+                batch_vec = batch.loc[idx] if hasattr(batch, 'loc') else batch
+            elif isinstance(batch, np.ndarray):
+                batch_vec = pd.Series(batch, index=idx)
+            else:
+                batch_vec = pd.Series(batch, index=idx)
+
+        batch_labels = np.array(batch_vec)
+
+        X_before = X.values
+        X_after = self.transform(X).values
+
+        n_samples, n_features = X_before.shape
+        if kbet_k0 is None:
+            kbet_k0 = max(10, int(0.10 * n_samples))
+
+        # Validate and apply PCA if requested
+        if pca_components is not None:
+            max_components = min(n_samples, n_features)
+            if pca_components >= max_components:
+                raise ValueError(
+                    f"pca_components={pca_components} must be less than "
+                    f"min(n_samples, n_features)={max_components}."
+                )
+            X_before_pca, X_after_pca, _ = _compute_pca_embedding(
+                X_before, X_after, pca_components
+            )
+        else:
+            X_before_pca = X_before
+            X_after_pca = X_after
+
+        silhouette_before = _silhouette_batch(X_before_pca, batch_labels)
+        silhouette_after = _silhouette_batch(X_after_pca, batch_labels)
+
+        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)
+
+        lisi_before = _lisi_score(X_before_pca, batch_labels, lisi_perplexity)
+        lisi_after = _lisi_score(X_after_pca, batch_labels, lisi_perplexity)
+
+        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)
+        dist_corr = _pairwise_distance_correlation(X_before_pca, X_after_pca)
+
+        centroid_before = _mean_centroid_distance(X_before_pca, batch_labels)
+        centroid_after = _mean_centroid_distance(X_after_pca, batch_labels)
+
+        levene_before = _levene_median_statistic(X_before, batch_labels)
+        levene_after = _levene_median_statistic(X_after, batch_labels)
+
+        n_batches = len(np.unique(batch_labels))
+
+        metrics = {
+            'batch_effect': {
+                'silhouette': {
+                    'before': silhouette_before,
+                    'after': silhouette_after,
+                },
+                'davies_bouldin': {
+                    'before': db_before,
+                    'after': db_after,
+                },
+                'kbet': {
+                    'before': kbet_before,
+                    'after': kbet_after,
+                },
+                'lisi': {
+                    'before': lisi_before,
+                    'after': lisi_after,
+                    'max_value': n_batches,
+                },
+                'variance_ratio': {
+                    'before': var_ratio_before,
+                    'after': var_ratio_after,
+                },
+            },
+            'preservation': {
+                'knn': knn_results,
+                'distance_correlation': dist_corr,
+            },
+            'alignment': {
+                'centroid_distance': {
+                    'before': centroid_before,
+                    'after': centroid_after,
+                },
+                'levene_statistic': {
+                    'before': levene_before,
+                    'after': levene_after,
+                },
+            },
+        }
+
+        return metrics
+
+    def fit_transform(
+        self,
+        X: ArrayLike,
+        y: Optional[ArrayLike] = 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 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 plot_transformation(
             self,
             X: ArrayLike, *, 
@@ -759,7 +1503,8 @@ class ComBat(BaseEstimator, TransformerMixin):
             - `'original'`: embedding of original data
             - `'transformed'`: embedding of ComBat-transformed data
         """
-        check_is_fitted(self._model, ["_gamma_star"])
+        if not hasattr(self._model, "_gamma_star"):
+            raise ValueError("This ComBat instance is not fitted yet. Call 'fit' before 'plot_transformation'.")
 
         if n_components not in [2, 3]:
             raise ValueError(f"n_components must be 2 or 3, got {n_components}")
@@ -768,11 +1513,6 @@ class ComBat(BaseEstimator, TransformerMixin):
         if plot_type not in ['static', 'interactive']:
             raise ValueError(f"plot_type must be 'static' or 'interactive', got '{plot_type}'")
 
-        if reduction_method == 'umap' and not UMAP_AVAILABLE:
-            raise ImportError("UMAP is not installed. Install with: pip install umap-learn")
-        if plot_type == 'interactive' and not PLOTLY_AVAILABLE:
-            raise ImportError("Plotly is not installed. Install with: pip install plotly")
-
         if not isinstance(X, pd.DataFrame):
             X = pd.DataFrame(X)
 
@@ -797,8 +1537,8 @@ class ComBat(BaseEstimator, TransformerMixin):
         else:
             umap_params = {'random_state': 42}
             umap_params.update(reduction_kwargs)
-            reducer_orig = umap.UMAP(n_components=n_components, **reduction_kwargs)
-            reducer_trans = umap.UMAP(n_components=n_components, **reduction_kwargs)
+            reducer_orig = umap.UMAP(n_components=n_components, **umap_params)
+            reducer_trans = umap.UMAP(n_components=n_components, **umap_params)
 
         X_embedded_orig = reducer_orig.fit_transform(X_np)
         X_embedded_trans = reducer_trans.fit_transform(X_trans_np)
@@ -845,9 +1585,9 @@ class ComBat(BaseEstimator, TransformerMixin):
         n_batches = len(unique_batches)
 
         if n_batches <= 10:
-            colors = plt.cm.get_cmap(cmap)(np.linspace(0, 1, n_batches))
+            colors = matplotlib.colormaps.get_cmap(cmap)(np.linspace(0, 1, n_batches))
         else:
-            colors = plt.cm.get_cmap('tab20')(np.linspace(0, 1, n_batches))
+            colors = matplotlib.colormaps.get_cmap('tab20')(np.linspace(0, 1, n_batches))
 
         if n_components == 2:
             ax1 = plt.subplot(1, 2, 1)
@@ -956,7 +1696,7 @@ class ComBat(BaseEstimator, TransformerMixin):
         unique_batches = batch_labels.drop_duplicates()
 
         n_batches = len(unique_batches)
-        cmap_func = plt.cm.get_cmap(cmap)
+        cmap_func = matplotlib.colormaps.get_cmap(cmap)
         color_list = [mcolors.to_hex(cmap_func(i / max(n_batches - 1, 1))) for i in range(n_batches)]
 
         batch_to_color = dict(zip(unique_batches, color_list))

{combatlearn-0.2.2.dist-info → combatlearn-1.1.0.dist-info}/METADATA

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: combatlearn
-Version: 0.2.2
+Version: 1.1.0
 Summary: Batch-effect harmonization for machine learning frameworks.
 Author-email: Ettore Rocchi <ettoreroc@gmail.com>
-License-Expression: MIT
+License: MIT
 Keywords: machine-learning,harmonization,combat,preprocessing
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Science/Research
@@ -19,13 +19,23 @@ Requires-Dist: matplotlib>=3.4
 Requires-Dist: plotly>=5.0
 Requires-Dist: nbformat>=4.2
 Requires-Dist: umap-learn>=0.5
-Requires-Dist: pytest>=7
+Provides-Extra: dev
+Requires-Dist: pytest>=7; extra == "dev"
+Requires-Dist: pytest-cov>=4.0; extra == "dev"
+Requires-Dist: ruff>=0.1; extra == "dev"
+Requires-Dist: mypy>=1.0; extra == "dev"
+Provides-Extra: docs
+Requires-Dist: mkdocs>=1.5.0; extra == "docs"
+Requires-Dist: mkdocs-material>=9.0.0; extra == "docs"
+Requires-Dist: mkdocstrings[python]>=0.24.0; extra == "docs"
+Requires-Dist: pymdown-extensions>=10.0; extra == "docs"
 Dynamic: license-file
 
 # **combatlearn**
 
 [![Python versions](https://img.shields.io/badge/python-%3E%3D3.10-blue?logo=python)](https://www.python.org/)
 [![Test](https://github.com/EttoreRocchi/combatlearn/actions/workflows/test.yaml/badge.svg)](https://github.com/EttoreRocchi/combatlearn/actions/workflows/test.yaml)
+[![Documentation](https://readthedocs.org/projects/combatlearn/badge/?version=latest)](https://combatlearn.readthedocs.io)
 [![PyPI Downloads](https://static.pepy.tech/badge/combatlearn)](https://pepy.tech/projects/combatlearn)
 [![PyPI Version](https://img.shields.io/pypi/v/combatlearn?cacheSeconds=300)](https://pypi.org/project/combatlearn/)
 [![License](https://img.shields.io/github/license/EttoreRocchi/combatlearn)](https://github.com/EttoreRocchi/combatlearn/blob/main/LICENSE)
@@ -47,8 +57,21 @@ Dynamic: license-file
 pip install combatlearn
 ```
 
+## Documentation
+
+**Full documentation is available at [combatlearn.readthedocs.io](https://combatlearn.readthedocs.io)**
+
+The documentation includes:
+- [Installation Guide](https://combatlearn.readthedocs.io/en/latest/installation/)
+- [Quick Start Tutorial](https://combatlearn.readthedocs.io/en/latest/quickstart/)
+- [User Guide](https://combatlearn.readthedocs.io/en/latest/user-guide/overview/)
+- [API Reference](https://combatlearn.readthedocs.io/en/latest/api/)
+- [Examples](https://combatlearn.readthedocs.io/en/latest/examples/basic-usage/)
+
 ## Quick start
 
+For more details, see the [Quick Start Tutorial](https://combatlearn.readthedocs.io/en/latest/quickstart/).
+
 ```python
 import pandas as pd
 from sklearn.pipeline import Pipeline
@@ -97,7 +120,7 @@ For a full example of how to use **combatlearn** see the [notebook demo](https:/
 
 ## `ComBat` parameters
 
-The following section provides a detailed explanation of all parameters available in the scikit-learn-compatible `ComBat` class.
+The following section provides a detailed explanation of all parameters available in the scikit-learn-compatible `ComBat` class. For complete API documentation, see the [API Reference](https://combatlearn.readthedocs.io/en/latest/api/).
 
 ### Main Parameters
 
@@ -119,11 +142,17 @@ The following section provides a detailed explanation of all parameters availabl
 | `eps` | float | `1e-8` | Small jitter value added to variances to prevent divide-by-zero errors during standardization. |
 
 
-### Batch Effect Correction Visualization 
+### Batch Effect Correction Visualization
 
 The `plot_transformation` method allows to visualize the **ComBat** transformation effect using dimensionality reduction, showing the before/after comparison of data transformed by `ComBat` using PCA, t-SNE, or UMAP to reduce dimensions for visualization.
 
-For further details see the [notebook demo](https://github.com/EttoreRocchi/combatlearn/blob/main/docs/demo/combatlearn_demo.ipynb).
+For further details see the [Visualization Guide](https://combatlearn.readthedocs.io/en/latest/user-guide/visualization/) and the [notebook demo](https://github.com/EttoreRocchi/combatlearn/blob/main/docs/demo/combatlearn_demo.ipynb).
+
+### Batch Effect Metrics
+
+The `compute_batch_metrics` method provides quantitative assessment of batch correction quality. It computes metrics including Silhouette coefficient, Davies-Bouldin index, kBET, LISI, and variance ratio for batch effect quantification, as well as k-NN preservation and distance correlation for structure preservation.
+
+For further details see the [Metrics Guide](https://combatlearn.readthedocs.io/en/latest/user-guide/metrics/) and the [notebook demo](https://github.com/EttoreRocchi/combatlearn/blob/main/docs/demo/combatlearn_demo.ipynb).
 
 ## Contributing
 
@@ -146,8 +175,7 @@ We gratefully acknowledge:
 
 ## Citation
 
-If **combatlearn** is useful in your research, please cite the original
-papers:
+If **combatlearn** is useful in your research, please cite the original papers:
 
 - Johnson WE, Li C, Rabinovic A. Adjusting batch effects in microarray expression data using empirical Bayes methods. _Biostatistics_. 2007 Jan;8(1):118-27. doi: [10.1093/biostatistics/kxj037](https://doi.org/10.1093/biostatistics/kxj037)
 

combatlearn-1.1.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+combatlearn/__init__.py,sha256=L4sPJuJzLJIODAuSXdNQECVoFJXcmVss7SzoqX6MlYg,99
+combatlearn/combat.py,sha256=yLoppVuLBvqO-0a01xWQve-8ZEMkEICabcGPf5u1goI,59309
+combatlearn-1.1.0.dist-info/licenses/LICENSE,sha256=O34CBRTmdL59PxDYOa6nq1N0-2A9xyXGkBXKbsL1NeY,1070
+combatlearn-1.1.0.dist-info/METADATA,sha256=s4Y3G0ou1TnNZJpu1eRC7VN2lDVan_BqEm_qpVQf2lk,9446
+combatlearn-1.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+combatlearn-1.1.0.dist-info/top_level.txt,sha256=3cFQv4oj2sh_NKra45cPy8Go0v8W9x9-zkkUibqZCMk,12
+combatlearn-1.1.0.dist-info/RECORD,,

combatlearn-0.2.2.dist-info/RECORD

@@ -1,7 +0,0 @@
-combatlearn/__init__.py,sha256=qZK8xAUibzM9TQJ-xho1cjMYmTGkdWvpFRTXOokNvMY,98
-combatlearn/combat.py,sha256=pVauFEgZ7wiYRimGZe7ZhBWZN7sGQ67A3o_SrBUtoJ8,38126
-combatlearn-0.2.2.dist-info/licenses/LICENSE,sha256=O34CBRTmdL59PxDYOa6nq1N0-2A9xyXGkBXKbsL1NeY,1070
-combatlearn-0.2.2.dist-info/METADATA,sha256=CNm0pbXPVVWORk4pI97WS1DohjWOu7fB88JS1JZ-3-A,7491
-combatlearn-0.2.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-combatlearn-0.2.2.dist-info/top_level.txt,sha256=3cFQv4oj2sh_NKra45cPy8Go0v8W9x9-zkkUibqZCMk,12
-combatlearn-0.2.2.dist-info/RECORD,,