""" Forecasting HiPPO (online, statistics-only) + interpretability plot (3 histories x 2 horizons) What this script does --------------------- 1) Builds THREE HiPPO systems per horizon: - S1: recent past memory on window length H (Legendre/HiPPO, ZOH discretized). - S2: "past-before-that" memory, driven by the lagged value reconstructed from S1. - S3: identical dynamics to S2, but driven by the true input x(t) (so S3 encodes the recent past in the same coordinates as S2, allowing us to learn the shift operator). 2) Learns a linear map T (RRR bottleneck d) ONLINE using only 1st/2nd order stats: - X := S2(t), Y := S3(t) (predict recent-window state from preceding-window state) 3) Input signal: mixture of independent 1D RBF GP draws with different lengthscales (FFT/circulant approx). Weights are normalized so Var[x_t] ~= 1. 4) Runs TWO forecasters on the same x(t): - "short horizon" H_short - "long horizon" H_long 5) After training, makes a 3x2 plot: rows = {true history, HiPPO reconstructed history, HiPPO predictive history} cols = {short horizon, long horizon} where the plotted history is the last 2H time units ([-2H, 0]): - True: actual x[t-2H+1 ... t] - HiPPO reconstructed: stitch (S2 recon on [-2H,-H]) + (S3 recon on [-H,0]) - HiPPO predictive: stitch (S2 recon on [-2H,-H]) + (T_d S2 recon on [-H,0]) Notes ----- - "Exact ZOH" here means we discretize continuous-time HiPPO with expm(A*dt) and the exact ZOH b_d. - The RRR "metric in coefficient space" is Q = T^T T (or T^T G T with a non-identity metric). In this experiment we keep the future OP coefficients orthonormal under the implicit continuous uniform measure induced by HiPPO evaluation; the downstream interpretability uses T only to generate predictive histories. """ from dataclasses import dataclass from typing import Sequence, Tuple, Dict import jax import jax.numpy as jnp import jax.random as jr from jax.scipy.linalg import expm import numpy as np import matplotlib.pyplot as plt from tqdm import tqdm from mpm import get_system_params, get_output_vector # ---------------------------- # Config # ---------------------------- @dataclass class Config: # HiPPO / forecasting measure: str = "legt" n: int = 64 bottleneck_d: int = 12 ridge: float = 1e-4 # Two horizons (in "time units"; dt=1 step) horizon_short: float = 4.0 horizon_long: float = 32.0 horizon_23: float = 32.0 # shared history timescale for systems 2 and 3 # Simulation T_train: int = 20000 T_total: int = T_train + int(max(horizon_short, horizon_long)) + 10 burnin: int = 0 dt: float = 1.0 # GP mixture (RBF) gp_lengthscales: Tuple[float, ...] = (0.1, 3.0, 16.0, 32.0, 64.0) gp_weights: Tuple[float, ...] | None = (0.5, 2, 5, 5, 5) gp_circulant_pad: int = 0 # if >0, uses length (T+pad) for FFT embedding # Plot plot_nt: int = 400 out_png: str = "forecasting_hippo_interpretability.png" seed: int = 0 # ---------------------------- # GP mixture sampling (RBF) on integer grid via FFT/circulant approx # ---------------------------- def _rbf_cov_first_row(n: int, ell: float) -> jnp.ndarray: """First row of Toeplitz covariance [C(0), C(1), ..., C(n-1)] for RBF on Z.""" k = jnp.arange(n) return jnp.exp(-(k * k) / (2.0 * ell * ell)) def sample_rbf_gp_fft(key: jax.Array, T: int, ell: float, pad: int = 0) -> jnp.ndarray: """ Approx sample from stationary GP on {0,...,T-1} with RBF covariance using circulant embedding. Returns approx zero-mean unit-variance series (variance ~ 1). """ # Use embedding length N >= 2*(T+pad)+1 to fit [c0..c_{T+pad}, 0..0, c_{T+pad}..c1] N = int(2 * (T + pad) + 1) # Build Toeplitz first-row values up to lag (T+pad) c_toe = _rbf_cov_first_row(T + pad + 1, ell) # length (T+pad+1): lags 0..T+pad # Construct length-N circulant row: # [c(0), c(1), ..., c(T+pad), 0, ..., 0, c(T+pad), ..., c(1)] left = c_toe # length L = T+pad+1 right = c_toe[1:][::-1] # length L-1 = T+pad mid_len = N - (left.shape[0] + right.shape[0]) # should be >= 0 (often 0) mid = jnp.zeros((mid_len,), dtype=left.dtype) c_circ = jnp.concatenate([left, mid, right], axis=0) # length N lam = jnp.real(jnp.fft.fft(c_circ)) lam = jnp.maximum(lam, 0.0) z = jr.normal(key, (N,)) zf = jnp.fft.fft(z) x = jnp.real(jnp.fft.ifft(jnp.sqrt(lam) * zf)) x = x[:T] x = x - jnp.mean(x) x = x / (jnp.std(x) + 1e-8) return x def sample_gp_mixture(key: jax.Array, cfg: Config) -> jnp.ndarray: """Mixture of independent RBF GPs with weights normalized so marginal variance ~ 1.""" ells = list(cfg.gp_lengthscales) m = len(ells) if cfg.gp_weights is None: w = jnp.ones((m,)) else: w = jnp.array(cfg.gp_weights, dtype=jnp.float32) if w.shape != (m,): raise ValueError("gp_weights must match gp_lengthscales length.") # Normalize so sum_i w_i^2 = 1 (independent unit-var components => total var ~ sum w_i^2) w = w / jnp.sqrt(jnp.sum(w * w) + 1e-12) keys = jr.split(key, m) xs = [] for i, ell in enumerate(ells): xs.append(sample_rbf_gp_fft(keys[i], cfg.T_total, float(ell), pad=cfg.gp_circulant_pad)) X = jnp.stack(xs, axis=0) # (m, T) x = jnp.tensordot(w, X, axes=([0], [0])) # (T,) # final normalize (numerical) x = x - jnp.mean(x) x = x / (jnp.std(x) + 1e-8) return x # ---------------------------- # HiPPO discretization (exact ZOH) # ---------------------------- def discretize_hippo_zoh(A: jnp.ndarray, b: jnp.ndarray, dt: float, horizon: float) -> Tuple[jnp.ndarray, jnp.ndarray]: """ Continuous-time HiPPO: dS/dt = (A/h) S + (b/h) u(t) (scaling by horizon h) Discretize with exact ZOH at step dt: S_{t+dt} = Ad S_t + bd u_t """ Ah = A / horizon bh = b / horizon Ad = expm(Ah * dt) # bd = \int_0^dt exp(Ah * t) bh dt = Ah^{-1}(Ad - I) bh I = jnp.eye(A.shape[0], dtype=A.dtype) # Solve Ah X = (Ad - I) bh => X = Ah^{-1}(Ad-I)bh bd = jnp.linalg.solve(Ah, (Ad - I) @ bh) return Ad, bd # ---------------------------- # Reduced-rank regression from streaming covariances # ---------------------------- def _sym_sqrt_and_invsqrt(S: jnp.ndarray, ridge: float) -> Tuple[jnp.ndarray, jnp.ndarray]: """Return (S^{1/2}, S^{-1/2}) for PSD S via eigendecomposition with ridge.""" # Add ridge to stabilize S = 0.5 * (S + S.T) + ridge * jnp.eye(S.shape[0], dtype=S.dtype) evals, evecs = jnp.linalg.eigh(S) evals = jnp.maximum(evals, 1e-12) sqrt = (evecs * jnp.sqrt(evals)) @ evecs.T invsqrt = (evecs * (1.0 / jnp.sqrt(evals))) @ evecs.T return sqrt, invsqrt def rrr_map_and_projector_from_covs( Sigma_xx: jnp.ndarray, Sigma_yy: jnp.ndarray, Sigma_yx: jnp.ndarray, d: int, ridge: float, ): """ y_hat = W_d x (rank-d RRR) also returns P_x: rank-d projector on x-space that preserves the bottleneck subspace. Uses whitening: C = \Sigma_yy^{-1/2} \Sigma_yx \Sigma_xx^{-1/2} = U diag(s) V^T W_d = \Sigma_yy^{1/2} U_d diag(s_d) V_d^T \Sigma_xx^{-1/2} P_x = \Sigma_xx^{1/2} V_d V_d^T \Sigma_xx^{-1/2} """ Sy_sqrt, Sy_invsqrt = _sym_sqrt_and_invsqrt(Sigma_yy, ridge) Sx_sqrt, Sx_invsqrt = _sym_sqrt_and_invsqrt(Sigma_xx, ridge) C = Sy_invsqrt @ Sigma_yx @ Sx_invsqrt U, s, Vt = jnp.linalg.svd(C, full_matrices=False) U_d = U[:, :d] s_d = s[:d] V_d = Vt[:d, :].T # (nx, d) W_d = Sy_sqrt @ (U_d * s_d) @ V_d.T @ Sx_invsqrt P_x = Sx_sqrt @ (V_d @ V_d.T) @ Sx_invsqrt return W_d, P_x # ---------------------------- # Forecasting HiPPO rollout (two horizons in parallel) # ---------------------------- def run_two_forecasters(cfg: Config) -> Dict[str, Dict[str, jnp.ndarray]]: """ Returns dict with keys {"short","long"} each containing: - "S1","S2","S3": final states - "Sigma_xx","Sigma_yy","Sigma_yx": covariances accumulated (x=S2, y=S3) - "T_d": learned RRR map (rank d) - "eval_func": evaluation function for reconstruction - "M": output vector for lag readout (oldest value) - "horizon": horizon float - "x": full input series (shared) """ key = jr.PRNGKey(cfg.seed) key_x, key_init = jr.split(key, 2) x = sample_gp_mixture(key_x, cfg) (A, b), eval_func, _ = get_system_params(cfg.measure, cfg.n) M = get_output_vector(cfg.measure, cfg.n) # used as lagged readout (oldest endpoint) # Two horizons horizons = {"short": float(cfg.horizon_short), "long": float(cfg.horizon_long)} # Precompute discretizations per horizon # Shared dynamics for systems 2 and 3 (same for short/long) Ad23, bd23 = discretize_hippo_zoh(A, b, cfg.dt, cfg.horizon_23) # Task-specific dynamics for system 1 Ad1_short, bd1_short = discretize_hippo_zoh(A, b, cfg.dt, cfg.horizon_short) Ad1_long, bd1_long = discretize_hippo_zoh(A, b, cfg.dt, cfg.horizon_long) discs = dict( short=(Ad1_short, bd1_short), long=(Ad1_long, bd1_long), ) def init_pack(): return dict( S1=jnp.zeros((cfg.n,)), S2=jnp.zeros((cfg.n,)), S3=jnp.zeros((cfg.n,)), # streaming covs for RRR: x=S2, y=S3 Sigma_xx=jnp.zeros((cfg.n, cfg.n)), Sigma_yy=jnp.zeros((cfg.n, cfg.n)), Sigma_yx=jnp.zeros((cfg.n, cfg.n)), count=jnp.array(0.0), ) packs = {k: init_pack() for k in horizons.keys()} def step_one(pack, u_t, Ad1, bd1): # Update S1 with true input S1 = Ad1 @ pack["S1"] + bd1 * u_t # Lagged scalar from oldest endpoint of S1 (approx x(t-H)) u_lag = jnp.inner(S1, M) # Update S2 with lagged input (encodes window before the recent one) S2 = Ad23 @ pack["S2"] + bd23 * u_lag # Update S3 with true input (same dynamics as S2, but on recent window) S3 = Ad23 @ pack["S3"] + bd23 * u_t count = pack["count"] Sigma_xx = pack["Sigma_xx"] Sigma_yy = pack["Sigma_yy"] Sigma_yx = pack["Sigma_yx"] return S1, S2, S3, Sigma_xx, Sigma_yy, Sigma_yx, count for t in tqdm(range(cfg.T_train)): u_t = x[t] for name, H in horizons.items(): Ad, bd = discs[name] pack = packs[name] S1, S2, S3, Sigma_xx, Sigma_yy, Sigma_yx, count = step_one(pack, u_t, Ad, bd) if t >= cfg.burnin: # online second-order stats Sigma_xx = Sigma_xx + jnp.outer(S2, S2) Sigma_yy = Sigma_yy + jnp.outer(S1, S1) Sigma_yx = Sigma_yx + jnp.outer(S1, S2) count = count + 1.0 packs[name] = dict( S1=S1, S2=S2, S3=S3, Sigma_xx=Sigma_xx, Sigma_yy=Sigma_yy, Sigma_yx=Sigma_yx, count=count ) out = {} for name, H in horizons.items(): pack = packs[name] count = jnp.maximum(pack["count"], 1.0) Sigma_xx = pack["Sigma_xx"] / count Sigma_yy = pack["Sigma_yy"] / count Sigma_yx = pack["Sigma_yx"] / count T_d, P_x = rrr_map_and_projector_from_covs(Sigma_xx, Sigma_yy, Sigma_yx, cfg.bottleneck_d, cfg.ridge) out[name] = dict( horizon=jnp.array(H), S1=pack["S1"], S2=pack["S2"], S3=pack["S3"], Sigma_xx=Sigma_xx, Sigma_yy=Sigma_yy, Sigma_yx=Sigma_yx, T_d=T_d, P_x=P_x, eval_func=eval_func, M=M, x=x, ) return out # ---------------------------- # Reconstruction helpers + plotting # ---------------------------- def eval_matrix_from_eval_func(eval_func, u: jnp.ndarray, n: int) -> jnp.ndarray: """ Build L (len(u) x n) such that for any state S (n,), eval_func(u, S) == L @ S """ I = jnp.eye(n) # columns: eval_func(u, e_j) cols = [eval_func(1-u, I[j]) for j in range(n)] L = jnp.stack(cols, axis=1) # (len(u), n) return L def rbf_kernel(delta: jnp.ndarray, ell: float) -> jnp.ndarray: return jnp.exp(-(delta * delta) / (2.0 * ell * ell)) def mixture_kernel(delta: jnp.ndarray, ells: jnp.ndarray, w: jnp.ndarray) -> jnp.ndarray: # total covariance C_tot(delta) = sum (w_m^2 * exp(-delta^2 / (2 ell_m^2))) # (weights w assumed already normalized so sum w^2 = 1) out = jnp.zeros_like(delta, dtype=jnp.float32) for ell, wi in zip(list(ells), list(w)): out = out + (wi * wi) * rbf_kernel(delta, float(ell)) return out def gp_posterior_mean_mixture( tau_query: jnp.ndarray, # (Q,) future times (relative, continuous), e.g. in [0, H] t_obs: jnp.ndarray, # (N,) observed times (relative), e.g. [-N,...,-1] y_obs: jnp.ndarray, # (N,) observed values x(t_obs) ells: jnp.ndarray, # (M,) w: jnp.ndarray, # (M,) with sum w^2 = 1 ridge: float = 1e-6, ) -> jnp.ndarray: """ GP posterior mean for a zero-mean stationary GP with C_tot induced by an RBF mixture. \mu(\tau) = K(\tau, t_obs) [K(t_obs, t_obs) + ridge I]^{-1} y_obs """ # Kxx D_xx = t_obs[:, None] - t_obs[None, :] Kxx = mixture_kernel(D_xx, ells, w) + ridge * jnp.eye(t_obs.shape[0]) # Kqx D_qx = tau_query[:, None] - t_obs[None, :] Kqx = mixture_kernel(D_qx, ells, w) alpha = jnp.linalg.solve(Kxx, y_obs) return Kqx @ alpha def gp_posterior_mean_std_from_alpha( tau_query: jnp.ndarray, # (Q,) t_obs: jnp.ndarray, # (N,) Kxx: jnp.ndarray, # (N,N) already includes ridge alpha: jnp.ndarray, # (N,) = solve(Kxx, y_obs) ells: jnp.ndarray, # (M,) w: jnp.ndarray, # (M,) sum w^2 = 1 ) -> Tuple[jnp.ndarray, jnp.ndarray]: D_qx = tau_query[:, None] - t_obs[None, :] Kqx = mixture_kernel(D_qx, ells, w) # (Q,N) mean = Kqx @ alpha # (Q,) # Solve Kxx^{-1} Kxq via Cholesky L = jnp.linalg.cholesky(Kxx) # (N,N) # v = L^{-1} Kxq, where Kxq = Kqx^T v = jax.scipy.linalg.solve_triangular(L, Kqx.T, lower=True) # (N,Q) # var = kqq - ||v||^2 kqq = mixture_kernel(jnp.zeros_like(tau_query), ells, w) # (Q,) = C_tot(0)=1 var = jnp.maximum(0.0, kqq - jnp.sum(v * v, axis=0)) # (Q,) std = jnp.sqrt(var + 1e-12) return mean, std def reconstruct_window(eval_func, state: jnp.ndarray, nt: int) -> Tuple[jnp.ndarray, jnp.ndarray]: """Return (u_grid in [0,1], f(u)) for a single HiPPO state reconstruction.""" u = jnp.linspace(0.0, 1.0, nt) f = eval_func(u, state) return u, f def eval_matrix_on_common_lags(eval_func, t_common: jnp.ndarray, H: float, n: int) -> jnp.ndarray: """ Build L_common (mQ x n) on a common lag grid t_common in [-Hmax, 0]. For lags t < -H (outside this system's window), rows are zero. For lags in [-H,0], use normalized u=(t+H)/H and evaluate with correct orientation. """ mQ = t_common.shape[0] u = (t_common + H) / H # maps [-H,0] -> [0,1] valid = (u >= 0.0) & (u <= 1.0) # We'll build columns by evaluating basis vectors at the valid u's. I = jnp.eye(n) L = jnp.zeros((mQ, n), dtype=jnp.float32) u_valid = u[valid] # orientation fix: eval_func(1-u, *) cols = [eval_func(1.0 - u_valid, I[j]) for j in range(n)] # each is (num_valid,) L_valid = jnp.stack(cols, axis=1) # (num_valid, n) L = L.at[valid, :].set(L_valid) return L def plot_gp_oracle_zoh( ax, Hn: int, t_obs: jnp.ndarray, # (N,) observed times, e.g. [-Hctx, ..., -1] y_obs: jnp.ndarray, # (N,) observed values ells: jnp.ndarray, # (M,) w: jnp.ndarray, # (M,) sum w^2 = 1 ridge: float = 1e-6, alpha_fill: float = 0.25, label_prefix: str = "oracle GP", ): """ Plot GP posterior mean and +/-1 std for future in ZOH style on [0,Hn]: mean is constant on [k,k+1) equal to \mu(k), k=0..Hn-1 band is constant on [k,k+1) equal to \mu(k)+/-\sigma(k) """ # Build Kxx once D_xx = t_obs[:, None] - t_obs[None, :] Kxx = mixture_kernel(D_xx, ells, w) + ridge * jnp.eye(t_obs.shape[0]) alpha = jnp.linalg.solve(Kxx, y_obs) Lchol = jnp.linalg.cholesky(Kxx) # Query at integer times 0..Hn-1 (one per ZOH interval) t_int = jnp.arange(0, Hn, dtype=jnp.float32) # (Hn,) D_qx = t_int[:, None] - t_obs[None, :] Kqx = mixture_kernel(D_qx, ells, w) mu = Kqx @ alpha # Posterior variance at integer times v = jax.scipy.linalg.solve_triangular(Lchol, Kqx.T, lower=True) # (N,Hn) kqq = mixture_kernel(jnp.zeros_like(t_int), ells, w) # (Hn,) var = jnp.maximum(0.0, kqq - jnp.sum(v * v, axis=0)) std = jnp.sqrt(var + 1e-12) # Convert to step plotting: # edges: 0..Hn, value on [k,k+1) is mu[k] edges = jnp.arange(0, Hn + 1, dtype=jnp.float32) mu_extended = jnp.concatenate([jnp.array(mu), mu[-1]*jnp.ones(1)], 0) std_extended = jnp.concatenate([jnp.array(std), std[-1]*jnp.ones(1)], 0) ax.step( jnp.array(edges), mu_extended, where="post", linestyle="--", linewidth=2.0, label=f"{label_prefix} mean (ZOH)", ) ax.fill_between( jnp.array(edges), jnp.array(mu_extended - std_extended), jnp.array(mu_extended + std_extended), step="post", alpha=alpha_fill, linewidth=0, label=f"{label_prefix} ±1 std (ZOH)", ) def plot_overlay_with_eigfns( cfg: Config, results: Dict[str, Dict[str, jnp.ndarray]], use_full_rank_Q: bool = True, k_eigs: int = 4, layout: str = "double", # "single" or "double" ) -> None: """ 2x2 figure: Left col: overlay (short on top, long on bottom) Right col: top eigenfunctions of Q_hist (short on top, long on bottom) Styling: - Legend only in top-left - Left plots: grayscale (no Tableau colors) - Right plots: blue ramp for eigenfunctions (dark->light) - Left plots share y-limits - Two layout presets: single-column (skinny-ish) or two-column (wide & short) """ # ------------------------- # Figure size presets # ------------------------- if layout == "single": # ~single column: skinny, roughly square figsize = (3.35, 3.35) # inches (common single-column width ~3.3") width_ratios = [2.2, 1.0] hspace = 0.25 wspace = 0.35 fontsize = 8 elif layout == "double": # ~two columns: short, squat figsize = (5.5, 2.9) # inches (two-column width ~6.9") width_ratios = [1.7, 1.0] hspace = 0.18 wspace = 0.30 fontsize = 8 else: raise ValueError("layout must be 'single' or 'double'") plt.rcParams.update({ "font.size": fontsize, "axes.titlesize": fontsize, "axes.labelsize": fontsize, "legend.fontsize": fontsize - 1, "xtick.labelsize": fontsize - 1, "ytick.labelsize": fontsize - 1, "axes.linewidth": 0.7, "figure.dpi": 200, "savefig.dpi": 300, }) fig, axes = plt.subplots( 2, 2, figsize=figsize, gridspec_kw={"width_ratios": width_ratios}, sharex=False, sharey=False ) # ------------------------- # Mixture kernel params (weights normalized so sum w^2 = 1) # ------------------------- ells = jnp.array(cfg.gp_lengthscales, dtype=jnp.float32) if cfg.gp_weights is None: w = jnp.ones((len(cfg.gp_lengthscales),), dtype=jnp.float32) else: w = jnp.array(cfg.gp_weights, dtype=jnp.float32) w = w / jnp.sqrt(jnp.sum(w * w) + 1e-12) # Common axes setup Hmax = float(max(cfg.horizon_short, cfg.horizon_long)) Hmax_int = int(round(Hmax)) t0 = cfg.T_train # ZOH: after ingesting x[t0-1], current time is t0 # Condition ONLY on past samples occupying [-Hmax, -1] t_obs = jnp.arange(-Hmax_int, 0, dtype=jnp.float32) y_obs = jnp.array(results["long"]["x"][t0 - Hmax_int : t0], dtype=jnp.float32) cols = [("short", float(cfg.horizon_short)), ("long", float(cfg.horizon_long))] # For consistent y-lims across left panels, collect plotted y-extents left_ymins, left_ymaxs = [], [] # Blue ramp for eigenfunctions: dark -> light # (use Matplotlib's "Blues" colormap but choose a range that avoids near-white) blues = plt.cm.Blues(np.linspace(0.85, 0.35, max(k_eigs, 1))) # Grayscale styles for left overlay curves (no Tableau colors) # Order: GP mean, GP band, true past, HiPPO hist, pred hist, forecast, forecast (pred) # We'll keep band as a light gray fill. style_true = dict(color="k", linewidth=0.5, alpha=1.0) style_hippo_hist = dict(color="maroon", linewidth=1.2, alpha=0.8) style_pred_hist = dict(color="firebrick", linewidth=1.2, alpha=0.8) style_forecast = dict(color="lightcoral", linewidth=1.2, alpha=0.8) style_forecast2 = dict(color="0.35", linewidth=1.2, linestyle=":") for row, (name, _) in enumerate(cols): r = results[name] x = r["x"] eval_func = r["eval_func"] S3 = r["S3"] T_d = r["T_d"] P_x = r["P_x"] H_hist = float(cfg.horizon_23) H_fut = float(r["horizon"]) Hn = int(round(H_fut)) # ------------------------- # LEFT: overlay # ------------------------- ax = axes[row, 0] # Oracle GP (ZOH mean + uncertainty) in grayscale ridge = 1e-6 D_xx = t_obs[:, None] - t_obs[None, :] Kxx = mixture_kernel(D_xx, ells, w) + ridge * jnp.eye(t_obs.shape[0]) alpha = jnp.linalg.solve(Kxx, y_obs) Lchol = jnp.linalg.cholesky(Kxx) t_int = jnp.arange(0, Hn, dtype=jnp.float32) D_qx = t_int[:, None] - t_obs[None, :] Kqx = mixture_kernel(D_qx, ells, w) mu = Kqx @ alpha v = jax.scipy.linalg.solve_triangular(Lchol, Kqx.T, lower=True) kqq = mixture_kernel(jnp.zeros_like(t_int), ells, w) var = jnp.maximum(0.0, kqq - jnp.sum(v * v, axis=0)) std = jnp.sqrt(var + 1e-12) edges = jnp.arange(0, Hn + 1, dtype=jnp.float32) mu_extended = np.concatenate([np.array(mu), mu[-1]*np.ones(1)], 0) std_extended = np.concatenate([np.array(std), std[-1]*np.ones(1)], 0) ax.step(np.array(edges), mu_extended, where="post", color="0.15", linewidth=1.4, linestyle="-", label="Oracle GP mean") ax.fill_between( np.array(edges), np.array(mu_extended - std_extended), np.array(mu_extended + std_extended), step="post", color="0.85", alpha=0.8, linewidth=0, label="Oracle GP ±1σ" ) # True past (ZOH) over [-Hmax, 0] past_vals_full = jnp.array(x[t0 - Hmax_int : t0]) past_edges = jnp.arange(-Hmax_int, 1) * cfg.dt ax.step(np.array(past_edges[:-1]), np.array(past_vals_full), where="post", label="True past", **style_true) # HiPPO reconstructed history from S3 over [-H_hist, 0] nt = cfg.plot_nt t_hist = jnp.linspace(-H_hist, 0.0, nt) _, f_hist = reconstruct_window(eval_func, S3, nt) ax.plot(np.array(t_hist[::-1]), np.array(f_hist), label="HiPPO memory", **style_hippo_hist) # Predictive history: P_d S3 S3_predhist = P_x @ S3 _, f_predhist = reconstruct_window(eval_func, S3_predhist, nt) ax.plot(np.array(t_hist[::-1]), np.array(f_predhist), label="Predictive HiPPO memory", **style_pred_hist) # Forecast via predictive history t_fut = jnp.linspace(0.0, H_fut, nt) S_future2 = T_d @ S3_predhist _, f_fut2 = reconstruct_window(eval_func, S_future2, nt) ax.plot(np.array(t_fut[::-1]), np.array(f_fut2), label="HiPPO forecast", **style_forecast) # Formatting ax.axvline(0.0, linewidth=0.8, color="0.2") ax.set_xlim([-Hmax, Hmax]) if row == 0: ax.set_title("Forecasts") # Only label y-axis on left column if row == 0: ax.set_ylabel(r"Short horizon ($\mathbf{H=4}$)") else: ax.set_ylabel(r"Long horizon ($\mathbf{H=32}$)") # Only bottom-left gets x-label ax.set_xlabel("Relative time" if row == 1 else "") # Remove legend from bottom-left if row == 0: ax.legend(loc="upper right", frameon=False, ncol=1, handlelength=2.5) else: ax.legend_.remove() if ax.get_legend() is not None else None # Collect y-lims for later syncing ylo, yhi = ax.get_ylim() left_ymins.append(ylo) left_ymaxs.append(yhi) # Make spines subtle for spine in ["top", "right"]: ax.spines[spine].set_visible(False) # ------------------------- # RIGHT: eigenfunctions of Q_hist # ------------------------- axr = axes[row, 1] # Choose T for Q if use_full_rank_Q: Sigma_xx = r["Sigma_xx"] Sigma_yx = r["Sigma_yx"] n = Sigma_xx.shape[0] T_full = Sigma_yx @ jnp.linalg.solve(Sigma_xx + cfg.ridge * jnp.eye(n), jnp.eye(n)) T_for_Q = T_full else: T_for_Q = T_d Q = jnp.array(T_for_Q.T @ T_for_Q) # Lag grid and eval matrix mQ = 220 t_common = jnp.linspace(-H_hist, 0.0, mQ) L_common = eval_matrix_on_common_lags(eval_func, t_common, H_hist, cfg.n) Q_hist = L_common @ Q @ L_common.T Q_hist = 0.5 * (Q_hist + Q_hist.T) # Eigs evals, evecs = jnp.linalg.eigh(Q_hist) # ascending k = int(min(k_eigs, evecs.shape[1])) idx = jnp.argsort(evals)[::-1][:k] top_evals = evals[idx] top_evecs = evecs[:, idx] # sign fix: make value at lag 0 nonnegative signs = jnp.sign(top_evecs[-1, :] + 1e-12) top_evecs = top_evecs * signs # Plot eigenfunctions with blue ramp, no legend by default (cleaner) for i in range(k): axr.plot(np.array(t_common), np.array(top_evecs[:, i]), color=blues[i], alpha=0.8, linewidth=1.4) axr.axvline(0.0, linewidth=0.8, color="k", alpha=0.6) axr.set_xlim([-H_hist, 0.0]) axr.set_title(r"Top eigfns of $Q$") # Label only bottom-right x-axis axr.set_xlabel("lag" if row == 1 else "") axr.set_ylabel("" if row == 0 else "") # Subtle spines for spine in ["top", "right"]: axr.spines[spine].set_visible(False) # ------------------------- # Sync y-limits for left panels # ------------------------- ylo = float(min(left_ymins)) yhi = float(max(left_ymaxs)) # Add a tiny padding pad = 0.03 * (yhi - ylo + 1e-12) ylo -= pad yhi += pad axes[0, 0].set_ylim([ylo, yhi]) axes[1, 0].set_ylim([ylo, yhi]) # Tight layout control tuned for paper fig.subplots_adjust(left=0.10, right=0.98, bottom=0.12, top=0.92, wspace=wspace, hspace=hspace) plt.savefig(cfg.out_png, bbox_inches="tight") plt.close(fig) # ---------------------------- # Main # ---------------------------- def main(): cfg = Config() cfg.T_total = cfg.T_train + int(max(cfg.horizon_short, cfg.horizon_long)) + 10 results = run_two_forecasters(cfg) plot_overlay_with_eigfns(cfg, results) print(f"Saved: {cfg.out_png}") if __name__ == "__main__": main()