CelerisLab/tests/run_inlet_scenario_fields.py

588 lines
19 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# CelerisLab/tests/run_inlet_scenario_fields.py
"""Three-scenario inlet field export: final or last-stable step at 5000 LBM steps.
Scenarios:
- empty_channel: zou_he_local × {SRT,MRT} × {free_slip,bounce_back}
- empty_channel_inlet_matrix: all inlet schemes × {SRT,MRT}, bounce_back only
- kan99b: zou_he_local × {SRT,MRT} × {free_slip,bounce_back}, Re=100, domain M
- sah04_case9: channel_stabilized × {SRT,MRT}, high-blockage case 9 geometry
Outputs per run (under ``--out-dir``):
fields/final_{rho,ux,vort}.png, fields/final.npz
lines/lines_ux_rho.png, lines/lines.npz
run_meta.json
Usage::
conda run -n pycuda_3_10 python tests/run_inlet_scenario_fields.py
conda run -n pycuda_3_10 python tests/run_inlet_scenario_fields.py --scenario empty_channel
conda run -n pycuda_3_10 python tests/run_inlet_scenario_fields.py --scenario empty_channel_inlet_matrix
"""
from __future__ import annotations
import argparse
import csv
import json
import os
import sys
import tempfile
from dataclasses import dataclass, replace
from typing import Any, Dict, List, Optional, Sequence, Tuple
import numpy as np
_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
_DEFAULT_LBM = os.path.join(_REPO, "src", "CelerisLab", "configs", "config_lbm.json")
# Kan99b lattice contract
_KAN_U_INF = 0.03
_KAN_D = 30.0
_KAN_R = 15.0
_KAN_RE = 100.0
_KAN_ALPHA = 1.0
# Sah04 case 9 (high tier)
_SAH_D = 30
_SAH_NX = 80 * _SAH_D + 2
_SAH_NY = 35
_SAH_CX = 40.0 * _SAH_D + 0.5
_SAH_CY = 17.0
_SAH_RE = 200.0
_SAH_U_MAX = 0.1
@dataclass(frozen=True)
class RunSpec:
"""One simulation run specification."""
scenario: str
run_id: str
label: str
nx: int
ny: int
collision: str
inlet_scheme: str
inlet_profile: str
y_wall_bc: str
outlet_mode: str
velocity: float
viscosity: float
steps: int
has_cylinder: bool
cylinder_center: Tuple[float, float] = (0.0, 0.0)
cylinder_radius: float = 0.0
cylinder_omega: float = 0.0
def _load_json(path: str) -> dict:
with open(path, "r", encoding="utf-8") as f:
return json.load(f)
def _write_json(path: str, payload: dict) -> None:
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
with open(path, "w", encoding="utf-8") as f:
json.dump(payload, f, indent=2)
def vorticity_z(ux: np.ndarray, uy: np.ndarray) -> np.ndarray:
ux = np.asarray(ux, dtype=np.float64)
uy = np.asarray(uy, dtype=np.float64)
return np.gradient(uy, axis=1) - np.gradient(ux, axis=0)
def _line_y_indices(ny: int) -> List[Tuple[int, str]]:
return [(1, "y1"), (ny // 2, f"y{ny // 2}"), (ny - 2, f"y{ny - 2}")]
_INLET_SCHEMES = (
"zou_he_local",
"channel_stabilized",
"equilibrium",
"regularized",
)
def _empty_channel_inlet_matrix_specs() -> List[RunSpec]:
"""All inlet schemes on empty channel, bounce_back, SRT/MRT (5000-step field export)."""
specs: List[RunSpec] = []
for coll in ("SRT", "MRT"):
for scheme in _INLET_SCHEMES:
run_id = f"{coll.lower()}_{scheme}"
specs.append(
RunSpec(
scenario="empty_channel_inlet_matrix",
run_id=run_id,
label=f"empty {scheme} {coll} bounce_back",
nx=401,
ny=201,
collision=coll,
inlet_scheme=scheme,
inlet_profile="uniform",
y_wall_bc="bounce_back",
outlet_mode="neq_extrap",
velocity=0.03,
viscosity=0.009,
steps=5000,
has_cylinder=False,
)
)
return specs
def _all_specs() -> List[RunSpec]:
specs: List[RunSpec] = []
for coll in ("SRT", "MRT"):
for wall in ("free_slip", "bounce_back"):
wid = f"{coll.lower()}_{wall}"
specs.append(
RunSpec(
scenario="empty_channel",
run_id=wid,
label=f"empty zou_he {coll} {wall}",
nx=401,
ny=201,
collision=coll,
inlet_scheme="zou_he_local",
inlet_profile="uniform",
y_wall_bc=wall,
outlet_mode="neq_extrap",
velocity=0.03,
viscosity=0.009,
steps=5000,
has_cylinder=False,
)
)
dom_m = (1351, 601, (450.0, 300.0))
nu_k = _KAN_U_INF * _KAN_D / _KAN_RE
omega = 2.0 * _KAN_ALPHA * _KAN_U_INF / _KAN_D
for coll in ("SRT", "MRT"):
for wall in ("free_slip", "bounce_back"):
wid = f"{coll.lower()}_{wall}"
specs.append(
RunSpec(
scenario="kan99b",
run_id=wid,
label=f"kan99b zou_he {coll} {wall}",
nx=dom_m[0],
ny=dom_m[1],
collision=coll,
inlet_scheme="zou_he_local",
inlet_profile="uniform",
y_wall_bc=wall,
outlet_mode="neq_extrap",
velocity=_KAN_U_INF,
viscosity=nu_k,
steps=5000,
has_cylinder=True,
cylinder_center=dom_m[2],
cylinder_radius=_KAN_R,
cylinder_omega=omega,
)
)
u0_mean = _SAH_U_MAX / 1.5
nu_s = _SAH_U_MAX * _SAH_D / _SAH_RE
for coll in ("SRT", "MRT"):
wid = f"{coll.lower()}_channel_stab"
specs.append(
RunSpec(
scenario="sah04_case9",
run_id=wid,
label=f"sah04 case9 channel_stab {coll}",
nx=_SAH_NX,
ny=_SAH_NY,
collision=coll,
inlet_scheme="channel_stabilized",
inlet_profile="parabolic",
y_wall_bc="bounce_back",
outlet_mode="neq_extrap",
velocity=u0_mean,
viscosity=nu_s,
steps=5000,
has_cylinder=True,
cylinder_center=(_SAH_CX, _SAH_CY),
cylinder_radius=0.5 * _SAH_D,
cylinder_omega=0.0,
)
)
return specs
def _build_cfg(base: dict, spec: RunSpec) -> dict:
cfg = json.loads(json.dumps(base))
cfg["grid"]["nx"] = spec.nx
cfg["grid"]["ny"] = spec.ny
cfg["grid"]["nz"] = 1
cfg["physics"]["velocity"] = float(spec.velocity)
cfg["physics"]["viscosity"] = float(spec.viscosity)
cfg["physics"]["rho"] = 1.0
cfg["method"]["collision"] = spec.collision.upper()
cfg["method"]["streaming"] = "double_buffer"
cfg["method"]["store_precision"] = "FP32"
cfg["method"]["les"]["enabled"] = False
cfg["method"]["inlet"]["profile"] = spec.inlet_profile
cfg["method"]["inlet"]["scheme"] = spec.inlet_scheme
cfg["method"]["y_wall_bc"] = spec.y_wall_bc
cfg["method"]["outlet"]["mode"] = spec.outlet_mode
return cfg
def _body_doc(spec: RunSpec) -> dict:
if not spec.has_cylinder:
return {"objects": []}
return {
"objects": [
{
"type": "cylinder",
"center": [float(spec.cylinder_center[0]), float(spec.cylinder_center[1])],
"radius": float(spec.cylinder_radius),
"omega": float(spec.cylinder_omega),
}
]
}
def _save_field_pngs(
out_dir: str,
prefix: str,
*,
rho: np.ndarray,
ux: np.ndarray,
vort: np.ndarray,
title: str,
) -> List[str]:
try:
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
except ImportError:
return []
os.makedirs(out_dir, exist_ok=True)
ny, nx = rho.shape
extent = (0, nx - 1, 0, ny - 1)
paths: List[str] = []
def _one(arr: np.ndarray, name: str, cmap: str, sym: bool) -> None:
a = np.asarray(arr, dtype=np.float64)
fin = a[np.isfinite(a)]
if fin.size == 0:
vmin, vmax = -1.0, 1.0
elif sym:
v = float(np.percentile(np.abs(fin), 99.5)) or 1.0
vmin, vmax = -v, v
else:
vmin = float(np.percentile(fin, 0.5))
vmax = float(np.percentile(fin, 99.5))
if vmax <= vmin:
vmax = vmin + 1.0
fw = min(18.0, max(8.0, nx / 70.0))
fh = min(10.0, max(3.0, ny / 45.0))
fig, ax = plt.subplots(figsize=(fw, fh))
im = ax.imshow(a, origin="lower", aspect="auto", cmap=cmap, vmin=vmin, vmax=vmax, extent=extent)
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_title(f"{title}{name}")
fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
fig.tight_layout()
p = os.path.join(out_dir, f"{prefix}_{name}.png")
fig.savefig(p, dpi=150, bbox_inches="tight")
plt.close(fig)
paths.append(p)
_one(rho, "rho", "viridis", sym=False)
_one(ux, "ux", "RdBu_r", sym=True)
_one(vort, "vort", "RdBu_r", sym=True)
return paths
def _save_line_plots(
path: str,
*,
rho: np.ndarray,
ux: np.ndarray,
step: int,
label: str,
y_rows: Sequence[Tuple[int, str]],
) -> None:
try:
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
except ImportError:
return
ny, nx = rho.shape
x = np.arange(nx, dtype=np.float64)
fig, axes = plt.subplots(2, 1, figsize=(min(14.0, max(8.0, nx / 55.0)), 7.0), sharex=True)
for y_idx, y_lab in y_rows:
yi = int(np.clip(y_idx, 0, ny - 1))
axes[0].plot(x, ux[yi, :], label=y_lab, linewidth=1.0)
axes[1].plot(x, rho[yi, :], label=y_lab, linewidth=1.0)
axes[0].set_ylabel("u_x")
axes[0].legend(loc="best", fontsize=8)
axes[0].grid(True, alpha=0.3)
axes[1].set_ylabel("rho")
axes[1].set_xlabel("x (lattice)")
axes[1].legend(loc="best", fontsize=8)
axes[1].grid(True, alpha=0.3)
fig.suptitle(f"{label} — ux/rho lines at step {step}")
fig.tight_layout()
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
fig.savefig(path, dpi=150, bbox_inches="tight")
plt.close(fig)
def _snapshot_from_sim(sim) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
macro = sim.get_macroscopic()
rho = np.asarray(macro["rho"], dtype=np.float64)
ux = np.asarray(macro["ux"], dtype=np.float64)
uy = np.asarray(macro["uy"], dtype=np.float64)
vort = vorticity_z(ux, uy)
return rho, ux, uy, vort
def _is_stable_fields(
rho: np.ndarray,
ux: np.ndarray,
*,
rho_lo: float = 0.85,
rho_hi: float = 1.25,
ux_cap: float = 0.15,
) -> bool:
"""Finite fields within a physically plausible band (reject pre-blow-up states)."""
if not (np.isfinite(rho).all() and np.isfinite(ux).all()):
return False
r0 = float(np.min(rho))
r1 = float(np.max(rho))
umax = float(np.max(np.abs(ux)))
return (rho_lo <= r0) and (r1 <= rho_hi) and (umax <= ux_cap)
def run_one(spec: RunSpec, base_cfg: dict, out_root: str) -> Dict[str, Any]:
sys.path.insert(0, os.path.join(_REPO, "src"))
import pycuda.driver as cuda
from CelerisLab import Simulation # noqa: WPS433
run_dir = os.path.join(out_root, spec.scenario, spec.run_id)
field_dir = os.path.join(run_dir, "fields")
line_dir = os.path.join(run_dir, "lines")
os.makedirs(field_dir, exist_ok=True)
os.makedirs(line_dir, exist_ok=True)
cfg = _build_cfg(base_cfg, spec)
tmpd = tempfile.mkdtemp(prefix="inlet_scenario_")
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(spec))
sim = Simulation(lbm_config_path=lbm_tmp, body_config_path=body_tmp)
if spec.has_cylinder and spec.cylinder_omega != 0.0:
sim.bodies.get(0).state.omega = np.float32(spec.cylinder_omega)
sim.initialize()
stream = cuda.Stream()
y_rows = _line_y_indices(spec.ny)
last_good: Optional[Dict[str, Any]] = None
first_bad_step: Optional[int] = None
force_bad_step: Optional[int] = None
print(f" [{spec.scenario}/{spec.run_id}] {spec.label} steps={spec.steps}", flush=True)
for step in range(1, spec.steps + 1):
if spec.has_cylinder:
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 % 100 == 0 or step == spec.steps:
stream.synchronize()
sim.bodies.download_obs_full_async(stream)
stream.synchronize()
fvec = sim.bodies.read_force(0)
if not (np.isfinite(fvec[0]) and np.isfinite(fvec[1])):
if force_bad_step is None:
force_bad_step = step
else:
sim.step(1)
rho, ux, uy, vort = _snapshot_from_sim(sim)
if _is_stable_fields(rho, ux):
last_good = {
"step": step,
"rho": rho.copy(),
"ux": ux.copy(),
"uy": uy.copy(),
"vort": vort.copy(),
}
elif first_bad_step is None:
first_bad_step = step
sim.close()
if last_good is None:
raise RuntimeError(f"No finite snapshot for {spec.run_id}")
out_step = int(last_good["step"])
rho = last_good["rho"]
ux = last_good["ux"]
uy = last_good["uy"]
vort = last_good["vort"]
requested_final = spec.steps
used_last_stable = out_step < requested_final
meta = {
"scenario": spec.scenario,
"run_id": spec.run_id,
"label": spec.label,
"nx": spec.nx,
"ny": spec.ny,
"collision": spec.collision,
"inlet_scheme": spec.inlet_scheme,
"inlet_profile": spec.inlet_profile,
"y_wall_bc": spec.y_wall_bc,
"outlet_mode": spec.outlet_mode,
"velocity": spec.velocity,
"viscosity": spec.viscosity,
"requested_steps": requested_final,
"output_step": out_step,
"used_last_stable": used_last_stable,
"first_nonfinite_step": first_bad_step,
"first_force_nonfinite_step": force_bad_step,
"rho_min": float(np.min(rho)),
"rho_max": float(np.max(rho)),
"ux_max": float(np.max(np.abs(ux))),
"vort_max": float(np.max(np.abs(vort[np.isfinite(vort)]))) if np.isfinite(vort).any() else float("nan"),
}
_write_json(os.path.join(run_dir, "run_meta.json"), meta)
stem = f"step_{out_step:06d}"
np.savez_compressed(
os.path.join(field_dir, "final.npz"),
rho=rho.astype(np.float32),
ux=ux.astype(np.float32),
uy=uy.astype(np.float32),
vort=vort.astype(np.float32),
step=np.int32(out_step),
)
title = f"{spec.label} (step {out_step}" + (", last stable" if used_last_stable else ", final") + ")"
pngs = _save_field_pngs(field_dir, "final", rho=rho, ux=ux, vort=vort, title=title)
_save_line_plots(
os.path.join(line_dir, "lines_ux_rho.png"),
rho=rho,
ux=ux,
step=out_step,
label=spec.label,
y_rows=y_rows,
)
line_payload: Dict[str, Any] = {"x": np.arange(spec.nx, dtype=np.float32), "step": np.int32(out_step)}
for y_idx, y_lab in y_rows:
yi = int(np.clip(y_idx, 0, spec.ny - 1))
line_payload[f"ux_{y_lab}"] = ux[yi, :].astype(np.float32)
line_payload[f"rho_{y_lab}"] = rho[yi, :].astype(np.float32)
np.savez_compressed(os.path.join(line_dir, "lines.npz"), **line_payload)
status = "last_stable" if used_last_stable else "final"
print(
f" -> {status} step {out_step} rho=[{meta['rho_min']:.4f},{meta['rho_max']:.4f}] "
f"ux_max={meta['ux_max']:.4f} force_bad={force_bad_step}",
flush=True,
)
return {**meta, "field_pngs": pngs, "run_dir": run_dir}
def main() -> int:
ap = argparse.ArgumentParser(description="Three-scenario inlet field export (5000 steps)")
ap.add_argument(
"--scenario",
choices=(
"empty_channel",
"empty_channel_inlet_matrix",
"kan99b",
"sah04_case9",
"all",
),
default="all",
)
ap.add_argument(
"--collision",
default="",
help="Optional filter: SRT or MRT only",
)
ap.add_argument("--steps", type=int, default=5000)
ap.add_argument(
"--out-dir",
default=os.path.join(_REPO, "tests", "output", "inlet_scenario_fields"),
)
args = ap.parse_args()
if not os.path.isfile(_DEFAULT_LBM):
print(f"Missing {_DEFAULT_LBM}", file=sys.stderr)
return 2
base = _load_json(_DEFAULT_LBM)
if args.scenario == "empty_channel_inlet_matrix":
specs = _empty_channel_inlet_matrix_specs()
elif args.scenario == "all":
specs = _all_specs() + _empty_channel_inlet_matrix_specs()
else:
specs = _all_specs()
if args.scenario != "all":
specs = [s for s in specs if s.scenario == args.scenario]
if args.collision.strip():
coll = args.collision.strip().upper()
specs = [s for s in specs if s.collision.upper() == coll]
specs = [replace(s, steps=int(args.steps)) for s in specs]
out_dir = os.path.abspath(args.out_dir)
os.makedirs(out_dir, exist_ok=True)
rows: List[Dict[str, Any]] = []
for spec in specs:
try:
row = run_one(spec, base, out_dir)
rows.append(row)
except Exception as e: # noqa: BLE001
print(f" FAILED {spec.scenario}/{spec.run_id}: {e}", flush=True)
rows.append(
{
"scenario": spec.scenario,
"run_id": spec.run_id,
"label": spec.label,
"error": str(e),
}
)
summary_path = os.path.join(out_dir, "summary.csv")
if rows:
keys: List[str] = []
for r in rows:
for k in r:
if k not in keys and k != "field_pngs":
keys.append(k)
with open(summary_path, "w", encoding="utf-8", newline="") as f:
w = csv.DictWriter(f, fieldnames=keys, extrasaction="ignore")
w.writeheader()
w.writerows(rows)
_write_json(
os.path.join(out_dir, "manifest.json"),
{"steps": args.steps, "scenario_filter": args.scenario, "runs": [s.run_id for s in specs]},
)
print(f"Wrote: {summary_path}", flush=True)
print(f"Output: {out_dir}", flush=True)
return 0
if __name__ == "__main__":
raise SystemExit(main())