MaldiBatchKit 0.1.0__py3-none-any.whl

maldibatchkit/diagnostics/generic.py

@@ -0,0 +1,327 @@
+"""Generic batch-mixing metrics.
+
+All functions in this module take a feature matrix ``X`` and a vector of
+batch labels and return a single scalar (or an array keyed by k for
+kBET) summarising how much batch structure remains after correction.
+
+These metrics are intentionally batch-effect-specific: they say nothing
+about whether biological signal is preserved. Pair them with supervised
+metrics (AMR classifier AUROC, VME, etc.) in any real comparison.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+from scipy.stats import chi2
+from sklearn.metrics import silhouette_score
+from sklearn.neighbors import NearestNeighbors
+
+from .._utils import ArrayLike
+
+__all__ = [
+    "kbet",
+    "lisi",
+    "lisi_max",
+    "lisi_normalized",
+    "silhouette_batch",
+    "species_preservation",
+]
+
+
+def _to_ndarray(X: ArrayLike) -> npt.NDArray[Any]:
+    if hasattr(X, "to_numpy"):
+        return np.asarray(X.to_numpy())
+    return np.asarray(X)
+
+
+def silhouette_batch(
+    X: ArrayLike, batch: ArrayLike, *, metric: str = "euclidean"
+) -> float:
+    """Silhouette coefficient using batch labels as clusters.
+
+    Values close to 0 indicate good mixing; values close to 1 indicate
+    strong batch separation; values close to -1 indicate batches sitting
+    inside each others' clusters (typically an artefact).
+
+    Parameters
+    ----------
+    X : array-like of shape (n_samples, n_features)
+        Feature matrix.
+    batch : array-like of shape (n_samples,)
+        Batch labels.
+    metric : str, default='euclidean'
+        Distance metric forwarded to ``sklearn.metrics.silhouette_score``.
+
+    Returns
+    -------
+    float
+        Silhouette coefficient, or ``0.0`` if there is fewer than two
+        distinct batches (silhouette is undefined in that case).
+    """
+    Xa = _to_ndarray(X).astype(float)
+    batch_arr = np.asarray(batch)
+    if len(np.unique(batch_arr)) < 2:
+        return 0.0
+    return float(silhouette_score(Xa, batch_arr, metric=metric))
+
+
+def kbet(
+    X: ArrayLike,
+    batch: ArrayLike,
+    *,
+    k: int | None = None,
+    alpha: float = 0.05,
+) -> dict[str, float]:
+    """k-nearest-neighbours Batch Effect Test (kBET; Büttner et al. 2019).
+
+    For each sample we compute a chi-square statistic testing whether
+    its k-nearest-neighbour batch composition matches the global batch
+    frequencies. The reported statistics are the acceptance rate (the
+    fraction of samples whose p-value exceeds ``alpha`` - higher is
+    better) and the mean chi-square statistic (lower is better).
+
+    Parameters
+    ----------
+    X : array-like of shape (n_samples, n_features)
+        Feature matrix.
+    batch : array-like of shape (n_samples,)
+        Batch labels.
+    k : int, optional
+        Number of nearest neighbours. Defaults to
+        ``max(10, int(0.1 * n_samples))``.
+    alpha : float, default=0.05
+        Significance threshold used to compute the acceptance rate.
+
+    Returns
+    -------
+    dict
+        ``{"acceptance_rate": float, "mean_chi2": float, "k": int}``.
+    """
+    Xa = _to_ndarray(X).astype(float)
+    batch_arr = np.asarray(batch)
+    n_samples = Xa.shape[0]
+    levels, counts = np.unique(batch_arr, return_counts=True)
+    n_batches = len(levels)
+    if n_batches < 2:
+        return {"acceptance_rate": 1.0, "mean_chi2": 0.0, "k": int(k or 0)}
+    if k is None:
+        k = max(10, int(0.1 * n_samples))
+    k = min(k, n_samples - 1)
+    global_freq = counts / counts.sum()
+    level_to_idx = {lvl: i for i, lvl in enumerate(levels)}
+
+    nn = NearestNeighbors(n_neighbors=k + 1).fit(Xa)
+    _, neigh_idx = nn.kneighbors(Xa)
+    neigh_idx = neigh_idx[:, 1:]
+
+    p_values = np.empty(n_samples)
+    chi2_stats = np.empty(n_samples)
+    for s in range(n_samples):
+        local_counts = np.zeros(n_batches)
+        for j in neigh_idx[s]:
+            local_counts[level_to_idx[batch_arr[j]]] += 1
+        expected = k * global_freq
+        chi2_stat = np.sum(
+            (local_counts - expected) ** 2 / np.clip(expected, 1e-9, None)
+        )
+        chi2_stats[s] = chi2_stat
+        p_values[s] = float(1.0 - chi2.cdf(chi2_stat, df=n_batches - 1))
+
+    return {
+        "acceptance_rate": float(np.mean(p_values > alpha)),
+        "mean_chi2": float(chi2_stats.mean()),
+        "k": int(k),
+    }
+
+
+def lisi(
+    X: ArrayLike,
+    batch: ArrayLike,
+    *,
+    perplexity: float = 30.0,
+) -> float:
+    """Local Inverse Simpson's Index for batch mixing.
+
+    LISI is the effective number of batches represented in each sample's
+    local neighbourhood (Gaussian-kernel weighted to the requested
+    perplexity). The returned value is the median LISI across samples.
+    It lies in ``[1, n_batches]``; values close to ``n_batches`` indicate
+    strong mixing, values close to 1 indicate batch-segregated
+    neighbourhoods.
+
+    Parameters
+    ----------
+    X : array-like of shape (n_samples, n_features)
+        Feature matrix.
+    batch : array-like of shape (n_samples,)
+        Batch labels.
+    perplexity : float, default=30.0
+        Target perplexity of the Gaussian kernel - matches the
+        convention used in Korsunsky et al. (2019) and combatlearn.
+
+    Returns
+    -------
+    float
+        Median LISI across samples.
+    """
+    Xa = _to_ndarray(X).astype(float)
+    batch_arr = np.asarray(batch)
+    n_samples = Xa.shape[0]
+    levels = np.unique(batch_arr)
+    n_batches = len(levels)
+    if n_batches < 2:
+        return 1.0
+    level_to_idx = {lvl: i for i, lvl in enumerate(levels)}
+    k = int(min(max(3 * perplexity, 10), n_samples - 1))
+    nn = NearestNeighbors(n_neighbors=k + 1).fit(Xa)
+    distances, neigh_idx = nn.kneighbors(Xa)
+    distances = distances[:, 1:]
+    neigh_idx = neigh_idx[:, 1:]
+
+    target = np.log(perplexity)
+    lisi_values = np.empty(n_samples)
+    for s in range(n_samples):
+        d = distances[s]
+        # Binary search for beta matching target perplexity
+        lo, hi = 1e-10, 1e10
+        beta = 1.0
+        for _ in range(50):
+            weights = np.exp(-d * beta)
+            Z = max(weights.sum(), 1e-12)
+            h = np.log(Z) + beta * np.sum(d * weights) / Z
+            if h > target:
+                lo = beta
+                beta = beta * 2 if hi == 1e10 else (beta + hi) / 2
+            else:
+                hi = beta
+                beta = (beta + lo) / 2
+            if abs(h - target) < 1e-5:
+                break
+        weights = np.exp(-d * beta)
+        weights /= max(weights.sum(), 1e-12)
+        probs = np.zeros(n_batches)
+        for w, j in zip(weights, neigh_idx[s], strict=True):
+            probs[level_to_idx[batch_arr[j]]] += w
+        simpson = float(np.sum(probs**2))
+        lisi_values[s] = 1.0 / max(simpson, 1e-12)
+    return float(np.median(lisi_values))
+
+
+def lisi_max(labels: ArrayLike) -> float:
+    r"""Theoretical maximum LISI under perfect mixing.
+
+    LISI is the inverse Simpson index of the local label distribution
+    in a sample's neighbourhood, so under *perfect* mixing (every
+    neighbourhood matches the global label frequencies) it equals the
+    inverse Simpson index of the global frequencies:
+
+    .. math::
+
+        \text{lisi}_{\max} = \frac{1}{\sum_i p_i^2}
+
+    where :math:`p_i` is the global fraction of label :math:`i`. For
+    perfectly balanced designs this equals the number of unique labels;
+    for **imbalanced** designs (e.g. DRIAMS-A dominating in a multi-site
+    cohort) the ceiling is strictly less than ``n_levels``, which is
+    why raw LISI numbers can mislead readers who expect them to
+    approach the level count.
+
+    Parameters
+    ----------
+    labels : array-like of shape (n_samples,)
+        Label vector (e.g. batch or species).
+
+    Returns
+    -------
+    float
+        Inverse Simpson index of the global label frequencies, in the
+        range ``[1, n_levels]``. Returns ``1.0`` when the input has
+        zero or one unique level (LISI is degenerate in that case).
+    """
+    arr = np.asarray(labels)
+    if arr.size == 0:
+        return 1.0
+    _, counts = np.unique(arr, return_counts=True)
+    if len(counts) < 2:
+        return 1.0
+    p = counts / counts.sum()
+    return float(1.0 / float((p**2).sum()))
+
+
+def lisi_normalized(
+    X: ArrayLike, labels: ArrayLike, *, perplexity: float = 30.0
+) -> float:
+    """LISI divided by its theoretical maximum.
+
+    Reports how close the local mixing is to *perfect* mixing given
+    the global label frequencies, on a ``[0, 1]`` scale where higher
+    means better mixed. This is the right way to compare LISI across
+    cohorts with different label-imbalance profiles -- raw LISI is
+    bounded by ``lisi_max(labels)`` (the inverse Simpson index of the
+    global label frequencies), not by the number of unique labels.
+
+    Parameters
+    ----------
+    X : array-like of shape (n_samples, n_features)
+        Feature matrix.
+    labels : array-like of shape (n_samples,)
+        Label vector (typically batch).
+    perplexity : float, default=30.0
+        Forwarded to :func:`lisi`.
+
+    Returns
+    -------
+    float
+        ``lisi(X, labels) / lisi_max(labels)``, or ``nan`` when there
+        are fewer than two unique labels (normalisation is undefined).
+    """
+    mx = lisi_max(labels)
+    if mx <= 1.0:
+        return float("nan")
+    return float(lisi(X, labels, perplexity=perplexity) / mx)
+
+
+def species_preservation(
+    X: ArrayLike, species: ArrayLike, *, perplexity: float = 30.0
+) -> float:
+    r"""Score on ``[0, 1]`` where 1 = species perfectly preserved, 0 = blurred.
+
+    The ideal post-correction LISI for *species* (or any biological
+    grouping you want to keep apart) is :math:`1` -- each
+    neighbourhood is dominated by one species. The *worst* value under
+    perfect mixing is :math:`\text{lisi}_{\max}(\text{species})`. We
+    rescale the gap to ``[0, 1]``:
+
+    .. math::
+
+        \text{species\_preservation}
+            = \frac{\text{lisi}_{\max} - \text{lisi}(X, \text{species})}
+                   {\text{lisi}_{\max} - 1}
+
+    so the metric reads naturally as "fraction of biological structure
+    retained" and is invariant to species-imbalance (unlike raw LISI).
+
+    Parameters
+    ----------
+    X : array-like of shape (n_samples, n_features)
+        Feature matrix.
+    species : array-like of shape (n_samples,)
+        Species (or other biology) labels.
+    perplexity : float, default=30.0
+        Forwarded to :func:`lisi`.
+
+    Returns
+    -------
+    float
+        Preservation score in ``[0, 1]`` (``1.0`` = perfectly preserved,
+        ``0.0`` = blurred at the imbalance ceiling). Returns ``nan``
+        when fewer than two unique species are present.
+    """
+    mx = lisi_max(species)
+    if mx <= 1.0:
+        return float("nan")
+    raw = lisi(X, species, perplexity=perplexity)
+    return float((mx - raw) / (mx - 1.0))

maldibatchkit/diagnostics/maldi.py

@@ -0,0 +1,165 @@
+"""MALDI-specific batch diagnostics.
+
+These functions operate directly on binned-intensity matrices (rows =
+samples, columns = m/z bins) and produce tidy per-batch summaries.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
+
+from .._utils import ArrayLike, _resolve_mz_axis
+
+__all__ = [
+    "peak_position_drift",
+    "per_batch_spectrum_count",
+    "tic_cov_per_batch",
+]
+
+
+def _to_dataframe(X: ArrayLike) -> pd.DataFrame:
+    if isinstance(X, pd.DataFrame):
+        return X
+    arr = np.asarray(X)
+    return pd.DataFrame(arr)
+
+
+def per_batch_spectrum_count(batch: ArrayLike) -> pd.Series:
+    """Return the number of spectra per batch as a sorted Series."""
+    s = pd.Series(np.asarray(batch)).value_counts().sort_index()
+    s.name = "n_spectra"
+    s.index.name = "batch"
+    return s
+
+
+def tic_cov_per_batch(X: ArrayLike, batch: ArrayLike) -> pd.Series:
+    """Per-batch coefficient of variation of the Total Ion Count.
+
+    High CoV within a batch suggests substantial within-batch intensity
+    fluctuation, which in MALDI-TOF tends to be driven by acquisition
+    issues rather than biology.
+
+    Parameters
+    ----------
+    X : array-like of shape (n_samples, n_features)
+        Binned intensities.
+    batch : array-like of shape (n_samples,)
+        Batch labels.
+
+    Returns
+    -------
+    pd.Series
+        CoV (``std/mean``) of TIC for each batch, indexed by batch level.
+    """
+    df = _to_dataframe(X)
+    tic = df.sum(axis=1).to_numpy()
+    b = np.asarray(batch)
+    out: dict[Any, float] = {}
+    for lvl in np.unique(b):
+        vals = tic[b == lvl]
+        mean = float(vals.mean())
+        std = float(vals.std(ddof=0))
+        out[lvl] = std / mean if mean > 0 else float("nan")
+    s = pd.Series(out).sort_index()
+    s.name = "tic_cov"
+    s.index.name = "batch"
+    return s
+
+
+def peak_position_drift(
+    X: ArrayLike,
+    batch: ArrayLike,
+    *,
+    mz_values: ArrayLike | None = None,
+    top_k: int = 50,
+) -> pd.DataFrame:
+    """Per-batch peak-position drift relative to a global reference.
+
+    For each batch we compute its median spectrum, identify the top
+    ``top_k`` peaks of the **global** median spectrum, and locate the
+    nearest local maximum of the per-batch median to each of those
+    global peaks. The returned table summarises the distribution of
+    (signed) position shifts per batch.
+
+    Parameters
+    ----------
+    X : array-like of shape (n_samples, n_features)
+        Binned intensities.
+    batch : array-like of shape (n_samples,)
+        Batch labels.
+    mz_values : array-like of shape (n_features,), optional
+        m/z values for the columns of ``X``. If None, integer column
+        positions are used, and the returned ``mean_delta_mz`` column is
+        in bin units.
+    top_k : int, default=50
+        Number of global peaks to track.
+
+    Returns
+    -------
+    pd.DataFrame
+        One row per batch, with columns
+        ``mean_delta_mz``, ``median_delta_mz``, ``max_abs_delta_mz``.
+    """
+    df = _to_dataframe(X).astype(float)
+    b = np.asarray(batch)
+    mz, _ = _resolve_mz_axis(df, mz_values)
+
+    global_median = df.median(axis=0).to_numpy()
+    # Local maxima of the global median
+    peaks_idx = _local_maxima(global_median)
+    if peaks_idx.size == 0:
+        # Degenerate case: flat / monotonic spectrum
+        return pd.DataFrame(
+            {
+                "mean_delta_mz": [],
+                "median_delta_mz": [],
+                "max_abs_delta_mz": [],
+            },
+        )
+    peak_intensities = global_median[peaks_idx]
+    order = np.argsort(-peak_intensities)[:top_k]
+    tracked = peaks_idx[order]
+
+    records = []
+    for lvl in np.unique(b):
+        batch_median = df.loc[b == lvl].median(axis=0).to_numpy()
+        batch_peaks = _local_maxima(batch_median)
+        if batch_peaks.size == 0:
+            records.append(
+                {
+                    "batch": lvl,
+                    "mean_delta_mz": float("nan"),
+                    "median_delta_mz": float("nan"),
+                    "max_abs_delta_mz": float("nan"),
+                }
+            )
+            continue
+        deltas = []
+        for p in tracked:
+            nearest = batch_peaks[int(np.argmin(np.abs(batch_peaks - p)))]
+            deltas.append(mz[nearest] - mz[p])
+        deltas_arr = np.asarray(deltas)
+        records.append(
+            {
+                "batch": lvl,
+                "mean_delta_mz": float(deltas_arr.mean()),
+                "median_delta_mz": float(np.median(deltas_arr)),
+                "max_abs_delta_mz": float(np.max(np.abs(deltas_arr))),
+            }
+        )
+    out = pd.DataFrame(records).set_index("batch")
+    out = out.sort_index()
+    return out
+
+
+def _local_maxima(x: npt.NDArray[np.float64]) -> npt.NDArray[np.int64]:
+    """Return indices of strict local maxima of a 1-D array."""
+    if x.size < 3:
+        return np.asarray([], dtype=np.int64)
+    greater_prev = x[1:-1] > x[:-2]
+    greater_next = x[1:-1] > x[2:]
+    return np.where(greater_prev & greater_next)[0] + 1

maldibatchkit/diagnostics/report.py

@@ -0,0 +1,151 @@
+"""Combined before/after diagnostic report."""
+
+from __future__ import annotations
+
+import pandas as pd
+
+from .._utils import ArrayLike
+from .generic import kbet, lisi, silhouette_batch
+from .maldi import peak_position_drift, per_batch_spectrum_count, tic_cov_per_batch
+
+__all__ = ["diagnostic_report"]
+
+
+def diagnostic_report(
+    before: ArrayLike,
+    after: ArrayLike,
+    batch: ArrayLike,
+    *,
+    mz_values: ArrayLike | None = None,
+    k: int | None = None,
+    lisi_perplexity: float = 30.0,
+    top_k_peaks: int = 50,
+) -> pd.DataFrame:
+    """Run every diagnostic on a (before, after) pair.
+
+    Parameters
+    ----------
+    before : array-like of shape (n_samples, n_features)
+        Feature matrix prior to batch correction.
+    after : array-like of shape (n_samples, n_features)
+        Feature matrix after batch correction. Must have the same shape
+        as ``before``.
+    batch : array-like of shape (n_samples,)
+        Batch labels.
+    mz_values : array-like, optional
+        m/z positions for the feature columns (passed to
+        :func:`peak_position_drift`).
+    k : int, optional
+        Neighbourhood size for kBET.
+    lisi_perplexity : float, default=30.0
+        Perplexity for LISI.
+    top_k_peaks : int, default=50
+        Number of peaks tracked for drift.
+
+    Returns
+    -------
+    pd.DataFrame
+        Tidy report with columns ``metric``, ``scope``, ``value_before``,
+        ``value_after`` (and ``delta`` where both columns are numeric and
+        the metric's improvement direction is well-defined).
+    """
+    rows: list[dict] = []
+
+    # Generic batch mixing (smaller = better for silhouette / chi2)
+    sil_b = silhouette_batch(before, batch)
+    sil_a = silhouette_batch(after, batch)
+    rows.append(
+        {
+            "metric": "silhouette_batch",
+            "scope": "overall",
+            "value_before": sil_b,
+            "value_after": sil_a,
+            "delta": sil_a - sil_b,
+            "better": "lower",
+        }
+    )
+
+    kbet_b = kbet(before, batch, k=k)
+    kbet_a = kbet(after, batch, k=k)
+    rows.append(
+        {
+            "metric": "kbet_acceptance_rate",
+            "scope": "overall",
+            "value_before": kbet_b["acceptance_rate"],
+            "value_after": kbet_a["acceptance_rate"],
+            "delta": kbet_a["acceptance_rate"] - kbet_b["acceptance_rate"],
+            "better": "higher",
+        }
+    )
+    rows.append(
+        {
+            "metric": "kbet_mean_chi2",
+            "scope": "overall",
+            "value_before": kbet_b["mean_chi2"],
+            "value_after": kbet_a["mean_chi2"],
+            "delta": kbet_a["mean_chi2"] - kbet_b["mean_chi2"],
+            "better": "lower",
+        }
+    )
+
+    lisi_b = lisi(before, batch, perplexity=lisi_perplexity)
+    lisi_a = lisi(after, batch, perplexity=lisi_perplexity)
+    rows.append(
+        {
+            "metric": "lisi",
+            "scope": "overall",
+            "value_before": lisi_b,
+            "value_after": lisi_a,
+            "delta": lisi_a - lisi_b,
+            "better": "higher",
+        }
+    )
+
+    # TIC CoV per batch
+    tic_b = tic_cov_per_batch(before, batch)
+    tic_a = tic_cov_per_batch(after, batch)
+    for lvl in tic_b.index:
+        rows.append(
+            {
+                "metric": "tic_cov",
+                "scope": str(lvl),
+                "value_before": float(tic_b.loc[lvl]),
+                "value_after": float(tic_a.loc[lvl]),
+                "delta": float(tic_a.loc[lvl] - tic_b.loc[lvl]),
+                "better": "lower",
+            }
+        )
+
+    # Peak-position drift per batch (mean |delta_mz|)
+    drift_b = peak_position_drift(before, batch, mz_values=mz_values, top_k=top_k_peaks)
+    drift_a = peak_position_drift(after, batch, mz_values=mz_values, top_k=top_k_peaks)
+    for lvl in drift_b.index:
+        rows.append(
+            {
+                "metric": "peak_drift_mean",
+                "scope": str(lvl),
+                "value_before": float(drift_b.loc[lvl, "mean_delta_mz"]),
+                "value_after": float(drift_a.loc[lvl, "mean_delta_mz"]),
+                "delta": float(
+                    drift_a.loc[lvl, "mean_delta_mz"]
+                    - drift_b.loc[lvl, "mean_delta_mz"]
+                ),
+                "better": "zero",
+            }
+        )
+
+    # Spectrum counts (unchanged by correction; still useful in reports)
+    counts = per_batch_spectrum_count(batch)
+    for lvl, n in counts.items():
+        rows.append(
+            {
+                "metric": "n_spectra",
+                "scope": str(lvl),
+                "value_before": int(n),
+                "value_after": int(n),
+                "delta": 0,
+                "better": "n/a",
+            }
+        )
+
+    return pd.DataFrame(rows)

maldibatchkit/integrations/__init__.py

@@ -0,0 +1,9 @@
+"""Integrations with external MALDI toolkits.
+
+Currently exposes :class:`MaldiSetAdapter`, a bridge to
+:class:`maldiamrkit.MaldiSet`.
+"""
+
+from .maldiset import MaldiSetAdapter
+
+__all__ = ["MaldiSetAdapter"]