diff --git a/experiments/abc_adaptive_ark_study.py b/experiments/abc_adaptive_ark_study.py new file mode 100644 index 0000000..940e2fe --- /dev/null +++ b/experiments/abc_adaptive_ark_study.py @@ -0,0 +1,673 @@ +#!/usr/bin/env python3 +""" +abc_adaptive_ark_study.py + +Adaptive step-size study for: + - Strang-3 splitting (2nd order) + - Yoshida-3 splitting (4th order) + - Embedded RK baseline using pythOS butcher_tableau embedded pairs + +Important note: +`additive_rk.py` in pythOS imports and uses tableau objects, but the actual +embedded RK pair definitions live in `butcher_tableau.py` under `embedded_pairs`. +So this script: + - verifies that `additive_rk.py` is importable from pythOS + - uses `butcher_tableau.embedded_pairs[...]` to obtain the pair + - runs the adaptive RK baseline on the full RHS y' = (A+B+C)y +""" + +from __future__ import annotations + +import argparse +import csv +import os +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import List, Optional, Tuple + +import numpy as np + +# --------------------------------------------------------------------- +# Repository-local imports +# --------------------------------------------------------------------- +THIS = Path(__file__).resolve() +REPO = THIS.parents[1] +if str(REPO) not in sys.path: + sys.path.insert(0, str(REPO)) + + +# ---------------------------- +# Utility helpers +# ---------------------------- + +def safe_norm(x: np.ndarray) -> float: + return float(np.linalg.norm(x)) + + +def clamp(x: float, lo: float, hi: float) -> float: + return float(max(lo, min(hi, x))) + + +def adapt_dt(dt: float, err_est: float, tol: float, p: int, safety: float = 0.9) -> float: + """ + Standard adaptive controller: + dt_new = dt * safety * (tol / err_est)^(1/(p+1)) + + where p is the order of the high solution. + """ + if err_est <= 0.0 or not np.isfinite(err_est): + return min(2.0 * dt, 10.0 * dt) + + fac = safety * (tol / err_est) ** (1.0 / (p + 1.0)) + fac = clamp(fac, 0.2, 5.0) + return dt * fac + + +# ---------------------------- +# ABC toy problem definition +# ---------------------------- + +def make_abc_matrices(Bmult: float) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Linear test problem: + y' = (A + B + C) y + with tunable stiffness in B via Bmult. + """ + A = np.array([ + [-1.0, 0.2, 0.0], + [ 0.0, -0.5, 0.1], + [ 0.0, 0.0, -0.2], + ], dtype=float) + + B = Bmult * np.array([ + [-2.0, 1.0, 0.0], + [-1.0, -2.0, 0.0], + [ 0.0, 0.0, -1.0], + ], dtype=float) + + C = np.array([ + [-0.1, 0.0, 0.0], + [ 0.0, 0.0, 0.3], + [ 0.0, -0.3, 0.0], + ], dtype=float) + + return A, B, C + + +def exact_solution(A: np.ndarray, B: np.ndarray, C: np.ndarray, y0: np.ndarray, T: float) -> np.ndarray: + """ + Exact solution using matrix exponential. + Uses scipy if available; otherwise falls back to eigendecomposition. + """ + M = A + B + C + try: + from scipy.linalg import expm + return expm(M * T) @ y0 + except Exception: + w, V = np.linalg.eig(M) + Vinv = np.linalg.inv(V) + expw = np.diag(np.exp(w * T)) + y = (V @ expw @ Vinv) @ y0 + return np.real_if_close(y, tol=1000) + + +# ---------------------------- +# Splitting building blocks +# ---------------------------- + +def flow_linear(M: np.ndarray, y: np.ndarray, h: float) -> np.ndarray: + """ + Exact flow for y' = M y over time h. + """ + try: + from scipy.linalg import expm + return expm(M * h) @ y + except Exception: + w, V = np.linalg.eig(M) + Vinv = np.linalg.inv(V) + expw = np.diag(np.exp(w * h)) + out = (V @ expw @ Vinv) @ y + return np.real_if_close(out, tol=1000) + + +def apply_ordering(A: np.ndarray, B: np.ndarray, C: np.ndarray, ordering: str) -> List[np.ndarray]: + """ + ordering is a string like 'CAB' specifying subflow order. + """ + mp = {"A": A, "B": B, "C": C} + return [mp[ch] for ch in ordering] + + +def strang_step(A: np.ndarray, B: np.ndarray, C: np.ndarray, y: np.ndarray, h: float, ordering: str) -> np.ndarray: + """ + Strang splitting for 3 operators: + phi_X(h/2) phi_Y(h/2) phi_Z(h) phi_Y(h/2) phi_X(h/2) + """ + X, Y, Z = apply_ordering(A, B, C, ordering) + y1 = flow_linear(X, y, 0.5 * h) + y2 = flow_linear(Y, y1, 0.5 * h) + y3 = flow_linear(Z, y2, h) + y4 = flow_linear(Y, y3, 0.5 * h) + y5 = flow_linear(X, y4, 0.5 * h) + return y5 + + +def yoshida_4th_from_strang(A: np.ndarray, B: np.ndarray, C: np.ndarray, y: np.ndarray, h: float, ordering: str) -> np.ndarray: + """ + 4th-order Yoshida composition from Strang: + S(w1 h) S(w2 h) S(w1 h) + """ + cbrt2 = 2.0 ** (1.0 / 3.0) + w1 = 1.0 / (2.0 - cbrt2) + w2 = -cbrt2 / (2.0 - cbrt2) + + y1 = strang_step(A, B, C, y, w1 * h, ordering) + y2 = strang_step(A, B, C, y1, w2 * h, ordering) + y3 = strang_step(A, B, C, y2, w1 * h, ordering) + return y3 + + +# ---------------------------- +# Adaptive CSV logging +# ---------------------------- + +@dataclass +class AdaptiveLogRow: + t: float + dt: float + accepted: int + err_est: float + y_norm_ratio: float + step_ratio: float + + +def write_adaptive_csv(path: str, rows: List[AdaptiveLogRow], summary_lines: List[str]) -> None: + """ + Writes a rectangular CSV with comment-header summary lines. + """ + with open(path, "w", newline="") as fcsv: + for line in summary_lines: + if not line.startswith("#"): + line = "# " + line + fcsv.write(line.rstrip() + "\n") + + w = csv.writer(fcsv) + w.writerow(["t", "dt", "accepted", "err_est", "y_norm_ratio", "step_ratio"]) + for r in rows: + w.writerow([ + f"{r.t:.16e}", + f"{r.dt:.16e}", + int(r.accepted), + f"{r.err_est:.16e}", + f"{r.y_norm_ratio:.16e}", + f"{r.step_ratio:.16e}", + ]) + + +# ---------------------------- +# Adaptive drivers for splitting methods +# ---------------------------- + +def step_with_two_dts( + method: str, + A: np.ndarray, + B: np.ndarray, + C: np.ndarray, + t: float, + y: np.ndarray, + dt: float, + ordering: str, +) -> Tuple[np.ndarray, np.ndarray]: + """ + One dt step and two half steps for local error estimation. + """ + if method == "Strang3": + y_big = strang_step(A, B, C, y, dt, ordering) + y_half = strang_step(A, B, C, y, 0.5 * dt, ordering) + y_half2 = strang_step(A, B, C, y_half, 0.5 * dt, ordering) + return y_big, y_half2 + + if method == "Yoshida3": + y_big = yoshida_4th_from_strang(A, B, C, y, dt, ordering) + y_half = yoshida_4th_from_strang(A, B, C, y, 0.5 * dt, ordering) + y_half2 = yoshida_4th_from_strang(A, B, C, y_half, 0.5 * dt, ordering) + return y_big, y_half2 + + raise ValueError(f"Unknown method={method}") + + +def run_adaptive_splitting( + method: str, + Bmult: float, + y0: np.ndarray, + T: float, + dt0: float, + tol: float, + ordering: str, + outdir: str, +): + A, B, C = make_abc_matrices(Bmult) + y_exact_T = exact_solution(A, B, C, y0, T) + exact_norm = max(safe_norm(y_exact_T), 1e-300) + y0_norm = max(safe_norm(y0), 1e-300) + + p = 2 if method == "Strang3" else 4 + + t = 0.0 + dt = float(dt0) + y = y0.copy() + + accept = 0 + reject = 0 + max_y_ratio = 1.0 + max_step_ratio = 1.0 + rows: List[AdaptiveLogRow] = [] + + while t < T - 1e-15: + dt = clamp(dt, 1e-12, 0.5) + if t + dt > T: + dt = T - t + + y_big, y_ref = step_with_two_dts(method, A, B, C, t, y, dt, ordering) + err_est = safe_norm(y_ref - y_big) + step_ratio = safe_norm(y_big) / max(safe_norm(y), 1e-300) + + if err_est <= tol: + t += dt + y = y_ref + accept += 1 + + y_ratio = safe_norm(y) / y0_norm + max_y_ratio = max(max_y_ratio, y_ratio) + max_step_ratio = max(max_step_ratio, step_ratio) + + rows.append( + AdaptiveLogRow( + t=t, + dt=dt, + accepted=1, + err_est=err_est, + y_norm_ratio=y_ratio, + step_ratio=step_ratio, + ) + ) + dt = adapt_dt(dt, err_est, tol, p=p) + else: + reject += 1 + rows.append( + AdaptiveLogRow( + t=t, + dt=dt, + accepted=0, + err_est=err_est, + y_norm_ratio=safe_norm(y) / y0_norm, + step_ratio=step_ratio, + ) + ) + dt = adapt_dt(dt, err_est, tol, p=p) + if reject > 20000: + break + + final_err = safe_norm(y - y_exact_T) / exact_norm + final_y_ratio = safe_norm(y) / y0_norm + + os.makedirs(outdir, exist_ok=True) + out_path = os.path.join(outdir, f"adaptive_{method}_Bmult{int(Bmult)}_{ordering}.csv") + + summary = [ + f"{method} Bmult={Bmult} T={T} dt0={dt0} tol={tol} ordering={ordering}", + f"accept={accept} reject={reject}", + f"final_err={final_err:.16e}", + f"max_y_norm_ratio={max_y_ratio:.16e} final_y_norm_ratio={final_y_ratio:.16e}", + f"max_step_ratio={max_step_ratio:.16e}", + ] + write_adaptive_csv(out_path, rows, summary) + + return out_path, accept, reject, final_err, max_y_ratio, final_y_ratio, max_step_ratio + + +# ---------------------------- +# Embedded RK baseline +# ---------------------------- + +def get_builtin_heun_euler_pair(): + """ + Fallback explicit embedded pair. + """ + A = np.array([ + [0.0, 0.0], + [1.0, 0.0], + ], dtype=float) + b_high = np.array([0.5, 0.5], dtype=float) + b_low = np.array([1.0, 0.0], dtype=float) + c = np.array([0.0, 1.0], dtype=float) + order_high = 2 + return A, b_high, b_low, c, order_high + + +def get_embedded_pair_from_pythos(pair_name: str): + """ + Get an embedded pair from butcher_tableau.py. + + We also import additive_rk.py first to verify the repo wiring is correct. + """ + import additive_rk + from butcher_tableau import embedded_pairs + + print("[ARK] Imported additive_rk from:", getattr(additive_rk, "__file__", "unknown")) + + if pair_name not in embedded_pairs: + available = sorted(list(embedded_pairs.keys())) + raise ValueError(f"Embedded pair '{pair_name}' not found. Available: {available}") + + pair = embedded_pairs[pair_name] + + A_rk = np.array(pair._a, dtype=float) + b_hi = np.array(pair._b, dtype=float) + b_lo = np.array(pair.b_aux, dtype=float) + c_rk = np.array(pair._c, dtype=float) + order_high = int(pair.order) + + print(f"[ARK] Using embedded pair from butcher_tableau: {pair_name}") + return A_rk, b_hi, b_lo, c_rk, order_high + + +def run_ark_baseline( + Bmult: float, + y0: np.ndarray, + T: float, + dt0: float, + tol: float, + outdir: str, + pair_name: str = "Dormand-Prince", +) -> Optional[str]: + """ + Embedded RK baseline on the full RHS: + y' = (A + B + C) y + + Uses pythOS butcher_tableau embedded pairs if available. + Falls back to builtin Heun-Euler only if something fails. + """ + A, B, C = make_abc_matrices(Bmult) + + try: + A_rk, b_hi, b_lo, c_rk, order_high = get_embedded_pair_from_pythos(pair_name) + except Exception as e: + print(f"[ARK] Could not use pythOS embedded pair: {e}") + print("[ARK] Falling back to Heun-Euler (builtin).") + A_rk, b_hi, b_lo, c_rk, order_high = get_builtin_heun_euler_pair() + + s = len(c_rk) + + def f(t: float, y: np.ndarray) -> np.ndarray: + return (A + B + C) @ y + + def rk_step_embedded(t: float, y: np.ndarray, h: float) -> Tuple[np.ndarray, float, float]: + K = np.zeros((s, len(y)), dtype=float) + + for i in range(s): + yi = y.copy() + for j in range(i): + aij = A_rk[i, j] + if aij != 0.0: + yi = yi + h * aij * K[j] + K[i] = f(t + c_rk[i] * h, yi) + + y_hi = y + h * np.sum(b_hi[:, None] * K, axis=0) + y_lo = y + h * np.sum(b_lo[:, None] * K, axis=0) + err = safe_norm(y_hi - y_lo) + step_ratio = safe_norm(y_hi) / max(safe_norm(y), 1e-300) + return y_hi, err, float(step_ratio) + + os.makedirs(outdir, exist_ok=True) + out_path = os.path.join(outdir, f"ark_Bmult{int(Bmult)}.csv") + + y_exact_T = exact_solution(A, B, C, y0, T) + exact_norm = max(safe_norm(y_exact_T), 1e-300) + y0_norm = max(safe_norm(y0), 1e-300) + + t = 0.0 + dt = float(dt0) + y = y0.copy() + + accept = 0 + reject = 0 + max_y_ratio = 1.0 + max_step_ratio = 1.0 + rows: List[AdaptiveLogRow] = [] + + while t < T - 1e-15: + dt = clamp(dt, 1e-12, 0.5) + if t + dt > T: + dt = T - t + + y_hi, err_est, step_ratio = rk_step_embedded(t, y, dt) + + if err_est <= tol: + t += dt + y = y_hi + accept += 1 + + y_ratio = safe_norm(y) / y0_norm + max_y_ratio = max(max_y_ratio, y_ratio) + max_step_ratio = max(max_step_ratio, step_ratio) + + rows.append( + AdaptiveLogRow( + t=t, + dt=dt, + accepted=1, + err_est=err_est, + y_norm_ratio=y_ratio, + step_ratio=step_ratio, + ) + ) + dt = adapt_dt(dt, err_est, tol, p=order_high) + else: + reject += 1 + rows.append( + AdaptiveLogRow( + t=t, + dt=dt, + accepted=0, + err_est=err_est, + y_norm_ratio=safe_norm(y) / y0_norm, + step_ratio=step_ratio, + ) + ) + dt = adapt_dt(dt, err_est, tol, p=order_high) + if reject > 20000: + break + + final_err = safe_norm(y - y_exact_T) / exact_norm + final_y_ratio = safe_norm(y) / y0_norm + + summary = [ + f"ARK embedded baseline Bmult={Bmult} T={T} dt0={dt0} tol={tol}", + f"pair_name={pair_name}", + f"accept={accept} reject={reject}", + f"final_err={final_err:.16e}", + f"max_y_norm_ratio={max_y_ratio:.16e} final_y_norm_ratio={final_y_ratio:.16e}", + f"max_step_ratio={max_step_ratio:.16e}", + ] + write_adaptive_csv(out_path, rows, summary) + + print(f"[ARK] saved: {out_path}") + return out_path + + +# ---------------------------- +# Plotting +# ---------------------------- + +def load_dt_series(csv_path: str) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + ts: List[float] = [] + dts: List[float] = [] + acc: List[int] = [] + + with open(csv_path, "r") as f: + reader = csv.reader(line for line in f if not line.lstrip().startswith("#")) + header = next(reader, None) + if header is None: + return np.array([]), np.array([]), np.array([]) + + for row in reader: + if len(row) < 3: + continue + try: + t = float(row[0]) + dt = float(row[1]) + accepted = int(row[2]) + except Exception: + continue + ts.append(t) + dts.append(dt) + acc.append(accepted) + + return np.array(ts), np.array(dts), np.array(acc) + + +def make_dt_plot(csv_path: str, title: str, out_png: str) -> None: + import matplotlib.pyplot as plt + + t, dt, accepted = load_dt_series(csv_path) + if len(t) == 0: + print(f"[plots] no data in {csv_path}") + return + + mask = accepted == 1 + tA = t[mask] + dtA = dt[mask] + + plt.figure() + plt.plot(tA, dtA, marker="o", linewidth=1.5) + plt.yscale("log") + plt.xlabel("t") + plt.ylabel("dt") + plt.title(title) + plt.tight_layout() + plt.savefig(out_png, dpi=150) + plt.close() + + +# ---------------------------- +# CLI +# ---------------------------- + +def parse_list_floats(s: str) -> List[float]: + out: List[float] = [] + for part in s.split(","): + part = part.strip() + if part: + out.append(float(part)) + return out + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--Bmults", type=str, default="1,20,50,100", help="comma-separated list") + ap.add_argument("--T", type=float, default=0.5) + ap.add_argument("--dt0", type=float, default=1e-2) + ap.add_argument("--tol", type=float, default=1e-7) + ap.add_argument("--ordering", type=str, default="CAB") + ap.add_argument("--outdir", type=str, default="experiments/abc_adaptive_outputs") + ap.add_argument("--do_ark", action="store_true") + ap.add_argument("--ark_pair", type=str, default="Dormand-Prince") + ap.add_argument("--plots", action="store_true") + args = ap.parse_args() + + Bmults = parse_list_floats(args.Bmults) + + print("=" * 72) + print( + f"Adaptive study: Bmults={Bmults} T={args.T} dt0={args.dt0} " + f"tol={args.tol} ordering={args.ordering}" + ) + print("=" * 72) + + os.makedirs(args.outdir, exist_ok=True) + plot_dir = os.path.join(args.outdir, "plots") + if args.plots: + os.makedirs(plot_dir, exist_ok=True) + + y0 = np.array([1.0, 0.2, -0.1], dtype=float) + + for Bmult in Bmults: + print("\n" + "=" * 72) + print( + f"Adaptive study: Bmult={int(Bmult)} T={args.T} dt0={args.dt0} " + f"tol={args.tol} ordering={args.ordering}" + ) + print("=" * 72) + + path_strang, aS, rS, eS, maxyS, finalyS, maxstepS = run_adaptive_splitting( + method="Strang3", + Bmult=Bmult, + y0=y0, + T=args.T, + dt0=args.dt0, + tol=args.tol, + ordering=args.ordering, + outdir=args.outdir, + ) + print( + f"[Strang3] accept={aS} reject={rS} final_err={eS:.6e} " + f"max||y||/||y0||={maxyS:.6e} final||y||/||y0||={finalyS:.6e} " + f"max_step={maxstepS:.6e}" + ) + print(f" saved: {path_strang}") + + path_yosh, aY, rY, eY, maxyY, finalyY, maxstepY = run_adaptive_splitting( + method="Yoshida3", + Bmult=Bmult, + y0=y0, + T=args.T, + dt0=args.dt0, + tol=args.tol, + ordering=args.ordering, + outdir=args.outdir, + ) + print( + f"[Yoshida3] accept={aY} reject={rY} final_err={eY:.6e} " + f"max||y||/||y0||={maxyY:.6e} final||y||/||y0||={finalyY:.6e} " + f"max_step={maxstepY:.6e}" + ) + print(f" saved: {path_yosh}") + + ark_path = None + if args.do_ark: + ark_path = run_ark_baseline( + Bmult=Bmult, + y0=y0, + T=args.T, + dt0=args.dt0, + tol=args.tol, + outdir=args.outdir, + pair_name=args.ark_pair, + ) + + if args.plots: + make_dt_plot( + path_strang, + f"Adaptive dt vs t | Strang-3 | Bmult={int(Bmult)} | {args.ordering}", + os.path.join(plot_dir, f"dt_vs_t_Strang3_Bmult{int(Bmult)}_{args.ordering}.png"), + ) + make_dt_plot( + path_yosh, + f"Adaptive dt vs t | Yoshida-3 | Bmult={int(Bmult)} | {args.ordering}", + os.path.join(plot_dir, f"dt_vs_t_Yoshida3_Bmult{int(Bmult)}_{args.ordering}.png"), + ) + if ark_path is not None: + make_dt_plot( + ark_path, + f"Adaptive dt vs t | ARK ({args.ark_pair}) | Bmult={int(Bmult)}", + os.path.join(plot_dir, f"dt_vs_t_ARK_Bmult{int(Bmult)}.png"), + ) + + if args.plots: + print(f"Plots saved to: {plot_dir}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/experiments/abc_adaptive_higher_order_splitting.py b/experiments/abc_adaptive_higher_order_splitting.py new file mode 100644 index 0000000..83b3c84 --- /dev/null +++ b/experiments/abc_adaptive_higher_order_splitting.py @@ -0,0 +1,486 @@ +""" +Adaptive and fixed-step higher-order splitting study for a three-operator ABC test problem. + +The script supports adaptive integration through step doubling, fixed-step +comparisons for the same methods, and CSV summaries suitable for plots or +tables. +""" + +from __future__ import annotations + +import argparse +import csv +import math +import os +import time +from dataclasses import dataclass +from typing import Callable, Dict, List, Tuple, Any + +import numpy as np + +try: + from experiments.utils.adaptive_splitting import ( + AdaptiveOptions, + adaptive_integrate, + l2_norm, + ) +except ModuleNotFoundError: + import sys + + THIS_DIR = os.path.dirname(os.path.abspath(__file__)) + REPO_ROOT = os.path.dirname(THIS_DIR) + if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + + from experiments.utils.adaptive_splitting import ( + AdaptiveOptions, + adaptive_integrate, + l2_norm, + ) + + +# ============================================================================= +# Problem definition +# ============================================================================= + +@dataclass +class ABCProblem: + Bmult: float + + def full_rhs(self, y: np.ndarray, t: float) -> np.ndarray: + y1, y2, y3 = y + + Ay = np.array([ + y2, + -y1, + 0.0, + ]) + + lam = self.Bmult + By = np.array([ + 0.0, + -lam * y2, + -0.5 * lam * y3, + ]) + + Cy = np.array([ + -0.2 * y1 * y3, + 0.1 * y1 * y1, + 0.3 * y1 - 0.1 * y3, + ]) + + return Ay + By + Cy + + def phi_A(self, y: np.ndarray, dt: float) -> np.ndarray: + y1, y2, y3 = y + c = math.cos(dt) + s = math.sin(dt) + return np.array([ + c * y1 + s * y2, + -s * y1 + c * y2, + y3, + ], dtype=float) + + def phi_B(self, y: np.ndarray, dt: float) -> np.ndarray: + y1, y2, y3 = y + lam = self.Bmult + return np.array([ + y1, + math.exp(-lam * dt) * y2, + math.exp(-0.5 * lam * dt) * y3, + ], dtype=float) + + def phi_C(self, y: np.ndarray, dt: float) -> np.ndarray: + def rhs_c(z: np.ndarray) -> np.ndarray: + z1, _, z3 = z + return np.array([ + -0.2 * z1 * z3, + 0.1 * z1 * z1, + 0.3 * z1 - 0.1 * z3, + ], dtype=float) + + z = y.copy() + k1 = rhs_c(z) + k2 = rhs_c(z + 0.5 * dt * k1) + k3 = rhs_c(z + 0.5 * dt * k2) + k4 = rhs_c(z + dt * k3) + return z + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4) + + +def make_abc_problem(Bmult: float) -> Tuple[ABCProblem, np.ndarray]: + problem = ABCProblem(Bmult=Bmult) + y0 = np.array([1.0, -0.25, 0.5], dtype=float) + return problem, y0 + + +# ============================================================================= +# Splitting methods +# ============================================================================= + +def strang_abc_step(problem: ABCProblem, y: np.ndarray, t: float, dt: float) -> np.ndarray: + z = problem.phi_A(y, 0.5 * dt) + z = problem.phi_B(z, 0.5 * dt) + z = problem.phi_C(z, dt) + z = problem.phi_B(z, 0.5 * dt) + z = problem.phi_A(z, 0.5 * dt) + return z + + +def yoshida4_from_strang( + base2: Callable[[np.ndarray, float, float], np.ndarray], +) -> Callable[[np.ndarray, float, float], np.ndarray]: + cbrt2 = 2.0 ** (1.0 / 3.0) + w1 = 1.0 / (2.0 - cbrt2) + w0 = -cbrt2 / (2.0 - cbrt2) + + def step(y: np.ndarray, t: float, dt: float) -> np.ndarray: + z = base2(y, t, w1 * dt) + z = base2(z, t + w1 * dt, w0 * dt) + z = base2(z, t + (w1 + w0) * dt, w1 * dt) + return z + + return step + + +def make_base_stepper(problem: ABCProblem, method: str) -> Tuple[Callable[[np.ndarray, float, float], np.ndarray], int]: + method = method.lower() + + if method == "strang2": + return (lambda y, t, dt: strang_abc_step(problem, y, t, dt), 2) + + if method == "yoshida4": + base2 = lambda y, t, dt: strang_abc_step(problem, y, t, dt) + return yoshida4_from_strang(base2), 4 + + raise ValueError(f"Unknown method '{method}'. Supported: strang2, yoshida4") + + +# ============================================================================= +# Reference solver +# ============================================================================= + +def rk4_step(rhs: Callable[[np.ndarray, float], np.ndarray], y: np.ndarray, t: float, dt: float) -> np.ndarray: + k1 = rhs(y, t) + k2 = rhs(y + 0.5 * dt * k1, t + 0.5 * dt) + k3 = rhs(y + 0.5 * dt * k2, t + 0.5 * dt) + k4 = rhs(y + dt * k3, t + dt) + return y + (dt / 6.0) * (k1 + 2*k2 + 2*k3 + k4) + + +def integrate_reference( + rhs: Callable[[np.ndarray, float], np.ndarray], + y0: np.ndarray, + T: float, + dt_ref: float, +) -> np.ndarray: + t = 0.0 + y = y0.copy() + while t < T - 1.0e-15: + dt = min(dt_ref, T - t) + y = rk4_step(rhs, y, t, dt) + t += dt + return y + + +# ============================================================================= +# Fixed-step solver +# ============================================================================= + +def integrate_fixed_step( + stepper: Callable[[np.ndarray, float, float], np.ndarray], + y0: np.ndarray, + T: float, + dt: float, +) -> Tuple[np.ndarray, int, float]: + t = 0.0 + y = y0.copy() + n_steps = 0 + tic = time.perf_counter() + + while t < T - 1.0e-15: + h = min(dt, T - t) + y = stepper(y, t, h) + t += h + n_steps += 1 + + wall = time.perf_counter() - tic + return y, n_steps, wall + + +# ============================================================================= +# IO helpers +# ============================================================================= + +def ensure_dir(path: str) -> None: + os.makedirs(path, exist_ok=True) + + +def save_adaptive_history_csv(path: str, result) -> None: + with open(path, "w", newline="") as f: + writer = csv.writer(f) + writer.writerow([ + "attempt_index", + "t", + "dt_try", + "dt_suggested", + "accepted", + "err_est", + "cpu_seconds", + ]) + for i, rec in enumerate(result.records): + writer.writerow([ + i, + rec.t, + rec.dt_try, + rec.dt_suggested, + int(rec.accepted), + rec.err_est, + rec.cpu_seconds, + ]) + + +def append_summary_csv(path: str, row: Dict[str, Any]) -> None: + file_exists = os.path.exists(path) + with open(path, "a", newline="") as f: + writer = csv.DictWriter(f, fieldnames=list(row.keys())) + if not file_exists: + writer.writeheader() + writer.writerow(row) + + +def parse_float_list(s: str) -> List[float]: + return [float(x.strip()) for x in s.split(",") if x.strip()] + + +# ============================================================================= +# Study runners +# ============================================================================= + +def run_one_adaptive_case( + Bmult: float, + method: str, + T: float, + dt0: float, + atol: float, + rtol: float, + dt_ref: float, + outdir: str, +) -> Dict[str, Any]: + problem, y0 = make_abc_problem(Bmult=Bmult) + stepper, order = make_base_stepper(problem, method=method) + + opts = AdaptiveOptions( + order=order, + atol=atol, + rtol=rtol, + dt_min=1.0e-12, + dt_max=max(dt0, T), + safety=0.9, + growth_max=2.0, + shrink_min=0.2, + max_reject=20, + use_weighted_rms=True, + ) + + result = adaptive_integrate( + stepper=stepper, + y0=y0, + t0=0.0, + tfinal=T, + dt0=dt0, + opts=opts, + ) + + y_ref = integrate_reference(problem.full_rhs, y0, T, dt_ref=dt_ref) + err_global = l2_norm(result.final_state - y_ref) + + history_name = ( + f"abc_history_mode-adaptive_method-{method}_B{Bmult:g}_" + f"atol{atol:.0e}_rtol{rtol:.0e}.csv" + ) + save_adaptive_history_csv(os.path.join(outdir, history_name), result) + + return { + "mode": "adaptive", + "method": method, + "order": order, + "Bmult": Bmult, + "T": T, + "dt0": dt0, + "fixed_dt": "", + "atol": atol, + "rtol": rtol, + "dt_ref": dt_ref, + "success": int(result.success), + "message": result.message, + "n_steps": result.n_accept, + "n_accept": result.n_accept, + "n_reject": result.n_reject, + "dt_min_used": result.dt_min_used, + "dt_max_used": result.dt_max_used, + "dt_avg_used": result.dt_avg_used, + "wall_seconds": result.wall_seconds, + "global_l2_error": err_global, + } + + +def run_one_fixed_case( + Bmult: float, + method: str, + T: float, + fixed_dt: float, + dt_ref: float, +) -> Dict[str, Any]: + problem, y0 = make_abc_problem(Bmult=Bmult) + stepper, order = make_base_stepper(problem, method=method) + + y_final, n_steps, wall = integrate_fixed_step( + stepper=stepper, + y0=y0, + T=T, + dt=fixed_dt, + ) + + y_ref = integrate_reference(problem.full_rhs, y0, T, dt_ref=dt_ref) + err_global = l2_norm(y_final - y_ref) + + return { + "mode": "fixed", + "method": method, + "order": order, + "Bmult": Bmult, + "T": T, + "dt0": "", + "fixed_dt": fixed_dt, + "atol": "", + "rtol": "", + "dt_ref": dt_ref, + "success": 1, + "message": "Fixed-step integration completed successfully.", + "n_steps": n_steps, + "n_accept": n_steps, + "n_reject": 0, + "dt_min_used": fixed_dt, + "dt_max_used": fixed_dt, + "dt_avg_used": fixed_dt, + "wall_seconds": wall, + "global_l2_error": err_global, + } + + +# ============================================================================= +# Main +# ============================================================================= + +def main() -> None: + parser = argparse.ArgumentParser(description="ABC higher-order splitting study: adaptive vs fixed.") + parser.add_argument("--mode", type=str, default="adaptive", + choices=["adaptive", "fixed", "both"], + help="Run adaptive, fixed, or both.") + parser.add_argument("--method", type=str, default="yoshida4", + choices=["strang2", "yoshida4"], + help="Base splitting method.") + parser.add_argument("--Bmults", type=str, default="100,200,500", + help="Comma-separated stiffness multipliers.") + parser.add_argument("--T", type=float, default=0.5, + help="Final time.") + parser.add_argument("--dt0", type=float, default=1.0e-2, + help="Initial adaptive step size guess.") + parser.add_argument("--atols", type=str, default="1e-6,1e-8", + help="Comma-separated absolute tolerances for adaptive mode.") + parser.add_argument("--rtols", type=str, default="1e-4,1e-6", + help="Comma-separated relative tolerances for adaptive mode.") + parser.add_argument("--fixed-dts", type=str, default="1e-2,5e-3,2e-3,1e-3,5e-4", + help="Comma-separated fixed time steps for fixed mode.") + parser.add_argument("--dt-ref", type=float, default=1.0e-6, + help="Reference RK4 step size.") + parser.add_argument("--outdir", type=str, default="experiments/outputs/abc_adaptive_higher_order", + help="Output directory.") + parser.add_argument("--overwrite-summary", action="store_true", + help="Overwrite summary CSV if it already exists.") + args = parser.parse_args() + + ensure_dir(args.outdir) + + Bmults = parse_float_list(args.Bmults) + atols = parse_float_list(args.atols) + rtols = parse_float_list(args.rtols) + fixed_dts = parse_float_list(args.fixed_dts) + + summary_path = os.path.join(args.outdir, "abc_higher_order_comparison_summary.csv") + if args.overwrite_summary and os.path.exists(summary_path): + os.remove(summary_path) + + print("=" * 72) + print("ABC HIGHER-ORDER SPLITTING COMPARISON STUDY") + print(f"mode = {args.mode}") + print(f"method = {args.method}") + print(f"Bmults = {Bmults}") + print(f"T = {args.T}") + print(f"dt0 = {args.dt0}") + print(f"atols = {atols}") + print(f"rtols = {rtols}") + print(f"fixed_dts = {fixed_dts}") + print(f"dt_ref = {args.dt_ref}") + print(f"outdir = {args.outdir}") + print("=" * 72) + + if args.mode in ("adaptive", "both"): + for Bmult in Bmults: + for atol in atols: + for rtol in rtols: + print( + f"\nAdaptive: method={args.method}, " + f"Bmult={Bmult:g}, atol={atol:.0e}, rtol={rtol:.0e}" + ) + + row = run_one_adaptive_case( + Bmult=Bmult, + method=args.method, + T=args.T, + dt0=args.dt0, + atol=atol, + rtol=rtol, + dt_ref=args.dt_ref, + outdir=args.outdir, + ) + append_summary_csv(summary_path, row) + + print( + f" success={row['success']} " + f"accept={row['n_accept']} reject={row['n_reject']} " + f"dt[min,avg,max]=({row['dt_min_used']:.3e}, " + f"{row['dt_avg_used']:.3e}, {row['dt_max_used']:.3e}) " + f"error={row['global_l2_error']:.6e} " + f"wall={row['wall_seconds']:.3f}s" + ) + + if args.mode in ("fixed", "both"): + for Bmult in Bmults: + for fixed_dt in fixed_dts: + print( + f"\nFixed: method={args.method}, " + f"Bmult={Bmult:g}, dt={fixed_dt:.3e}" + ) + + row = run_one_fixed_case( + Bmult=Bmult, + method=args.method, + T=args.T, + fixed_dt=fixed_dt, + dt_ref=args.dt_ref, + ) + append_summary_csv(summary_path, row) + + print( + f" steps={row['n_steps']} " + f"error={row['global_l2_error']:.6e} " + f"wall={row['wall_seconds']:.3f}s" + ) + + print("\nDone.") + print(f"Summary written to: {summary_path}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/experiments/abc_higher_order_stability_scan.py b/experiments/abc_higher_order_stability_scan.py new file mode 100644 index 0000000..1fd1a9e --- /dev/null +++ b/experiments/abc_higher_order_stability_scan.py @@ -0,0 +1,547 @@ +""" +abc_higher_order_stability_scan.py + +Phase 1C stability-window scan for higher-order 3-operator splitting methods. + +Purpose +------- +For each method and each stiffness regime, scan over a list of dt values and +record whether the method remains stable under a chosen criterion. + +This is designed to answer: +- Which methods have the largest stable timestep window? +- Do higher-order methods lose robustness earlier than Strang-type methods? +- How does that depend on the problem mode (decay vs growth) and stiffness? + +Model +----- +We solve y' = (A + B + C) y on R^2 with analytic subflows for A, B, C, using +pythOS fractional_step(...) and built-in alpha tables. + +Modes +----- +1. decay: + stiff dissipative B +2. growth: + one expanding and one contracting stiff direction + +Outputs +------- +For each (mode, Bmult, method, dt), writes: +- stable flag +- final error (if stable) +- max ||y_n|| / ||y_0|| +- final ||y_T|| / ||y_0|| +- max step ratio ||y_{n+1}|| / ||y_n|| +- exploded_at_step +- largest stable dt per method/Bmult summary + +Examples +-------- +python experiments/abc_higher_order_stability_scan.py +python experiments/abc_higher_order_stability_scan.py --mode decay --Bmults 20,50,100 +python experiments/abc_higher_order_stability_scan.py --mode growth --Bmults 20,50 +python experiments/abc_higher_order_stability_scan.py --dtmax 2e-2 --levels 10 +""" + +from __future__ import annotations + +import argparse +import csv +import math +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Dict, List, Sequence, Tuple + +import numpy as np + +# --------------------------------------------------------------------- +# Repository-local imports +# --------------------------------------------------------------------- +THIS = Path(__file__).resolve() +REPO = THIS.parents[1] +if str(REPO) not in sys.path: + sys.path.insert(0, str(REPO)) + +from fractional_step import fractional_step # noqa: E402 + +try: + from scipy.linalg import expm # type: ignore +except Exception: + expm = None + + +# --------------------------------------------------------------------- +# Utilities +# --------------------------------------------------------------------- +def parse_csv_floats(s: str) -> List[float]: + return [float(x.strip()) for x in s.split(",") if x.strip()] + + +def parse_csv_strings(s: str) -> List[str]: + return [x.strip() for x in s.split(",") if x.strip()] + + +def norm2(v: np.ndarray) -> float: + return float(np.linalg.norm(v)) + + +def safe_ratio(a: float, b: float) -> float: + if b == 0.0: + return float("inf") if a != 0.0 else 1.0 + return a / b + + +def expm_fallback(M: np.ndarray) -> np.ndarray: + w, V = np.linalg.eig(M) + Vinv = np.linalg.inv(V) + E = V @ np.diag(np.exp(w)) @ Vinv + if np.max(np.abs(E.imag)) < 1e-12: + return E.real + return E + + +def mat_expm(M: np.ndarray) -> np.ndarray: + if expm is not None: + return expm(M) + return expm_fallback(M) + + +def geometric_dt_list(dtmin: float, dtmax: float, levels: int) -> List[float]: + if levels < 2: + return [dtmax] + vals = np.geomspace(dtmax, dtmin, num=levels) + return [float(v) for v in vals] + + +# --------------------------------------------------------------------- +# Problem definition +# --------------------------------------------------------------------- +def make_abc_matrices(Bmult: float, mode: str) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + A = np.array( + [ + [0.0, 2.0], + [-2.0, 0.0], + ], + dtype=float, + ) + + if mode == "decay": + B = np.array( + [ + [-1.0, 0.0], + [0.0, -0.05], + ], + dtype=float, + ) * Bmult + elif mode == "growth": + B = np.array( + [ + [1.0, 0.0], + [0.0, -0.10], + ], + dtype=float, + ) * Bmult + else: + raise ValueError("mode must be 'decay' or 'growth'") + + C = np.array( + [ + [0.0, 1.5], + [0.4, 0.0], + ], + dtype=float, + ) + + return A, B, C + + +def exact_solution(y0: np.ndarray, T: float, A: np.ndarray, B: np.ndarray, C: np.ndarray) -> np.ndarray: + return mat_expm(T * (A + B + C)) @ y0 + + +def make_linear_flow(L: np.ndarray) -> Callable[[float, float, np.ndarray], np.ndarray]: + def flow(_t: float, h: float, y: np.ndarray) -> np.ndarray: + return mat_expm(h * L) @ y + return flow + + +# --------------------------------------------------------------------- +# Diagnostics +# --------------------------------------------------------------------- +@dataclass +class RunMetrics: + err: float + stable: bool + max_norm_ratio: float + final_norm_ratio: float + max_step_ratio: float + final_norm: float + exploded_at_step: int + steps: int + dt_used: float + + +def advance_one_step( + y: np.ndarray, + t: float, + dt: float, + flows: List[Callable[[float, float, np.ndarray], np.ndarray]], + alpha_name: str, +) -> np.ndarray: + y1 = fractional_step( + functions=flows, + delta_t=dt, + initial_y=np.array(y, dtype=float), + initial_t=t, + final_t=t + dt, + alpha=alpha_name, + methods={(1,): "ANALYTIC", (2,): "ANALYTIC", (3,): "ANALYTIC"}, + ) + return np.array(y1, dtype=np.complex128 if np.iscomplexobj(y1) else float) + + +def simulate_method( + alpha_name: str, + dt: float, + T: float, + y0: np.ndarray, + flows: List[Callable[[float, float, np.ndarray], np.ndarray]], + y_exact: np.ndarray, + explosion_factor: float, +) -> RunMetrics: + steps = max(1, int(round(T / dt))) + dt = T / steps + + y = np.array(y0, dtype=float) + t = 0.0 + + n0 = norm2(y0) + max_norm_ratio = 1.0 + max_step_ratio = 1.0 + stable = True + exploded_at_step = -1 + + for n in range(steps): + y_new = advance_one_step(y, t, dt, flows, alpha_name) + + if not np.all(np.isfinite(y_new)): + stable = False + exploded_at_step = n + 1 + y = y_new + break + + ny = norm2(np.asarray(y, dtype=np.complex128)) + ny_new = norm2(np.asarray(y_new, dtype=np.complex128)) + + if ny > 0.0: + max_step_ratio = max(max_step_ratio, ny_new / ny) + else: + max_step_ratio = float("inf") + + if n0 > 0.0: + max_norm_ratio = max(max_norm_ratio, ny_new / n0) + + if ny_new > explosion_factor * max(1.0, n0): + stable = False + exploded_at_step = n + 1 + y = y_new + break + + y = np.array(y_new, copy=True) + t += dt + + final_norm = norm2(np.asarray(y, dtype=np.complex128)) + final_norm_ratio = safe_ratio(final_norm, n0) + + if stable and np.all(np.isfinite(y)): + err = norm2(np.asarray(y, dtype=np.complex128) - np.asarray(y_exact, dtype=np.complex128)) + else: + err = float("inf") + + return RunMetrics( + err=err, + stable=stable, + max_norm_ratio=max_norm_ratio, + final_norm_ratio=final_norm_ratio, + max_step_ratio=max_step_ratio, + final_norm=final_norm, + exploded_at_step=exploded_at_step, + steps=steps, + dt_used=dt, + ) + + +# --------------------------------------------------------------------- +# Scan runner +# --------------------------------------------------------------------- +def run_scan( + mode: str, + Bmults: Sequence[float], + dts: Sequence[float], + methods: Sequence[str], + T: float, + y0: np.ndarray, + outdir: Path, + explosion_factor: float, +) -> None: + outdir.mkdir(parents=True, exist_ok=True) + + detail_rows: List[Dict[str, object]] = [] + summary_rows: List[Dict[str, object]] = [] + + for Bmult in Bmults: + A, B, C = make_abc_matrices(Bmult, mode) + flows = [make_linear_flow(A), make_linear_flow(B), make_linear_flow(C)] + y_exact = exact_solution(y0, T, A, B, C) + + print("\n" + "=" * 80) + print(f"STABILITY SCAN | mode = {mode} | Bmult = {Bmult:g} | T = {T:.6g}") + print("=" * 80) + + for method_name in methods: + print(f"\nMethod: {method_name}") + + largest_stable_dt = float("nan") + smallest_unstable_dt = float("nan") + stable_count = 0 + + for dt in dts: + metrics = simulate_method( + alpha_name=method_name, + dt=dt, + T=T, + y0=y0, + flows=flows, + y_exact=y_exact, + explosion_factor=explosion_factor, + ) + + print( + f" dt_req={dt:10.4e} | " + f"dt_used={metrics.dt_used:10.4e} | " + f"stable={metrics.stable!s:5s} | " + f"err={metrics.err:12.5e} | " + f"max||y||/||y0||={metrics.max_norm_ratio:10.4e} | " + f"max step ratio={metrics.max_step_ratio:10.4e}" + ) + + detail_rows.append( + { + "mode": mode, + "Bmult": Bmult, + "method": method_name, + "dt_requested": dt, + "dt_used": metrics.dt_used, + "steps": metrics.steps, + "T": T, + "y0_0": float(y0[0]), + "y0_1": float(y0[1]), + "stable": int(metrics.stable), + "err": metrics.err, + "max_norm_ratio": metrics.max_norm_ratio, + "final_norm_ratio": metrics.final_norm_ratio, + "max_step_ratio": metrics.max_step_ratio, + "final_norm": metrics.final_norm, + "exploded_at_step": metrics.exploded_at_step, + } + ) + + if metrics.stable: + stable_count += 1 + if np.isnan(largest_stable_dt) or metrics.dt_used > largest_stable_dt: + largest_stable_dt = metrics.dt_used + else: + if np.isnan(smallest_unstable_dt) or metrics.dt_used < smallest_unstable_dt: + smallest_unstable_dt = metrics.dt_used + + summary_rows.append( + { + "mode": mode, + "Bmult": Bmult, + "method": method_name, + "stable_count": stable_count, + "num_dts_tested": len(dts), + "largest_stable_dt": largest_stable_dt, + "smallest_unstable_dt": smallest_unstable_dt, + } + ) + + print("\nLargest stable dt by method:") + for row in [r for r in summary_rows if r["mode"] == mode and r["Bmult"] == Bmult]: + print( + f" {str(row['method']):20s} -> " + f"largest_stable_dt = {row['largest_stable_dt']!s}" + ) + + detail_path = outdir / f"abc_stability_scan_{mode}_detail.csv" + summary_path = outdir / f"abc_stability_scan_{mode}_summary.csv" + + with detail_path.open("w", newline="") as f: + writer = csv.DictWriter( + f, + fieldnames=[ + "mode", + "Bmult", + "method", + "dt_requested", + "dt_used", + "steps", + "T", + "y0_0", + "y0_1", + "stable", + "err", + "max_norm_ratio", + "final_norm_ratio", + "max_step_ratio", + "final_norm", + "exploded_at_step", + ], + ) + writer.writeheader() + writer.writerows(detail_rows) + + with summary_path.open("w", newline="") as f: + writer = csv.DictWriter( + f, + fieldnames=[ + "mode", + "Bmult", + "method", + "stable_count", + "num_dts_tested", + "largest_stable_dt", + "smallest_unstable_dt", + ], + ) + writer.writeheader() + writer.writerows(summary_rows) + + print(f"\nSaved detail CSV: {detail_path}") + print(f"Saved summary CSV: {summary_path}") + + +# --------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------- +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Stability-window scan for higher-order ABC splitting methods." + ) + + parser.add_argument( + "--mode", + type=str, + choices=["decay", "growth"], + default="growth", + help="Which regime to scan.", + ) + parser.add_argument( + "--Bmults", + type=str, + default="20,50,100", + help="Comma-separated stiffness multipliers.", + ) + parser.add_argument( + "--dts", + type=str, + default="", + help="Optional explicit comma-separated dt list. If omitted, use geometric grid from dtmax..dtmin.", + ) + parser.add_argument( + "--dtmax", + type=float, + default=2e-2, + help="Largest requested dt when --dts is not provided.", + ) + parser.add_argument( + "--dtmin", + type=float, + default=3.125e-4, + help="Smallest requested dt when --dts is not provided.", + ) + parser.add_argument( + "--levels", + type=int, + default=7, + help="Number of geometric dt levels when --dts is not provided.", + ) + parser.add_argument( + "--methods", + type=str, + default="Strang-3,OS32_7op_minLEM-3,PP3_4A-3,Yoshida-3", + help="Comma-separated pythOS alpha names to compare.", + ) + parser.add_argument( + "--T", + type=float, + default=0.5, + help="Final time.", + ) + parser.add_argument( + "--y0", + type=str, + default="1.0,-0.5", + help="Initial condition as 'y0,y1'.", + ) + parser.add_argument( + "--explosion_factor", + type=float, + default=1e10, + help="Declare unstable if ||y|| exceeds explosion_factor * max(1, ||y0||).", + ) + parser.add_argument( + "--outdir", + type=str, + default="experiments/abc_higher_order_stability_outputs", + help="Directory for CSV outputs.", + ) + + return parser + + +def main() -> None: + parser = build_parser() + args = parser.parse_args() + + if args.dts.strip(): + dts = parse_csv_floats(args.dts) + else: + dts = geometric_dt_list(args.dtmin, args.dtmax, args.levels) + + Bmults = parse_csv_floats(args.Bmults) + methods = parse_csv_strings(args.methods) + y0_vals = parse_csv_floats(args.y0) + + if len(y0_vals) != 2: + raise ValueError("--y0 must have exactly two components, e.g. --y0 1.0,-0.5") + + y0 = np.array(y0_vals, dtype=float) + outdir = REPO / args.outdir + + print("Running higher-order stability scan with:") + print(f" mode = {args.mode}") + print(f" Bmults = {Bmults}") + print(f" dts = {dts}") + print(f" methods = {methods}") + print(f" T = {args.T}") + print(f" y0 = {y0}") + print(f" explosion_factor = {args.explosion_factor}") + print(f" outdir = {outdir}") + + run_scan( + mode=args.mode, + Bmults=Bmults, + dts=dts, + methods=methods, + T=args.T, + y0=y0, + outdir=outdir, + explosion_factor=args.explosion_factor, + ) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/experiments/abc_higher_order_stiff_study.py b/experiments/abc_higher_order_stiff_study.py new file mode 100644 index 0000000..a95c4c0 --- /dev/null +++ b/experiments/abc_higher_order_stiff_study.py @@ -0,0 +1,499 @@ +""" +abc_higher_order_stiff_study.py + +A tougher Phase 1B three-operator splitting study. + +Purpose +------- +This script is designed to stress higher-order 3-operator splitting methods, +especially methods with negative coefficients / backward substeps such as +Yoshida-3 and PP3_4A-3. + +Compared with abc_higher_order_study.py, this script: +- uses a much harsher dissipative-stiff operator B +- includes a "growth" mode where backward substeps in B can amplify strongly +- tracks per-step and global growth diagnostics more carefully +- allows longer final times and larger stiffness multipliers + +Model +----- +We solve y' = (A + B + C) y on R^2 with analytic subflows for A, B, C. + +Modes: +1. decay + B is strongly dissipative; backward B-substeps can become dangerous because + exp(h B) is very contracting for h>0 but can amplify for h<0. + +2. growth + B has one strongly positive eigenvalue and one negative one; this is not + "physical dissipation", but it is useful as a stress test to reveal how + compositions behave when backward/forward substeps interact badly. + +Outputs +------- +For each (Bmult, mode, method, dt), computes: +- final error vs exact expm((A+B+C)T)y0 +- observed order from dt sweep +- max ||y_n|| / ||y_0|| +- final ||y_T|| / ||y_0|| +- max step ratio ||y_{n+1}|| / ||y_n|| +- unstable flag if nonfinite or explosive growth occurs + +Example runs +------------ +python experiments/abc_higher_order_stiff_study.py +python experiments/abc_higher_order_stiff_study.py --mode decay --Bmults 20,50,100 --T 1.0 +python experiments/abc_higher_order_stiff_study.py --mode growth --Bmults 20,50 --T 0.5 +python experiments/abc_higher_order_stiff_study.py --methods Strang-3,OS32_7op_minLEM-3,PP3_4A-3,Yoshida-3 +""" + +from __future__ import annotations + +import argparse +import csv +import math +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Dict, List, Sequence, Tuple + +import numpy as np + +# --------------------------------------------------------------------- +# Repository-local imports +# --------------------------------------------------------------------- +THIS = Path(__file__).resolve() +REPO = THIS.parents[1] +if str(REPO) not in sys.path: + sys.path.insert(0, str(REPO)) + +from fractional_step import fractional_step # noqa: E402 + +try: + from scipy.linalg import expm # type: ignore +except Exception: + expm = None + + +# --------------------------------------------------------------------- +# Utilities +# --------------------------------------------------------------------- +def parse_csv_floats(s: str) -> List[float]: + return [float(x.strip()) for x in s.split(",") if x.strip()] + + +def parse_csv_strings(s: str) -> List[str]: + return [x.strip() for x in s.split(",") if x.strip()] + + +def norm2(v: np.ndarray) -> float: + return float(np.linalg.norm(v)) + + +def safe_ratio(a: float, b: float) -> float: + if b == 0.0: + return float("inf") if a != 0.0 else 1.0 + return a / b + + +def expm_fallback(M: np.ndarray) -> np.ndarray: + w, V = np.linalg.eig(M) + Vinv = np.linalg.inv(V) + E = V @ np.diag(np.exp(w)) @ Vinv + if np.max(np.abs(E.imag)) < 1e-12: + return E.real + return E + + +def mat_expm(M: np.ndarray) -> np.ndarray: + if expm is not None: + return expm(M) + return expm_fallback(M) + + +def estimate_order(errs: Sequence[float], dts: Sequence[float]) -> float: + xs: List[float] = [] + ys: List[float] = [] + for e, dt in zip(errs, dts): + if np.isfinite(e) and e > 0.0 and dt > 0.0: + xs.append(math.log(dt)) + ys.append(math.log(e)) + + if len(xs) < 2: + return float("nan") + + x = np.array(xs, dtype=float) + y = np.array(ys, dtype=float) + A = np.vstack([x, np.ones_like(x)]).T + p, _ = np.linalg.lstsq(A, y, rcond=None)[0] + return float(p) + + +# --------------------------------------------------------------------- +# Problem definition +# --------------------------------------------------------------------- +def make_abc_matrices(Bmult: float, mode: str) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Build a tougher noncommuting ABC test. + + A: rotation/shear mix + B: stiff diagonal operator whose behavior depends on mode + C: additional noncommuting shear/coupling + """ + A = np.array( + [ + [0.0, 2.0], + [-2.0, 0.0], + ], + dtype=float, + ) + + if mode == "decay": + B = np.array( + [ + [-1.0, 0.0], + [0.0, -0.05], + ], + dtype=float, + ) * Bmult + elif mode == "growth": + B = np.array( + [ + [1.0, 0.0], + [0.0, -0.10], + ], + dtype=float, + ) * Bmult + else: + raise ValueError("mode must be 'decay' or 'growth'") + + C = np.array( + [ + [0.0, 1.5], + [0.4, 0.0], + ], + dtype=float, + ) + + return A, B, C + + +def exact_solution(y0: np.ndarray, T: float, A: np.ndarray, B: np.ndarray, C: np.ndarray) -> np.ndarray: + return mat_expm(T * (A + B + C)) @ y0 + + +def make_linear_flow(L: np.ndarray) -> Callable[[float, float, np.ndarray], np.ndarray]: + def flow(_t: float, h: float, y: np.ndarray) -> np.ndarray: + return mat_expm(h * L) @ y + return flow + + +# --------------------------------------------------------------------- +# Diagnostics +# --------------------------------------------------------------------- +@dataclass +class RunMetrics: + err: float + stable: bool + max_norm_ratio: float + final_norm_ratio: float + max_step_ratio: float + final_norm: float + exploded_at_step: int + + +def advance_one_step( + y: np.ndarray, + t: float, + dt: float, + flows: List[Callable[[float, float, np.ndarray], np.ndarray]], + alpha_name: str, +) -> np.ndarray: + y1 = fractional_step( + functions=flows, + delta_t=dt, + initial_y=np.array(y, dtype=float), + initial_t=t, + final_t=t + dt, + alpha=alpha_name, + methods={(1,): "ANALYTIC", (2,): "ANALYTIC", (3,): "ANALYTIC"}, + ) + return np.array(y1, dtype=np.complex128 if np.iscomplexobj(y1) else float) + + +def simulate_method( + alpha_name: str, + dt: float, + T: float, + y0: np.ndarray, + flows: List[Callable[[float, float, np.ndarray], np.ndarray]], + y_exact: np.ndarray, + explosion_factor: float, +) -> RunMetrics: + steps = max(1, int(round(T / dt))) + dt = T / steps + + y = np.array(y0, dtype=float) + t = 0.0 + + n0 = norm2(y0) + max_norm_ratio = 1.0 + max_step_ratio = 1.0 + stable = True + exploded_at_step = -1 + + for n in range(steps): + y_new = advance_one_step(y, t, dt, flows, alpha_name) + + if not np.all(np.isfinite(y_new)): + stable = False + exploded_at_step = n + 1 + y = y_new + break + + ny = norm2(np.asarray(y, dtype=np.complex128)) + ny_new = norm2(np.asarray(y_new, dtype=np.complex128)) + + if ny > 0.0: + max_step_ratio = max(max_step_ratio, ny_new / ny) + else: + max_step_ratio = float("inf") + + if n0 > 0.0: + max_norm_ratio = max(max_norm_ratio, ny_new / n0) + + if ny_new > explosion_factor * max(1.0, n0): + stable = False + exploded_at_step = n + 1 + y = y_new + break + + y = np.array(y_new, copy=True) + t += dt + + final_norm = norm2(np.asarray(y, dtype=np.complex128)) + final_norm_ratio = safe_ratio(final_norm, n0) + + if stable and np.all(np.isfinite(y)): + err = norm2(np.asarray(y, dtype=np.complex128) - np.asarray(y_exact, dtype=np.complex128)) + else: + err = float("inf") + + return RunMetrics( + err=err, + stable=stable, + max_norm_ratio=max_norm_ratio, + final_norm_ratio=final_norm_ratio, + max_step_ratio=max_step_ratio, + final_norm=final_norm, + exploded_at_step=exploded_at_step, + ) + + +# --------------------------------------------------------------------- +# Study runner +# --------------------------------------------------------------------- +def run_study( + mode: str, + Bmults: Sequence[float], + dts: Sequence[float], + methods: Sequence[str], + T: float, + y0: np.ndarray, + outdir: Path, + explosion_factor: float, +) -> None: + outdir.mkdir(parents=True, exist_ok=True) + + for Bmult in Bmults: + A, B, C = make_abc_matrices(Bmult, mode) + flows = [make_linear_flow(A), make_linear_flow(B), make_linear_flow(C)] + y_exact = exact_solution(y0, T, A, B, C) + + rows: List[Dict[str, object]] = [] + err_by_method: Dict[str, List[float]] = {m: [] for m in methods} + + print("\n" + "=" * 78) + print(f"STIFF ABC STUDY | mode = {mode} | Bmult = {Bmult:g} | T = {T:.6g}") + print("=" * 78) + + for method_name in methods: + print(f"\nMethod: {method_name}") + + for dt in dts: + metrics = simulate_method( + alpha_name=method_name, + dt=dt, + T=T, + y0=y0, + flows=flows, + y_exact=y_exact, + explosion_factor=explosion_factor, + ) + err_by_method[method_name].append(metrics.err) + + print( + f" dt={dt:10.4e} | " + f"err={metrics.err:12.5e} | " + f"stable={metrics.stable!s:5s} | " + f"max||y||/||y0||={metrics.max_norm_ratio:10.4e} | " + f"final||y||/||y0||={metrics.final_norm_ratio:10.4e} | " + f"max step ratio={metrics.max_step_ratio:10.4e}" + ) + + rows.append( + { + "mode": mode, + "Bmult": Bmult, + "method": method_name, + "dt": dt, + "T": T, + "y0_0": float(y0[0]), + "y0_1": float(y0[1]), + "err": metrics.err, + "stable": int(metrics.stable), + "max_norm_ratio": metrics.max_norm_ratio, + "final_norm_ratio": metrics.final_norm_ratio, + "max_step_ratio": metrics.max_step_ratio, + "final_norm": metrics.final_norm, + "exploded_at_step": metrics.exploded_at_step, + } + ) + + p_by_method: Dict[str, float] = { + m: estimate_order(err_by_method[m], dts) for m in methods + } + + print("\nObserved orders:") + for m in methods: + print(f" {m:20s} -> p ≈ {p_by_method[m]:.4f}") + + for row in rows: + row["observed_order_fit"] = p_by_method[str(row["method"])] + + csv_path = outdir / f"abc_stiff_{mode}_Bmult{str(Bmult).replace('.', 'p')}.csv" + with csv_path.open("w", newline="") as f: + writer = csv.DictWriter( + f, + fieldnames=[ + "mode", + "Bmult", + "method", + "dt", + "T", + "y0_0", + "y0_1", + "err", + "stable", + "max_norm_ratio", + "final_norm_ratio", + "max_step_ratio", + "final_norm", + "exploded_at_step", + "observed_order_fit", + ], + ) + writer.writeheader() + writer.writerows(rows) + + print(f"\nSaved: {csv_path}") + + +# --------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------- +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Tougher stiff ABC splitting study using pythOS fractional_step.") + + parser.add_argument( + "--mode", + type=str, + choices=["decay", "growth"], + default="decay", + help="Which stiff test to run.", + ) + parser.add_argument( + "--Bmults", + type=str, + default="20,50,100", + help="Comma-separated stiffness multipliers for B.", + ) + parser.add_argument( + "--dts", + type=str, + default="1e-2,5e-3,2.5e-3,1.25e-3", + help="Comma-separated time steps for the study.", + ) + parser.add_argument( + "--methods", + type=str, + default="Strang-3,OS32_7op_minLEM-3,PP3_4A-3,Yoshida-3", + help="Comma-separated pythOS alpha names to compare.", + ) + parser.add_argument( + "--T", + type=float, + default=1.0, + help="Final time.", + ) + parser.add_argument( + "--y0", + type=str, + default="1.0,-0.5", + help="Initial condition as 'y0,y1'.", + ) + parser.add_argument( + "--explosion_factor", + type=float, + default=1e10, + help="Declare unstable if ||y|| exceeds explosion_factor * max(1, ||y0||).", + ) + parser.add_argument( + "--outdir", + type=str, + default="experiments/abc_higher_order_stiff_outputs", + help="Directory for CSV files.", + ) + + return parser + + +def main() -> None: + parser = build_parser() + args = parser.parse_args() + + Bmults = parse_csv_floats(args.Bmults) + dts = parse_csv_floats(args.dts) + methods = parse_csv_strings(args.methods) + y0_vals = parse_csv_floats(args.y0) + + if len(y0_vals) != 2: + raise ValueError("--y0 must have exactly two components, e.g. --y0 1.0,-0.5") + + y0 = np.array(y0_vals, dtype=float) + outdir = REPO / args.outdir + + print("Running stiff ABC study with:") + print(f" mode = {args.mode}") + print(f" Bmults = {Bmults}") + print(f" dts = {dts}") + print(f" methods = {methods}") + print(f" T = {args.T}") + print(f" y0 = {y0}") + print(f" explosion_factor = {args.explosion_factor}") + print(f" outdir = {outdir}") + + run_study( + mode=args.mode, + Bmults=Bmults, + dts=dts, + methods=methods, + T=args.T, + y0=y0, + outdir=outdir, + explosion_factor=args.explosion_factor, + ) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/experiments/abc_higher_order_study.py b/experiments/abc_higher_order_study.py new file mode 100644 index 0000000..9414fcc --- /dev/null +++ b/experiments/abc_higher_order_study.py @@ -0,0 +1,454 @@ +""" +abc_higher_order_study.py + +Phase 1 higher-order 3-operator splitting study on a small linear ABC test. + +What it does +------------ +- Defines a linear 2x2 system y' = (A + B + C) y with tunable stiffness via Bmult. +- Uses pythOS fractional_step(...) with ANALYTIC subflows for A, B, C. +- Compares several built-in 3-operator splitting methods already available in + fractional_step.py, including higher-order candidates. +- For each (method, dt), computes: + * final-time error vs exact expm solution + * observed order from a dt sweep + * simple stability / growth diagnostics: + - max ||y_n|| / ||y_0|| + - final ||y_T|| / ||y_0|| + - max step ratio ||y_{n+1}|| / ||y_n|| +- Saves one CSV per Bmult regime. + +Examples +-------- +python experiments/abc_higher_order_study.py +python experiments/abc_higher_order_study.py --Bmults 1,20,50,100 +python experiments/abc_higher_order_study.py --dts 1e-2,5e-3,2.5e-3,1.25e-3 +python experiments/abc_higher_order_study.py --methods Godunov-3,Strang-3,Yoshida-3,PP3_4A-3 +""" + +from __future__ import annotations + +import argparse +import csv +import math +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Dict, List, Sequence, Tuple + +import numpy as np + +# --------------------------------------------------------------------- +# Repository-local imports +# --------------------------------------------------------------------- +THIS = Path(__file__).resolve() +REPO = THIS.parents[1] +if str(REPO) not in sys.path: + sys.path.insert(0, str(REPO)) + +from fractional_step import fractional_step # noqa: E402 + +try: + from scipy.linalg import expm # type: ignore +except Exception: + expm = None + + +# --------------------------------------------------------------------- +# Utilities +# --------------------------------------------------------------------- +def parse_csv_floats(s: str) -> List[float]: + return [float(x.strip()) for x in s.split(",") if x.strip()] + + +def parse_csv_strings(s: str) -> List[str]: + return [x.strip() for x in s.split(",") if x.strip()] + + +def norm2(v: np.ndarray) -> float: + return float(np.linalg.norm(v)) + + +def safe_ratio(a: float, b: float) -> float: + if b == 0.0: + return float("inf") if a != 0.0 else 1.0 + return a / b + + +def expm_fallback(M: np.ndarray) -> np.ndarray: + """ + Small-matrix expm fallback using eigendecomposition. + Fine for this 2x2 experiment. + """ + w, V = np.linalg.eig(M) + Vinv = np.linalg.inv(V) + E = V @ np.diag(np.exp(w)) @ Vinv + if np.max(np.abs(E.imag)) < 1e-12: + return E.real + return E + + +def mat_expm(M: np.ndarray) -> np.ndarray: + if expm is not None: + return expm(M) + return expm_fallback(M) + + +def estimate_order(errs: Sequence[float], dts: Sequence[float]) -> float: + """ + Least-squares fit to log(err) = p log(dt) + c, ignoring non-finite values. + """ + xs: List[float] = [] + ys: List[float] = [] + for e, dt in zip(errs, dts): + if np.isfinite(e) and e > 0.0 and dt > 0.0: + xs.append(math.log(dt)) + ys.append(math.log(e)) + + if len(xs) < 2: + return float("nan") + + x = np.array(xs, dtype=float) + y = np.array(ys, dtype=float) + A = np.vstack([x, np.ones_like(x)]).T + p, _ = np.linalg.lstsq(A, y, rcond=None)[0] + return float(p) + + +# --------------------------------------------------------------------- +# Problem definition +# --------------------------------------------------------------------- +def make_abc_matrices(Bmult: float) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + A, B, C are chosen so that: + - operators do not commute + - B controls stiffness / damping scale + """ + A = np.array( + [ + [0.0, 1.0], + [-1.0, 0.0], + ], + dtype=float, + ) # rotation-like + + B = np.array( + [ + [-1.0, 0.0], + [0.0, -0.2], + ], + dtype=float, + ) * Bmult # damping / stiffness + + C = np.array( + [ + [0.0, 0.6], + [0.0, 0.0], + ], + dtype=float, + ) # shear-like + + return A, B, C + + +def exact_solution(y0: np.ndarray, T: float, A: np.ndarray, B: np.ndarray, C: np.ndarray) -> np.ndarray: + return mat_expm(T * (A + B + C)) @ y0 + + +# --------------------------------------------------------------------- +# Analytic subflows for fractional_step(..., methods={(i,): "ANALYTIC"}) +# --------------------------------------------------------------------- +def make_linear_flow(L: np.ndarray) -> Callable[[float, float, np.ndarray], np.ndarray]: + """ + fractional_step's analytic callback expects signature: + flow(current_t, delta_t, y) + """ + def flow(_t: float, h: float, y: np.ndarray) -> np.ndarray: + return mat_expm(h * L) @ y + + return flow + + +# --------------------------------------------------------------------- +# Diagnostics +# --------------------------------------------------------------------- +@dataclass +class RunMetrics: + err: float + stable: bool + max_norm_ratio: float + final_norm_ratio: float + max_step_ratio: float + final_norm: float + + +def advance_one_step( + y: np.ndarray, + t: float, + dt: float, + flows: List[Callable[[float, float, np.ndarray], np.ndarray]], + alpha_name: str, +) -> np.ndarray: + """ + Advance exactly one global splitting step of size dt using fractional_step. + """ + y1 = fractional_step( + functions=flows, + delta_t=dt, + initial_y=np.array(y, dtype=float), + initial_t=t, + final_t=t + dt, + alpha=alpha_name, + methods={(1,): "ANALYTIC", (2,): "ANALYTIC", (3,): "ANALYTIC"}, + ) + return np.array(y1, dtype=np.complex128 if np.iscomplexobj(y1) else float) + + +def simulate_method( + alpha_name: str, + dt: float, + T: float, + y0: np.ndarray, + flows: List[Callable[[float, float, np.ndarray], np.ndarray]], + y_exact: np.ndarray, + explosion_factor: float = 1e12, +) -> RunMetrics: + """ + Repeatedly calls fractional_step one global step at a time so that we can + track growth diagnostics across the trajectory. + """ + steps = max(1, int(round(T / dt))) + dt = T / steps # force exact final time hit + + y = np.array(y0, dtype=float) + t = 0.0 + + n0 = norm2(y0) + max_norm_ratio = 1.0 + max_step_ratio = 1.0 + stable = True + + for _ in range(steps): + y_new = advance_one_step(y, t, dt, flows, alpha_name) + + if not np.all(np.isfinite(y_new)): + stable = False + y = y_new + break + + ny = norm2(np.asarray(y, dtype=np.complex128)) + ny_new = norm2(np.asarray(y_new, dtype=np.complex128)) + + if ny > 0.0: + max_step_ratio = max(max_step_ratio, ny_new / ny) + else: + max_step_ratio = float("inf") + + if n0 > 0.0: + max_norm_ratio = max(max_norm_ratio, ny_new / n0) + + if ny_new > explosion_factor * max(1.0, n0): + stable = False + y = y_new + break + + y = np.array(y_new, copy=True) + t += dt + + final_norm = norm2(np.asarray(y, dtype=np.complex128)) + final_norm_ratio = safe_ratio(final_norm, n0) + + if stable and np.all(np.isfinite(y)): + err = norm2(np.asarray(y, dtype=np.complex128) - np.asarray(y_exact, dtype=np.complex128)) + else: + err = float("inf") + + return RunMetrics( + err=err, + stable=stable, + max_norm_ratio=max_norm_ratio, + final_norm_ratio=final_norm_ratio, + max_step_ratio=max_step_ratio, + final_norm=final_norm, + ) + + +# --------------------------------------------------------------------- +# Main study +# --------------------------------------------------------------------- +def run_study( + Bmults: Sequence[float], + dts: Sequence[float], + methods: Sequence[str], + T: float, + y0: np.ndarray, + outdir: Path, +) -> None: + outdir.mkdir(parents=True, exist_ok=True) + + for Bmult in Bmults: + A, B, C = make_abc_matrices(Bmult) + flows = [make_linear_flow(A), make_linear_flow(B), make_linear_flow(C)] + y_exact = exact_solution(y0, T, A, B, C) + + rows: List[Dict[str, object]] = [] + err_by_method: Dict[str, List[float]] = {m: [] for m in methods} + + print("\n" + "=" * 72) + print(f"HIGHER-ORDER ABC STUDY | Bmult = {Bmult:g} | T = {T:.6g}") + print("=" * 72) + + for method_name in methods: + print(f"\nMethod: {method_name}") + + for dt in dts: + metrics = simulate_method( + alpha_name=method_name, + dt=dt, + T=T, + y0=y0, + flows=flows, + y_exact=y_exact, + ) + err_by_method[method_name].append(metrics.err) + + print( + f" dt={dt:10.4e} | " + f"err={metrics.err:12.5e} | " + f"stable={metrics.stable!s:5s} | " + f"max||y||/||y0||={metrics.max_norm_ratio:10.4e} | " + f"max step ratio={metrics.max_step_ratio:10.4e}" + ) + + rows.append( + { + "Bmult": Bmult, + "method": method_name, + "dt": dt, + "T": T, + "y0_0": float(y0[0]), + "y0_1": float(y0[1]), + "err": metrics.err, + "stable": int(metrics.stable), + "max_norm_ratio": metrics.max_norm_ratio, + "final_norm_ratio": metrics.final_norm_ratio, + "max_step_ratio": metrics.max_step_ratio, + "final_norm": metrics.final_norm, + } + ) + + p_by_method: Dict[str, float] = { + m: estimate_order(err_by_method[m], dts) for m in methods + } + + print("\nObserved orders:") + for m in methods: + print(f" {m:20s} -> p ≈ {p_by_method[m]:.4f}") + + for row in rows: + row["observed_order_fit"] = p_by_method[str(row["method"])] + + csv_path = outdir / f"abc_higher_order_Bmult{str(Bmult).replace('.', 'p')}.csv" + with csv_path.open("w", newline="") as f: + writer = csv.DictWriter( + f, + fieldnames=[ + "Bmult", + "method", + "dt", + "T", + "y0_0", + "y0_1", + "err", + "stable", + "max_norm_ratio", + "final_norm_ratio", + "max_step_ratio", + "final_norm", + "observed_order_fit", + ], + ) + writer.writeheader() + writer.writerows(rows) + + print(f"\nSaved: {csv_path}") + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Higher-order ABC splitting study using pythOS fractional_step.") + + parser.add_argument( + "--Bmults", + type=str, + default="1,20,50,100", + help="Comma-separated stiffness multipliers for B.", + ) + parser.add_argument( + "--dts", + type=str, + default="1e-2,5e-3,2.5e-3,1.25e-3", + help="Comma-separated time steps for the order study.", + ) + parser.add_argument( + "--methods", + type=str, + default="Godunov-3,Strang-3,StrangACB-3,StrangBAC-3,StrangBCA-3,StrangCAB-3,StrangCBA-3,OS32_7op_minLEM-3,Yoshida-3,PP3_4A-3", + help="Comma-separated pythOS alpha names to compare.", + ) + parser.add_argument( + "--T", + type=float, + default=0.5, + help="Final time.", + ) + parser.add_argument( + "--y0", + type=str, + default="1.0,-0.5", + help="Initial condition as 'y0,y1'.", + ) + parser.add_argument( + "--outdir", + type=str, + default="experiments/abc_higher_order_outputs", + help="Directory where CSV files will be saved.", + ) + + return parser + + +def main() -> None: + parser = build_parser() + args = parser.parse_args() + + Bmults = parse_csv_floats(args.Bmults) + dts = parse_csv_floats(args.dts) + methods = parse_csv_strings(args.methods) + y0_vals = parse_csv_floats(args.y0) + + if len(y0_vals) != 2: + raise ValueError("--y0 must have exactly two components, e.g. --y0 1.0,-0.5") + + y0 = np.array(y0_vals, dtype=float) + outdir = REPO / args.outdir + + print("Running higher-order ABC study with:") + print(f" Bmults = {Bmults}") + print(f" dts = {dts}") + print(f" methods = {methods}") + print(f" T = {args.T}") + print(f" y0 = {y0}") + print(f" outdir = {outdir}") + + run_study( + Bmults=Bmults, + dts=dts, + methods=methods, + T=args.T, + y0=y0, + outdir=outdir, + ) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/experiments/abc_make_plots.py b/experiments/abc_make_plots.py new file mode 100644 index 0000000..f1a2a19 --- /dev/null +++ b/experiments/abc_make_plots.py @@ -0,0 +1,288 @@ +#!/usr/bin/env python3 +""" +Make plots from experiments/abc_outputs/abc_study_Bmult*.csv + +Outputs (per Bmult): + - err_vs_order__.png (log-scale error by ordering) + - maxstep_vs_order__.png (max_step by ordering) + +No pandas required (uses csv + matplotlib). +""" + +from __future__ import annotations + +import argparse +import csv +import glob +import os +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import numpy as np +import matplotlib.pyplot as plt + + +# ----------------------------- +# Helpers to be robust to header names +# ----------------------------- +def _first_present(row: Dict[str, str], keys: List[str]) -> Optional[str]: + for k in keys: + if k in row and row[k] != "": + return row[k] + return None + + +def _as_float(x: Optional[str]) -> float: + if x is None: + return float("nan") + x = x.strip() + if x.lower() == "inf": + return float("inf") + if x.lower() == "nan": + return float("nan") + return float(x) + + +def _as_bool(x: Optional[str]) -> bool: + if x is None: + return False + s = x.strip().lower() + return s in ("1", "true", "t", "yes", "y", "ok", "stable") + + +@dataclass +class Row: + bmult: float + method: str + ordering: str + err_dt_min: float + p_obs: float + max_norm_ratio: float + final_norm_ratio: float + max_step: float + stable: bool + + +def read_csv_rows(path: str) -> List[Row]: + rows: List[Row] = [] + bmult_guess = _infer_bmult_from_filename(path) + + with open(path, "r", newline="") as f: + reader = csv.DictReader(f) + headers = reader.fieldnames or [] + # print("CSV headers:", headers) + + for r in reader: + method = _first_present(r, ["method", "Method"]) + ordering = _first_present(r, ["ordering", "Ordering", "perm", "Permutation"]) + if method is None or ordering is None: + # Skip malformed lines + continue + + # Error at smallest dt (dt_min) + err_dt_min = _as_float( + _first_present(r, ["err_dt_min", "err(dt_min)", "err_dtmin", "err_min", "err"]) + ) + # Observed order + p_obs = _as_float(_first_present(r, ["p", "p_obs", "observed_p", "p≈"])) + + # Norm ratios + max_norm_ratio = _as_float( + _first_present(r, ["max_norm_ratio", "max||y||/||y0||", "max_ratio"]) + ) + final_norm_ratio = _as_float( + _first_present(r, ["final_norm_ratio", "final||y||/||y0||", "final_ratio"]) + ) + + # New metric + max_step = _as_float(_first_present(r, ["max_step", "max_step_ratio"])) + + # Stability flag + stable_str = _first_present(r, ["stable", "ok", "status"]) + if stable_str is None: + # Sometimes there's an "UNSTABLE" status column + status = _first_present(r, ["Status", "status", "label"]) + if status is not None and status.strip().upper() == "UNSTABLE": + stable = False + else: + # fall back: unstable if err is inf or max_norm_ratio huge + stable = np.isfinite(err_dt_min) and (not (np.isfinite(max_norm_ratio) and max_norm_ratio > 1e3)) + else: + # "ok" is stable; "UNSTABLE" is not + s = stable_str.strip().lower() + if "unstable" in s: + stable = False + elif s in ("ok", "stable", "true", "1", "yes", "y"): + stable = True + else: + stable = _as_bool(stable_str) + + rows.append( + Row( + bmult=bmult_guess, + method=method.strip(), + ordering=ordering.strip(), + err_dt_min=err_dt_min, + p_obs=p_obs, + max_norm_ratio=max_norm_ratio, + final_norm_ratio=final_norm_ratio, + max_step=max_step, + stable=stable, + ) + ) + + return rows + + +def _infer_bmult_from_filename(path: str) -> float: + base = os.path.basename(path) + # expected like abc_study_Bmult50.csv + for token in base.replace(".csv", "").split("_"): + if token.lower().startswith("bmult"): + try: + return float(token[len("Bmult"):]) + except Exception: + pass + return float("nan") + + +def ensure_dir(d: str) -> None: + os.makedirs(d, exist_ok=True) + + +def sanitize(s: str) -> str: + return "".join(c if c.isalnum() or c in ("-", "_") else "_" for c in s) + + +# ----------------------------- +# Plotting +# ----------------------------- +def plot_err_vs_order(rows: List[Row], outdir: str, tag: str) -> None: + """ + One plot per method: x=ordering, y=err_dt_min (log scale) + Stable: circle marker, Unstable: x marker + """ + ensure_dir(outdir) + methods = sorted(set(r.method for r in rows)) + for method in methods: + sub = [r for r in rows if r.method == method] + + # sort by err (stable first), then ordering name + sub.sort(key=lambda r: (not r.stable, r.ordering)) + + xlabels = [r.ordering for r in sub] + x = np.arange(len(sub)) + y = np.array([r.err_dt_min for r in sub], dtype=float) + + # Replace non-finite with a large sentinel for plotting + finite = np.isfinite(y) & (y > 0) + if np.any(finite): + y_max = float(np.max(y[finite])) + y_plot = y.copy() + y_plot[~finite] = y_max * 10.0 + else: + y_plot = np.ones_like(y) + + fig = plt.figure() + ax = plt.gca() + ax.set_yscale("log") + + stable_idx = [i for i, r in enumerate(sub) if r.stable and np.isfinite(r.err_dt_min) and r.err_dt_min > 0] + unstable_idx = [i for i, r in enumerate(sub) if not r.stable or not (np.isfinite(r.err_dt_min) and r.err_dt_min > 0)] + + if stable_idx: + ax.plot(x[stable_idx], y_plot[stable_idx], linestyle="None", marker="o", label="stable") + if unstable_idx: + ax.plot(x[unstable_idx], y_plot[unstable_idx], linestyle="None", marker="x", label="unstable / inf") + + # annotate unstable points with max_step (if available) + for i in unstable_idx: + ms = sub[i].max_step + if np.isfinite(ms): + ax.annotate(f"{ms:.2g}", (x[i], y_plot[i]), textcoords="offset points", xytext=(0, 8), ha="center") + + ax.set_xticks(x) + ax.set_xticklabels(xlabels, rotation=0) + ax.set_xlabel("Ordering") + ax.set_ylabel("abs_err at smallest dt (log scale)") + ax.set_title(f"{tag} | {method} | error vs ordering") + ax.legend() + + fig.tight_layout() + outpath = os.path.join(outdir, f"err_vs_order_{sanitize(tag)}_{sanitize(method)}.png") + fig.savefig(outpath, dpi=200) + plt.close(fig) + + +def plot_maxstep_vs_order(rows: List[Row], outdir: str, tag: str) -> None: + """ + One plot per method: x=ordering, y=max_step (linear) + """ + ensure_dir(outdir) + methods = sorted(set(r.method for r in rows)) + for method in methods: + sub = [r for r in rows if r.method == method] + sub.sort(key=lambda r: (not r.stable, r.ordering)) + + xlabels = [r.ordering for r in sub] + x = np.arange(len(sub)) + y = np.array([r.max_step for r in sub], dtype=float) + + fig = plt.figure() + ax = plt.gca() + + stable_idx = [i for i, r in enumerate(sub) if r.stable and np.isfinite(r.max_step)] + unstable_idx = [i for i, r in enumerate(sub) if (not r.stable) and np.isfinite(r.max_step)] + + if stable_idx: + ax.plot(x[stable_idx], y[stable_idx], linestyle="None", marker="o", label="stable") + if unstable_idx: + ax.plot(x[unstable_idx], y[unstable_idx], linestyle="None", marker="x", label="unstable") + + ax.axhline(1.0, linewidth=1) + + ax.set_xticks(x) + ax.set_xticklabels(xlabels, rotation=0) + ax.set_xlabel("Ordering") + ax.set_ylabel("max_step = max ||y_{n+1}||/||y_n||") + ax.set_title(f"{tag} | {method} | max_step vs ordering") + ax.legend() + + fig.tight_layout() + outpath = os.path.join(outdir, f"maxstep_vs_order_{sanitize(tag)}_{sanitize(method)}.png") + fig.savefig(outpath, dpi=200) + plt.close(fig) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument( + "--csv_glob", + default="experiments/abc_outputs/abc_study_Bmult*.csv", + help="Glob for input CSVs", + ) + ap.add_argument( + "--outdir", + default="experiments/abc_outputs", + help="Where to save PNGs", + ) + args = ap.parse_args() + + paths = sorted(glob.glob(args.csv_glob)) + if not paths: + raise SystemExit(f"No CSVs matched: {args.csv_glob}") + + for p in paths: + rows = read_csv_rows(p) + if not rows: + print(f"Skipping (no rows): {p}") + continue + bmult = rows[0].bmult + tag = f"Bmult={bmult:g}" + plot_err_vs_order(rows, args.outdir, tag) + plot_maxstep_vs_order(rows, args.outdir, tag) + print(f"Saved plots for {tag} -> {args.outdir}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/experiments/abc_splitting_study.py b/experiments/abc_splitting_study.py new file mode 100644 index 0000000..1853632 --- /dev/null +++ b/experiments/abc_splitting_study.py @@ -0,0 +1,512 @@ +#!/usr/bin/env python3 +""" +ABC splitting ordering study. + +What it does +------------ +- Defines a linear 2x2 system y' = (A + B + C) y with tunable stiffness via Bmult. +- Compares orderings of 3-operator splittings for: + * Strang-3 (2nd order; symmetric) + * Yoshida-3 (4th order composition built from Strang; has negative substeps) +- For each (method, ordering, dt) computes error vs exact expm solution at final time. +- Reports observed order p from multiple dt values. +- Tracks stability/step-growth metrics: + * max||y||/||y0|| + * final||y||/||y0|| + * max_step_ratio = max_n ||y_{n+1}||/||y_n|| +- Saves per-regime CSV and generates plots: + * error vs ordering + * max_step_ratio vs ordering + +Run examples +------------ +python3 experiments/abc_splitting_study.py +python3 experiments/abc_splitting_study.py --Bmults 1,20,50,100 +python3 experiments/abc_splitting_study.py --dts 1e-2,5e-3,2.5e-3 --T 0.5 +""" + +from __future__ import annotations + +import argparse +import csv +import math +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Tuple + +import numpy as np + +# --------------------------------------------------------------------- +# Repository-local imports +# --------------------------------------------------------------------- +THIS = Path(__file__).resolve() +REPO = THIS.parents[1] +if str(REPO) not in sys.path: + sys.path.insert(0, str(REPO)) + +try: + from scipy.linalg import expm +except Exception: + expm = None + + +def expm_fallback(M: np.ndarray) -> np.ndarray: + """ + Small-matrix expm fallback using eigendecomposition. + Works well for 2x2 (our use case). + """ + w, V = np.linalg.eig(M) + Vinv = np.linalg.inv(V) + return (V @ np.diag(np.exp(w)) @ Vinv).astype(complex) + + +def mat_expm(M: np.ndarray) -> np.ndarray: + if expm is not None: + return expm(M) + E = expm_fallback(M) + # If imaginary noise is tiny, cast back to real. + if np.max(np.abs(E.imag)) < 1e-12: + return E.real + return E + + +def norm2(v: np.ndarray) -> float: + return float(np.linalg.norm(v)) + + +def safe_ratio(a: float, b: float) -> float: + if b == 0.0: + return float("inf") if a != 0.0 else 1.0 + return a / b + + +@dataclass +class RunMetrics: + err: float + max_norm_ratio: float + final_norm_ratio: float + max_step_ratio: float + stable: bool + + +# ------------------------------- +# Problem definition (A, B, C) +# ------------------------------- +def make_abc_matrices(Bmult: float) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + A, B, C are chosen so ordering matters (non-commuting operators), + and stiffness is introduced by scaling B. + + You can tweak these matrices if you want a stronger/weaker ordering effect. + """ + A = np.array([[0.0, 1.0], + [-1.0, 0.0]], dtype=float) # rotation-like + B = np.array([[-1.0, 0.0], + [0.0, -0.2]], dtype=float) * Bmult # damping/stiffness + C = np.array([[0.0, 0.6], + [0.0, 0.0]], dtype=float) # shear-like + return A, B, C + + +def exact_solution(y0: np.ndarray, T: float, A: np.ndarray, B: np.ndarray, C: np.ndarray) -> np.ndarray: + L = A + B + C + return mat_expm(T * L) @ y0 + + +# ------------------------------- +# Splitting building blocks +# ------------------------------- +def flow_linear(L: np.ndarray, h: float, y: np.ndarray) -> np.ndarray: + return mat_expm(h * L) @ y + + +def strang_step(ordering: str, h: float, y: np.ndarray, ops: Dict[str, np.ndarray]) -> np.ndarray: + """ + Strang splitting for 3 operators using an explicit ordering like 'ABC'. + + For 'ABC', we use: + exp(h/2 A) exp(h/2 B) exp(h C) exp(h/2 B) exp(h/2 A) + """ + a, b, c = ordering[0], ordering[1], ordering[2] + y1 = flow_linear(ops[a], 0.5 * h, y) + y2 = flow_linear(ops[b], 0.5 * h, y1) + y3 = flow_linear(ops[c], 1.0 * h, y2) + y4 = flow_linear(ops[b], 0.5 * h, y3) + y5 = flow_linear(ops[a], 0.5 * h, y4) + return y5 + + +def yoshida4_step(ordering: str, h: float, y: np.ndarray, ops: Dict[str, np.ndarray]) -> np.ndarray: + """ + 4th-order Yoshida composition built from Strang: + S4(h) = S2(w1 h) S2(w0 h) S2(w1 h) + where + w1 = 1/(2 - 2^(1/3)), w0 = -2^(1/3)/(2 - 2^(1/3)) + Note w0 is negative => potential stability issues in stiff regimes. + """ + cbrt2 = 2.0 ** (1.0 / 3.0) + w1 = 1.0 / (2.0 - cbrt2) + w0 = -cbrt2 / (2.0 - cbrt2) + + y1 = strang_step(ordering, w1 * h, y, ops) + y2 = strang_step(ordering, w0 * h, y1, ops) + y3 = strang_step(ordering, w1 * h, y2, ops) + return y3 + + +def simulate(method_name: str, ordering: str, dt: float, T: float, y0: np.ndarray, ops: Dict[str, np.ndarray]) -> RunMetrics: + steps = int(round(T / dt)) + if steps <= 0: + raise ValueError("dt too large for given T") + # Make dt so N*dt exactly hits T (avoid drift) + dt = T / steps + + y = y0.copy() + n0 = norm2(y0) + max_norm_ratio = 1.0 + max_step_ratio = 1.0 + stable = True + + stepper = strang_step if method_name == "Strang-3" else yoshida4_step + + for _ in range(steps): + y_new = stepper(ordering, dt, y, ops) + + ny = norm2(y) + ny_new = norm2(y_new) + if not np.all(np.isfinite(y_new)): + stable = False + break + + # step ratio ||y_{n+1}|| / ||y_n|| + if ny > 0: + max_step_ratio = max(max_step_ratio, ny_new / ny) + else: + max_step_ratio = max(max_step_ratio, float("inf")) + + # global ratio ||y||/||y0|| + if n0 > 0: + max_norm_ratio = max(max_norm_ratio, ny_new / n0) + + # crude stability guard: if it explodes massively, declare unstable + if ny_new > 1e12 * max(1.0, n0): + stable = False + break + + y = y_new + + yT = y + # If unstable, set err=inf + return RunMetrics( + err=float("inf"), + max_norm_ratio=max_norm_ratio, + final_norm_ratio=safe_ratio(norm2(yT), n0), + max_step_ratio=max_step_ratio, + stable=stable, + ) + + +def compute_error(y_approx: np.ndarray, y_exact: np.ndarray) -> float: + return norm2(y_approx - y_exact) + + +def run_one(method_name: str, ordering: str, dt: float, T: float, y0: np.ndarray, ops: Dict[str, np.ndarray]) -> Tuple[np.ndarray, RunMetrics]: + steps = int(round(T / dt)) + steps = max(1, steps) + dt = T / steps + + y = y0.copy() + n0 = norm2(y0) + max_norm_ratio = 1.0 + max_step_ratio = 1.0 + stable = True + + stepper = strang_step if method_name == "Strang-3" else yoshida4_step + + for _ in range(steps): + y_new = stepper(ordering, dt, y, ops) + + if not np.all(np.isfinite(y_new)): + stable = False + break + + ny = norm2(y) + ny_new = norm2(y_new) + if ny > 0: + max_step_ratio = max(max_step_ratio, ny_new / ny) + else: + max_step_ratio = max(max_step_ratio, float("inf")) + + if n0 > 0: + max_norm_ratio = max(max_norm_ratio, ny_new / n0) + + if ny_new > 1e12 * max(1.0, n0): + stable = False + break + + y = y_new + + metrics = RunMetrics( + err=float("inf"), + max_norm_ratio=max_norm_ratio, + final_norm_ratio=safe_ratio(norm2(y), n0), + max_step_ratio=max_step_ratio, + stable=stable, + ) + return y, metrics + + +def estimate_order(errs: List[float], dts: List[float]) -> float: + """ + Estimate observed order from (dt1,dt2,dt3...) and errs at final time. + Uses least-squares fit on log(err)=p log(dt)+const, ignoring non-finite. + """ + xs, ys = [], [] + for e, dt in zip(errs, dts): + if np.isfinite(e) and e > 0 and dt > 0: + xs.append(math.log(dt)) + ys.append(math.log(e)) + if len(xs) < 2: + return float("nan") + x = np.array(xs) + y = np.array(ys) + A = np.vstack([x, np.ones_like(x)]).T + p, _ = np.linalg.lstsq(A, y, rcond=None)[0] + return float(p) + + +def all_orderings() -> List[str]: + return ["ABC", "ACB", "BAC", "BCA", "CAB", "CBA"] + + +# ------------------------------- +# Plotting +# ------------------------------- +def make_plots_for_csv(csv_path: Path, outdir: Path) -> None: + import matplotlib.pyplot as plt + + rows = [] + with csv_path.open("r", newline="") as f: + reader = csv.DictReader(f) + for r in reader: + rows.append(r) + + if not rows: + print(f"[plot] No rows in {csv_path}") + return + + # infer Bmult and method from filename + stem = csv_path.stem # abc_study_Bmult20 + Bmult_label = stem.replace("abc_study_", "") + + # group by method + by_method: Dict[str, List[dict]] = {} + for r in rows: + by_method.setdefault(r["method"], []).append(r) + + for method, mr in by_method.items(): + # sort by fixed ordering list + ords = all_orderings() + mr_map = {r["ordering"]: r for r in mr} + errs = [float(mr_map[o]["err_dtmin"]) if o in mr_map else float("nan") for o in ords] + msteps = [float(mr_map[o]["max_step_ratio"]) if o in mr_map else float("nan") for o in ords] + stable = [mr_map[o]["status"] if o in mr_map else "?" for o in ords] + + # Error bar plot + plt.figure() + x = np.arange(len(ords)) + plt.bar(x, errs) + plt.xticks(x, ords) + plt.yscale("log") + plt.xlabel("Ordering") + plt.ylabel("Error at dt_min (log scale)") + plt.title(f"Error vs ordering — {Bmult_label} — {method}") + for i, st in enumerate(stable): + if st != "ok": + plt.text(i, errs[i], st, rotation=90, va="bottom", ha="center") + pth = outdir / f"err_vs_order_{Bmult_label}_{method}.png" + plt.tight_layout() + plt.savefig(pth, dpi=200) + plt.close() + + # max_step_ratio plot + plt.figure() + plt.bar(x, msteps) + plt.xticks(x, ords) + plt.yscale("log") + plt.xlabel("Ordering") + plt.ylabel("max_step_ratio = max ||y_{n+1}||/||y_n|| (log scale)") + plt.title(f"max_step_ratio vs ordering — {Bmult_label} — {method}") + pth = outdir / f"maxstep_vs_order_{Bmult_label}_{method}.png" + plt.tight_layout() + plt.savefig(pth, dpi=200) + plt.close() + + +# ------------------------------- +# Main experiment +# ------------------------------- +def run_regime(Bmult: float, T: float, dts: List[float], outdir: Path) -> Path: + A, B, C = make_abc_matrices(Bmult) + ops = {"A": A, "B": B, "C": C} + + y0 = np.array([1.0, 0.25], dtype=float) + y_exact = exact_solution(y0, T, A, B, C) + exact_norm = norm2(y_exact) + + print("\n" + "=" * 72) + print(f"Regime: Bmult={Bmult:g} T={T:g} dts={dts}") + print(f"Exact ||y(T)|| = {exact_norm:.12g} ||y0|| = {norm2(y0):.12g}") + print("=" * 72 + "\n") + + methods = ["Strang-3", "Yoshida-3"] + ords = all_orderings() + + outdir.mkdir(parents=True, exist_ok=True) + csv_path = outdir / f"abc_study_Bmult{int(Bmult) if float(Bmult).is_integer() else Bmult}.csv" + + with csv_path.open("w", newline="") as f: + writer = csv.writer(f) + writer.writerow([ + "Bmult", "T", "method", "ordering", + "err_dtmin", "p_est", + "max_norm_ratio", "final_norm_ratio", + "max_step_ratio", + "status" + ]) + + for method in methods: + print(f"Method: {method}") + print("-" * 72) + + # compute per ordering errors at each dt + results: Dict[str, Dict[str, float]] = {} + + for ordering in ords: + errs = [] + max_norms = [] + final_norms = [] + max_steps = [] + statuses = [] + + for dt in dts: + y_num, m = run_one(method, ordering, dt, T, y0, ops) + if not m.stable: + errs.append(float("inf")) + statuses.append("UNSTABLE") + else: + errs.append(compute_error(y_num, y_exact)) + statuses.append("ok") + max_norms.append(m.max_norm_ratio) + final_norms.append(m.final_norm_ratio) + max_steps.append(m.max_step_ratio) + + # dt_min is smallest dt + dt_min_idx = int(np.argmin(np.array(dts))) + err_dtmin = float(errs[dt_min_idx]) + + # order estimate from errs vs dt + p_est = estimate_order(errs, dts) + + # use dt_min metrics (consistent with err_dtmin) + max_norm_ratio = float(max_norms[dt_min_idx]) + final_norm_ratio = float(final_norms[dt_min_idx]) + max_step_ratio = float(max_steps[dt_min_idx]) + + status = "ok" + if not np.isfinite(err_dtmin): + status = "UNSTABLE" + + results[ordering] = { + "err_dtmin": err_dtmin, + "p_est": p_est, + "max_norm_ratio": max_norm_ratio, + "final_norm_ratio": final_norm_ratio, + "max_step_ratio": max_step_ratio, + "status": status + } + + print( + f"{ordering:3s} " + f"err(dt_min)={err_dtmin:.6e} " + f"p={p_est:.2f} " + f"max||y||/||y0||={max_norm_ratio:.6e} " + f"final||y||/||y0||={final_norm_ratio:.6e} " + f"max_step={max_step_ratio:.6e} " + f"{status}" + ) + + writer.writerow([ + f"{Bmult:g}", f"{T:g}", method, ordering, + f"{err_dtmin:.16e}", f"{p_est:.8g}", + f"{max_norm_ratio:.16e}", f"{final_norm_ratio:.16e}", + f"{max_step_ratio:.16e}", + status + ]) + + # best/worst based on err_dtmin among stable + stable_errs = [(o, results[o]["err_dtmin"]) for o in ords if np.isfinite(results[o]["err_dtmin"])] + if stable_errs: + best_o, best_e = min(stable_errs, key=lambda x: x[1]) + worst_o, worst_e = max(stable_errs, key=lambda x: x[1]) + spread = worst_e / best_e if best_e > 0 else float("inf") + print("\nBest ordering:", best_o, "err=", f"{best_e:.3e}") + print("Worst ordering:", worst_o, "err=", f"{worst_e:.3e}") + print(f"Spread (worst/best): {spread:.2f}x\n") + else: + print("\nAll orderings unstable for this method/regime.\n") + + print(f"CSV saved: {csv_path}") + return csv_path + + +def parse_csv_floats(s: str) -> List[float]: + out = [] + for part in s.split(","): + part = part.strip() + if not part: + continue + out.append(float(part)) + return out + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--T", type=float, default=0.5) + ap.add_argument("--dts", type=str, default="1e-2,5e-3,2.5e-3") + ap.add_argument("--Bmults", type=str, default="1,20") + ap.add_argument("--outdir", type=str, default="experiments/abc_outputs") + ap.add_argument("--plots", action="store_true", help="Generate plots after running.") + ap.add_argument("--plot_only", action="store_true", help="Only generate plots from existing CSVs in outdir.") + args = ap.parse_args() + + outdir = Path(args.outdir) + dts = parse_csv_floats(args.dts) + Bmults = parse_csv_floats(args.Bmults) + + if args.plot_only: + if not outdir.exists(): + print(f"[plot_only] outdir not found: {outdir}") + return + csvs = sorted(outdir.glob("abc_study_Bmult*.csv")) + if not csvs: + print(f"[plot_only] no CSVs found in {outdir}") + return + for c in csvs: + make_plots_for_csv(c, outdir) + print(f"[plot_only] Plots saved to: {outdir}") + return + + csv_paths = [] + for Bmult in Bmults: + csv_paths.append(run_regime(Bmult, args.T, dts, outdir)) + + if args.plots: + for c in csv_paths: + make_plots_for_csv(c, outdir) + print(f"Plots saved to: {outdir}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/experiments/bgk_adaptive_higher_order_splitting.py b/experiments/bgk_adaptive_higher_order_splitting.py new file mode 100644 index 0000000..037ed24 --- /dev/null +++ b/experiments/bgk_adaptive_higher_order_splitting.py @@ -0,0 +1,642 @@ +""" +Adaptive higher-order splitting study for a BGK-style kinetic test problem. + +This script reuses the generic adaptive step-doubling driver, records BGK-specific +diagnostics such as min(f), and compares adaptive and fixed-step runs. + +The file is structured so it can be connected to an existing three-operator BGK +implementation through make_bgk_problem(...), make_bgk_stepper(...), and +make_reference_solver(...). +""" + +from __future__ import annotations + +import argparse +import csv +import os +import time +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Tuple + +import numpy as np + +try: + from experiments.utils.adaptive_splitting import ( + AdaptiveOptions, + adaptive_integrate, + l2_norm, + ) +except ModuleNotFoundError: + import sys + + THIS_DIR = os.path.dirname(os.path.abspath(__file__)) + REPO_ROOT = os.path.dirname(THIS_DIR) + if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + + from experiments.utils.adaptive_splitting import ( + AdaptiveOptions, + adaptive_integrate, + l2_norm, + ) + + +Array = np.ndarray + + +# ============================================================================= +# Problem wrapper +# ============================================================================= + +@dataclass +class BGKProblem: + """ + Generic wrapper for a BGK-type problem. + + Required capabilities: + - initial state y0 + - optional reference integration + - one-step higher-order splitting method through make_bgk_stepper(...) + - diagnostics on state, especially min(f) + + State convention: + - by default we assume y is an ndarray containing f over phase space + - if your code uses a custom object, adapt diagnostics_fn and norm_fn + """ + case: str + Nx: int + Nv: int + xmax: float + vmax: float + eps0: float + transport: str + eta_model: str + init: str + + def initial_state(self) -> Array: + """ + Placeholder initial state used for repository structure testing. + """ + x = np.linspace(-self.xmax, self.xmax, self.Nx, endpoint=False) + v = np.linspace(-self.vmax, self.vmax, self.Nv, endpoint=False) + X, V = np.meshgrid(x, v, indexing="ij") + + rho = 1.0 + 0.2 * np.sin(np.pi * X / max(self.xmax, 1.0)) + u = 0.3 * np.cos(np.pi * X / max(self.xmax, 1.0)) + T = 1.0 + 0.1 * np.sin(2.0 * np.pi * X / max(self.xmax, 1.0)) + + f = rho / np.sqrt(2.0 * np.pi * T) * np.exp(-0.5 * ((V - u) ** 2) / T) + return f.astype(float) + + def diagnostics(self, y: Array) -> Dict[str, Any]: + """ + BGK-specific diagnostics stored per attempted adaptive step. + """ + arr = np.asarray(y, dtype=float) + return { + "min_f": float(np.min(arr)), + "max_f": float(np.max(arr)), + "l2_state": float(l2_norm(arr)), + } + + def error_norm(self, diff: Array) -> float: + """ + Error norm used when not using weighted RMS. + """ + return float(l2_norm(diff)) + + +# ============================================================================= +# Hook points to existing BGK implementation +# ============================================================================= + +def make_bgk_problem( + case: str, + Nx: int, + Nv: int, + xmax: float, + vmax: float, + eps0: float, + transport: str, + eta_model: str, + init: str, +) -> BGKProblem: + """ + Construct the BGK problem. + """ + return BGKProblem( + case=case, + Nx=Nx, + Nv=Nv, + xmax=xmax, + vmax=vmax, + eps0=eps0, + transport=transport, + eta_model=eta_model, + init=init, + ) + + +def make_bgk_stepper(problem: BGKProblem, method: str) -> Tuple[Callable[[Array, float, float], Array], int]: + """ + Return (stepper, nominal_order). + + This is the main hook for an existing higher-order BGK splitting implementation. + + Expected signature: + stepper(y, t, dt) -> y_next + + Supported placeholder methods: + - strang2 + - yoshida4 + + Note: + The placeholder below is only a structural stand-in for the project-specific + BGK operator splitting routine. + """ + method = method.lower() + + def demo_collision_relax(y: Array, dt: float) -> Array: + # mild relaxation toward local copy; structural placeholder only + lam = max(problem.eps0, 1.0e-12) + return np.exp(-dt / lam) * y + (1.0 - np.exp(-dt / lam)) * np.maximum(y, 0.0) + + def demo_transport(y: Array, dt: float) -> Array: + # simple smoothing/roll placeholder; replace with actual transport solve + z = np.asarray(y, dtype=float) + rolled = np.roll(z, shift=1, axis=0) + return z - 0.1 * dt * (z - rolled) + + def demo_aux(y: Array, dt: float) -> Array: + # placeholder third operator + z = np.asarray(y, dtype=float) + return np.maximum(z * (1.0 - 0.02 * dt), 0.0) + + def strang_step(y: Array, t: float, dt: float) -> Array: + z = demo_transport(y, 0.5 * dt) + z = demo_collision_relax(z, 0.5 * dt) + z = demo_aux(z, dt) + z = demo_collision_relax(z, 0.5 * dt) + z = demo_transport(z, 0.5 * dt) + return z + + if method == "strang2": + return strang_step, 2 + + if method == "yoshida4": + cbrt2 = 2.0 ** (1.0 / 3.0) + w1 = 1.0 / (2.0 - cbrt2) + w0 = -cbrt2 / (2.0 - cbrt2) + + def yoshida4_step(y: Array, t: float, dt: float) -> Array: + z = strang_step(y, t, w1 * dt) + z = strang_step(z, t + w1 * dt, w0 * dt) + z = strang_step(z, t + (w1 + w0) * dt, w1 * dt) + return z + + return yoshida4_step, 4 + + raise ValueError(f"Unknown method '{method}'. Supported: strang2, yoshida4") + + +def make_reference_solver(problem: BGKProblem, reference_method: str) -> Callable[[Array, float, float], Array]: + """ + Return a one-step reference integrator step(y, t, dt) -> y_next. + """ + reference_method = reference_method.lower() + + def ref_step(y: Array, t: float, dt: float) -> Array: + z = np.asarray(y, dtype=float) + z = np.roll(z, shift=1, axis=0) * (0.02 * dt) + z * (1.0 - 0.02 * dt) + z = np.maximum(z * np.exp(-0.2 * dt), 0.0) + return z + + if reference_method in ("rk4", "exprk2", "strang2", "yoshida4"): + return ref_step + + raise ValueError(f"Unknown reference method '{reference_method}'.") + + +# ============================================================================= +# Integration helpers +# ============================================================================= + +def integrate_fixed_step( + stepper: Callable[[Array, float, float], Array], + y0: Array, + T: float, + dt: float, +) -> Tuple[Array, int, float]: + t = 0.0 + y = np.array(y0, copy=True) + n_steps = 0 + tic = time.perf_counter() + + while t < T - 1.0e-15: + h = min(dt, T - t) + y = stepper(y, t, h) + t += h + n_steps += 1 + + wall = time.perf_counter() - tic + return y, n_steps, wall + + +def integrate_reference( + ref_stepper: Callable[[Array, float, float], Array], + y0: Array, + T: float, + dt_ref: float, +) -> Array: + t = 0.0 + y = np.array(y0, copy=True) + + while t < T - 1.0e-15: + h = min(dt_ref, T - t) + y = ref_stepper(y, t, h) + t += h + + return y + + +# ============================================================================= +# CSV helpers +# ============================================================================= + +def ensure_dir(path: str) -> None: + os.makedirs(path, exist_ok=True) + + +def append_summary_csv(path: str, row: Dict[str, Any]) -> None: + exists = os.path.exists(path) + with open(path, "a", newline="") as f: + writer = csv.DictWriter(f, fieldnames=list(row.keys())) + if not exists: + writer.writeheader() + writer.writerow(row) + + +def save_adaptive_history_csv(path: str, result) -> None: + with open(path, "w", newline="") as f: + writer = csv.writer(f) + writer.writerow([ + "attempt_index", + "t", + "dt_try", + "dt_suggested", + "accepted", + "err_est", + "cpu_seconds", + "min_f", + "max_f", + "l2_state", + ]) + for i, rec in enumerate(result.records): + extra = rec.extra or {} + writer.writerow([ + i, + rec.t, + rec.dt_try, + rec.dt_suggested, + int(rec.accepted), + rec.err_est, + rec.cpu_seconds, + extra.get("min_f", ""), + extra.get("max_f", ""), + extra.get("l2_state", ""), + ]) + + +def parse_float_list(s: str) -> List[float]: + return [float(x.strip()) for x in s.split(",") if x.strip()] + + +# ============================================================================= +# Case runners +# ============================================================================= + +def run_one_adaptive_case( + case: str, + method: str, + reference_method: str, + T: float, + dt0: float, + atol: float, + rtol: float, + dt_ref: float, + Nx: int, + Nv: int, + xmax: float, + vmax: float, + eps0: float, + transport: str, + eta_model: str, + init: str, + outdir: str, +) -> Dict[str, Any]: + problem = make_bgk_problem( + case=case, + Nx=Nx, + Nv=Nv, + xmax=xmax, + vmax=vmax, + eps0=eps0, + transport=transport, + eta_model=eta_model, + init=init, + ) + y0 = problem.initial_state() + + stepper, order = make_bgk_stepper(problem, method=method) + ref_stepper = make_reference_solver(problem, reference_method=reference_method) + + opts = AdaptiveOptions( + order=order, + atol=atol, + rtol=rtol, + dt_min=1.0e-12, + dt_max=max(dt0, T), + safety=0.9, + growth_max=2.0, + shrink_min=0.2, + max_reject=20, + use_weighted_rms=True, + ) + + result = adaptive_integrate( + stepper=stepper, + y0=y0, + t0=0.0, + tfinal=T, + dt0=dt0, + opts=opts, + diagnostics_fn=problem.diagnostics, + ) + + y_ref = integrate_reference(ref_stepper, y0, T=T, dt_ref=dt_ref) + err_global = l2_norm(result.final_state - y_ref) + + final_diag = problem.diagnostics(result.final_state) + + history_name = ( + f"bgk_history_mode-adaptive_method-{method}_case-{case}_" + f"eps0-{eps0:.0e}_atol-{atol:.0e}_rtol-{rtol:.0e}.csv" + ) + save_adaptive_history_csv(os.path.join(outdir, history_name), result) + + return { + "mode": "adaptive", + "case": case, + "method": method, + "reference_method": reference_method, + "order": order, + "T": T, + "dt0": dt0, + "fixed_dt": "", + "atol": atol, + "rtol": rtol, + "dt_ref": dt_ref, + "Nx": Nx, + "Nv": Nv, + "xmax": xmax, + "vmax": vmax, + "eps0": eps0, + "transport": transport, + "eta_model": eta_model, + "init": init, + "success": int(result.success), + "message": result.message, + "n_steps": result.n_accept, + "n_accept": result.n_accept, + "n_reject": result.n_reject, + "dt_min_used": result.dt_min_used, + "dt_max_used": result.dt_max_used, + "dt_avg_used": result.dt_avg_used, + "wall_seconds": result.wall_seconds, + "global_l2_error": err_global, + "final_min_f": final_diag["min_f"], + "final_max_f": final_diag["max_f"], + "final_l2_state": final_diag["l2_state"], + } + + +def run_one_fixed_case( + case: str, + method: str, + reference_method: str, + T: float, + fixed_dt: float, + dt_ref: float, + Nx: int, + Nv: int, + xmax: float, + vmax: float, + eps0: float, + transport: str, + eta_model: str, + init: str, +) -> Dict[str, Any]: + problem = make_bgk_problem( + case=case, + Nx=Nx, + Nv=Nv, + xmax=xmax, + vmax=vmax, + eps0=eps0, + transport=transport, + eta_model=eta_model, + init=init, + ) + y0 = problem.initial_state() + + stepper, order = make_bgk_stepper(problem, method=method) + ref_stepper = make_reference_solver(problem, reference_method=reference_method) + + y_final, n_steps, wall = integrate_fixed_step(stepper, y0, T=T, dt=fixed_dt) + y_ref = integrate_reference(ref_stepper, y0, T=T, dt_ref=dt_ref) + err_global = l2_norm(y_final - y_ref) + final_diag = problem.diagnostics(y_final) + + return { + "mode": "fixed", + "case": case, + "method": method, + "reference_method": reference_method, + "order": order, + "T": T, + "dt0": "", + "fixed_dt": fixed_dt, + "atol": "", + "rtol": "", + "dt_ref": dt_ref, + "Nx": Nx, + "Nv": Nv, + "xmax": xmax, + "vmax": vmax, + "eps0": eps0, + "transport": transport, + "eta_model": eta_model, + "init": init, + "success": 1, + "message": "Fixed-step integration completed successfully.", + "n_steps": n_steps, + "n_accept": n_steps, + "n_reject": 0, + "dt_min_used": fixed_dt, + "dt_max_used": fixed_dt, + "dt_avg_used": fixed_dt, + "wall_seconds": wall, + "global_l2_error": err_global, + "final_min_f": final_diag["min_f"], + "final_max_f": final_diag["max_f"], + "final_l2_state": final_diag["l2_state"], + } + + +# ============================================================================= +# Command-line interface +# ============================================================================= + +def main() -> None: + parser = argparse.ArgumentParser(description="Adaptive BGK higher-order splitting study.") + parser.add_argument("--mode", type=str, default="adaptive", + choices=["adaptive", "fixed", "both"]) + parser.add_argument("--case", type=str, default="mixed", + help="BGK case label, e.g. mixed, smooth, stiff.") + parser.add_argument("--method", type=str, default="yoshida4", + choices=["strang2", "yoshida4"]) + parser.add_argument("--reference-method", type=str, default="rk4", + help="Reference integrator label.") + parser.add_argument("--T", type=float, default=0.5) + parser.add_argument("--dt0", type=float, default=1.0e-3) + parser.add_argument("--atols", type=str, default="1e-6,1e-8") + parser.add_argument("--rtols", type=str, default="1e-4,1e-6") + parser.add_argument("--fixed-dts", type=str, default="2e-3,1e-3,5e-4") + parser.add_argument("--dt-ref", type=float, default=1.0e-5) + + parser.add_argument("--Nx", type=int, default=40) + parser.add_argument("--Nv", type=int, default=100) + parser.add_argument("--xmax", type=float, default=2.0) + parser.add_argument("--vmax", type=float, default=15.0) + parser.add_argument("--eps0-list", type=str, default="1e-2,1e-3,1e-4", + help="Comma-separated stiffness / regime parameters.") + parser.add_argument("--transport", type=str, default="muscl2") + parser.add_argument("--eta-model", type=str, default="constant1") + parser.add_argument("--init", type=str, default="mixed") + + parser.add_argument("--outdir", type=str, + default="experiments/outputs/bgk_adaptive_higher_order") + parser.add_argument("--overwrite-summary", action="store_true") + args = parser.parse_args() + + ensure_dir(args.outdir) + + atols = parse_float_list(args.atols) + rtols = parse_float_list(args.rtols) + fixed_dts = parse_float_list(args.fixed_dts) + eps0_list = parse_float_list(args.eps0_list) + + summary_path = os.path.join(args.outdir, "bgk_higher_order_comparison_summary.csv") + if args.overwrite_summary and os.path.exists(summary_path): + os.remove(summary_path) + + print("=" * 72) + print("BGK HIGHER-ORDER SPLITTING COMPARISON STUDY") + print(f"mode = {args.mode}") + print(f"case = {args.case}") + print(f"method = {args.method}") + print(f"reference_method = {args.reference_method}") + print(f"T = {args.T}") + print(f"dt0 = {args.dt0}") + print(f"atols = {atols}") + print(f"rtols = {rtols}") + print(f"fixed_dts = {fixed_dts}") + print(f"dt_ref = {args.dt_ref}") + print(f"Nx, Nv = {args.Nx}, {args.Nv}") + print(f"xmax, vmax = {args.xmax}, {args.vmax}") + print(f"eps0_list = {eps0_list}") + print(f"transport = {args.transport}") + print(f"eta_model = {args.eta_model}") + print(f"init = {args.init}") + print(f"outdir = {args.outdir}") + print("=" * 72) + + if args.mode in ("adaptive", "both"): + for eps0 in eps0_list: + for atol in atols: + for rtol in rtols: + print( + f"\nAdaptive: case={args.case}, method={args.method}, " + f"eps0={eps0:.0e}, atol={atol:.0e}, rtol={rtol:.0e}" + ) + + row = run_one_adaptive_case( + case=args.case, + method=args.method, + reference_method=args.reference_method, + T=args.T, + dt0=args.dt0, + atol=atol, + rtol=rtol, + dt_ref=args.dt_ref, + Nx=args.Nx, + Nv=args.Nv, + xmax=args.xmax, + vmax=args.vmax, + eps0=eps0, + transport=args.transport, + eta_model=args.eta_model, + init=args.init, + outdir=args.outdir, + ) + append_summary_csv(summary_path, row) + + print( + f" success={row['success']} " + f"accept={row['n_accept']} reject={row['n_reject']} " + f"dt[min,avg,max]=({row['dt_min_used']:.3e}, " + f"{row['dt_avg_used']:.3e}, {row['dt_max_used']:.3e}) " + f"error={row['global_l2_error']:.6e} " + f"min(f)={row['final_min_f']:.6e} " + f"wall={row['wall_seconds']:.3f}s" + ) + + if args.mode in ("fixed", "both"): + for eps0 in eps0_list: + for fixed_dt in fixed_dts: + print( + f"\nFixed: case={args.case}, method={args.method}, " + f"eps0={eps0:.0e}, dt={fixed_dt:.3e}" + ) + + row = run_one_fixed_case( + case=args.case, + method=args.method, + reference_method=args.reference_method, + T=args.T, + fixed_dt=fixed_dt, + dt_ref=args.dt_ref, + Nx=args.Nx, + Nv=args.Nv, + xmax=args.xmax, + vmax=args.vmax, + eps0=eps0, + transport=args.transport, + eta_model=args.eta_model, + init=args.init, + ) + append_summary_csv(summary_path, row) + + print( + f" steps={row['n_steps']} " + f"error={row['global_l2_error']:.6e} " + f"min(f)={row['final_min_f']:.6e} " + f"wall={row['wall_seconds']:.3f}s" + ) + + print("\nDone.") + print(f"Summary written to: {summary_path}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/experiments/bgk_three_operator_phase2_adaptive_study.py b/experiments/bgk_three_operator_phase2_adaptive_study.py new file mode 100644 index 0000000..f0251f2 --- /dev/null +++ b/experiments/bgk_three_operator_phase2_adaptive_study.py @@ -0,0 +1,677 @@ +from __future__ import annotations + +import argparse +import csv +import math +import sys +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Sequence, Tuple + +import numpy as np + +THIS = Path(__file__).resolve() +REPO = THIS.parents[1] if len(THIS.parents) > 1 else THIS.parent +if str(REPO) not in sys.path: + sys.path.insert(0, str(REPO)) + +from fractional_step import fractional_step # noqa: E402 + + +def parse_csv_floats(s: str) -> List[float]: + return [float(x.strip()) for x in s.split(",") if x.strip()] + + +def parse_csv_strings(s: str) -> List[str]: + return [x.strip() for x in s.split(",") if x.strip()] + + +def trapz_weights_uniform(n: int, h: float) -> np.ndarray: + return np.ones(n, dtype=float) * h + + +def periodic_shift_1d(values: np.ndarray, shift_cells: float) -> np.ndarray: + n = values.shape[0] + idx = np.arange(n, dtype=float) + src = (idx - shift_cells) % n + i0 = np.floor(src).astype(int) + i1 = (i0 + 1) % n + theta = src - i0 + return (1.0 - theta) * values[i0] + theta * values[i1] + + +def shift_along_x_periodic(f: np.ndarray, shift_cells_per_v: np.ndarray) -> np.ndarray: + Nx, Nv = f.shape + out = np.empty_like(f) + for j in range(Nv): + out[:, j] = periodic_shift_1d(f[:, j], shift_cells_per_v[j]) + return out + + +def shift_along_v_periodic(f: np.ndarray, shift_cells: float) -> np.ndarray: + Nx, Nv = f.shape + out = np.empty_like(f) + for i in range(Nx): + out[i, :] = periodic_shift_1d(f[i, :], shift_cells) + return out + + +def compute_moments( + f: np.ndarray, + v: np.ndarray, + dv: float, + rho_floor: float = 1e-14, + T_floor: float = 1e-12, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + rho = np.sum(f, axis=1) * dv + rho_safe = np.maximum(rho, rho_floor) + mom1 = np.sum(f * v[None, :], axis=1) * dv + u = mom1 / rho_safe + centered = v[None, :] - u[:, None] + energy = np.sum(f * centered * centered, axis=1) * dv + T = np.maximum(energy / rho_safe, T_floor) + return rho, u, T + + +def maxwellian_from_moments( + rho: np.ndarray, + u: np.ndarray, + T: np.ndarray, + v: np.ndarray, + rho_floor: float = 1e-14, + T_floor: float = 1e-12, +) -> np.ndarray: + rho_safe = np.maximum(rho, rho_floor) + T_safe = np.maximum(T, T_floor) + pref = rho_safe[:, None] / np.sqrt(2.0 * np.pi * T_safe[:, None]) + z = (v[None, :] - u[:, None]) ** 2 / (2.0 * T_safe[:, None]) + return pref * np.exp(-z) + + +def initial_condition(x: np.ndarray, v: np.ndarray, profile: str = "mixed") -> np.ndarray: + X = x[:, None] + V = v[None, :] + + if profile == "mixed": + rho = 1.0 + 0.2 * np.sin(2.0 * np.pi * X) + u = 0.5 * np.cos(2.0 * np.pi * X) + T = 0.8 + 0.1 * np.sin(4.0 * np.pi * X) + M1 = rho / np.sqrt(2.0 * np.pi * T) * np.exp(-((V - u) ** 2) / (2.0 * T)) + M2 = 0.35 * rho / np.sqrt(2.0 * np.pi * T) * np.exp(-((V + 0.7 * u) ** 2) / (2.0 * (1.25 * T))) + f0 = 0.75 * M1 + 0.25 * M2 + return np.maximum(f0, 1e-14) + + if profile == "simple": + rho = 1.0 + 0.1 * np.sin(2.0 * np.pi * X) + u = 0.25 * np.cos(2.0 * np.pi * X) + T = 1.0 + 0.05 * np.sin(2.0 * np.pi * X) + f0 = rho / np.sqrt(2.0 * np.pi * T) * np.exp(-((V - u) ** 2) / (2.0 * T)) + return np.maximum(f0, 1e-14) + + raise ValueError("profile must be 'mixed' or 'simple'") + + +@dataclass +class BGKPhase2Problem: + x: np.ndarray + v: np.ndarray + dx: float + dv: float + force: float + eps: float + + +def build_problem(Nx: int, Nv: int, xmax: float, vmax: float, force: float, eps: float) -> BGKPhase2Problem: + x = np.linspace(0.0, xmax, Nx, endpoint=False) + v = np.linspace(-vmax, vmax, Nv, endpoint=False) + dx = xmax / Nx + dv = 2.0 * vmax / Nv + return BGKPhase2Problem(x=x, v=v, dx=dx, dv=dv, force=force, eps=eps) + + +def make_transport_flow(problem: BGKPhase2Problem): + v = problem.v + dx = problem.dx + + def flow(_t: float, h: float, f: np.ndarray) -> np.ndarray: + shift_cells = (v * h) / dx + return shift_along_x_periodic(f, shift_cells) + + return flow + + +def make_force_flow(problem: BGKPhase2Problem): + dv = problem.dv + F = problem.force + + def flow(_t: float, h: float, f: np.ndarray) -> np.ndarray: + shift_cells = (F * h) / dv + return shift_along_v_periodic(f, shift_cells) + + return flow + + +def make_bgk_relaxation_flow(problem: BGKPhase2Problem): + v = problem.v + dv = problem.dv + eps = problem.eps + + def flow(_t: float, h: float, f: np.ndarray) -> np.ndarray: + rho, u, T = compute_moments(f, v, dv) + M = maxwellian_from_moments(rho, u, T, v) + E = np.exp(-h / eps) + return E * f + (1.0 - E) * M + + return flow + + +def total_mass(f: np.ndarray, dx: float, dv: float) -> float: + return float(np.sum(f) * dx * dv) + + +def l1_norm(f: np.ndarray, dx: float, dv: float) -> float: + return float(np.sum(np.abs(f)) * dx * dv) + + +def make_flows(problem: BGKPhase2Problem): + return [ + make_transport_flow(problem), + make_bgk_relaxation_flow(problem), + make_force_flow(problem), + ] + + +def bgk_split_step(method_name: str, problem: BGKPhase2Problem, f: np.ndarray, t: float, dt: float) -> np.ndarray: + flows = make_flows(problem) + f_new = fractional_step( + functions=flows, + delta_t=dt, + initial_y=f, + initial_t=t, + final_t=t + dt, + alpha=method_name, + methods={(1,): "ANALYTIC", (2,): "ANALYTIC", (3,): "ANALYTIC"}, + ) + return np.array(f_new, dtype=float, copy=False) + + +def compute_reference(ref_method: str, ref_dt: float, T: float, f0: np.ndarray, problem: BGKPhase2Problem) -> np.ndarray: + flows = make_flows(problem) + f_ref = fractional_step( + functions=flows, + delta_t=ref_dt, + initial_y=f0, + initial_t=0.0, + final_t=T, + alpha=ref_method, + methods={(1,): "ANALYTIC", (2,): "ANALYTIC", (3,): "ANALYTIC"}, + ) + return np.array(f_ref, dtype=float) + + +@dataclass +class AdaptiveOptions: + order: int + atol: float + rtol: float + dt_min: float = 1e-12 + dt_max: float = 1.0 + safety: float = 0.9 + growth_max: float = 2.0 + shrink_min: float = 0.2 + max_reject: int = 20 + positivity_tol: float = -1e-12 + explosion_factor: float = 1e8 + + +@dataclass +class AttemptRecord: + t: float + dt_try: float + dt_suggested: float + accepted: bool + err_est: float + cpu_seconds: float + min_f: float + max_f: float + l1_ratio: float + stable: bool + + +@dataclass +class AdaptiveRunMetrics: + err: float + min_f_over_run: float + final_min_f: float + final_mass: float + mass_rel_drift: float + max_l1_ratio: float + positive: bool + stable: bool + n_accept: int + n_reject: int + dt_min_used: float + dt_avg_used: float + dt_max_used: float + wall_seconds: float + message: str + history: List[AttemptRecord] + + +def weighted_rms_error(y_high: np.ndarray, y_low: np.ndarray, atol: float, rtol: float) -> float: + scale = atol + rtol * np.maximum(np.abs(y_high), np.abs(y_low)) + diff = (y_high - y_low) / scale + return float(np.sqrt(np.mean(diff * diff))) + + +def propose_new_dt(err: float, dt: float, opts: AdaptiveOptions) -> float: + if not np.isfinite(err): + factor = opts.shrink_min + elif err <= 0.0: + factor = opts.growth_max + else: + factor = opts.safety * (err ** (-1.0 / (opts.order + 1.0))) + factor = min(opts.growth_max, max(opts.shrink_min, factor)) + return min(opts.dt_max, max(opts.dt_min, dt * factor)) + + +def bgk_diagnostics(f: np.ndarray, problem: BGKPhase2Problem, l10: float) -> Tuple[float, float, float, bool]: + min_f = float(np.min(f)) + max_f = float(np.max(f)) + l1r = l1_norm(f, problem.dx, problem.dv) / max(l10, 1e-30) + stable = bool(np.all(np.isfinite(f)) and l1r <= 1e8) + return min_f, max_f, l1r, stable + + +def run_adaptive_method( + method_name: str, + order: int, + dt0: float, + T: float, + f0: np.ndarray, + problem: BGKPhase2Problem, + reference: np.ndarray, + opts: AdaptiveOptions, +) -> AdaptiveRunMetrics: + f = np.array(f0, dtype=float, copy=True) + t = 0.0 + dt = min(max(dt0, opts.dt_min), opts.dt_max) + m0 = total_mass(f0, problem.dx, problem.dv) + l10 = l1_norm(f0, problem.dx, problem.dv) + min_f_over_run = float(np.min(f)) + max_l1_ratio = 1.0 + accepted_dts: List[float] = [] + history: List[AttemptRecord] = [] + n_accept = 0 + n_reject = 0 + start = time.perf_counter() + + while t < T - 1e-15: + dt = min(dt, T - t) + local_rejects = 0 + + while True: + tic = time.perf_counter() + full = bgk_split_step(method_name, problem, f, t, dt) + half = bgk_split_step(method_name, problem, f, t, 0.5 * dt) + two_half = bgk_split_step(method_name, problem, half, t + 0.5 * dt, 0.5 * dt) + cpu = time.perf_counter() - tic + + min_f, max_f, l1r, stable = bgk_diagnostics(two_half, problem, l10) + err = weighted_rms_error(two_half, full, opts.atol, opts.rtol) + if not np.isfinite(err): + stable = False + if min_f < opts.positivity_tol: + stable = False + err = float("inf") + + dt_new = propose_new_dt(err, dt, opts) + accepted = bool(stable and err <= 1.0) + + history.append(AttemptRecord( + t=t, + dt_try=dt, + dt_suggested=dt_new, + accepted=accepted, + err_est=err, + cpu_seconds=cpu, + min_f=min_f, + max_f=max_f, + l1_ratio=l1r, + stable=stable, + )) + + if accepted: + f = two_half + t += dt + n_accept += 1 + accepted_dts.append(dt) + min_f_over_run = min(min_f_over_run, min_f) + max_l1_ratio = max(max_l1_ratio, l1r) + dt = dt_new + break + + n_reject += 1 + local_rejects += 1 + dt = dt_new + if dt <= opts.dt_min + 1e-30: + wall = time.perf_counter() - start + final_min_f = float(np.min(f)) + final_mass = total_mass(f, problem.dx, problem.dv) + mass_rel_drift = abs(final_mass - m0) / max(abs(m0), 1e-30) + err_final = l1_norm(f - reference, problem.dx, problem.dv) + return AdaptiveRunMetrics( + err=err_final, + min_f_over_run=min_f_over_run, + final_min_f=final_min_f, + final_mass=final_mass, + mass_rel_drift=mass_rel_drift, + max_l1_ratio=max_l1_ratio, + positive=bool(final_min_f >= opts.positivity_tol), + stable=False, + n_accept=n_accept, + n_reject=n_reject, + dt_min_used=min(accepted_dts) if accepted_dts else 0.0, + dt_avg_used=float(np.mean(accepted_dts)) if accepted_dts else 0.0, + dt_max_used=max(accepted_dts) if accepted_dts else 0.0, + wall_seconds=wall, + message="Step size reached dt_min during rejection.", + history=history, + ) + if local_rejects >= opts.max_reject: + wall = time.perf_counter() - start + final_min_f = float(np.min(f)) + final_mass = total_mass(f, problem.dx, problem.dv) + mass_rel_drift = abs(final_mass - m0) / max(abs(m0), 1e-30) + err_final = l1_norm(f - reference, problem.dx, problem.dv) + return AdaptiveRunMetrics( + err=err_final, + min_f_over_run=min_f_over_run, + final_min_f=final_min_f, + final_mass=final_mass, + mass_rel_drift=mass_rel_drift, + max_l1_ratio=max_l1_ratio, + positive=bool(final_min_f >= opts.positivity_tol), + stable=False, + n_accept=n_accept, + n_reject=n_reject, + dt_min_used=min(accepted_dts) if accepted_dts else 0.0, + dt_avg_used=float(np.mean(accepted_dts)) if accepted_dts else 0.0, + dt_max_used=max(accepted_dts) if accepted_dts else 0.0, + wall_seconds=wall, + message="Exceeded maximum rejects for one advance.", + history=history, + ) + + wall = time.perf_counter() - start + final_min_f = float(np.min(f)) + final_mass = total_mass(f, problem.dx, problem.dv) + mass_rel_drift = abs(final_mass - m0) / max(abs(m0), 1e-30) + err_final = l1_norm(f - reference, problem.dx, problem.dv) + return AdaptiveRunMetrics( + err=err_final, + min_f_over_run=min_f_over_run, + final_min_f=final_min_f, + final_mass=final_mass, + mass_rel_drift=mass_rel_drift, + max_l1_ratio=max_l1_ratio, + positive=bool(final_min_f >= opts.positivity_tol), + stable=True, + n_accept=n_accept, + n_reject=n_reject, + dt_min_used=min(accepted_dts) if accepted_dts else 0.0, + dt_avg_used=float(np.mean(accepted_dts)) if accepted_dts else 0.0, + dt_max_used=max(accepted_dts) if accepted_dts else 0.0, + wall_seconds=wall, + message="Adaptive integration completed successfully.", + history=history, + ) + + +def run_fixed_method(method_name: str, dt: float, T: float, f0: np.ndarray, problem: BGKPhase2Problem, reference: np.ndarray, + positivity_tol: float = -1e-12, explosion_factor: float = 1e8): + steps = max(1, int(round(T / dt))) + dt = T / steps + f = np.array(f0, dtype=float, copy=True) + t = 0.0 + m0 = total_mass(f0, problem.dx, problem.dv) + l10 = l1_norm(f0, problem.dx, problem.dv) + min_f_over_run = float(np.min(f)) + max_l1_ratio = 1.0 + stable = True + tic = time.perf_counter() + for _ in range(steps): + f_new = bgk_split_step(method_name, problem, f, t, dt) + if not np.all(np.isfinite(f_new)): + stable = False + f = f_new + break + min_f_over_run = min(min_f_over_run, float(np.min(f_new))) + l1r = l1_norm(f_new, problem.dx, problem.dv) / max(l10, 1e-30) + max_l1_ratio = max(max_l1_ratio, l1r) + if l1r > explosion_factor: + stable = False + f = f_new + break + f = f_new + t += dt + wall = time.perf_counter() - tic + final_min_f = float(np.min(f)) if np.all(np.isfinite(f)) else float("nan") + final_mass = total_mass(f, problem.dx, problem.dv) if np.all(np.isfinite(f)) else float("nan") + mass_rel_drift = abs(final_mass - m0) / max(abs(m0), 1e-30) if np.isfinite(final_mass) else float("inf") + positive = bool(final_min_f >= positivity_tol) if np.isfinite(final_min_f) else False + err = l1_norm(f - reference, problem.dx, problem.dv) if stable and np.all(np.isfinite(f)) else float("inf") + return { + "err": err, + "min_f_over_run": min_f_over_run, + "final_min_f": final_min_f, + "final_mass": final_mass, + "mass_rel_drift": mass_rel_drift, + "max_l1_ratio": max_l1_ratio, + "positive": positive, + "stable": stable, + "n_steps": steps, + "wall_seconds": wall, + "dt": dt, + } + + +def method_order(method_name: str) -> int: + name = method_name.lower() + if "strang" in name: + return 2 + return 3 + + +def save_history_csv(path: Path, history: List[AttemptRecord]) -> None: + with path.open("w", newline="") as f: + writer = csv.writer(f) + writer.writerow([ + "attempt_index", "t", "dt_try", "dt_suggested", "accepted", "err_est", "cpu_seconds", + "min_f", "max_f", "l1_ratio", "stable" + ]) + for i, rec in enumerate(history): + writer.writerow([ + i, rec.t, rec.dt_try, rec.dt_suggested, int(rec.accepted), rec.err_est, rec.cpu_seconds, + rec.min_f, rec.max_f, rec.l1_ratio, int(rec.stable) + ]) + + +def build_parser() -> argparse.ArgumentParser: + ap = argparse.ArgumentParser(description="Adaptive BGK-inspired three-operator higher-order study.") + ap.add_argument("--mode", type=str, default="adaptive", choices=["adaptive", "fixed", "both"]) + ap.add_argument("--eps-list", type=str, default="1e-1,1e-2,1e-3") + ap.add_argument("--fixed-dts", type=str, default="1e-2,5e-3,2.5e-3") + ap.add_argument("--methods", type=str, default="Strang-3,OS32_7op_minLEM-3,PP3_4A-3,Yoshida-3") + ap.add_argument("--reference-method", type=str, default="Yoshida-3") + ap.add_argument("--reference-refine", type=int, default=8) + ap.add_argument("--T", type=float, default=0.2) + ap.add_argument("--dt0", type=float, default=1e-3) + ap.add_argument("--atols", type=str, default="1e-6,1e-8") + ap.add_argument("--rtols", type=str, default="1e-4,1e-6") + ap.add_argument("--Nx", type=int, default=40) + ap.add_argument("--Nv", type=int, default=80) + ap.add_argument("--xmax", type=float, default=1.0) + ap.add_argument("--vmax", type=float, default=8.0) + ap.add_argument("--force", type=float, default=0.5) + ap.add_argument("--profile", type=str, choices=["mixed", "simple"], default="mixed") + ap.add_argument("--outdir", type=str, default="experiments/bgk_three_operator_adaptive_outputs") + ap.add_argument("--overwrite-summary", action="store_true") + return ap + + +def main() -> None: + ap = build_parser() + args = ap.parse_args() + + eps_list = parse_csv_floats(args.eps_list) + fixed_dts = parse_csv_floats(args.fixed_dts) + methods = parse_csv_strings(args.methods) + atols = parse_csv_floats(args.atols) + rtols = parse_csv_floats(args.rtols) + + outdir = REPO / args.outdir + outdir.mkdir(parents=True, exist_ok=True) + summary_path = outdir / "bgk_three_operator_adaptive_comparison_summary.csv" + if args.overwrite_summary and summary_path.exists(): + summary_path.unlink() + + print("Running adaptive BGK-inspired three-operator study with:") + print(f" mode = {args.mode}") + print(f" eps_list = {eps_list}") + print(f" fixed_dts = {fixed_dts}") + print(f" methods = {methods}") + print(f" reference_method = {args.reference_method}") + print(f" reference_refine = {args.reference_refine}") + print(f" T = {args.T}") + print(f" dt0 = {args.dt0}") + print(f" atols = {atols}") + print(f" rtols = {rtols}") + print(f" Nx, Nv = {args.Nx}, {args.Nv}") + print(f" xmax, vmax = {args.xmax}, {args.vmax}") + print(f" force = {args.force}") + print(f" profile = {args.profile}") + print(f" outdir = {outdir}") + + wrote_header = False + for eps in eps_list: + problem = build_problem(args.Nx, args.Nv, args.xmax, args.vmax, args.force, eps) + f0 = initial_condition(problem.x, problem.v, profile=args.profile) + dt_min = min(fixed_dts + [args.dt0]) + ref_dt = dt_min / args.reference_refine + reference = compute_reference(args.reference_method, ref_dt, args.T, f0, problem) + + print("\n" + "=" * 88) + print(f"REGIME eps = {eps:.3e} | T = {args.T:g} | ref_dt = {ref_dt:.4e}") + print("=" * 88) + + rows: List[dict] = [] + for method_name in methods: + if args.mode in ("adaptive", "both"): + order = method_order(method_name) + for atol in atols: + for rtol in rtols: + opts = AdaptiveOptions(order=order, atol=atol, rtol=rtol, dt_max=max(args.dt0, args.T)) + metrics = run_adaptive_method(method_name, order, args.dt0, args.T, f0, problem, reference, opts) + print( + f"Adaptive | {method_name:18s} | atol={atol:.0e} rtol={rtol:.0e} | " + f"accept={metrics.n_accept:4d} reject={metrics.n_reject:4d} | " + f"dt[min,avg,max]=({metrics.dt_min_used:.3e}, {metrics.dt_avg_used:.3e}, {metrics.dt_max_used:.3e}) | " + f"err={metrics.err:.5e} | stable={metrics.stable!s:5s} | positive={metrics.positive!s:5s} | " + f"min_f={metrics.min_f_over_run:.5e}" + ) + hist_name = f"bgk_phase2_history_eps{eps:.0e}_{method_name}_atol{atol:.0e}_rtol{rtol:.0e}.csv".replace("/", "-") + save_history_csv(outdir / hist_name, metrics.history) + rows.append({ + "mode": "adaptive", + "eps": eps, + "T": args.T, + "method": method_name, + "fixed_dt": "", + "dt0": args.dt0, + "atol": atol, + "rtol": rtol, + "ref_method": args.reference_method, + "ref_dt": ref_dt, + "Nx": args.Nx, + "Nv": args.Nv, + "xmax": args.xmax, + "vmax": args.vmax, + "force": args.force, + "profile": args.profile, + "err_l1": metrics.err, + "stable": int(metrics.stable), + "positive": int(metrics.positive), + "min_f_over_run": metrics.min_f_over_run, + "final_min_f": metrics.final_min_f, + "final_mass": metrics.final_mass, + "mass_rel_drift": metrics.mass_rel_drift, + "max_l1_ratio": metrics.max_l1_ratio, + "n_steps": metrics.n_accept, + "n_accept": metrics.n_accept, + "n_reject": metrics.n_reject, + "dt_min_used": metrics.dt_min_used, + "dt_avg_used": metrics.dt_avg_used, + "dt_max_used": metrics.dt_max_used, + "wall_seconds": metrics.wall_seconds, + "message": metrics.message, + }) + + if args.mode in ("fixed", "both"): + for dt in fixed_dts: + metrics = run_fixed_method(method_name, dt, args.T, f0, problem, reference) + print( + f"Fixed | {method_name:18s} | dt={metrics['dt']:.3e} | " + f"steps={metrics['n_steps']:4d} | err={metrics['err']:.5e} | stable={metrics['stable']!s:5s} | " + f"positive={metrics['positive']!s:5s} | min_f={metrics['min_f_over_run']:.5e}" + ) + rows.append({ + "mode": "fixed", + "eps": eps, + "T": args.T, + "method": method_name, + "fixed_dt": metrics["dt"], + "dt0": "", + "atol": "", + "rtol": "", + "ref_method": args.reference_method, + "ref_dt": ref_dt, + "Nx": args.Nx, + "Nv": args.Nv, + "xmax": args.xmax, + "vmax": args.vmax, + "force": args.force, + "profile": args.profile, + "err_l1": metrics["err"], + "stable": int(metrics["stable"]), + "positive": int(metrics["positive"]), + "min_f_over_run": metrics["min_f_over_run"], + "final_min_f": metrics["final_min_f"], + "final_mass": metrics["final_mass"], + "mass_rel_drift": metrics["mass_rel_drift"], + "max_l1_ratio": metrics["max_l1_ratio"], + "n_steps": metrics["n_steps"], + "n_accept": metrics["n_steps"], + "n_reject": 0, + "dt_min_used": metrics["dt"], + "dt_avg_used": metrics["dt"], + "dt_max_used": metrics["dt"], + "wall_seconds": metrics["wall_seconds"], + "message": "Fixed-step integration completed successfully.", + }) + + with summary_path.open("a", newline="") as f: + fieldnames = list(rows[0].keys()) if rows else [] + writer = csv.DictWriter(f, fieldnames=fieldnames) + if not wrote_header and fieldnames: + writer.writeheader() + wrote_header = True + for row in rows: + writer.writerow(row) + + print(f"\nSaved summary: {summary_path}") + + +if __name__ == "__main__": + main() diff --git a/experiments/bgk_three_operator_phase2_study.py b/experiments/bgk_three_operator_phase2_study.py new file mode 100644 index 0000000..69435fb --- /dev/null +++ b/experiments/bgk_three_operator_phase2_study.py @@ -0,0 +1,636 @@ +""" +bgk_three_operator_phase2_study.py + +Phase 2: BGK-inspired three-operator splitting study. + +Purpose +------- +Move beyond linear ABC matrix tests and compare higher-order 3-operator +splitting methods on a small kinetic-style model with: + A = transport in x + B = BGK relaxation toward a local Maxwellian + C = force/advection in v + +Model +----- +We evolve f(x, v, t) on a small Cartesian grid according to + + f_t + v f_x + F f_v = ( M[f] - f ) / eps + +and split it into three operators: + A: f_t = - v f_x + B: f_t = ( M[f] - f ) / eps + C: f_t = - F f_v + +Subflows +-------- +A: exact shift in x by distance v * h using periodic interpolation +B: exact homogeneous BGK relaxation: + f(h) = exp(-h/eps) f0 + (1 - exp(-h/eps)) M[f0] + where M[f0] uses the local moments of f0 in x +C: exact shift in v by distance F * h using periodic interpolation + +What we measure +--------------- +For each method and dt: +- error against a fine-reference solution +- minimum value of f over the run +- final min(f) +- mass drift +- max L1 norm ratio +- simple positivity flag + +Recommended methods +------------------- +- Strang-3 +- OS32_7op_minLEM-3 +- PP3_4A-3 +- Yoshida-3 + +Examples +-------- +python experiments/bgk_three_operator_phase2_study.py +python experiments/bgk_three_operator_phase2_study.py --eps-list 1e-1,1e-2,1e-3 --T 0.2 +python experiments/bgk_three_operator_phase2_study.py --methods Strang-3,PP3_4A-3,Yoshida-3 +""" + +from __future__ import annotations + +import argparse +import csv +import math +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Dict, List, Sequence, Tuple + +import numpy as np + +# --------------------------------------------------------------------- +# Robust repo-root imports +# --------------------------------------------------------------------- +THIS = Path(__file__).resolve() +REPO = THIS.parents[1] +if str(REPO) not in sys.path: + sys.path.insert(0, str(REPO)) + +from fractional_step import fractional_step # noqa: E402 + + +# --------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------- +def parse_csv_floats(s: str) -> List[float]: + return [float(x.strip()) for x in s.split(",") if x.strip()] + + +def parse_csv_strings(s: str) -> List[str]: + return [x.strip() for x in s.split(",") if x.strip()] + + +def estimate_order(errs: Sequence[float], dts: Sequence[float]) -> float: + xs: List[float] = [] + ys: List[float] = [] + for e, dt in zip(errs, dts): + if np.isfinite(e) and e > 0.0 and dt > 0.0: + xs.append(math.log(dt)) + ys.append(math.log(e)) + if len(xs) < 2: + return float("nan") + x = np.array(xs, dtype=float) + y = np.array(ys, dtype=float) + A = np.vstack([x, np.ones_like(x)]).T + p, _ = np.linalg.lstsq(A, y, rcond=None)[0] + return float(p) + + +def trapz_weights_uniform(n: int, h: float) -> np.ndarray: + w = np.ones(n, dtype=float) * h + return w + + +def periodic_shift_1d(values: np.ndarray, shift_cells: float) -> np.ndarray: + """ + Periodic linear interpolation for a 1D array. + Positive shift_cells means evaluate values(x - shift). + """ + n = values.shape[0] + idx = np.arange(n, dtype=float) + src = (idx - shift_cells) % n + i0 = np.floor(src).astype(int) + i1 = (i0 + 1) % n + theta = src - i0 + return (1.0 - theta) * values[i0] + theta * values[i1] + + +def shift_along_x_periodic(f: np.ndarray, shift_cells_per_v: np.ndarray) -> np.ndarray: + """ + f has shape (Nx, Nv). For each velocity column j, shift in x. + """ + Nx, Nv = f.shape + out = np.empty_like(f) + for j in range(Nv): + out[:, j] = periodic_shift_1d(f[:, j], shift_cells_per_v[j]) + return out + + +def shift_along_v_periodic(f: np.ndarray, shift_cells: float) -> np.ndarray: + """ + Shift in v for each x-row. + """ + Nx, Nv = f.shape + out = np.empty_like(f) + for i in range(Nx): + out[i, :] = periodic_shift_1d(f[i, :], shift_cells) + return out + + +# --------------------------------------------------------------------- +# Moments / Maxwellian +# --------------------------------------------------------------------- +def compute_moments( + f: np.ndarray, + v: np.ndarray, + dv: float, + rho_floor: float = 1e-14, + T_floor: float = 1e-12, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Moments per x: + rho(x) = ∫ f dv + u(x) = (1/rho) ∫ v f dv + T(x) = (1/rho) ∫ (v-u)^2 f dv + """ + rho = np.sum(f, axis=1) * dv + rho_safe = np.maximum(rho, rho_floor) + + mom1 = np.sum(f * v[None, :], axis=1) * dv + u = mom1 / rho_safe + + centered = v[None, :] - u[:, None] + energy = np.sum(f * centered * centered, axis=1) * dv + T = np.maximum(energy / rho_safe, T_floor) + + return rho, u, T + + +def maxwellian_from_moments( + rho: np.ndarray, + u: np.ndarray, + T: np.ndarray, + v: np.ndarray, + rho_floor: float = 1e-14, + T_floor: float = 1e-12, +) -> np.ndarray: + rho_safe = np.maximum(rho, rho_floor) + T_safe = np.maximum(T, T_floor) + pref = rho_safe[:, None] / np.sqrt(2.0 * np.pi * T_safe[:, None]) + z = (v[None, :] - u[:, None]) ** 2 / (2.0 * T_safe[:, None]) + return pref * np.exp(-z) + + +# --------------------------------------------------------------------- +# Initial data +# --------------------------------------------------------------------- +def initial_condition( + x: np.ndarray, + v: np.ndarray, + profile: str = "mixed", +) -> np.ndarray: + """ + Build a positive, nontrivial initial condition. + """ + X = x[:, None] + V = v[None, :] + + if profile == "mixed": + rho = 1.0 + 0.2 * np.sin(2.0 * np.pi * X) + u = 0.5 * np.cos(2.0 * np.pi * X) + T = 0.8 + 0.1 * np.sin(4.0 * np.pi * X) + M1 = rho / np.sqrt(2.0 * np.pi * T) * np.exp(-((V - u) ** 2) / (2.0 * T)) + M2 = 0.35 * rho / np.sqrt(2.0 * np.pi * T) * np.exp(-((V + 0.7 * u) ** 2) / (2.0 * (1.25 * T))) + f0 = 0.75 * M1 + 0.25 * M2 + return np.maximum(f0, 1e-14) + + if profile == "simple": + rho = 1.0 + 0.1 * np.sin(2.0 * np.pi * X) + u = 0.25 * np.cos(2.0 * np.pi * X) + T = 1.0 + 0.05 * np.sin(2.0 * np.pi * X) + f0 = rho / np.sqrt(2.0 * np.pi * T) * np.exp(-((V - u) ** 2) / (2.0 * T)) + return np.maximum(f0, 1e-14) + + raise ValueError("profile must be 'mixed' or 'simple'") + + +# --------------------------------------------------------------------- +# Problem builder +# --------------------------------------------------------------------- +@dataclass +class BGKPhase2Problem: + x: np.ndarray + v: np.ndarray + dx: float + dv: float + force: float + eps: float + + +def build_problem( + Nx: int, + Nv: int, + xmax: float, + vmax: float, + force: float, + eps: float, +) -> BGKPhase2Problem: + x = np.linspace(0.0, xmax, Nx, endpoint=False) + v = np.linspace(-vmax, vmax, Nv, endpoint=False) + dx = xmax / Nx + dv = 2.0 * vmax / Nv + return BGKPhase2Problem(x=x, v=v, dx=dx, dv=dv, force=force, eps=eps) + + +# --------------------------------------------------------------------- +# Analytic subflows for fractional_step +# --------------------------------------------------------------------- +def make_transport_flow(problem: BGKPhase2Problem) -> Callable[[float, float, np.ndarray], np.ndarray]: + v = problem.v + dx = problem.dx + + def flow(_t: float, h: float, f: np.ndarray) -> np.ndarray: + shift_cells = (v * h) / dx + return shift_along_x_periodic(f, shift_cells) + + return flow + + +def make_force_flow(problem: BGKPhase2Problem) -> Callable[[float, float, np.ndarray], np.ndarray]: + dv = problem.dv + F = problem.force + + def flow(_t: float, h: float, f: np.ndarray) -> np.ndarray: + shift_cells = (F * h) / dv + return shift_along_v_periodic(f, shift_cells) + + return flow + + +def make_bgk_relaxation_flow(problem: BGKPhase2Problem) -> Callable[[float, float, np.ndarray], np.ndarray]: + v = problem.v + dv = problem.dv + eps = problem.eps + + def flow(_t: float, h: float, f: np.ndarray) -> np.ndarray: + rho, u, T = compute_moments(f, v, dv) + M = maxwellian_from_moments(rho, u, T, v) + + # exact homogeneous BGK update; negative h is allowed algebraically, + # but may amplify and/or damage positivity + E = np.exp(-h / eps) + return E * f + (1.0 - E) * M + + return flow + + +# --------------------------------------------------------------------- +# Diagnostics +# --------------------------------------------------------------------- +@dataclass +class RunMetrics: + err: float + min_f_over_run: float + final_min_f: float + final_mass: float + mass_rel_drift: float + max_l1_ratio: float + positive: bool + stable: bool + + +def total_mass(f: np.ndarray, dx: float, dv: float) -> float: + return float(np.sum(f) * dx * dv) + + +def l1_norm(f: np.ndarray, dx: float, dv: float) -> float: + return float(np.sum(np.abs(f)) * dx * dv) + + +def run_method( + method_name: str, + dt: float, + T: float, + f0: np.ndarray, + problem: BGKPhase2Problem, + reference: np.ndarray, + positivity_tol: float = -1e-12, + explosion_factor: float = 1e8, +) -> RunMetrics: + flows = [ + make_transport_flow(problem), + make_bgk_relaxation_flow(problem), + make_force_flow(problem), + ] + + steps = max(1, int(round(T / dt))) + dt = T / steps + + f = np.array(f0, dtype=float, copy=True) + t = 0.0 + + m0 = total_mass(f0, problem.dx, problem.dv) + l10 = l1_norm(f0, problem.dx, problem.dv) + + min_f_over_run = float(np.min(f)) + max_l1_ratio = 1.0 + stable = True + + for _ in range(steps): + f_new = fractional_step( + functions=flows, + delta_t=dt, + initial_y=f, + initial_t=t, + final_t=t + dt, + alpha=method_name, + methods={(1,): "ANALYTIC", (2,): "ANALYTIC", (3,): "ANALYTIC"}, + ) + + f_new = np.array(f_new, dtype=float, copy=False) + + if not np.all(np.isfinite(f_new)): + stable = False + f = f_new + break + + min_f_over_run = min(min_f_over_run, float(np.min(f_new))) + + l1r = l1_norm(f_new, problem.dx, problem.dv) / max(l10, 1e-30) + max_l1_ratio = max(max_l1_ratio, l1r) + + if l1r > explosion_factor: + stable = False + f = f_new + break + + f = f_new + t += dt + + final_min_f = float(np.min(f)) if np.all(np.isfinite(f)) else float("nan") + final_mass = total_mass(f, problem.dx, problem.dv) if np.all(np.isfinite(f)) else float("nan") + mass_rel_drift = abs(final_mass - m0) / max(abs(m0), 1e-30) if np.isfinite(final_mass) else float("inf") + positive = bool(final_min_f >= positivity_tol) if np.isfinite(final_min_f) else False + + if stable and np.all(np.isfinite(f)): + err = l1_norm(f - reference, problem.dx, problem.dv) + else: + err = float("inf") + + return RunMetrics( + err=err, + min_f_over_run=min_f_over_run, + final_min_f=final_min_f, + final_mass=final_mass, + mass_rel_drift=mass_rel_drift, + max_l1_ratio=max_l1_ratio, + positive=positive, + stable=stable, + ) + + +def compute_reference( + ref_method: str, + ref_dt: float, + T: float, + f0: np.ndarray, + problem: BGKPhase2Problem, +) -> np.ndarray: + flows = [ + make_transport_flow(problem), + make_bgk_relaxation_flow(problem), + make_force_flow(problem), + ] + f_ref = fractional_step( + functions=flows, + delta_t=ref_dt, + initial_y=f0, + initial_t=0.0, + final_t=T, + alpha=ref_method, + methods={(1,): "ANALYTIC", (2,): "ANALYTIC", (3,): "ANALYTIC"}, + ) + return np.array(f_ref, dtype=float) + + +# --------------------------------------------------------------------- +# Main study +# --------------------------------------------------------------------- +def run_regime( + eps: float, + T: float, + dts: Sequence[float], + methods: Sequence[str], + problem_base: Dict[str, float], + profile: str, + ref_method: str, + ref_refine: int, + outdir: Path, +) -> Path: + problem = build_problem( + Nx=int(problem_base["Nx"]), + Nv=int(problem_base["Nv"]), + xmax=float(problem_base["xmax"]), + vmax=float(problem_base["vmax"]), + force=float(problem_base["force"]), + eps=eps, + ) + + f0 = initial_condition(problem.x, problem.v, profile=profile) + + dt_min = min(dts) + ref_dt = dt_min / ref_refine + reference = compute_reference(ref_method, ref_dt, T, f0, problem) + + outdir.mkdir(parents=True, exist_ok=True) + eps_tag = str(eps).replace(".", "p").replace("-", "m") + csv_path = outdir / f"bgk_three_operator_eps{eps_tag}.csv" + + rows: List[Dict[str, object]] = [] + err_by_method: Dict[str, List[float]] = {m: [] for m in methods} + + print("\n" + "=" * 88) + print(f"BGK-INSPIRED THREE-OPERATOR STUDY | eps = {eps:.3e} | T = {T:g}") + print("=" * 88) + print(f"Reference method = {ref_method} | ref_dt = {ref_dt:.4e}") + + for method_name in methods: + print(f"\nMethod: {method_name}") + for dt in dts: + metrics = run_method( + method_name=method_name, + dt=dt, + T=T, + f0=f0, + problem=problem, + reference=reference, + ) + err_by_method[method_name].append(metrics.err) + + print( + f" dt={dt:10.4e} | " + f"err={metrics.err:12.5e} | " + f"stable={metrics.stable!s:5s} | " + f"positive={metrics.positive!s:5s} | " + f"min_f_run={metrics.min_f_over_run:12.5e} | " + f"mass_drift={metrics.mass_rel_drift:10.4e}" + ) + + rows.append( + { + "eps": eps, + "T": T, + "method": method_name, + "dt": dt, + "ref_method": ref_method, + "ref_dt": ref_dt, + "Nx": int(problem_base["Nx"]), + "Nv": int(problem_base["Nv"]), + "xmax": float(problem_base["xmax"]), + "vmax": float(problem_base["vmax"]), + "force": float(problem_base["force"]), + "profile": profile, + "err_l1": metrics.err, + "stable": int(metrics.stable), + "positive": int(metrics.positive), + "min_f_over_run": metrics.min_f_over_run, + "final_min_f": metrics.final_min_f, + "final_mass": metrics.final_mass, + "mass_rel_drift": metrics.mass_rel_drift, + "max_l1_ratio": metrics.max_l1_ratio, + } + ) + + p_by_method = {m: estimate_order(err_by_method[m], dts) for m in methods} + + print("\nObserved order fits:") + for m in methods: + print(f" {m:20s} -> p ≈ {p_by_method[m]:.4f}") + + for row in rows: + row["observed_order_fit"] = p_by_method[str(row["method"])] + + with csv_path.open("w", newline="") as f: + writer = csv.DictWriter( + f, + fieldnames=[ + "eps", + "T", + "method", + "dt", + "ref_method", + "ref_dt", + "Nx", + "Nv", + "xmax", + "vmax", + "force", + "profile", + "err_l1", + "stable", + "positive", + "min_f_over_run", + "final_min_f", + "final_mass", + "mass_rel_drift", + "max_l1_ratio", + "observed_order_fit", + ], + ) + writer.writeheader() + writer.writerows(rows) + + print(f"\nSaved: {csv_path}") + return csv_path + + +# --------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------- +def build_parser() -> argparse.ArgumentParser: + ap = argparse.ArgumentParser(description="BGK-inspired three-operator higher-order study.") + + ap.add_argument("--eps-list", type=str, default="1e-1,1e-2,1e-3") + ap.add_argument("--dts", type=str, default="1e-2,5e-3,2.5e-3") + ap.add_argument( + "--methods", + type=str, + default="Strang-3,OS32_7op_minLEM-3,PP3_4A-3,Yoshida-3", + ) + ap.add_argument("--reference-method", type=str, default="Yoshida-3") + ap.add_argument("--reference-refine", type=int, default=8) + ap.add_argument("--T", type=float, default=0.2) + + ap.add_argument("--Nx", type=int, default=40) + ap.add_argument("--Nv", type=int, default=80) + ap.add_argument("--xmax", type=float, default=1.0) + ap.add_argument("--vmax", type=float, default=8.0) + ap.add_argument("--force", type=float, default=0.5) + + ap.add_argument("--profile", type=str, choices=["mixed", "simple"], default="mixed") + ap.add_argument("--outdir", type=str, default="experiments/bgk_three_operator_outputs") + + return ap + + +def main() -> None: + ap = build_parser() + args = ap.parse_args() + + eps_list = parse_csv_floats(args.eps_list) + dts = parse_csv_floats(args.dts) + methods = parse_csv_strings(args.methods) + outdir = REPO / args.outdir + + problem_base = { + "Nx": args.Nx, + "Nv": args.Nv, + "xmax": args.xmax, + "vmax": args.vmax, + "force": args.force, + } + + print("Running BGK-inspired three-operator study with:") + print(f" eps_list = {eps_list}") + print(f" dts = {dts}") + print(f" methods = {methods}") + print(f" reference_method = {args.reference_method}") + print(f" reference_refine = {args.reference_refine}") + print(f" T = {args.T}") + print(f" Nx, Nv = {args.Nx}, {args.Nv}") + print(f" xmax, vmax = {args.xmax}, {args.vmax}") + print(f" force = {args.force}") + print(f" profile = {args.profile}") + print(f" outdir = {outdir}") + + csvs: List[Path] = [] + for eps in eps_list: + csvs.append( + run_regime( + eps=eps, + T=args.T, + dts=dts, + methods=methods, + problem_base=problem_base, + profile=args.profile, + ref_method=args.reference_method, + ref_refine=args.reference_refine, + outdir=outdir, + ) + ) + + print("\nCompleted CSVs:") + for c in csvs: + print(f" {c}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/experiments/bgk_three_operator_stability_scan.py b/experiments/bgk_three_operator_stability_scan.py new file mode 100644 index 0000000..f61f7f0 --- /dev/null +++ b/experiments/bgk_three_operator_stability_scan.py @@ -0,0 +1,641 @@ +""" +bgk_three_operator_stability_scan.py + +Stability / positivity scan for the BGK-inspired three-operator experiment. + +Purpose +------- +This script refines the Phase 2 study by scanning a list of dt values for a +fixed stiff regime, typically eps = 1e-3, and recording: + +- stable / unstable +- positive / not positive +- minimum value of f over the run +- final minimum of f +- mass drift +- L1 error vs fine reference +- max L1 growth ratio + +Model +----- +We split + + f_t + v f_x + F f_v = ( M[f] - f ) / eps + +into three operators: + A: f_t = - v f_x + B: f_t = ( M[f] - f ) / eps + C: f_t = - F f_v + +using pythOS fractional_step(...) with built-in 3-operator alpha tables. + +Recommended use +--------------- +Start with: + eps = 1e-3 + methods = Strang-3, OS32_7op_minLEM-3, PP3_4A-3, Yoshida-3 + +Examples +-------- +python experiments/bgk_three_operator_stability_scan.py +python experiments/bgk_three_operator_stability_scan.py --eps 1e-3 --T 0.2 +python experiments/bgk_three_operator_stability_scan.py --dts 1e-2,8e-3,6e-3,5e-3,4e-3,3e-3,2.5e-3,2e-3,1.5e-3,1e-3 +""" + +from __future__ import annotations + +import argparse +import csv +import math +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Dict, List, Sequence + +import numpy as np + +# --------------------------------------------------------------------- +# Robust repo-root imports +# --------------------------------------------------------------------- +THIS = Path(__file__).resolve() +REPO = THIS.parents[1] +if str(REPO) not in sys.path: + sys.path.insert(0, str(REPO)) + +from fractional_step import fractional_step # noqa: E402 + + +# --------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------- +def parse_csv_floats(s: str) -> List[float]: + return [float(x.strip()) for x in s.split(",") if x.strip()] + + +def parse_csv_strings(s: str) -> List[str]: + return [x.strip() for x in s.split(",") if x.strip()] + + +def periodic_shift_1d(values: np.ndarray, shift_cells: float) -> np.ndarray: + """ + Periodic linear interpolation for a 1D array. + Positive shift_cells means evaluate values(x - shift). + """ + n = values.shape[0] + idx = np.arange(n, dtype=float) + src = (idx - shift_cells) % n + i0 = np.floor(src).astype(int) + i1 = (i0 + 1) % n + theta = src - i0 + return (1.0 - theta) * values[i0] + theta * values[i1] + + +def shift_along_x_periodic(f: np.ndarray, shift_cells_per_v: np.ndarray) -> np.ndarray: + Nx, Nv = f.shape + out = np.empty_like(f) + for j in range(Nv): + out[:, j] = periodic_shift_1d(f[:, j], shift_cells_per_v[j]) + return out + + +def shift_along_v_periodic(f: np.ndarray, shift_cells: float) -> np.ndarray: + Nx, Nv = f.shape + out = np.empty_like(f) + for i in range(Nx): + out[i, :] = periodic_shift_1d(f[i, :], shift_cells) + return out + + +# --------------------------------------------------------------------- +# Moments / Maxwellian +# --------------------------------------------------------------------- +def compute_moments( + f: np.ndarray, + v: np.ndarray, + dv: float, + rho_floor: float = 1e-14, + T_floor: float = 1e-12, +): + rho = np.sum(f, axis=1) * dv + rho_safe = np.maximum(rho, rho_floor) + + mom1 = np.sum(f * v[None, :], axis=1) * dv + u = mom1 / rho_safe + + centered = v[None, :] - u[:, None] + energy = np.sum(f * centered * centered, axis=1) * dv + T = np.maximum(energy / rho_safe, T_floor) + + return rho, u, T + + +def maxwellian_from_moments( + rho: np.ndarray, + u: np.ndarray, + T: np.ndarray, + v: np.ndarray, + rho_floor: float = 1e-14, + T_floor: float = 1e-12, +) -> np.ndarray: + rho_safe = np.maximum(rho, rho_floor) + T_safe = np.maximum(T, T_floor) + pref = rho_safe[:, None] / np.sqrt(2.0 * np.pi * T_safe[:, None]) + z = (v[None, :] - u[:, None]) ** 2 / (2.0 * T_safe[:, None]) + return pref * np.exp(-z) + + +# --------------------------------------------------------------------- +# Initial condition +# --------------------------------------------------------------------- +def initial_condition(x: np.ndarray, v: np.ndarray, profile: str = "mixed") -> np.ndarray: + X = x[:, None] + V = v[None, :] + + if profile == "mixed": + rho = 1.0 + 0.2 * np.sin(2.0 * np.pi * X) + u = 0.5 * np.cos(2.0 * np.pi * X) + T = 0.8 + 0.1 * np.sin(4.0 * np.pi * X) + + M1 = rho / np.sqrt(2.0 * np.pi * T) * np.exp(-((V - u) ** 2) / (2.0 * T)) + M2 = 0.35 * rho / np.sqrt(2.0 * np.pi * T) * np.exp(-((V + 0.7 * u) ** 2) / (2.0 * (1.25 * T))) + f0 = 0.75 * M1 + 0.25 * M2 + return np.maximum(f0, 1e-14) + + if profile == "simple": + rho = 1.0 + 0.1 * np.sin(2.0 * np.pi * X) + u = 0.25 * np.cos(2.0 * np.pi * X) + T = 1.0 + 0.05 * np.sin(2.0 * np.pi * X) + f0 = rho / np.sqrt(2.0 * np.pi * T) * np.exp(-((V - u) ** 2) / (2.0 * T)) + return np.maximum(f0, 1e-14) + + raise ValueError("profile must be 'mixed' or 'simple'") + + +# --------------------------------------------------------------------- +# Problem definition +# --------------------------------------------------------------------- +@dataclass +class BGKPhase2Problem: + x: np.ndarray + v: np.ndarray + dx: float + dv: float + force: float + eps: float + + +def build_problem( + Nx: int, + Nv: int, + xmax: float, + vmax: float, + force: float, + eps: float, +) -> BGKPhase2Problem: + x = np.linspace(0.0, xmax, Nx, endpoint=False) + v = np.linspace(-vmax, vmax, Nv, endpoint=False) + dx = xmax / Nx + dv = 2.0 * vmax / Nv + return BGKPhase2Problem(x=x, v=v, dx=dx, dv=dv, force=force, eps=eps) + + +# --------------------------------------------------------------------- +# Analytic subflows +# --------------------------------------------------------------------- +def make_transport_flow(problem: BGKPhase2Problem) -> Callable[[float, float, np.ndarray], np.ndarray]: + v = problem.v + dx = problem.dx + + def flow(_t: float, h: float, f: np.ndarray) -> np.ndarray: + shift_cells = (v * h) / dx + return shift_along_x_periodic(f, shift_cells) + + return flow + + +def make_force_flow(problem: BGKPhase2Problem) -> Callable[[float, float, np.ndarray], np.ndarray]: + dv = problem.dv + F = problem.force + + def flow(_t: float, h: float, f: np.ndarray) -> np.ndarray: + shift_cells = (F * h) / dv + return shift_along_v_periodic(f, shift_cells) + + return flow + + +def make_bgk_relaxation_flow(problem: BGKPhase2Problem) -> Callable[[float, float, np.ndarray], np.ndarray]: + v = problem.v + dv = problem.dv + eps = problem.eps + + def flow(_t: float, h: float, f: np.ndarray) -> np.ndarray: + rho, u, T = compute_moments(f, v, dv) + M = maxwellian_from_moments(rho, u, T, v) + E = np.exp(-h / eps) + return E * f + (1.0 - E) * M + + return flow + + +# --------------------------------------------------------------------- +# Diagnostics +# --------------------------------------------------------------------- +def total_mass(f: np.ndarray, dx: float, dv: float) -> float: + return float(np.sum(f) * dx * dv) + + +def l1_norm(f: np.ndarray, dx: float, dv: float) -> float: + return float(np.sum(np.abs(f)) * dx * dv) + + +@dataclass +class RunMetrics: + stable: bool + positive: bool + err_l1: float + min_f_over_run: float + final_min_f: float + final_mass: float + mass_rel_drift: float + max_l1_ratio: float + exploded_at_step: int + dt_used: float + steps: int + + +def run_method( + method_name: str, + dt: float, + T: float, + f0: np.ndarray, + problem: BGKPhase2Problem, + reference: np.ndarray, + positivity_tol: float = -1e-12, + explosion_factor: float = 1e8, +) -> RunMetrics: + flows = [ + make_transport_flow(problem), + make_bgk_relaxation_flow(problem), + make_force_flow(problem), + ] + + steps = max(1, int(round(T / dt))) + dt = T / steps + + f = np.array(f0, dtype=float, copy=True) + m0 = total_mass(f0, problem.dx, problem.dv) + l10 = l1_norm(f0, problem.dx, problem.dv) + + min_f_over_run = float(np.min(f)) + max_l1_ratio = 1.0 + stable = True + exploded_at_step = -1 + t = 0.0 + + for n in range(steps): + f_new = fractional_step( + functions=flows, + delta_t=dt, + initial_y=f, + initial_t=t, + final_t=t + dt, + alpha=method_name, + methods={(1,): "ANALYTIC", (2,): "ANALYTIC", (3,): "ANALYTIC"}, + ) + f_new = np.array(f_new, dtype=float, copy=False) + + if not np.all(np.isfinite(f_new)): + stable = False + exploded_at_step = n + 1 + f = f_new + break + + min_f_over_run = min(min_f_over_run, float(np.min(f_new))) + + l1r = l1_norm(f_new, problem.dx, problem.dv) / max(l10, 1e-30) + max_l1_ratio = max(max_l1_ratio, l1r) + + if l1r > explosion_factor: + stable = False + exploded_at_step = n + 1 + f = f_new + break + + f = f_new + t += dt + + if stable and np.all(np.isfinite(f)): + err_l1 = l1_norm(f - reference, problem.dx, problem.dv) + final_min_f = float(np.min(f)) + final_mass = total_mass(f, problem.dx, problem.dv) + else: + err_l1 = float("inf") + final_min_f = float("nan") + final_mass = float("nan") + + mass_rel_drift = abs(final_mass - m0) / max(abs(m0), 1e-30) if np.isfinite(final_mass) else float("inf") + positive = bool(np.isfinite(final_min_f) and final_min_f >= positivity_tol) + + return RunMetrics( + stable=stable, + positive=positive, + err_l1=err_l1, + min_f_over_run=min_f_over_run, + final_min_f=final_min_f, + final_mass=final_mass, + mass_rel_drift=mass_rel_drift, + max_l1_ratio=max_l1_ratio, + exploded_at_step=exploded_at_step, + dt_used=dt, + steps=steps, + ) + + +def compute_reference( + ref_method: str, + ref_dt: float, + T: float, + f0: np.ndarray, + problem: BGKPhase2Problem, +) -> np.ndarray: + flows = [ + make_transport_flow(problem), + make_bgk_relaxation_flow(problem), + make_force_flow(problem), + ] + f_ref = fractional_step( + functions=flows, + delta_t=ref_dt, + initial_y=f0, + initial_t=0.0, + final_t=T, + alpha=ref_method, + methods={(1,): "ANALYTIC", (2,): "ANALYTIC", (3,): "ANALYTIC"}, + ) + return np.array(f_ref, dtype=float) + + +# --------------------------------------------------------------------- +# Scan runner +# --------------------------------------------------------------------- +def run_scan( + eps: float, + T: float, + dts: Sequence[float], + methods: Sequence[str], + problem_base: Dict[str, float], + profile: str, + ref_method: str, + ref_refine: int, + outdir: Path, + positivity_tol: float, + explosion_factor: float, +) -> None: + problem = build_problem( + Nx=int(problem_base["Nx"]), + Nv=int(problem_base["Nv"]), + xmax=float(problem_base["xmax"]), + vmax=float(problem_base["vmax"]), + force=float(problem_base["force"]), + eps=eps, + ) + + f0 = initial_condition(problem.x, problem.v, profile=profile) + + dt_min = min(dts) + ref_dt = dt_min / ref_refine + reference = compute_reference(ref_method, ref_dt, T, f0, problem) + + outdir.mkdir(parents=True, exist_ok=True) + detail_path = outdir / f"bgk_three_operator_scan_eps{str(eps).replace('.', 'p').replace('-', 'm')}_detail.csv" + summary_path = outdir / f"bgk_three_operator_scan_eps{str(eps).replace('.', 'p').replace('-', 'm')}_summary.csv" + + detail_rows: List[Dict[str, object]] = [] + summary_rows: List[Dict[str, object]] = [] + + print("\n" + "=" * 92) + print(f"BGK THREE-OPERATOR STABILITY SCAN | eps = {eps:.3e} | T = {T:g}") + print("=" * 92) + print(f"Reference method = {ref_method} | ref_dt = {ref_dt:.4e}") + + for method_name in methods: + print(f"\nMethod: {method_name}") + + largest_stable_dt = float("nan") + largest_positive_dt = float("nan") + stable_count = 0 + positive_count = 0 + + for dt in dts: + metrics = run_method( + method_name=method_name, + dt=dt, + T=T, + f0=f0, + problem=problem, + reference=reference, + positivity_tol=positivity_tol, + explosion_factor=explosion_factor, + ) + + print( + f" dt_req={dt:10.4e} | " + f"dt_used={metrics.dt_used:10.4e} | " + f"stable={metrics.stable!s:5s} | " + f"positive={metrics.positive!s:5s} | " + f"err={metrics.err_l1:12.5e} | " + f"min_f_run={metrics.min_f_over_run:12.5e} | " + f"mass_drift={metrics.mass_rel_drift:10.4e}" + ) + + detail_rows.append( + { + "eps": eps, + "T": T, + "method": method_name, + "dt_requested": dt, + "dt_used": metrics.dt_used, + "steps": metrics.steps, + "ref_method": ref_method, + "ref_dt": ref_dt, + "Nx": int(problem_base["Nx"]), + "Nv": int(problem_base["Nv"]), + "xmax": float(problem_base["xmax"]), + "vmax": float(problem_base["vmax"]), + "force": float(problem_base["force"]), + "profile": profile, + "stable": int(metrics.stable), + "positive": int(metrics.positive), + "err_l1": metrics.err_l1, + "min_f_over_run": metrics.min_f_over_run, + "final_min_f": metrics.final_min_f, + "final_mass": metrics.final_mass, + "mass_rel_drift": metrics.mass_rel_drift, + "max_l1_ratio": metrics.max_l1_ratio, + "exploded_at_step": metrics.exploded_at_step, + } + ) + + if metrics.stable: + stable_count += 1 + if np.isnan(largest_stable_dt) or metrics.dt_used > largest_stable_dt: + largest_stable_dt = metrics.dt_used + + if metrics.stable and metrics.positive: + positive_count += 1 + if np.isnan(largest_positive_dt) or metrics.dt_used > largest_positive_dt: + largest_positive_dt = metrics.dt_used + + summary_rows.append( + { + "eps": eps, + "method": method_name, + "stable_count": stable_count, + "positive_count": positive_count, + "num_dts_tested": len(dts), + "largest_stable_dt": largest_stable_dt, + "largest_positive_dt": largest_positive_dt, + } + ) + + with detail_path.open("w", newline="") as f: + writer = csv.DictWriter( + f, + fieldnames=[ + "eps", + "T", + "method", + "dt_requested", + "dt_used", + "steps", + "ref_method", + "ref_dt", + "Nx", + "Nv", + "xmax", + "vmax", + "force", + "profile", + "stable", + "positive", + "err_l1", + "min_f_over_run", + "final_min_f", + "final_mass", + "mass_rel_drift", + "max_l1_ratio", + "exploded_at_step", + ], + ) + writer.writeheader() + writer.writerows(detail_rows) + + with summary_path.open("w", newline="") as f: + writer = csv.DictWriter( + f, + fieldnames=[ + "eps", + "method", + "stable_count", + "positive_count", + "num_dts_tested", + "largest_stable_dt", + "largest_positive_dt", + ], + ) + writer.writeheader() + writer.writerows(summary_rows) + + print("\nSummary:") + for row in summary_rows: + print( + f" {str(row['method']):20s} | " + f"largest_stable_dt={row['largest_stable_dt']} | " + f"largest_positive_dt={row['largest_positive_dt']}" + ) + + print(f"\nSaved detail CSV: {detail_path}") + print(f"Saved summary CSV: {summary_path}") + + +# --------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------- +def build_parser() -> argparse.ArgumentParser: + ap = argparse.ArgumentParser(description="BGK-inspired three-operator stability and positivity scan.") + + ap.add_argument("--eps", type=float, default=1e-3) + ap.add_argument("--dts", type=str, default="1e-2,8e-3,6e-3,5e-3,4e-3,3e-3,2.5e-3,2e-3,1.5e-3,1e-3") + ap.add_argument( + "--methods", + type=str, + default="Strang-3,OS32_7op_minLEM-3,PP3_4A-3,Yoshida-3", + ) + ap.add_argument("--reference-method", type=str, default="Yoshida-3") + ap.add_argument("--reference-refine", type=int, default=8) + ap.add_argument("--T", type=float, default=0.2) + + ap.add_argument("--Nx", type=int, default=40) + ap.add_argument("--Nv", type=int, default=80) + ap.add_argument("--xmax", type=float, default=1.0) + ap.add_argument("--vmax", type=float, default=8.0) + ap.add_argument("--force", type=float, default=0.5) + + ap.add_argument("--profile", type=str, choices=["mixed", "simple"], default="mixed") + ap.add_argument("--positivity-tol", type=float, default=-1e-12) + ap.add_argument("--explosion-factor", type=float, default=1e8) + ap.add_argument("--outdir", type=str, default="experiments/bgk_three_operator_scan_outputs") + + return ap + + +def main() -> None: + ap = build_parser() + args = ap.parse_args() + + dts = parse_csv_floats(args.dts) + methods = parse_csv_strings(args.methods) + + problem_base = { + "Nx": args.Nx, + "Nv": args.Nv, + "xmax": args.xmax, + "vmax": args.vmax, + "force": args.force, + } + + outdir = REPO / args.outdir + + print("Running BGK three-operator stability scan with:") + print(f" eps = {args.eps}") + print(f" dts = {dts}") + print(f" methods = {methods}") + print(f" reference_method = {args.reference_method}") + print(f" reference_refine = {args.reference_refine}") + print(f" T = {args.T}") + print(f" Nx, Nv = {args.Nx}, {args.Nv}") + print(f" xmax, vmax = {args.xmax}, {args.vmax}") + print(f" force = {args.force}") + print(f" profile = {args.profile}") + print(f" positivity_tol = {args.positivity_tol}") + print(f" explosion_factor = {args.explosion_factor}") + print(f" outdir = {outdir}") + + run_scan( + eps=args.eps, + T=args.T, + dts=dts, + methods=methods, + problem_base=problem_base, + profile=args.profile, + ref_method=args.reference_method, + ref_refine=args.reference_refine, + outdir=outdir, + positivity_tol=args.positivity_tol, + explosion_factor=args.explosion_factor, + ) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/experiments/hu_shu_adaptive_higher_order_study.py b/experiments/hu_shu_adaptive_higher_order_study.py new file mode 100644 index 0000000..898d294 --- /dev/null +++ b/experiments/hu_shu_adaptive_higher_order_study.py @@ -0,0 +1,776 @@ +""" +Adaptive higher-order study for the Hu–Shu BGK setting. + +This script is designed to live alongside the validated +experiments/hu_shu_expk2_bgk.py baseline rather than replace it. It uses the +exported helpers from that file and runtime signature inspection so that calls +remain compatible with the available function signatures. +""" + +from __future__ import annotations + +import argparse +import csv +import inspect +import os +import time +import traceback +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Tuple + +import numpy as np + +try: + from experiments.utils.adaptive_splitting import ( + AdaptiveOptions, + adaptive_integrate, + l2_norm, + ) +except ModuleNotFoundError: + import sys + THIS_DIR = os.path.dirname(os.path.abspath(__file__)) + REPO_ROOT = os.path.dirname(THIS_DIR) + if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + from experiments.utils.adaptive_splitting import ( + AdaptiveOptions, + adaptive_integrate, + l2_norm, + ) + +try: + import experiments.hu_shu_expk2_bgk as hs +except Exception as exc: + raise ImportError( + "Could not import experiments.hu_shu_expk2_bgk. " + "Make sure you are running from the repo root or with the repo on PYTHONPATH." + ) from exc + +try: + from fractional_step import fractional_step +except Exception: + fractional_step = None + +Array = np.ndarray + + +def ensure_dir(path: str) -> None: + os.makedirs(path, exist_ok=True) + + +def parse_float_list(s: str) -> List[float]: + return [float(x.strip()) for x in s.split(",") if x.strip()] + + +def append_summary_csv(path: str, row: Dict[str, Any]) -> None: + exists = os.path.exists(path) + with open(path, "a", newline="") as f: + writer = csv.DictWriter(f, fieldnames=list(row.keys())) + if not exists: + writer.writeheader() + writer.writerow(row) + + +def save_adaptive_history_csv(path: str, result) -> None: + with open(path, "w", newline="") as f: + writer = csv.writer(f) + writer.writerow([ + "attempt_index", + "t", + "dt_try", + "dt_suggested", + "accepted", + "err_est", + "cpu_seconds", + "min_f", + "rho_min", + "u_min", + "T_min", + ]) + for i, rec in enumerate(result.records): + extra = rec.extra or {} + writer.writerow([ + i, + rec.t, + rec.dt_try, + rec.dt_suggested, + int(rec.accepted), + rec.err_est, + rec.cpu_seconds, + extra.get("min_f", ""), + extra.get("rho_min", ""), + extra.get("u_min", ""), + extra.get("T_min", ""), + ]) + + +def call_with_known_kwargs(func: Callable, pool: Dict[str, Any]) -> Any: + sig = inspect.signature(func) + kwargs = {} + accepts_var_kw = False + for name, param in sig.parameters.items(): + if param.kind == inspect.Parameter.VAR_KEYWORD: + accepts_var_kw = True + elif param.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY): + if name in pool: + kwargs[name] = pool[name] + if accepts_var_kw: + return func(**pool) + return func(**kwargs) + + +def extract_attr(obj: Any, candidates: List[str], default: Any = None) -> Any: + for name in candidates: + if hasattr(obj, name): + return getattr(obj, name) + return default + + +@dataclass +class HuShuProblem: + case: str + Nx: int + Nv: int + xmax: float + vmax: float + eps0: float + transport: str + eta_model: str + init: str + + grid: Any = None + x: Optional[Array] = None + v: Optional[Array] = None + f0: Optional[Array] = None + eps_x: Optional[Array] = None + + def __post_init__(self) -> None: + if fractional_step is None: + raise ImportError("Could not import pythOS fractional_step. Make sure the pythOS environment is set correctly.") + + if not hasattr(hs, "make_grid"): + raise AttributeError("hu_shu_expk2_bgk.py does not expose make_grid.") + + self.grid = call_with_known_kwargs( + hs.make_grid, + { + "Nx": self.Nx, + "Nv": self.Nv, + "xmax": self.xmax, + "vmax": self.vmax, + "x_min": -self.xmax, + "x_max": self.xmax, + "v_min": -self.vmax, + "v_max": self.vmax, + }, + ) + + self.x = extract_attr(self.grid, ["x", "xc", "x_grid", "x_nodes"]) + self.v = extract_attr(self.grid, ["v", "vc", "v_grid", "v_nodes"]) + + if self.x is None or self.v is None: + self.x = np.linspace(-self.xmax, self.xmax, self.Nx, endpoint=False) + self.v = np.linspace(-self.vmax, self.vmax, self.Nv, endpoint=False) + + if not hasattr(hs, "initial_condition"): + raise AttributeError("hu_shu_expk2_bgk.py does not expose initial_condition.") + + self.f0 = call_with_known_kwargs( + hs.initial_condition, + { + "grid": self.grid, + "x": self.x, + "v": self.v, + "Nx": self.Nx, + "Nv": self.Nv, + "kind": self.init, + "init": self.init, + "init_name": self.init, + "case": self.case, + "eps0": self.eps0, + "eta_model": self.eta_model, + }, + ) + + if self.case == "mixed": + if hasattr(hs, "mixed_regime_eps"): + self.eps_x = call_with_known_kwargs( + hs.mixed_regime_eps, + { + "grid": self.grid, + "x": self.x, + "eps0": self.eps0, + "eps": self.eps0, + }, + ) + else: + self.eps_x = self.eps0 * np.ones_like(self.x) + else: + self.eps_x = self.eps0 * np.ones_like(self.x) + + def diagnostics(self, y: Array) -> Dict[str, Any]: + out: Dict[str, Any] = {"min_f": float(np.min(y))} + try: + rho, u, T = call_with_known_kwargs( + hs.moments, + { + "f": y, + "v": self.v, + "grid": self.grid, + }, + ) + out["rho_min"] = float(np.min(rho)) + out["u_min"] = float(np.min(u)) + out["T_min"] = float(np.min(T)) + except Exception: + pass + return out + + def macro_errors(self, y: Array, y_ref: Array) -> Tuple[float, float, float]: + try: + rho, u, T = call_with_known_kwargs(hs.moments, {"f": y, "v": self.v, "grid": self.grid}) + rho_ref, u_ref, T_ref = call_with_known_kwargs(hs.moments, {"f": y_ref, "v": self.v, "grid": self.grid}) + return ( + float(np.mean(np.abs(rho - rho_ref))), + float(np.mean(np.abs(u - u_ref))), + float(np.mean(np.abs(T - T_ref))), + ) + except Exception: + return (float("nan"), float("nan"), float("nan")) + + def full_rhs(self, y: Array, t: float) -> Array: + if not hasattr(hs, "full_rhs"): + raise AttributeError("hu_shu_expk2_bgk.py does not expose full_rhs.") + return call_with_known_kwargs( + hs.full_rhs, + { + "f": y, + "y": y, + "t": t, + "grid": self.grid, + "x": self.x, + "v": self.v, + "eps0": self.eps0, + "eps": self.eps0, + "eps_x": self.eps_x, + "case": self.case, + "transport": self.transport, + "transport_name": self.transport, + "eta_model": self.eta_model, + }, + ) + + def exprk2_step(self, y: Array, t: float, dt: float) -> Array: + if not hasattr(hs, "hu_shu_exprk2_step"): + raise AttributeError("hu_shu_expk2_bgk.py does not expose hu_shu_exprk2_step.") + return call_with_known_kwargs( + hs.hu_shu_exprk2_step, + { + "f": y, + "y": y, + "t": t, + "dt": dt, + "grid": self.grid, + "x": self.x, + "v": self.v, + "eps0": self.eps0, + "eps": self.eps0, + "eps_x": self.eps_x, + "case": self.case, + "transport": self.transport, + "transport_name": self.transport, + "eta_model": self.eta_model, + }, + ) + + def transport_rhs(self, y: Array) -> Array: + if not hasattr(hs, "get_transport_rhs"): + raise AttributeError("hu_shu_expk2_bgk.py does not expose get_transport_rhs.") + rhs_fun = call_with_known_kwargs( + hs.get_transport_rhs, + { + "name": self.transport, + "transport": self.transport, + "transport_name": self.transport, + "grid": self.grid, + }, + ) + return call_with_known_kwargs(rhs_fun, {"f": y, "grid": self.grid}) + + def transport_flow(self) -> Callable[[float, float, Array], Array]: + def _flow(t: float, dt: float, y: Array) -> Array: + rhs = self.transport_rhs(y) + return y + dt * rhs + return _flow + + def collision_flow(self) -> Callable[[float, float, Array], Array]: + if not hasattr(hs, "bgk_homogeneous_exact"): + raise AttributeError("hu_shu_expk2_bgk.py does not expose bgk_homogeneous_exact.") + if not hasattr(hs, "moments"): + raise AttributeError("hu_shu_expk2_bgk.py does not expose moments.") + if not hasattr(hs, "maxwellian"): + raise AttributeError("hu_shu_expk2_bgk.py does not expose maxwellian.") + + def _flow(t: float, dt: float, y: Array) -> Array: + # Build local equilibrium g from the current state y + rho, u, T = call_with_known_kwargs( + hs.moments, + { + "f": y, + "v": self.v, + "grid": self.grid, + }, + ) + + g = call_with_known_kwargs( + hs.maxwellian, + { + "rho": rho, + "u": u, + "T": T, + "v": self.v, + "grid": self.grid, + }, + ) + + return call_with_known_kwargs( + hs.bgk_homogeneous_exact, + { + "f": y, + "g": g, + "dt_collision": dt, + "dt": dt, + "grid": self.grid, + "x": self.x, + "v": self.v, + "eps0": self.eps0, + "eps": self.eps0, + "eps_x": self.eps_x, + "case": self.case, + "eta_model": self.eta_model, + }, + ) + + return _flow + + def null_third_flow(self) -> Callable[[float, float, Array], Array]: + def _flow(t: float, dt: float, y: Array) -> Array: + return np.array(y, copy=True) + return _flow + + +def rk4_step(rhs: Callable[[Array, float], Array], y: Array, t: float, dt: float) -> Array: + k1 = rhs(y, t) + k2 = rhs(y + 0.5 * dt * k1, t + 0.5 * dt) + k3 = rhs(y + 0.5 * dt * k2, t + 0.5 * dt) + k4 = rhs(y + dt * k3, t + dt) + return y + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4) + + +def integrate_reference(rhs: Callable[[Array, float], Array], y0: Array, T: float, dt_ref: float) -> Array: + y = np.array(y0, copy=True) + t = 0.0 + while t < T - 1.0e-15: + h = min(dt_ref, T - t) + y = rk4_step(rhs, y, t, h) + t += h + return y + + +def make_three_operator_stepper(problem: HuShuProblem, method_name: str) -> Tuple[Callable[[Array, float, float], Array], int]: + methods = {(1,): "ANALYTIC", (2,): "ANALYTIC", (3,): "ANALYTIC"} + flows = [problem.transport_flow(), problem.collision_flow(), problem.null_third_flow()] + order_guess = { + "Strang-3": 2, + "OS32_7op_minLEM-3": 2, + "PP3_4A-3": 3, + "Yoshida-3": 4, + }.get(method_name, 2) + + def _step(y: Array, t: float, dt: float) -> Array: + return fractional_step( + functions=flows, + delta_t=dt, + initial_y=y, + initial_t=t, + final_t=t + dt, + alpha=method_name, + methods=methods, + ) + + return _step, order_guess + + +def integrate_fixed_step(stepper: Callable[[Array, float, float], Array], y0: Array, T: float, dt: float) -> Tuple[Array, int, float]: + y = np.array(y0, copy=True) + t = 0.0 + n_steps = 0 + tic = time.perf_counter() + while t < T - 1.0e-15: + h = min(dt, T - t) + y = stepper(y, t, h) + t += h + n_steps += 1 + wall = time.perf_counter() - tic + return y, n_steps, wall + + +def run_one_adaptive_case(case: str, method: str, T: float, dt0: float, atol: float, rtol: float, + dt_ref: float, Nx: int, Nv: int, xmax: float, vmax: float, eps0: float, + transport: str, eta_model: str, init: str, outdir: str) -> Dict[str, Any]: + problem = HuShuProblem(case, Nx, Nv, xmax, vmax, eps0, transport, eta_model, init) + y0 = np.array(problem.f0, copy=True) + stepper, order = make_three_operator_stepper(problem, method) + + opts = AdaptiveOptions( + order=order, + atol=atol, + rtol=rtol, + dt_min=1.0e-12, + dt_max=max(dt0, T), + safety=0.9, + growth_max=2.0, + shrink_min=0.2, + max_reject=20, + use_weighted_rms=True, + ) + + result = adaptive_integrate( + stepper=stepper, + y0=y0, + t0=0.0, + tfinal=T, + dt0=dt0, + opts=opts, + diagnostics_fn=problem.diagnostics, + ) + + y_ref = integrate_reference(problem.full_rhs, y0, T, dt_ref) + f_l2 = l2_norm(result.final_state - y_ref) + rho_l1, u_l1, T_l1 = problem.macro_errors(result.final_state, y_ref) + diag = problem.diagnostics(result.final_state) + + history_path = os.path.join( + outdir, + f"hu_shu_history_mode-adaptive_method-{method}_case-{case}_eps0-{eps0:.0e}_atol-{atol:.0e}_rtol-{rtol:.0e}.csv", + ) + save_adaptive_history_csv(history_path, result) + + return { + "mode": "adaptive", + "case": case, + "method": method, + "T": T, + "dt0": dt0, + "fixed_dt": "", + "atol": atol, + "rtol": rtol, + "dt_ref": dt_ref, + "Nx": Nx, + "Nv": Nv, + "xmax": xmax, + "vmax": vmax, + "eps0": eps0, + "transport": transport, + "eta_model": eta_model, + "init": init, + "success": int(result.success), + "message": result.message, + "n_steps": result.n_accept, + "n_accept": result.n_accept, + "n_reject": result.n_reject, + "dt_min_used": result.dt_min_used, + "dt_max_used": result.dt_max_used, + "dt_avg_used": result.dt_avg_used, + "wall_seconds": result.wall_seconds, + "f_l2": f_l2, + "rho_l1": rho_l1, + "u_l1": u_l1, + "T_l1": T_l1, + "min_f": diag.get("min_f", float("nan")), + "rho_min": diag.get("rho_min", float("nan")), + "u_min": diag.get("u_min", float("nan")), + "T_min": diag.get("T_min", float("nan")), + } + + +def run_one_fixed_case(case: str, method: str, T: float, fixed_dt: float, dt_ref: float, + Nx: int, Nv: int, xmax: float, vmax: float, eps0: float, + transport: str, eta_model: str, init: str) -> Dict[str, Any]: + problem = HuShuProblem(case, Nx, Nv, xmax, vmax, eps0, transport, eta_model, init) + y0 = np.array(problem.f0, copy=True) + stepper, _ = make_three_operator_stepper(problem, method) + + y_final, n_steps, wall = integrate_fixed_step(stepper, y0, T, fixed_dt) + y_ref = integrate_reference(problem.full_rhs, y0, T, dt_ref) + f_l2 = l2_norm(y_final - y_ref) + rho_l1, u_l1, T_l1 = problem.macro_errors(y_final, y_ref) + diag = problem.diagnostics(y_final) + + return { + "mode": "fixed", + "case": case, + "method": method, + "T": T, + "dt0": "", + "fixed_dt": fixed_dt, + "atol": "", + "rtol": "", + "dt_ref": dt_ref, + "Nx": Nx, + "Nv": Nv, + "xmax": xmax, + "vmax": vmax, + "eps0": eps0, + "transport": transport, + "eta_model": eta_model, + "init": init, + "success": 1, + "message": "Fixed-step integration completed successfully.", + "n_steps": n_steps, + "n_accept": n_steps, + "n_reject": 0, + "dt_min_used": fixed_dt, + "dt_max_used": fixed_dt, + "dt_avg_used": fixed_dt, + "wall_seconds": wall, + "f_l2": f_l2, + "rho_l1": rho_l1, + "u_l1": u_l1, + "T_l1": T_l1, + "min_f": diag.get("min_f", float("nan")), + "rho_min": diag.get("rho_min", float("nan")), + "u_min": diag.get("u_min", float("nan")), + "T_min": diag.get("T_min", float("nan")), + } + + +def run_exprk2_baseline(case: str, T: float, fixed_dt: float, dt_ref: float, + Nx: int, Nv: int, xmax: float, vmax: float, eps0: float, + transport: str, eta_model: str, init: str) -> Dict[str, Any]: + problem = HuShuProblem(case, Nx, Nv, xmax, vmax, eps0, transport, eta_model, init) + y0 = np.array(problem.f0, copy=True) + + y_final, n_steps, wall = integrate_fixed_step(problem.exprk2_step, y0, T, fixed_dt) + y_ref = integrate_reference(problem.full_rhs, y0, T, dt_ref) + f_l2 = l2_norm(y_final - y_ref) + rho_l1, u_l1, T_l1 = problem.macro_errors(y_final, y_ref) + diag = problem.diagnostics(y_final) + + return { + "mode": "baseline", + "case": case, + "method": "ExpRK2", + "T": T, + "dt0": "", + "fixed_dt": fixed_dt, + "atol": "", + "rtol": "", + "dt_ref": dt_ref, + "Nx": Nx, + "Nv": Nv, + "xmax": xmax, + "vmax": vmax, + "eps0": eps0, + "transport": transport, + "eta_model": eta_model, + "init": init, + "success": 1, + "message": "ExpRK2 baseline completed successfully.", + "n_steps": n_steps, + "n_accept": n_steps, + "n_reject": 0, + "dt_min_used": fixed_dt, + "dt_max_used": fixed_dt, + "dt_avg_used": fixed_dt, + "wall_seconds": wall, + "f_l2": f_l2, + "rho_l1": rho_l1, + "u_l1": u_l1, + "T_l1": T_l1, + "min_f": diag.get("min_f", float("nan")), + "rho_min": diag.get("rho_min", float("nan")), + "u_min": diag.get("u_min", float("nan")), + "T_min": diag.get("T_min", float("nan")), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description="Phase 3: Adaptive higher-order Hu–Shu / BGK study.") + parser.add_argument("--mode", type=str, default="both", choices=["adaptive", "fixed", "baseline", "both"]) + parser.add_argument("--case", type=str, default="mixed", choices=["mixed", "uniform"]) + parser.add_argument("--methods", type=str, default="Strang-3,PP3_4A-3,Yoshida-3") + parser.add_argument("--T", type=float, default=0.5) + parser.add_argument("--dt0", type=float, default=1.0e-4) + parser.add_argument("--baseline-dt", type=float, default=3.0e-5) + parser.add_argument("--atols", type=str, default="1e-6,1e-8") + parser.add_argument("--rtols", type=str, default="1e-4,1e-6") + parser.add_argument("--fixed-dts", type=str, default="2.5e-4,1.25e-4,6.25e-5,3.125e-5") + parser.add_argument("--dt-ref", type=float, default=1.0e-5) + parser.add_argument("--Nx", type=int, default=40) + parser.add_argument("--Nv", type=int, default=150) + parser.add_argument("--xmax", type=float, default=2.0) + parser.add_argument("--vmax", type=float, default=15.0) + parser.add_argument("--eps0", type=float, default=1.0e-5) + parser.add_argument("--transport", type=str, default="muscl2") + parser.add_argument("--eta-model", type=str, default="constant1") + parser.add_argument("--init", type=str, default="mixed") + parser.add_argument("--outdir", type=str, default="experiments/outputs/hu_shu_adaptive_higher_order") + parser.add_argument("--overwrite-summary", action="store_true") + args = parser.parse_args() + + ensure_dir(args.outdir) + methods = [m.strip() for m in args.methods.split(",") if m.strip()] + atols = parse_float_list(args.atols) + rtols = parse_float_list(args.rtols) + fixed_dts = parse_float_list(args.fixed_dts) + + summary_path = os.path.join(args.outdir, "hu_shu_adaptive_higher_order_summary.csv") + if args.overwrite_summary and os.path.exists(summary_path): + os.remove(summary_path) + + print("=" * 72) + print("PHASE 3: HU–SHU ADAPTIVE HIGHER-ORDER STUDY") + print(f"mode = {args.mode}") + print(f"case = {args.case}") + print(f"methods = {methods}") + print(f"T = {args.T}") + print(f"dt0 = {args.dt0}") + print(f"baseline_dt = {args.baseline_dt}") + print(f"atols = {atols}") + print(f"rtols = {rtols}") + print(f"fixed_dts = {fixed_dts}") + print(f"dt_ref = {args.dt_ref}") + print(f"Nx, Nv = {args.Nx}, {args.Nv}") + print(f"xmax, vmax = {args.xmax}, {args.vmax}") + print(f"eps0 = {args.eps0}") + print(f"transport = {args.transport}") + print(f"eta_model = {args.eta_model}") + print(f"init = {args.init}") + print(f"outdir = {args.outdir}") + print("=" * 72) + + if args.mode in ("baseline", "both"): + print("\nBaseline: ExpRK2") + row = run_exprk2_baseline( + args.case, args.T, args.baseline_dt, args.dt_ref, + args.Nx, args.Nv, args.xmax, args.vmax, args.eps0, + args.transport, args.eta_model, args.init + ) + append_summary_csv(summary_path, row) + print( + f" steps={row['n_steps']} " + f"f_l2={row['f_l2']:.6e} " + f"rho_l1={row['rho_l1']:.6e} " + f"u_l1={row['u_l1']:.6e} " + f"T_l1={row['T_l1']:.6e} " + f"min(f)={row['min_f']:.6e} " + f"wall={row['wall_seconds']:.3f}s" + ) + + if args.mode in ("adaptive", "both"): + for method in methods: + for atol in atols: + for rtol in rtols: + print(f"\nAdaptive: {method}, atol={atol:.0e}, rtol={rtol:.0e}") + try: + row = run_one_adaptive_case( + args.case, method, args.T, args.dt0, atol, rtol, args.dt_ref, + args.Nx, args.Nv, args.xmax, args.vmax, args.eps0, + args.transport, args.eta_model, args.init, args.outdir + ) + except Exception as exc: + tb = traceback.format_exc() + row = { + "mode": "adaptive", + "case": args.case, + "method": method, + "T": args.T, + "dt0": args.dt0, + "fixed_dt": "", + "atol": atol, + "rtol": rtol, + "dt_ref": args.dt_ref, + "Nx": args.Nx, + "Nv": args.Nv, + "xmax": args.xmax, + "vmax": args.vmax, + "eps0": args.eps0, + "transport": args.transport, + "eta_model": args.eta_model, + "init": args.init, + "success": 0, + "message": tb, + "n_steps": 0, + "n_accept": 0, + "n_reject": 0, + "dt_min_used": float("nan"), + "dt_max_used": float("nan"), + "dt_avg_used": float("nan"), + "wall_seconds": 0.0, + "f_l2": float("inf"), + "rho_l1": float("inf"), + "u_l1": float("inf"), + "T_l1": float("inf"), + "min_f": float("nan"), + "rho_min": float("nan"), + "u_min": float("nan"), + "T_min": float("nan"), + } + append_summary_csv(summary_path, row) + if row["success"] == 0: + print(f" FAILED: {row['message']}") + else: + print( + f" success={row['success']} " + f"accept={row['n_accept']} reject={row['n_reject']} " + f"dt_avg={row['dt_avg_used']:.3e} " + f"f_l2={row['f_l2']:.6e} " + f"rho_l1={row['rho_l1']:.6e} " + f"u_l1={row['u_l1']:.6e} " + f"T_l1={row['T_l1']:.6e} " + f"min(f)={row['min_f']:.6e}" + ) + + if args.mode in ("fixed", "both"): + for method in methods: + for fixed_dt in fixed_dts: + print(f"\nFixed: {method}, dt={fixed_dt:.3e}") + try: + row = run_one_fixed_case( + args.case, method, args.T, fixed_dt, args.dt_ref, + args.Nx, args.Nv, args.xmax, args.vmax, args.eps0, + args.transport, args.eta_model, args.init + ) + except Exception as exc: + row = { + "mode": "fixed", "case": args.case, "method": method, "T": args.T, + "dt0": "", "fixed_dt": fixed_dt, "atol": "", "rtol": "", "dt_ref": args.dt_ref, + "Nx": args.Nx, "Nv": args.Nv, "xmax": args.xmax, "vmax": args.vmax, + "eps0": args.eps0, "transport": args.transport, "eta_model": args.eta_model, "init": args.init, + "success": 0, "message": repr(exc), "n_steps": 0, "n_accept": 0, "n_reject": 0, + "dt_min_used": fixed_dt, "dt_max_used": fixed_dt, "dt_avg_used": fixed_dt, + "wall_seconds": 0.0, "f_l2": float("inf"), "rho_l1": float("inf"), + "u_l1": float("inf"), "T_l1": float("inf"), "min_f": float("nan"), + "rho_min": float("nan"), "u_min": float("nan"), "T_min": float("nan"), + } + append_summary_csv(summary_path, row) + if row["success"] == 0: + print(f" FAILED: {row['message']}") + else: + print( + f" success={row['success']} " + f"steps={row['n_steps']} " + f"f_l2={row['f_l2']:.6e} " + f"rho_l1={row['rho_l1']:.6e} " + f"u_l1={row['u_l1']:.6e} " + f"T_l1={row['T_l1']:.6e} " + f"min(f)={row['min_f']:.6e}" + ) + + print("\nDone.") + print(f"Summary written to: {summary_path}") + + +if __name__ == "__main__": + main() diff --git a/experiments/hu_shu_expk2_bgk.py b/experiments/hu_shu_expk2_bgk.py new file mode 100644 index 0000000..5db1006 --- /dev/null +++ b/experiments/hu_shu_expk2_bgk.py @@ -0,0 +1,663 @@ +#!/usr/bin/env python3 +""" +hu_shu_expk2_bgk.py + +Reconstructed BGK experiment driver using the actual Hu–Shu ExpRK2 scheme: + + f^(0) = exp((a0 dt / eps) Q) f^n + f^(1) = exp((a1 dt / eps) Q) ( f^(0) + b1 dt T(f^(0)) ) + f^(2) = f^(1) + b2 dt T(f^(1)) + f^(n+1) = exp((a2 dt / eps) Q) [ w f^(2) + (1-w) exp(((1-a2) dt / eps) Q) f^n ] + +with the Hu–Shu coefficients: + w = 1/2, b1 = b2 = 1, a0 = a1 = a2 = 1/3 + +For BGK: + Q(f) = eta ( M[f] - f ) + +and the homogeneous solve is exact: + exp(s Q) g = exp(-eta s) g + (1 - exp(-eta s)) M[g]. + +Paper references: +- scheme form: Eq. (2.6) +- recommended coefficients: Eqs. (2.23) and (2.26) +- BGK homogeneous solve: Eq. (4.5) +- accuracy test initial data: Eqs. (7.1)–(7.3) +- mixed-regime epsilon(x): Eq. (7.6) + +Practical note: +- This script reconstructs the *time integrator* from the paper. +- The paper's spatial transport discretization used WENO5 + positivity limiter. + Here we keep simpler upwind / MUSCL2 options for experimentation. +""" + +import argparse +import math +import time +from dataclasses import dataclass +from typing import Callable, Dict, List, Tuple, Union + +import numpy as np + + +ArrayLike = np.ndarray + + +@dataclass +class Grid: + Nx: int + Nv: int + xmax: float + vmax: float + x: np.ndarray + v: np.ndarray + dx: float + dv: float + + +# ----------------------------------------------------------------------------- +# Grid / moments / Maxwellian +# ----------------------------------------------------------------------------- + +def make_grid(Nx: int, Nv: int, xmax: float = 2.0, vmax: float = 15.0) -> Grid: + dx = xmax / Nx + dv = 2.0 * vmax / Nv + x = (np.arange(Nx) + 0.5) * dx + v = -vmax + (np.arange(Nv) + 0.5) * dv + return Grid(Nx=Nx, Nv=Nv, xmax=xmax, vmax=vmax, x=x, v=v, dx=dx, dv=dv) + + +def moments(f: ArrayLike, grid: Grid, min_temperature: float = 1e-14) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + 1D velocity moments: + rho = ∫ f dv + u = (1/rho) ∫ v f dv + T = (1/rho) ∫ (v-u)^2 f dv + """ + v = grid.v[None, :] + rho = np.sum(f, axis=1) * grid.dv + rho_safe = np.maximum(rho, 1e-14) + + mom = np.sum(f * v, axis=1) * grid.dv + u = mom / rho_safe + + centered2 = (v - u[:, None]) ** 2 + T = np.sum(f * centered2, axis=1) * grid.dv / rho_safe + T = np.maximum(T, min_temperature) + return rho, u, T + + +def maxwellian(rho: np.ndarray, u: np.ndarray, T: np.ndarray, grid: Grid) -> ArrayLike: + """ + 1D Maxwellian: + M = rho / sqrt(2 pi T) * exp(-(v-u)^2 / (2T)) + """ + v = grid.v[None, :] + rho = rho[:, None] + u = u[:, None] + T = np.maximum(T[:, None], 1e-14) + + coef = rho / np.sqrt(2.0 * np.pi * T) + expo = np.exp(-0.5 * (v - u) ** 2 / T) + return coef * expo + + +def collision_frequency(rho: np.ndarray, T: np.ndarray, model: str = "constant1") -> np.ndarray: + """ + BGK collision frequency eta(rho, T). + + The paper allows eta to be any positive function of rho and T. + For a clean baseline, default to eta = 1. + + Options: + - constant1 : eta = 1 + - rho : eta = rho + - rho_sqrtT : eta = rho * sqrt(T) + """ + if model == "constant1": + return np.ones_like(rho) + if model == "rho": + return np.maximum(rho, 1e-14) + if model == "rho_sqrtT": + return np.maximum(rho * np.sqrt(np.maximum(T, 1e-14)), 1e-14) + raise ValueError(f"Unknown eta model: {model}") + + +# ----------------------------------------------------------------------------- +# Paper-aligned initial data / epsilon profiles +# ----------------------------------------------------------------------------- + +def initial_condition_accuracy(grid: Grid) -> ArrayLike: + """ + Paper accuracy-test initial data, Eqs. (7.1)-(7.2): + f(0,x,v) = 0.5 M_{rho,u,T} + 0.3 M_{rho,-0.5u,T} + rho = 1 + 0.2 sin(pi x), u = 1, T = 1 / (1 + 0.2 sin(pi x)) + """ + x = grid.x + rho = 1.0 + 0.2 * np.sin(np.pi * x) + u = np.ones_like(x) + T = 1.0 / (1.0 + 0.2 * np.sin(np.pi * x)) + + M1 = maxwellian(rho, u, T, grid) + M2 = maxwellian(rho, -0.5 * u, T, grid) + f0 = 0.5 * M1 + 0.3 * M2 + return np.maximum(f0, 1e-14) + + +def initial_condition_simple(grid: Grid) -> ArrayLike: + """ + A fallback smooth non-equilibrium initial condition. + """ + x = grid.x + rho = 1.0 + 0.1 * np.sin(2.0 * np.pi * x / grid.xmax) + u = 0.5 * np.sin(2.0 * np.pi * x / grid.xmax) + T = 1.0 + 0.1 * np.cos(2.0 * np.pi * x / grid.xmax) + M = maxwellian(rho, u, T, grid) + + v = grid.v[None, :] + perturb = 0.05 * np.sin(2.0 * np.pi * x / grid.xmax)[:, None] * (v ** 2 - 1.0) * np.exp(-0.1 * v ** 2) + f0 = M * (1.0 + perturb) + return np.maximum(f0, 1e-14) + + +def initial_condition(grid: Grid, kind: str) -> ArrayLike: + if kind in ("accuracy", "paper", "mixed"): + return initial_condition_accuracy(grid) + if kind == "simple": + return initial_condition_simple(grid) + raise ValueError(f"Unknown initial condition kind: {kind}") + + +def constant_eps(grid: Grid, eps_value: float) -> np.ndarray: + return np.full(grid.Nx, float(eps_value)) + + +def mixed_regime_eps(grid: Grid, eps0: float = 1e-5) -> np.ndarray: + """ + Paper mixed-regime epsilon, Eq. (7.6): + eps(x) = eps0 + tanh(1 - 11(x-1)) + tanh(1 + 11(x-1)) + """ + x = grid.x + return eps0 + np.tanh(1.0 - 11.0 * (x - 1.0)) + np.tanh(1.0 + 11.0 * (x - 1.0)) + + +# ----------------------------------------------------------------------------- +# Transport operators +# ----------------------------------------------------------------------------- + +def minmod(a: np.ndarray, b: np.ndarray) -> np.ndarray: + out = np.zeros_like(a) + mask = (a * b) > 0.0 + out[mask] = np.sign(a[mask]) * np.minimum(np.abs(a[mask]), np.abs(b[mask])) + return out + + +def mc_limiter(dl: np.ndarray, dr: np.ndarray) -> np.ndarray: + return minmod(0.5 * (dl + dr), minmod(2.0 * dl, 2.0 * dr)) + + +def transport_rhs_upwind(f: ArrayLike, grid: Grid) -> ArrayLike: + """ + First-order upwind discretization of T(f) = -v d_x f. + Periodic in x. + """ + rhs = np.zeros_like(f) + for j, vj in enumerate(grid.v): + q = f[:, j] + if vj >= 0.0: + dfdx = (q - np.roll(q, 1)) / grid.dx + else: + dfdx = (np.roll(q, -1) - q) / grid.dx + rhs[:, j] = -vj * dfdx + return rhs + + +def transport_rhs_muscl2(f: ArrayLike, grid: Grid) -> ArrayLike: + """ + Second-order MUSCL finite-volume discretization of T(f) = -v d_x f. + Periodic in x. + """ + rhs = np.zeros_like(f) + + for j, vj in enumerate(grid.v): + q = f[:, j] + + qm1 = np.roll(q, 1) + qp1 = np.roll(q, -1) + + dl = q - qm1 + dr = qp1 - q + slope = mc_limiter(dl, dr) + + qL = q + 0.5 * slope + qR = q - 0.5 * slope + + # interface i+1/2 + qL_ip = qL + qR_ip = np.roll(qR, -1) + + if vj >= 0.0: + flux_ip = vj * qL_ip + else: + flux_ip = vj * qR_ip + + flux_im = np.roll(flux_ip, 1) + rhs[:, j] = -(flux_ip - flux_im) / grid.dx + + return rhs + + +def get_transport_rhs(name: str) -> Callable[[ArrayLike, Grid], ArrayLike]: + if name == "upwind": + return transport_rhs_upwind + if name == "muscl2": + return transport_rhs_muscl2 + raise ValueError(f"Unknown transport method: {name}") + + +# ----------------------------------------------------------------------------- +# BGK pieces +# ----------------------------------------------------------------------------- + +def bgk_homogeneous_exact(g: ArrayLike, dt_collision: float, grid: Grid, eps_x: np.ndarray, eta_model: str) -> ArrayLike: + """ + Exact homogeneous BGK solve for: + d_t f = (1/eps) Q(f), Q(f) = eta (M[f] - f) + + Since rho,u,T are conserved by the homogeneous BGK flow, M[g] and eta[g] + stay constant over this substep, giving: + exp((dt/eps)Q) g = e^{-eta dt / eps} g + (1 - e^{-eta dt / eps}) M[g] + """ + rho, u, T = moments(g, grid) + M = maxwellian(rho, u, T, grid) + eta = collision_frequency(rho, T, model=eta_model) + alpha = np.exp(-eta * dt_collision / np.maximum(eps_x, 1e-14)) + return alpha[:, None] * g + (1.0 - alpha)[:, None] * M + + +def collision_rhs(f: ArrayLike, grid: Grid, eps_x: np.ndarray, eta_model: str) -> ArrayLike: + rho, u, T = moments(f, grid) + eta = collision_frequency(rho, T, model=eta_model) + M = maxwellian(rho, u, T, grid) + return (eta[:, None] / np.maximum(eps_x, 1e-14)[:, None]) * (M - f) + + +def full_rhs(f: ArrayLike, grid: Grid, eps_x: np.ndarray, transport_name: str, eta_model: str) -> ArrayLike: + T_rhs = get_transport_rhs(transport_name) + return T_rhs(f, grid) + collision_rhs(f, grid, eps_x, eta_model) + + +# ----------------------------------------------------------------------------- +# Hu–Shu ExpRK2 from paper Eq. (2.6) with Eqs. (2.23), (2.26) +# ----------------------------------------------------------------------------- + +def hu_shu_exprk2_step( + f: ArrayLike, + dt: float, + grid: Grid, + eps_x: np.ndarray, + transport_name: str, + eta_model: str = "constant1", +) -> ArrayLike: + """ + Reconstructed Hu–Shu ExpRK2 BGK step. + + Paper coefficients: + w = 1/2, b1 = b2 = 1, a0 = a1 = a2 = 1/3 + """ + a0 = 1.0 / 3.0 + a1 = 1.0 / 3.0 + a2 = 1.0 / 3.0 + b1 = 1.0 + b2 = 1.0 + w = 0.5 + + T_rhs = get_transport_rhs(transport_name) + + # f^(0) = exp((a0 dt / eps)Q) f^n + f0 = bgk_homogeneous_exact(f, a0 * dt, grid, eps_x, eta_model) + + # f^(1) = exp((a1 dt / eps)Q) ( f^(0) + b1 dt T(f^(0)) ) + g1 = f0 + b1 * dt * T_rhs(f0, grid) + f1 = bgk_homogeneous_exact(g1, a1 * dt, grid, eps_x, eta_model) + + # f^(2) = f^(1) + b2 dt T(f^(1)) + f2 = f1 + b2 * dt * T_rhs(f1, grid) + + # fn+1 = exp((a2 dt / eps)Q) [ w f^(2) + (1-w) exp(((1-a2)dt/eps)Q) f^n ] + tail = bgk_homogeneous_exact(f, (1.0 - a2) * dt, grid, eps_x, eta_model) + combo = w * f2 + (1.0 - w) * tail + fn1 = bgk_homogeneous_exact(combo, a2 * dt, grid, eps_x, eta_model) + + return fn1 + + +# ----------------------------------------------------------------------------- +# Explicit reference methods +# ----------------------------------------------------------------------------- + +def ssp_rk2_step( + f: ArrayLike, + dt: float, + grid: Grid, + eps_x: np.ndarray, + transport_name: str, + eta_model: str, +) -> ArrayLike: + rhs1 = full_rhs(f, grid, eps_x, transport_name, eta_model) + f1 = f + dt * rhs1 + rhs2 = full_rhs(f1, grid, eps_x, transport_name, eta_model) + return 0.5 * f + 0.5 * (f1 + dt * rhs2) + + +def rk4_step( + f: ArrayLike, + dt: float, + grid: Grid, + eps_x: np.ndarray, + transport_name: str, + eta_model: str, +) -> ArrayLike: + k1 = full_rhs(f, grid, eps_x, transport_name, eta_model) + k2 = full_rhs(f + 0.5 * dt * k1, grid, eps_x, transport_name, eta_model) + k3 = full_rhs(f + 0.5 * dt * k2, grid, eps_x, transport_name, eta_model) + k4 = full_rhs(f + dt * k3, grid, eps_x, transport_name, eta_model) + return f + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4) + + +def advance_one_step( + f: ArrayLike, + dt: float, + grid: Grid, + eps_x: np.ndarray, + method: str, + transport_name: str, + eta_model: str, +) -> ArrayLike: + if method == "exprk2": + return hu_shu_exprk2_step(f, dt, grid, eps_x, transport_name, eta_model) + if method == "ssp_rk2": + return ssp_rk2_step(f, dt, grid, eps_x, transport_name, eta_model) + if method == "rk4": + return rk4_step(f, dt, grid, eps_x, transport_name, eta_model) + raise ValueError(f"Unknown method: {method}") + + +def evolve( + f0: ArrayLike, + Tfinal: float, + dt: float, + grid: Grid, + eps_x: np.ndarray, + method: str, + transport_name: str, + eta_model: str, +) -> ArrayLike: + f = f0.copy() + t = 0.0 + while t < Tfinal - 1e-15: + dt_step = min(dt, Tfinal - t) + f = advance_one_step(f, dt_step, grid, eps_x, method, transport_name, eta_model) + t += dt_step + return f + + +# ----------------------------------------------------------------------------- +# Diagnostics +# ----------------------------------------------------------------------------- + +def l1_x(a: np.ndarray, b: np.ndarray, dx: float) -> float: + return float(np.sum(np.abs(a - b)) * dx) + + +def l2_xv(a: np.ndarray, b: np.ndarray, dx: float, dv: float) -> float: + return float(np.sqrt(np.sum((a - b) ** 2) * dx * dv)) + + +def compare_solution(f: ArrayLike, fref: ArrayLike, grid: Grid) -> Dict[str, float]: + rho, u, T = moments(f, grid) + rho_ref, u_ref, T_ref = moments(fref, grid) + + return { + "f_L2": l2_xv(f, fref, grid.dx, grid.dv), + "rho_L1": l1_x(rho, rho_ref, grid.dx), + "u_L1": l1_x(u, u_ref, grid.dx), + "T_L1": l1_x(T, T_ref, grid.dx), + "min_f": float(np.min(f)), + "max_f": float(np.max(f)), + } + + +def observed_orders(errors: List[float]) -> List[float]: + out = [np.nan] + for i in range(1, len(errors)): + if errors[i] <= 0.0 or errors[i - 1] <= 0.0: + out.append(np.nan) + else: + out.append(math.log(errors[i - 1] / errors[i], 2.0)) + return out + + +def estimate_transport_cfl_dt(grid: Grid) -> float: + vmax = np.max(np.abs(grid.v)) + return grid.dx / max(vmax, 1e-14) + + +def print_run_summary(method: str, transport: str, dt: float, elapsed: float, errs: Dict[str, float]) -> None: + print(f"method={method}, transport={transport}, dt={dt:.6e}") + print(f"runtime = {elapsed:.3f} s") + print(f"L2 f = {errs['f_L2']:.6e}") + print(f"L1 rho = {errs['rho_L1']:.6e}") + print(f"L1 u = {errs['u_L1']:.6e}") + print(f"L1 T = {errs['T_L1']:.6e}") + print(f"min(f) = {errs['min_f']:.6e}") + print(f"max(f) = {errs['max_f']:.6e}") + if errs["min_f"] < -1e-12: + print("WARNING: negative values detected.") + print("-" * 72) + + +# ----------------------------------------------------------------------------- +# Experiment cases +# ----------------------------------------------------------------------------- + +def run_timeorder(args: argparse.Namespace) -> None: + grid = make_grid(args.Nx, args.Nv, xmax=args.xmax, vmax=args.vmax) + f0 = initial_condition(grid, args.init) + eps_x = constant_eps(grid, args.eps) + + dt_list = [float(s.strip()) for s in args.dt_list.split(",")] + dt_ref = min(dt_list) / args.refine_factor + + cfl_dt = estimate_transport_cfl_dt(grid) + print("=" * 72) + print("TIME-ORDER STUDY") + print(f"reference method = {args.reference_method}") + print(f"reference dt = {dt_ref:.6e}") + print(f"eps = {args.eps:.3e}") + print(f"Tfinal = {args.T:.3f}") + print(f"heuristic dx/|v|max CFL dt ≈ {cfl_dt:.6e}") + print("=" * 72) + + t0 = time.perf_counter() + fref = evolve(f0, args.T, dt_ref, grid, eps_x, args.reference_method, args.transport, args.eta_model) + tref = time.perf_counter() - t0 + print(f"Reference done in {tref:.3f} s") + print("-" * 72) + + f_errors = [] + rho_errors = [] + u_errors = [] + T_errors = [] + + for dt in dt_list: + if dt > cfl_dt: + print(f"WARNING: dt={dt:.3e} exceeds heuristic transport CFL {cfl_dt:.3e}") + + t1 = time.perf_counter() + f = evolve(f0, args.T, dt, grid, eps_x, args.method, args.transport, args.eta_model) + elapsed = time.perf_counter() - t1 + + errs = compare_solution(f, fref, grid) + print_run_summary(args.method, args.transport, dt, elapsed, errs) + + f_errors.append(errs["f_L2"]) + rho_errors.append(errs["rho_L1"]) + u_errors.append(errs["u_L1"]) + T_errors.append(errs["T_L1"]) + + ord_f = observed_orders(f_errors) + ord_rho = observed_orders(rho_errors) + ord_u = observed_orders(u_errors) + ord_T = observed_orders(T_errors) + + print("Observed orders") + print(f"{'dt':>12} {'f_L2':>14} {'ord':>8} {'rho_L1':>14} {'ord':>8} {'u_L1':>14} {'ord':>8} {'T_L1':>14} {'ord':>8}") + for i, dt in enumerate(dt_list): + print( + f"{dt:12.4e} " + f"{f_errors[i]:14.6e} {ord_f[i]:8.3f} " + f"{rho_errors[i]:14.6e} {ord_rho[i]:8.3f} " + f"{u_errors[i]:14.6e} {ord_u[i]:8.3f} " + f"{T_errors[i]:14.6e} {ord_T[i]:8.3f}" + ) + + +def run_table(args: argparse.Namespace) -> None: + grid = make_grid(args.Nx, args.Nv, xmax=args.xmax, vmax=args.vmax) + f0 = initial_condition(grid, args.init) + eps_list = [float(s.strip()) for s in args.eps_list.split(",")] + + dt = args.dt + dt_ref = dt / args.refine_factor + + print("=" * 72) + print("EPSILON TABLE STUDY") + print(f"method = {args.method}") + print(f"reference method = {args.reference_method}") + print(f"dt = {dt:.3e}") + print(f"Tfinal = {args.T:.3f}") + print("=" * 72) + print(f"{'eps':>12} {'f_L2':>14} {'rho_L1':>14} {'u_L1':>14} {'T_L1':>14} {'min(f)':>14} {'runtime(s)':>12}") + + for eps_val in eps_list: + eps_x = constant_eps(grid, eps_val) + fref = evolve(f0, args.T, dt_ref, grid, eps_x, args.reference_method, args.transport, args.eta_model) + + t0 = time.perf_counter() + f = evolve(f0, args.T, dt, grid, eps_x, args.method, args.transport, args.eta_model) + elapsed = time.perf_counter() - t0 + + errs = compare_solution(f, fref, grid) + print( + f"{eps_val:12.4e} " + f"{errs['f_L2']:14.6e} " + f"{errs['rho_L1']:14.6e} " + f"{errs['u_L1']:14.6e} " + f"{errs['T_L1']:14.6e} " + f"{errs['min_f']:14.6e} " + f"{elapsed:12.3f}" + ) + + +def run_mixed(args: argparse.Namespace) -> None: + grid = make_grid(args.Nx, args.Nv, xmax=args.xmax, vmax=args.vmax) + f0 = initial_condition(grid, "mixed") + eps_x = mixed_regime_eps(grid, eps0=args.eps0) + + print("=" * 72) + print("MIXED-REGIME STUDY") + print(f"method = {args.method}") + print(f"reference method = {args.reference_method}") + print(f"transport = {args.transport}") + print(f"Tfinal = {args.T:.3f}") + print(f"eps0 = {args.eps0:.3e}") + print(f"dt = {args.dt:.3e}") + print("=" * 72) + + dt_ref = args.dt / args.refine_factor + t0 = time.perf_counter() + fref = evolve(f0, args.T, dt_ref, grid, eps_x, args.reference_method, args.transport, args.eta_model) + tref = time.perf_counter() - t0 + print(f"Reference done in {tref:.3f} s") + print("-" * 72) + + t1 = time.perf_counter() + f = evolve(f0, args.T, args.dt, grid, eps_x, args.method, args.transport, args.eta_model) + elapsed = time.perf_counter() - t1 + + errs = compare_solution(f, fref, grid) + print_run_summary(args.method, args.transport, args.dt, elapsed, errs) + + rho, u, T = moments(f, grid) + rho_ref, u_ref, T_ref = moments(fref, grid) + + print("Ranges and max-errors vs reference:") + print(f"rho range = [{rho.min():.6e}, {rho.max():.6e}]") + print(f"u range = [{u.min():.6e}, {u.max():.6e}]") + print(f"T range = [{T.min():.6e}, {T.max():.6e}]") + print(f"max |rho-rho_ref| = {np.max(np.abs(rho - rho_ref)):.6e}") + print(f"max |u-u_ref| = {np.max(np.abs(u - u_ref)):.6e}") + print(f"max |T-T_ref| = {np.max(np.abs(T - T_ref)):.6e}") + + +# ----------------------------------------------------------------------------- +# CLI +# ----------------------------------------------------------------------------- + +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description="Hu–Shu ExpRK2 BGK experiments reconstructed from the paper.") + + p.add_argument("--case", choices=["timeorder", "table", "mixed"], default="timeorder") + p.add_argument("--method", choices=["exprk2", "ssp_rk2", "rk4"], default="exprk2") + p.add_argument("--reference-method", choices=["exprk2", "ssp_rk2", "rk4"], default="exprk2") + p.add_argument("--transport", choices=["upwind", "muscl2"], default="muscl2") + p.add_argument("--eta-model", choices=["constant1", "rho", "rho_sqrtT"], default="constant1") + + p.add_argument("--Nx", type=int, default=160, help="Number of x cells.") + p.add_argument("--Nv", type=int, default=150, help="Number of v cells.") + p.add_argument("--xmax", type=float, default=2.0, help="Domain length in x.") + p.add_argument("--vmax", type=float, default=15.0, help="Velocity truncation.") + p.add_argument("--T", type=float, default=0.1, help="Final time.") + p.add_argument("--dt", type=float, default=2.5e-4, help="Main dt for table/mixed.") + p.add_argument("--dt-list", dest="dt_list", type=str, default="4e-4,2e-4,1e-4,5e-5", + help="Comma-separated dt list for timeorder.") + p.add_argument("--refine-factor", type=int, default=4, + help="Reference dt = min(dt_list)/refine_factor or dt/refine_factor.") + + p.add_argument("--eps", type=float, default=1e-2, help="Uniform epsilon for timeorder.") + p.add_argument("--eps-list", type=str, default="1e0,1e-1,1e-2,1e-3,1e-4", + help="Comma-separated epsilon list for table.") + p.add_argument("--eps0", type=float, default=1e-5, help="Minimum epsilon for mixed regime.") + p.add_argument("--init", choices=["accuracy", "paper", "mixed", "simple"], default="accuracy") + + return p + + +def main() -> None: + args = build_parser().parse_args() + + print("Running BGK experiment with:") + print(f" case = {args.case}") + print(f" method = {args.method}") + print(f" reference_method = {args.reference_method}") + print(f" transport = {args.transport}") + print(f" eta_model = {args.eta_model}") + print(f" Nx, Nv = {args.Nx}, {args.Nv}") + print(f" xmax, vmax = {args.xmax}, {args.vmax}") + print(f" T = {args.T}") + print(f" dt = {args.dt}") + print() + + if args.case == "timeorder": + run_timeorder(args) + elif args.case == "table": + run_table(args) + elif args.case == "mixed": + run_mixed(args) + else: + raise ValueError(f"Unknown case: {args.case}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/experiments/hu_shu_three_operator_mixed_scan.py b/experiments/hu_shu_three_operator_mixed_scan.py new file mode 100644 index 0000000..58d0ec6 --- /dev/null +++ b/experiments/hu_shu_three_operator_mixed_scan.py @@ -0,0 +1,334 @@ +#!/usr/bin/env python3 +""" +hu_shu_three_operator_mixed_scan.py + +Focused timestep scan for the mixed-regime Hu–Shu three-operator screening study. + +Purpose +------- +Starting from the mixed-regime failure observed at dt = 2.5e-4, this script +scans smaller timesteps to determine when each shortlisted built-in three- +operator method becomes usable again. + +We keep the same three-operator screening decomposition used in +hu_shu_three_operator_study.py: + A : transport in x + B : exact homogeneous BGK relaxation + C : identity / null operator + +Outputs +------- +For each method and dt: +- L2 error in f vs reference +- L1 errors in rho, u, T +- min(f), max(f) +- positivity flag +- finite/nonfinite flag +- mass drift vs reference +- runtime + +Examples +-------- +python experiments/hu_shu_three_operator_mixed_scan.py +python experiments/hu_shu_three_operator_mixed_scan.py --dts 2.5e-4,1.25e-4,6.25e-5,3.125e-5 +python experiments/hu_shu_three_operator_mixed_scan.py --reference-method Strang-3 --refine-factor 8 +""" + +from __future__ import annotations + +import argparse +import csv +import time +from pathlib import Path +from typing import Dict, List + +import numpy as np +import sys + +# --------------------------------------------------------------------- +# Robust repo-root imports +# --------------------------------------------------------------------- +THIS = Path(__file__).resolve() +REPO = THIS.parents[1] +if str(REPO) not in sys.path: + sys.path.insert(0, str(REPO)) + +from experiments.hu_shu_three_operator_study import ( + evolve_three_operator, + compare_solution, + total_mass, + positivity_flag, +) +from experiments.hu_shu_expk2_bgk import ( + make_grid, + initial_condition, + mixed_regime_eps, +) + +# --------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------- +def parse_csv_floats(s: str) -> List[float]: + return [float(x.strip()) for x in s.split(",") if x.strip()] + + +def parse_csv_strings(s: str) -> List[str]: + return [x.strip() for x in s.split(",") if x.strip()] + + +# --------------------------------------------------------------------- +# Main mixed scan +# --------------------------------------------------------------------- +def run_scan(args: argparse.Namespace) -> None: + grid = make_grid(args.Nx, args.Nv, xmax=args.xmax, vmax=args.vmax) + f0 = initial_condition(grid, "mixed") + eps_x = mixed_regime_eps(grid, eps0=args.eps0) + + methods = parse_csv_strings(args.methods) + dts = parse_csv_floats(args.dts) + + outdir = REPO / args.outdir + outdir.mkdir(parents=True, exist_ok=True) + csv_path = outdir / "hu_shu_three_operator_mixed_scan.csv" + + rows: List[Dict[str, object]] = [] + + dt_min = min(dts) + dt_ref = dt_min / args.refine_factor + + print("=" * 92) + print("HU–SHU THREE-OPERATOR MIXED-REGIME TIMESTEP SCAN") + print(f"methods = {methods}") + print(f"dts = {dts}") + print(f"reference method = {args.reference_method}") + print(f"reference dt = {dt_ref:.6e}") + print(f"Tfinal = {args.T:.6f}") + print(f"eps0 = {args.eps0:.3e}") + print("=" * 92) + + # One fine reference for the mixed regime + print("\nComputing reference solution...") + t_ref0 = time.perf_counter() + fref = evolve_three_operator( + f0=f0, + Tfinal=args.T, + dt=dt_ref, + grid=grid, + eps_x=eps_x, + alpha_name=args.reference_method, + eta_model=args.eta_model, + ) + ref_elapsed = time.perf_counter() - t_ref0 + ref_mass = total_mass(fref, grid) + print(f"Reference done in {ref_elapsed:.3f}s") + + # Scan all methods / dts + for method_name in methods: + print(f"\nMethod: {method_name}") + + for dt in dts: + t0 = time.perf_counter() + f = evolve_three_operator( + f0=f0, + Tfinal=args.T, + dt=dt, + grid=grid, + eps_x=eps_x, + alpha_name=method_name, + eta_model=args.eta_model, + ) + elapsed = time.perf_counter() - t0 + + finite = bool(np.all(np.isfinite(f))) + + if finite: + errs = compare_solution(f, fref, grid) + mass = total_mass(f, grid) + mass_drift = abs(mass - ref_mass) / max(abs(ref_mass), 1e-30) + positive = positivity_flag(errs["min_f"], tol=args.positivity_tol) + else: + errs = { + "f_L2": float("inf"), + "rho_L1": float("inf"), + "u_L1": float("inf"), + "T_L1": float("inf"), + "min_f": float("nan"), + "max_f": float("nan"), + } + mass = float("nan") + mass_drift = float("inf") + positive = False + + usable = bool( + finite + and positive + and np.isfinite(errs["f_L2"]) + and errs["f_L2"] <= args.usable_f_l2_threshold + ) + + print( + f" dt={dt:10.4e} | " + f"finite={finite!s:5s} | " + f"positive={positive!s:5s} | " + f"usable={usable!s:5s} | " + f"f_L2={errs['f_L2']:12.5e} | " + f"rho_L1={errs['rho_L1']:12.5e} | " + f"min(f)={errs['min_f']:12.5e} | " + f"runtime={elapsed:8.3f}s" + ) + + rows.append( + { + "case": "mixed_scan", + "method": method_name, + "dt": dt, + "reference_method": args.reference_method, + "reference_dt": dt_ref, + "T": args.T, + "Nx": args.Nx, + "Nv": args.Nv, + "eta_model": args.eta_model, + "eps0": args.eps0, + "finite": int(finite), + "positive": int(positive), + "usable": int(usable), + "f_L2": errs["f_L2"], + "rho_L1": errs["rho_L1"], + "u_L1": errs["u_L1"], + "T_L1": errs["T_L1"], + "min_f": errs["min_f"], + "max_f": errs["max_f"], + "final_mass": mass, + "mass_rel_drift_vs_ref": mass_drift, + "runtime_sec": elapsed, + } + ) + + with csv_path.open("w", newline="") as fcsv: + writer = csv.DictWriter( + fcsv, + fieldnames=[ + "case", + "method", + "dt", + "reference_method", + "reference_dt", + "T", + "Nx", + "Nv", + "eta_model", + "eps0", + "finite", + "positive", + "usable", + "f_L2", + "rho_L1", + "u_L1", + "T_L1", + "min_f", + "max_f", + "final_mass", + "mass_rel_drift_vs_ref", + "runtime_sec", + ], + ) + writer.writeheader() + writer.writerows(rows) + + print(f"\nSaved: {csv_path}") + + # Console summary + print("\nSummary by method:") + for method_name in methods: + subset = [r for r in rows if r["method"] == method_name] + largest_finite = None + largest_positive = None + largest_usable = None + + for r in subset: + dt = float(r["dt"]) + if int(r["finite"]) == 1: + largest_finite = dt if largest_finite is None else max(largest_finite, dt) + if int(r["positive"]) == 1: + largest_positive = dt if largest_positive is None else max(largest_positive, dt) + if int(r["usable"]) == 1: + largest_usable = dt if largest_usable is None else max(largest_usable, dt) + + print( + f" {method_name:20s} | " + f"largest_finite_dt={largest_finite} | " + f"largest_positive_dt={largest_positive} | " + f"largest_usable_dt={largest_usable}" + ) + + +# --------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------- +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser( + description="Mixed-regime timestep scan for Hu–Shu three-operator methods." + ) + + p.add_argument( + "--methods", + type=str, + default="Strang-3,OS32_7op_minLEM-3,PP3_4A-3,Yoshida-3", + help="Comma-separated pythOS alpha names.", + ) + p.add_argument( + "--reference-method", + type=str, + default="Yoshida-3", + help="Reference built-in alpha name.", + ) + p.add_argument( + "--dts", + type=str, + default="2.5e-4,1.25e-4,6.25e-5,3.125e-5", + help="Comma-separated dt values to scan.", + ) + + p.add_argument("--eta-model", choices=["constant1", "rho", "rho_sqrtT"], default="constant1") + + p.add_argument("--Nx", type=int, default=80) + p.add_argument("--Nv", type=int, default=120) + p.add_argument("--xmax", type=float, default=2.0) + p.add_argument("--vmax", type=float, default=15.0) + p.add_argument("--T", type=float, default=0.1) + p.add_argument("--refine-factor", type=int, default=4) + + p.add_argument("--eps0", type=float, default=1e-5) + p.add_argument("--positivity-tol", type=float, default=-1e-12) + p.add_argument( + "--usable-f-l2-threshold", + type=float, + default=1e-1, + help="Simple threshold used for a practical 'usable' flag.", + ) + p.add_argument("--outdir", type=str, default="experiments/hu_shu_three_operator_outputs") + + return p + + +def main() -> None: + args = build_parser().parse_args() + + print("Running Hu–Shu mixed-regime three-operator scan with:") + print(f" methods = {args.methods}") + print(f" reference_method = {args.reference_method}") + print(f" dts = {args.dts}") + print(f" eta_model = {args.eta_model}") + print(f" Nx, Nv = {args.Nx}, {args.Nv}") + print(f" xmax, vmax = {args.xmax}, {args.vmax}") + print(f" T = {args.T}") + print(f" refine_factor = {args.refine_factor}") + print(f" eps0 = {args.eps0}") + print(f" usable_f_l2_threshold= {args.usable_f_l2_threshold}") + print() + + run_scan(args) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/experiments/hu_shu_three_operator_study.py b/experiments/hu_shu_three_operator_study.py new file mode 100644 index 0000000..86cb0e7 --- /dev/null +++ b/experiments/hu_shu_three_operator_study.py @@ -0,0 +1,551 @@ +#!/usr/bin/env python3 +""" +hu_shu_three_operator_study.py + +Screening study for built-in pythOS three-operator splitting methods on the +Hu–Shu-style BGK problem. + +Purpose +------- +Compare shortlisted three-operator methods on the same BGK setting used by the +current Hu–Shu driver, but without overwriting the existing script. + +We use a three-operator split: + A : transport in x + B : exact homogeneous BGK relaxation + C : identity / null operator + +Why include C? +-------------- +The goal here is to test the built-in three-operator coefficient tables +(Strang-3, OS32_7op_minLEM-3, PP3_4A-3, Yoshida-3) on the Hu–Shu/BGK problem +with minimal additional modeling assumptions. Using a null third operator is a +clean screening step before attempting a more elaborate genuine three-way kinetic +decomposition. + +Outputs +------- +For each eps and dt, the script records: +- L2 error in f vs a fine reference +- L1 errors in rho, u, T +- min(f), max(f) +- positivity flag +- mass drift +- runtime + +Examples +-------- +python experiments/hu_shu_three_operator_study.py --case table +python experiments/hu_shu_three_operator_study.py --case mixed --dt 2.5e-4 +python experiments/hu_shu_three_operator_study.py --case table --eps-list 1e0,1e-1,1e-2,1e-3 --dt 2.5e-4 +""" + +from __future__ import annotations + +import argparse +import csv +import time +from pathlib import Path +from typing import Callable, Dict, List + +import numpy as np +import sys + +# --------------------------------------------------------------------- +# Repository-local imports +# --------------------------------------------------------------------- +THIS = Path(__file__).resolve() +REPO = THIS.parents[1] +if str(REPO) not in sys.path: + sys.path.insert(0, str(REPO)) + +from fractional_step import fractional_step + +# Reuse existing Hu–Shu BGK building blocks from the current driver. +from experiments.hu_shu_expk2_bgk import ( + Grid, + make_grid, + moments, + maxwellian, + collision_frequency, + initial_condition, + constant_eps, + mixed_regime_eps, + get_transport_rhs, + l1_x, + l2_xv, +) + + +ArrayLike = np.ndarray + + +# --------------------------------------------------------------------- +# Exact / analytic subflows +# --------------------------------------------------------------------- +def bgk_homogeneous_exact( + g: ArrayLike, + dt_collision: float, + grid: Grid, + eps_x: np.ndarray, + eta_model: str, +) -> ArrayLike: + """ + Exact homogeneous BGK solve: + d_t f = (eta/eps) (M[f] - f) + + Since rho,u,T are conserved under the homogeneous BGK flow, M[g] remains + fixed over the substep, giving: + f(h) = exp(-eta h / eps) g + (1 - exp(-eta h / eps)) M[g] + """ + rho, u, T = moments(g, grid) + M = maxwellian(rho, u, T, grid) + eta = collision_frequency(rho, T, model=eta_model) + alpha = np.exp(-eta * dt_collision / np.maximum(eps_x, 1e-14)) + return alpha[:, None] * g + (1.0 - alpha)[:, None] * M + + +def periodic_shift_1d(values: np.ndarray, shift_cells: float) -> np.ndarray: + """ + Periodic linear interpolation for a 1D array. + Positive shift_cells means evaluate values(x - shift). + """ + n = values.shape[0] + idx = np.arange(n, dtype=float) + src = (idx - shift_cells) % n + i0 = np.floor(src).astype(int) + i1 = (i0 + 1) % n + theta = src - i0 + return (1.0 - theta) * values[i0] + theta * values[i1] + + +def transport_exact_shift( + f: ArrayLike, + dt: float, + grid: Grid, +) -> ArrayLike: + """ + Analytic transport subflow for: + f_t + v f_x = 0 + i.e. + f(x, v, t + dt) = f(x - v dt, v, t) + using periodic interpolation in x. + """ + out = np.empty_like(f) + shift_cells = (grid.v * dt) / grid.dx + for j in range(grid.Nv): + out[:, j] = periodic_shift_1d(f[:, j], shift_cells[j]) + return out + + +def identity_flow( + f: ArrayLike, + dt: float, + grid: Grid, +) -> ArrayLike: + _ = dt, grid + return np.array(f, copy=True) + + +def make_A_flow(grid: Grid) -> Callable[[float, float, ArrayLike], ArrayLike]: + def flow(_t: float, h: float, f: ArrayLike) -> ArrayLike: + return transport_exact_shift(f, h, grid) + return flow + + +def make_B_flow(grid: Grid, eps_x: np.ndarray, eta_model: str) -> Callable[[float, float, ArrayLike], ArrayLike]: + def flow(_t: float, h: float, f: ArrayLike) -> ArrayLike: + return bgk_homogeneous_exact(f, h, grid, eps_x, eta_model) + return flow + + +def make_C_flow(grid: Grid) -> Callable[[float, float, ArrayLike], ArrayLike]: + def flow(_t: float, h: float, f: ArrayLike) -> ArrayLike: + return identity_flow(f, h, grid) + return flow + + +# --------------------------------------------------------------------- +# Diagnostics +# --------------------------------------------------------------------- +def total_mass(f: ArrayLike, grid: Grid) -> float: + return float(np.sum(f) * grid.dx * grid.dv) + + +def compare_solution(f: ArrayLike, fref: ArrayLike, grid: Grid) -> Dict[str, float]: + rho, u, T = moments(f, grid) + rho_ref, u_ref, T_ref = moments(fref, grid) + + return { + "f_L2": l2_xv(f, fref, grid.dx, grid.dv), + "rho_L1": l1_x(rho, rho_ref, grid.dx), + "u_L1": l1_x(u, u_ref, grid.dx), + "T_L1": l1_x(T, T_ref, grid.dx), + "min_f": float(np.min(f)), + "max_f": float(np.max(f)), + } + + +def positivity_flag(min_f: float, tol: float = -1e-12) -> bool: + return bool(min_f >= tol) + + +# --------------------------------------------------------------------- +# Time evolution +# --------------------------------------------------------------------- +def evolve_three_operator( + f0: ArrayLike, + Tfinal: float, + dt: float, + grid: Grid, + eps_x: np.ndarray, + alpha_name: str, + eta_model: str, +) -> ArrayLike: + f = np.array(f0, copy=True) + t = 0.0 + + flows = [ + make_A_flow(grid), + make_B_flow(grid, eps_x, eta_model), + make_C_flow(grid), + ] + + while t < Tfinal - 1e-15: + dt_step = min(dt, Tfinal - t) + f = fractional_step( + functions=flows, + delta_t=dt_step, + initial_y=f, + initial_t=t, + final_t=t + dt_step, + alpha=alpha_name, + methods={(1,): "ANALYTIC", (2,): "ANALYTIC", (3,): "ANALYTIC"}, + ) + f = np.array(f, dtype=float, copy=False) + t += dt_step + + return f + + +# --------------------------------------------------------------------- +# Experiment cases +# --------------------------------------------------------------------- +def run_table(args: argparse.Namespace) -> None: + grid = make_grid(args.Nx, args.Nv, xmax=args.xmax, vmax=args.vmax) + f0 = initial_condition(grid, args.init) + + eps_list = [float(s.strip()) for s in args.eps_list.split(",")] + methods = [s.strip() for s in args.methods.split(",")] + + dt = args.dt + dt_ref = dt / args.refine_factor + + outdir = REPO / args.outdir + outdir.mkdir(parents=True, exist_ok=True) + csv_path = outdir / "hu_shu_three_operator_table.csv" + + rows: List[Dict[str, object]] = [] + + print("=" * 84) + print("HU–SHU THREE-OPERATOR EPSILON TABLE STUDY") + print(f"methods = {methods}") + print(f"dt = {dt:.3e}") + print(f"reference method = {args.reference_method}") + print(f"reference dt = {dt_ref:.3e}") + print(f"Tfinal = {args.T:.3f}") + print("=" * 84) + + for eps_val in eps_list: + eps_x = constant_eps(grid, eps_val) + + fref = evolve_three_operator( + f0=f0, + Tfinal=args.T, + dt=dt_ref, + grid=grid, + eps_x=eps_x, + alpha_name=args.reference_method, + eta_model=args.eta_model, + ) + + ref_mass = total_mass(fref, grid) + + for method_name in methods: + t0 = time.perf_counter() + f = evolve_three_operator( + f0=f0, + Tfinal=args.T, + dt=dt, + grid=grid, + eps_x=eps_x, + alpha_name=method_name, + eta_model=args.eta_model, + ) + elapsed = time.perf_counter() - t0 + + errs = compare_solution(f, fref, grid) + mass = total_mass(f, grid) + mass_drift = abs(mass - ref_mass) / max(abs(ref_mass), 1e-30) + positive = positivity_flag(errs["min_f"], tol=args.positivity_tol) + + print( + f"eps={eps_val:10.3e} | " + f"method={method_name:18s} | " + f"f_L2={errs['f_L2']:10.4e} | " + f"rho_L1={errs['rho_L1']:10.4e} | " + f"min(f)={errs['min_f']:10.4e} | " + f"positive={positive!s:5s} | " + f"runtime={elapsed:8.3f}s" + ) + + rows.append( + { + "case": "table", + "eps": eps_val, + "method": method_name, + "dt": dt, + "reference_method": args.reference_method, + "reference_dt": dt_ref, + "T": args.T, + "Nx": args.Nx, + "Nv": args.Nv, + "eta_model": args.eta_model, + "init": args.init, + "f_L2": errs["f_L2"], + "rho_L1": errs["rho_L1"], + "u_L1": errs["u_L1"], + "T_L1": errs["T_L1"], + "min_f": errs["min_f"], + "max_f": errs["max_f"], + "positive": int(positive), + "mass_rel_drift_vs_ref": mass_drift, + "runtime_sec": elapsed, + } + ) + + with csv_path.open("w", newline="") as fcsv: + writer = csv.DictWriter( + fcsv, + fieldnames=[ + "case", + "eps", + "method", + "dt", + "reference_method", + "reference_dt", + "T", + "Nx", + "Nv", + "eta_model", + "init", + "f_L2", + "rho_L1", + "u_L1", + "T_L1", + "min_f", + "max_f", + "positive", + "mass_rel_drift_vs_ref", + "runtime_sec", + ], + ) + writer.writeheader() + writer.writerows(rows) + + print(f"\nSaved: {csv_path}") + + +def run_mixed(args: argparse.Namespace) -> None: + grid = make_grid(args.Nx, args.Nv, xmax=args.xmax, vmax=args.vmax) + f0 = initial_condition(grid, "mixed") + eps_x = mixed_regime_eps(grid, eps0=args.eps0) + methods = [s.strip() for s in args.methods.split(",")] + + dt = args.dt + dt_ref = dt / args.refine_factor + + outdir = REPO / args.outdir + outdir.mkdir(parents=True, exist_ok=True) + csv_path = outdir / "hu_shu_three_operator_mixed.csv" + + rows: List[Dict[str, object]] = [] + + print("=" * 84) + print("HU–SHU THREE-OPERATOR MIXED-REGIME STUDY") + print(f"methods = {methods}") + print(f"dt = {dt:.3e}") + print(f"reference method = {args.reference_method}") + print(f"reference dt = {dt_ref:.3e}") + print(f"Tfinal = {args.T:.3f}") + print(f"eps0 = {args.eps0:.3e}") + print("=" * 84) + + fref = evolve_three_operator( + f0=f0, + Tfinal=args.T, + dt=dt_ref, + grid=grid, + eps_x=eps_x, + alpha_name=args.reference_method, + eta_model=args.eta_model, + ) + + ref_mass = total_mass(fref, grid) + + for method_name in methods: + t0 = time.perf_counter() + f = evolve_three_operator( + f0=f0, + Tfinal=args.T, + dt=dt, + grid=grid, + eps_x=eps_x, + alpha_name=method_name, + eta_model=args.eta_model, + ) + elapsed = time.perf_counter() - t0 + + errs = compare_solution(f, fref, grid) + mass = total_mass(f, grid) + mass_drift = abs(mass - ref_mass) / max(abs(ref_mass), 1e-30) + positive = positivity_flag(errs["min_f"], tol=args.positivity_tol) + + rho, u, T = moments(f, grid) + rho_ref, u_ref, T_ref = moments(fref, grid) + + print( + f"method={method_name:18s} | " + f"f_L2={errs['f_L2']:10.4e} | " + f"rho_L1={errs['rho_L1']:10.4e} | " + f"min(f)={errs['min_f']:10.4e} | " + f"positive={positive!s:5s} | " + f"runtime={elapsed:8.3f}s" + ) + + rows.append( + { + "case": "mixed", + "method": method_name, + "dt": dt, + "reference_method": args.reference_method, + "reference_dt": dt_ref, + "T": args.T, + "Nx": args.Nx, + "Nv": args.Nv, + "eta_model": args.eta_model, + "eps0": args.eps0, + "f_L2": errs["f_L2"], + "rho_L1": errs["rho_L1"], + "u_L1": errs["u_L1"], + "T_L1": errs["T_L1"], + "min_f": errs["min_f"], + "max_f": errs["max_f"], + "positive": int(positive), + "mass_rel_drift_vs_ref": mass_drift, + "rho_max_abs_err": float(np.max(np.abs(rho - rho_ref))), + "u_max_abs_err": float(np.max(np.abs(u - u_ref))), + "T_max_abs_err": float(np.max(np.abs(T - T_ref))), + "runtime_sec": elapsed, + } + ) + + with csv_path.open("w", newline="") as fcsv: + writer = csv.DictWriter( + fcsv, + fieldnames=[ + "case", + "method", + "dt", + "reference_method", + "reference_dt", + "T", + "Nx", + "Nv", + "eta_model", + "eps0", + "f_L2", + "rho_L1", + "u_L1", + "T_L1", + "min_f", + "max_f", + "positive", + "mass_rel_drift_vs_ref", + "rho_max_abs_err", + "u_max_abs_err", + "T_max_abs_err", + "runtime_sec", + ], + ) + writer.writeheader() + writer.writerows(rows) + + print(f"\nSaved: {csv_path}") + + +# --------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------- +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description="Three-operator screening study on the Hu–Shu BGK problem.") + + p.add_argument("--case", choices=["table", "mixed"], default="table") + p.add_argument( + "--methods", + type=str, + default="Strang-3,OS32_7op_minLEM-3,PP3_4A-3,Yoshida-3", + help="Comma-separated pythOS alpha names.", + ) + p.add_argument( + "--reference-method", + type=str, + default="Yoshida-3", + help="Reference built-in alpha name.", + ) + + p.add_argument("--eta-model", choices=["constant1", "rho", "rho_sqrtT"], default="constant1") + + p.add_argument("--Nx", type=int, default=80) + p.add_argument("--Nv", type=int, default=120) + p.add_argument("--xmax", type=float, default=2.0) + p.add_argument("--vmax", type=float, default=15.0) + p.add_argument("--T", type=float, default=0.1) + p.add_argument("--dt", type=float, default=2.5e-4) + p.add_argument("--refine-factor", type=int, default=4) + + p.add_argument("--eps-list", type=str, default="1e0,1e-1,1e-2,1e-3") + p.add_argument("--eps0", type=float, default=1e-5) + p.add_argument("--init", choices=["accuracy", "paper", "mixed", "simple"], default="accuracy") + + p.add_argument("--positivity-tol", type=float, default=-1e-12) + p.add_argument("--outdir", type=str, default="experiments/hu_shu_three_operator_outputs") + + return p + + +def main() -> None: + args = build_parser().parse_args() + + print("Running Hu–Shu three-operator study with:") + print(f" case = {args.case}") + print(f" methods = {args.methods}") + print(f" reference_method = {args.reference_method}") + print(f" eta_model = {args.eta_model}") + print(f" Nx, Nv = {args.Nx}, {args.Nv}") + print(f" xmax, vmax = {args.xmax}, {args.vmax}") + print(f" T = {args.T}") + print(f" dt = {args.dt}") + print(f" refine_factor = {args.refine_factor}") + print() + + if args.case == "table": + run_table(args) + elif args.case == "mixed": + run_mixed(args) + else: + raise ValueError(f"Unknown case: {args.case}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/experiments/utils/adaptive_splitting.py b/experiments/utils/adaptive_splitting.py new file mode 100644 index 0000000..3cdd110 --- /dev/null +++ b/experiments/utils/adaptive_splitting.py @@ -0,0 +1,358 @@ +""" +Reusable adaptive driver for operator splitting methods. + +The core idea is step-doubling: + +- One full step of size h +- Two half-steps of size h/2 +- Use the difference as an error estimate +- Accept/reject the step and adapt h + +This is intentionally generic so it can wrap ABC, BGK, or other split steppers. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Callable, Dict, List, Optional, Tuple, Any + +import math +import time +import numpy as np + + +ArrayLike = Any + + +def _to_numpy(x: ArrayLike) -> np.ndarray: + """ + Convert state to numpy array for error calculations. + + Supports: + - numpy arrays + - scalars + - lists/tuples + """ + if isinstance(x, np.ndarray): + return x + return np.asarray(x, dtype=float) + + +def l2_norm(x: ArrayLike) -> float: + arr = _to_numpy(x) + return float(np.sqrt(np.sum(arr * arr))) + + +def linf_norm(x: ArrayLike) -> float: + arr = _to_numpy(x) + return float(np.max(np.abs(arr))) + + +def weighted_rms_error( + y_high: ArrayLike, + y_low: ArrayLike, + atol: float, + rtol: float, +) -> float: + """ + Weighted RMS error often used in adaptive time stepping. + + err = sqrt( mean( ((y_high - y_low) / (atol + rtol*max(|y_high|,|y_low|)))^2 ) ) + + Returns a dimensionless scalar. Accept if err <= 1. + """ + yh = _to_numpy(y_high) + yl = _to_numpy(y_low) + + scale = atol + rtol * np.maximum(np.abs(yh), np.abs(yl)) + diff = (yh - yl) / scale + return float(np.sqrt(np.mean(diff * diff))) + + +@dataclass +class AdaptiveOptions: + """ + Configuration for adaptive step control. + """ + order: int + atol: float = 1.0e-8 + rtol: float = 1.0e-6 + dt_min: float = 1.0e-12 + dt_max: float = 1.0 + safety: float = 0.9 + growth_max: float = 2.0 + shrink_min: float = 0.2 + max_reject: int = 20 + use_weighted_rms: bool = True + accept_on_nan_error: bool = False + + +@dataclass +class StepRecord: + """ + One accepted or rejected attempt. + """ + t: float + dt_try: float + dt_suggested: float + accepted: bool + err_est: float + cpu_seconds: float + extra: Dict[str, Any] = field(default_factory=dict) + + +@dataclass +class AdaptiveResult: + """ + Full adaptive integration result. + """ + t_values: List[float] + y_values: List[ArrayLike] + records: List[StepRecord] + n_accept: int + n_reject: int + wall_seconds: float + success: bool + message: str + + @property + def final_time(self) -> float: + return self.t_values[-1] + + @property + def final_state(self) -> ArrayLike: + return self.y_values[-1] + + @property + def dt_min_used(self) -> float: + if not self.records: + return 0.0 + return min(r.dt_try for r in self.records if r.accepted) + + @property + def dt_max_used(self) -> float: + if not self.records: + return 0.0 + return max(r.dt_try for r in self.records if r.accepted) + + @property + def dt_avg_used(self) -> float: + accepted = [r.dt_try for r in self.records if r.accepted] + if not accepted: + return 0.0 + return float(sum(accepted) / len(accepted)) + + +def _default_error_estimate( + y_two_half: ArrayLike, + y_full: ArrayLike, + atol: float, + rtol: float, + use_weighted_rms: bool, + error_norm_fn: Optional[Callable[[ArrayLike], float]] = None, +) -> float: + """ + Returns a dimensionless error estimate. + + If weighted RMS is used, accept if err <= 1. + Otherwise, uses user-supplied norm on the raw difference and scales by atol+rtol. + """ + if use_weighted_rms: + return weighted_rms_error(y_two_half, y_full, atol=atol, rtol=rtol) + + diff = _to_numpy(y_two_half) - _to_numpy(y_full) + if error_norm_fn is None: + err_raw = l2_norm(diff) + else: + err_raw = float(error_norm_fn(diff)) + + scale = atol + rtol * max(l2_norm(y_two_half), l2_norm(y_full), 1.0) + return err_raw / scale + + +def _propose_new_dt(err: float, dt: float, opts: AdaptiveOptions) -> float: + """ + Standard adaptive controller: + + dt_new = safety * dt * err^(-1/(p+1)) + + Because the local error estimator from step-doubling behaves like O(h^(p+1)). + """ + if not math.isfinite(err): + factor = opts.shrink_min + elif err <= 0.0: + factor = opts.growth_max + else: + exponent = -1.0 / (opts.order + 1.0) + factor = opts.safety * (err ** exponent) + factor = min(opts.growth_max, max(opts.shrink_min, factor)) + + dt_new = dt * factor + dt_new = min(opts.dt_max, max(opts.dt_min, dt_new)) + return dt_new + + +def adaptive_integrate( + stepper: Callable[[ArrayLike, float, float], ArrayLike], + y0: ArrayLike, + t0: float, + tfinal: float, + dt0: float, + opts: AdaptiveOptions, + error_norm_fn: Optional[Callable[[ArrayLike], float]] = None, + post_accept_callback: Optional[Callable[[float, ArrayLike], None]] = None, + diagnostics_fn: Optional[Callable[[ArrayLike], Dict[str, Any]]] = None, +) -> AdaptiveResult: + """ + Adaptive integrator using step-doubling around a user-supplied one-step method. + + Parameters + ---------- + stepper + Function stepper(y, t, dt) -> y_next + This should perform one step of the splitting method of nominal order `opts.order`. + y0 + Initial state. + t0 + Initial time. + tfinal + Final time. + dt0 + Initial time step guess. + opts + Adaptive options. + error_norm_fn + Optional norm function for raw differences if not using weighted RMS. + post_accept_callback + Optional hook called after each accepted step. + diagnostics_fn + Optional function diagnostics_fn(y) -> dict to store per-step diagnostics. + + Returns + ------- + AdaptiveResult + """ + if tfinal <= t0: + raise ValueError("tfinal must be greater than t0.") + if dt0 <= 0.0: + raise ValueError("dt0 must be positive.") + if opts.order < 1: + raise ValueError("opts.order must be at least 1.") + + start_all = time.perf_counter() + + t = float(t0) + y = y0 + dt = min(max(dt0, opts.dt_min), opts.dt_max) + + t_values: List[float] = [t] + y_values: List[ArrayLike] = [y] + records: List[StepRecord] = [] + + n_accept = 0 + n_reject = 0 + + while t < tfinal: + dt = min(dt, tfinal - t) + reject_count_this_step = 0 + + while True: + tic = time.perf_counter() + + # One full step + y_full = stepper(y, t, dt) + + # Two half-steps + y_half = stepper(y, t, 0.5 * dt) + y_two_half = stepper(y_half, t + 0.5 * dt, 0.5 * dt) + + err = _default_error_estimate( + y_two_half=y_two_half, + y_full=y_full, + atol=opts.atol, + rtol=opts.rtol, + use_weighted_rms=opts.use_weighted_rms, + error_norm_fn=error_norm_fn, + ) + + cpu = time.perf_counter() - tic + dt_new = _propose_new_dt(err, dt, opts) + + if not math.isfinite(err) and not opts.accept_on_nan_error: + accepted = False + else: + accepted = (err <= 1.0) + + extra = {} + if diagnostics_fn is not None: + try: + extra = diagnostics_fn(y_two_half if accepted else y_full) + except Exception as exc: + extra = {"diagnostics_error": str(exc)} + + records.append( + StepRecord( + t=t, + dt_try=dt, + dt_suggested=dt_new, + accepted=accepted, + err_est=err, + cpu_seconds=cpu, + extra=extra, + ) + ) + + if accepted: + t = t + dt + y = y_two_half # higher quality solution than single full step + t_values.append(t) + y_values.append(y) + n_accept += 1 + + if post_accept_callback is not None: + post_accept_callback(t, y) + + dt = dt_new + break + + n_reject += 1 + reject_count_this_step += 1 + dt = dt_new + + if dt <= opts.dt_min + 1.0e-30: + wall = time.perf_counter() - start_all + return AdaptiveResult( + t_values=t_values, + y_values=y_values, + records=records, + n_accept=n_accept, + n_reject=n_reject, + wall_seconds=wall, + success=False, + message="Step size reached dt_min during rejection.", + ) + + if reject_count_this_step >= opts.max_reject: + wall = time.perf_counter() - start_all + return AdaptiveResult( + t_values=t_values, + y_values=y_values, + records=records, + n_accept=n_accept, + n_reject=n_reject, + wall_seconds=wall, + success=False, + message="Exceeded maximum rejects for a single advance.", + ) + + wall = time.perf_counter() - start_all + return AdaptiveResult( + t_values=t_values, + y_values=y_values, + records=records, + n_accept=n_accept, + n_reject=n_reject, + wall_seconds=wall, + success=True, + message="Adaptive integration completed successfully.", + ) \ No newline at end of file