MaldiDeepKit 0.1.0__py3-none-any.whl

maldideepkit/__init__.py

@@ -0,0 +1,54 @@
+"""MaldiDeepKit -- deep learning classifiers for MALDI-TOF binned spectra.
+
+Provides a catalog of PyTorch architectures (MLP, CNN, ResNet,
+Transformer) adapted to 1-D binned MALDI-TOF spectra, each wrapped in
+a scikit-learn compatible estimator with sensible defaults.
+
+Subpackages
+-----------
+- ``maldideepkit.base`` -- ``BaseSpectralClassifier``, ``SpectralDataset``,
+  ``make_loaders``.
+- ``maldideepkit.attention`` -- ``MaldiMLPClassifier`` (MLP with optional
+  sigmoid-gated attention).
+- ``maldideepkit.cnn`` -- ``MaldiCNNClassifier`` (Conv1D blocks).
+- ``maldideepkit.resnet`` -- ``MaldiResNetClassifier`` (1-D ResNet-18).
+- ``maldideepkit.transformer`` -- ``MaldiTransformerClassifier`` (1-D ViT).
+- ``maldideepkit.blocks`` -- re-exports of every backbone and
+  composable primitive for users embedding components into their own
+  networks.
+- ``maldideepkit.utils`` -- reproducibility helpers and shared
+  training primitives.
+
+Examples
+--------
+>>> import numpy as np
+>>> from maldideepkit import MaldiMLPClassifier
+>>> rng = np.random.default_rng(0)
+>>> X = rng.standard_normal((64, 256)).astype("float32")
+>>> y = rng.integers(0, 2, size=64)
+>>> clf = MaldiMLPClassifier(epochs=2, batch_size=16, random_state=0)
+>>> _ = clf.fit(X, y)
+>>> proba = clf.predict_proba(X)
+"""
+
+from .attention.mlp import MaldiMLPClassifier
+from .base.classifier import BaseSpectralClassifier
+from .base.data import SpectralDataset, make_loaders
+from .cnn.cnn import MaldiCNNClassifier
+from .resnet.resnet import MaldiResNetClassifier
+from .transformer.transformer import MaldiTransformerClassifier
+
+__version__ = "0.1.0"
+__author__ = "Ettore Rocchi"
+
+__all__ = [
+    "BaseSpectralClassifier",
+    "MaldiCNNClassifier",
+    "MaldiMLPClassifier",
+    "MaldiResNetClassifier",
+    "MaldiTransformerClassifier",
+    "SpectralDataset",
+    "__author__",
+    "__version__",
+    "make_loaders",
+]

maldideepkit/_bin_scaling.py

@@ -0,0 +1,52 @@
+"""Helpers to scale architectural hyperparameters with spectrum layout.
+
+The package defaults are calibrated for ``bin_width=3`` over the
+2000-20000 Da range (``input_dim=6000``). For other bin widths,
+:func:`scale_odd_kernel` adjusts the conv kernel size inversely so
+the receptive field in Daltons stays comparable.
+
+See ``docs/source/spectrum_scaling.rst`` for the per-species
+peak-width statistics underpinning the defaults.
+"""
+
+from __future__ import annotations
+
+REFERENCE_BIN_WIDTH = 3
+REFERENCE_CONV_KERNEL = 7
+
+
+def scale_odd_kernel(
+    bin_width: int,
+    *,
+    reference_kernel: int = REFERENCE_CONV_KERNEL,
+    reference_bin_width: int = REFERENCE_BIN_WIDTH,
+    min_kernel: int = 3,
+) -> int:
+    """Return an odd conv kernel scaled inversely with ``bin_width``.
+
+    Used by :class:`~maldideepkit.MaldiCNNClassifier` and
+    :class:`~maldideepkit.MaldiResNetClassifier` to adjust the first
+    convolutional layer's receptive field for non-default bin widths.
+
+    Parameters
+    ----------
+    bin_width : int
+        The target bin width in Daltons.
+    reference_kernel : int, default=7
+        Kernel size used at ``reference_bin_width`` (the package default).
+    reference_bin_width : int, default=3
+        Reference bin width at which ``reference_kernel`` was chosen.
+    min_kernel : int, default=3
+        Lower bound on the returned kernel size.
+
+    Returns
+    -------
+    int
+        Odd kernel size ``>= min_kernel`` closest to
+        ``reference_kernel * reference_bin_width / bin_width``.
+    """
+    target = reference_kernel * reference_bin_width / bin_width
+    k = max(min_kernel, round(target))
+    if k % 2 == 0:
+        k += 1
+    return k

maldideepkit/_blocks.py

@@ -0,0 +1,80 @@
+"""Shared ``nn.Module`` primitives used by more than one architecture."""
+
+from __future__ import annotations
+
+import torch
+from einops import rearrange
+from torch import nn
+
+
+class PatchEmbed1D(nn.Module):
+    """1-D patch embedding via a strided Conv1D + LayerNorm.
+
+    Maps ``(B, C_in, L)`` to ``(B, L // patch_size, embed_dim)`` by
+    convolving with a non-overlapping kernel of width ``patch_size``
+    and stride ``patch_size``.
+
+    Parameters
+    ----------
+    patch_size : int
+        Non-overlapping patch width (and stride).
+    in_channels : int
+        Input channels. ``1`` for raw binned spectra after an
+        ``unsqueeze(1)``.
+    embed_dim : int
+        Output embedding dimension.
+    """
+
+    def __init__(self, patch_size: int, in_channels: int, embed_dim: int) -> None:
+        super().__init__()
+        self.proj = nn.Conv1d(
+            in_channels, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
+        self.norm = nn.LayerNorm(embed_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Map ``(B, C, L)`` to ``(B, L // patch_size, embed_dim)``."""
+        x = self.proj(x)
+        x = rearrange(x, "b c l -> b l c")
+        return self.norm(x)
+
+
+class DropPath(nn.Module):
+    """Stochastic depth (per-sample residual drop).
+
+    During training, zero out a fraction ``drop_prob`` of samples
+    uniformly at random and rescale the survivors by
+    ``1 / (1 - drop_prob)`` so the expectation is unchanged. At
+    inference time or when ``drop_prob == 0`` this is a no-op.
+
+    Parameters
+    ----------
+    drop_prob : float, default=0.0
+        Probability of dropping a sample's residual. Must be in
+        ``[0, 1)``.
+    generator : torch.Generator or None, default=None
+        Optional explicit RNG for the Bernoulli mask. When ``None``
+        PyTorch's global generator is used.
+    """
+
+    def __init__(
+        self,
+        drop_prob: float = 0.0,
+        generator: torch.Generator | None = None,
+    ) -> None:
+        super().__init__()
+        if not 0.0 <= drop_prob < 1.0:
+            raise ValueError(f"drop_prob must be in [0, 1); got {drop_prob!r}.")
+        self.drop_prob = float(drop_prob)
+        self.generator = generator
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.drop_prob == 0.0 or not self.training:
+            return x
+        keep_prob = 1.0 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        mask = x.new_empty(shape).bernoulli_(keep_prob, generator=self.generator)
+        return x * mask / keep_prob
+
+    def extra_repr(self) -> str:
+        return f"drop_prob={self.drop_prob}"

maldideepkit/attention/__init__.py

@@ -0,0 +1,5 @@
+"""Attention-based MLP classifier for MALDI-TOF spectra."""
+
+from .mlp import MaldiMLPClassifier, SpectralAttentionMLP
+
+__all__ = ["MaldiMLPClassifier", "SpectralAttentionMLP"]

maldideepkit/attention/mlp.py

@@ -0,0 +1,319 @@
+"""MLP classifier with optional sigmoid-gated attention.
+
+The architecture is a dense network with a learned per-feature gate on
+the first hidden layer that doubles as an interpretable attention map
+over the projected bin representation.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from sklearn.utils.validation import check_is_fitted
+from torch import nn
+
+from ..base.classifier import BaseSpectralClassifier
+from ..base.data import _to_numpy
+
+
+class SpectralAttentionMLP(nn.Module):
+    """Projection + optional sigmoid-gated attention + deep MLP head.
+
+    Parameters
+    ----------
+    input_dim : int
+        Number of input bins. The first linear layer projects this down
+        to :attr:`hidden_dim`.
+    n_classes : int, default=2
+        Number of output logits.
+    hidden_dim : int, default=512
+        Width of the projection layer and attention gate.
+    head_dims : sequence of int, default=(256, 128)
+        Widths of the hidden layers between the gated representation and
+        the output logits.
+    use_attention : bool, default=True
+        If ``True``, apply a sigmoid-gated element-wise attention on the
+        projected features. If ``False``, the model reduces to a plain
+        MLP of the same depth.
+    dropout_high : float, default=0.3
+        Dropout applied after the projection and the first dense layer.
+    dropout_low : float, default=0.2
+        Dropout applied before the output logits.
+
+    Attributes
+    ----------
+    last_attention : torch.Tensor or None
+        Attention weights from the most recent forward pass
+        (``(batch, hidden_dim)``). ``None`` when ``use_attention=False``.
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        n_classes: int = 2,
+        hidden_dim: int = 512,
+        head_dims: tuple[int, ...] = (256, 128),
+        use_attention: bool = True,
+        dropout_high: float = 0.3,
+        dropout_low: float = 0.2,
+    ) -> None:
+        super().__init__()
+        self.use_attention = use_attention
+        self.hidden_dim = hidden_dim
+
+        self.proj = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.BatchNorm1d(hidden_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout_high),
+        )
+        if use_attention:
+            self.attn: nn.Module = nn.Sequential(
+                nn.Linear(hidden_dim, hidden_dim), nn.Sigmoid()
+            )
+        else:
+            self.attn = nn.Identity()
+
+        head_layers: list[nn.Module] = []
+        prev = hidden_dim
+        for i, width in enumerate(head_dims):
+            head_layers += [
+                nn.Linear(prev, width),
+                nn.BatchNorm1d(width),
+                nn.ReLU(),
+                nn.Dropout(dropout_high if i == 0 else dropout_low),
+            ]
+            prev = width
+        head_layers.append(nn.Linear(prev, n_classes))
+        self.head = nn.Sequential(*head_layers)
+
+        self.last_attention: torch.Tensor | None = None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Map ``(batch, input_dim)`` to ``(batch, n_classes)`` logits."""
+        projected = self.proj(x)
+        if self.use_attention:
+            weights = self.attn(projected)
+            self.last_attention = weights.detach()
+            gated = projected * weights
+        else:
+            self.last_attention = None
+            gated = projected
+        return self.head(gated)
+
+
+class MaldiMLPClassifier(BaseSpectralClassifier):
+    """sklearn-compatible MLP classifier with optional attention gating.
+
+    Parameters
+    ----------
+    hidden_dim : int, default=512
+        Width of the projection and attention gate.
+    head_dims : sequence of int, default=(256, 128)
+        Widths of the hidden layers of the classification head.
+    use_attention : bool, default=True
+        Toggle the sigmoid-gated attention. When ``False``, the model
+        is a plain MLP of the same depth.
+    dropout_high : float, default=0.3
+        Dropout after the projection and first head layer.
+    dropout_low : float, default=0.2
+        Dropout before the output logits.
+    **kwargs
+        Forwarded to :class:`~maldideepkit.base.classifier.BaseSpectralClassifier`:
+        ``input_dim``, ``n_classes``, ``learning_rate``, ``batch_size``,
+        ``epochs``, ``early_stopping_patience``, ``val_fraction``,
+        ``standardize``, ``class_weight``, ``device``, ``random_state``,
+        ``verbose``.
+
+    Attributes
+    ----------
+    attention_weights_ : ndarray or None
+        Attention weights from the last :meth:`fit` or :meth:`predict`
+        forward pass. Shape ``(n_samples_last_call, hidden_dim)``. Set to
+        ``None`` when ``use_attention=False``.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from maldideepkit import MaldiMLPClassifier
+    >>> rng = np.random.default_rng(0)
+    >>> X = rng.standard_normal((64, 256)).astype("float32")
+    >>> y = rng.integers(0, 2, size=64)
+    >>> clf = MaldiMLPClassifier(epochs=2, batch_size=16, random_state=0).fit(X, y)
+    >>> clf.predict(X).shape
+    (64,)
+    >>> weights = clf.get_attention_weights(X[:4])
+    >>> weights.shape
+    (4, 512)
+    """
+
+    def __init__(
+        self,
+        input_dim: int | None = None,
+        n_classes: int = 2,
+        hidden_dim: int = 512,
+        head_dims: tuple[int, ...] = (256, 128),
+        use_attention: bool = True,
+        dropout_high: float = 0.3,
+        dropout_low: float = 0.2,
+        learning_rate: float = 1e-3,
+        weight_decay: float = 0.0,
+        grad_clip_norm: float | None = None,
+        label_smoothing: float = 0.0,
+        loss: str = "cross_entropy",
+        focal_gamma: float = 2.0,
+        use_amp: bool = False,
+        swa_start_epoch: int | None = None,
+        tune_threshold: bool = False,
+        threshold_metric: str = "balanced_accuracy",
+        calibrate_temperature: bool = False,
+        min_val_auroc_for_threshold_tune: float = 0.6,
+        use_sam: bool = False,
+        sam_rho: float = 0.05,
+        batch_size: int = 32,
+        epochs: int = 100,
+        early_stopping_patience: int = 10,
+        val_fraction: float = 0.1,
+        warmup_epochs: int = 0,
+        standardize: bool = False,
+        input_transform: str | None = None,
+        warping: Any | None = None,
+        metrics_log_path: str | Path | None = None,
+        track_train_metrics: bool = False,
+        augment: Any | None = None,
+        mixup_alpha: float = 0.0,
+        cutmix_alpha: float = 0.0,
+        ema_decay: float | None = None,
+        retry_on_val_auroc_below: float | None = None,
+        max_retries: int = 2,
+        class_weight: str | np.ndarray | list | None = None,
+        device: str | torch.device = "auto",
+        random_state: int = 0,
+        verbose: bool = False,
+    ) -> None:
+        super().__init__(
+            input_dim=input_dim,
+            n_classes=n_classes,
+            learning_rate=learning_rate,
+            weight_decay=weight_decay,
+            grad_clip_norm=grad_clip_norm,
+            label_smoothing=label_smoothing,
+            loss=loss,
+            focal_gamma=focal_gamma,
+            use_amp=use_amp,
+            swa_start_epoch=swa_start_epoch,
+            tune_threshold=tune_threshold,
+            threshold_metric=threshold_metric,
+            calibrate_temperature=calibrate_temperature,
+            min_val_auroc_for_threshold_tune=min_val_auroc_for_threshold_tune,
+            use_sam=use_sam,
+            sam_rho=sam_rho,
+            batch_size=batch_size,
+            epochs=epochs,
+            early_stopping_patience=early_stopping_patience,
+            val_fraction=val_fraction,
+            warmup_epochs=warmup_epochs,
+            standardize=standardize,
+            input_transform=input_transform,
+            warping=warping,
+            metrics_log_path=metrics_log_path,
+            track_train_metrics=track_train_metrics,
+            augment=augment,
+            mixup_alpha=mixup_alpha,
+            cutmix_alpha=cutmix_alpha,
+            ema_decay=ema_decay,
+            retry_on_val_auroc_below=retry_on_val_auroc_below,
+            max_retries=max_retries,
+            class_weight=class_weight,
+            device=device,
+            random_state=random_state,
+            verbose=verbose,
+        )
+        self.hidden_dim = hidden_dim
+        self.head_dims = head_dims
+        self.use_attention = use_attention
+        self.dropout_high = dropout_high
+        self.dropout_low = dropout_low
+        self.attention_weights_: np.ndarray | None = None
+
+    def _build_model(self) -> nn.Module:
+        return SpectralAttentionMLP(
+            input_dim=self.input_dim_,
+            n_classes=self.n_classes_,
+            hidden_dim=int(self.hidden_dim),
+            head_dims=tuple(self.head_dims),
+            use_attention=bool(self.use_attention),
+            dropout_high=float(self.dropout_high),
+            dropout_low=float(self.dropout_low),
+        )
+
+    def _forward_logits(self, X: Any) -> np.ndarray:
+        logits = super()._forward_logits(X)
+        if self.use_attention and self.model_.last_attention is not None:
+            self.attention_weights_ = self.model_.last_attention.detach().cpu().numpy()
+        else:
+            self.attention_weights_ = None
+        return logits
+
+    def fit(self, X: Any, y: Any) -> MaldiMLPClassifier:  # type: ignore[override]
+        """Fit the model and cache attention weights from the final batch.
+
+        See :meth:`BaseSpectralClassifier.fit` for shared parameters.
+        """
+        super().fit(X, y)
+        if self.use_attention:
+            X_np = _to_numpy(X)
+            tail = X_np[: min(len(X_np), 64)]
+            self._forward_logits(tail)
+        else:
+            self.attention_weights_ = None
+        return self
+
+    def get_attention_weights(self, X: Any) -> np.ndarray:
+        """Return attention weights for ``X`` of shape ``(len(X), hidden_dim)``.
+
+        Parameters
+        ----------
+        X : array-like or MaldiSet of shape (n_samples, n_bins)
+            Spectra to inspect. Must match ``input_dim_``.
+
+        Returns
+        -------
+        ndarray of shape (n_samples, hidden_dim)
+            Sigmoid-gated attention weights.
+
+        Raises
+        ------
+        RuntimeError
+            If the classifier was built with ``use_attention=False``.
+        """
+        check_is_fitted(self, "model_")
+        if not self.use_attention:
+            raise RuntimeError(
+                "get_attention_weights is only available when use_attention=True."
+            )
+        self._forward_logits(X)
+        if self.attention_weights_ is None:
+            raise RuntimeError(
+                "Attention weights were not captured during forward; "
+                "ensure the model was built with use_attention=True."
+            )
+        return self.attention_weights_
+
+    @classmethod
+    def from_spectrum(
+        cls, bin_width: int, input_dim: int, **overrides
+    ) -> "MaldiMLPClassifier":
+        """Construct a classifier for a given ``(bin_width, input_dim)`` layout.
+
+        The MLP is architecturally scale-agnostic, so this factory
+        only forwards ``input_dim`` and any ``**overrides``. Provided
+        for API symmetry with the other classifiers.
+        """
+        del bin_width
+        kwargs: dict[str, Any] = {"input_dim": input_dim}
+        kwargs.update(overrides)
+        return cls(**kwargs)

maldideepkit/augment/__init__.py

@@ -0,0 +1,15 @@
+"""Data-augmentation utilities for binned MALDI-TOF spectra.
+
+All augmentations are callables that transform a training-batch tensor
+of shape ``(batch, n_bins)`` and return a tensor of the same shape.
+Wire them into :class:`~maldideepkit.BaseSpectralClassifier` via the
+``augment=`` kwarg; they apply to training batches only and are
+bypassed during validation and inference.
+"""
+
+from __future__ import annotations
+
+from .mixing import apply_cutmix, apply_mixup, to_one_hot
+from .spectra import SpectrumAugment
+
+__all__ = ["SpectrumAugment", "apply_cutmix", "apply_mixup", "to_one_hot"]

maldideepkit/augment/mixing.py

@@ -0,0 +1,113 @@
+"""MixUp and CutMix for 1-D binned MALDI-TOF spectra.
+
+Both transforms operate on a batch of features ``x`` with shape
+``(batch, n_bins)`` and one-hot targets ``y_oh`` with shape
+``(batch, n_classes)``, and return a mixed ``(x, y_soft)`` pair.
+
+- MixUp: ``x = lam * x_i + (1 - lam) * x_j``.
+- CutMix: splice a contiguous m/z window from a shuffled sample
+  into the original; labels mixed by window fraction.
+
+Both draw the mix coefficient from ``Beta(alpha, alpha)``.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+def to_one_hot(y: torch.Tensor, n_classes: int) -> torch.Tensor:
+    """Return ``y`` as a float one-hot tensor of shape ``(batch, n_classes)``."""
+    return F.one_hot(y.long(), num_classes=n_classes).to(dtype=torch.float32)
+
+
+def _sample_beta(alpha: float, generator: torch.Generator | None = None) -> float:
+    """Draw a single ``Beta(alpha, alpha)`` sample."""
+    if generator is None:
+        return float(np.random.beta(alpha, alpha))
+    a = torch.tensor([float(alpha)], dtype=torch.float64)
+    x = torch._standard_gamma(a, generator=generator)
+    y = torch._standard_gamma(a, generator=generator)
+    return float((x / (x + y)).item())
+
+
+def apply_mixup(
+    x: torch.Tensor,
+    y_oh: torch.Tensor,
+    alpha: float,
+    generator: torch.Generator | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Mixup: convex-combine two random permutations of the batch.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Feature tensor of shape ``(batch, n_bins)``.
+    y_oh : torch.Tensor
+        One-hot (or soft) target tensor of shape ``(batch, n_classes)``.
+    alpha : float
+        Beta-distribution parameter (``Beta(alpha, alpha)``). Typical
+        values 0.1-0.4 for tabular-ish inputs. Must be ``> 0``.
+    generator : torch.Generator or None, default=None
+        Seeded RNG for reproducibility.
+
+    Returns
+    -------
+    tuple of torch.Tensor
+        ``(x_mixed, y_mixed)`` with the same shapes as the inputs.
+    """
+    if alpha <= 0:
+        raise ValueError(f"mixup alpha must be > 0; got {alpha!r}.")
+    lam = _sample_beta(alpha, generator)
+    perm = torch.randperm(x.shape[0], generator=generator).to(x.device)
+    x_mixed = lam * x + (1.0 - lam) * x[perm]
+    y_mixed = lam * y_oh + (1.0 - lam) * y_oh[perm]
+    return x_mixed, y_mixed
+
+
+def apply_cutmix(
+    x: torch.Tensor,
+    y_oh: torch.Tensor,
+    alpha: float,
+    generator: torch.Generator | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """CutMix on 1-D spectra: splice a contiguous m/z window.
+
+    A window of length ``w = round(n_bins * (1 - lam))`` is drawn
+    uniformly along the m/z axis and copied from the shuffled sample
+    into the original. Labels are mixed by the window fraction.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Feature tensor of shape ``(batch, n_bins)``.
+    y_oh : torch.Tensor
+        One-hot (or soft) target tensor of shape ``(batch, n_classes)``.
+    alpha : float
+        Beta-distribution parameter (``Beta(alpha, alpha)``). Typical
+        value 1.0 (uniform over window fractions). Must be ``> 0``.
+    generator : torch.Generator or None, default=None
+        Seeded RNG for reproducibility.
+
+    Returns
+    -------
+    tuple of torch.Tensor
+        ``(x_mixed, y_mixed)`` with the same shapes as the inputs.
+    """
+    if alpha <= 0:
+        raise ValueError(f"cutmix alpha must be > 0; got {alpha!r}.")
+    batch, n_bins = x.shape
+    lam = _sample_beta(alpha, generator)
+    window = int(round(n_bins * (1.0 - lam)))
+    window = max(0, min(window, n_bins))
+    if window == 0:
+        return x.clone(), y_oh.clone()
+    start = int(torch.randint(0, n_bins - window + 1, (1,), generator=generator).item())
+    perm = torch.randperm(batch, generator=generator).to(x.device)
+    x_mixed = x.clone()
+    x_mixed[:, start : start + window] = x[perm][:, start : start + window]
+    effective_lam = 1.0 - window / n_bins
+    y_mixed = effective_lam * y_oh + (1.0 - effective_lam) * y_oh[perm]
+    return x_mixed, y_mixed