phasmanet 0.1.0__py3-none-any.whl

phasmanet/estimators.py

@@ -0,0 +1,680 @@
+"""scikit-learn style estimators around the PhasmaNet architecture.
+
+``PhasmaNetRegressor`` and ``PhasmaNetClassifier`` follow the fit / predict convention,
+so they work with scikit-learn tools such as ``cross_val_score``, ``GridSearchCV`` and
+``Pipeline`` when scikit-learn is installed (it is optional). Every hyperparameter of the
+architecture and of the training loop is a constructor argument; the defaults are the
+configuration used in the paper.
+"""
+from __future__ import annotations
+
+import copy
+import inspect
+import math
+from typing import Any, Dict, Optional
+
+import numpy as np
+import torch
+from torch import nn
+
+from .calibration import MonotoneCalibrator
+from .model import PhasmaNet
+from ._version import __version__
+
+try:  # scikit-learn is optional
+    from sklearn.base import BaseEstimator as _SkBase, ClassifierMixin as _SkClf, RegressorMixin as _SkReg
+except ImportError:  # pragma: no cover - exercised only without scikit-learn
+
+    class _SkBase:  # minimal get_params / set_params compatible with sklearn.clone
+        @classmethod
+        def _get_param_names(cls):
+            sig = inspect.signature(cls.__init__)
+            return sorted(p for p in sig.parameters if p != "self")
+
+        def get_params(self, deep: bool = True) -> Dict[str, Any]:
+            return {k: getattr(self, k) for k in self._get_param_names()}
+
+        def set_params(self, **params):
+            valid = self._get_param_names()
+            for k, v in params.items():
+                if k not in valid:
+                    raise ValueError(f"Invalid parameter {k!r} for {type(self).__name__}")
+                setattr(self, k, v)
+            return self
+
+    class _SkReg:
+        _estimator_type = "regressor"
+
+    class _SkClf:
+        _estimator_type = "classifier"
+
+
+__all__ = ["PhasmaNetRegressor", "PhasmaNetClassifier"]
+
+_FORMAT = "phasmanet-estimator"
+_ARCH_PARAMS = ("hidden_dim", "num_tokens", "num_stages", "num_feedback", "heads", "dropout",
+                "num_refractor_subspaces", "use_refractor", "ffn_mult", "activation", "attn_dropout")
+
+
+def _resolve_device(device: Optional[str]) -> torch.device:
+    if device is None or device == "auto":
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return torch.device(device)
+
+
+def _as_2d_float(X) -> np.ndarray:
+    X = np.asarray(X, dtype=np.float32)
+    if X.ndim == 1:
+        X = X.reshape(-1, 1)
+    if X.ndim != 2:
+        raise ValueError(f"X must be 2-dimensional, got shape {X.shape}")
+    if not np.all(np.isfinite(X)):
+        raise ValueError("X contains NaN or infinite values")
+    return X
+
+
+class _PhasmaNetBase(_SkBase):
+    """Shared training loop. Subclasses define the targets, loss and outputs.
+
+    Subclasses must store every constructor argument as an attribute of the same name.
+    """
+
+    # ----- hooks implemented by subclasses -----
+    def _prepare_targets(self, y, refit: bool = True) -> np.ndarray:
+        raise NotImplementedError
+
+    def _fit_target_scaler(self, Y_train: np.ndarray) -> None:
+        pass
+
+    def _scale_targets(self, Y: np.ndarray) -> np.ndarray:
+        return Y
+
+    def _scheme_weights(self, Y_train: np.ndarray) -> Optional[np.ndarray]:
+        """Extra per-sample weights from a weighting scheme (None for uniform)."""
+        return None
+
+    def _output_dim(self) -> int:
+        raise NotImplementedError
+
+    def _loss(self, out: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        """Per-sample loss, shape (B,)."""
+        raise NotImplementedError
+
+    def _after_fit(self, X_val: Optional[np.ndarray], y_val: Optional[np.ndarray]) -> None:
+        pass
+
+    # ----- common machinery -----
+    def _architecture(self) -> Dict[str, Any]:
+        return {k: getattr(self, k) for k in _ARCH_PARAMS}
+
+    def _build_module(self) -> PhasmaNet:
+        return PhasmaNet(input_dim=self.n_features_in_, output_dim=self._output_dim(), **self._architecture())
+
+    def _make_optimizer(self, params):
+        name = str(self.optimizer).lower()
+        if name == "adamw":
+            return torch.optim.AdamW(params, lr=self.lr, weight_decay=self.weight_decay)
+        if name == "adam":
+            return torch.optim.Adam(params, lr=self.lr, weight_decay=self.weight_decay)
+        if name == "sgd":
+            return torch.optim.SGD(params, lr=self.lr, weight_decay=self.weight_decay, momentum=self.momentum)
+        raise ValueError("optimizer must be 'adamw', 'adam' or 'sgd'")
+
+    def _make_scheduler(self, opt):
+        name = str(self.scheduler).lower()
+        if name == "plateau":
+            return torch.optim.lr_scheduler.ReduceLROnPlateau(opt, mode="min", factor=self.lr_factor,
+                                                              patience=self.lr_patience)
+        if name == "cosine":
+            return torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(1, self.max_epochs))
+        if name == "none":
+            return None
+        raise ValueError("scheduler must be 'plateau', 'cosine' or 'none'")
+
+    def _check_X(self, X, reset: bool) -> np.ndarray:
+        names = getattr(X, "columns", None)
+        X = _as_2d_float(X)
+        if reset:
+            self.n_features_in_ = X.shape[1]
+            if names is not None:
+                self.feature_names_in_ = np.asarray([str(c) for c in names], dtype=object)
+        elif X.shape[1] != self.n_features_in_:
+            raise ValueError(f"X has {X.shape[1]} features, but the model was fitted with {self.n_features_in_}")
+        return X
+
+    def _scale_X(self, X: np.ndarray) -> np.ndarray:
+        return ((X - self.x_mean_) / self.x_std_).astype(np.float32)
+
+    def _check_is_fitted(self) -> None:
+        if not hasattr(self, "module_"):
+            raise RuntimeError(f"This {type(self).__name__} instance is not fitted yet; call fit first.")
+
+    def fit(self, X, y, sample_weight=None, validation_data=None):
+        """Train the network with early stopping on a validation set.
+
+        Args:
+            X: array-like of shape (n_samples, n_features).
+            y: targets.
+            sample_weight: optional per-sample loss weights (training split only).
+            validation_data: optional ``(X_val, y_val)``. If omitted, a random
+                ``validation_fraction`` of the training data is held out.
+        """
+        X = self._check_X(X, reset=True)
+        Y = self._prepare_targets(y)
+        if len(X) != len(Y):
+            raise ValueError("X and y have different numbers of samples")
+        w = None if sample_weight is None else np.asarray(sample_weight, dtype=np.float64).ravel()
+        if w is not None and len(w) != len(X):
+            raise ValueError("sample_weight has the wrong length")
+
+        rng = np.random.default_rng(self.random_state)
+        if self.random_state is not None:
+            torch.manual_seed(self.random_state)
+
+        if validation_data is not None:
+            X_tr, Y_tr, w_tr = X, Y, w
+            X_val = self._check_X(validation_data[0], reset=False)
+            Y_val = self._prepare_targets(validation_data[1], refit=False)
+        elif self.validation_fraction and self.validation_fraction > 0:
+            n_val = max(1, int(round(self.validation_fraction * len(X))))
+            if n_val >= len(X):
+                raise ValueError("validation_fraction leaves no training samples")
+            perm = rng.permutation(len(X))
+            val_idx, tr_idx = perm[:n_val], perm[n_val:]
+            X_tr, Y_tr, X_val, Y_val = X[tr_idx], Y[tr_idx], X[val_idx], Y[val_idx]
+            w_tr = None if w is None else w[tr_idx]
+        else:
+            X_tr, Y_tr, w_tr, X_val, Y_val = X, Y, w, None, None
+
+        # Scalers from the training split only (float64 accumulation, then stored as float32).
+        self.x_mean_ = X_tr.mean(axis=0, dtype=np.float64).astype(np.float32)
+        self.x_std_ = (X_tr.std(axis=0, dtype=np.float64) + 1e-9).astype(np.float32)
+        self._fit_target_scaler(Y_tr)
+        Y_tr_t = self._scale_targets(Y_tr)
+        Y_val_t = None if Y_val is None else self._scale_targets(Y_val)
+
+        scheme = self._scheme_weights(Y_tr)
+        if scheme is not None:
+            w_tr = scheme if w_tr is None else w_tr * scheme
+        Wt = None if w_tr is None else torch.from_numpy((w_tr / (w_tr.mean() + 1e-12)).astype(np.float32))
+
+        device = _resolve_device(self.device)
+        self.module_ = self._build_module().to(device)
+        opt = self._make_optimizer(self.module_.parameters())
+        sched = self._make_scheduler(opt)
+
+        Xt = torch.from_numpy(self._scale_X(X_tr))
+        Yt = torch.from_numpy(Y_tr_t)
+        n = len(Xt)
+
+        best_loss, best_state, bad = math.inf, None, 0
+        self.history_ = {"train_loss": [], "val_loss": [], "lr": []}
+        epoch = 0
+        for epoch in range(1, self.max_epochs + 1):
+            self.module_.train()
+            order = torch.from_numpy(rng.permutation(n))
+            total = 0.0
+            for start in range(0, n, self.batch_size):
+                idx = order[start:start + self.batch_size]
+                xb, yb = Xt[idx].to(device), Yt[idx].to(device)
+                opt.zero_grad()
+                loss_vec = self._loss(self.module_(xb), yb)
+                if Wt is not None:
+                    loss_vec = loss_vec * Wt[idx].to(device)
+                loss = loss_vec.mean()
+                loss.backward()
+                if self.grad_clip:
+                    nn.utils.clip_grad_norm_(self.module_.parameters(), self.grad_clip)
+                opt.step()
+                total += loss.item() * len(idx)
+            train_loss = total / n
+            self.history_["train_loss"].append(train_loss)
+            self.history_["lr"].append(opt.param_groups[0]["lr"])
+
+            if X_val is not None:
+                monitor = self._eval_loss(self._scale_X(X_val), Y_val_t, device)
+                self.history_["val_loss"].append(monitor)
+            else:
+                monitor = train_loss
+            if isinstance(sched, torch.optim.lr_scheduler.ReduceLROnPlateau):
+                sched.step(monitor)
+            elif sched is not None:
+                sched.step()
+
+            if self.verbose:
+                msg = f"epoch {epoch:4d}  train {train_loss:.6f}"
+                if X_val is not None:
+                    msg += f"  val {monitor:.6f}"
+                print(msg)
+
+            if monitor + self.min_delta < best_loss:
+                best_loss, bad = monitor, 0
+                best_state = copy.deepcopy(self.module_.state_dict())
+            else:
+                bad += 1
+                if self.patience and bad >= self.patience:
+                    if self.verbose:
+                        print("early stopping")
+                    break
+
+        if best_state is not None:
+            self.module_.load_state_dict(best_state)
+        self.module_.eval()
+        self.n_epochs_ = epoch
+        self.best_loss_ = best_loss
+        self._after_fit(X_val, Y_val)
+        return self
+
+    def _eval_loss(self, Xs: np.ndarray, Ys: np.ndarray, device: torch.device) -> float:
+        self.module_.eval()
+        total = 0.0
+        with torch.no_grad():
+            for start in range(0, len(Xs), 4096):
+                xb = torch.from_numpy(Xs[start:start + 4096]).to(device)
+                yb = torch.from_numpy(Ys[start:start + 4096]).to(device)
+                total += self._loss(self.module_(xb), yb).sum().item()
+        return total / len(Xs)
+
+    def _forward_numpy(self, X, batch_size: int = 4096) -> np.ndarray:
+        self._check_is_fitted()
+        Xs = self._scale_X(self._check_X(X, reset=False))
+        device = next(self.module_.parameters()).device
+        self.module_.eval()
+        outs = []
+        with torch.inference_mode():
+            for start in range(0, len(Xs), batch_size):
+                outs.append(self.module_(torch.from_numpy(Xs[start:start + batch_size]).to(device)).cpu().numpy())
+        return np.concatenate(outs, axis=0) if outs else np.zeros((0, self._output_dim()), dtype=np.float32)
+
+    # ----- persistence -----
+    def _extra_state(self) -> Dict[str, Any]:
+        return {}
+
+    def _load_extra_state(self, state: Dict[str, Any]) -> None:
+        pass
+
+    def save(self, path) -> None:
+        """Save a fitted estimator (hyperparameters, weights and scalers) to ``path``."""
+        self._check_is_fitted()
+        payload = {
+            "format": _FORMAT,
+            "phasmanet_version": __version__,
+            "class": type(self).__name__,
+            "params": self.get_params(),
+            "n_features_in": int(self.n_features_in_),
+            "feature_names_in": None if not hasattr(self, "feature_names_in_") else [str(c) for c in self.feature_names_in_],
+            "x_mean": torch.from_numpy(self.x_mean_),
+            "x_std": torch.from_numpy(self.x_std_),
+            "state_dict": {k: v.detach().cpu() for k, v in self.module_.state_dict().items()},
+            "extra": self._extra_state(),
+        }
+        torch.save(payload, path)
+
+    @classmethod
+    def load(cls, path, device: Optional[str] = "auto"):
+        """Load an estimator saved with :meth:`save`."""
+        payload = torch.load(path, map_location="cpu", weights_only=True)
+        if payload.get("format") != _FORMAT:
+            raise ValueError(f"{path} is not a saved PhasmaNet estimator")
+        if payload["class"] != cls.__name__:
+            raise ValueError(f"{path} holds a {payload['class']}, not a {cls.__name__}")
+        valid = set(inspect.signature(cls.__init__).parameters) - {"self"}
+        est = cls(**{k: v for k, v in payload["params"].items() if k in valid})
+        est.device = device
+        est.n_features_in_ = payload["n_features_in"]
+        if payload.get("feature_names_in"):
+            est.feature_names_in_ = np.asarray(payload["feature_names_in"], dtype=object)
+        est.x_mean_ = payload["x_mean"].numpy()
+        est.x_std_ = payload["x_std"].numpy()
+        est._load_extra_state(payload["extra"])
+        est.module_ = est._build_module()
+        est.module_.load_state_dict(payload["state_dict"])
+        est.module_.to(_resolve_device(device)).eval()
+        return est
+
+
+_COMMON_DOC = """
+    Architecture:
+        hidden_dim (128), num_tokens (8), num_stages (3), num_feedback (2), heads (4),
+        dropout (0.1), num_refractor_subspaces (4), use_refractor (True), ffn_mult (4),
+        activation ("gelu"; also "relu", "silu", "tanh", "elu", "leaky_relu"),
+        attn_dropout (0.0).
+    Training:
+        optimizer ("adamw", "adam" or "sgd"), lr (1e-3), weight_decay (1e-4),
+        momentum (0.9, SGD only), batch_size (128), max_epochs (150), patience (20,
+        early stopping; 0 disables it), min_delta (1e-6), scheduler ("plateau",
+        "cosine" or "none"), lr_patience (5) and lr_factor (0.5) for "plateau",
+        grad_clip (None, max gradient norm), validation_fraction (0.15).
+    Other:
+        device ("auto", "cpu", "cuda", ...), random_state (None), verbose (0).
+"""
+
+
+class PhasmaNetRegressor(_SkReg, _PhasmaNetBase):
+    """PhasmaNet for (multi-output) regression.
+
+    Targets are standardised internally. Regression options:
+        loss ("mse" or "huber"), huber_delta (1.0, in units of the standardised target),
+        weighting ("none", "by_hist" or "inverse_freq"; see below), hist_bins (20),
+        wmin (0.2), wmax (5.0), calibrate (False; monotonic PCHIP calibration fitted on
+        the validation split, single-output only), calibration_bins (60),
+        calibration_min_bin (20).
+
+    Weighting schemes over a histogram of the (first) training target:
+        "by_hist": the formula of the paper's training script, w = clip(1 / bin_count,
+        wmin, wmax), then normalised to mean 1.
+        "inverse_freq": w proportional to 1 / bin_count normalised to mean 1 first, then
+        clipped to [wmin, wmax] and renormalised, so rare target ranges are up-weighted.
+
+    Example:
+        >>> from phasmanet import PhasmaNetRegressor
+        >>> reg = PhasmaNetRegressor(num_stages=4, max_epochs=300, random_state=0).fit(X_train, y_train)
+        >>> y_pred = reg.predict(X_test)
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int = 128,
+        num_tokens: int = 8,
+        num_stages: int = 3,
+        num_feedback: int = 2,
+        heads: int = 4,
+        dropout: float = 0.1,
+        num_refractor_subspaces: int = 4,
+        use_refractor: bool = True,
+        ffn_mult: int = 4,
+        activation: str = "gelu",
+        attn_dropout: float = 0.0,
+        loss: str = "mse",
+        huber_delta: float = 1.0,
+        weighting: str = "none",
+        hist_bins: int = 20,
+        wmin: float = 0.2,
+        wmax: float = 5.0,
+        calibrate: bool = False,
+        calibration_bins: int = 60,
+        calibration_min_bin: int = 20,
+        optimizer: str = "adamw",
+        lr: float = 1e-3,
+        weight_decay: float = 1e-4,
+        momentum: float = 0.9,
+        batch_size: int = 128,
+        max_epochs: int = 150,
+        patience: int = 20,
+        min_delta: float = 1e-6,
+        scheduler: str = "plateau",
+        lr_patience: int = 5,
+        lr_factor: float = 0.5,
+        grad_clip: Optional[float] = None,
+        validation_fraction: float = 0.15,
+        device: Optional[str] = "auto",
+        random_state: Optional[int] = None,
+        verbose: int = 0,
+    ):
+        self.hidden_dim = hidden_dim
+        self.num_tokens = num_tokens
+        self.num_stages = num_stages
+        self.num_feedback = num_feedback
+        self.heads = heads
+        self.dropout = dropout
+        self.num_refractor_subspaces = num_refractor_subspaces
+        self.use_refractor = use_refractor
+        self.ffn_mult = ffn_mult
+        self.activation = activation
+        self.attn_dropout = attn_dropout
+        self.loss = loss
+        self.huber_delta = huber_delta
+        self.weighting = weighting
+        self.hist_bins = hist_bins
+        self.wmin = wmin
+        self.wmax = wmax
+        self.calibrate = calibrate
+        self.calibration_bins = calibration_bins
+        self.calibration_min_bin = calibration_min_bin
+        self.optimizer = optimizer
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.momentum = momentum
+        self.batch_size = batch_size
+        self.max_epochs = max_epochs
+        self.patience = patience
+        self.min_delta = min_delta
+        self.scheduler = scheduler
+        self.lr_patience = lr_patience
+        self.lr_factor = lr_factor
+        self.grad_clip = grad_clip
+        self.validation_fraction = validation_fraction
+        self.device = device
+        self.random_state = random_state
+        self.verbose = verbose
+
+    def _prepare_targets(self, y, refit: bool = True) -> np.ndarray:
+        Y = np.asarray(y, dtype=np.float32)
+        if refit:
+            self._y_was_1d = Y.ndim == 1
+        Y = Y.reshape(len(Y), -1)
+        if not np.all(np.isfinite(Y)):
+            raise ValueError("y contains NaN or infinite values")
+        if refit:
+            self.n_outputs_ = Y.shape[1]
+        elif Y.shape[1] != self.n_outputs_:
+            raise ValueError("validation targets have a different number of outputs")
+        return Y
+
+    def _fit_target_scaler(self, Y_train: np.ndarray) -> None:
+        self.y_mean_ = Y_train.mean(axis=0, dtype=np.float64).astype(np.float32)
+        self.y_std_ = (Y_train.std(axis=0, dtype=np.float64) + 1e-9).astype(np.float32)
+
+    def _scale_targets(self, Y: np.ndarray) -> np.ndarray:
+        return ((Y - self.y_mean_) / self.y_std_).astype(np.float32)
+
+    def _scheme_weights(self, Y_train: np.ndarray) -> Optional[np.ndarray]:
+        scheme = str(self.weighting).lower()
+        if scheme == "none":
+            return None
+        if scheme not in ("by_hist", "inverse_freq"):
+            raise ValueError("weighting must be 'none', 'by_hist' or 'inverse_freq'")
+        y = Y_train[:, 0].astype(np.float64)
+        hist, edges = np.histogram(y, bins=self.hist_bins)
+        hist = np.maximum(hist, 1)
+        idx = np.clip(np.searchsorted(edges, y, side="right") - 1, 0, len(hist) - 1)
+        w = 1.0 / hist[idx]
+        if scheme == "inverse_freq":
+            w = w / (w.mean() + 1e-12)
+        w = np.clip(w, self.wmin, self.wmax)
+        return w / (w.mean() + 1e-12)
+
+    def fit(self, X, y, sample_weight=None, validation_data=None):
+        if self.loss not in ("mse", "huber"):
+            raise ValueError("loss must be 'mse' or 'huber'")
+        return super().fit(X, y, sample_weight=sample_weight, validation_data=validation_data)
+
+    def _output_dim(self) -> int:
+        return self.n_outputs_
+
+    def _loss(self, out, target):
+        if self.loss == "mse":
+            per = (out - target) ** 2
+        else:
+            per = nn.functional.smooth_l1_loss(out, target, beta=self.huber_delta, reduction="none")
+        return per.mean(dim=1)
+
+    def _after_fit(self, X_val, Y_val) -> None:
+        self.calibrator_ = MonotoneCalibrator()
+        if self.calibrate:
+            if X_val is None:
+                raise ValueError("calibrate=True needs a validation set (validation_fraction > 0 or validation_data)")
+            if self.n_outputs_ != 1:
+                raise ValueError("calibrate=True supports single-output regression only")
+            raw = self._predict_raw(X_val)[:, 0]
+            # The requested number of quantile bins, reduced for small validation sets.
+            bins = int(min(self.calibration_bins, max(3, len(raw) // max(1, self.calibration_min_bin))))
+            self.calibrator_.fit(raw, Y_val[:, 0], bins=bins, min_bin=self.calibration_min_bin)
+
+    def _predict_raw(self, X) -> np.ndarray:
+        out = self._forward_numpy(X)
+        return out * self.y_std_ + self.y_mean_
+
+    def predict(self, X, calibrated: bool = True) -> np.ndarray:
+        """Predict targets. Set ``calibrated=False`` to skip a fitted calibrator."""
+        Y = self._predict_raw(X)
+        cal = getattr(self, "calibrator_", None)
+        if calibrated and cal is not None and not cal.is_identity:
+            Y = cal(Y[:, 0]).reshape(-1, 1)
+        return Y[:, 0] if getattr(self, "_y_was_1d", True) and Y.shape[1] == 1 else Y
+
+    def score(self, X, y, sample_weight=None) -> float:
+        """Coefficient of determination R^2 (uniform average over outputs)."""
+        y = np.asarray(y, dtype=np.float64).reshape(len(np.asarray(y)), -1)
+        p = np.asarray(self.predict(X), dtype=np.float64).reshape(len(y), -1)
+        w = None if sample_weight is None else np.asarray(sample_weight, dtype=np.float64).reshape(-1, 1)
+        if w is None:
+            ss_res = ((y - p) ** 2).sum(axis=0)
+            ss_tot = ((y - y.mean(axis=0)) ** 2).sum(axis=0)
+        else:
+            mean = (w * y).sum(axis=0) / w.sum()
+            ss_res = (w * (y - p) ** 2).sum(axis=0)
+            ss_tot = (w * (y - mean) ** 2).sum(axis=0)
+        r2 = np.where(ss_tot > 0, 1.0 - ss_res / np.where(ss_tot > 0, ss_tot, 1.0), 0.0)
+        return float(r2.mean())
+
+    def _extra_state(self):
+        return {
+            "y_mean": torch.from_numpy(self.y_mean_),
+            "y_std": torch.from_numpy(self.y_std_),
+            "n_outputs": int(self.n_outputs_),
+            "y_was_1d": bool(self._y_was_1d),
+            "calibrator": self.calibrator_.to_dict(),
+        }
+
+    def _load_extra_state(self, state):
+        self.y_mean_ = state["y_mean"].numpy()
+        self.y_std_ = state["y_std"].numpy()
+        self.n_outputs_ = state["n_outputs"]
+        self._y_was_1d = state["y_was_1d"]
+        self.calibrator_ = MonotoneCalibrator.from_dict(state["calibrator"])
+
+
+class PhasmaNetClassifier(_SkClf, _PhasmaNetBase):
+    """PhasmaNet for binary and multiclass classification (cross-entropy on logits).
+
+    Classification option:
+        class_weight (None or "balanced": weights inversely proportional to class
+        frequencies in the training split).
+
+    Example:
+        >>> from phasmanet import PhasmaNetClassifier
+        >>> clf = PhasmaNetClassifier(max_epochs=50, random_state=0).fit(X_train, y_train)
+        >>> proba = clf.predict_proba(X_test)
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int = 128,
+        num_tokens: int = 8,
+        num_stages: int = 3,
+        num_feedback: int = 2,
+        heads: int = 4,
+        dropout: float = 0.1,
+        num_refractor_subspaces: int = 4,
+        use_refractor: bool = True,
+        ffn_mult: int = 4,
+        activation: str = "gelu",
+        attn_dropout: float = 0.0,
+        class_weight: Optional[str] = None,
+        optimizer: str = "adamw",
+        lr: float = 1e-3,
+        weight_decay: float = 1e-4,
+        momentum: float = 0.9,
+        batch_size: int = 128,
+        max_epochs: int = 150,
+        patience: int = 20,
+        min_delta: float = 1e-6,
+        scheduler: str = "plateau",
+        lr_patience: int = 5,
+        lr_factor: float = 0.5,
+        grad_clip: Optional[float] = None,
+        validation_fraction: float = 0.15,
+        device: Optional[str] = "auto",
+        random_state: Optional[int] = None,
+        verbose: int = 0,
+    ):
+        self.hidden_dim = hidden_dim
+        self.num_tokens = num_tokens
+        self.num_stages = num_stages
+        self.num_feedback = num_feedback
+        self.heads = heads
+        self.dropout = dropout
+        self.num_refractor_subspaces = num_refractor_subspaces
+        self.use_refractor = use_refractor
+        self.ffn_mult = ffn_mult
+        self.activation = activation
+        self.attn_dropout = attn_dropout
+        self.class_weight = class_weight
+        self.optimizer = optimizer
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.momentum = momentum
+        self.batch_size = batch_size
+        self.max_epochs = max_epochs
+        self.patience = patience
+        self.min_delta = min_delta
+        self.scheduler = scheduler
+        self.lr_patience = lr_patience
+        self.lr_factor = lr_factor
+        self.grad_clip = grad_clip
+        self.validation_fraction = validation_fraction
+        self.device = device
+        self.random_state = random_state
+        self.verbose = verbose
+
+    def _prepare_targets(self, y, refit: bool = True) -> np.ndarray:
+        y = np.asarray(y).ravel()
+        if refit:
+            self.classes_ = np.unique(y)
+            if len(self.classes_) < 2:
+                raise ValueError("need at least two classes")
+        idx = np.clip(np.searchsorted(self.classes_, y), 0, len(self.classes_) - 1)
+        if not np.all(self.classes_[idx] == y):
+            raise ValueError("y contains labels not seen during fit")
+        return idx.astype(np.int64)
+
+    def _scheme_weights(self, Y_train: np.ndarray) -> Optional[np.ndarray]:
+        if self.class_weight is None:
+            return None
+        if self.class_weight != "balanced":
+            raise ValueError("class_weight must be None or 'balanced'")
+        counts = np.bincount(Y_train, minlength=len(self.classes_)).astype(np.float64)
+        per_class = len(Y_train) / (len(self.classes_) * np.maximum(counts, 1.0))
+        return per_class[Y_train]
+
+    def _output_dim(self) -> int:
+        return len(self.classes_)
+
+    def _loss(self, out, target):
+        return nn.functional.cross_entropy(out, target, reduction="none")
+
+    def predict_proba(self, X) -> np.ndarray:
+        logits = torch.from_numpy(self._forward_numpy(X))
+        return torch.softmax(logits, dim=1).numpy()
+
+    def predict(self, X) -> np.ndarray:
+        return self.classes_[np.argmax(self._forward_numpy(X), axis=1)]
+
+    def score(self, X, y, sample_weight=None) -> float:
+        """Mean accuracy."""
+        correct = (self.predict(X) == np.asarray(y)).astype(np.float64)
+        if sample_weight is None:
+            return float(correct.mean())
+        w = np.asarray(sample_weight, dtype=np.float64)
+        return float((w * correct).sum() / w.sum())
+
+    def _extra_state(self):
+        return {"classes": self.classes_.tolist()}
+
+    def _load_extra_state(self, state):
+        self.classes_ = np.asarray(state["classes"])
+
+
+PhasmaNetRegressor.__doc__ += _COMMON_DOC
+PhasmaNetClassifier.__doc__ += _COMMON_DOC