Weak Koopman Latent Dynamics / weak_koopman_experiment.py

Failed on benchmark

Raw ⬇ ZIP
  1"""Weak Koopman latent-dynamics toy verification.
  2
  3Run with: python3 weak_koopman_experiment.py
  4Outputs results.json with analytic predictions and measured values.
  5"""
  6import json
  7from pathlib import Path
  8import numpy as np
  9
 10SEED = 1425
 11rng = np.random.default_rng(SEED)
 12
 13
 14def psi_and_derivative(t, a, b):
 15    """Compactly supported Hann test function and its physical derivative."""
 16    u = (t - a) / (b - a)
 17    inside = (u >= 0) & (u <= 1)
 18    psi = np.zeros_like(t, dtype=float)
 19    dpsi = np.zeros_like(t, dtype=float)
 20    psi[inside] = np.sin(np.pi * u[inside]) ** 2
 21    dpsi[inside] = (np.pi / (b - a)) * np.sin(2 * np.pi * u[inside])
 22    return psi, dpsi
 23
 24
 25def weak_weights(n, dt, width):
 26    t = np.arange(n) * dt
 27    psi, dpsi = psi_and_derivative(t, 0.0, width * dt)
 28    # B = - integral psi'(t) z(t) dt; sign is irrelevant for variance.
 29    return psi, -dt * dpsi
 30
 31
 32def analytic_variances(sigma, dt, width):
 33    _, wprime = weak_weights(width, dt, width)
 34    weak = sigma**2 * np.sum(wprime**2)
 35    fd = 2.0 * sigma**2 / dt**2
 36    return weak, fd
 37
 38
 39def identity_check():
 40    # Smooth oscillator observable z=[sin(t), cos(t)] and exact generator.
 41    dt, n = 0.002, 1501
 42    t = np.arange(n) * dt
 43    z = np.column_stack([np.sin(t), np.cos(t)])
 44    A = np.array([[0.0, 1.0], [-1.0, 0.0]])
 45    psi, wp = weak_weights(n, dt, n - 1)
 46    G = np.sum((dt * psi)[:, None] * z, axis=0)
 47    B = np.sum(wp[:, None] * z, axis=0)
 48    # B_l = integral psi * (A z)_l = (G @ A.T)_l.
 49    rhs = G @ A.T
 50    return float(np.max(np.abs(B - rhs)) / (np.max(np.abs(rhs)) + 1e-12))
 51
 52
 53def variance_sweep():
 54    # Prediction 1: empirical weak noise variance / analytic variance ~= 1,
 55    # and FD variance / analytic variance ~= 1, over sigma.
 56    dt, width, reps = 0.01, 2.0, 12000
 57    sigmas = [0.02, 0.05, 0.10, 0.20]
 58    _, ww = weak_weights(int(round(width / dt)) + 1, dt, int(round(width / dt)))
 59    # use exactly the samples represented by the returned weights
 60    rows_sigma = []
 61    for sigma in sigmas:
 62        weak_samples = sigma * (rng.standard_normal((reps, len(ww))) @ ww)
 63        fd_samples = sigma * (rng.standard_normal((reps, 2))) @ np.array([-1.0, 1.0]) / dt
 64        weak_emp = float(np.var(weak_samples, ddof=1))
 65        fd_emp = float(np.var(fd_samples, ddof=1))
 66        weak_pred = float(sigma**2 * np.sum(ww**2))
 67        fd_pred = float(2 * sigma**2 / dt**2)
 68        rows_sigma.append({"sigma": sigma, "weak_emp": weak_emp, "weak_pred": weak_pred,
 69                           "fd_emp": fd_emp, "fd_pred": fd_pred,
 70                           "weak_ratio": weak_emp / weak_pred, "fd_ratio": fd_emp / fd_pred})
 71
 72    # Prediction 2: at fixed physical window, weak variance is proportional to
 73    # sigma^2 and approximately dt, while pointwise FD is proportional dt^-2.
 74    # Prediction 3: with fixed dt, widening a Hann test function gives W^-3.
 75    rows_dt, rows_width = [], []
 76    sigma, physical_width = 0.10, 2.0
 77    for dt_i in [0.005, 0.01, 0.02, 0.04, 0.08]:
 78        m = int(round(physical_width / dt_i))
 79        weak_pred, fd_pred = analytic_variances(sigma, dt_i, m)
 80        _, ww_i = weak_weights(m + 1, dt_i, m)
 81        weak_emp = float(np.var(sigma * (rng.standard_normal((reps, len(ww_i))) @ ww_i), ddof=1))
 82        fd_emp = float(np.var(sigma * (rng.standard_normal((reps, 2))) @ np.array([-1., 1.]) / dt_i, ddof=1))
 83        rows_dt.append({"dt": dt_i, "weak_emp": weak_emp, "weak_pred": weak_pred,
 84                        "fd_emp": fd_emp, "fd_pred": fd_pred})
 85    for m in [50, 75, 100, 150, 200, 300, 400]:
 86        dt_i = 0.01
 87        weak_pred, fd_pred = analytic_variances(sigma, dt_i, m)
 88        _, ww_i = weak_weights(m + 1, dt_i, m)
 89        weak_emp = float(np.var(sigma * (rng.standard_normal((reps, len(ww_i))) @ ww_i), ddof=1))
 90        rows_width.append({"width": m * dt_i, "weak_emp": weak_emp, "weak_pred": weak_pred,
 91                           "fd_pred": fd_pred})
 92    # Log-log slopes quantify the parameter scalings.
 93    dt_arr = np.array([r["dt"] for r in rows_dt])
 94    weak_dt_slope = float(np.polyfit(np.log(dt_arr), np.log([r["weak_emp"] for r in rows_dt]), 1)[0])
 95    fd_dt_slope = float(np.polyfit(np.log(dt_arr), np.log([r["fd_emp"] for r in rows_dt]), 1)[0])
 96    w_arr = np.array([r["width"] for r in rows_width])
 97    weak_w_slope = float(np.polyfit(np.log(w_arr), np.log([r["weak_emp"] for r in rows_width]), 1)[0])
 98    return {"sigma": rows_sigma, "dt": rows_dt, "width": rows_width,
 99            "slopes": {"weak_vs_dt": weak_dt_slope, "fd_vs_dt": fd_dt_slope,
100                        "weak_vs_window": weak_w_slope}}
101
102
103def generator_mini_experiment():
104    """Compare FD and weak regression for the exact 2D oscillator generator."""
105    dt, n, trajectories, sigma = 0.02, 101, 300, 0.15
106    t = np.arange(n) * dt
107    A_true = np.array([[0., 1.], [-1., 0.]])
108    errs_fd, errs_weak = [], []
109    width = (n - 1) * dt
110    psi, wp = weak_weights(n, dt, n - 1)
111    # Multiple random phase trajectories make a small EDMD-like regression.
112    for _ in range(trajectories):
113        phase = rng.uniform(0, 2*np.pi)
114        clean = np.column_stack([np.sin(t + phase), np.cos(t + phase)])
115        noisy = clean + sigma * rng.standard_normal(clean.shape)
116        # Standard pointwise derivative matching, central differences.
117        dz = (noisy[2:] - noisy[:-2]) / (2*dt)
118        Zmid = noisy[1:-1]
119        Afd_T = np.linalg.lstsq(Zmid, dz, rcond=None)[0]
120        Afd = Afd_T.T
121        # One weak equation per trajectory is insufficient for 2D A, so use
122        # several overlapping windows as rows in G and B.
123        Gs, Bs = [], []
124        win = 51
125        for start in range(0, n - win + 1, 10):
126            zwin = noisy[start:start+win]
127            ps, ws = weak_weights(win, dt, win - 1)
128            Gs.append(np.sum((dt * ps)[:, None] * zwin, axis=0))
129            Bs.append(np.sum(ws[:, None] * zwin, axis=0))
130        G = np.asarray(Gs); B = np.asarray(Bs)
131        Aw_T = np.linalg.solve(G.T @ G + 1e-5*np.eye(2), G.T @ B)
132        Aw = Aw_T.T
133        errs_fd.append(np.linalg.norm(Afd - A_true))
134        errs_weak.append(np.linalg.norm(Aw - A_true))
135    return {"sigma": sigma, "fd_generator_error": float(np.mean(errs_fd)),
136            "weak_generator_error": float(np.mean(errs_weak)),
137            "weak_over_fd": float(np.mean(errs_weak) / np.mean(errs_fd))}
138
139
140def main():
141    result = {"seed": SEED, "identity_relative_error": identity_check(),
142              "variance_sweeps": variance_sweep(),
143              "generator_mini_experiment": generator_mini_experiment()}
144    Path("results.json").write_text(json.dumps(result, indent=2))
145    print(json.dumps(result, indent=2))
146
147if __name__ == "__main__":
148    main()