MaldiBatchKit 0.1.0__py3-none-any.whl

maldibatchkit/integrations/maldiset.py

@@ -0,0 +1,184 @@
+"""``MaldiSet`` <-> MaldiBatchKit bridge."""
+
+from __future__ import annotations
+
+import inspect
+from copy import copy
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import pandas as pd
+
+if TYPE_CHECKING:  # pragma: no cover - type-checking imports only
+    from maldiamrkit import MaldiSet
+
+__all__ = ["MaldiSetAdapter"]
+
+
+def _require_maldiset():
+    try:
+        from maldiamrkit import MaldiSet
+    except ImportError as exc:
+        raise ImportError(
+            "MaldiSetAdapter requires `maldiamrkit` (normally installed "
+            "alongside MaldiBatchKit). Reinstall with "
+            "`pip install -U maldibatchkit` or `pip install maldiamrkit`."
+        ) from exc
+    return MaldiSet
+
+
+class MaldiSetAdapter:
+    """Bridge :class:`maldiamrkit.MaldiSet` and MaldiBatchKit correctors.
+
+    The adapter does three things:
+
+    1. Extracts the feature matrix ``ds.X`` and any metadata columns
+       (batch / species / quality) from a ``MaldiSet``.
+    2. Delegates to a MaldiBatchKit corrector (or any sklearn-compatible
+       transformer with the right constructor signature) to produce a
+       corrected matrix.
+    3. Returns a **new** ``MaldiSet`` whose ``X`` property yields the
+       corrected matrix, leaving the original dataset untouched.
+
+    Batch / covariate slicing follows the rest of the package: pass the
+    metadata column names at construction time; the adapter pulls the
+    aligned series from ``ds.meta`` itself so users do not have to
+    rebuild arrays manually.
+
+    Parameters
+    ----------
+    batch_column : str
+        Column in ``ds.meta`` with the batch labels.
+    species_column : str, optional
+        Column in ``ds.meta`` with species labels (used when the chosen
+        corrector needs a species covariate).
+    quality_column : str, optional
+        Column in ``ds.meta`` with per-sample quality scores (used by
+        :class:`maldibatchkit.QualityWeightedComBat`).
+
+    Examples
+    --------
+    >>> from maldiamrkit import MaldiSet
+    >>> from maldibatchkit.integrations import MaldiSetAdapter
+    >>> from maldibatchkit import SpeciesAwareComBat
+    >>> adapter = MaldiSetAdapter(batch_column="Batch", species_column="Species")
+    >>> corrected_ds = adapter.correct(ds, SpeciesAwareComBat)
+    >>> corrected_ds.X.head()   # corrected feature matrix
+    """
+
+    def __init__(
+        self,
+        *,
+        batch_column: str,
+        species_column: str | None = None,
+        quality_column: str | None = None,
+    ) -> None:
+        self.batch_column = batch_column
+        self.species_column = species_column
+        self.quality_column = quality_column
+
+    def extract(self, ds: MaldiSet) -> dict[str, Any]:
+        """Pull ``X``, batch, species, quality from a ``MaldiSet``.
+
+        Returns
+        -------
+        dict
+            Dictionary with keys ``X`` (DataFrame), ``batch``,
+            ``species`` (or ``None``), ``quality`` (or ``None``). The
+            series are aligned to ``X.index``.
+        """
+        _require_maldiset()
+        X = ds.X
+        idx = X.index
+        meta = ds.meta.loc[idx]
+        if self.batch_column not in meta.columns:
+            raise KeyError(
+                f"batch_column={self.batch_column!r} not found in MaldiSet metadata."
+            )
+        batch = meta[self.batch_column]
+        species = (
+            meta[self.species_column]
+            if self.species_column and self.species_column in meta.columns
+            else None
+        )
+        quality = (
+            meta[self.quality_column]
+            if self.quality_column and self.quality_column in meta.columns
+            else None
+        )
+        return {"X": X, "batch": batch, "species": species, "quality": quality}
+
+    def correct(
+        self,
+        ds: MaldiSet,
+        transformer_cls,
+        *,
+        transformer_kwargs: dict[str, Any] | None = None,
+    ) -> MaldiSet:
+        """Run ``transformer_cls(batch=..., ...)`` and return a new MaldiSet.
+
+        Parameters
+        ----------
+        ds : MaldiSet
+            Source dataset.
+        transformer_cls : type
+            A MaldiBatchKit corrector class (or any transformer whose
+            constructor takes ``batch=`` and, optionally, a
+            species-style protected-covariate argument named
+            ``species`` / ``discrete_covariates`` / ``design`` /
+            ``covariates`` and/or a ``quality=`` argument).
+        transformer_kwargs : dict, optional
+            Extra keyword arguments forwarded to ``transformer_cls``.
+
+        Returns
+        -------
+        MaldiSet
+            Shallow-copied dataset with its ``_X_cache`` replaced by the
+            corrected feature matrix. The ``spectra`` list is the same
+            object in the returned ``MaldiSet``; labels / metadata are
+            unchanged.
+        """
+        _require_maldiset()
+        extracted = self.extract(ds)
+        X = extracted["X"]
+        batch = extracted["batch"]
+        species = extracted["species"]
+        quality = extracted["quality"]
+
+        kwargs: dict[str, Any] = {"batch": batch}
+        try:
+            init_params = set(inspect.signature(transformer_cls.__init__).parameters)
+        except (TypeError, ValueError):  # pragma: no cover - exotic subclasses
+            init_params = set()
+
+        # Route species metadata to a *categorical* protected-covariate
+        # slot, in decreasing order of specificity: dedicated
+        # ``species=``, then ComBat / Fortin's ``discrete_covariates=``,
+        # then Harmony's ``covariates=``.  We intentionally do NOT auto-
+        # route to Limma's ``design=`` (that slot is a numeric design
+        # matrix, not a label); Limma users who want to protect species
+        # should pass e.g. ``transformer_kwargs={"design": species_dummies}``.
+        if species is not None:
+            for name in ("species", "discrete_covariates", "covariates"):
+                if name in init_params:
+                    kwargs[name] = species
+                    break
+        if quality is not None and "quality" in init_params:
+            kwargs["quality"] = quality
+        if transformer_kwargs:
+            kwargs.update(transformer_kwargs)
+
+        transformer = transformer_cls(**kwargs)
+        X_corrected = transformer.fit_transform(X)
+        if not isinstance(X_corrected, pd.DataFrame):
+            X_corrected = pd.DataFrame(
+                np.asarray(X_corrected), index=X.index, columns=X.columns
+            )
+
+        new_ds = copy(ds)
+        # MaldiSet lazily rebuilds X from spectra unless _X_cache is set;
+        # inject the corrected matrix there so every downstream property
+        # (``.y``, ``.filter()``...) keeps working on the corrected data.
+        new_ds._X_cache = X_corrected
+        new_ds.meta = ds.meta.loc[X_corrected.index].copy()
+        return new_ds

maldibatchkit/viz/__init__.py

@@ -0,0 +1,11 @@
+"""Visualization helpers for batch-effect inspection."""
+
+from .peaks import plot_peak_shift
+from .summary import plot_diagnostic_summary
+from .umap import plot_batch_umap
+
+__all__ = [
+    "plot_batch_umap",
+    "plot_diagnostic_summary",
+    "plot_peak_shift",
+]

maldibatchkit/viz/peaks.py

@@ -0,0 +1,82 @@
+"""Visualize per-batch peak position drift."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+
+from .._utils import ArrayLike, _resolve_mz_axis
+
+__all__ = ["plot_peak_shift"]
+
+
+def _to_dataframe(X: ArrayLike) -> pd.DataFrame:
+    if isinstance(X, pd.DataFrame):
+        return X
+    return pd.DataFrame(np.asarray(X))
+
+
+def plot_peak_shift(
+    batches: ArrayLike,
+    X: ArrayLike,
+    reference: ArrayLike | None = None,
+    *,
+    mz_values: ArrayLike | None = None,
+    ax: Any | None = None,
+    max_batches: int | None = 6,
+) -> tuple[Any, Any]:
+    """Overlay per-batch median spectra against a reference spectrum.
+
+    Parameters
+    ----------
+    batches : array-like of shape (n_samples,)
+        Batch labels.
+    X : array-like of shape (n_samples, n_features)
+        Binned intensities.
+    reference : array-like of shape (n_features,), optional
+        Reference spectrum. If None, the global median across all rows
+        is used.
+    mz_values : array-like of shape (n_features,), optional
+        m/z coordinates for the x-axis. Defaults to column positions.
+    ax : matplotlib Axes, optional
+        Axis to draw on. If None, a new figure is created.
+    max_batches : int or None, default=6
+        Maximum number of batches to overlay. Oldest ties broken
+        alphabetically. Set to None to draw every batch (slow on
+        many-batch studies).
+
+    Returns
+    -------
+    fig, ax : matplotlib Figure and Axes
+    """
+    df = _to_dataframe(X).astype(float)
+    b = np.asarray(batches)
+    mz, mz_is_real = _resolve_mz_axis(df, mz_values)
+
+    if reference is None:
+        ref = df.median(axis=0).to_numpy()
+    else:
+        ref = np.asarray(reference, dtype=float).ravel()
+
+    levels = sorted(set(b.tolist()))
+    if max_batches is not None:
+        levels = levels[:max_batches]
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(10, 4))
+    else:
+        fig = ax.figure
+
+    ax.plot(mz, ref, color="black", lw=1.3, label="reference", zorder=5)
+    cmap = plt.get_cmap("tab10")
+    for i, lvl in enumerate(levels):
+        med = df.loc[b == lvl].median(axis=0).to_numpy()
+        ax.plot(mz, med, lw=0.9, color=cmap(i % 10), label=str(lvl), alpha=0.85)
+    ax.set_xlabel("m/z" if mz_is_real else "bin index")
+    ax.set_ylabel("intensity (median per batch)")
+    ax.legend(title="batch", fontsize=7, loc="best")
+    ax.set_title("Per-batch peak-shape comparison")
+    return fig, ax

maldibatchkit/viz/summary.py

@@ -0,0 +1,155 @@
+"""Grid bar plot of before/after diagnostic metrics.
+
+Different diagnostic metrics live on very different scales
+(silhouette_batch ~ 0, LISI ~ 1-n_batches, kBET mean chi-square in the
+tens to thousands...), so sharing a single y-axis makes the smaller
+metrics unreadable.  We render one subplot per metric, each with its
+own y-axis.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Iterable
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+
+__all__ = ["plot_diagnostic_summary"]
+
+
+_BAR_BEFORE = "#b0b0b0"
+_BAR_AFTER = "#2b8cbe"
+
+
+def plot_diagnostic_summary(
+    report_df: pd.DataFrame,
+    *,
+    scope: str | Iterable[str] = "overall",
+    ncols: int | None = None,
+    figsize_per_plot: tuple[float, float] = (3.2, 3.0),
+    axes: Any | None = None,
+) -> tuple[Any, np.ndarray]:
+    """Plot before/after diagnostic values, one subplot per metric.
+
+    Parameters
+    ----------
+    report_df : pd.DataFrame
+        Output of :func:`maldibatchkit.diagnostics.diagnostic_report`.
+    scope : str or iterable of str, default='overall'
+        Which slice(s) of the report to plot. Pass a single scope name
+        (``"overall"``, ``"batch_00"``, ...) to render one pair of bars
+        per metric, or a list (``["batch_00", "batch_01"]``) to group
+        bars by scope inside each metric's subplot.
+    ncols : int, optional
+        Number of subplot columns. Defaults to ``min(n_metrics, 4)``.
+    figsize_per_plot : (float, float), default=(3.2, 3.0)
+        Width, height of each metric's subplot (inches). The returned
+        figure has size ``(ncols * w, nrows * h)``.
+    axes : array of matplotlib Axes, optional
+        Pre-built axes grid. Must have at least ``n_metrics`` entries
+        when flattened.
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+    axes : np.ndarray of matplotlib.axes.Axes
+        Flattened array of the axes actually used. Unused slots from a
+        non-rectangular grid are turned off.
+    """
+    required = {"metric", "scope", "value_before", "value_after"}
+    missing = required - set(report_df.columns)
+    if missing:
+        raise ValueError(f"report_df is missing required columns: {sorted(missing)}.")
+
+    scopes = [scope] if isinstance(scope, str) else list(scope)
+    if not scopes:
+        raise ValueError("`scope` must name at least one scope.")
+
+    df = report_df[report_df["scope"].isin(scopes)].copy()
+    if df.empty:
+        raise ValueError(
+            f"No rows in report_df with scope in {scopes!r}. "
+            f"Available scopes: {sorted(report_df['scope'].unique())}."
+        )
+
+    # Preserve the metric order as it appears in the report
+    metrics: list[str] = []
+    for m in df["metric"].tolist():
+        if m not in metrics:
+            metrics.append(m)
+
+    n_metrics = len(metrics)
+    if ncols is None:
+        ncols = min(n_metrics, 4)
+    ncols = max(1, ncols)
+    nrows = int(np.ceil(n_metrics / ncols))
+
+    if axes is None:
+        fig, ax_grid = plt.subplots(
+            nrows,
+            ncols,
+            figsize=(ncols * figsize_per_plot[0], nrows * figsize_per_plot[1]),
+        )
+        ax_flat = np.atleast_1d(ax_grid).ravel()
+    else:
+        ax_flat = np.atleast_1d(np.asarray(axes)).ravel()
+        if ax_flat.size < n_metrics:
+            raise ValueError(
+                f"Provided axes has {ax_flat.size} slots but {n_metrics} "
+                f"metrics need plotting."
+            )
+        fig = ax_flat[0].figure
+
+    used_axes = []
+    for i, metric in enumerate(metrics):
+        ax = ax_flat[i]
+        sub = df[df["metric"] == metric]
+        scope_order = [s for s in scopes if s in sub["scope"].unique()]
+        if not scope_order:
+            ax.set_axis_off()
+            continue
+        sub = sub.set_index("scope").loc[scope_order]
+        n = len(scope_order)
+        xs = np.arange(n)
+        width = 0.38
+        ax.bar(
+            xs - width / 2,
+            sub["value_before"].astype(float).to_numpy(),
+            width,
+            label="before",
+            color=_BAR_BEFORE,
+        )
+        ax.bar(
+            xs + width / 2,
+            sub["value_after"].astype(float).to_numpy(),
+            width,
+            label="after",
+            color=_BAR_AFTER,
+        )
+        ax.set_xticks(xs)
+        if n > 1:
+            ax.set_xticklabels(scope_order, rotation=30, ha="right")
+        else:
+            ax.set_xticks([])
+        ax.set_title(metric, fontsize=10)
+        ax.grid(axis="y", alpha=0.3)
+        # Hairline guide at 0 for signed metrics
+        if sub[["value_before", "value_after"]].astype(float).min().min() < 0:
+            ax.axhline(0.0, color="k", lw=0.5)
+        if i == 0:
+            ax.legend(fontsize=8, loc="best")
+        used_axes.append(ax)
+
+    # Turn off any leftover subplot slots
+    for j in range(n_metrics, ax_flat.size):
+        ax_flat[j].set_axis_off()
+
+    title_suffix = scopes[0] if len(scopes) == 1 else f"{len(scopes)} scopes"
+    fig.suptitle(
+        f"Batch-effect diagnostics ({title_suffix})",
+        fontsize=12,
+        y=1.02,
+    )
+    fig.tight_layout()
+    return fig, ax_flat[:n_metrics]

maldibatchkit/viz/umap.py

@@ -0,0 +1,118 @@
+"""Side-by-side UMAP before/after batch correction."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+from sklearn.decomposition import PCA
+
+from .._utils import ArrayLike
+
+__all__ = ["plot_batch_umap"]
+
+
+def _require_umap():
+    try:
+        import umap  # noqa: F401
+    except ImportError as exc:
+        raise ImportError(
+            "plot_batch_umap requires `umap-learn`. Install the MaldiBatchKit "
+            "viz extras with `pip install maldibatchkit[viz]`."
+        ) from exc
+    return umap
+
+
+def _embed(X: np.ndarray, random_state: int) -> np.ndarray:
+    umap_mod = _require_umap()
+    reducer = umap_mod.UMAP(n_components=2, random_state=random_state)
+    return np.asarray(reducer.fit_transform(X))
+
+
+def _to_ndarray(X: ArrayLike) -> np.ndarray:
+    if hasattr(X, "to_numpy"):
+        return np.asarray(X.to_numpy())
+    return np.asarray(X)
+
+
+def plot_batch_umap(
+    before: ArrayLike,
+    after: ArrayLike,
+    batch: ArrayLike,
+    *,
+    color_by: str = "batch",
+    species: ArrayLike | None = None,
+    random_state: int = 42,
+    pca_preprocess: int | None = 50,
+    ax: tuple[Any, Any] | None = None,
+) -> tuple[Any, Any]:
+    """Plot UMAP embeddings of ``X`` before and after batch correction.
+
+    Parameters
+    ----------
+    before : array-like of shape (n_samples, n_features)
+        Feature matrix prior to correction.
+    after : array-like of shape (n_samples, n_features)
+        Feature matrix after correction. Must share ``before``'s shape.
+    batch : array-like of shape (n_samples,)
+        Batch labels.
+    color_by : {'batch', 'species'}, default='batch'
+        Which label drives colouring. ``'species'`` requires passing
+        ``species`` explicitly.
+    species : array-like of shape (n_samples,), optional
+        Species labels. Required when ``color_by='species'``.
+    random_state : int, default=42
+        Seed for UMAP.
+    pca_preprocess : int or None, default=50
+        If not None, reduce to this many PCs before UMAP. Gives a large
+        speed-up on typical MALDI-TOF matrices without hurting the plot
+        qualitatively.
+    ax : tuple of matplotlib Axes, optional
+        Two axes to draw on. If None, a new 1x2 figure is created.
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        The figure used (or the parent figure of the provided axes).
+    axes : tuple of matplotlib.axes.Axes
+        The two axes that were drawn on.
+    """
+    before_arr = _to_ndarray(before).astype(float)
+    after_arr = _to_ndarray(after).astype(float)
+    if before_arr.shape != after_arr.shape:
+        raise ValueError(
+            f"before and after must share shape; got {before_arr.shape} vs {after_arr.shape}."
+        )
+    if color_by == "species" and species is None:
+        raise ValueError("color_by='species' requires the `species` argument.")
+    labels = np.asarray(batch if color_by == "batch" else species)
+
+    def _maybe_pca(X: np.ndarray) -> np.ndarray:
+        if pca_preprocess is None:
+            return X
+        n = min(pca_preprocess, X.shape[0], X.shape[1])
+        if n < 2:
+            return X
+        return np.asarray(
+            PCA(n_components=n, random_state=random_state).fit_transform(X)
+        )
+
+    b_emb = _embed(_maybe_pca(before_arr), random_state)
+    a_emb = _embed(_maybe_pca(after_arr), random_state)
+
+    if ax is None:
+        fig, axes = plt.subplots(1, 2, figsize=(12, 5), sharex=False, sharey=False)
+    else:
+        axes = ax
+        fig = axes[0].figure
+
+    for title, emb, ax_ in (("Before", b_emb, axes[0]), ("After", a_emb, axes[1])):
+        for lvl in np.unique(labels):
+            mask = labels == lvl
+            ax_.scatter(emb[mask, 0], emb[mask, 1], s=8, label=str(lvl), alpha=0.7)
+        ax_.set_title(f"UMAP - {title}")
+        ax_.set_xlabel("UMAP1")
+        ax_.set_ylabel("UMAP2")
+    axes[0].legend(title=color_by, loc="best", fontsize=7)
+    return fig, tuple(axes) if isinstance(axes, (list, np.ndarray)) else axes