loss.py

"""Cross-entropy and GPD negative log-likelihood (JAX)."""

from __future__ import annotations

import jax
import jax.numpy as jnp
from flax import struct


@struct.dataclass
class LossConfig:
    """Training loss configuration."""

    alpha: float = 0.5
    gpd_threshold_u: float = 1e-4
    label_smoothing: float = 0.0


def gpd_log_pdf(exceedances: jnp.ndarray, xi: jnp.ndarray, sigma: jnp.ndarray) -> jnp.ndarray:
    """
    Log PDF of the Generalized Pareto distribution for exceedances z > 0.

    z: (...,) positive exceedances over threshold u (already subtracted).
    xi, sigma: GPD shape and scale; broadcast with z.

    Uses exponential limit at xi=0; enforces support 1 + xi*z/sigma > 0.
    """
    z = jnp.maximum(exceedances, jnp.asarray(1e-10, dtype=exceedances.dtype))
    sigma = jnp.maximum(sigma, jnp.asarray(1e-8, dtype=sigma.dtype))
    inner = 1.0 + xi * z / sigma
    inner = jnp.maximum(inner, jnp.asarray(1e-12, dtype=inner.dtype))
    general = -jnp.log(sigma) - (1.0 + 1.0 / xi) * jnp.log(inner)
    exponential = -jnp.log(sigma) - z / sigma
    return jnp.where(jnp.abs(xi) < jnp.asarray(1e-4, dtype=xi.dtype), exponential, general)


def gpd_nll(
    exceedances: jnp.ndarray,
    xi: jnp.ndarray,
    sigma: jnp.ndarray,
    mask: jnp.ndarray | None = None,
) -> jnp.ndarray:
    """
    Mean negative log-likelihood over batch (or masked elements).

    mask: optional boolean same shape as exceedances; True = include in mean.
    """
    ll = gpd_log_pdf(exceedances, xi, sigma)
    nll = -ll
    if mask is None:
        return jnp.mean(nll)
    m = mask.astype(nll.dtype)
    denom = jnp.maximum(jnp.sum(m), 1.0)
    return jnp.sum(nll * m) / denom


def cross_entropy_logits(logits: jnp.ndarray, labels: jnp.ndarray, label_smoothing: float = 0.0) -> jnp.ndarray:
    """Softmax cross-entropy with optional label smoothing."""
    num_classes = logits.shape[-1]
    if label_smoothing > 0:
        one_hot = jax.nn.one_hot(labels, num_classes)
        smooth = label_smoothing / max(num_classes - 1, 1)
        targets = (1.0 - label_smoothing) * one_hot + smooth * (1.0 - one_hot)
        logp = jax.nn.log_softmax(logits)
        return -jnp.mean(jnp.sum(targets * logp, axis=-1))
    logp = jax.nn.log_softmax(logits)
    return -jnp.mean(jnp.take_along_axis(logp, labels[:, None], axis=1).squeeze(-1))


def combined_loss(
    logits: jnp.ndarray,
    class_labels: jnp.ndarray,
    exceedances: jnp.ndarray,
    xi_pred: jnp.ndarray,
    sigma_pred: jnp.ndarray,
    evt_mask: jnp.ndarray | None,
    *,
    alpha: float,
    label_smoothing: float = 0.0,
) -> jnp.ndarray:
    """
    Total_Loss = alpha * CE + (1 - alpha) * GPD_NLL.

    xi_pred, sigma_pred: per-sample outputs matching exceedances leading dimension.
    evt_mask: include sample in GPD term when True (e.g. positive-class or POT samples).
    """
    ce = cross_entropy_logits(logits, class_labels, label_smoothing=label_smoothing)
    gpd_term = gpd_nll(exceedances, xi_pred, sigma_pred, mask=evt_mask)
    a = jnp.asarray(alpha, dtype=logits.dtype)
    return a * ce + (1.0 - a) * gpd_term