MaldiBatchKit 0.1.0__py3-none-any.whl

maldibatchkit/__init__.py

@@ -0,0 +1,56 @@
+"""MaldiBatchKit - batch-effect correction for MALDI-TOF AMR workflows.
+
+Subpackage guide
+----------------
+- ``maldibatchkit.corrections``  - ComBat variants, Limma, Harmony,
+  simple baselines, and MALDI-specific corrections
+  (``BatchAwareWarping``, ``QualityWeightedComBat``,
+  ``SpeciesAwareComBat``).
+- ``maldibatchkit.diagnostics``  - generic batch-mixing metrics (kBET, LISI,
+  silhouette) and MALDI-specific metrics (peak drift, per-batch TIC CoV,
+  spectrum counts), plus a ``diagnostic_report`` helper.
+- ``maldibatchkit.viz``          - ``plot_batch_umap``, ``plot_peak_shift``,
+  ``plot_diagnostic_summary``.
+- ``maldibatchkit.integrations`` - ``MaldiSetAdapter`` for seamless use with
+  ``maldiamrkit.MaldiSet``.
+
+Examples
+--------
+>>> from maldibatchkit import ComBat, QualityWeightedComBat
+>>> from maldibatchkit.diagnostics import diagnostic_report
+>>> from maldibatchkit.integrations import MaldiSetAdapter
+
+Extending MaldiBatchKit: subclass :class:`BaseBatchCorrector` and implement
+``_fit_impl`` / ``_transform_impl`` to ship a custom corrector that plugs
+into sklearn pipelines without re-solving index alignment or leakage.
+"""
+
+from ._base import BaseBatchCorrector
+from .corrections.baselines import (
+    MedianCentering,
+    ReferenceScaling,
+    ZScorePerBatch,
+)
+from .corrections.combat import ComBat, SpeciesAwareComBat
+from .corrections.harmony import Harmony
+from .corrections.limma import Limma
+from .corrections.maldi import BatchAwareWarping
+from .corrections.quality_weighted import QualityWeightedComBat
+
+__version__ = "0.1.0"
+__author__ = "Ettore Rocchi"
+
+__all__ = [
+    "BaseBatchCorrector",
+    "BatchAwareWarping",
+    "ComBat",
+    "Harmony",
+    "Limma",
+    "MedianCentering",
+    "QualityWeightedComBat",
+    "ReferenceScaling",
+    "SpeciesAwareComBat",
+    "ZScorePerBatch",
+    "__author__",
+    "__version__",
+]

maldibatchkit/_base.py

@@ -0,0 +1,176 @@
+"""Shared scaffolding for MaldiBatchKit batch correctors.
+
+Every corrector in this package stores ``batch`` (and any covariates) at
+``__init__`` time and aligns the stored arrays to ``X.index`` at
+``fit`` / ``transform``. This mirrors the ``combatlearn`` convention and
+lets sklearn ``Pipeline`` / cross-validation pass only ``X`` around while
+keeping train and test folds correctly labelled.
+
+The contract implemented here avoids the two common data-leakage traps for
+batch correction:
+
+1. ``fit`` learns correction parameters **only** from the training rows it
+   actually receives; no statistic is derived from rows the user has
+   already held back as a test set.
+2. ``transform`` never re-fits. Fitted attributes (those ending in ``_``)
+   are used verbatim on unseen rows, even when a new batch label appears
+   in the test set.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
+from sklearn.base import BaseEstimator, TransformerMixin
+from sklearn.utils.validation import check_array, check_is_fitted
+
+from ._utils import ArrayLike, _ensure_index, _subset
+
+
+def _check_array_finite(X: Any) -> None:
+    """Run :func:`sklearn.utils.validation.check_array` with finite values.
+
+    The kwarg that controls this check was renamed from
+    ``force_all_finite`` (sklearn < 1.6) to ``ensure_all_finite``
+    (sklearn >= 1.6). We support both transparently so MaldiBatchKit
+    installs cleanly on either.
+    """
+    try:
+        check_array(X, ensure_all_finite=True, dtype="numeric")
+    except TypeError:
+        check_array(X, force_all_finite=True, dtype="numeric")
+
+
+class BaseBatchCorrector(BaseEstimator, TransformerMixin):  # type: ignore[misc]
+    """Base class for batch correctors that store ``batch`` at construction.
+
+    Subclasses must implement :meth:`_fit_impl` and :meth:`_transform_impl`.
+    They should store fitted attributes as ``..._`` members so sklearn's
+    ``check_is_fitted`` works out of the box.
+
+    Parameters
+    ----------
+    batch : array-like of shape (n_samples,)
+        Batch labels for each sample. Passed once at construction; the
+        correct subset is aligned to ``X.index`` on ``fit`` / ``transform``.
+
+    Notes
+    -----
+    Subclasses typically override ``__init__`` to add method-specific
+    hyperparameters but should always call ``super().__init__(batch=batch)``
+    and leave ``batch`` stored under ``self.batch`` so
+    :meth:`get_params` / :meth:`set_params` work.
+    """
+
+    def __init__(self, batch: ArrayLike) -> None:
+        self.batch = batch
+
+    def _validate_X(self, X: ArrayLike) -> None:
+        """Validate the feature matrix."""
+        _check_array_finite(X)
+
+    def _prepare_fit(
+        self, X: ArrayLike
+    ) -> tuple[pd.DataFrame, npt.NDArray[Any], pd.Index]:
+        """Run common ``fit``-time bookkeeping.
+
+        Returns
+        -------
+        X_df : pd.DataFrame
+            ``X`` as a DataFrame with the inferred / preserved index.
+        batch_arr : ndarray
+            Batch labels aligned to ``X.index``.
+        idx : pd.Index
+            The shared index used by ``X`` and ``batch``.
+        """
+        self._validate_X(X)
+        idx = _ensure_index(X)
+        if isinstance(X, pd.DataFrame):
+            self.feature_names_in_ = np.asarray(X.columns, dtype=object)
+            X_df = X
+        else:
+            self.feature_names_in_ = np.asarray(
+                [f"x{i}" for i in range(np.asarray(X).shape[1])], dtype=object
+            )
+            X_df = pd.DataFrame(np.asarray(X), index=idx)
+        batch_vec = _subset(self.batch, idx)
+        if batch_vec is None:
+            raise ValueError("`batch` must not be None for a batch corrector.")
+        if isinstance(batch_vec, pd.DataFrame):
+            if batch_vec.shape[1] != 1:
+                raise ValueError(f"`batch` must be 1-D; got shape {batch_vec.shape}.")
+            batch_vec = batch_vec.iloc[:, 0]
+        batch_ser = pd.Series(batch_vec, index=idx)
+        if batch_ser.isna().any():
+            raise ValueError(
+                f"batch contains {int(batch_ser.isna().sum())} NaN value(s); "
+                f"all batch labels must be non-null."
+            )
+        return X_df, batch_ser.to_numpy(), idx
+
+    def _prepare_transform(
+        self, X: ArrayLike
+    ) -> tuple[pd.DataFrame, npt.NDArray[Any], pd.Index, bool]:
+        """Run common ``transform``-time bookkeeping."""
+        check_is_fitted(self)
+        self._validate_X(X)
+        idx = _ensure_index(X)
+        was_df = isinstance(X, pd.DataFrame)
+        X_df = X if was_df else pd.DataFrame(np.asarray(X), index=idx)
+        batch_vec = _subset(self.batch, idx)
+        if isinstance(batch_vec, pd.DataFrame):
+            batch_vec = batch_vec.iloc[:, 0]
+        batch_ser = pd.Series(batch_vec, index=idx)
+        return X_df, batch_ser.to_numpy(), idx, was_df
+
+    def _fit_impl(
+        self, X_df: pd.DataFrame, batch: npt.NDArray[Any]
+    ) -> None:  # pragma: no cover - abstract
+        raise NotImplementedError
+
+    def _transform_impl(
+        self, X_df: pd.DataFrame, batch: npt.NDArray[Any]
+    ) -> pd.DataFrame:  # pragma: no cover - abstract
+        raise NotImplementedError
+
+    def fit(self, X: ArrayLike, y: ArrayLike | None = None) -> BaseBatchCorrector:
+        """Fit the corrector on the training rows supplied.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Training features.
+        y : None
+            Ignored. Present for sklearn API compatibility.
+
+        Returns
+        -------
+        self : BaseBatchCorrector
+            Fitted estimator.
+        """
+        X_df, batch_arr, _ = self._prepare_fit(X)
+        self._fit_impl(X_df, batch_arr)
+        self.n_features_in_ = X_df.shape[1]
+        return self
+
+    def transform(self, X: ArrayLike) -> Any:
+        """Apply the fitted correction to ``X``.
+
+        Returns a :class:`pandas.DataFrame` when the input was a DataFrame,
+        or a :class:`numpy.ndarray` otherwise.
+        """
+        X_df, batch_arr, idx, was_df = self._prepare_transform(X)
+        X_corrected = self._transform_impl(X_df, batch_arr)
+        X_corrected.index = idx
+        X_corrected.columns = X_df.columns
+        return X_corrected if was_df else X_corrected.to_numpy()
+
+    def get_feature_names_out(
+        self, input_features: ArrayLike | None = None
+    ) -> npt.NDArray[Any]:
+        """Return the feature names seen during :meth:`fit`."""
+        check_is_fitted(self, "feature_names_in_")
+        return self.feature_names_in_

maldibatchkit/_cli_io.py

@@ -0,0 +1,211 @@
+"""Shared I/O helpers for the MaldiBatchKit CLI.
+
+Every ``correct`` subcommand and ``diagnose`` goes through this module so
+the file-format contract is documented once.
+
+Supported input / output containers
+-----------------------------------
+* **CSV**: first column is treated as the sample index; remaining
+  columns are features.  Labels (batch, species, quality, covariates)
+  live in sidecar CSVs with a sample-id index column and one or more
+  data columns.
+* **NPZ**: a ``np.savez`` archive containing at least ``X``; optionally
+  ``columns``, ``index``, ``batch``, and arbitrary extra arrays kept as
+  passthrough metadata.  This is the MaldiSet-friendly format.
+
+All auxiliary loaders align the loaded object to ``X.index`` via
+``.reindex`` when an index is available in the CSV, or positionally
+otherwise.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import typer
+
+__all__ = [
+    "load_matrix",
+    "load_dataframe_csv",
+    "load_labels_csv",
+    "load_mz_csv",
+    "save_matrix",
+]
+
+
+def _read_csv_indexed(path: Path) -> pd.DataFrame:
+    """Read a CSV assuming the first column is the sample index."""
+    return pd.read_csv(path, index_col=0)
+
+
+def load_matrix(
+    path: Path, batch_csv: Path | None
+) -> tuple[pd.DataFrame, pd.Series, dict[str, np.ndarray]]:
+    """Load a feature matrix and its aligned batch labels.
+
+    Parameters
+    ----------
+    path : Path
+        CSV or NPZ file. CSV first column is the sample index. NPZ must
+        contain an ``X`` array and may also contain ``columns``,
+        ``index`` and ``batch`` arrays.
+    batch_csv : Path or None
+        Required when ``path`` is a CSV or when an NPZ does not include
+        ``batch``. A single-column CSV with a sample-id index column.
+
+    Returns
+    -------
+    X : pd.DataFrame
+    batch : pd.Series
+        Indexed by ``X.index``.
+    extras : dict
+        Any additional arrays found in an NPZ (beyond X / index /
+        columns / batch).  Empty for CSV inputs.
+    """
+    if path.suffix.lower() == ".npz":
+        with np.load(path, allow_pickle=True) as npz:
+            if "X" not in npz.files:
+                raise typer.BadParameter(
+                    f"{path} is missing the 'X' array. Expected "
+                    f"np.savez(..., X=, [columns=, index=, batch=])."
+                )
+            X_arr = np.asarray(npz["X"])
+            columns = list(npz["columns"]) if "columns" in npz.files else None
+            index = list(npz["index"]) if "index" in npz.files else None
+            npz_batch = (
+                pd.Series(np.asarray(npz["batch"])) if "batch" in npz.files else None
+            )
+            extras = {
+                k: np.asarray(npz[k])
+                for k in npz.files
+                if k not in {"X", "columns", "index", "batch"}
+            }
+        X = pd.DataFrame(X_arr, index=index, columns=columns)
+        if npz_batch is not None:
+            if len(npz_batch) != len(X):
+                raise typer.BadParameter(
+                    f"NPZ batch length {len(npz_batch)} does not match X rows {len(X)}."
+                )
+            npz_batch.index = X.index
+            batch = npz_batch
+        else:
+            if batch_csv is None:
+                raise typer.BadParameter(
+                    f"No 'batch' array in {path}; pass --batch-csv."
+                )
+            batch = load_labels_csv(batch_csv, X.index)
+        return X, batch, extras
+
+    # CSV path
+    X = _read_csv_indexed(path)
+    if batch_csv is None:
+        raise typer.BadParameter(
+            "CSV input requires --batch-csv with a single-column file of "
+            "batch labels (first column is the sample index)."
+        )
+    batch = load_labels_csv(batch_csv, X.index)
+    return X, batch, {}
+
+
+def load_labels_csv(path: Path, index: pd.Index) -> pd.Series:
+    """Load a single-column sidecar CSV as a Series aligned to ``index``.
+
+    The CSV's first column is interpreted as the sample index; the
+    second column is the label value. If the CSV has no explicit sample
+    index (i.e. a single data column read with the default
+    ``RangeIndex``), the labels are assumed to already be in ``X`` order
+    and are aligned positionally.
+    """
+    df = pd.read_csv(path, index_col=0)
+    if df.shape[1] == 0:
+        # Only an index column was present -> data is the index itself
+        raise typer.BadParameter(
+            f"{path} has no data columns. Expected <sample_id>,<label>."
+        )
+    if df.shape[1] > 1:
+        raise typer.BadParameter(
+            f"{path} has {df.shape[1]} data columns; expected exactly one. "
+            f"For multi-column covariates use a dedicated "
+            f"--*-covariates-csv / --design-csv flag."
+        )
+    series = df.iloc[:, 0]
+    return _align_series(series, index, source=path)
+
+
+def load_dataframe_csv(path: Path, index: pd.Index) -> pd.DataFrame:
+    """Load a multi-column sidecar CSV and align it to ``index``."""
+    df = pd.read_csv(path, index_col=0)
+    if df.shape[1] == 0:
+        raise typer.BadParameter(f"{path} has no data columns.")
+    return _align_frame(df, index, source=path)
+
+
+def load_mz_csv(path: Path, n_features: int) -> np.ndarray:
+    """Load an m/z axis CSV (single column, length = n_features)."""
+    series = pd.read_csv(path).iloc[:, -1]
+    values = series.to_numpy(dtype=float)
+    if values.shape[0] != n_features:
+        raise typer.BadParameter(
+            f"--mz-csv length {values.shape[0]} does not match n_features {n_features}."
+        )
+    return values
+
+
+def save_matrix(
+    X: pd.DataFrame,
+    path: Path,
+    *,
+    batch: pd.Series | None = None,
+    extras: dict[str, np.ndarray] | None = None,
+) -> None:
+    """Write ``X`` to CSV (index + features) or NPZ (bundled arrays)."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    if path.suffix.lower() == ".npz":
+        payload: dict[str, np.ndarray] = {
+            "X": X.to_numpy(),
+            "columns": np.asarray(X.columns),
+            "index": np.asarray(X.index),
+        }
+        if batch is not None:
+            payload["batch"] = np.asarray(batch.to_numpy())
+        if extras:
+            for k, v in extras.items():
+                if k in payload:
+                    continue
+                payload[k] = np.asarray(v)
+        np.savez(path, **payload)
+    else:
+        X.to_csv(path)
+
+
+def _align_series(series: pd.Series, index: pd.Index, *, source: Path) -> pd.Series:
+    if series.index.equals(index):
+        return series
+    if set(index).issubset(set(series.index)):
+        return series.reindex(index)
+    if len(series) == len(index):
+        # Positional fallback: the CSV did not carry the X sample IDs.
+        aligned = series.copy()
+        aligned.index = index
+        return aligned
+    raise typer.BadParameter(
+        f"{source}: cannot align labels to X.index. "
+        f"Labels have {len(series)} rows, X has {len(index)} rows, "
+        f"and label indices do not cover X.index."
+    )
+
+
+def _align_frame(df: pd.DataFrame, index: pd.Index, *, source: Path) -> pd.DataFrame:
+    if df.index.equals(index):
+        return df
+    if set(index).issubset(set(df.index)):
+        return df.reindex(index)
+    if len(df) == len(index):
+        aligned = df.copy()
+        aligned.index = index
+        return aligned
+    raise typer.BadParameter(
+        f"{source}: cannot align to X.index. Got {len(df)} rows, X has {len(index)}."
+    )

maldibatchkit/_utils.py

@@ -0,0 +1,122 @@
+"""Internal helpers shared across MaldiBatchKit corrections."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
+
+ArrayLike = pd.DataFrame | pd.Series | npt.NDArray[Any]
+
+
+def _as_dataframe(X: ArrayLike) -> tuple[pd.DataFrame, bool]:
+    """Coerce ``X`` to a ``pd.DataFrame``.
+
+    Returns
+    -------
+    df : pd.DataFrame
+        The input as a DataFrame (copy if necessary so mutation is safe).
+    was_dataframe : bool
+        ``True`` if the caller passed a DataFrame (so caller can choose to
+        return a DataFrame); ``False`` otherwise.
+    """
+    if isinstance(X, pd.DataFrame):
+        return X, True
+    arr = np.asarray(X)
+    if arr.ndim == 1:
+        arr = arr.reshape(-1, 1)
+    return pd.DataFrame(arr), False
+
+
+def _subset(obj: ArrayLike | None, idx: pd.Index) -> pd.Series | pd.DataFrame | None:
+    """Align a batch/covariate array-like to ``idx``.
+
+    Mirrors ``combatlearn._utils._subset`` so all MaldiBatchKit correctors
+    behave the same way: the caller stores ``batch`` at ``__init__`` time
+    and we select the rows matching ``X.index`` at ``fit`` / ``transform``.
+
+    When ``obj`` is a pandas object it is sliced with ``.loc[idx]`` (which
+    naturally supports train/test splits). A numpy array or list is wrapped
+    in a Series/DataFrame using ``idx`` directly - the user is responsible
+    for making sure the ordering matches ``X``.
+    """
+    if obj is None:
+        return None
+    if isinstance(obj, pd.Series | pd.DataFrame):
+        try:
+            return obj.loc[idx]
+        except KeyError as exc:
+            missing = [k for k in idx if k not in obj.index]
+            raise ValueError(
+                f"batch/covariate index does not cover all of X.index. "
+                f"{len(missing)} missing key(s) (first 3): {missing[:3]}. "
+                f"Ensure your batch labels are indexed by the same sample IDs "
+                f"used by X, or pass a bare numpy array when X has a default "
+                f"RangeIndex."
+            ) from exc
+    arr = np.asarray(obj)
+    if arr.ndim == 1:
+        return pd.Series(arr, index=idx)
+    return pd.DataFrame(arr, index=idx)
+
+
+def _ensure_index(X: ArrayLike) -> pd.Index:
+    """Return ``X.index`` if ``X`` is a pandas object, else a default RangeIndex."""
+    if isinstance(X, pd.DataFrame | pd.Series):
+        return X.index
+    return pd.RangeIndex(len(X))
+
+
+def _restore_output(X_corrected: pd.DataFrame, was_dataframe: bool) -> Any:
+    """Return a DataFrame if the input was a DataFrame, otherwise a numpy array."""
+    if was_dataframe:
+        return X_corrected
+    return X_corrected.to_numpy()
+
+
+def _resolve_mz_axis(
+    df: pd.DataFrame, mz_values: ArrayLike | None
+) -> tuple[npt.NDArray[np.float64], bool]:
+    """Resolve the m/z axis for an ``(n_samples, n_features)`` feature matrix.
+
+    Resolution order:
+
+    1. If the caller passed ``mz_values`` explicitly, use it verbatim
+       (length is validated against ``df.shape[1]``). Return
+       ``(mz, True)``.
+    2. Else, try to coerce ``df.columns`` to ``float`` - this is the
+       case for matrices coming straight from
+       :class:`maldiamrkit.MaldiSet`, whose columns are the binned m/z
+       values (string-typed in the current MaldiAMRKit release, but
+       numerically convertible, e.g. ``'2000.0'`` ...). Return
+       ``(mz, True)`` when it works.
+    3. Fall back to ``np.arange(df.shape[1])`` so ndarray inputs still
+       produce a plot - but signal to the caller that the axis is
+       bin indices rather than m/z, by returning ``(mz, False)``.
+    """
+    n_features = df.shape[1]
+    if mz_values is not None:
+        mz = np.asarray(mz_values, dtype=float)
+        if mz.shape[0] != n_features:
+            raise ValueError(
+                f"mz_values has length {mz.shape[0]} but X has {n_features} columns."
+            )
+        return mz, True
+    # pandas.DataFrame(ndarray) yields a default RangeIndex(0, n) -
+    # that is positional, not an m/z axis. Treat it as bin indices.
+    if isinstance(df.columns, pd.RangeIndex):
+        return np.arange(n_features, dtype=float), False
+    try:
+        mz = np.asarray(df.columns, dtype=float)
+    except (TypeError, ValueError):
+        return np.arange(n_features, dtype=float), False
+    if np.any(np.isnan(mz)):
+        return np.arange(n_features, dtype=float), False
+    # If the coerced axis is exactly ``[0, 1, ..., n-1]`` it is an
+    # integer-labelled DataFrame that happens to parse as floats;
+    # the caller almost certainly did not mean those as m/z.
+    if np.array_equal(mz, np.arange(n_features, dtype=float)):
+        return mz, False
+    return mz, True