575 lines
18 KiB
Python
575 lines
18 KiB
Python
# CelerisLab/tests/run_inlet_channel_diagnostic.py
|
||
"""Empty-channel inlet diagnostic: field snapshots and line profiles.
|
||
|
||
Runs no-cylinder channel flows to isolate inlet / wall / outlet effects before
|
||
adding a body. See user matrix in module docstring sections A–C.
|
||
|
||
Usage::
|
||
|
||
conda run -n pycuda_3_10 python tests/run_inlet_channel_diagnostic.py --part all
|
||
conda run -n pycuda_3_10 python tests/run_inlet_channel_diagnostic.py --part a
|
||
conda run -n pycuda_3_10 python tests/run_inlet_channel_diagnostic.py --part b --nx 401 --ny 201
|
||
|
||
Output (default ``tests/output/inlet_channel_diag/``)::
|
||
|
||
A_baseline/{SRT,MRT}/snapshots/step_XXXXXX.{npz,png}
|
||
B_matrix/caseNN_.../snapshots/...
|
||
B_matrix/caseNN_.../lines/ux_lines_stepXXXXXX.png
|
||
summary.csv
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import csv
|
||
import json
|
||
import os
|
||
import sys
|
||
import tempfile
|
||
from dataclasses import dataclass
|
||
from typing import Any, Dict, Iterable, 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")
|
||
|
||
# Default snapshot steps for parts A and B.
|
||
DEFAULT_SNAPSHOT_STEPS: Tuple[int, ...] = (100, 500, 1000, 1500, 2000)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class CaseSpec:
|
||
"""One empty-channel configuration."""
|
||
|
||
case_id: str
|
||
label: str
|
||
inlet_scheme: str
|
||
y_wall_bc: str
|
||
outlet_mode: str
|
||
collision: str = "MRT"
|
||
|
||
|
||
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:
|
||
"""ωz = ∂uy/∂x − ∂ux/∂y on (ny, nx) arrays."""
|
||
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}"),
|
||
]
|
||
|
||
|
||
def _build_cfg(
|
||
base: dict,
|
||
*,
|
||
nx: int,
|
||
ny: int,
|
||
collision: str,
|
||
inlet_scheme: str,
|
||
inlet_profile: str,
|
||
y_wall_bc: str,
|
||
outlet_mode: str,
|
||
velocity: float,
|
||
viscosity: float,
|
||
) -> dict:
|
||
cfg = json.loads(json.dumps(base))
|
||
cfg["grid"]["nx"] = int(nx)
|
||
cfg["grid"]["ny"] = int(ny)
|
||
cfg["grid"]["nz"] = 1
|
||
cfg["physics"]["velocity"] = float(velocity)
|
||
cfg["physics"]["viscosity"] = float(viscosity)
|
||
cfg["physics"]["rho"] = 1.0
|
||
cfg["method"]["collision"] = str(collision).upper()
|
||
cfg["method"]["streaming"] = "double_buffer"
|
||
cfg["method"]["store_precision"] = "FP32"
|
||
cfg["method"]["les"]["enabled"] = False
|
||
cfg["method"]["inlet"]["profile"] = str(inlet_profile)
|
||
cfg["method"]["inlet"]["scheme"] = str(inlet_scheme)
|
||
cfg["method"]["y_wall_bc"] = str(y_wall_bc)
|
||
cfg["method"]["outlet"]["mode"] = str(outlet_mode)
|
||
return cfg
|
||
|
||
|
||
def _field_stats(rho: np.ndarray, ux: np.ndarray, vort: np.ndarray) -> Dict[str, float]:
|
||
def _f(a: np.ndarray) -> float:
|
||
b = a[np.isfinite(a)]
|
||
return float("nan") if b.size == 0 else float(np.max(np.abs(b)))
|
||
|
||
return {
|
||
"rho_min": float(np.nanmin(rho)) if np.isfinite(rho).any() else float("nan"),
|
||
"rho_max": float(np.nanmax(rho)) if np.isfinite(rho).any() else float("nan"),
|
||
"ux_max": _f(ux),
|
||
"vort_max": _f(vort),
|
||
"finite": bool(np.isfinite(rho).all() and np.isfinite(ux).all()),
|
||
}
|
||
|
||
|
||
def _save_snapshot_npz(
|
||
path: str,
|
||
*,
|
||
step: int,
|
||
rho: np.ndarray,
|
||
ux: np.ndarray,
|
||
uy: np.ndarray,
|
||
vort: np.ndarray,
|
||
meta: dict,
|
||
) -> None:
|
||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||
np.savez_compressed(
|
||
path,
|
||
rho=np.asarray(rho, dtype=np.float32),
|
||
ux=np.asarray(ux, dtype=np.float32),
|
||
uy=np.asarray(uy, dtype=np.float32),
|
||
vort=np.asarray(vort, dtype=np.float32),
|
||
step=np.int32(step),
|
||
meta_json=np.asarray(json.dumps(meta)),
|
||
)
|
||
|
||
|
||
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
|
||
fig, ax = plt.subplots(figsize=(min(16.0, max(8.0, nx / 80.0)), min(8.0, max(3.0, ny / 50.0))))
|
||
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=140, 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,
|
||
case_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 / 60.0)), 7.0), sharex=True)
|
||
for y_idx, y_lab in y_rows:
|
||
y_idx = int(np.clip(y_idx, 0, ny - 1))
|
||
axes[0].plot(x, ux[y_idx, :], label=y_lab, linewidth=1.2)
|
||
axes[1].plot(x, rho[y_idx, :], label=y_lab, linewidth=1.2)
|
||
|
||
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"{case_label} — line profiles 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 _run_channel(
|
||
case: CaseSpec,
|
||
*,
|
||
base_cfg: dict,
|
||
nx: int,
|
||
ny: int,
|
||
velocity: float,
|
||
viscosity: float,
|
||
out_root: str,
|
||
snapshot_steps: Sequence[int],
|
||
max_step: int,
|
||
save_png: bool,
|
||
save_lines: bool,
|
||
stop_on_nan: bool,
|
||
) -> List[Dict[str, Any]]:
|
||
"""Run one case; write snapshots and optional line plots. Return summary rows."""
|
||
sys.path.insert(0, os.path.join(_REPO, "src"))
|
||
from CelerisLab import Simulation # noqa: WPS433
|
||
|
||
cfg = _build_cfg(
|
||
base_cfg,
|
||
nx=nx,
|
||
ny=ny,
|
||
collision=case.collision,
|
||
inlet_scheme=case.inlet_scheme,
|
||
inlet_profile="uniform",
|
||
y_wall_bc=case.y_wall_bc,
|
||
outlet_mode=case.outlet_mode,
|
||
velocity=velocity,
|
||
viscosity=viscosity,
|
||
)
|
||
bdoc = {"objects": []}
|
||
|
||
run_dir = os.path.join(out_root, case.case_id)
|
||
snap_dir = os.path.join(run_dir, "snapshots")
|
||
line_dir = os.path.join(run_dir, "lines")
|
||
os.makedirs(snap_dir, exist_ok=True)
|
||
if save_lines:
|
||
os.makedirs(line_dir, exist_ok=True)
|
||
|
||
tmpd = tempfile.mkdtemp(prefix="inlet_diag_")
|
||
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, bdoc)
|
||
|
||
meta_base = {
|
||
"case_id": case.case_id,
|
||
"label": case.label,
|
||
"nx": nx,
|
||
"ny": ny,
|
||
"inlet_scheme": case.inlet_scheme,
|
||
"inlet_profile": "uniform",
|
||
"y_wall_bc": case.y_wall_bc,
|
||
"outlet_mode": case.outlet_mode,
|
||
"collision": case.collision,
|
||
"velocity": velocity,
|
||
"viscosity": viscosity,
|
||
}
|
||
_write_json(os.path.join(run_dir, "case_meta.json"), meta_base)
|
||
|
||
sim = Simulation(lbm_config_path=lbm_tmp, body_config_path=body_tmp)
|
||
sim.initialize()
|
||
|
||
y_rows = _line_y_indices(ny)
|
||
want_steps = sorted({int(s) for s in snapshot_steps if 0 < int(s) <= max_step})
|
||
next_snap = 0
|
||
rows: List[Dict[str, Any]] = []
|
||
|
||
for step in range(1, max_step + 1):
|
||
sim.step(1)
|
||
if next_snap < len(want_steps) and step == want_steps[next_snap]:
|
||
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)
|
||
stats = _field_stats(rho, ux, vort)
|
||
stem = f"step_{step:06d}"
|
||
meta = {**meta_base, "step": step, **stats}
|
||
npz_path = os.path.join(snap_dir, f"{stem}.npz")
|
||
_save_snapshot_npz(
|
||
npz_path,
|
||
step=step,
|
||
rho=rho,
|
||
ux=ux,
|
||
uy=uy,
|
||
vort=vort,
|
||
meta=meta,
|
||
)
|
||
if save_png:
|
||
_save_field_pngs(
|
||
snap_dir,
|
||
stem,
|
||
rho=rho,
|
||
ux=ux,
|
||
vort=vort,
|
||
title=f"{case.label} step {step}",
|
||
)
|
||
if save_lines:
|
||
line_png = os.path.join(line_dir, f"lines_{stem}.png")
|
||
_save_line_plots(
|
||
line_png,
|
||
rho=rho,
|
||
ux=ux,
|
||
step=step,
|
||
case_label=case.label,
|
||
y_rows=y_rows,
|
||
)
|
||
# Also save raw 1D data for replotting.
|
||
line_npz = os.path.join(line_dir, f"lines_{stem}.npz")
|
||
payload = {"x": np.arange(nx, dtype=np.float32)}
|
||
for y_idx, y_lab in y_rows:
|
||
payload[f"ux_{y_lab}"] = ux[y_idx, :].astype(np.float32)
|
||
payload[f"rho_{y_lab}"] = rho[y_idx, :].astype(np.float32)
|
||
payload["step"] = np.int32(step)
|
||
np.savez_compressed(line_npz, **payload)
|
||
|
||
rows.append(
|
||
{
|
||
"case_id": case.case_id,
|
||
"label": case.label,
|
||
"collision": case.collision,
|
||
"inlet_scheme": case.inlet_scheme,
|
||
"y_wall_bc": case.y_wall_bc,
|
||
"outlet_mode": case.outlet_mode,
|
||
"step": step,
|
||
**stats,
|
||
"npz": npz_path,
|
||
}
|
||
)
|
||
if stop_on_nan and not stats["finite"]:
|
||
sim.close()
|
||
rows[-1]["stopped_early"] = True
|
||
return rows
|
||
next_snap += 1
|
||
|
||
sim.close()
|
||
return rows
|
||
|
||
|
||
def _part_a_cases() -> List[CaseSpec]:
|
||
# Kan99b-style baseline: zou_he + free_slip + neq_extrap; SRT and MRT.
|
||
base = CaseSpec(
|
||
case_id="",
|
||
label="",
|
||
inlet_scheme="zou_he_local",
|
||
y_wall_bc="free_slip",
|
||
outlet_mode="neq_extrap",
|
||
)
|
||
out: List[CaseSpec] = []
|
||
for coll in ("SRT", "MRT"):
|
||
cid = f"A_{coll.lower()}_zouhe_fs_neq"
|
||
out.append(
|
||
CaseSpec(
|
||
case_id=cid,
|
||
label=f"A baseline {coll} zou_he / free_slip / neq_extrap",
|
||
inlet_scheme=base.inlet_scheme,
|
||
y_wall_bc=base.y_wall_bc,
|
||
outlet_mode=base.outlet_mode,
|
||
collision=coll,
|
||
)
|
||
)
|
||
return out
|
||
|
||
|
||
def _part_b_cases() -> List[CaseSpec]:
|
||
return [
|
||
CaseSpec(
|
||
case_id="B_case01_zouhe_fs_neq",
|
||
label="1 zou_he / free_slip / neq_extrap",
|
||
inlet_scheme="zou_he_local",
|
||
y_wall_bc="free_slip",
|
||
outlet_mode="neq_extrap",
|
||
collision="MRT",
|
||
),
|
||
CaseSpec(
|
||
case_id="B_case02_zouhe_bb_neq",
|
||
label="2 zou_he / bounce_back / neq_extrap",
|
||
inlet_scheme="zou_he_local",
|
||
y_wall_bc="bounce_back",
|
||
outlet_mode="neq_extrap",
|
||
collision="MRT",
|
||
),
|
||
CaseSpec(
|
||
case_id="B_case03_zouhe_fs_zgrad",
|
||
label="3 zou_he / free_slip / zero_gradient",
|
||
inlet_scheme="zou_he_local",
|
||
y_wall_bc="free_slip",
|
||
outlet_mode="zero_gradient",
|
||
collision="MRT",
|
||
),
|
||
CaseSpec(
|
||
case_id="B_case04_stab_fs_neq",
|
||
label="4 channel_stabilized / free_slip / neq_extrap",
|
||
inlet_scheme="channel_stabilized",
|
||
y_wall_bc="free_slip",
|
||
outlet_mode="neq_extrap",
|
||
collision="MRT",
|
||
),
|
||
]
|
||
|
||
|
||
def _write_summary_csv(path: str, rows: Sequence[Dict[str, Any]]) -> None:
|
||
if not rows:
|
||
return
|
||
keys = [
|
||
"case_id",
|
||
"label",
|
||
"collision",
|
||
"inlet_scheme",
|
||
"y_wall_bc",
|
||
"outlet_mode",
|
||
"step",
|
||
"finite",
|
||
"rho_min",
|
||
"rho_max",
|
||
"ux_max",
|
||
"vort_max",
|
||
"stopped_early",
|
||
]
|
||
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
|
||
with open(path, "w", encoding="utf-8", newline="") as f:
|
||
w = csv.DictWriter(f, fieldnames=keys, extrasaction="ignore")
|
||
w.writeheader()
|
||
for r in rows:
|
||
w.writerow(r)
|
||
|
||
|
||
def main() -> int:
|
||
ap = argparse.ArgumentParser(description="Empty-channel inlet diagnostic (fields + line profiles)")
|
||
ap.add_argument(
|
||
"--part",
|
||
choices=("a", "b", "all"),
|
||
default="all",
|
||
help="A=baseline SRT/MRT; B=4-case matrix; all=both",
|
||
)
|
||
ap.add_argument("--nx", type=int, default=401, help="Channel length (lattice)")
|
||
ap.add_argument("--ny", type=int, default=201, help="Channel height (lattice)")
|
||
ap.add_argument("--velocity", type=float, default=0.03, help="Uniform inlet U0")
|
||
ap.add_argument("--viscosity", type=float, default=0.009, help="Kinematic viscosity")
|
||
ap.add_argument(
|
||
"--out-dir",
|
||
type=str,
|
||
default=os.path.join(_REPO, "tests", "output", "inlet_channel_diag"),
|
||
)
|
||
ap.add_argument(
|
||
"--steps",
|
||
type=str,
|
||
default="",
|
||
help="Comma-separated snapshot steps (default: 100,500,1000,1500,2000)",
|
||
)
|
||
ap.add_argument("--no-png", action="store_true", help="Skip rho/ux/vort PNG maps")
|
||
ap.add_argument("--no-lines", action="store_true", help="Skip ux/rho line-profile plots")
|
||
ap.add_argument(
|
||
"--continue-on-nan",
|
||
action="store_true",
|
||
help="Keep stepping after non-finite fields (default: stop case early)",
|
||
)
|
||
args = ap.parse_args()
|
||
|
||
if not os.path.isfile(_DEFAULT_LBM):
|
||
print(f"Missing config: {_DEFAULT_LBM}", file=sys.stderr)
|
||
return 2
|
||
|
||
if args.steps.strip():
|
||
snap_steps = tuple(int(s.strip()) for s in args.steps.split(",") if s.strip())
|
||
else:
|
||
snap_steps = DEFAULT_SNAPSHOT_STEPS
|
||
max_step = max(snap_steps)
|
||
|
||
base_cfg = _load_json(_DEFAULT_LBM)
|
||
out_dir = os.path.abspath(args.out_dir)
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
|
||
cases: List[CaseSpec] = []
|
||
if args.part in ("a", "all"):
|
||
cases.extend(_part_a_cases())
|
||
if args.part in ("b", "all"):
|
||
cases.extend(_part_b_cases())
|
||
|
||
save_png = not args.no_png
|
||
# Part A: field maps only; Part B: fields + line plots.
|
||
all_rows: List[Dict[str, Any]] = []
|
||
|
||
for case in cases:
|
||
part = "A_baseline" if case.case_id.startswith("A_") else "B_matrix"
|
||
root = os.path.join(out_dir, part)
|
||
save_lines = not args.no_lines and part == "B_matrix"
|
||
print(f"--- {case.case_id}: {case.label} ({case.collision}) ---", flush=True)
|
||
rows = _run_channel(
|
||
case,
|
||
base_cfg=base_cfg,
|
||
nx=args.nx,
|
||
ny=args.ny,
|
||
velocity=args.velocity,
|
||
viscosity=args.viscosity,
|
||
out_root=root,
|
||
snapshot_steps=snap_steps,
|
||
max_step=max_step,
|
||
save_png=save_png,
|
||
save_lines=save_lines,
|
||
stop_on_nan=not args.continue_on_nan,
|
||
)
|
||
all_rows.extend(rows)
|
||
for r in rows:
|
||
fin = "OK" if r.get("finite") else "NONFINITE"
|
||
print(
|
||
f" step {r['step']:5d} {fin} rho=[{r['rho_min']:.4f},{r['rho_max']:.4f}] "
|
||
f"ux_max={r['ux_max']:.4f} vort_max={r['vort_max']:.4f}",
|
||
flush=True,
|
||
)
|
||
if rows and rows[-1].get("stopped_early"):
|
||
print(" (stopped early: non-finite fields)", flush=True)
|
||
|
||
summary_path = os.path.join(out_dir, "summary.csv")
|
||
_write_summary_csv(summary_path, all_rows)
|
||
manifest = {
|
||
"snapshot_steps": list(snap_steps),
|
||
"nx": args.nx,
|
||
"ny": args.ny,
|
||
"velocity": args.velocity,
|
||
"viscosity": args.viscosity,
|
||
"line_y_indices": [{"y": y, "label": lab} for y, lab in _line_y_indices(args.ny)],
|
||
"cases": [c.case_id for c in cases],
|
||
}
|
||
_write_json(os.path.join(out_dir, "manifest.json"), manifest)
|
||
|
||
print(f"Wrote: {summary_path}", flush=True)
|
||
print(f"Wrote: {os.path.join(out_dir, 'manifest.json')}", flush=True)
|
||
print(f"Output root: {out_dir}", flush=True)
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|