train.py

"""
FFCC training — preconditioned L-BFGS in the Fourier latent space.

This is the training procedure that produces the benchmark numbers reported in
the README.  It follows the MATLAB reference (``train/TrainModel.m`` plus the
``lib_fft`` helpers): the filters and bias are optimized in a whitened latent
space where the Fourier-domain total-variation regularizer becomes an identity
prior, which is what makes the L-BFGS problem well conditioned.

Optimizing the raw real/imaginary FFT coefficients directly (without the
preconditioner) converges to a far worse solution — that is why the latent-space
formulation, rather than a plain parameterization, is used here.

Reference MATLAB code: https://github.com/google/ffcc
"""

from __future__ import annotations

import numpy as np
from scipy.optimize import minimize as scipy_minimize

from ffcc.core import (
    BIN_SIZE,
    NUM_BINS,
    FFCCModel,
    angular_error_uv,
    fit_bivariate_von_mises,
    idx_to_uv,
    ll_multivariate_normal,
    make_gt_soft_label,
    softmax_backward,
    softmax_forward,
)


# ====================================================================
# Hyperparameters
# ====================================================================
# Tuned values shipped with the MATLAB reference implementation for the
# Gehler/Shi thumbnail project:
#   projects/GehlerShiThumb/GehlerShiThumbHyperparams.m
# These are the values used for every benchmark number in the README.  They are
# copied verbatim from that file — do not "round" them, the model is sensitive
# to the regularizer scale over several orders of magnitude.
GEHLER_HYPERPARAMS = {
    "VONMISES_MULTIPLIER": 2.0 ** 1.75,
    "CROSSENT_MULTIPLIER": 2.0 ** -8.25,
    "FILTER_MULTIPLIERS": np.array([2.0 ** -26.25, 2.0 ** -23.5]),
    "FILTER_SHIFTS": np.array([2.0 ** -60, 2.0 ** -71.75]),
    "BIAS_MULTIPLIER": 2.0 ** -23,
    "BIAS_SHIFT": 2.0 ** -93,
    "VON_MISES_DIAGONAL_EPS": 2.0 ** 0,
}

DEFAULT_HYPERPARAMS = GEHLER_HYPERPARAMS

# Annealing schedule: L-BFGS iteration count grows log-linearly across stages.
LBFGS_ITERS_FIRST = 16
LBFGS_ITERS_LAST = 64


# ====================================================================
# FFT latent-space utilities (MATLAB lib_fft/)
# ====================================================================
def fft2_to_vec(F, apply_scaling=True):
    """Pack a Hermitian-symmetric 2D FFT into its real free parameters.

    Matches MATLAB ``lib_fft/Fft2ToVec.m``.
    """
    F = np.asarray(F, dtype=np.complex128).copy()
    n = F.shape[0]
    Fr = F.real.copy()
    Fi = F.imag.copy()
    h = n // 2
    if apply_scaling:
        Fr[np.ix_([0, h], [0, h])] /= np.sqrt(2.0)
    parts = [
        Fr[:h + 1, 0], Fr[:h + 1, h],
        Fr[:, 1:h].reshape(-1, order="F"),
        Fi[1:h, 0], Fi[1:h, h],
        Fi[:, 1:h].reshape(-1, order="F"),
    ]
    V = np.concatenate(parts).astype(np.float64)
    if apply_scaling:
        V *= np.sqrt(2.0)
    return V


def vec_to_fft2(V, apply_scaling=True):
    """Inverse of :func:`fft2_to_vec`.

    Matches MATLAB ``lib_fft/VecToFft2.m``.
    """
    V = np.asarray(V, dtype=np.float64).copy().reshape(-1)
    n = int(round(np.sqrt(V.size)))
    h = n // 2
    if apply_scaling:
        V /= np.sqrt(2.0)

    Fr = np.zeros((n, n), dtype=np.float64)
    Fi = np.zeros((n, n), dtype=np.float64)

    off = 0
    l1 = h + 1
    Fr[:h + 1, 0] = V[off:off + l1]; off += l1
    Fr[:h + 1, h] = V[off:off + l1]; off += l1
    l3 = n * (h - 1)
    Fr[:, 1:h] = V[off:off + l3].reshape((n, h - 1), order="F"); off += l3
    if apply_scaling:
        Fr[np.ix_([0, h], [0, h])] *= np.sqrt(2.0)
    l4 = h - 1
    Fi[1:h, 0] = V[off:off + l4]; off += l4
    Fi[1:h, h] = V[off:off + l4]; off += l4
    l6 = n * (h - 1)
    Fi[:, 1:h] = V[off:off + l6].reshape((n, h - 1), order="F"); off += l6

    F = Fr + 1j * Fi
    F[0, h + 1:] = np.conj(F[0, 1:h][::-1])
    F[1:, h + 1:] = np.rot90(np.conj(F[1:, 1:h]), 2)
    F[h + 1:, [0, h]] = np.flipud(np.conj(F[1:h, [0, h]]))
    return F


def fft2_regularizer_to_preconditioner(F_reg):
    """Turn a Fourier-domain regularizer into a whitening preconditioner.

    Matches MATLAB ``lib_fft/Fft2RegularizerToPreconditioner.m``.
    """
    F_reg = np.asarray(F_reg, dtype=np.float64)
    n = F_reg.shape[0]
    v = fft2_to_vec(F_reg + 1j * F_reg, apply_scaling=True)
    v_pre = np.sqrt(np.sqrt(2.0) / v)
    idx = np.array([0, n // 2, n // 2 + 1, n + 1], dtype=np.int64)
    v_pre[idx] /= np.sqrt(np.sqrt(2.0))
    return v_pre


def total_variation_filter(n=NUM_BINS):
    """Squared magnitude of the 2D total-variation filter in Fourier space."""
    tv_u = np.abs(np.fft.fft2(np.array([[-1.0], [1.0]]) / np.sqrt(8.0), s=(n, n))) ** 2
    tv_v = np.abs(np.fft.fft2(np.array([[-1.0, 1.0]]) / np.sqrt(8.0), s=(n, n))) ** 2
    return tv_u + tv_v


# ====================================================================
# Loss + gradient (in the preconditioned latent space)
# ====================================================================
def _loss_and_grad(z_pre_vec, X_fft_all, P_gt_all, gt_uv_all,
                   precond_F, precond_B, n_channels,
                   ce_mult, vm_mult, vm_diag_eps):
    """Combined cross-entropy + Von Mises loss and its gradient.

    Matches MATLAB ``train/TrainModelLossfun.m``.  The L2 penalty on the latent
    parameters is the whitened form of the Fourier TV regularizer.
    """
    n = NUM_BINS
    n2 = n * n
    N = X_fft_all.shape[0]

    off = 0
    zF_pre = z_pre_vec[off:off + n2 * n_channels].reshape(n2, n_channels)
    off += n2 * n_channels
    zB_pre = z_pre_vec[off:off + n2]

    zF_lat = precond_F * zF_pre
    zB_lat = precond_B * zB_pre
    F_fft = np.zeros((n, n, n_channels), dtype=np.complex128)
    for c in range(n_channels):
        F_fft[:, :, c] = vec_to_fft2(zF_lat[:, c])
    B = np.real(np.fft.ifft2(vec_to_fft2(zB_lat)))

    total_loss = 0.0
    grad_F_fft = np.zeros_like(F_fft)
    grad_B = np.zeros((n, n), dtype=np.float64)

    for i in range(N):
        FX_fft = np.sum(X_fft_all[i] * F_fft, axis=2)
        H = np.real(np.fft.ifft2(FX_fft)) + B
        P, P_meta = softmax_forward(H)
        d_loss_H = np.zeros_like(H)

        if ce_mult > 0:
            log_P = P_meta['H_shifted'] - np.log(P_meta['expH_sum'])
            ce_loss = max(0.0, -np.sum(P_gt_all[i] * log_P))
            total_loss += ce_mult * ce_loss
            d_loss_H += ce_mult * (P - P_gt_all[i])

        if vm_mult > 0:
            mu_idx, Sigma_idx, dmu_P, dSigma_P = \
                fit_bivariate_von_mises(P, compute_grad=True)
            Sigma_padded = Sigma_idx + vm_diag_eps * np.eye(2)
            mu_uv = idx_to_uv(mu_idx)
            Sigma_uv = Sigma_padded * (BIN_SIZE ** 2)

            LL, dLL_mu, dLL_Sigma = ll_multivariate_normal(
                gt_uv_all[i], mu_uv, Sigma_uv)
            vm_loss = max(0.0, -LL - np.log(2 * np.pi * vm_diag_eps * BIN_SIZE ** 2))
            total_loss += vm_mult * vm_loss

            dvm_mu_idx = -dLL_mu * BIN_SIZE
            dvm_Sigma_idx = -dLL_Sigma * (BIN_SIZE ** 2)
            dvm1_P1, dvm2_P2 = dmu_P
            dS1_P1, dS12_P, dS2_P2 = dSigma_P

            d_vm_P = np.zeros((n, n), dtype=np.float64)
            d_vm_P += dvm_mu_idx[0] * dvm1_P1[:, np.newaxis]
            d_vm_P += dvm_mu_idx[1] * dvm2_P2[np.newaxis, :]
            d_vm_P += dvm_Sigma_idx[0, 0] * dS1_P1[:, np.newaxis]
            d_vm_P += dvm_Sigma_idx[1, 1] * dS2_P2[np.newaxis, :]
            d_vm_P += 2.0 * dvm_Sigma_idx[0, 1] * dS12_P
            d_loss_H += vm_mult * softmax_backward(d_vm_P, P_meta)

        # 1/n^2 comes from the ifft2 in the forward pass
        d_loss_H_fft = np.fft.fft2(d_loss_H / (n * n))
        for c in range(n_channels):
            grad_F_fft[:, :, c] += np.conj(X_fft_all[i, :, :, c]) * d_loss_H_fft
        grad_B += d_loss_H

    grad_B_fft = np.fft.fft2(grad_B / (n * n))
    grad_F_lat = np.zeros((n2, n_channels), dtype=np.float64)
    for c in range(n_channels):
        grad_F_lat[:, c] = fft2_to_vec(grad_F_fft[:, :, c])
    grad_B_lat = fft2_to_vec(grad_B_fft)

    grad_F_pre = precond_F * grad_F_lat
    grad_B_pre = precond_B * grad_B_lat

    data_mass = float(N)
    total_loss += 0.5 * data_mass * (np.sum(zF_pre ** 2) + np.sum(zB_pre ** 2))
    grad_F_pre += data_mass * zF_pre
    grad_B_pre += data_mass * zB_pre

    grad_vec = np.concatenate([grad_F_pre.reshape(-1), grad_B_pre.reshape(-1)])
    return float(total_loss), grad_vec


# ====================================================================
# Public training API
# ====================================================================
def train_ffcc(train_data: list, val_data: list | None = None,
               hyperparams: dict | None = None, n_anneal_stages: int = 2,
               verbose: bool = True) -> FFCCModel:
    """Train an FFCC model with preconditioned L-BFGS and loss annealing.

    The loss is annealed from pure cross-entropy (convex warm-up) to pure Von
    Mises negative log-likelihood (non-convex, but what the evaluation metric
    actually cares about), while the L-BFGS iteration budget grows from
    ``LBFGS_ITERS_FIRST`` to ``LBFGS_ITERS_LAST``.

    Args:
        train_data: list of ``(X, gt_uv)`` tuples, where ``X`` is a
            ``(NUM_BINS, NUM_BINS, n_channels)`` histogram from
            :func:`ffcc.featurize_image` and ``gt_uv`` comes from
            :func:`ffcc.gt_rgb_to_uv`.
        val_data: optional data in the same format; if given, the mean angular
            error after each annealing stage is recorded on the returned model
            as ``model.val_history``.  It does not influence training (there is
            no early stopping — the final stage is the model).
        hyperparams: regularizer / loss multipliers.  Defaults to
            :data:`GEHLER_HYPERPARAMS`, the values tuned for Gehler/Shi in the
            MATLAB reference.  Retune these for a different sensor or dataset.
        n_anneal_stages: number of annealing stages (>= 1).  Two stages
            reproduce the reference results; more stages cost proportionally
            more time for a marginal change.
        verbose: print per-stage progress.

    Returns:
        A trained :class:`ffcc.FFCCModel`.
    """
    if not train_data:
        raise ValueError("train_data is empty")
    if n_anneal_stages < 1:
        raise ValueError("n_anneal_stages must be >= 1")

    hp = dict(DEFAULT_HYPERPARAMS if hyperparams is None else hyperparams)

    n = NUM_BINS
    n2 = n * n
    N = len(train_data)
    n_channels = int(np.asarray(train_data[0][0]).shape[-1])

    X_all = np.array([d[0] for d in train_data], dtype=np.float64)
    gt_uv_all = np.array([d[1] for d in train_data], dtype=np.float64)
    P_gt_all = np.array([make_gt_soft_label(gt) for gt in gt_uv_all])
    X_fft_all = np.zeros((N, n, n, n_channels), dtype=np.complex128)
    for c in range(n_channels):
        X_fft_all[:, :, :, c] = np.fft.fft2(X_all[:, :, :, c], axes=(1, 2))

    tv = total_variation_filter(n)

    f_mult = np.asarray(hp["FILTER_MULTIPLIERS"], dtype=np.float64).reshape(-1)
    f_shift = np.asarray(hp["FILTER_SHIFTS"], dtype=np.float64).reshape(-1)
    if f_mult.size < n_channels or f_shift.size < n_channels:
        raise ValueError(
            f"FILTER_MULTIPLIERS/FILTER_SHIFTS must have >= {n_channels} entries")
    b_mult = float(hp["BIAS_MULTIPLIER"])
    b_shift = float(hp["BIAS_SHIFT"])
    ce_mult_base = float(hp["CROSSENT_MULTIPLIER"])
    vm_mult_base = float(hp["VONMISES_MULTIPLIER"])
    vm_diag_eps = float(hp.get("VON_MISES_DIAGONAL_EPS", 1.0))

    reg_F = np.stack([f_mult[c] * tv + f_shift[c] for c in range(n_channels)], axis=2)
    reg_B = b_mult * tv + b_shift
    precond_F = np.stack([fft2_regularizer_to_preconditioner(reg_F[:, :, c])
                          for c in range(n_channels)], axis=1)
    precond_B = fft2_regularizer_to_preconditioner(reg_B)

    if verbose:
        print(f"Training FFCC on {N} samples ({n_channels} channels, "
              f"{n_anneal_stages} annealing stages)...")

    # Zero init: in the whitened space this is the prior mean, so the first
    # (convex) stage always starts from the same well-defined point.
    z_pre = np.zeros(n2 * n_channels + n2, dtype=np.float64)
    val_history = []

    for i_stage in range(n_anneal_stages):
        w = i_stage / (n_anneal_stages - 1) if n_anneal_stages > 1 else 1.0
        ce_mult = (1.0 - w) * ce_mult_base
        vm_mult = w * vm_mult_base
        num_iters = int(round(np.exp(
            np.log(LBFGS_ITERS_FIRST) * (1 - w) + np.log(LBFGS_ITERS_LAST) * w)))

        if verbose:
            print(f"  Stage {i_stage + 1}/{n_anneal_stages}: "
                  f"CE={ce_mult:.6g}, VM={vm_mult:.6g}, L-BFGS iters={num_iters}")

        result = scipy_minimize(
            lambda x: _loss_and_grad(
                x, X_fft_all, P_gt_all, gt_uv_all,
                precond_F, precond_B, n_channels,
                ce_mult, vm_mult, vm_diag_eps),
            z_pre, method="L-BFGS-B", jac=True,
            options={"maxiter": num_iters, "maxcor": num_iters,
                     "maxfun": 4 + 2 * num_iters, "ftol": 0, "gtol": 0})
        z_pre = result.x.astype(np.float64, copy=True)

        if verbose:
            print(f"    Loss: {result.fun:.4f}")
        if val_data:
            model = _model_from_latent(z_pre, precond_F, precond_B, n_channels)
            errs = [angular_error_uv(model.forward(X)[0], gt) for X, gt in val_data]
            val_history.append(float(np.mean(errs)))
            if verbose:
                print(f"    Val mean angular error: {val_history[-1]:.4f} deg")

    model = _model_from_latent(z_pre, precond_F, precond_B, n_channels)
    model.val_history = val_history
    return model


def _model_from_latent(z_pre, precond_F, precond_B, n_channels) -> FFCCModel:
    """Unwhiten the latent parameters into a usable :class:`FFCCModel`."""
    n = NUM_BINS
    n2 = n * n
    zF_pre = z_pre[:n2 * n_channels].reshape(n2, n_channels)
    zB_pre = z_pre[n2 * n_channels:]
    zF_lat = precond_F * zF_pre
    zB_lat = precond_B * zB_pre

    model = FFCCModel(n_channels=n_channels, init_mode="zeros")
    for c in range(n_channels):
        model.F_fft[:, :, c] = vec_to_fft2(zF_lat[:, c])
        model.F[:, :, c] = np.real(np.fft.ifft2(model.F_fft[:, :, c]))
    model.B = np.real(np.fft.ifft2(vec_to_fft2(zB_lat)))
    return model