MaldiBatchKit 0.1.0__py3-none-any.whl

maldibatchkit/corrections/maldi.py

@@ -0,0 +1,167 @@
+"""MALDI-specific batch corrections that leverage ``maldiamrkit`` internals."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
+
+from .._base import BaseBatchCorrector
+from .._utils import ArrayLike
+
+__all__ = ["BatchAwareWarping"]
+
+
+def _try_import_maldiamrkit():
+    try:
+        import maldiamrkit  # noqa: F401
+    except ImportError as exc:
+        raise ImportError(
+            "This transformer requires 'maldiamrkit' (normally installed "
+            "alongside MaldiBatchKit). Reinstall with "
+            "`pip install -U maldibatchkit` or "
+            "`pip install maldiamrkit>=0.13.0`."
+        ) from exc
+    return maldiamrkit
+
+
+class BatchAwareWarping(BaseBatchCorrector):
+    r"""Per-batch m/z warping to a single global reference.
+
+    This transformer reuses :class:`maldiamrkit.alignment.Warping` under
+    the hood, but fits one warper **per batch**. At ``fit`` time we
+    learn a global reference spectrum from all training rows, then fit a
+    dedicated ``Warping`` instance per batch with that shared reference.
+    At ``transform`` time each sample is routed through the warper of
+    its batch.
+
+    Why this matters
+    ----------------
+    Mass-calibration drift is usually a per-acquisition-session
+    phenomenon. Running a single global ``Warping`` can either be too
+    aggressive (squashing real biological peaks in well-calibrated
+    batches) or too conservative (leaving large shifts in a miscalibrated
+    batch). Per-batch warpers share the reference but learn a
+    batch-specific transformation.
+
+    Parameters
+    ----------
+    batch : array-like of shape (n_samples,)
+        Batch labels.
+    reference : str, default="median"
+        How to build the global reference.
+        ``"median"`` uses the median spectrum across all training rows.
+        ``"batch_mean"`` uses the mean of per-batch medians (reduces
+        bias towards large batches).
+    method : str, default="shift"
+        Warping method passed through to
+        ``maldiamrkit.alignment.Warping``.
+    n_segments : int, default=5
+        Piecewise warping segments.
+    max_shift : int, default=50
+        Maximum shift in bins.
+    dtw_radius : int, default=10
+        DTW radius constraint.
+    smooth_sigma : float, default=2.0
+        Gaussian smoothing for piecewise shifts.
+    n_jobs : int, default=1
+        Parallel jobs forwarded to ``Warping``.
+
+    Attributes
+    ----------
+    reference\_ : np.ndarray
+        The global reference spectrum.
+    warpers\_ : dict[Any, Warping]
+        One fitted ``Warping`` per batch level.
+
+    Notes
+    -----
+    For batches that do not appear at ``fit`` time but show up during
+    ``transform``, we fall back to a global warper trained on all
+    training rows (stored as :attr:`warpers_` under the ``"__global__"``
+    key). This keeps the transformer defined on test folds without
+    peeking at test data.
+
+    Examples
+    --------
+    >>> from maldibatchkit import BatchAwareWarping
+    >>> warper = BatchAwareWarping(batch=batches, method="piecewise")
+    >>> X_aligned = warper.fit_transform(X)
+    """
+
+    def __init__(
+        self,
+        batch: ArrayLike,
+        *,
+        reference: str = "median",
+        method: str = "shift",
+        n_segments: int = 5,
+        max_shift: int = 50,
+        dtw_radius: int = 10,
+        smooth_sigma: float = 2.0,
+        n_jobs: int = 1,
+    ) -> None:
+        super().__init__(batch=batch)
+        self.reference = reference
+        self.method = method
+        self.n_segments = n_segments
+        self.max_shift = max_shift
+        self.dtw_radius = dtw_radius
+        self.smooth_sigma = smooth_sigma
+        self.n_jobs = n_jobs
+
+    def _make_warper(self):
+        _try_import_maldiamrkit()
+        from maldiamrkit.alignment import Warping  # type: ignore[import-not-found]
+
+        return Warping(
+            reference="median",
+            method=self.method,
+            n_segments=self.n_segments,
+            max_shift=self.max_shift,
+            dtw_radius=self.dtw_radius,
+            smooth_sigma=self.smooth_sigma,
+            n_jobs=self.n_jobs,
+        )
+
+    def _build_reference(
+        self, X_df: pd.DataFrame, batch: npt.NDArray[Any]
+    ) -> np.ndarray:
+        if self.reference == "median":
+            return X_df.median(axis=0).to_numpy()
+        if self.reference == "batch_mean":
+            per_batch = []
+            for lvl in np.unique(batch):
+                per_batch.append(X_df.loc[batch == lvl].median(axis=0).to_numpy())
+            return np.mean(per_batch, axis=0)
+        raise ValueError(
+            f"Unknown reference={self.reference!r}; expected 'median' or 'batch_mean'."
+        )
+
+    def _fit_warper_on(self, df: pd.DataFrame, reference: np.ndarray):
+        w = self._make_warper()
+        w.fit(df)
+        w.ref_spec_ = reference  # swap in the shared global reference
+        return w
+
+    def _fit_impl(self, X_df: pd.DataFrame, batch: npt.NDArray[Any]) -> None:
+        self.reference_ = self._build_reference(X_df, batch)
+        warpers: dict[Any, Any] = {}
+        for lvl in np.unique(batch):
+            df = X_df.loc[batch == lvl]
+            warpers[lvl] = self._fit_warper_on(df, self.reference_)
+        warpers["__global__"] = self._fit_warper_on(X_df, self.reference_)
+        self.warpers_ = warpers
+
+    def _transform_impl(
+        self, X_df: pd.DataFrame, batch: npt.NDArray[Any]
+    ) -> pd.DataFrame:
+        out = X_df.copy().astype(float)
+        for lvl in np.unique(batch):
+            mask = batch == lvl
+            chunk = X_df.loc[mask]
+            warper = self.warpers_.get(lvl, self.warpers_["__global__"])
+            out.loc[mask] = warper.transform(chunk).to_numpy()
+        return out

maldibatchkit/corrections/quality_weighted.py

@@ -0,0 +1,338 @@
+r"""Quality-weighted ComBat.
+
+This module implements a *weighted* empirical-Bayes variant of the
+classical Johnson-et-al. (2007) ComBat algorithm. Samples with higher
+acquisition quality (typically a signal-to-noise ratio, or any
+user-supplied scalar score) contribute proportionally more to the
+batch-level moment estimates, and the shrinkage prior on each batch is
+scaled by the effective (weighted) batch size rather than the raw
+sample count.
+
+Mathematical formulation
+------------------------
+For feature :math:`g` in batch :math:`i`, let :math:`w_{ij} \\ge 0` be the
+per-sample weight (weights are normalised so the weight vector sums to
+the number of samples). Write :math:`\\bar w_i = \\sum_j w_{ij}` for the
+effective batch size.
+
+*Grand mean and pooled variance* are estimated with weights:
+
+.. math::
+
+    \\hat\\alpha_g &= \\frac{\\sum_{i,j} w_{ij} X_{ij,g}}{\\sum_{i,j} w_{ij}}
+
+    \\hat\\sigma_g^2 &= \\frac{\\sum_{i,j} w_{ij} (X_{ij,g} - \\hat\\alpha_g)^2}
+                              {\\sum_{i,j} w_{ij}}
+
+*Per-batch L/S model.* After standardisation
+:math:`Z_{ij,g} = (X_{ij,g} - \\hat\\alpha_g) / \\hat\\sigma_g`, the batch-
+and-feature effects are estimated with weighted sums:
+
+.. math::
+
+    \\hat\\gamma_{i,g} &= \\frac{\\sum_j w_{ij} Z_{ij,g}}{\\bar w_i}
+
+    \\hat\\delta_{i,g}^2 &= \\frac{\\sum_j w_{ij} (Z_{ij,g} - \\hat\\gamma_{i,g})^2}
+                                {\\bar w_i - 1}
+
+(``max(bar_w_i - 1, eps)`` is used to protect against tiny effective
+batch sizes.)
+
+*Empirical-Bayes priors.* The Normal-Inverse-Gamma hyperparameters
+are fit by method-of-moments on the per-feature
+:math:`\\hat\\gamma_{i,g}` (mean, variance) and per-feature
+:math:`\\hat\\delta_{i,g}^2` (inverse-gamma :math:`\\alpha, \\beta`),
+exactly as in classical ComBat.
+
+*Posteriors.* The posterior mean and variance are then the usual
+weighted combination of prior and data, but with the **effective** batch
+size :math:`\\bar w_i` replacing :math:`n_i`:
+
+.. math::
+
+    \\gamma^*_{i,g} &= \\frac{\\bar w_i \\bar\\tau_i^2 \\hat\\gamma_{i,g}
+                               + \\delta^*_{i,g} \\bar\\gamma_i}
+                             {\\bar w_i \\bar\\tau_i^2 + \\delta^*_{i,g}}
+
+    \\delta^{*2}_{i,g} &= \\frac{\\bar\\theta_i + 0.5 \\sum_j w_{ij}
+                                 (Z_{ij,g} - \\gamma^*_{i,g})^2}
+                               {\\bar w_i / 2 + \\bar\\lambda_i - 1}
+
+with :math:`\\bar\\gamma_i, \\bar\\tau_i^2, \\bar\\lambda_i, \\bar\\theta_i`
+the prior hyperparameters fit as in Johnson et al. (2007). The
+iterative scheme follows Johnson's original fixed-point updates; only
+the summation weights differ.
+
+Rationale
+---------
+In MALDI-TOF AMR studies, not every spectrum is equally informative:
+low-SNR acquisitions or spectra flagged by the lab QC are often kept in
+the analysis to avoid discarding hard-won samples, but they should not
+drive the per-batch correction. Down-weighting low-quality spectra in
+the empirical-Bayes shrinkage keeps the posterior pulled towards the
+estimate obtained from the high-quality spectra.
+
+This corrector is **not** a replacement for more general covariate
+modelling: like classical Johnson ComBat it does not preserve
+biological signal encoded in a covariate. If you need that, either
+residualise your covariate first or use
+:class:`maldibatchkit.ComBat` with ``method='fortin'``.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
+
+from .._base import BaseBatchCorrector
+from .._utils import ArrayLike, _subset
+
+__all__ = ["QualityWeightedComBat"]
+
+
+def _normalise_weights(
+    w: npt.NDArray[np.float64], *, eps: float
+) -> npt.NDArray[np.float64]:
+    w = np.asarray(w, dtype=float)
+    if np.any(w < 0):
+        raise ValueError("quality weights must be non-negative.")
+    if not np.any(w > 0):
+        raise ValueError("at least one quality weight must be positive.")
+    # Normalise so weights sum to the number of samples: equal weights
+    # should recover the unweighted estimates exactly.
+    w = w * (w.size / max(w.sum(), eps))
+    return w
+
+
+def _aprior(gamma_hat: np.ndarray) -> np.ndarray:
+    """Method-of-moments hyperparameter for the gamma prior (per batch)."""
+    m = gamma_hat.mean(axis=1)
+    s2 = gamma_hat.var(axis=1, ddof=1)
+    return (2 * s2 + m**2) / np.clip(s2, 1e-12, None)
+
+
+def _bprior(gamma_hat: np.ndarray) -> np.ndarray:
+    m = gamma_hat.mean(axis=1)
+    s2 = gamma_hat.var(axis=1, ddof=1)
+    return (m * s2 + m**3) / np.clip(s2, 1e-12, None)
+
+
+def _postmean(
+    g_hat: np.ndarray,
+    g_bar: float,
+    n_bar: float,
+    t2: float,
+    d_star: np.ndarray,
+) -> np.ndarray:
+    return (n_bar * t2 * g_hat + d_star * g_bar) / (n_bar * t2 + d_star)
+
+
+def _postvar(
+    sum_sq: np.ndarray,
+    n_bar: float,
+    a: float,
+    b: float,
+) -> np.ndarray:
+    return (0.5 * sum_sq + b) / (n_bar / 2.0 + a - 1.0)
+
+
+class QualityWeightedComBat(BaseBatchCorrector):
+    r"""Weighted empirical-Bayes extension of Johnson-ComBat.
+
+    Parameters
+    ----------
+    batch : array-like of shape (n_samples,)
+        Batch labels.
+    quality : array-like of shape (n_samples,)
+        Non-negative per-sample quality scores (typically SNR). Higher
+        values mean a sample should contribute more to per-batch
+        estimation. Internally rescaled to sum to ``n_samples``.
+    parametric : bool, default=True
+        Use the parametric EB fixed-point iteration (as in Johnson
+        2007). ``False`` falls back to the data estimates
+        :math:`\\hat\\gamma` / :math:`\\hat\\delta^2` directly - this is
+        useful as a sanity check.
+    reference_batch : Any, optional
+        Batch level to leave unchanged.
+    eps : float, default=1e-8
+        Numerical jitter.
+    max_iter : int, default=50
+        Hard cap on the parametric fixed-point iterations.
+    tol : float, default=1e-4
+        Convergence tolerance on the max absolute change in
+        :math:`\\gamma^*, \\delta^*`.
+
+    Attributes
+    ----------
+    batch_levels_ : np.ndarray
+        Batch levels observed at fit time.
+    grand_mean_ : np.ndarray
+        Weighted feature-wise grand mean (:math:`\\hat\\alpha`).
+    pooled_var_ : np.ndarray
+        Weighted feature-wise pooled variance (:math:`\\hat\\sigma^2`).
+    gamma_star_ : np.ndarray
+        Posterior batch means, shape ``(n_batches, n_features)``.
+    delta_star_ : np.ndarray
+        Posterior batch variances, shape ``(n_batches, n_features)``.
+    effective_batch_sizes_ : np.ndarray
+        :math:`\\bar w_i` per batch (weighted sample counts).
+    n_iter_ : int
+        Iterations actually taken by the parametric EB loop.
+
+    Examples
+    --------
+    >>> from maldibatchkit import QualityWeightedComBat
+    >>> corrector = QualityWeightedComBat(batch=batches, quality=snr)
+    >>> X_corrected = corrector.fit_transform(X)
+
+    References
+    ----------
+    Johnson, W.E., Li, C. and Rabinovic, A. (2007). "Adjusting batch
+    effects in microarray expression data using empirical Bayes
+    methods." *Biostatistics* 8(1): 118-127. - the unweighted
+    predecessor of the scheme implemented here.
+    """
+
+    def __init__(
+        self,
+        batch: ArrayLike,
+        *,
+        quality: ArrayLike,
+        parametric: bool = True,
+        reference_batch: Any | None = None,
+        eps: float = 1e-8,
+        max_iter: int = 50,
+        tol: float = 1e-4,
+    ) -> None:
+        super().__init__(batch=batch)
+        self.quality = quality
+        self.parametric = parametric
+        self.reference_batch = reference_batch
+        self.eps = eps
+        self.max_iter = max_iter
+        self.tol = tol
+
+    def _weights_at(self, idx: pd.Index) -> np.ndarray:
+        w = _subset(self.quality, idx)
+        if isinstance(w, pd.DataFrame):
+            if w.shape[1] != 1:
+                raise ValueError("quality must be 1-D.")
+            w = w.iloc[:, 0]
+        if w is None:
+            raise ValueError("quality weights must be provided.")
+        w = np.asarray(w, dtype=float)
+        return _normalise_weights(w, eps=self.eps)
+
+    def _fit_impl(self, X_df: pd.DataFrame, batch: npt.NDArray[Any]) -> None:
+        X = X_df.to_numpy(dtype=float)
+        n_samples, n_features = X.shape
+
+        w = self._weights_at(X_df.index)  # shape (n,)
+        levels = np.unique(batch)
+        self.batch_levels_ = levels
+
+        if self.reference_batch is not None and self.reference_batch not in levels:
+            raise ValueError(
+                f"reference_batch={self.reference_batch!r} not found among "
+                f"training batch levels {list(levels)}."
+            )
+        self.reference_batch_ = self.reference_batch
+
+        W = w.sum()
+        self.grand_mean_ = (w[:, None] * X).sum(axis=0) / max(W, self.eps)
+        centred = X - self.grand_mean_
+        self.pooled_var_ = (w[:, None] * centred**2).sum(axis=0) / max(W, self.eps)
+        std = np.sqrt(self.pooled_var_ + self.eps)
+        Z = centred / std
+
+        n_batches = len(levels)
+        eff = np.zeros(n_batches)
+        gamma_hat = np.zeros((n_batches, n_features))
+        delta_hat2 = np.zeros((n_batches, n_features))
+
+        for i, lvl in enumerate(levels):
+            mask = batch == lvl
+            w_i = w[mask]
+            w_bar = float(w_i.sum())
+            eff[i] = w_bar
+            if w_bar < self.eps:
+                raise ValueError(
+                    f"Effective batch size for {lvl!r} is ~0 "
+                    f"(all quality weights are zero in this batch)."
+                )
+            g_hat = (w_i[:, None] * Z[mask]).sum(axis=0) / w_bar
+            gamma_hat[i] = g_hat
+            resid = Z[mask] - g_hat
+            denom = max(w_bar - 1.0, self.eps)
+            delta_hat2[i] = (w_i[:, None] * resid**2).sum(axis=0) / denom
+
+        self.effective_batch_sizes_ = eff
+
+        gamma_bar = gamma_hat.mean(axis=1)
+        t2 = gamma_hat.var(axis=1, ddof=1)
+        a_prior = _aprior(delta_hat2)
+        b_prior = _bprior(delta_hat2)
+
+        if not self.parametric:
+            gamma_star = gamma_hat
+            delta_star = delta_hat2
+            self.n_iter_ = 0
+        else:
+            gamma_star = gamma_hat.copy()
+            delta_star = delta_hat2.copy()
+            n_iter = 0
+            for _ in range(self.max_iter):
+                n_iter += 1
+                prev_g = gamma_star.copy()
+                prev_d = delta_star.copy()
+                for i, lvl in enumerate(levels):
+                    mask = batch == lvl
+                    w_i = w[mask]
+                    w_bar = eff[i]
+                    g_new = _postmean(
+                        gamma_hat[i], gamma_bar[i], w_bar, t2[i], delta_star[i]
+                    )
+                    resid = Z[mask] - g_new
+                    sum_sq = (w_i[:, None] * resid**2).sum(axis=0)
+                    d_new = _postvar(sum_sq, w_bar, a_prior[i], b_prior[i])
+                    gamma_star[i] = g_new
+                    delta_star[i] = np.clip(d_new, self.eps, None)
+                change = max(
+                    np.max(np.abs(gamma_star - prev_g)),
+                    np.max(np.abs(delta_star - prev_d)),
+                )
+                if change < self.tol:
+                    break
+            self.n_iter_ = n_iter
+
+        self.gamma_star_ = gamma_star
+        self.delta_star_ = delta_star
+
+    def _transform_impl(
+        self, X_df: pd.DataFrame, batch: npt.NDArray[Any]
+    ) -> pd.DataFrame:
+        X = X_df.to_numpy(dtype=float)
+        std = np.sqrt(self.pooled_var_ + self.eps)
+        Z = (X - self.grand_mean_) / std
+        out = np.empty_like(X)
+        level_to_idx = {lvl: i for i, lvl in enumerate(self.batch_levels_)}
+        for lvl in np.unique(batch):
+            mask = batch == lvl
+            if lvl == self.reference_batch_:
+                out[mask] = X[mask]
+                continue
+            if lvl in level_to_idx:
+                i = level_to_idx[lvl]
+                gamma = self.gamma_star_[i]
+                delta = self.delta_star_[i]
+            else:
+                # Unknown batch at transform time: correct using zero
+                # per-batch effect (identity in standardised space).
+                gamma = np.zeros(X.shape[1])
+                delta = np.ones(X.shape[1])
+            z_corr = (Z[mask] - gamma) / np.sqrt(np.clip(delta, self.eps, None))
+            out[mask] = z_corr * std + self.grand_mean_
+        return pd.DataFrame(out, index=X_df.index, columns=X_df.columns)

maldibatchkit/diagnostics/__init__.py

@@ -0,0 +1,39 @@
+"""Batch-effect diagnostics.
+
+This subpackage exposes two families of metrics plus a convenience
+"report" helper:
+
+- :mod:`maldibatchkit.diagnostics.generic` - classic batch-mixing
+  metrics (k-nearest-neighbour batch test / kBET, local inverse
+  Simpson index / LISI, silhouette coefficient by batch).
+- :mod:`maldibatchkit.diagnostics.maldi` - MALDI-specific summaries
+  (per-batch peak-position drift, per-batch TIC coefficient of
+  variation, per-batch spectrum count).
+- :func:`maldibatchkit.diagnostics.diagnostic_report` - run every
+  metric on a ``(before, after)`` pair and summarise into a tidy
+  DataFrame.
+"""
+
+from .generic import (
+    kbet,
+    lisi,
+    lisi_max,
+    lisi_normalized,
+    silhouette_batch,
+    species_preservation,
+)
+from .maldi import peak_position_drift, per_batch_spectrum_count, tic_cov_per_batch
+from .report import diagnostic_report
+
+__all__ = [
+    "diagnostic_report",
+    "kbet",
+    "lisi",
+    "lisi_max",
+    "lisi_normalized",
+    "peak_position_drift",
+    "per_batch_spectrum_count",
+    "silhouette_batch",
+    "species_preservation",
+    "tic_cov_per_batch",
+]