239 lines
7.7 KiB
Python
239 lines
7.7 KiB
Python
# tests/run_sah04_case9_flow_diagnostics.py
|
|
"""Sah04 matrix case 9: velocity profiles at several x and full-field vorticity.
|
|
|
|
Loads geometry and setup from ``tests/tests/run_sah04_st_matrix.py``, runs one
|
|
simulation to the last LBM step, then writes PNG + JSON under
|
|
``tests/output/sah04_case9_flow_diag/`` (default).
|
|
|
|
Usage::
|
|
conda run -n pycuda_3_10 python tests/run_sah04_case9_flow_diagnostics.py
|
|
conda run -n pycuda_3_10 python tests/run_sah04_case9_flow_diagnostics.py --steps 60000 --burn 20000
|
|
|
|
Design::
|
|
Intended for inlet / channel diagnostics after boundary fixes; not part of
|
|
the St matrix gate.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
import shutil
|
|
import sys
|
|
import tempfile
|
|
from typing import Any, Dict, List, Sequence
|
|
|
|
import numpy as np
|
|
import pycuda.driver as cuda
|
|
|
|
_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
|
_mtx_path = os.path.join(_REPO, "tests", "tests", "run_sah04_st_matrix.py")
|
|
_spec = importlib.util.spec_from_file_location("sah04_mtx", _mtx_path)
|
|
if _spec is None or _spec.loader is None:
|
|
raise RuntimeError(f"Cannot load {_mtx_path}")
|
|
sah = importlib.util.module_from_spec(_spec)
|
|
sys.modules["sah04_mtx"] = sah
|
|
_spec.loader.exec_module(sah)
|
|
|
|
|
|
def _inlet_parabolic_ref_u(y: np.ndarray, ny: int, u0_mean: float) -> np.ndarray:
|
|
"""Match ``inlet_target_u`` in ``inlet_outlet.cuh`` (parabolic, mean ``u0_mean``)."""
|
|
h = max(float(ny - 2), 1.0)
|
|
yc = np.clip(y.astype(np.float64), 1.0, float(ny - 2))
|
|
eta = (yc - 0.5) / h
|
|
shape = np.maximum(0.0, 4.0 * eta * (1.0 - eta))
|
|
return u0_mean * 1.5 * shape
|
|
|
|
|
|
def _default_x_indices(nx: int) -> List[int]:
|
|
raw = [1, 2, 3, 5, 10, 20, 40, 80, 160, 320, 640, 960, 1200, 1600, 2000, 2200]
|
|
out: List[int] = []
|
|
for x in raw:
|
|
if 1 <= x < nx - 1:
|
|
out.append(int(x))
|
|
mid = max(1, (nx - 2) // 2)
|
|
if mid not in out:
|
|
out.append(mid)
|
|
out = sorted(set(out))
|
|
return out
|
|
|
|
|
|
def _plot_profiles(
|
|
out_png: str,
|
|
*,
|
|
ux: np.ndarray,
|
|
uy: np.ndarray,
|
|
x_indices: Sequence[int],
|
|
ny: int,
|
|
u0_mean: float,
|
|
title: str,
|
|
) -> None:
|
|
import matplotlib
|
|
|
|
matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt
|
|
|
|
y = np.arange(ny, dtype=np.float64)
|
|
u_ref = _inlet_parabolic_ref_u(y, ny, u0_mean)
|
|
|
|
fig, axes = plt.subplots(1, 2, figsize=(12, 6), sharey=True)
|
|
for xi in x_indices:
|
|
axes[0].plot(ux[:, xi], y, label=f"x={xi}", linewidth=1.2)
|
|
axes[1].plot(uy[:, xi], y, label=f"x={xi}", linewidth=1.2)
|
|
axes[0].plot(u_ref, y, "k--", linewidth=1.5, label="inlet parabolic ref")
|
|
axes[0].set_xlabel("u_x")
|
|
axes[0].set_ylabel("y (lattice)")
|
|
axes[0].set_title("u_x (y) at selected x")
|
|
axes[0].legend(loc="best", fontsize=7)
|
|
axes[0].grid(True, alpha=0.3)
|
|
axes[1].set_xlabel("u_y")
|
|
axes[1].set_title("u_y (y) at selected x")
|
|
axes[1].legend(loc="best", fontsize=7)
|
|
axes[1].grid(True, alpha=0.3)
|
|
fig.suptitle(title)
|
|
fig.tight_layout()
|
|
os.makedirs(os.path.dirname(os.path.abspath(out_png)) or ".", exist_ok=True)
|
|
fig.savefig(out_png, dpi=150, bbox_inches="tight")
|
|
plt.close(fig)
|
|
|
|
|
|
def main() -> int:
|
|
ap = argparse.ArgumentParser(description="Sah04 case 9 flow: profiles + vorticity")
|
|
ap.add_argument("--collision", default="MRT", help="SRT, TRT, or MRT")
|
|
ap.add_argument("--steps", type=int, default=60_000)
|
|
ap.add_argument("--burn", type=int, default=20_000)
|
|
ap.add_argument("--record-every", type=int, default=5)
|
|
ap.add_argument("--outlet", default="neq_extrap")
|
|
ap.add_argument(
|
|
"--out-dir",
|
|
default=os.path.join(_REPO, "tests", "output", "sah04_case9_flow_diag"),
|
|
help="Output directory for PNG + JSON",
|
|
)
|
|
ap.add_argument(
|
|
"--x-list",
|
|
type=str,
|
|
default="",
|
|
help='Comma-separated x indices (default: built-in list); e.g. "1,2,5,40,200"',
|
|
)
|
|
args = ap.parse_args()
|
|
|
|
mc = next(c for c in sah.MATRIX if c.case_id == 9)
|
|
tier = sah._TIERS[mc.tier]
|
|
ny = int(tier["ny"])
|
|
nx = int(sah._NX)
|
|
center = (sah._CX, float(tier["center_y"]))
|
|
u_max = 0.1
|
|
u0_mean = u_max / 1.5
|
|
nu = u_max * float(sah._D) / float(mc.re)
|
|
|
|
if args.x_list.strip():
|
|
x_indices = [int(s.strip()) for s in args.x_list.split(",") if s.strip()]
|
|
x_indices = [x for x in x_indices if 1 <= x < nx - 1]
|
|
else:
|
|
x_indices = _default_x_indices(nx)
|
|
|
|
lbm_path = sah._DEFAULT_LBM
|
|
cfg = sah._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"] = str(args.collision).upper()
|
|
cfg["method"]["streaming"] = "double_buffer"
|
|
cfg["method"]["les"]["enabled"] = False
|
|
cfg["method"]["outlet"]["mode"] = args.outlet
|
|
|
|
body_doc = {
|
|
"objects": [
|
|
{
|
|
"type": "cylinder",
|
|
"center": list(center),
|
|
"radius": float(sah._R_CYL),
|
|
}
|
|
]
|
|
}
|
|
|
|
tmpd = tempfile.mkdtemp(prefix="celeris_sah09_diag_")
|
|
lbm_tmp = os.path.join(tmpd, "config_lbm.json")
|
|
body_tmp = os.path.join(tmpd, "config_body.json")
|
|
sah._write_json(lbm_tmp, cfg)
|
|
sah._write_json(body_tmp, body_doc)
|
|
|
|
os.makedirs(args.out_dir, exist_ok=True)
|
|
png_profiles = os.path.join(args.out_dir, "case9_profiles_ux_uy.png")
|
|
png_vort = os.path.join(args.out_dir, "case9_omega_z_full.png")
|
|
json_path = os.path.join(args.out_dir, "case9_profiles.json")
|
|
|
|
try:
|
|
from CelerisLab import Simulation # noqa: WPS433
|
|
|
|
sim = Simulation(lbm_config_path=lbm_tmp, body_config_path=body_tmp)
|
|
sim.initialize()
|
|
stream = cuda.Stream()
|
|
rec_every = max(1, int(args.record_every))
|
|
steps = int(args.steps)
|
|
|
|
for step in range(1, 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,
|
|
)
|
|
|
|
stream.synchronize()
|
|
macro = sim.get_macroscopic()
|
|
ux = np.asarray(macro["ux"], dtype=np.float64)
|
|
uy = np.asarray(macro["uy"], dtype=np.float64)
|
|
sim.close()
|
|
|
|
y = np.arange(ny, dtype=np.float64)
|
|
u_ref = _inlet_parabolic_ref_u(y, ny, u0_mean)
|
|
|
|
title = (
|
|
f"Sah04 case 9 {args.collision} Re={mc.re} ny={ny} "
|
|
f"steps={steps} (last step)"
|
|
)
|
|
_plot_profiles(png_profiles, ux=ux, uy=uy, x_indices=x_indices, ny=ny, u0_mean=u0_mean, title=title)
|
|
sah.save_final_vorticity_png(png_vort, macro["ux"], macro["uy"], title=title + " omega_z")
|
|
|
|
profiles: Dict[str, Any] = {
|
|
"case_id": 9,
|
|
"collision": str(args.collision).upper(),
|
|
"Re": mc.re,
|
|
"nx": nx,
|
|
"ny": ny,
|
|
"steps": steps,
|
|
"burn_requested": int(args.burn),
|
|
"u_max": u_max,
|
|
"u0_mean_parabolic": u0_mean,
|
|
"x_indices": list(x_indices),
|
|
"y_lattice": y.tolist(),
|
|
"inlet_parabolic_ref_ux": u_ref.tolist(),
|
|
"profiles": {},
|
|
}
|
|
for xi in x_indices:
|
|
profiles["profiles"][str(xi)] = {
|
|
"ux": ux[:, xi].tolist(),
|
|
"uy": uy[:, xi].tolist(),
|
|
}
|
|
with open(json_path, "w", encoding="utf-8") as f:
|
|
json.dump(profiles, f, indent=2)
|
|
|
|
print("Wrote:", png_profiles, flush=True)
|
|
print("Wrote:", png_vort, flush=True)
|
|
print("Wrote:", json_path, flush=True)
|
|
finally:
|
|
shutil.rmtree(tmpd, ignore_errors=True)
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|