1"""Residual-screened Koopman latent bottleneck utilities.
 2
 3Small NumPy implementation of the empirical residual criterion from the idea.
 4"""
 5import numpy as np
 6
 7
 8def fit_koopman(z):
 9    """Fit K in z[t+1] ~= K z[t], with z shaped [T, m]."""
10    zm, zp = z[:-1].T, z[1:].T
11    return zp @ np.linalg.pinv(zm)
12
13
14def eig_residuals(K, z, eps=1e-12):
15    """Return eigenvalues, right eigenvectors, and normalized residuals.
16
17    This follows the displayed idea formula using u_i^* z_t. For the
18    diagonal/normal systems used in the sanity check this is also the usual
19    modal coordinate. Complex arithmetic is retained until final metrics.
20    """
21    vals, vecs = np.linalg.eig(K)
22    residuals = []
23    for i, lam in enumerate(vals):
24        u = vecs[:, i]
25        a0 = np.conjugate(u) @ z[:-1].T
26        a1 = np.conjugate(u) @ z[1:].T
27        residuals.append(np.sqrt(np.sum(np.abs(a1 - lam * a0)**2) /
28                                 (np.sum(np.abs(a0)**2) + eps)))
29    return vals, vecs, np.asarray(residuals).real
30
31
32def screen_mask(residuals, quantile=0.75):
33    tau = float(np.quantile(residuals, quantile))
34    return residuals <= tau + 1e-12, tau
35
36
37def spectral_forecast(K, z0, steps, mask=None):
38    """Forecast with K's spectral modes, optionally zeroing screened modes."""
39    vals, vecs = np.linalg.eig(K)
40    inv = np.linalg.pinv(vecs)
41    if mask is None:
42        mask = np.ones(len(vals), dtype=bool)
43    amps = inv @ z0
44    amps = amps * mask
45    out = []
46    for h in range(1, steps + 1):
47        out.append(np.real(vecs @ ((vals ** h) * amps)))
48    return np.asarray(out)