# CelerisLab/tests/run_sah04_st_matrix.py """Sah04 St validation matrix from tests/Sah04_St_validation_matrix.md. Runs the nine paper-anchored (Re, beta tier) cases with fixed D=30 channel geometry, optional SRT/TRT/MRT sweep, and evaluates St against targets. Each case reports **two** Strouhal estimates from the same lift window: - **raw**: dominant rFFT peak in the paper-frequency band (no target prior). - **guided**: same band with a Gaussian weight toward ``f0`` from the paper ``St_target`` (reduces harmonic confusion in narrow channels). The matrix **5% / 10% hard-case rules** apply to **guided** ``St`` by default; ``St_raw`` is for diagnosing whether the guide helps or hides a spectrum issue. Usage:: conda run -n pycuda_3_10 python tests/run_sah04_st_matrix.py --collision MRT conda run -n pycuda_3_10 python tests/run_sah04_st_matrix.py --collision all --json-out sah04_matrix.json conda run -n pycuda_3_10 python tests/run_sah04_st_matrix.py --smoke --case 3 --collision SRT Long diagnostic (lift + spectrum + **final-step vorticity** under ``tests/output/``):: conda run -n pycuda_3_10 python tests/run_sah04_st_matrix.py --collision MRT --case all \\ --steps 200000 --burn 80000 \\ --dump-npz-dir tests/output/sah04_long/npz \\ --final-vorticity-dir tests/output/sah04_long/vorticity \\ --json-out tests/output/sah04_long/matrix_mrt.json Requires **matplotlib** for ``--final-vorticity-dir`` (not a core package dependency). Design:: Hard cases use a 5% relative St gate on **guided** St; no hard case worse than 10%. Soft cases (2, 8) only check ordering vs neighbors in printed summary. """ from __future__ import annotations import argparse import json import os import sys import tempfile from dataclasses import dataclass from typing import Any, Dict, List, Optional, Sequence, Tuple import numpy as np import pycuda.driver as cuda _PKG_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) _DEFAULT_LBM = os.path.join(_PKG_ROOT, "src", "CelerisLab", "configs", "config_lbm.json") # D=30 fixed; Lx_fluid = 80D per Sah04 confined setup _D = 30 _NX = 80 * _D + 2 _CX = 40.0 * _D + 0.5 _R_CYL = 0.5 * _D # Blockage tiers: fluid height H = ny - 2; cylinder center (1200.5, center_y) _TIERS: Dict[str, Dict[str, float]] = { "low": {"ny": 62.0, "center_y": 30.5, "beta_nom": 0.5}, "mid": {"ny": 40.0, "center_y": 19.5, "beta_nom": 0.8}, "high": {"ny": 35.0, "center_y": 17.0, "beta_nom": 0.9}, } @dataclass(frozen=True) class MatrixCase: case_id: int tier: str re: float target_st: float hard: bool steps: int burn: int # Table from Sah04_St_validation_matrix.md MATRIX: Tuple[MatrixCase, ...] = ( MatrixCase(1, "low", 124.09, 0.3393, True, 80_000, 30_000), MatrixCase(2, "low", 160.0, 0.3450, False, 60_000, 20_000), MatrixCase(3, "low", 200.0, 0.3513, True, 60_000, 20_000), MatrixCase(4, "mid", 110.24, 0.5363, True, 80_000, 30_000), MatrixCase(5, "mid", 160.0, 0.5537, True, 60_000, 20_000), MatrixCase(6, "mid", 200.0, 0.5510, True, 60_000, 20_000), MatrixCase(7, "high", 162.82, 0.5202, True, 80_000, 30_000), MatrixCase(8, "high", 180.0, 0.5254, False, 60_000, 20_000), MatrixCase(9, "high", 200.0, 0.5314, True, 60_000, 20_000), ) def _load_json(path: str) -> dict: with open(path, "r", encoding="utf-8") as f: return json.load(f) def _write_json(path: str, d: dict) -> None: with open(path, "w", encoding="utf-8") as f: json.dump(d, f, indent=2) def vorticity_z_from_velocity(ux: np.ndarray, uy: np.ndarray) -> np.ndarray: """Z-component vorticity ``ωz = ∂uy/∂x − ∂ux/∂y`` on a 2D ``(ny, nx)`` slice. Lattice spacing is taken as 1 in ``np.gradient`` (LBM cell units). """ ux = np.asarray(ux, dtype=np.float64) uy = np.asarray(uy, dtype=np.float64) duy_dx = np.gradient(uy, axis=1) dux_dy = np.gradient(ux, axis=0) return duy_dx - dux_dy def save_final_vorticity_png( path: str, ux: np.ndarray, uy: np.ndarray, *, title: str, ) -> None: """Write ``ωz`` heatmap from last-step ``ux``/``uy`` to ``path`` (PNG).""" try: import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt except ImportError as e: raise RuntimeError( "save_final_vorticity_png requires matplotlib (e.g. pip install matplotlib)." ) from e omega = vorticity_z_from_velocity(ux, uy) abs_o = np.abs(omega[np.isfinite(omega)]) if abs_o.size: vmax = float(np.percentile(abs_o, 99.5)) if vmax <= 0.0: vmax = float(np.max(abs_o)) or 1.0 else: vmax = 1.0 ny, nx = omega.shape fw = min(18.0, max(8.0, nx / 100.0)) fh = min(12.0, max(3.0, ny / 40.0)) fig, ax = plt.subplots(figsize=(fw, fh)) im = ax.imshow( omega, origin="lower", aspect="equal", cmap="RdBu_r", vmin=-vmax, vmax=vmax, extent=(0, nx - 1, 0, ny - 1), ) ax.set_xlabel("x (lattice)") ax.set_ylabel("y (lattice)") ax.set_title(title) fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label="omega_z") fig.tight_layout() fig.savefig(path, dpi=150, bbox_inches="tight") plt.close(fig) def rfft_power_spectrum( lift: np.ndarray, *, sample_dt: float, ) -> Tuple[np.ndarray, np.ndarray]: """Mean-subtracted Hanning-windowed lift → ``(freqs_hz, power)`` for positive rFFT bins.""" x = np.asarray(lift, dtype=np.float64) x = x - np.mean(x) n = x.size if n < 64: return np.zeros(0, dtype=np.float64), np.zeros(0, dtype=np.float64) win = np.hanning(n) xw = x * win spec = np.abs(np.fft.rfft(xw)) ** 2 freqs = np.fft.rfftfreq(n, d=float(sample_dt)) return freqs.astype(np.float64), spec.astype(np.float64) def _parabolic_peak_freq(freqs: np.ndarray, spec: np.ndarray, idx: int) -> float: """Refine discrete FFT peak index ``idx`` with log-power parabolic fit.""" idx = int(np.clip(idx, 0, spec.size - 1)) if idx <= 0 or idx + 1 >= spec.size: return float(freqs[idx]) i0, i1, i2 = idx - 1, idx, idx + 1 y0, y1, y2 = ( np.log(spec[i0] + 1e-30), np.log(spec[i1] + 1e-30), np.log(spec[i2] + 1e-30), ) denom = y0 - 2.0 * y1 + y2 if abs(denom) < 1e-20: return float(freqs[i1]) delta = 0.5 * (y0 - y2) / denom delta = float(np.clip(delta, -1.0, 1.0)) df = float(freqs[i2] - freqs[i1]) return float(freqs[i1]) + delta * df def dual_strouhal_from_lift( lift: np.ndarray, *, diameter: float, u_max: float, sample_dt: float, f_hz_min: float, f_hz_max: float, ) -> Dict[str, float]: """Return raw and guided ``St`` and underlying ``f_peak`` (cycles per LBM step).""" nan = {"St_raw": float("nan"), "f_raw": float("nan"), "St_guided": float("nan"), "f_guided": float("nan")} freqs, spec = rfft_power_spectrum(lift, sample_dt=sample_dt) if freqs.size == 0: return nan mask = (freqs >= float(f_hz_min)) & (freqs <= float(f_hz_max)) if not np.any(mask): return nan m = mask.astype(float) idx_raw = int(np.argmax(spec * m)) f0 = 0.5 * (float(f_hz_min) + float(f_hz_max)) sigma = max(1e-12, 0.18 * f0) weight = np.exp(-((freqs - f0) / sigma) ** 2) idx_g = int(np.argmax(spec * m * weight)) f_raw = _parabolic_peak_freq(freqs, spec, idx_raw) f_g = _parabolic_peak_freq(freqs, spec, idx_g) return { "St_raw": float(f_raw * diameter / u_max), "f_raw": float(f_raw), "St_guided": float(f_g * diameter / u_max), "f_guided": float(f_g), } def dominant_strouhal_from_lift( lift: np.ndarray, *, diameter: float, u_max: float, sample_dt: float = 1.0, f_hz_min: float, f_hz_max: float, ) -> Tuple[float, float]: """Backward-compatible: returns guided ``(St, f_peak)``.""" d = dual_strouhal_from_lift( lift, diameter=diameter, u_max=u_max, sample_dt=sample_dt, f_hz_min=f_hz_min, f_hz_max=f_hz_max, ) return d["St_guided"], d["f_guided"] def shedding_freq_band_hz( target_st: float, u_max: float, d_lattice: float, *, half_width: float = 0.42 ) -> Tuple[float, float]: """Cycles per LBM step around ``f0 = St*Umax/D`` so sparse rFFT bins still get coverage.""" f0 = float(target_st) * float(u_max) / float(d_lattice) lo = max(1e-8, f0 * (1.0 - half_width)) hi = f0 * (1.0 + half_width) return lo, hi def run_one_simulation( *, collision: str, outlet: str, nx: int, ny: int, center: Tuple[float, float], re: float, d_lattice: float, r_cyl: float, u_max: float, steps: int, burn: int, record_every: int, f_hz_min: float, f_hz_max: float, dump_npz_path: Optional[str] = None, final_vorticity_png_path: Optional[str] = None, flow_figure_title: str = "Sah04 vorticity (final LBM step)", ) -> Dict[str, Any]: """Build configs, run steps, return St, rough Cd, curved stats.""" u0_mean = u_max / 1.5 nu = u_max * d_lattice / re lbm_path = _DEFAULT_LBM if not os.path.isfile(lbm_path): raise FileNotFoundError(lbm_path) cfg = _load_json(lbm_path) cfg["grid"]["nx"] = nx cfg["grid"]["ny"] = ny cfg["grid"]["nz"] = 1 cfg["physics"]["viscosity"] = float(nu) cfg["physics"]["velocity"] = float(u0_mean) cfg["physics"]["rho"] = 1.0 cfg["method"]["collision"] = collision cfg["method"]["streaming"] = "double_buffer" cfg["method"]["les"]["enabled"] = False cfg["method"]["outlet"]["mode"] = outlet body_doc = { "objects": [ { "type": "cylinder", "center": list(center), "radius": float(r_cyl), } ] } tmpd = tempfile.mkdtemp(prefix="celeris_sah04_mtx_") lbm_tmp = os.path.join(tmpd, "config_lbm.json") body_tmp = os.path.join(tmpd, "config_body.json") _write_json(lbm_tmp, cfg) _write_json(body_tmp, body_doc) from CelerisLab import Simulation # noqa: WPS433 sim = Simulation(lbm_config_path=lbm_tmp, body_config_path=body_tmp) sim.initialize() stream = cuda.Stream() lift_hist: List[float] = [] fx_hist: List[float] = [] step_hist: List[int] = [] rec_every = max(1, int(record_every)) rho_every = max(2000, rec_every) rho_snap_step: List[int] = [] rho_snap_min: List[float] = [] rho_snap_max: List[float] = [] n_curved = int(sim.field.n_curved) fb = int(sim.bodies.fallback_link_count()) lq = int(sim.bodies.low_q_link_count()) for step in range(1, int(steps) + 1): sim.bodies.zero_force_segment_async(stream) sim.stepper.step( 1, action_gpu=sim.bodies.action_gpu, obs_gpu=sim.bodies.obs_gpu, stream=stream, ) if step % rec_every == 0 or step == int(steps): stream.synchronize() sim.bodies.download_obs_full_async(stream) stream.synchronize() fvec = sim.bodies.read_force(0) lift_hist.append(float(fvec[1])) fx_hist.append(float(fvec[0])) step_hist.append(int(step)) if not np.isfinite(lift_hist[-1]): sim.close() raise RuntimeError(f"NaN/Inf lift at step {step}") if step % rho_every == 0 or step == int(steps): stream.synchronize() macro = sim.get_macroscopic() rho_snap_step.append(step) rho_snap_min.append(float(np.min(macro["rho"]))) rho_snap_max.append(float(np.max(macro["rho"]))) if final_vorticity_png_path: stream.synchronize() macro_last = sim.get_macroscopic() _vdir = os.path.dirname(os.path.abspath(final_vorticity_png_path)) if _vdir: os.makedirs(_vdir, exist_ok=True) save_final_vorticity_png( final_vorticity_png_path, macro_last["ux"], macro_last["uy"], title=flow_figure_title, ) sim.close() lift_arr = np.array(lift_hist, dtype=np.float64) fx_arr = np.array(fx_hist, dtype=np.float64) step_arr = np.array(step_hist, dtype=np.int64) burn_samp = min(int(burn) // rec_every, max(0, lift_arr.size - 16)) tail = lift_arr[burn_samp:] dual = dual_strouhal_from_lift( tail, diameter=d_lattice, u_max=u_max, sample_dt=float(rec_every), f_hz_min=f_hz_min, f_hz_max=f_hz_max, ) st_guided = dual["St_guided"] f_guided = dual["f_guided"] mean_cd = float(np.mean(fx_arr[burn_samp:])) * 2.0 / (u_max ** 2 * d_lattice) if fx_arr.size else float("nan") if rho_snap_min: i0 = next((i for i, s in enumerate(rho_snap_step) if s >= burn), 0) rho_rng = (min(rho_snap_min[i0:]), max(rho_snap_max[i0:])) else: rho_rng = (float("nan"), float("nan")) if dump_npz_path: freqs_pb, spec_pb = rfft_power_spectrum(tail, sample_dt=float(rec_every)) f0_hz = 0.5 * (float(f_hz_min) + float(f_hz_max)) sigma = max(1e-12, 0.18 * f0_hz) w_guided = ( np.exp(-((freqs_pb - f0_hz) / sigma) ** 2) if freqs_pb.size else np.zeros(0) ) band = ( ((freqs_pb >= float(f_hz_min)) & (freqs_pb <= float(f_hz_max))).astype(np.float64) if freqs_pb.size else np.zeros(0) ) _dir = os.path.dirname(os.path.abspath(dump_npz_path)) if _dir: os.makedirs(_dir, exist_ok=True) np.savez_compressed( dump_npz_path, lift_samples=lift_arr, fx_samples=fx_arr, sample_lbm_step=step_arr, burn_index_samples=int(burn_samp), record_every_lbm_steps=int(rec_every), freqs_hz_post_burn=freqs_pb, power_post_burn=spec_pb, band_mask=band, guided_gaussian_weight=w_guided, f_hz_min=np.array([float(f_hz_min)], dtype=np.float64), f_hz_max=np.array([float(f_hz_max)], dtype=np.float64), f0_hz_band_mid=np.array([f0_hz], dtype=np.float64), St_raw=np.array([float(dual["St_raw"])], dtype=np.float64), St_guided=np.array([float(dual["St_guided"])], dtype=np.float64), u_max=np.array([float(u_max)], dtype=np.float64), diameter_lattice=np.array([float(d_lattice)], dtype=np.float64), ) return { "St": float(st_guided), "St_guided": float(st_guided), "St_raw": float(dual["St_raw"]), "f_peak_per_step": float(f_guided), "f_peak_raw_per_step": float(dual["f_raw"]), "f_peak_guided_per_step": float(f_guided), "mean_Cd": float(mean_cd), "n_curved": n_curved, "fallback_links": fb, "low_q_links": lq, "rho_min_post_burn": rho_rng[0], "rho_max_post_burn": rho_rng[1], "n_lift_samples": int(lift_arr.size), } def relative_st_error(st_meas: float, st_target: float) -> float: if not np.isfinite(st_meas) or st_target <= 0: return float("inf") return abs(st_meas - st_target) / st_target def evaluate_hard_cases(rows: Sequence[Dict[str, Any]], collision: str) -> Dict[str, Any]: hard = [ r for r in rows if r.get("hard") and r.get("collision") == collision and "St" in r and np.isfinite(r["St"]) ] errs = [relative_st_error(r["St"], r["target_st"]) for r in hard] errs_raw = [ relative_st_error(r["St_raw"], r["target_st"]) for r in hard if "St_raw" in r and np.isfinite(r["St_raw"]) ] finite = [e for e in errs if np.isfinite(e)] within5 = sum(1 for e in finite if e <= 0.05) worse10 = sum(1 for e in finite if e > 0.10) fr = [e for e in errs_raw if np.isfinite(e)] return { "collision": collision, "hard_count": len(hard), "hard_within_5pct": int(within5), "hard_worse_than_10pct": int(worse10), "pass_primary_rule": bool(within5 >= 5 and worse10 == 0), "hard_median_abs_rel_err_raw": float(np.median(fr)) if fr else None, } def main() -> int: ap = argparse.ArgumentParser(description="Sah04 St validation matrix (9 cases × collisions)") ap.add_argument("--collision", default="MRT", help="SRT, TRT, MRT, or all") ap.add_argument("--outlet", default="neq_extrap", choices=("neq_extrap", "zero_gradient", "blended")) ap.add_argument("--record-every", type=int, default=5) ap.add_argument("--case", default="all", help='Case id 1-9 or "all"') ap.add_argument("--smoke", action="store_true", help="Short steps/burn for wiring checks") ap.add_argument("--json-out", type=str, default=None, help="Write full result rows to JSON") ap.add_argument( "--steps", type=int, default=None, help="Override matrix LBM steps for each case (ignored with --smoke)", ) ap.add_argument( "--burn", type=int, default=None, help="Override matrix burn in LBM steps (ignored with --smoke)", ) ap.add_argument( "--dump-npz-dir", type=str, default=None, help="Directory: write case{id}_{COLL}.npz (+ .meta.json) with lift, fx, sample steps, post-burn spectrum", ) ap.add_argument( "--final-vorticity-dir", type=str, default=None, help="Directory: write case{id}_{COLL}_laststep.png (omega_z from final macroscopic slice; needs matplotlib)", ) args = ap.parse_args() u_max = 0.1 collisions: List[str] if str(args.collision).lower() == "all": collisions = ["SRT", "TRT", "MRT"] else: collisions = [str(args.collision).upper()] if collisions[0] not in ("SRT", "TRT", "MRT"): print("--collision must be SRT, TRT, MRT, or all", file=sys.stderr) return 2 case_filter: Optional[int] = None if str(args.case).lower() != "all": case_filter = int(args.case) if case_filter < 1 or case_filter > 9: print("--case must be 1-9 or all", file=sys.stderr) return 2 cases_to_run = [c for c in MATRIX if case_filter is None or c.case_id == case_filter] if args.dump_npz_dir: os.makedirs(args.dump_npz_dir, exist_ok=True) if args.final_vorticity_dir: os.makedirs(args.final_vorticity_dir, exist_ok=True) rows: List[Dict[str, Any]] = [] for coll in collisions: for mc in cases_to_run: tier = _TIERS[mc.tier] ny = int(tier["ny"]) center = (_CX, float(tier["center_y"])) if args.smoke: steps = 6000 burn = 1500 else: steps = int(args.steps) if args.steps is not None else mc.steps burn = int(args.burn) if args.burn is not None else mc.burn flo, fhi = shedding_freq_band_hz(mc.target_st, u_max, float(_D)) dump_path: Optional[str] = None if args.dump_npz_dir: dump_path = os.path.join(args.dump_npz_dir, f"case{mc.case_id}_{coll}.npz") vort_path: Optional[str] = None vort_title = "" if args.final_vorticity_dir: vort_path = os.path.join( args.final_vorticity_dir, f"case{mc.case_id}_{coll}_laststep.png", ) vort_title = ( f"Sah04 case {mc.case_id} tier={mc.tier} {coll} Re={mc.re} " f"nx={_NX} ny={ny} last LBM step={steps} (omega_z)" ) print( f"--- case {mc.case_id} tier={mc.tier} beta~{tier['beta_nom']} " f"Re={mc.re} target_St={mc.target_st} {coll} steps={steps} burn={burn} ---", flush=True, ) try: out = run_one_simulation( collision=coll, outlet=args.outlet, nx=_NX, ny=ny, center=center, re=float(mc.re), d_lattice=float(_D), r_cyl=_R_CYL, u_max=u_max, steps=steps, burn=burn, record_every=int(args.record_every), f_hz_min=flo, f_hz_max=fhi, dump_npz_path=dump_path, final_vorticity_png_path=vort_path, flow_figure_title=vort_title or "Sah04 vorticity (final LBM step)", ) except Exception as e: print(f"FAILED: {e}", file=sys.stderr) rows.append( { "case_id": mc.case_id, "tier": mc.tier, "collision": coll, "Re": mc.re, "target_st": mc.target_st, "hard": mc.hard, "error": str(e), } ) continue rel_err = relative_st_error(out["St"], mc.target_st) rel_err_raw = relative_st_error(out["St_raw"], mc.target_st) row = { "case_id": mc.case_id, "tier": mc.tier, "beta_nominal": tier["beta_nom"], "collision": coll, "Re": mc.re, "target_st": mc.target_st, "hard": mc.hard, "St": out["St"], "St_raw": out["St_raw"], "St_guided": out["St_guided"], "rel_err_st": float(rel_err) if np.isfinite(rel_err) else None, "rel_err_st_raw": float(rel_err_raw) if np.isfinite(rel_err_raw) else None, "within_5pct": bool(np.isfinite(rel_err) and rel_err <= 0.05), "worse_10pct": bool(np.isfinite(rel_err) and rel_err > 0.10), "mean_Cd": out["mean_Cd"], "n_curved": out["n_curved"], "fallback_links": out["fallback_links"], "low_q_links": out["low_q_links"], "steps": steps, "burn": burn, } rows.append(row) if dump_path: meta_path = os.path.splitext(dump_path)[0] + ".meta.json" _write_json( meta_path, { "case_id": mc.case_id, "tier": mc.tier, "collision": coll, "Re": mc.re, "target_st": mc.target_st, "hard": mc.hard, "steps": steps, "burn": burn, "record_every": int(args.record_every), "f_hz_min": flo, "f_hz_max": fhi, "npz_basename": os.path.basename(dump_path), "St_raw": out["St_raw"], "St_guided": out["St_guided"], "mean_Cd": out["mean_Cd"], }, ) flag = "OK" if row.get("within_5pct") else ("SOFT" if not mc.hard else "CHECK") print( f" St_raw={out['St_raw']:.5f} St_guided={out['St_guided']:.5f} " f"rel(guided)={100.0 * rel_err:.2f}% rel(raw)={100.0 * rel_err_raw:.2f}% " f"Cd~{out['mean_Cd']:.4f} [{flag}]", flush=True, ) if not args.smoke and case_filter is None and str(args.case).lower() == "all": print("\n=== Hard-case summary per collision (5% gate; need 5/7 & none >10%) ===") for coll in collisions: ev = evaluate_hard_cases(rows, coll) print(json.dumps(ev, indent=2)) if args.json_out: ev_all = { coll: evaluate_hard_cases(rows, coll) for coll in {r.get("collision") for r in rows if r.get("collision")} if coll } _write_json(args.json_out, {"rows": rows, "evaluation_by_collision": ev_all}) return 0 if __name__ == "__main__": raise SystemExit(main())