source_code/scripts/salience_hippo.py

"""
Online (TBPTT) training: Shared-salience multi-channel HiPPO on vector-token selective copying.

Training uses one sampled episode per iteration, splits each episode into fixed TBPTT chunks, and applies one optimizer update per chunk. The carried state is stop_gradient'ed at chunk boundaries.
"""
__date__ = "January 2026"

import numpy as np
import matplotlib.pyplot as plt

import jax
import jax.numpy as jnp
import jax.random as jr
from jax import jit, value_and_grad
from jax.lax import scan

from tqdm import tqdm

import optax

from mpm import get_system_params
from mpm.models import SalienceHiPPO
from mpm.polynomials import legsval
from mpm.tasks.selective_recall import make_selective_copying_task_tokens


# ----------------------------
# TBPTT helpers
# ----------------------------
def pad_chunk(X, Y, t0: int, L: int):
    """
    Extract a chunk starting at t0 of length <= L and pad to exactly L.
    Returns Xc, Yc, mask where mask[t]=1 for valid timesteps.
    """
    T = X.shape[0]
    t1 = min(t0 + L, T)
    ell = t1 - t0

    Xc = jnp.zeros((L, X.shape[1]), dtype=X.dtype)
    Yc = jnp.zeros((L, Y.shape[1]), dtype=Y.dtype)
    mask = jnp.zeros((L,), dtype=X.dtype)

    Xc = Xc.at[:ell].set(X[t0:t1])
    Yc = Yc.at[:ell].set(Y[t0:t1])
    mask = mask.at[:ell].set(1.0)
    return Xc, Yc, mask


def main():
    # ----------------------------
    # Config
    # ----------------------------
    seed = 42
    key = jr.PRNGKey(seed)

    # task
    episode_len = 30
    num_episodes = 10000

    # online training steps (episodes seen)
    num_steps = 3000

    # TBPTT
    tbptt_steps = 30

    # token dims / model dims
    d_in = 32
    d_model = 64

    # hippo
    n = 256 # 64
    measure = "legs"
    base_timescale = episode_len

    # salience range
    g_max = 5.0

    # opt
    lr = 1e-3

    # ----------------------------
    # Dataset
    # ----------------------------
    X_np, Y_np, token_table, ids_np = make_selective_copying_task_tokens(
        episode_len=episode_len,
        num_episodes=num_episodes,
        d_in=d_in,
        seed=seed,
    )
    X_all = jnp.array(X_np, dtype=jnp.float32)  # (N,T,d_in)
    Y_all = jnp.array(Y_np, dtype=jnp.float32)
    ids_all = np.array(ids_np)

    # ----------------------------
    # HiPPO system
    # ----------------------------
    (A, b), _, _ = get_system_params(measure, n)
    A, b = A / base_timescale, b / base_timescale
    A = jnp.array(A, dtype=jnp.float32)
    b = jnp.array(b, dtype=jnp.float32)

    # ----------------------------
    # Init model + optimizer
    # ----------------------------
    model = SalienceHiPPO(
        d_in=d_in,
        d_model=d_model,
        n=n,
        A=A,
        b=b,
        g_max=g_max,
        sal_hidden=128,
        outvec_hidden=128,
        mem_dim=32,
    )

    key, k1 = jr.split(key)
    S0 = jnp.zeros((d_model, n), dtype=jnp.float32)
    x0 = jnp.zeros((d_in,), dtype=jnp.float32)
    params = model.init(k1, x0, S0)["params"]

    optimizer = optax.adamw(learning_rate=lr)
    opt_state = optimizer.init(params)

    # ----------------------------
    # Chunk rollout + TBPTT update
    # ----------------------------
    def chunk_loss_and_state(params, S_init, Xc, Yc, mask):
        """
        One TBPTT chunk.
          S_init: (d_model,n)
          Xc,Yc:  (L,d_in) padded
          mask:   (L,) in {0,1}
        Returns: (loss_scalar, S_final, aux)
        """
        def step_fn(S, inputs):
            x_t, y_t, m_t = inputs
            S_next, yhat, g, _ = model.apply({"params": params}, x_t, S)
            mse_t = jnp.mean((yhat - y_t) ** 2)
            # masked mean contribution (mask will be normalized at end)
            return S_next, (m_t * mse_t, m_t, yhat, g)

        S_final, (mse_masked, m_sum, yhat_ts, g_ts) = scan(
            step_fn,
            S_init,
            (Xc, Yc, mask),
        )
        denom = jnp.maximum(jnp.sum(m_sum), 1.0)
        loss = jnp.sum(mse_masked) / denom
        aux = (yhat_ts, g_ts)
        return loss, S_final, aux

    @jit
    def tbptt_chunk_train_step(params, opt_state, S_init, Xc, Yc, mask):
        """
        One optimizer step for one chunk (gradients flow only within the chunk).
        """

        def loss_fn(p):
            loss, S_final, aux = chunk_loss_and_state(p, S_init, Xc, Yc, mask)
            # pack S_final into aux so has_aux=True is satisfied
            return loss, (S_final, aux)

        (loss, (S_final, aux)), grads = value_and_grad(loss_fn, has_aux=True)(params)

        updates, opt_state = optimizer.update(grads, opt_state, params)
        params = optax.apply_updates(params, updates)
        return params, opt_state, S_final, loss, aux


    # ----------------------------
    # Online training loop (episode-by-episode) with TBPTT
    # ----------------------------
    losses = []
    key = jr.PRNGKey(seed + 1)
    N = X_all.shape[0]

    pbar = tqdm(range(num_steps))
    smooth_loss = None

    # state resets per episode (typical for this task)
    S = jnp.zeros((d_model, n), dtype=jnp.float32)

    for step in pbar:
        key, sub = jr.split(key)
        idx = jr.randint(sub, (), 0, N)
        X_ep = X_all[idx]  # (T,d_in)
        Y_ep = Y_all[idx]  # (T,d_in)

        # TBPTT over chunks
        T = X_ep.shape[0]
        n_chunks = (T + tbptt_steps - 1) // tbptt_steps
        loss_ep = 0.0

        for c in range(n_chunks):
            t0 = c * tbptt_steps
            Xc, Yc, mask = pad_chunk(X_ep, Y_ep, t0=t0, L=tbptt_steps)

            params, opt_state, S_end, loss_chunk, _ = tbptt_chunk_train_step(
                params, opt_state, S, Xc, Yc, mask
            )
            loss_ep = loss_ep + loss_chunk

            # CRITICAL: stop gradient at TBPTT boundary
            S = jax.lax.stop_gradient(S_end)

        loss_ep = loss_ep / n_chunks
        losses.append(float(loss_ep))

        if smooth_loss is None:
            smooth_loss = losses[-1]
        else:
            smooth_loss = 0.98 * smooth_loss + 0.02 * losses[-1]
        pbar.set_description(f"loss: {smooth_loss:.7f}")

    # ----------------------------
    # Evaluate + plot: last episode
    # ----------------------------
    X_eval = X_all[-1]  # (T,d_in)
    Y_eval = Y_all[-1]
    ids_eval = ids_all[-1]


    def full_rollout_collect(params, X, Y):
        S = jnp.zeros((d_model, n), dtype=jnp.float32)

        def step_fn(S, inputs):
            x_t, y_t = inputs
            S_next, yhat, g, out_vec = model.apply({"params": params}, x_t, S)
            loss_t = jnp.mean((yhat - y_t) ** 2)
            return S_next, (loss_t, yhat, g, S_next, out_vec)

        _, (loss_ts, yhat_ts, g_ts, S_ts, outvec_ts) = scan(step_fn, S, (X, Y))
        return jnp.mean(loss_ts), yhat_ts, g_ts, S_ts, outvec_ts  # outvec_ts: (T, n)


    loss_eval, preds_ts, gs_ts, S_ts, outvec_ts = full_rollout_collect(params, X_eval, Y_eval)

    preds = np.array(preds_ts)  # (T,d_in)
    gs = np.array(gs_ts)        # (T,)
    Xp = np.array(X_eval)
    Yp = np.array(Y_eval)

    def cos_sim(a, b, eps=1e-8):
        na = np.linalg.norm(a, axis=-1)
        nb = np.linalg.norm(b, axis=-1)
        return np.sum(a * b, axis=-1) / (na * nb + eps)

    cos = cos_sim(preds, Yp)

    t = np.arange(episode_len)

    fig, axarr = plt.subplots(7, 1, figsize=(10, 14))

    axarr[0].set_title(f"Online TBPTT selective copying (eval loss={float(loss_eval):.4g})")
    axarr[0].plot(np.arange(len(losses)), losses)
    axarr[0].set_ylabel("Train loss")

    dims_to_plot = [0, 1, 2]
    for k in dims_to_plot:
        axarr[1].plot(t, Yp[:, k], label=f"Y dim{k}", linewidth=2)
        axarr[1].plot(t, preds[:, k], linestyle="--", label=f"pred dim{k}", alpha=0.8)
    axarr[1].set_ylabel("Selected dims")
    axarr[1].legend(loc="upper left", ncol=2)

    axarr[2].plot(t, cos)
    axarr[2].set_ylabel("cos(pred, target)")
    axarr[2].set_ylim(-0.05, 1.05)

    axarr[3].plot(t, gs)
    axarr[3].set_ylabel("Salience g")
    axarr[3].set_xlabel("Timestep")
    axarr[3].set_ylim(0, None)


    # ----------------------------
    # Interpretability plot (sub-timestep): input + decoded memories, and induced measures
    # ----------------------------
    # Convert to numpy
    S_np = np.array(S_ts)          # (T, d_model, n)
    g_np = np.array(gs_ts)         # (T,)
    X_np_ep = np.array(X_eval)     # (T, d_in)

    T = episode_len

    # Pull learned W_in and compute projected inputs x_proj[t, j]
    W_in = np.array(params["W_in"]["kernel"])      # (d_in, d_model)
    xproj_np = X_np_ep @ W_in                      # (T, d_model)

    # Choose a channel to visualize: largest energy in this episode
    ch = int(np.argmax(np.sum(xproj_np**2, axis=0)))

    # --- Sub-timestep warp phi(t) with piecewise-constant g(t) on [k,k+1)
    # phi(t) = \int_0^t g(s)/base ds
    g_scaled = g_np / float(base_timescale)  # (T,)

    # prefix[k] = sum_{i<k} g_scaled[i], length T+1
    prefix = np.zeros((T + 1,), dtype=np.float32)
    prefix[1:] = np.cumsum(g_scaled).astype(np.float32)

    def phi_cont(t_cont: np.ndarray) -> np.ndarray:
        """
        Piecewise-linear integral of g_scaled on [0,T], assuming g constant on [k,k+1).
        t_cont can be float array in [0,T]. Returns same shape.
        """
        t = np.clip(t_cont, 0.0, float(T))
        k = np.floor(t).astype(np.int32)
        # handle t==T: clamp k to T-1 and set frac=1, then phi=prefix[T]
        k_clamped = np.clip(k, 0, T - 1)
        frac = (t - k).astype(np.float32)
        frac = np.where(t >= float(T), 1.0, frac)  # ensures exact endpoint
        return prefix[k_clamped] + frac * g_scaled[k_clamped].astype(np.float32)

    # Sub-timestep time grid for plotting (absolute time u)
    # (dense enough to look continuous)
    u_grid = np.linspace(0.0, float(T - 1), 600, dtype=np.float32)

    # g(u) piecewise constant
    ku = np.floor(u_grid).astype(np.int32)
    ku = np.clip(ku, 0, T - 1)
    g_u = g_np[ku].astype(np.float32)
    g_u_scaled = g_u / float(base_timescale)

    # True projected input on that dense grid (piecewise constant)
    true_u = xproj_np[ku, ch].astype(np.float32)

    # Choose three snapshot times (integers; you can change these)
    first_len = (2 * T) // 3
    t_snaps = [first_len - 1, first_len + 2, T - 1]  # 3 timepoints

    def decoded_memory_on_u(t_idx: int):
        """
        For snapshot t_idx (integer), return decoded memory as a function of absolute time u.
        Memory curve is defined only for u <= t_idx, else NaN.
        """
        # coefficients at the snapshot (channel ch)
        c = S_np[t_idx, ch, :].astype(np.float32)[None, :]  # (1, n)

        # tau1(u;t) = phi(t) - phi(u)
        phi_t = float(phi_cont(np.array([float(t_idx) + 1.0], dtype=np.float32))[0])
        tau1 = (phi_t - phi_cont(u_grid)).astype(np.float32)  # (len(u_grid),)

        # valid only for u <= t_idx
        valid = (u_grid <= float(t_idx))
        tau1_eval = np.where(valid, tau1, 0.0).astype(np.float32)

        # decode \hat f(tau1)
        hat = legsval(tau1_eval, c)[0].astype(np.float32)  # (len(u_grid),)
        hat = np.where(valid, hat, np.nan).astype(np.float32)
        return hat, valid

    def induced_measure_on_u(t_idx: int):
        """
        \omega0(u|t) propto exp(-(phi(t)-phi(u))) * g(u)/base_timescale
        Defined for u <= t_idx, else NaN. Normalized to integrate to 1 over u<=t.
        """
        phi_t = float(phi_cont(np.array([float(t_idx) + 1.0], dtype=np.float32))[0])
        tau1 = (phi_t - phi_cont(u_grid)).astype(np.float32)

        valid = (u_grid <= float(t_idx))
        w = np.exp(-tau1) * g_u_scaled  # (len(u_grid),)

        w = np.where(valid, w, 0.0).astype(np.float32)

        # normalize as a density over u (continuous approx)
        Z = np.trapz(w, u_grid) + 1e-8
        w = w / Z
        w = np.where(valid, w, np.nan).astype(np.float32)
        return w, valid

    # Build curves
    mem_curves = []
    meas_curves = []
    for t_idx in t_snaps:
        hat, _ = decoded_memory_on_u(t_idx)
        w, _ = induced_measure_on_u(t_idx)
        mem_curves.append((t_idx, hat))
        meas_curves.append((t_idx, w))


    # (1) Input (projected) + decoded memories (flipped horizontally because x-axis is absolute time u)
    axarr[4].plot(u_grid, true_u, linewidth=2, label=f"true input proj (channel {ch})")
    for (t_idx, hat) in mem_curves:
        axarr[4].plot(u_grid, hat, linestyle="--", linewidth=2, label=f"decoded memory @ t={t_idx}")
    axarr[4].set_ylabel("value (proj)")
    axarr[4].set_title("Input (projected) and decoded memories (absolute time axis)")
    axarr[4].legend(loc="upper left", ncol=2)

    # (2) Measures for three timepoints on the same axis
    for (t_idx, w) in meas_curves:
        axarr[5].fill_between(u_grid, 0.0, w, alpha=0.25, label=f"\omega_0(u|t={t_idx})")
        axarr[5].plot(u_grid, w, linewidth=2)  # optional outline for readability

    axarr[5].set_ylim(0, None)
    axarr[5].set_ylabel("density")
    axarr[5].set_xlabel("absolute time u (unwarped)")
    axarr[5].set_title("Induced history measure \omega_0(u|t) (normalized densities)")
    axarr[5].legend(loc="upper left", ncol=2)


    # (3) Linear functionals (kernels) for each prediction timestep in the last third
    outvec_np = np.array(outvec_ts)  # (T, n)

    write_start = first_len
    t_preds = list(range(write_start, T))  # each prediction in last third

    for t_idx in t_preds:
        # post-step time corresponds to t_idx + 1
        phi_t = float(phi_cont(np.array([float(t_idx) + 1.0], dtype=np.float32))[0])
        tau1 = (phi_t - phi_cont(u_grid)).astype(np.float32)

        valid = (u_grid <= float(t_idx))
        tau1_eval = np.where(valid, tau1, 0.0).astype(np.float32)

        # Decode warped kernel k1(tau1)
        c = outvec_np[t_idx].astype(np.float32)[None, :]  # (1, n)
        k1_u = legsval(tau1_eval, c)[0].astype(np.float32)

        # Measure factor omega0(u|t) propto exp(-tau1) * g(u)/base_timescale
        omega_u = np.exp(-tau1_eval) * g_u_scaled  # g_u_scaled already = g(u)/base_timescale

        # Effective absolute-time functional
        k_eff = k1_u * omega_u

        # Mask future
        k_eff = np.where(valid, k_eff, np.nan).astype(np.float32)

        axarr[-1].plot(u_grid, k_eff, alpha=0.5, linewidth=1.5)


    axarr[-1].set_ylabel("kernel value")
    axarr[-1].set_xlabel("absolute time u (unwarped)")
    axarr[-1].set_title("Time-varying linear functionals (decoded kernels) for predictions in last third")


    plt.tight_layout()
    plt.savefig("temp.png")
    plt.show()



if __name__ == "__main__":
    main()