| from __future__ import annotations |
|
|
| import numpy as np |
| import matplotlib.pyplot as plt |
|
|
| from forcing import lorenz63_regime_switch, Lorenz63Params |
|
|
|
|
| |
| |
| |
|
|
| def standardize_train_only(F_train: np.ndarray, F_test: np.ndarray, eps: float = 1e-12): |
| mu = F_train.mean(axis=0, keepdims=True) |
| sd = F_train.std(axis=0, keepdims=True) + eps |
| return (F_train - mu) / sd, (F_test - mu) / sd, mu, sd |
|
|
|
|
| def make_task(F_test_z: np.ndarray, a_start: int, ell_ctx: int, steps: int): |
| """ |
| prefix_z: (ell_ctx,3) |
| truth_z : (steps,3) |
| """ |
| prefix = F_test_z[a_start - ell_ctx : a_start, :] |
| truth = F_test_z[a_start : a_start + steps, :] |
| return prefix, truth |
|
|
|
|
| def decimate(X: np.ndarray, max_points: int = 8000) -> np.ndarray: |
| X = np.asarray(X) |
| n = X.shape[0] |
| if n <= max_points: |
| return X |
| stride = max(1, n // max_points) |
| return X[::stride] |
|
|
|
|
| def set_equal_3d_limits(ax, all_xyz: np.ndarray, pad: float = 0.05): |
| mins = all_xyz.min(axis=0) |
| maxs = all_xyz.max(axis=0) |
| ctr = 0.5 * (mins + maxs) |
| span = float(np.max(maxs - mins)) * (1.0 + pad) |
|
|
| ax.set_xlim(ctr[0] - span/2, ctr[0] + span/2) |
| ax.set_ylim(ctr[1] - span/2, ctr[1] + span/2) |
| ax.set_zlim(ctr[2] - span/2, ctr[2] + span/2) |
|
|
|
|
| def plot_context_grid_3d( |
| traj_B_u: np.ndarray, |
| preds_lisa_u: dict[int, np.ndarray], |
| preds_alsa_u: dict[int, np.ndarray], |
| *, |
| title: str, |
| elev: int = 20, |
| azim: int = -60, |
| max_bg_points: int = 12000, |
| ): |
| """ |
| 2 x M grid: |
| row 0: LISA predictions |
| row 1: ALSA predictions |
| Each panel has regime B attractor as gray background. |
| """ |
| multipliers = sorted(preds_lisa_u.keys()) |
| M = len(multipliers) |
|
|
| B = decimate(traj_B_u, max_bg_points) |
|
|
| |
| all_parts = [B] |
| for m in multipliers: |
| all_parts.append(preds_lisa_u[m]) |
| all_parts.append(preds_alsa_u[m]) |
| all_xyz = np.concatenate(all_parts, axis=0) |
|
|
| fig = plt.figure(figsize=(4.0 * M, 9.0)) |
| fig.suptitle(title, fontsize=14) |
|
|
| |
| def panel(ax, pred_u: np.ndarray, name: str, color: str): |
| |
| ax.plot(B[:, 0], B[:, 1], B[:, 2], color="0.80", lw=0.6, alpha=0.55) |
|
|
| |
| P = np.asarray(pred_u) |
| ax.plot(P[:, 0], P[:, 1], P[:, 2], color=color, lw=2.2, alpha=0.95) |
|
|
| |
| ax.scatter(P[0, 0], P[0, 1], P[0, 2], s=25, color=color, alpha=0.95, marker="o") |
| ax.scatter(P[-1, 0], P[-1, 1], P[-1, 2], s=25, color=color, alpha=0.95, marker="^") |
|
|
| ax.set_title(name) |
| ax.set_xlabel("x"); ax.set_ylabel("y"); ax.set_zlabel("z") |
| ax.view_init(elev=elev, azim=azim) |
| set_equal_3d_limits(ax, all_xyz) |
|
|
| |
| ax.grid(False) |
| ax.set_xticks([]) |
| ax.set_yticks([]) |
| ax.set_zticks([]) |
|
|
| |
| for j, m in enumerate(multipliers): |
| ax = fig.add_subplot(2, M, 1 + j, projection="3d") |
| panel(ax, preds_lisa_u[m], f"LISA ℓ={m}L", color="C1") |
|
|
| |
| for j, m in enumerate(multipliers): |
| ax = fig.add_subplot(2, M, M + 1 + j, projection="3d") |
| panel(ax, preds_alsa_u[m], f"ALSA ℓ={m}L", color="C3") |
|
|
| plt.tight_layout() |
| plt.savefig("forcing_test.png") |
| plt.show() |
| return None |
|
|
| def unstandardize(Z: np.ndarray, sd) -> np.ndarray: |
| return Z * sd + mu |
