310 lines
10 KiB
Python
310 lines
10 KiB
Python
# CelerisLab/tests/run_kan99b_streakline.py
|
|
"""Kan99b streakline demo using CelerisLab common streakline engine.
|
|
|
|
Modes:
|
|
- online: run CFD and compute streakline in memory (no velocity dump).
|
|
- offline: load velocity snapshots and reconstruct streakline.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import tempfile
|
|
from dataclasses import dataclass
|
|
from typing import List, Tuple
|
|
|
|
import numpy as np
|
|
|
|
from CelerisLab import Simulation
|
|
from CelerisLab.common.streakline import (
|
|
FlowFrame,
|
|
IntegratorConfig,
|
|
ReleaseConfig,
|
|
build_release_points,
|
|
estimate_sampling_plan,
|
|
render_streakline_density,
|
|
run_streakline_offline,
|
|
run_streakline_online,
|
|
)
|
|
|
|
_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
|
_DEFAULT_LBM = os.path.join(_REPO, "src", "CelerisLab", "configs", "config_lbm.json")
|
|
|
|
U_INF = 0.03
|
|
D_LATTICE = 30.0
|
|
R_LATTICE = 15.0
|
|
KAN99B_ST_REF = 0.1655
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DomainSpec:
|
|
key: str
|
|
nx: int
|
|
ny: int
|
|
center: Tuple[float, float]
|
|
|
|
|
|
def _domain_specs() -> dict:
|
|
return {
|
|
"S": DomainSpec("S", 1081, 481, (360.0, 240.0)),
|
|
"M": DomainSpec("M", 1351, 601, (450.0, 300.0)),
|
|
"L": DomainSpec("L", 1801, 721, (600.0, 360.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:
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
json.dump(payload, f, indent=2)
|
|
|
|
|
|
def _nu_from_re(reynolds: float) -> float:
|
|
return U_INF * D_LATTICE / float(reynolds)
|
|
|
|
|
|
def _omega_body(alpha: float) -> float:
|
|
return 2.0 * float(alpha) * U_INF / D_LATTICE
|
|
|
|
|
|
def _build_cfg(base_cfg: dict, *, nx: int, ny: int, re: float, inlet_scheme: str) -> dict:
|
|
cfg = json.loads(json.dumps(base_cfg))
|
|
cfg["grid"]["nx"] = int(nx)
|
|
cfg["grid"]["ny"] = int(ny)
|
|
cfg["grid"]["nz"] = 1
|
|
cfg["physics"]["velocity"] = float(U_INF)
|
|
cfg["physics"]["viscosity"] = float(_nu_from_re(re))
|
|
cfg["physics"]["rho"] = 1.0
|
|
cfg["method"]["collision"] = "MRT"
|
|
cfg["method"]["streaming"] = "double_buffer"
|
|
cfg["method"]["store_precision"] = "FP32"
|
|
cfg["method"]["ddf_shifting"] = False
|
|
cfg["method"]["les"]["enabled"] = False
|
|
cfg["method"]["inlet"]["profile"] = "uniform"
|
|
cfg["method"]["inlet"]["scheme"] = str(inlet_scheme)
|
|
cfg["method"]["outlet"]["mode"] = "neq_extrap"
|
|
cfg["method"]["y_wall_bc"] = "free_slip"
|
|
return cfg
|
|
|
|
|
|
def _build_simulation(*, domain: DomainSpec, re: float, alpha: float, inlet_scheme: str) -> Simulation:
|
|
base_cfg = _load_json(_DEFAULT_LBM)
|
|
cfg = _build_cfg(base_cfg, nx=domain.nx, ny=domain.ny, re=re, inlet_scheme=inlet_scheme)
|
|
body_doc = {
|
|
"objects": [
|
|
{
|
|
"type": "cylinder",
|
|
"center": [float(domain.center[0]), float(domain.center[1])],
|
|
"radius": float(R_LATTICE),
|
|
"omega": float(_omega_body(alpha)),
|
|
}
|
|
]
|
|
}
|
|
tmpd = tempfile.mkdtemp(prefix="celeris_streakline_online_")
|
|
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)
|
|
sim = Simulation(lbm_config_path=lbm_tmp, body_config_path=body_tmp)
|
|
sim.bodies.get(0).state.omega = np.float32(_omega_body(alpha))
|
|
sim.initialize()
|
|
return sim
|
|
|
|
|
|
def _default_base_release_points(center: Tuple[float, float]) -> np.ndarray:
|
|
x_rel = float(center[0] - 6.0 * D_LATTICE)
|
|
y0 = float(center[1])
|
|
return np.array(
|
|
[
|
|
[x_rel, y0 - 18.0],
|
|
[x_rel, y0 - 6.0],
|
|
[x_rel, y0 + 6.0],
|
|
[x_rel, y0 + 18.0],
|
|
],
|
|
dtype=np.float64,
|
|
)
|
|
|
|
|
|
def _load_offline_frames(snapshot_dir: str) -> List[FlowFrame]:
|
|
files = sorted(
|
|
[
|
|
os.path.join(snapshot_dir, name)
|
|
for name in os.listdir(snapshot_dir)
|
|
if name.endswith(".npz") and name.startswith("vel_step")
|
|
]
|
|
)
|
|
frames: List[FlowFrame] = []
|
|
for path in files:
|
|
data = np.load(path)
|
|
step = int(np.asarray(data["step"]).reshape(-1)[0])
|
|
frames.append(
|
|
FlowFrame(
|
|
step=step,
|
|
ux=np.asarray(data["ux"], dtype=np.float64),
|
|
uy=np.asarray(data["uy"], dtype=np.float64),
|
|
)
|
|
)
|
|
return frames
|
|
|
|
|
|
def main() -> int:
|
|
ap = argparse.ArgumentParser(description="Kan99b streakline demo (online/offline)")
|
|
ap.add_argument("--mode", default="online", choices=("online", "offline"))
|
|
ap.add_argument("--domain", default="M", choices=("S", "M", "L"))
|
|
ap.add_argument("--re", type=float, default=100.0)
|
|
ap.add_argument("--alpha", type=float, default=1.0)
|
|
ap.add_argument("--inlet-scheme", default="regularized", choices=("regularized", "zou_he_local"))
|
|
ap.add_argument("--start-step", type=int, default=60_000)
|
|
ap.add_argument("--sample-every", type=int, default=0, help="0 uses recommended value.")
|
|
ap.add_argument("--n-snapshots", type=int, default=0, help="0 uses recommended value.")
|
|
ap.add_argument("--snapshot-dir", type=str, default="", help="Required in offline mode.")
|
|
ap.add_argument("--release-mode", default="strip", choices=("point", "line", "strip"))
|
|
ap.add_argument("--line-span", type=float, default=0.0)
|
|
ap.add_argument("--line-count", type=int, default=1)
|
|
ap.add_argument(
|
|
"--downstream-count",
|
|
type=int,
|
|
default=5,
|
|
help="Particles per seed along +x (denser streak).",
|
|
)
|
|
ap.add_argument(
|
|
"--downstream-spacing",
|
|
type=float,
|
|
default=1.0,
|
|
help="Lattice spacing between x-staggered release points.",
|
|
)
|
|
ap.add_argument("--inject-per-seed", type=int, default=2)
|
|
ap.add_argument("--alpha-t", type=float, default=0.2)
|
|
ap.add_argument("--alpha-x", type=float, default=0.4)
|
|
ap.add_argument("--diffusion-coeff", type=float, default=0.0)
|
|
ap.add_argument(
|
|
"--out-dir",
|
|
type=str,
|
|
default=os.path.join(_REPO, "tests", "output", "streakline", "kan99b_k2_online"),
|
|
)
|
|
args = ap.parse_args()
|
|
|
|
domain = _domain_specs()[args.domain]
|
|
out_dir = os.path.abspath(args.out_dir)
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
|
|
plan = estimate_sampling_plan(st_ref=KAN99B_ST_REF, diameter=D_LATTICE, u_ref=U_INF)
|
|
sample_every = int(args.sample_every) if int(args.sample_every) > 0 else int(plan["save_every_recommended"])
|
|
n_snapshots = int(args.n_snapshots) if int(args.n_snapshots) > 0 else int(plan["snapshot_count_recommended"])
|
|
|
|
release_cfg = ReleaseConfig(
|
|
mode=args.release_mode,
|
|
line_span=float(args.line_span),
|
|
line_count=max(1, int(args.line_count)),
|
|
downstream_count=max(1, int(args.downstream_count)),
|
|
downstream_spacing=float(args.downstream_spacing),
|
|
inject_per_seed=max(1, int(args.inject_per_seed)),
|
|
)
|
|
integrator_cfg = IntegratorConfig(
|
|
alpha_t=float(args.alpha_t),
|
|
alpha_x=float(args.alpha_x),
|
|
diffusion_coeff=float(args.diffusion_coeff),
|
|
)
|
|
base_release_points = _default_base_release_points(domain.center)
|
|
dense_release_points = build_release_points(base_release_points, release_cfg)
|
|
|
|
if args.mode == "online":
|
|
sim = _build_simulation(
|
|
domain=domain,
|
|
re=float(args.re),
|
|
alpha=float(args.alpha),
|
|
inlet_scheme=args.inlet_scheme,
|
|
)
|
|
try:
|
|
particles, ages, diag = run_streakline_online(
|
|
sim,
|
|
start_step=int(args.start_step),
|
|
sample_every=int(sample_every),
|
|
n_samples=int(n_snapshots),
|
|
release_points=base_release_points,
|
|
release_cfg=release_cfg,
|
|
integrator_cfg=integrator_cfg,
|
|
solid_center=domain.center,
|
|
solid_radius=R_LATTICE,
|
|
)
|
|
finally:
|
|
sim.close()
|
|
frame_count = int(n_snapshots)
|
|
frame_source = "in-memory online sampling"
|
|
else:
|
|
if not args.snapshot_dir:
|
|
raise ValueError("--snapshot-dir is required in offline mode.")
|
|
frames = _load_offline_frames(os.path.abspath(args.snapshot_dir))
|
|
if len(frames) < 2:
|
|
raise RuntimeError("Offline mode needs at least two velocity snapshots.")
|
|
particles, ages, diag = run_streakline_offline(
|
|
frames,
|
|
nx=domain.nx,
|
|
ny=domain.ny,
|
|
release_points=base_release_points,
|
|
release_cfg=release_cfg,
|
|
integrator_cfg=integrator_cfg,
|
|
solid_center=domain.center,
|
|
solid_radius=R_LATTICE,
|
|
)
|
|
frame_count = len(frames)
|
|
frame_source = os.path.abspath(args.snapshot_dir)
|
|
|
|
render_info = render_streakline_density(
|
|
particles,
|
|
ages,
|
|
nx=domain.nx,
|
|
ny=domain.ny,
|
|
out_path=os.path.join(out_dir, "streakline.png"),
|
|
release_points=dense_release_points,
|
|
solid_center=domain.center,
|
|
solid_radius=R_LATTICE,
|
|
age_decay_steps=integrator_cfg.age_decay_steps,
|
|
blur_sigma=1.2,
|
|
title=f"Kan99b streakline ({args.mode}, {args.release_mode})",
|
|
)
|
|
|
|
meta = {
|
|
"case": {
|
|
"domain": args.domain,
|
|
"re": float(args.re),
|
|
"alpha": float(args.alpha),
|
|
"inlet_scheme": args.inlet_scheme,
|
|
"collision": "MRT",
|
|
},
|
|
"mode": args.mode,
|
|
"sampling_estimate": plan,
|
|
"sampling_used": {
|
|
"start_step": int(args.start_step),
|
|
"sample_every": int(sample_every),
|
|
"n_snapshots": int(n_snapshots),
|
|
"frame_source": frame_source,
|
|
"frames_used": int(frame_count),
|
|
},
|
|
"release": {
|
|
"config": vars(args),
|
|
"base_points": base_release_points.tolist(),
|
|
"dense_point_count": int(dense_release_points.shape[0]),
|
|
"dense_points_preview": dense_release_points[: min(12, dense_release_points.shape[0])].tolist(),
|
|
},
|
|
"diagnostics": diag,
|
|
"render": render_info,
|
|
}
|
|
_write_json(os.path.join(out_dir, "streakline_meta.json"), meta)
|
|
|
|
print(f"Recommended sample_every: {plan['save_every_recommended']} steps")
|
|
print(f"Recommended snapshots: {plan['snapshot_count_recommended']}")
|
|
print(f"Mode: {args.mode} | frames used: {frame_count}")
|
|
print(f"Dense release points: {dense_release_points.shape[0]}")
|
|
print(f"Output image: {render_info['image_path']}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|