baselines.py

# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 The Regents of the University of California, Mazin Lab

"""
Learned baselines for VENOM.

These exist so the paper can answer reviewer questions of the form
"is your fancy SSM actually doing anything?". Each baseline is matched to
VENOM's parameter budget (~3.3 k) and trained on the same data with the
same train/val split, optimizer, and Gaussian-NLL loss.

Currently provides:

- ``extract_pulse_features`` + ``FeatureMLP``
    A pulse-by-pulse MLP on hand-engineered features
    (peak height, rise/decay times, integral, pre-trigger RMS, FWHM,
    half-decay). This is what an MKID experimentalist actually deploys
    today via ``mkidcalculator``-style optimal filtering — anything we
    do should beat it.

Future additions:

- Non-selective S4D ablation (planned in mkid_ssm.py via a flag).
- Causal TCN backbone with a ring-buffer ``step`` (planned here).
"""
from __future__ import annotations

import math
from typing import List, Tuple

import numpy as np

from backend import BACKEND, mx, nn, no_grad, optim, to_numpy

# ---------------------------------------------------------------------------
# Feature extraction
# ---------------------------------------------------------------------------

FEATURE_NAMES: List[str] = [
    "peak_height",   # max(|IQ - baseline|) over the post-trigger window
    "peak_idx",      # absolute index of the peak within the windowed trace
    "rise_time",     # 10%-90% rise (samples)
    "decay_time",    # 90%-10% decay (samples)
    "integral",      # sum of max(|IQ| - baseline, 0) over the window
    "pre_rms",       # std of |IQ| in the pre-trigger baseline samples
    "fwhm",          # full width at half max (samples)
    "half_decay",    # samples from peak to 50% on the falling edge
]


def extract_pulse_features(
    iq_windows: np.ndarray,
    pre_trigger_n: int = 20,
    peak_search_n: int | None = None,
    smooth_n: int = 0,
) -> Tuple[np.ndarray, List[str]]:
    """Compute hand-engineered features from baseline-subtracted IQ traces.

    The input is expected to be the output of ``load_mkid_data`` /
    ``load_ptsi_data``: shape ``(n_pulses, seq_len, 2)``, with the
    cross-pulse mean baseline subtracted and the global IQ radius set to
    one. We use the post-baseline magnitude ``r(t) = sqrt(I^2 + Q^2)`` as
    the 1-D pulse trace.

    Parameters
    ----------
    iq_windows : np.ndarray, shape (n, T, 2)
    pre_trigger_n : int
        Leading samples treated as the pre-trigger baseline (used for
        the noise RMS estimate and as the lower bound of the peak
        search). 20 works for the windows produced by ``load_mkid_data``.
    peak_search_n : int or None
        Width of the peak-search window starting just before
        ``pre_trigger_n``. If ``None``, defaults to ``seq_len // 2`` —
        wide enough to cover the rising edge and the first few decay
        time-constants, narrow enough to keep tail noise from
        masquerading as the peak.
    smooth_n : int
        Boxcar-smoothing kernel (samples) applied to the magnitude
        trace before peak finding. 0 disables smoothing (default).
        Enabling smoothing biases the recovered peak amplitude low for
        sharply-rising pulses and is rarely worth it once the peak
        search is restricted to the trigger neighborhood.

    Returns
    -------
    features : np.ndarray, shape (n, 8), float32
    feature_names : list of str
    """
    iq = np.asarray(iq_windows, dtype=np.float32)
    if iq.ndim != 3 or iq.shape[-1] != 2:
        raise ValueError(f"expected (n, T, 2) IQ, got shape {iq.shape}")
    n, seq_len, _ = iq.shape
    if peak_search_n is None:
        peak_search_n = max(seq_len // 2, 1)

    m = np.sqrt(iq[..., 0] ** 2 + iq[..., 1] ** 2)
    pre = m[:, :pre_trigger_n]
    pre_mean = pre.mean(axis=1, keepdims=True)
    pre_rms = pre.std(axis=1).astype(np.float32)
    m_bs = m - pre_mean

    if smooth_n and smooth_n > 1:
        kernel = np.ones(smooth_n, dtype=np.float32) / smooth_n
        m_smooth = np.apply_along_axis(
            lambda row: np.convolve(row, kernel, mode="same"),
            axis=1, arr=m_bs,
        )
    else:
        m_smooth = m_bs

    peak_lo = max(pre_trigger_n - 2, 0)
    peak_hi = min(pre_trigger_n + peak_search_n, seq_len)
    search = m_smooth[:, peak_lo:peak_hi]
    peak_idx = (search.argmax(axis=1) + peak_lo).astype(np.int64)
    # Read the smoothed amplitude — using m_bs here would read a sample
    # offset by ~smooth_n/2 from the smoothed peak, biasing low.
    peak_height = m_smooth[np.arange(n), peak_idx].astype(np.float32)

    integral = np.maximum(m_bs, 0.0).sum(axis=1).astype(np.float32)

    thr10 = 0.1 * peak_height[:, None]
    thr50 = 0.5 * peak_height[:, None]
    thr90 = 0.9 * peak_height[:, None]

    above10 = m_bs >= thr10
    above50 = m_bs >= thr50
    above90 = m_bs >= thr90

    t_idx = np.arange(seq_len)
    before_peak = t_idx[None, :] <= peak_idx[:, None]
    after_peak = t_idx[None, :] >= peak_idx[:, None]

    # First rising crossing: argmax returns 0 when no True, masked below.
    rise10 = (above10 & before_peak).argmax(axis=1)
    rise50 = (above50 & before_peak).argmax(axis=1)
    rise90 = (above90 & before_peak).argmax(axis=1)

    below10_after = (~above10) & after_peak
    below50_after = (~above50) & after_peak
    below90_after = (~above90) & after_peak
    fall10 = below10_after.argmax(axis=1)
    fall50 = below50_after.argmax(axis=1)
    fall90 = below90_after.argmax(axis=1)

    rise_ok = above90.any(axis=1)
    decay_ok = below10_after.any(axis=1)
    fwhm_ok = above50.any(axis=1) & below50_after.any(axis=1)

    rise_time = np.where(rise_ok, np.maximum(rise90 - rise10, 0), 0).astype(np.float32)
    decay_time = np.where(
        decay_ok, np.maximum(fall10 - fall90, 0), 0
    ).astype(np.float32)
    fwhm = np.where(
        fwhm_ok, np.maximum(fall50 - rise50, 0), 0
    ).astype(np.float32)
    half_decay = np.where(
        below50_after.any(axis=1), np.maximum(fall50 - peak_idx, 0), 0
    ).astype(np.float32)

    features = np.stack(
        [
            peak_height,
            peak_idx.astype(np.float32),
            rise_time,
            decay_time,
            integral,
            pre_rms,
            fwhm,
            half_decay,
        ],
        axis=1,
    )
    return features, list(FEATURE_NAMES)


def standardize_features(
    train: np.ndarray,
    *others: np.ndarray,
) -> Tuple[np.ndarray, ...]:
    """Z-score features using train-set statistics; returns scaled arrays
    followed by ``(mean, std)``."""
    mean = train.mean(axis=0)
    std = train.std(axis=0) + 1e-8
    out = [(train - mean) / std]
    for x in others:
        out.append((x - mean) / std)
    out.append((mean.astype(np.float32), std.astype(np.float32)))
    return tuple(out)


# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------


class FeatureMLP(nn.Module):
    """Hand-feature MLP baseline.

    Outputs ``(batch, 2)`` matching ``MKIDEnergySSM`` so the same
    Gaussian-NLL loss and per-wavelength evaluation helpers apply.
    Defaults of (48, 56) put us at 3,290 trainable parameters — close
    to the 3,314-parameter VENOM winner config.
    """

    def __init__(
        self,
        n_features: int = 8,
        d_hidden1: int = 48,
        d_hidden2: int = 56,
    ):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(n_features, d_hidden1),
            nn.GELU(),
            nn.Linear(d_hidden1, d_hidden2),
            nn.GELU(),
            nn.Linear(d_hidden2, 2),
        )

    def __call__(self, features):
        return self.mlp(features)


# ---------------------------------------------------------------------------
# Training loop
# ---------------------------------------------------------------------------


def train_feature_mlp(
    model: FeatureMLP,
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_val: np.ndarray,
    y_val: np.ndarray,
    *,
    n_epochs: int = 100,
    batch_size: int = 256,
    lr: float = 3e-3,
    loss_type: str = "gaussian_nll",
    seed: int = 42,
    verbose: bool = True,
) -> dict:
    """Minimal trainer for the feature MLP.

    Mirrors ``mkid_ssm.train``'s recipe (AdamW, weight decay 1e-4, cosine
    schedule with 5-epoch warmup) but operates on pre-extracted feature
    matrices rather than IQ traces.
    """
    from mkid_ssm import energy_resolution_loss, gaussian_nll_loss

    rng = np.random.default_rng(seed)
    n_train = X_train.shape[0]
    warmup = 5

    def get_lr(epoch: int) -> float:
        if epoch < warmup:
            return lr * (epoch + 1) / warmup
        progress = (epoch - warmup) / max(n_epochs - warmup, 1)
        return lr * 0.5 * (1 + math.cos(math.pi * progress))

    optimizer = optim.AdamW(
        learning_rate=lr, weight_decay=1e-4, bias_correction=True
    )
    loss_core = (
        gaussian_nll_loss if loss_type == "gaussian_nll" else energy_resolution_loss
    )

    history = {"epoch": [], "train_loss": [], "val_loss": [], "lr": []}

    if BACKEND == "torch":
        import torch

        optimizer.init(model)

        device = next(model.parameters()).device
        x_train_t = torch.from_numpy(X_train.astype(np.float32)).to(device)
        y_train_t = torch.from_numpy(y_train.astype(np.float32)).to(device)
        x_val_t = torch.from_numpy(X_val.astype(np.float32)).to(device)
        y_val_t = torch.from_numpy(y_val.astype(np.float32)).to(device)

        for epoch in range(n_epochs):
            optimizer.learning_rate = get_lr(epoch)
            perm = rng.permutation(n_train)
            train_loss = 0.0
            n_batches = 0
            model.train()
            for i in range(0, n_train, batch_size):
                idx = perm[i : i + batch_size]
                idx_t = torch.from_numpy(idx).to(device)
                xb = x_train_t.index_select(0, idx_t)
                yb = y_train_t.index_select(0, idx_t)
                pred = model(xb)
                loss = loss_core(pred, yb)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                train_loss += float(loss.item())
                n_batches += 1
            train_loss /= max(n_batches, 1)

            model.eval()
            with torch.no_grad():
                val_pred = model(x_val_t)
                val_loss = float(loss_core(val_pred, y_val_t).item())

            history["epoch"].append(epoch)
            history["train_loss"].append(train_loss)
            history["val_loss"].append(val_loss)
            history["lr"].append(get_lr(epoch))
            if verbose and (epoch % max(n_epochs // 10, 1) == 0 or epoch == n_epochs - 1):
                print(
                    f"  epoch {epoch:>4d} | train {train_loss:.5f} | "
                    f"val {val_loss:.5f} | lr {get_lr(epoch):.1e}"
                )
    else:  # MLX
        def loss_fn(model, xb, yb):
            return loss_core(model(xb), yb)

        loss_and_grad = nn.value_and_grad(model, loss_fn)

        x_train_mx = mx.array(X_train.astype(np.float32))
        y_train_mx = mx.array(y_train.astype(np.float32))
        x_val_mx = mx.array(X_val.astype(np.float32))
        y_val_mx = mx.array(y_val.astype(np.float32))

        for epoch in range(n_epochs):
            optimizer.learning_rate = get_lr(epoch)
            perm = rng.permutation(n_train)
            train_loss = 0.0
            n_batches = 0
            for i in range(0, n_train, batch_size):
                idx = perm[i : i + batch_size]
                xb = x_train_mx[mx.array(idx)]
                yb = y_train_mx[mx.array(idx)]
                loss, grads = loss_and_grad(model, xb, yb)
                optimizer.update(model, grads)
                mx.eval(model.parameters(), optimizer.state)
                train_loss += float(loss.item())
                n_batches += 1
            train_loss /= max(n_batches, 1)

            with no_grad():
                val_pred = model(x_val_mx)
                mx.eval(val_pred)
                val_loss = float(loss_core(val_pred, y_val_mx).item())

            history["epoch"].append(epoch)
            history["train_loss"].append(train_loss)
            history["val_loss"].append(val_loss)
            history["lr"].append(get_lr(epoch))
            if verbose and (epoch % max(n_epochs // 10, 1) == 0 or epoch == n_epochs - 1):
                print(
                    f"  epoch {epoch:>4d} | train {train_loss:.5f} | "
                    f"val {val_loss:.5f} | lr {get_lr(epoch):.1e}"
                )

    return history


# ---------------------------------------------------------------------------
# Causal TCN baseline
# ---------------------------------------------------------------------------


class CausalDilatedConvBlock(nn.Module):
    """LayerNorm → causal dilated depthwise conv → GELU → 1×1 pointwise mix
    → residual.

    The depthwise + pointwise structure matches the FPGA-friendly footprint
    of the existing Mamba block (small per-step compute, easy to pipeline)
    and gives a fair "swap the backbone" comparison.
    """

    def __init__(self, channels: int, kernel_size: int = 3, dilation: int = 1):
        super().__init__()
        self.channels = channels
        self.kernel_size = kernel_size
        self.dilation = dilation
        self.recv_field = (kernel_size - 1) * dilation  # past samples needed

        self.norm = nn.LayerNorm(channels)
        # Causal padding: enough on the left so output[t] only depends on
        # input[<= t]. We append the padding on both sides via the conv's
        # padding arg (which adds it symmetrically in the underlying op)
        # and trim the trailing future samples after the conv.
        self.depthwise = nn.Conv1d(
            in_channels=channels,
            out_channels=channels,
            kernel_size=kernel_size,
            dilation=dilation,
            padding=self.recv_field,
            groups=channels,
            bias=True,
        )
        self.pointwise = nn.Linear(channels, channels, bias=True)

    def __call__(self, x):
        residual = x
        seq_len = x.shape[1]
        x = self.norm(x)
        x = self.depthwise(x)[:, :seq_len, :]
        x = nn.gelu(x)
        x = self.pointwise(x)
        return x + residual

    def step(self, x_t, conv_state):
        """Single-step recurrent.

        Parameters
        ----------
        x_t : (batch, channels)
            Current input sample.
        conv_state : (batch, channels, recv_field)
            Last ``recv_field = (k-1)*dilation`` post-norm samples. The
            current sample is normalized inside this method and consumed
            with the buffer on the right.

        Returns
        -------
        y_t : (batch, channels)
        new_state : (batch, channels, recv_field)
        """
        residual = x_t
        x_norm = self.norm(mx.expand_dims(x_t, 1))[:, 0, :]  # (batch, channels)

        # Build the kernel-length window: positions
        # [t - (k-1)*d, ..., t - d, t] sampled at stride `dilation`.
        # conv_state holds the last recv_field post-norm samples ordered
        # oldest-to-newest, so window indices into conv_state are
        # [0, dilation, 2*dilation, ..., (k-2)*dilation], plus x_norm at
        # the end for position t.
        if self.kernel_size > 1:
            buf = conv_state  # (batch, channels, recv_field)
            taps = [buf[:, :, i] for i in range(0, self.recv_field, self.dilation)]
            taps.append(x_norm)
            window = mx.stack(taps, axis=-1)  # (batch, channels, kernel_size)
        else:
            window = mx.expand_dims(x_norm, axis=-1)

        # Depthwise conv: each channel has its own (kernel,) filter.
        # backend.Conv1d.weight is (out_channels, kernel, in/groups=1) for
        # depthwise — squeeze the trailing 1 to get (channels, kernel).
        w = self.depthwise.weight.squeeze(-1)  # (channels, kernel)
        y = mx.sum(window * w[None, :, :], axis=-1)  # (batch, channels)
        if self.depthwise.bias is not None:
            y = y + self.depthwise.bias

        y = nn.gelu(y) if hasattr(nn, "gelu") else _gelu(y)
        y = self.pointwise(mx.expand_dims(y, 1))[:, 0, :]

        # Update the buffer: drop the oldest, append the new x_norm on
        # the right. conv_state stays at length recv_field.
        if self.recv_field > 0:
            new_state = mx.concatenate(
                [conv_state[:, :, 1:], mx.expand_dims(x_norm, axis=-1)], axis=-1
            )
        else:
            new_state = conv_state
        return y + residual, new_state


class TCNBackbone(nn.Module):
    """Stack of causal dilated conv blocks. Drop-in replacement for
    ``mkid_ssm.SSMBackbone``."""

    def __init__(
        self,
        d_input: int = 2,
        d_model: int = 16,
        n_layers: int = 4,
        kernel_size: int = 3,
        dilations: List[int] | None = None,
    ):
        super().__init__()
        if dilations is None:
            dilations = [2 ** i for i in range(n_layers)]
        if len(dilations) != n_layers:
            raise ValueError(
                f"dilations has length {len(dilations)} but n_layers={n_layers}"
            )
        self.d_model = d_model
        self.n_layers = n_layers
        self.kernel_size = kernel_size
        self.dilations = list(dilations)

        self.input_proj = nn.Linear(d_input, d_model, bias=True)
        self.layers = [
            CausalDilatedConvBlock(d_model, kernel_size=kernel_size, dilation=d)
            for d in self.dilations
        ]
        self.norm_out = nn.LayerNorm(d_model)

    def __call__(self, x):
        x = self.input_proj(x)
        for layer in self.layers:
            x = layer(x)
        return self.norm_out(x)

    def init_state(self, batch_size: int = 1):
        states = []
        for layer in self.layers:
            cs = mx.zeros((batch_size, layer.channels, max(layer.recv_field, 1)))
            states.append(cs)
        return states

    def step(self, x_t, states):
        x = self.input_proj(mx.expand_dims(x_t, 1))[:, 0, :]
        new_states = []
        for layer, cs in zip(self.layers, states):
            x, cs_new = layer.step(x, cs)
            new_states.append(cs_new)
        x = self.norm_out(mx.expand_dims(x, 1))[:, 0, :]
        return x, new_states


class MKIDEnergyTCN(nn.Module):
    """Full TCN baseline: raw IQ window → photon energy.

    Mirrors :class:`mkid_ssm.MKIDEnergySSM` so the existing ``train()``
    loop, evaluation helpers, and recurrent verifier all work unchanged.
    """

    def __init__(
        self,
        d_model: int = 16,
        n_layers: int = 4,
        kernel_size: int = 3,
        d_head_hidden: int = 32,
        dilations: List[int] | None = None,
    ):
        super().__init__()
        self.backbone = TCNBackbone(
            d_input=2,
            d_model=d_model,
            n_layers=n_layers,
            kernel_size=kernel_size,
            dilations=dilations,
        )
        # EnergyHead lives in mkid_ssm.py; import lazily so this module
        # has no top-level dependency on the SSM file.
        from mkid_ssm import EnergyHead
        self.head = EnergyHead(d_model=d_model, d_hidden=d_head_hidden)

    def __call__(self, iq_window):
        features = self.backbone(iq_window)
        return self.head(features)


def predict(model: FeatureMLP, features: np.ndarray) -> np.ndarray:
    """Run the model on a pre-extracted, pre-standardized feature matrix
    and return the energy estimate column ``mu`` as numpy."""
    if BACKEND == "torch":
        import torch

        device = next(model.parameters()).device
        with torch.no_grad():
            x = torch.from_numpy(features.astype(np.float32)).to(device)
            return model(x)[:, 0].detach().cpu().numpy()
    with no_grad():
        x = mx.array(features.astype(np.float32))
        pred = model(x)
        mx.eval(pred)
        return to_numpy(pred[:, 0])