File size: 3,611 Bytes
e0d6dea a6ddcdd e0d6dea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | from __future__ import annotations
import numpy as np
import matplotlib.pyplot as plt
from forcing import lorenz63_regime_switch, Lorenz63Params
# ============================================================
# Utilities
# ============================================================
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)
# Global axis limits based on everything (background + all predictions)
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)
# ---- helper panel ----
def panel(ax, pred_u: np.ndarray, name: str, color: str):
# regime B attractor background
ax.plot(B[:, 0], B[:, 1], B[:, 2], color="0.80", lw=0.6, alpha=0.55)
# prediction
P = np.asarray(pred_u)
ax.plot(P[:, 0], P[:, 1], P[:, 2], color=color, lw=2.2, alpha=0.95)
# start/end markers
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)
# cleaner look
ax.grid(False)
ax.set_xticks([])
ax.set_yticks([])
ax.set_zticks([])
# Row 1: LISA
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")
# Row 2: ALSA
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 |