MaldiDeepKit 0.1.0__py3-none-any.whl

maldideepkit/utils/ensemble.py

@@ -0,0 +1,132 @@
+"""Simple mean-of-``predict_proba`` ensemble for spectral classifiers.
+
+A thin wrapper that fits each member independently and averages
+``predict_proba`` across them.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+
+class SpectralEnsemble:
+    """Ensemble N fitted or unfitted spectral classifiers.
+
+    Parameters
+    ----------
+    classifiers : sequence of BaseSpectralClassifier
+        Unfitted classifier instances. :meth:`fit` calls each
+        member's own ``fit`` in order.
+
+    Attributes
+    ----------
+    classes_ : np.ndarray
+        Union of class labels reported by the members. Members must
+        agree on the label set after fitting.
+    """
+
+    def __init__(self, classifiers: list[Any]) -> None:
+        if not classifiers:
+            raise ValueError("SpectralEnsemble needs at least one classifier.")
+        self.classifiers = list(classifiers)
+
+    def fit(self, X: Any, y: Any) -> SpectralEnsemble:
+        """Fit every member on the same ``(X, y)``."""
+        first_classes: np.ndarray | None = None
+        for i, clf in enumerate(self.classifiers):
+            clf.fit(X, y)
+            if first_classes is None:
+                first_classes = np.asarray(clf.classes_)
+            elif not np.array_equal(clf.classes_, first_classes):
+                raise ValueError(
+                    f"Ensemble member {i} produced classes_={clf.classes_!r}; "
+                    f"expected {first_classes!r}. All members must see the same labels."
+                )
+        self.classes_ = first_classes
+        return self
+
+    def predict_proba(self, X: Any) -> np.ndarray:
+        """Return the mean of member ``predict_proba`` outputs."""
+        probas = [clf.predict_proba(X) for clf in self.classifiers]
+        stacked = np.stack(probas, axis=0)
+        return stacked.mean(axis=0)
+
+    def predict(self, X: Any) -> np.ndarray:
+        """Argmax of the averaged probabilities.
+
+        Per-member post-hoc calibration / thresholds are intentionally
+        not averaged.
+        """
+        proba = self.predict_proba(X)
+        idx = np.argmax(proba, axis=1)
+        return np.asarray(self.classes_)[idx]
+
+    def score(self, X: Any, y: Any) -> float:
+        """Mean accuracy against ``y``."""
+        preds = self.predict(X)
+        if hasattr(y, "to_numpy"):
+            y = y.to_numpy()
+        return float(np.mean(preds == np.asarray(y).ravel()))
+
+    def save(self, path: str | Path) -> None:
+        """Save each member under ``<path>_<i>``.
+
+        Example: ``SpectralEnsemble.save("my_ens")`` writes
+        ``my_ens_0.pt`` / ``my_ens_0.json`` / ... plus an index file
+        ``my_ens.ensemble.json`` recording the per-member classes.
+        """
+        import json
+
+        base = Path(path)
+        if base.suffix:
+            base = base.with_suffix("")
+        base.parent.mkdir(parents=True, exist_ok=True)
+        member_paths = []
+        for i, clf in enumerate(self.classifiers):
+            member_path = base.parent / f"{base.name}_{i}"
+            clf.save(member_path)
+            member_paths.append(str(member_path.name))
+        index_path = base.with_suffix(".ensemble.json")
+        with open(index_path, "w") as fh:
+            json.dump(
+                {
+                    "version": 1,
+                    "n_members": len(self.classifiers),
+                    "member_files": member_paths,
+                    "classes_": (
+                        np.asarray(self.classes_).tolist()
+                        if hasattr(self, "classes_") and self.classes_ is not None
+                        else None
+                    ),
+                    "member_class_names": [type(c).__name__ for c in self.classifiers],
+                },
+                fh,
+                indent=2,
+            )
+
+    @classmethod
+    def load(cls, path: str | Path) -> SpectralEnsemble:
+        """Inverse of :meth:`save`."""
+        import json
+
+        from ..base.classifier import BaseSpectralClassifier
+
+        base = Path(path)
+        if base.suffix:
+            base = base.with_suffix("")
+        index_path = base.with_suffix(".ensemble.json")
+        if not index_path.exists():
+            raise FileNotFoundError(index_path)
+        with open(index_path) as fh:
+            meta = json.load(fh)
+        members: list[Any] = []
+        for name in meta["member_files"]:
+            member_path = base.parent / name
+            members.append(BaseSpectralClassifier.load(member_path))
+        ens = cls(members)
+        if meta.get("classes_") is not None:
+            ens.classes_ = np.asarray(meta["classes_"])
+        return ens

maldideepkit/utils/loss.py

@@ -0,0 +1,138 @@
+"""Loss functions used by :class:`BaseSpectralClassifier`."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class FocalLoss(nn.Module):
+    r"""Multi-class focal loss with optional class weighting and label smoothing.
+
+    Implements
+
+    .. math::
+
+        L = - (1 - p_t)^\gamma \log p_t
+
+    where :math:`p_t` is the predicted probability of the true class.
+    At :math:`\gamma = 0` and ``label_smoothing=0`` this reduces to
+    :class:`~torch.nn.CrossEntropyLoss`.
+
+    Parameters
+    ----------
+    weight : torch.Tensor or None, default=None
+        Per-class weight tensor of shape ``(n_classes,)``. Applied to
+        every sample by gathering at its target index (matches the
+        :class:`CrossEntropyLoss` convention for ``weight``).
+    gamma : float, default=2.0
+        Focusing parameter. ``0`` degrades to cross-entropy; ``2`` is
+        the value used in Lin et al. 2017.
+    label_smoothing : float, default=0.0
+        Target smoothing in ``[0, 1)``. At ``0.0`` the target is a
+        one-hot vector; otherwise the target distribution becomes
+        ``(1 - eps) * one_hot + eps / n_classes``.
+    reduction : {"mean", "sum", "none"}, default="mean"
+        How to reduce the per-sample loss tensor.
+    """
+
+    def __init__(
+        self,
+        weight: torch.Tensor | None = None,
+        gamma: float = 2.0,
+        label_smoothing: float = 0.0,
+        reduction: str = "mean",
+    ) -> None:
+        super().__init__()
+        if gamma < 0:
+            raise ValueError(f"gamma must be >= 0; got {gamma!r}.")
+        if not 0.0 <= label_smoothing < 1.0:
+            raise ValueError(
+                f"label_smoothing must be in [0, 1); got {label_smoothing!r}."
+            )
+        if reduction not in {"mean", "sum", "none"}:
+            raise ValueError(
+                f"reduction must be 'mean', 'sum', or 'none'; got {reduction!r}."
+            )
+        self.register_buffer(
+            "class_weight",
+            weight.detach().clone() if weight is not None else None,
+            persistent=False,
+        )
+        self.gamma = float(gamma)
+        self.label_smoothing = float(label_smoothing)
+        self.reduction = reduction
+
+    def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        r"""Compute focal loss for ``(N, C)`` logits.
+
+        Accepts either integer targets of shape ``(N,)`` or a soft
+        probability distribution of shape ``(N, C)`` (as produced by
+        MixUp / CutMix). When soft targets are passed the loss
+        becomes
+
+        .. math::
+
+            L = - \sum_c t_c \, (1 - p_c)^\gamma \log p_c
+
+        ``label_smoothing`` is ignored on the soft-target path.
+
+        Class weighting follows the :class:`~torch.nn.CrossEntropyLoss`
+        convention with ``reduction="mean"``: the per-sample weight is
+        ``weight[y_i]`` (or ``Σ_c t_c · weight_c`` for soft targets),
+        and the mean reduction divides by ``Σ_i sample_weight_i``
+        rather than ``N``.
+        """
+        log_probs = F.log_softmax(logits, dim=-1)
+        probs = log_probs.exp()
+        n_classes = logits.shape[-1]
+        sample_weight: torch.Tensor | None = None
+
+        if target.dim() == 2:
+            smooth = target.to(dtype=log_probs.dtype)
+            focal_per_class = (1.0 - probs).clamp_min(1e-12).pow(self.gamma)
+            per_class_loss = -smooth * focal_per_class * log_probs
+            loss = per_class_loss.sum(dim=-1)
+            if self.class_weight is not None:
+                w = self.class_weight.to(loss.device)
+                sample_weight = (smooth * w).sum(dim=-1)
+                loss = loss * sample_weight
+        elif self.label_smoothing == 0.0:
+            logpt = log_probs.gather(1, target.unsqueeze(1)).squeeze(1)
+            pt = probs.gather(1, target.unsqueeze(1)).squeeze(1)
+            focal_term = (1.0 - pt).clamp_min(1e-12).pow(self.gamma)
+            loss = -focal_term * logpt
+            if self.class_weight is not None:
+                sample_weight = self.class_weight.to(loss.device).gather(0, target)
+                loss = loss * sample_weight
+        else:
+            eps = self.label_smoothing
+            smooth = torch.full_like(probs, eps / n_classes)
+            smooth.scatter_(
+                1,
+                target.unsqueeze(1),
+                smooth.gather(1, target.unsqueeze(1)) + (1.0 - eps),
+            )
+            focal_per_class = (1.0 - probs).clamp_min(1e-12).pow(self.gamma)
+            per_class_loss = -smooth * focal_per_class * log_probs
+            loss = per_class_loss.sum(dim=-1)
+            if self.class_weight is not None:
+                sample_weight = self.class_weight.to(loss.device).gather(0, target)
+                loss = loss * sample_weight
+
+        if self.reduction == "mean":
+            if sample_weight is not None:
+                denom = sample_weight.sum().clamp_min(1e-12)
+                return loss.sum() / denom
+            return loss.mean()
+        if self.reduction == "sum":
+            return loss.sum()
+        return loss
+
+    def extra_repr(self) -> str:
+        """Return a string with the focal-loss hyperparameters for ``repr``."""
+        return (
+            f"gamma={self.gamma}, label_smoothing={self.label_smoothing}, "
+            f"reduction={self.reduction!r}"
+        )

maldideepkit/utils/lr_finder.py

@@ -0,0 +1,173 @@
+"""Learning-rate finder.
+
+Sweeps the learning rate geometrically over a small training run and
+records the smoothed loss at each step. The minimum of the smoothed
+loss curve's gradient gives a reasonable starting point for the base
+learning rate.
+
+Diagnostic only; not wired into
+:meth:`~maldideepkit.BaseSpectralClassifier.fit`.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+from torch import nn
+
+from .loss import FocalLoss
+from .reproducibility import resolve_device, seed_everything
+
+if TYPE_CHECKING:
+    from ..base.classifier import BaseSpectralClassifier
+
+
+def find_lr(
+    classifier: "BaseSpectralClassifier",
+    X: Any,
+    y: Any,
+    *,
+    start_lr: float = 1e-8,
+    end_lr: float = 1.0,
+    num_iter: int = 200,
+    smoothing: float = 0.98,
+    divergence_factor: float = 4.0,
+    plot: bool = False,
+) -> dict[str, Any]:
+    """Sweep learning rate geometrically and return the LR / loss curve.
+
+    Parameters
+    ----------
+    classifier : BaseSpectralClassifier
+        An unfitted classifier configured with the desired architecture
+        / batch_size / loss. Its ``learning_rate`` is ignored (we drive
+        it manually across the sweep) and its weights are reset at the
+        start of every call.
+    X, y : array-like
+        Training data. Only enough batches to cover ``num_iter`` steps
+        are consumed.
+    start_lr, end_lr : float, default=1e-8, 1.0
+        Bounds of the geometric LR sweep.
+    num_iter : int, default=200
+        Number of steps in the sweep.
+    smoothing : float, default=0.98
+        Exponential-moving-average factor applied to the per-step loss
+        (0 = no smoothing, 0.99 = heavy smoothing).
+    divergence_factor : float, default=4.0
+        Stop early once the smoothed loss exceeds
+        ``divergence_factor * min_smoothed_loss``.
+    plot : bool, default=False
+        If ``True``, render a matplotlib plot. ``matplotlib`` is only
+        imported when this is true.
+
+    Returns
+    -------
+    dict
+        ``{"lrs": np.ndarray, "losses": np.ndarray, "suggested_lr": float}``.
+        ``suggested_lr`` is the LR at the steepest-descent point of the
+        smoothed loss curve.
+    """
+    from ..base.data import make_loaders
+
+    seed_everything(int(classifier.random_state))
+    device = resolve_device(classifier.device)
+
+    X_np, y_encoded = classifier._prepare_inputs(X, y)
+    train_loader, _, stats = make_loaders(
+        X_np,
+        y_encoded,
+        batch_size=int(classifier.batch_size),
+        val_size=float(classifier.val_fraction),
+        random_state=int(classifier.random_state),
+        standardize=bool(classifier.standardize),
+    )
+    classifier.feature_mean_ = stats["mean"]
+    classifier.feature_std_ = stats["std"]
+    model = classifier._build_model().to(device)
+
+    class_weight = classifier._compute_class_weight(y_encoded)
+    if class_weight is not None:
+        class_weight = class_weight.to(device)
+    if classifier.loss == "cross_entropy":
+        criterion: nn.Module = nn.CrossEntropyLoss(
+            weight=class_weight, label_smoothing=float(classifier.label_smoothing)
+        )
+    else:
+        criterion = FocalLoss(
+            weight=class_weight,
+            gamma=float(classifier.focal_gamma),
+            label_smoothing=float(classifier.label_smoothing),
+        )
+    optimizer = torch.optim.Adam(model.parameters(), lr=start_lr)
+
+    lrs: list[float] = []
+    losses: list[float] = []
+    best_smoothed = float("inf")
+    smoothed = 0.0
+    log_step = (np.log(end_lr) - np.log(start_lr)) / max(1, num_iter - 1)
+
+    model.train()
+    step = 0
+    data_iter = iter(train_loader)
+    while step < num_iter:
+        try:
+            xb, yb = next(data_iter)
+        except StopIteration:
+            data_iter = iter(train_loader)
+            xb, yb = next(data_iter)
+        xb = xb.to(device, non_blocking=True)
+        yb = yb.to(device, non_blocking=True)
+
+        lr = float(np.exp(np.log(start_lr) + log_step * step))
+        for pg in optimizer.param_groups:
+            pg["lr"] = lr
+
+        optimizer.zero_grad()
+        logits = model(xb)
+        loss = criterion(logits, yb)
+        loss.backward()
+        optimizer.step()
+
+        raw = float(loss.detach().item())
+        smoothed = smoothing * smoothed + (1 - smoothing) * raw
+        debiased = smoothed / (1 - smoothing ** (step + 1)) if smoothing > 0 else raw
+
+        lrs.append(lr)
+        losses.append(debiased)
+
+        best_smoothed = min(best_smoothed, debiased)
+        if debiased > divergence_factor * best_smoothed and step > 5:
+            break
+
+        step += 1
+
+    lrs_arr = np.asarray(lrs)
+    losses_arr = np.asarray(losses)
+    if len(lrs_arr) < 3:
+        suggested_idx = int(np.argmin(losses_arr))
+    else:
+        grads = np.gradient(losses_arr, np.log(lrs_arr))
+        suggested_idx = int(np.argmin(grads))
+    suggested_lr = float(lrs_arr[suggested_idx])
+
+    if plot:
+        import matplotlib.pyplot as plt
+
+        fig, ax = plt.subplots(figsize=(6, 4))
+        ax.plot(lrs_arr, losses_arr)
+        ax.set_xscale("log")
+        ax.set_xlabel("learning rate")
+        ax.set_ylabel("smoothed loss")
+        ax.axvline(
+            suggested_lr,
+            color="r",
+            linestyle="--",
+            label=f"suggested = {suggested_lr:.2e}",
+        )
+        ax.legend()
+        fig.tight_layout()
+        plt.show()
+
+    return {"lrs": lrs_arr, "losses": losses_arr, "suggested_lr": suggested_lr}

maldideepkit/utils/reproducibility.py

@@ -0,0 +1,70 @@
+"""Seeding and device-placement helpers for deterministic training."""
+
+from __future__ import annotations
+
+import os
+import random
+
+import numpy as np
+import torch
+
+
+def seed_everything(seed: int, deterministic: bool = False) -> None:
+    """Seed Python, NumPy, and PyTorch (CPU + CUDA) RNGs in one call.
+
+    Parameters
+    ----------
+    seed : int
+        Non-negative integer used for every RNG. Also fixes
+        ``PYTHONHASHSEED`` in the current process environment.
+    deterministic : bool, default=False
+        When ``True``, additionally enable PyTorch's deterministic
+        algorithm mode. Sets
+        ``torch.use_deterministic_algorithms(True, warn_only=True)``,
+        ``torch.backends.cudnn.deterministic = True``,
+        ``torch.backends.cudnn.benchmark = False``, and
+        ``CUBLAS_WORKSPACE_CONFIG=:4096:8``. The env-var must be set
+        before the first CUDA context is created. Once enabled,
+        determinism is **sticky** - subsequent plain
+        ``seed_everything(seed)`` calls do not turn it off.
+    """
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    if deterministic:
+        os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+        torch.use_deterministic_algorithms(True, warn_only=True)
+        if torch.cuda.is_available():
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
+
+
+def resolve_device(device: str | torch.device | None) -> torch.device:
+    """Resolve a user-facing device specifier to a :class:`torch.device`.
+
+    Parameters
+    ----------
+    device : {"auto", "cpu", "cuda"} or torch.device or None
+        ``"auto"`` (or ``None``) picks ``cuda`` when available and falls
+        back to ``cpu``.
+
+    Returns
+    -------
+    torch.device
+        The resolved device.
+
+    Raises
+    ------
+    ValueError
+        If ``device`` is an unknown string.
+    """
+    if device is None or device == "auto":
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if isinstance(device, torch.device):
+        return device
+    if isinstance(device, str):
+        return torch.device(device)
+    raise ValueError(f"Unsupported device specifier: {device!r}")

maldideepkit/utils/sam.py

@@ -0,0 +1,121 @@
+"""Sharpness-Aware Minimization (SAM) optimizer wrapper.
+
+SAM pushes weights toward flatter regions of the loss landscape, at
+the cost of ~2x the forward / backward compute per step.
+
+Usage
+-----
+
+.. code-block:: python
+
+    optimizer = SAMOptimizer(
+        model.parameters(), base_optimizer=torch.optim.AdamW,
+        rho=0.05, lr=1e-3, weight_decay=0.05,
+    )
+
+    loss = criterion(model(x), y)
+    loss.backward()
+    optimizer.first_step(zero_grad=True)
+
+    loss = criterion(model(x), y)
+    loss.backward()
+    optimizer.second_step(zero_grad=True)
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+
+class SAMOptimizer(torch.optim.Optimizer):
+    """Wrap a base optimizer in the SAM two-step update.
+
+    Parameters
+    ----------
+    params : iterable
+        Parameters or param-group dicts (as for any torch optimizer).
+    base_optimizer : type
+        The base optimizer **class** (e.g. :class:`torch.optim.AdamW`).
+        Instantiated internally against the same param groups.
+    rho : float, default=0.05
+        Size of the ascent step in parameter space. Paper default is
+        ``0.05``. Typical range: ``[0.01, 0.2]``.
+    **base_kwargs
+        Forwarded to the base optimizer (e.g. ``lr``, ``weight_decay``).
+    """
+
+    def __init__(
+        self,
+        params: Any,
+        base_optimizer: type[torch.optim.Optimizer],
+        rho: float = 0.05,
+        **base_kwargs: Any,
+    ) -> None:
+        if rho <= 0:
+            raise ValueError(f"rho must be > 0; got {rho!r}.")
+        defaults = {"rho": float(rho), **base_kwargs}
+        super().__init__(params, defaults)
+        self.base_optimizer = base_optimizer(self.param_groups, **base_kwargs)
+        self.param_groups = self.base_optimizer.param_groups
+        for g in self.param_groups:
+            g.setdefault("rho", float(rho))
+
+    @torch.no_grad()
+    def first_step(self, zero_grad: bool = False) -> None:
+        """Ascend to ``w + e`` using the current gradients."""
+        grad_norm = self._grad_norm()
+        eps = 1e-12
+        for group in self.param_groups:
+            scale = group["rho"] / (grad_norm + eps)
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                e_w = p.grad * scale
+                self.state[p]["e_w"] = e_w
+                p.add_(e_w)
+        if zero_grad:
+            self.zero_grad()
+
+    @torch.no_grad()
+    def second_step(self, zero_grad: bool = False) -> None:
+        """Undo the ascent and apply the base optimizer step from ``w``."""
+        for group in self.param_groups:
+            for p in group["params"]:
+                if "e_w" in self.state.get(p, {}):
+                    p.sub_(self.state[p]["e_w"])
+                    del self.state[p]["e_w"]
+        self.base_optimizer.step()
+        self._step_count = getattr(self.base_optimizer, "_step_count", 0)
+        self._opt_called = True
+        if zero_grad:
+            self.zero_grad()
+
+    def step(self, closure: Any = None) -> Any:
+        """Unsupported. Use ``first_step`` / ``second_step`` instead."""
+        raise RuntimeError(
+            "SAMOptimizer requires an explicit two-pass training loop. "
+            "Call first_step() after the first backward, then recompute "
+            "the loss and backward, then call second_step()."
+        )
+
+    @torch.no_grad()
+    def _grad_norm(self) -> torch.Tensor:
+        ref_device = None
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is not None:
+                    ref_device = p.grad.device
+                    break
+            if ref_device is not None:
+                break
+        if ref_device is None:
+            return torch.tensor(0.0)
+        norms = [
+            p.grad.norm(p=2).to(ref_device)
+            for group in self.param_groups
+            for p in group["params"]
+            if p.grad is not None
+        ]
+        return torch.norm(torch.stack(norms), p=2)