Upload 4 files

attractors.py

@@ -0,0 +1,336 @@
+# attractors.py
+"""
+A small zoo of chaotic attractors with a standardized API.
+
+All generators return a trajectory array of shape (T, D), where:
+  - T is the number of time points / steps.
+  - D is the state dimension (e.g., 3 for Lorenz-63).
+
+For continuous-time systems (ODEs), dt is the integration step (RK4).
+For discrete-time maps (Hénon, Ikeda), dt is accepted for API consistency but ignored.
+
+You can later standardize/split/chop trajectories for train/test externally.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable, Dict, Optional, Tuple
+
+import numpy as np
+
+
+Array = np.ndarray
+
+
+# -----------------------------
+# generic RK4 integrator
+# -----------------------------
+def rk4(f: Callable[[Array], Array], x0: Array, T: int, dt: float) -> Array:
+    """
+    Integrate x' = f(x) using fixed-step RK4.
+
+    Parameters
+    ----------
+    f  : function x -> dx/dt
+    x0 : (D,)
+    T  : number of time points (trajectory length)
+    dt : step size
+
+    Returns
+    -------
+    traj : (T, D)
+    """
+    x = np.asarray(x0, dtype=float).copy()
+    D = x.size
+    traj = np.zeros((int(T), D), dtype=float)
+    for t in range(int(T)):
+        traj[t] = x
+        k1 = f(x)
+        k2 = f(x + 0.5 * dt * k1)
+        k3 = f(x + 0.5 * dt * k2)
+        k4 = f(x + dt * k3)
+        x = x + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)
+    return traj
+
+
+def _rng_from_seed(seed: Optional[int]) -> np.random.Generator:
+    return np.random.default_rng(None if seed is None else int(seed))
+
+
+def _discard_burn_in(traj: Array, burn_in: int) -> Array:
+    burn_in = int(burn_in)
+    if burn_in <= 0:
+        return traj
+    if burn_in >= traj.shape[0]:
+        raise ValueError(f"burn_in={burn_in} must be < T={traj.shape[0]}")
+    return traj[burn_in:]
+
+
+# -----------------------------
+# 1) Lorenz-63 (3D ODE)
+# -----------------------------
+def lorenz63(
+    T: int,
+    dt: float = 0.01,
+    *,
+    sigma: float = 10.0,
+    rho: float = 28.0,
+    beta: float = 8.0 / 3.0,
+    x0: Optional[Array] = None,
+    seed: Optional[int] = 0,
+    burn_in: int = 0,
+) -> Array:
+    """
+    Lorenz-63 system (classic chaotic regime).
+    Returns (T-burn_in, 3).
+    """
+    rng = _rng_from_seed(seed)
+    if x0 is None:
+        x0 = np.array([1.0, 1.0, 1.0], dtype=float) + 0.01 * rng.standard_normal(3)
+    else:
+        x0 = np.asarray(x0, dtype=float).reshape(3)
+
+    def f(x: Array) -> Array:
+        return np.array(
+            [
+                sigma * (x[1] - x[0]),
+                x[0] * (rho - x[2]) - x[1],
+                x[0] * x[1] - beta * x[2],
+            ],
+            dtype=float,
+        )
+
+    traj = rk4(f, x0, T=T, dt=float(dt))
+    return _discard_burn_in(traj, burn_in)
+
+
+# -----------------------------
+# 2) Lorenz-96 (dD ODE)
+# -----------------------------
+def lorenz96(
+    T: int,
+    dt: float = 0.01,
+    *,
+    dim: int = 40,
+    F: float = 8.0,
+    x0: Optional[Array] = None,
+    seed: Optional[int] = 0,
+    burn_in: int = 0,
+) -> Array:
+    """
+    Lorenz-96 in the standard chaotic regime (F ~ 8, dim ~ 40).
+    Returns (T-burn_in, dim).
+    """
+    rng = _rng_from_seed(seed)
+    d = int(dim)
+    if x0 is None:
+        # common init: near-constant with small noise
+        x0 = F * np.ones(d, dtype=float)
+        x0 += 0.01 * rng.standard_normal(d)
+        # add a localized bump
+        x0[0] += 0.1
+    else:
+        x0 = np.asarray(x0, dtype=float).reshape(d)
+
+    def f(x: Array) -> Array:
+        # cyclic indices: x_{i+1}, x_{i-1}, x_{i-2}
+        xp1 = np.roll(x, -1)
+        xm1 = np.roll(x, 1)
+        xm2 = np.roll(x, 2)
+        return (xp1 - xm2) * xm1 - x + F
+
+    traj = rk4(f, x0, T=T, dt=float(dt))
+    return _discard_burn_in(traj, burn_in)
+
+
+# -----------------------------
+# 3) Rössler (3D ODE)
+# -----------------------------
+def rossler(
+    T: int,
+    dt: float = 0.02,
+    *,
+    a: float = 0.2,
+    b: float = 0.2,
+    c: float = 5.7,
+    x0: Optional[Array] = None,
+    seed: Optional[int] = 0,
+    burn_in: int = 0,
+) -> Array:
+    """
+    Rössler attractor (classic chaotic parameters).
+    Returns (T-burn_in, 3).
+    """
+    rng = _rng_from_seed(seed)
+    if x0 is None:
+        x0 = np.array([0.1, 0.0, 0.0], dtype=float) + 0.01 * rng.standard_normal(3)
+    else:
+        x0 = np.asarray(x0, dtype=float).reshape(3)
+
+    def f(x: Array) -> Array:
+        x1, y1, z1 = x
+        return np.array(
+            [
+                -y1 - z1,
+                x1 + a * y1,
+                b + z1 * (x1 - c),
+            ],
+            dtype=float,
+        )
+
+    traj = rk4(f, x0, T=T, dt=float(dt))
+    return _discard_burn_in(traj, burn_in)
+
+
+# -----------------------------
+# 4) Hénon map (2D discrete)
+# -----------------------------
+def henon(
+    T: int,
+    dt: float = 1.0,  # ignored (kept for API consistency)
+    *,
+    a: float = 1.4,
+    b: float = 0.3,
+    x0: Optional[Array] = None,
+    seed: Optional[int] = 0,
+    burn_in: int = 0,
+) -> Array:
+    """
+    Hénon map (discrete-time chaos). dt is ignored.
+    Returns (T-burn_in, 2).
+    """
+    rng = _rng_from_seed(seed)
+    if x0 is None:
+        x = np.array([0.1, 0.1], dtype=float) + 0.01 * rng.standard_normal(2)
+    else:
+        x = np.asarray(x0, dtype=float).reshape(2)
+
+    traj = np.zeros((int(T), 2), dtype=float)
+    for t in range(int(T)):
+        traj[t] = x
+        xn = 1.0 - a * x[0] ** 2 + x[1]
+        yn = b * x[0]
+        x = np.array([xn, yn], dtype=float)
+
+    return _discard_burn_in(traj, burn_in)
+
+
+# -----------------------------
+# 5) Ikeda map (2D discrete)
+# -----------------------------
+def ikeda(
+    T: int,
+    dt: float = 1.0,  # ignored
+    *,
+    u: float = 0.9,
+    k: float = 0.4,
+    p: float = 6.0,
+    x0: Optional[Array] = None,
+    seed: Optional[int] = 0,
+    burn_in: int = 0,
+) -> Array:
+    """
+    Ikeda map (discrete-time chaos). dt is ignored.
+    Standard real-form update:
+      t = k - p / (1 + x^2 + y^2)
+      x' = 1 + u * (x cos t - y sin t)
+      y' =     u * (x sin t + y cos t)
+    Returns (T-burn_in, 2).
+    """
+    rng = _rng_from_seed(seed)
+    if x0 is None:
+        x = np.array([0.1, 0.1], dtype=float) + 0.01 * rng.standard_normal(2)
+    else:
+        x = np.asarray(x0, dtype=float).reshape(2)
+
+    traj = np.zeros((int(T), 2), dtype=float)
+    for tstep in range(int(T)):
+        traj[tstep] = x
+        x1, y1 = x
+        t = k - p / (1.0 + x1 * x1 + y1 * y1)
+        ct = np.cos(t)
+        st = np.sin(t)
+        xn = 1.0 + u * (x1 * ct - y1 * st)
+        yn = u * (x1 * st + y1 * ct)
+        x = np.array([xn, yn], dtype=float)
+
+    return _discard_burn_in(traj, burn_in)
+
+
+# -----------------------------
+# 6) Duffing oscillator (forced, 2D ODE)
+# -----------------------------
+def duffing(
+    T: int,
+    dt: float = 0.01,
+    *,
+    delta: float = 0.2,
+    gamma: float = 0.3,
+    omega: float = 1.2,
+    # equation: x'' + delta x' - x + x^3 = gamma cos(omega t)
+    x0: Optional[Array] = None,   # (x, v)
+    seed: Optional[int] = 0,
+    burn_in: int = 0,
+) -> Array:
+    """
+    Forced Duffing oscillator in a commonly chaotic regime.
+    State is (x, v) where v = x'.
+    Returns (T-burn_in, 2).
+
+    Note: This is already nonautonomous due to cos(omega t).
+    """
+    rng = _rng_from_seed(seed)
+    if x0 is None:
+        x = np.array([0.1, 0.0], dtype=float) + 0.01 * rng.standard_normal(2)
+    else:
+        x = np.asarray(x0, dtype=float).reshape(2)
+
+    traj = np.zeros((int(T), 2), dtype=float)
+
+    # Nonautonomous RK4 (depends on time)
+    t = 0.0
+    for n in range(int(T)):
+        traj[n] = x
+
+        def f(tcur: float, s: Array) -> Array:
+            xx, vv = s
+            # x' = v
+            # v' = x - x^3 - delta v + gamma cos(omega t)
+            return np.array(
+                [vv, xx - xx**3 - delta * vv + gamma * np.cos(omega * tcur)],
+                dtype=float,
+            )
+
+        k1 = f(t, x)
+        k2 = f(t + 0.5 * dt, x + 0.5 * dt * k1)
+        k3 = f(t + 0.5 * dt, x + 0.5 * dt * k2)
+        k4 = f(t + dt, x + dt * k3)
+        x = x + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)
+        t += dt
+
+    return _discard_burn_in(traj, burn_in)
+
+
+# -----------------------------
+# registry / convenience
+# -----------------------------
+ATTRACTORS: Dict[str, Callable[..., Array]] = {
+    "lorenz63": lorenz63,
+    "lorenz96": lorenz96,
+    "rossler": rossler,
+    "henon": henon,
+    "ikeda": ikeda,
+    "duffing": duffing,
+}
+
+
+def generate(name: str, T: int, dt: float = 0.01, **kwargs) -> Array:
+    """
+    Convenience wrapper:
+      traj = generate("lorenz63", T=20000, dt=0.01, burn_in=2000)
+    """
+    key = name.lower()
+    if key not in ATTRACTORS:
+        raise KeyError(f"Unknown attractor '{name}'. Options: {sorted(ATTRACTORS.keys())}")
+    return ATTRACTORS[key](T=T, dt=dt, **kwargs)

datasets.py

@@ -0,0 +1,289 @@
+# datasets.py
+from __future__ import annotations
+
+import os
+import sys
+import time
+import urllib.request
+from dataclasses import dataclass
+from typing import Iterable, Optional, Sequence, Tuple, Union
+
+import numpy as np
+
+try:
+    import pandas as pd
+except Exception as e:  # pragma: no cover
+    raise ImportError("datasets.py requires pandas. Install with: pip install pandas") from e
+
+
+# -----------------------------
+# Config
+# -----------------------------
+ELECTRICITY_URL = "https://data.openei.org/files/8562/historic_load_hourly_2016_2023_county.h5"
+ELECTRICITY_H5 = "historic_load_hourly_2016_2023_county.h5"
+
+
+# -----------------------------
+# Helpers
+# -----------------------------
+def download_with_progress(url: str, out_path: str, chunk_size: int = 1024 * 1024) -> None:
+    """
+    Download a large file with a simple text progress bar (no extra deps).
+    Skips if file already exists.
+    """
+    if os.path.exists(out_path):
+        return
+
+    os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
+    print(f"[datasets] Downloading:\n  {url}\n  -> {out_path}")
+
+    with urllib.request.urlopen(url) as response, open(out_path, "wb") as out_file:
+        total = getattr(response, "length", None)
+        downloaded = 0
+        start_time = time.time()
+
+        while True:
+            chunk = response.read(chunk_size)
+            if not chunk:
+                break
+            out_file.write(chunk)
+            downloaded += len(chunk)
+
+            if total:
+                frac = downloaded / total
+                pct = frac * 100
+                bar_len = 30
+                filled = int(bar_len * frac)
+                bar = "#" * filled + "-" * (bar_len - filled)
+                rate = downloaded / (time.time() - start_time + 1e-9)
+                sys.stdout.write(
+                    f"\r[{bar}] {pct:6.2f}%  "
+                    f"{downloaded / 1e6:8.1f} / {total / 1e6:8.1f} MB  "
+                    f"{rate / 1e6:4.1f} MB/s"
+                )
+                sys.stdout.flush()
+
+    if total:
+        print("\n[datasets] Download complete.")
+    else:
+        print("[datasets] Download complete.")
+
+
+def _safe_cache_name(
+    *,
+    start: Optional[str],
+    stop: Optional[str],
+    resample: Optional[str],
+    agg: str,
+    n_features: Optional[int],
+    feature_method: str,
+    seed: int,
+    downcast_float32: bool,
+) -> str:
+    def clean(x: Optional[str]) -> str:
+        if x is None:
+            return "none"
+        return str(x).replace(":", "-").replace("/", "-").replace(" ", "_")
+
+    return (
+        "electricity"
+        f"_start-{clean(start)}"
+        f"_stop-{clean(stop)}"
+        f"_res-{clean(resample)}"
+        f"_agg-{agg}"
+        f"_nf-{n_features if n_features is not None else 'all'}"
+        f"_fsel-{feature_method}"
+        f"_seed-{int(seed)}"
+        f"_f32-{int(bool(downcast_float32))}"
+        ".npz"
+    )
+
+
+def _require_tables_hint(err: Exception) -> None:
+    msg = (
+        "\n[datasets] Failed to read .h5. This usually means the optional dependency `tables` is missing.\n"
+        "Install once to build a cache:\n"
+        "  pip install tables\n"
+        "or, if using extras:\n"
+        "  pip install .[electricity]\n"
+    )
+    raise ImportError(msg) from err
+
+
+def _coerce_datetime_index(df: pd.DataFrame) -> pd.DataFrame:
+    if not isinstance(df.index, pd.DatetimeIndex):
+        df.index = pd.to_datetime(df.index, utc=True)
+    if df.index.tz is None:
+        df.index = df.index.tz_localize("UTC")
+    return df.sort_index()
+
+
+def _fill_missing(df: pd.DataFrame) -> pd.DataFrame:
+    # Conservative: forward-fill then back-fill; if still NaN, fill 0.
+    df = df.ffill().bfill()
+    if df.isna().values.any():
+        df = df.fillna(0.0)
+    return df
+
+
+def _select_features(
+    df: pd.DataFrame,
+    *,
+    features: Optional[Sequence[str]] = None,
+    n_features: Optional[int] = None,
+    method: str = "variance",
+    seed: int = 0,
+) -> pd.DataFrame:
+    """
+    Feature selection on columns.
+    - features: explicit list of column labels
+    - n_features: choose subset
+      method: "variance" (top variance) or "random"
+    """
+    if features is not None:
+        missing = [c for c in features if c not in df.columns]
+        if missing:
+            raise KeyError(f"Requested features not found in columns: {missing[:10]}{'...' if len(missing)>10 else ''}")
+        return df.loc[:, list(features)]
+
+    if n_features is None:
+        return df
+
+    n_features = int(n_features)
+    if n_features <= 0:
+        raise ValueError("n_features must be positive")
+
+    if n_features >= df.shape[1]:
+        return df
+
+    method = method.lower()
+    if method == "random":
+        rng = np.random.default_rng(int(seed))
+        cols = list(df.columns)
+        idx = rng.choice(len(cols), size=n_features, replace=False)
+        sel = [cols[i] for i in idx]
+        return df.loc[:, sel]
+
+    if method == "variance":
+        # top variance across time (after missing fill)
+        v = df.var(axis=0, ddof=0)
+        sel = list(v.sort_values(ascending=False).index[:n_features])
+        return df.loc[:, sel]
+
+    raise ValueError("feature_method must be one of: {'variance','random'}")
+
+
+# -----------------------------
+# Main API: Electricity
+# -----------------------------
+@dataclass(frozen=True)
+class ElectricityData:
+    R_tX: np.ndarray
+    t_index: np.ndarray          # int64 UTC ns
+    X_index: np.ndarray          # dtype=object strings
+
+
+def load_electricity(
+    *,
+    data_dir: str = "data",
+    url: str = ELECTRICITY_URL,
+    h5_name: str = ELECTRICITY_H5,
+    # slicing
+    start: Optional[str] = None,      # e.g. "2019-01-01"
+    stop: Optional[str] = None,       # e.g. "2020-01-01"
+    # resampling
+    resample: Optional[str] = None,   # e.g. "D" for daily, "W" weekly; None keeps hourly
+    agg: str = "sum",                 # "sum" or "mean"
+    # feature downselect
+    features: Optional[Sequence[str]] = None,
+    n_features: Optional[int] = None,
+    feature_method: str = "variance", # "variance" or "random"
+    seed: int = 0,
+    # dtype / missing
+    downcast_float32: bool = True,
+    # caching
+    cache: bool = True,
+    cache_dir: Optional[str] = None,
+    force_rebuild_cache: bool = False,
+) -> ElectricityData:
+    """
+    Returns a manageable (T, X) array from the OpenEI historic county load dataset.
+
+    Behavior:
+      - If a matching .npz cache exists -> loads it (NO `tables` required).
+      - Else downloads .h5 and reads via pandas (requires `tables` installed once),
+        then applies slicing/resampling/feature-selection, and optionally writes cache.
+
+    Notes:
+      - index is returned as int64 UTC nanoseconds (portable, no pandas dependency later).
+      - columns are returned as strings (county IDs like 'p36041').
+    """
+    data_dir = str(data_dir)
+    os.makedirs(data_dir, exist_ok=True)
+
+    if cache_dir is None:
+        cache_dir = os.path.join(data_dir, "cache")
+    os.makedirs(cache_dir, exist_ok=True)
+
+    agg = agg.lower()
+    if agg not in ("sum", "mean"):
+        raise ValueError("agg must be 'sum' or 'mean'")
+
+    cache_name = _safe_cache_name(
+        start=start, stop=stop, resample=resample, agg=agg,
+        n_features=n_features, feature_method=feature_method,
+        seed=seed, downcast_float32=downcast_float32,
+    )
+    cache_path = os.path.join(cache_dir, cache_name)
+
+    # 1) Load cache if present
+    if cache and (not force_rebuild_cache) and os.path.exists(cache_path):
+        z = np.load(cache_path, allow_pickle=True)
+        R_tX = z["R_tX"]
+        t_index = z["t_index"]
+        X_index = z["X_index"]
+        return ElectricityData(R_tX=R_tX, t_index=t_index, X_index=X_index)
+
+    # 2) Ensure H5 present
+    h5_path = os.path.join(data_dir, h5_name)
+    download_with_progress(url, h5_path)
+
+    # 3) Read H5 (requires tables)
+    try:
+        df = pd.read_hdf(h5_path)
+    except Exception as e:
+        _require_tables_hint(e)
+
+    # 4) Basic preprocess
+    df = _coerce_datetime_index(df)
+    df = _fill_missing(df)
+
+    # 5) Time slice
+    if start is not None or stop is not None:
+        df = df.loc[start:stop]
+
+    # 6) Resample
+    if resample is not None:
+        if agg == "sum":
+            df = df.resample(resample).sum()
+        else:
+            df = df.resample(resample).mean()
+
+    # 7) Feature selection
+    df = _select_features(df, features=features, n_features=n_features, method=feature_method, seed=seed)
+
+    # 8) Downcast
+    if downcast_float32:
+        df = df.astype("float32", copy=False)
+
+    # 9) Build outputs
+    R_tX = df.to_numpy()
+    t_index = df.index.view("int64")          # UTC ns
+    X_index = np.asarray(df.columns.astype(str), dtype=object)
+
+    # 10) Cache
+    if cache:
+        np.savez_compressed(cache_path, R_tX=R_tX, t_index=t_index, X_index=X_index)
+        print(f"[datasets] Wrote cache: {cache_path}")
+
+    return ElectricityData(R_tX=R_tX, t_index=t_index, X_index=X_index)

forcing.py

@@ -0,0 +1,521 @@
+# forcing.py
+"""
+Forced / nonstationary variants of chaotic attractors.
+
+Standard output is always a single long trajectory:
+  R_tX : (T-burn_in, D)
+
+For ODEs, forcing MUST be applied inside the integrator (nonautonomous RK4),
+not post-hoc. This file provides three forcing mechanisms for Lorenz-63:
+
+  1) Additive forcing to the dynamics (adds kappa*u(t) to dx/dt, dy/dt, dz/dt)
+  2) Parameter drift (e.g., rho(t) varies with time / forcing)
+  3) Regime switching (piecewise-constant params over time)
+
+All functions accept (T, dt) as primary inputs, with reasonable defaults.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable, Dict, Optional, Sequence, Tuple, Union
+
+import numpy as np
+
+Array = np.ndarray
+
+
+# -----------------------------
+# small utils
+# -----------------------------
+def _rng_from_seed(seed: Optional[int]) -> np.random.Generator:
+    return np.random.default_rng(None if seed is None else int(seed))
+
+
+def _discard_burn_in(traj: Array, burn_in: int) -> Array:
+    burn_in = int(burn_in)
+    if burn_in <= 0:
+        return traj
+    if burn_in >= traj.shape[0]:
+        raise ValueError(f"burn_in={burn_in} must be < T={traj.shape[0]}")
+    return traj[burn_in:]
+
+
+def _as_1d(u: Union[None, float, Array]) -> Optional[Array]:
+    if u is None:
+        return None
+    if np.isscalar(u):
+        return np.asarray([float(u)], dtype=float)
+    u = np.asarray(u, dtype=float)
+    if u.ndim != 1:
+        raise ValueError(f"Expected 1D forcing array, got shape {u.shape}")
+    return u
+
+
+def _build_time_grid(T: int, dt: float) -> Array:
+    return np.arange(int(T), dtype=float) * float(dt)
+
+
+def _interp_u(u_t: Array, t_grid: Array, tcur: float) -> float:
+    """
+    Forcing evaluated at arbitrary time using linear interpolation.
+    Uses endpoint values outside the grid.
+    """
+    if u_t.size == 1:
+        return float(u_t[0])
+    return float(np.interp(tcur, t_grid, u_t, left=u_t[0], right=u_t[-1]))
+
+
+def _kappa_vec(kappa: Union[float, Sequence[float], Array], axis: Optional[int] = None) -> Array:
+    """
+    Convert kappa specification into a length-3 vector for Lorenz63.
+    - If axis is provided, interpret kappa as scalar magnitude applied to that axis.
+    - Else if kappa is scalar -> applied to x-axis by default.
+    - Else if kappa has length 3 -> used directly.
+    """
+    if axis is not None:
+        kv = np.zeros(3, dtype=float)
+        kv[int(axis)] = float(kappa)  # type: ignore[arg-type]
+        return kv
+
+    if np.isscalar(kappa):
+        # default: force x-equation
+        return np.array([float(kappa), 0.0, 0.0], dtype=float)
+
+    kv = np.asarray(kappa, dtype=float).reshape(-1)
+    if kv.size != 3:
+        raise ValueError("kappa must be scalar, length-3, or scalar+axis")
+    return kv
+
+
+# -----------------------------
+# forcing signal helpers (optional)
+# -----------------------------
+def forcing_sine(T: int, dt: float, *, amp: float = 1.0, freq_hz: float = 0.1, phase: float = 0.0, bias: float = 0.0) -> Array:
+    """
+    u(t) = bias + amp * sin(2*pi*freq_hz*t + phase)
+    Returns u_t of length T aligned with sample times n*dt.
+    """
+    t = _build_time_grid(T, dt)
+    return (bias + amp * np.sin(2.0 * np.pi * float(freq_hz) * t + float(phase))).astype(float)
+
+
+def forcing_piecewise_constant(T: int, *, values: Sequence[float], lengths: Sequence[int]) -> Array:
+    """
+    Build a piecewise-constant forcing u_t of length T.
+    Example: values=[0,1,0.5], lengths=[2000,3000,5000]
+    """
+    vals = list(values)
+    lens = list(map(int, lengths))
+    if len(vals) != len(lens):
+        raise ValueError("values and lengths must have same length")
+    u = np.concatenate([np.full(L, v, dtype=float) for v, L in zip(vals, lens)], axis=0)
+    if u.size < T:
+        u = np.pad(u, (0, T - u.size), mode="edge")
+    return u[:T]
+
+
+# -----------------------------
+# core integrator: nonautonomous RK4 for Lorenz63
+# -----------------------------
+def _lorenz63_step_rk4(
+    x: Array,
+    t: float,
+    dt: float,
+    *,
+    sigma_fn: Callable[[float], float],
+    rho_fn: Callable[[float], float],
+    beta_fn: Callable[[float], float],
+    force_fn: Callable[[float], Array],   # returns (3,) additive term to dx/dt
+) -> Array:
+    """
+    One RK4 step for x' = f(t,x) with Lorenz63 + forcing.
+    """
+    x = np.asarray(x, dtype=float)
+
+    def f(tcur: float, s: Array) -> Array:
+        sigma = float(sigma_fn(tcur))
+        rho = float(rho_fn(tcur))
+        beta = float(beta_fn(tcur))
+
+        fx = np.empty(3, dtype=float)
+        fx[0] = sigma * (s[1] - s[0])
+        fx[1] = s[0] * (rho - s[2]) - s[1]
+        fx[2] = s[0] * s[1] - beta * s[2]
+        fx += force_fn(tcur)
+        return fx
+
+    k1 = f(t, x)
+    k2 = f(t + 0.5 * dt, x + 0.5 * dt * k1)
+    k3 = f(t + 0.5 * dt, x + 0.5 * dt * k2)
+    k4 = f(t + dt, x + dt * k3)
+    return x + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)
+
+
+def _simulate_lorenz63_nonauto(
+    T: int,
+    dt: float,
+    *,
+    x0: Array,
+    sigma_fn: Callable[[float], float],
+    rho_fn: Callable[[float], float],
+    beta_fn: Callable[[float], float],
+    force_fn: Callable[[float], Array],
+) -> Array:
+    T = int(T)
+    dt = float(dt)
+    x = np.asarray(x0, dtype=float).reshape(3).copy()
+    traj = np.zeros((T, 3), dtype=float)
+    t = 0.0
+    for n in range(T):
+        traj[n] = x
+        x = _lorenz63_step_rk4(
+            x, t, dt,
+            sigma_fn=sigma_fn,
+            rho_fn=rho_fn,
+            beta_fn=beta_fn,
+            force_fn=force_fn,
+        )
+        t += dt
+    return traj
+
+
+# -----------------------------
+# 1) Additive forcing: x' = f(x) + kappa * u(t)
+# -----------------------------
+def lorenz63_forced_additive(
+    T: int,
+    dt: float = 0.01,
+    *,
+    sigma: float = 10.0,
+    rho: float = 28.0,
+    beta: float = 8.0 / 3.0,
+    u_t: Optional[Array] = None,
+    u_fn: Optional[Callable[[float], float]] = None,
+    kappa: Union[float, Sequence[float], Array] = 1.0,
+    axis: Optional[int] = 0,
+    x0: Optional[Array] = None,
+    seed: Optional[int] = 0,
+    burn_in: int = 0,
+    return_u: bool = False,
+) -> Union[Array, Tuple[Array, Array]]:
+    """
+    Additive forcing applied to dx/dt,dy/dt,dz/dt.
+
+    You may provide either:
+      - u_t: array of length T (sampled at n*dt), or
+      - u_fn: callable u(t)
+
+    kappa:
+      - scalar + axis (default) => applies to one coordinate derivative
+      - length-3 vector => applies to all derivatives
+      - scalar with axis=None => defaults to x-axis
+
+    Returns:
+      traj : (T-burn_in, 3)
+      and optionally u_used : (T,) if return_u=True
+    """
+    rng = _rng_from_seed(seed)
+    if x0 is None:
+        x0 = np.array([1.0, 1.0, 1.0], dtype=float) + 0.01 * rng.standard_normal(3)
+    else:
+        x0 = np.asarray(x0, dtype=float).reshape(3)
+
+    kv = _kappa_vec(kappa, axis=axis)
+
+    # build forcing evaluator
+    if u_fn is None:
+        ut = _as_1d(u_t)
+        if ut is None:
+            ut = np.zeros(int(T), dtype=float)
+        # ensure length T
+        if ut.size != int(T):
+            if ut.size < int(T):
+                ut = np.pad(ut, (0, int(T) - ut.size), mode="edge")
+            ut = ut[: int(T)]
+        t_grid = _build_time_grid(T, dt)
+
+        def u_eval(tcur: float) -> float:
+            return _interp_u(ut, t_grid, tcur)
+
+        u_used = ut
+    else:
+        def u_eval(tcur: float) -> float:
+            return float(u_fn(tcur))
+
+        # produce a sampled u for logging/return
+        t_grid = _build_time_grid(T, dt)
+        u_used = np.array([u_eval(ti) for ti in t_grid], dtype=float)
+
+    def sigma_fn(_t: float) -> float:
+        return float(sigma)
+
+    def rho_fn(_t: float) -> float:
+        return float(rho)
+
+    def beta_fn(_t: float) -> float:
+        return float(beta)
+
+    def force_fn(tcur: float) -> Array:
+        return kv * u_eval(tcur)
+
+    traj = _simulate_lorenz63_nonauto(
+        T=T, dt=dt, x0=x0,
+        sigma_fn=sigma_fn, rho_fn=rho_fn, beta_fn=beta_fn,
+        force_fn=force_fn
+    )
+    traj = _discard_burn_in(traj, burn_in)
+    u_used2 = _discard_burn_in(u_used[:, None], burn_in).reshape(-1)  # align
+    return (traj, u_used2) if return_u else traj
+
+
+# -----------------------------
+# 2) Parameter drift: rho(t) varies (optionally driven by u(t))
+# -----------------------------
+def lorenz63_param_drift(
+    T: int,
+    dt: float = 0.01,
+    *,
+    sigma: float = 10.0,
+    rho0: float = 28.0,
+    beta: float = 8.0 / 3.0,
+    # specify rho(t) either directly or via forcing
+    rho_t: Optional[Array] = None,                   # length T at n*dt
+    rho_fn: Optional[Callable[[float], float]] = None,
+    u_t: Optional[Array] = None,                     # used only if rho_t/rho_fn not provided
+    u_fn: Optional[Callable[[float], float]] = None,
+    kappa_rho: float = 1.0,                          # rho(t) = rho0 + kappa_rho * u(t)
+    # optional additive forcing simultaneously (set to 0 to disable)
+    kappa_force: Union[float, Sequence[float], Array] = 0.0,
+    axis_force: Optional[int] = 0,
+    x0: Optional[Array] = None,
+    seed: Optional[int] = 0,
+    burn_in: int = 0,
+    return_rho: bool = False,
+) -> Union[Array, Tuple[Array, Array]]:
+    """
+    Nonstationary Lorenz63 where rho(t) drifts.
+
+    Priority for defining rho(t):
+      1) rho_fn(t)
+      2) rho_t sampled at n*dt
+      3) rho0 + kappa_rho * u(t)  (u from u_fn or u_t; defaults to 0)
+
+    You can also add *simultaneous additive forcing* via kappa_force/axis_force.
+
+    Returns:
+      traj : (T-burn_in, 3)
+      and optionally rho_used : (T-burn_in,) if return_rho=True
+    """
+    rng = _rng_from_seed(seed)
+    if x0 is None:
+        x0 = np.array([1.0, 1.0, 1.0], dtype=float) + 0.01 * rng.standard_normal(3)
+    else:
+        x0 = np.asarray(x0, dtype=float).reshape(3)
+
+    # build u evaluator (only if needed)
+    ut = _as_1d(u_t)
+    if ut is None:
+        ut = np.zeros(int(T), dtype=float)
+    if ut.size != int(T):
+        if ut.size < int(T):
+            ut = np.pad(ut, (0, int(T) - ut.size), mode="edge")
+        ut = ut[: int(T)]
+    t_grid = _build_time_grid(T, dt)
+
+    if u_fn is None:
+        def u_eval(tcur: float) -> float:
+            return _interp_u(ut, t_grid, tcur)
+        u_used = ut
+    else:
+        def u_eval(tcur: float) -> float:
+            return float(u_fn(tcur))
+        u_used = np.array([u_eval(ti) for ti in t_grid], dtype=float)
+
+    # build rho evaluator
+    if rho_fn is not None:
+        def rho_eval(tcur: float) -> float:
+            return float(rho_fn(tcur))
+        rho_used = np.array([rho_eval(ti) for ti in t_grid], dtype=float)
+    elif rho_t is not None:
+        rt = _as_1d(rho_t)
+        if rt is None:
+            raise ValueError("rho_t must be a 1D array if provided")
+        if rt.size != int(T):
+            if rt.size < int(T):
+                rt = np.pad(rt, (0, int(T) - rt.size), mode="edge")
+            rt = rt[: int(T)]
+
+        def rho_eval(tcur: float) -> float:
+            return float(np.interp(tcur, t_grid, rt, left=rt[0], right=rt[-1]))
+        rho_used = rt
+    else:
+        def rho_eval(tcur: float) -> float:
+            return float(rho0 + kappa_rho * u_eval(tcur))
+        rho_used = np.array([rho_eval(ti) for ti in t_grid], dtype=float)
+
+    kv_force = _kappa_vec(kappa_force, axis=axis_force)
+
+    def sigma_fn(_t: float) -> float:
+        return float(sigma)
+
+    def rho_fn2(tcur: float) -> float:
+        return float(rho_eval(tcur))
+
+    def beta_fn(_t: float) -> float:
+        return float(beta)
+
+    def force_fn(tcur: float) -> Array:
+        # optional additive forcing in dynamics
+        return kv_force * u_eval(tcur) if np.any(kv_force != 0.0) else np.zeros(3, dtype=float)
+
+    traj = _simulate_lorenz63_nonauto(
+        T=T, dt=dt, x0=x0,
+        sigma_fn=sigma_fn, rho_fn=rho_fn2, beta_fn=beta_fn,
+        force_fn=force_fn
+    )
+
+    traj = _discard_burn_in(traj, burn_in)
+    rho_used2 = _discard_burn_in(rho_used[:, None], burn_in).reshape(-1)
+    return (traj, rho_used2) if return_rho else traj
+
+
+# -----------------------------
+# 3) Regime switching: piecewise-constant params (sigma,rho,beta)
+# -----------------------------
+@dataclass(frozen=True)
+class Lorenz63Params:
+    sigma: float = 10.0
+    rho: float = 28.0
+    beta: float = 8.0 / 3.0
+
+
+def lorenz63_regime_switch(
+    T: int,
+    dt: float = 0.01,
+    *,
+    # a schedule of regimes: (start_step, params)
+    schedule: Sequence[Tuple[int, Lorenz63Params]] = ((0, Lorenz63Params()),),
+    # optional additive forcing simultaneously
+    u_t: Optional[Array] = None,
+    u_fn: Optional[Callable[[float], float]] = None,
+    kappa: Union[float, Sequence[float], Array] = 0.0,
+    axis: Optional[int] = 0,
+    x0: Optional[Array] = None,
+    seed: Optional[int] = 0,
+    burn_in: int = 0,
+    return_params: bool = False,
+) -> Union[Array, Tuple[Array, Dict[str, Array]]]:
+    """
+    Piecewise-constant parameter switching. `schedule` is a list of
+    (start_step, Lorenz63Params). Example:
+
+      schedule = [
+        (0,    Lorenz63Params(rho=28.0)),
+        (8000, Lorenz63Params(rho=35.0)),
+        (14000,Lorenz63Params(rho=25.0)),
+      ]
+
+    Regime selection at time t uses idx = floor(t/dt), then picks the last
+    schedule entry with start_step <= idx.
+
+    Returns:
+      traj : (T-burn_in, 3)
+      and optionally a dict with sigma_t, rho_t, beta_t arrays (aligned).
+    """
+    rng = _rng_from_seed(seed)
+    if x0 is None:
+        x0 = np.array([1.0, 1.0, 1.0], dtype=float) + 0.01 * rng.standard_normal(3)
+    else:
+        x0 = np.asarray(x0, dtype=float).reshape(3)
+
+    # build regime arrays (per step)
+    T = int(T)
+    sigma_t = np.empty(T, dtype=float)
+    rho_t = np.empty(T, dtype=float)
+    beta_t = np.empty(T, dtype=float)
+
+    sched = sorted([(int(s), p) for s, p in schedule], key=lambda z: z[0])
+    if not sched or sched[0][0] != 0:
+        raise ValueError("schedule must start at step 0")
+
+    # fill by segments
+    for i, (s0, p0) in enumerate(sched):
+        s1 = sched[i + 1][0] if i + 1 < len(sched) else T
+        s0 = max(0, min(T, s0))
+        s1 = max(s0, min(T, s1))
+        sigma_t[s0:s1] = float(p0.sigma)
+        rho_t[s0:s1] = float(p0.rho)
+        beta_t[s0:s1] = float(p0.beta)
+
+    # forcing evaluator (optional)
+    kv = _kappa_vec(kappa, axis=axis)
+
+    if u_fn is None:
+        ut = _as_1d(u_t)
+        if ut is None:
+            ut = np.zeros(T, dtype=float)
+        if ut.size != T:
+            if ut.size < T:
+                ut = np.pad(ut, (0, T - ut.size), mode="edge")
+            ut = ut[:T]
+        t_grid = _build_time_grid(T, dt)
+
+        def u_eval(tcur: float) -> float:
+            return _interp_u(ut, t_grid, tcur)
+
+        u_used = ut
+    else:
+        t_grid = _build_time_grid(T, dt)
+
+        def u_eval(tcur: float) -> float:
+            return float(u_fn(tcur))
+
+        u_used = np.array([u_eval(ti) for ti in t_grid], dtype=float)
+
+    # parameter evaluators: stepwise constant
+    def _param_at(arr: Array, tcur: float) -> float:
+        idx = int(tcur / float(dt))
+        if idx < 0:
+            idx = 0
+        elif idx >= T:
+            idx = T - 1
+        return float(arr[idx])
+
+    def sigma_fn(tcur: float) -> float:
+        return _param_at(sigma_t, tcur)
+
+    def rho_fn(tcur: float) -> float:
+        return _param_at(rho_t, tcur)
+
+    def beta_fn(tcur: float) -> float:
+        return _param_at(beta_t, tcur)
+
+    def force_fn(tcur: float) -> Array:
+        if np.all(kv == 0.0):
+            return np.zeros(3, dtype=float)
+        return kv * u_eval(tcur)
+
+    traj = _simulate_lorenz63_nonauto(
+        T=T, dt=dt, x0=x0,
+        sigma_fn=sigma_fn, rho_fn=rho_fn, beta_fn=beta_fn,
+        force_fn=force_fn
+    )
+
+    traj = _discard_burn_in(traj, burn_in)
+    if not return_params:
+        return traj
+
+    # align parameter arrays with burn_in
+    sigma_used = _discard_burn_in(sigma_t[:, None], burn_in).reshape(-1)
+    rho_used = _discard_burn_in(rho_t[:, None], burn_in).reshape(-1)
+    beta_used = _discard_burn_in(beta_t[:, None], burn_in).reshape(-1)
+    u_used2 = _discard_burn_in(u_used[:, None], burn_in).reshape(-1)
+
+    info = {
+        "sigma_t": sigma_used,
+        "rho_t": rho_used,
+        "beta_t": beta_used,
+        "u_t": u_used2,
+    }
+    return traj, info

metrics.py

@@ -0,0 +1,494 @@
+# Minimal usage
+# from metrics import compare_all, mse_by_horizon
+# 
+# metrics = compare_all(truth, pred, burn_in=0, acf_max_lag=200, mmd_sample=1024, seed=0)
+# print(metrics)
+# 
+# mseh = mse_by_horizon(truth, pred)  # per-step MSE curve
+
+
+# metrics.py
+from __future__ import annotations
+
+import numpy as np
+from typing import Dict, Optional, Tuple
+
+
+def _as_2d(a: np.ndarray) -> np.ndarray:
+    a = np.asarray(a)
+    if a.ndim == 1:
+        return a[:, None]
+    if a.ndim != 2:
+        raise ValueError(f"Expected 1D or 2D array, got shape {a.shape}")
+    return a
+
+
+def _slice_time(
+    R: np.ndarray,
+    Q: np.ndarray,
+    *,
+    burn_in: int = 0,
+    T_max: Optional[int] = None,
+    stride: int = 1,
+) -> Tuple[np.ndarray, np.ndarray]:
+    R = _as_2d(R)
+    Q = _as_2d(Q)
+    if R.shape != Q.shape:
+        raise ValueError(f"R and Q must have same shape, got {R.shape} vs {Q.shape}")
+
+    b = int(burn_in)
+    s = int(stride)
+    if b < 0 or s <= 0:
+        raise ValueError("burn_in must be >=0 and stride must be >=1")
+
+    T = R.shape[0]
+    t0 = min(b, T)
+    t1 = T if T_max is None else min(T, t0 + int(T_max))
+    R2 = R[t0:t1:s]
+    Q2 = Q[t0:t1:s]
+    return R2, Q2
+
+
+# -----------------------------
+# 1) MSE
+# -----------------------------
+def mse(R_tX: np.ndarray, Q_tX: np.ndarray, *, burn_in: int = 0, T_max: Optional[int] = None, stride: int = 1) -> float:
+    R, Q = _slice_time(R_tX, Q_tX, burn_in=burn_in, T_max=T_max, stride=stride)
+    return float(np.mean((R - Q) ** 2))
+
+
+def mse_by_horizon(R_tX: np.ndarray, Q_tX: np.ndarray, *, burn_in: int = 0, T_max: Optional[int] = None, stride: int = 1) -> np.ndarray:
+    """
+    Per-step MSE across features: returns shape (T',)
+    """
+    R, Q = _slice_time(R_tX, Q_tX, burn_in=burn_in, T_max=T_max, stride=stride)
+    return np.mean((R - Q) ** 2, axis=1)
+
+
+def mse_per_feature(R_tX: np.ndarray, Q_tX: np.ndarray, *, burn_in: int = 0, T_max: Optional[int] = None, stride: int = 1) -> np.ndarray:
+    """
+    Per-feature MSE: returns shape (X,)
+    """
+    R, Q = _slice_time(R_tX, Q_tX, burn_in=burn_in, T_max=T_max, stride=stride)
+    return np.mean((R - Q) ** 2, axis=0)
+
+
+# -----------------------------
+# 2) PSD log-MSE (spectrum mismatch)
+# -----------------------------
+def _hann(n: int) -> np.ndarray:
+    if n <= 1:
+        return np.ones((n,), dtype=float)
+    i = np.arange(n, dtype=float)
+    return 0.5 - 0.5 * np.cos(2.0 * np.pi * i / (n - 1))
+
+
+def psd_logmse(
+    R_tX: np.ndarray,
+    Q_tX: np.ndarray,
+    *,
+    burn_in: int = 0,
+    T_max: Optional[int] = None,
+    stride: int = 1,
+    detrend: bool = True,
+    window: str = "hann",
+    eps: float = 1e-12,
+    normalize_per_feature: bool = True,
+) -> float:
+    """
+    Compare log power spectra of R and Q.
+    - computes one-sided PSD via rFFT along time for each feature
+    - returns mean squared error of log(PSD + eps) over (freq, feature)
+
+    normalize_per_feature=True rescales PSD of each feature to sum to 1
+    (so mismatch reflects shape, not total variance).
+    """
+    R, Q = _slice_time(R_tX, Q_tX, burn_in=burn_in, T_max=T_max, stride=stride)
+    T, X = R.shape
+    if T < 8:
+        return float("nan")
+
+    Rw = R.astype(float)
+    Qw = Q.astype(float)
+
+    if detrend:
+        Rw = Rw - np.mean(Rw, axis=0, keepdims=True)
+        Qw = Qw - np.mean(Qw, axis=0, keepdims=True)
+
+    if window is None or window.lower() == "none":
+        w = np.ones((T,), dtype=float)
+    elif window.lower() == "hann":
+        w = _hann(T)
+    else:
+        raise ValueError("window must be 'hann' or None/'none'")
+
+    # apply window (broadcast over features)
+    Rw = Rw * w[:, None]
+    Qw = Qw * w[:, None]
+
+    # rFFT
+    FR = np.fft.rfft(Rw, axis=0)
+    FQ = np.fft.rfft(Qw, axis=0)
+
+    # raw power
+    PR = (np.abs(FR) ** 2) / max(T, 1)
+    PQ = (np.abs(FQ) ** 2) / max(T, 1)
+
+    if normalize_per_feature:
+        PR = PR / (np.sum(PR, axis=0, keepdims=True) + eps)
+        PQ = PQ / (np.sum(PQ, axis=0, keepdims=True) + eps)
+
+    LR = np.log(PR + eps)
+    LQ = np.log(PQ + eps)
+    return float(np.mean((LR - LQ) ** 2))
+
+
+# -----------------------------
+# 3) ACF-MSE (autocorrelation mismatch)
+# -----------------------------
+def _acf_1d(x: np.ndarray, max_lag: int) -> np.ndarray:
+    """
+    Normalized autocorrelation via FFT, lags 0..max_lag.
+    """
+    x = np.asarray(x, float)
+    n = x.size
+    if n == 0:
+        return np.zeros((max_lag + 1,), dtype=float)
+    x = x - np.mean(x)
+    # next power of 2 for speed
+    nfft = 1 << int(np.ceil(np.log2(max(1, 2 * n - 1))))
+    fx = np.fft.rfft(x, n=nfft)
+    acov = np.fft.irfft(fx * np.conj(fx), n=nfft)[: max_lag + 1]
+    acov = acov / max(n, 1)
+    if acov[0] <= 0:
+        return np.zeros((max_lag + 1,), dtype=float)
+    return acov / acov[0]
+
+
+def acf_mse(
+    R_tX: np.ndarray,
+    Q_tX: np.ndarray,
+    *,
+    max_lag: int = 200,
+    burn_in: int = 0,
+    T_max: Optional[int] = None,
+    stride: int = 1,
+) -> float:
+    """
+    Mean squared error between autocorrelation functions of R and Q,
+    averaged over features and lags 1..max_lag (excluding lag 0).
+    """
+    R, Q = _slice_time(R_tX, Q_tX, burn_in=burn_in, T_max=T_max, stride=stride)
+    T, X = R.shape
+    if T < max_lag + 2:
+        max_lag = max(1, min(max_lag, T - 2))
+
+    errs = []
+    for j in range(X):
+        aR = _acf_1d(R[:, j], max_lag=max_lag)
+        aQ = _acf_1d(Q[:, j], max_lag=max_lag)
+        # exclude lag 0 (always 1)
+        errs.append(np.mean((aR[1:] - aQ[1:]) ** 2))
+    return float(np.mean(errs)) if errs else float("nan")
+
+
+# -----------------------------
+# 4) MMD² (distribution/attractor fidelity)
+# -----------------------------
+def _sq_dists(A: np.ndarray, B: np.ndarray) -> np.ndarray:
+    # ||a-b||^2 = ||a||^2 + ||b||^2 - 2a·b
+    AA = np.sum(A * A, axis=1, keepdims=True)
+    BB = np.sum(B * B, axis=1, keepdims=True).T
+    D2 = AA + BB - 2.0 * (A @ B.T)
+    return np.maximum(D2, 0.0)
+
+
+def _median_bandwidth(Z: np.ndarray, max_pairs: int = 4096, seed: int = 0, eps: float = 1e-12) -> float:
+    rng = np.random.default_rng(seed)
+    n = Z.shape[0]
+    if n < 2:
+        return 1.0
+    m = min(max_pairs, n * (n - 1) // 2)
+    i = rng.integers(0, n, size=m)
+    j = rng.integers(0, n, size=m)
+    mask = i != j
+    i, j = i[mask], j[mask]
+    if i.size == 0:
+        return 1.0
+    d2 = np.sum((Z[i] - Z[j]) ** 2, axis=1)
+    med = np.median(d2)
+    # kernel uses exp(-||x-y||^2/(2*sigma^2))
+    sigma = np.sqrt(max(med, 0.0) / 2.0 + eps)
+    return float(max(sigma, 1e-6))
+
+
+def mmd_rbf(
+    R_tX: np.ndarray,
+    Q_tX: np.ndarray,
+    *,
+    burn_in: int = 0,
+    T_max: Optional[int] = None,
+    stride: int = 1,
+    sample: int = 2048,
+    sigma: Optional[float] = None,
+    sigma_seed: int = 0,
+    unbiased: bool = True,
+    seed: int = 0,
+) -> float:
+    """
+    MMD^2 between the empirical distributions of state samples from R and Q.
+
+    Uses an RBF kernel: k(x,y) = exp(-||x-y||^2/(2*sigma^2))
+
+    Notes:
+      - O(sample^2) time/memory, so keep sample modest (e.g. 512–4096).
+      - If sigma is None, uses median heuristic on pooled samples.
+    """
+    R, Q = _slice_time(R_tX, Q_tX, burn_in=burn_in, T_max=T_max, stride=stride)
+    rng = np.random.default_rng(seed)
+
+    T = R.shape[0]
+    if T == 0:
+        return float("nan")
+
+    n = min(int(sample), T)
+    idxR = rng.choice(T, size=n, replace=False) if n < T else np.arange(T)
+    idxQ = rng.choice(T, size=n, replace=False) if n < T else np.arange(T)
+
+    Xs = R[idxR].astype(float)
+    Ys = Q[idxQ].astype(float)
+
+    if sigma is None:
+        Z = np.vstack([Xs, Ys])
+        sigma = _median_bandwidth(Z, seed=sigma_seed)
+
+    sig2 = float(sigma) ** 2
+    if sig2 <= 0:
+        raise ValueError("sigma must be > 0")
+
+    Kxx = np.exp(-_sq_dists(Xs, Xs) / (2.0 * sig2))
+    Kyy = np.exp(-_sq_dists(Ys, Ys) / (2.0 * sig2))
+    Kxy = np.exp(-_sq_dists(Xs, Ys) / (2.0 * sig2))
+
+    if unbiased and n > 1:
+        np.fill_diagonal(Kxx, 0.0)
+        np.fill_diagonal(Kyy, 0.0)
+        mmd2 = (Kxx.sum() / (n * (n - 1))) + (Kyy.sum() / (n * (n - 1))) - 2.0 * (Kxy.mean())
+    else:
+        mmd2 = Kxx.mean() + Kyy.mean() - 2.0 * Kxy.mean()
+
+    return float(max(mmd2, 0.0))
+
+
+# -----------------------------
+# convenience: compute all metrics at once
+# -----------------------------
+def compare_all(
+    R_tX: np.ndarray,
+    Q_tX: np.ndarray,
+    *,
+    burn_in: int = 0,
+    T_max: Optional[int] = None,
+    stride: int = 1,
+    acf_max_lag: int = 200,
+    mmd_sample: int = 2048,
+    mmd_sigma: Optional[float] = None,
+    seed: int = 0,
+) -> Dict[str, float]:
+    """
+    Returns a dict with scalar metrics.
+    """
+    out = {}
+    out["mse"] = mse(R_tX, Q_tX, burn_in=burn_in, T_max=T_max, stride=stride)
+    out["psd_logmse"] = psd_logmse(R_tX, Q_tX, burn_in=burn_in, T_max=T_max, stride=stride)
+    out["acf_mse"] = acf_mse(R_tX, Q_tX, max_lag=acf_max_lag, burn_in=burn_in, T_max=T_max, stride=stride)
+    out["mmd2_rbf"] = mmd_rbf(
+        R_tX, Q_tX,
+        burn_in=burn_in, T_max=T_max, stride=stride,
+        sample=mmd_sample, sigma=mmd_sigma,
+        seed=seed
+    )
+    return out
+
+
+
+from typing import Optional, Tuple, Dict
+
+def compare_all_with_curves(
+    R_tX: np.ndarray,
+    Q_tX: np.ndarray,
+    *,
+    burn_in: int = 0,
+    T_max: Optional[int] = None,
+    stride: int = 1,
+    acf_max_lag: int = 200,
+    mmd_sample: int = 2048,
+    mmd_sigma: Optional[float] = None,
+    seed: int = 0,
+    ) -> Tuple[Dict[str, float], Dict[str, np.ndarray]]:
+    """
+    Returns:
+      scalars: dict[str, float]
+      curves : dict[str, np.ndarray]
+    """
+    scalars = compare_all(
+        R_tX, Q_tX,
+        burn_in=burn_in, T_max=T_max, stride=stride,
+        acf_max_lag=acf_max_lag,
+        mmd_sample=mmd_sample,
+        mmd_sigma=mmd_sigma,
+        seed=seed,
+    )
+    curves = {
+        "mse_by_horizon": mse_by_horizon(R_tX, Q_tX, burn_in=burn_in, T_max=T_max, stride=stride),
+        "mse_per_feature": mse_per_feature(R_tX, Q_tX, burn_in=burn_in, T_max=T_max, stride=stride),
+    }
+    return scalars, curves
+
+
+
+
+
+
+
+
+def compare_all_fast(
+    R_tX: np.ndarray,
+    Q_tX: np.ndarray,
+    *,
+    burn_in: int = 0,
+    T_max: Optional[int] = None,
+    stride: int = 1,
+    acf_max_lag: int = 200,
+    include_mmd_rff: bool = False,
+    mmd_sample: int = 1024,
+    mmd_rff_dim: int = 512,
+    mmd_sigma: Optional[float] = None,
+    seed: int = 0,
+) -> Dict[str, float]:
+    out = {}
+    out["mse"] = mse(R_tX, Q_tX, burn_in=burn_in, T_max=T_max, stride=stride)
+    out["psd_logmse"] = psd_logmse(R_tX, Q_tX, burn_in=burn_in, T_max=T_max, stride=stride)
+    out["acf_mse"] = acf_mse_fast(
+        R_tX, Q_tX,
+        max_lag=acf_max_lag,
+        burn_in=burn_in, T_max=T_max, stride=stride,
+    )
+    if include_mmd_rff:
+        out["mmd2_rff"] = mmd_rff(
+            R_tX, Q_tX,
+            burn_in=burn_in, T_max=T_max, stride=stride,
+            sample=mmd_sample, rff_dim=mmd_rff_dim,
+            sigma=mmd_sigma,
+            seed=seed,
+        )
+    return out
+
+import numpy as np
+from typing import Optional, Dict, Tuple
+
+def acf_mse_fast(
+    R_tX: np.ndarray,
+    Q_tX: np.ndarray,
+    *,
+    max_lag: int = 200,
+    burn_in: int = 0,
+    T_max: Optional[int] = None,
+    stride: int = 1,
+    eps: float = 1e-12,
+    ) -> float:
+    R, Q = _slice_time(R_tX, Q_tX, burn_in=burn_in, T_max=T_max, stride=stride)
+    T, X = R.shape
+    if T < 8:
+        return float("nan")
+    if T < max_lag + 2:
+        max_lag = max(1, min(max_lag, T - 2))
+
+    def acf_all(Z: np.ndarray) -> np.ndarray:
+        Z = Z.astype(np.float32)
+        Z = Z - Z.mean(axis=0, keepdims=True)
+        n = Z.shape[0]
+        nfft = 1 << int(np.ceil(np.log2(max(1, 2*n - 1))))
+        F = np.fft.rfft(Z, n=nfft, axis=0)
+        acov = np.fft.irfft(F * np.conj(F), n=nfft, axis=0)[: max_lag + 1]
+        acov = acov / max(n, 1)
+        ac0 = acov[0:1, :]  # (1,X)
+        return np.where(ac0 > eps, acov / ac0, 0.0)
+
+    aR = acf_all(R)
+    aQ = acf_all(Q)
+    return float(np.mean((aR[1:] - aQ[1:]) ** 2))
+
+
+def mmd_rff(
+    R_tX: np.ndarray,
+    Q_tX: np.ndarray,
+    *,
+    burn_in: int = 0,
+    T_max: Optional[int] = None,
+    stride: int = 1,
+    sample: int = 2048,
+    rff_dim: int = 512,
+    sigma: Optional[float] = None,
+    seed: int = 0,
+    ) -> float:
+    R, Q = _slice_time(R_tX, Q_tX, burn_in=burn_in, T_max=T_max, stride=stride)
+    rng = np.random.default_rng(seed)
+    T, X = R.shape
+    if T == 0:
+        return float("nan")
+
+    n = min(int(sample), T)
+    idxR = rng.choice(T, size=n, replace=False) if n < T else np.arange(T)
+    idxQ = rng.choice(T, size=n, replace=False) if n < T else np.arange(T)
+
+    Xs = R[idxR].astype(np.float32)
+    Ys = Q[idxQ].astype(np.float32)
+
+    if sigma is None:
+        Z = np.vstack([Xs, Ys])
+        sigma = _median_bandwidth(Z, seed=seed)
+
+    sigma = float(sigma)
+    W = rng.normal(0.0, 1.0 / sigma, size=(X, rff_dim)).astype(np.float32)
+    b = rng.uniform(0.0, 2*np.pi, size=(rff_dim,)).astype(np.float32)
+
+    def phi(A):
+        return np.sqrt(2.0 / rff_dim) * np.cos(A @ W + b)
+
+    mx = phi(Xs).mean(axis=0)
+    my = phi(Ys).mean(axis=0)
+    return float(np.mean((mx - my) ** 2))
+
+
+
+
+def compare_all_fast_with_curves(
+    R_tX: np.ndarray,
+    Q_tX: np.ndarray,
+    *,
+    burn_in: int = 0,
+    T_max: Optional[int] = None,
+    stride: int = 1,
+    acf_max_lag: int = 200,
+    include_mmd_rff: bool = False,
+    mmd_sample: int = 1024,
+    mmd_rff_dim: int = 512,
+    mmd_sigma: Optional[float] = None,
+    seed: int = 0,
+    ) -> Tuple[Dict[str, float], Dict[str, np.ndarray]]:
+    scalars = compare_all_fast(
+        R_tX, Q_tX,
+        burn_in=burn_in, T_max=T_max, stride=stride,
+        acf_max_lag=acf_max_lag,
+        include_mmd_rff=include_mmd_rff,
+        mmd_sample=mmd_sample, mmd_rff_dim=mmd_rff_dim,
+        mmd_sigma=mmd_sigma,
+        seed=seed,
+    )
+    curves = {
+        "mse_by_horizon": mse_by_horizon(R_tX, Q_tX, burn_in=burn_in, T_max=T_max, stride=stride),
+        "mse_per_feature": mse_per_feature(R_tX, Q_tX, burn_in=burn_in, T_max=T_max, stride=stride),
+    }
+    return scalars, curves
+