CelerisLab/tests/postproc/run_exp_ctrl_matrix_vorticity.py

476 lines
15 KiB
Python

# CelerisLab/tests/postproc/run_exp_ctrl_matrix_vorticity.py
"""Batch vorticity images for three-cylinder control matrix (exp_ctrl_matrix.md).
Usage::
# Single case
conda run -n pycuda_3_10 python tests/postproc/run_exp_ctrl_matrix_vorticity.py \\
--cases C0 --batch 10 --device-id 0
# All cases
conda run -n pycuda_3_10 python tests/postproc/run_exp_ctrl_matrix_vorticity.py \\
--batch 10 --device-id 0
# Full 100k steps, no batching (default)
conda run -n pycuda_3_10 python tests/postproc/run_exp_ctrl_matrix_vorticity.py
"""
from __future__ import annotations
import argparse
import json
import math
import os
import sys
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Tuple
import numpy as np
_REPO = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(_REPO / "src"))
from CelerisLab import Simulation
from CelerisLab.common.render import (
compute_vorticity,
render_vorticity_field,
)
from CelerisLab.common.preprocess import cylinders_from_triangle_layout
INLET_U_PHYS_M_S = 0.009028
CYLINDER_DIAMETER_M = 0.010
CENTER_SPACING_M = 0.015
DIAMETER_CELLS = 20.0
RAMP_TIME_S = 5.0
INITIAL_ACTIONS_M_S = (0.0, 0.0, 0.0)
OMEGA_SIGN_FROM_ACTION = -1.0
VORT_VMIN = -0.003
VORT_VMAX = 0.003
CONFIG_PATH = _REPO / "src/CelerisLab/configs/config_lbm_three_cylinder_triangle.json"
DEFAULT_OUT = _REPO / "tests" / "output" / "exp_ctrl_matrix_vort_ny300"
FIXED_STEPS = 100000 # keep constant while grid changes
# Body order: 0=apex, 1=rear-lower(y_lower), 2=rear-upper(y_upper); swap action2/action3 targets.
SWAP_ACTION23_BODIES = True
# From tests/exp_ctrl_matrix.md (SIGNAL_FEATURES0 .. 6)
CONTROL_CASES: List[Tuple[str, str, Dict[str, Any]]] = [
(
"C0",
"no_ctrl",
{
"action1": {"mean": 0.0, "components": [(0.1354, 0.0, 1.600)]},
"action2": {"mean": 0.0, "components": [(0.1354, 0.0, 2.099)]},
"action3": {"mean": 0.0, "components": [(0.1354, 0.0, 1.639)]},
},
),
(
"C1",
"stealth",
{
"action1": {"mean": 0.0, "components": [(0.1354, 0.0, 1.600)]},
"action2": {"mean": -0.01806, "components": [(0.1354, 0.0, 2.099)]},
"action3": {"mean": 0.01806, "components": [(0.1354, 0.0, 1.639)]},
},
),
(
"C2",
"deceit",
{
"action1": {"mean": 0.0, "components": [(0.1354, 0.0026, 1.600)]},
"action2": {
"mean": -0.008730,
"components": [(0.1354, 0.0045, 2.099), (0.2708, 0.0010, 0.612)],
},
"action3": {
"mean": 0.008730,
"components": [(0.1354, 0.0045, 1.639), (0.2708, 0.0010, -2.962)],
},
},
),
(
"C3",
"deceit_multi",
{
"action1": {
"mean": 0.0,
"components": [(0.1354, 0.0029, -2.619), (0.2708, 0.0008, 2.856)],
},
"action2": {
"mean": -0.0140,
"components": [
(0.1354, 0.0050, -0.933),
(0.2708, 0.0010, 0.801),
(0.1806, 0.0003, 1.854),
],
},
"action3": {
"mean": 0.014,
"components": [
(0.1354, 0.0050, -1.398),
(0.2708, 0.0010, 2.208),
(0.1806, 0.0003, 1.810),
],
},
},
),
(
"C4",
"deceit_f1p5",
{
"action1": {"mean": 0.0, "components": [(0.2031, 0.0026, 1.600)]},
"action2": {
"mean": -0.008730,
"components": [(0.2031, 0.0045, 2.099), (0.4062, 0.0010, 0.612)],
},
"action3": {
"mean": 0.008730,
"components": [(0.2031, 0.0045, 1.639), (0.4062, 0.0010, -2.962)],
},
},
),
(
"C5",
"deceit_multi_f1p5",
{
"action1": {
"mean": 0.0,
"components": [(0.2031, 0.0029, -2.619), (0.4062, 0.0008, 2.856)],
},
"action2": {
"mean": -0.0140,
"components": [
(0.2031, 0.0050, -0.933),
(0.4062, 0.0010, 0.801),
(0.2709, 0.0003, 1.854),
],
},
"action3": {
"mean": 0.014,
"components": [
(0.2031, 0.0050, -1.398),
(0.4062, 0.0010, 2.208),
(0.2709, 0.0003, 1.810),
],
},
},
),
(
"C6",
"deceit_f2",
{
"action1": {
"mean": 0.0,
"components": [(0.2708, 0.0044, -2.619), (0.8124, 0.0012, 2.856)],
},
"action2": {
"mean": -0.014,
"components": [
(0.2708, 0.0075, -0.933),
(0.8124, 0.0015, 0.801),
(0.5418, 0.0005, 1.854),
],
},
"action3": {
"mean": 0.014,
"components": [
(0.2708, 0.0075, -1.398),
(0.8124, 0.0015, 2.208),
(0.5418, 0.0005, 1.810),
],
},
},
),
]
def _ensure_compat_config(config_path: Path, preferred_scheme: str = "regularized") -> str:
with config_path.open("r", encoding="utf-8") as f:
cfg = json.load(f)
method = cfg.setdefault("method", {})
inlet = method.setdefault("inlet", {})
outlet = method.setdefault("outlet", {})
changed = False
if "scheme" not in inlet:
inlet["scheme"] = preferred_scheme
changed = True
if "regularized_neq_damp" not in inlet:
inlet["regularized_neq_damp"] = 0.5
changed = True
if "blend_alpha" not in outlet:
outlet["blend_alpha"] = 0.7
changed = True
if "backflow_clamp" not in outlet:
outlet["backflow_clamp"] = True
changed = True
if not changed:
return str(config_path)
tmp = tempfile.NamedTemporaryFile(
mode="w", suffix="_compat_lbm.json", delete=False, encoding="utf-8"
)
with tmp:
json.dump(cfg, tmp, indent=4)
return tmp.name
def _triangle_layout(cfg) -> dict:
dx_phys = CYLINDER_DIAMETER_M / DIAMETER_CELLS
spacing_lb = CENTER_SPACING_M / dx_phys
radius_lb = DIAMETER_CELLS / 2.0
y_center = 0.5 * (cfg.ny - 1)
x_cluster_center = cfg.nx / 3.0
x_apex = x_cluster_center - (math.sqrt(3.0) / 3.0) * spacing_lb
x_rear = x_apex + (math.sqrt(3.0) / 2.0) * spacing_lb
return {
"x_apex": x_apex,
"x_rear": x_rear,
"y_center": y_center,
"y_upper": y_center + 0.5 * spacing_lb,
"y_lower": y_center - 0.5 * spacing_lb,
"radius_lb": radius_lb,
}
def _add_triangle_cylinders(sim: Simulation) -> dict:
layout = _triangle_layout(sim.lbm_cfg)
sim.add_body("circle", center=(layout["x_apex"], layout["y_center"]),
radius=layout["radius_lb"])
sim.add_body("circle", center=(layout["x_rear"], layout["y_lower"]),
radius=layout["radius_lb"])
sim.add_body("circle", center=(layout["x_rear"], layout["y_upper"]),
radius=layout["radius_lb"])
return layout
def _generate_signal(t_phys: float, feature: dict) -> float:
value = float(feature["mean"])
for freq_hz, amp, phase in feature["components"]:
value += amp * math.cos(2.0 * math.pi * freq_hz * t_phys + phase)
return value
def _ramp_factor(elapsed_s: float) -> float:
if elapsed_s <= 0.0:
return 0.0
if elapsed_s >= RAMP_TIME_S:
return 1.0
x = elapsed_s / RAMP_TIME_S
return 0.5 * (1.0 - math.cos(math.pi * x))
def _actions_at_time(t_phys: float, features: dict) -> Tuple[float, float, float]:
s1 = _generate_signal(t_phys, features["action1"])
s2 = _generate_signal(t_phys, features["action2"])
s3 = _generate_signal(t_phys, features["action3"])
r = _ramp_factor(t_phys)
a1 = INITIAL_ACTIONS_M_S[0] * (1.0 - r) + s1 * r
a2 = INITIAL_ACTIONS_M_S[1] * (1.0 - r) + s2 * r
a3 = INITIAL_ACTIONS_M_S[2] * (1.0 - r) + s3 * r
return a1, a2, a3
def _action_to_omega_lb(action_m_s: float, u_lb: float) -> float:
u_surf_lb = action_m_s * (u_lb / INLET_U_PHYS_M_S)
r_lb = DIAMETER_CELLS / 2.0
return OMEGA_SIGN_FROM_ACTION * (u_surf_lb / r_lb)
def _set_body_omegas(sim: Simulation, omega0: float, omega1: float, omega2: float) -> None:
"""Set all three body rotation speeds using new API (implicit GPU upload)."""
sim.set_body(0, omega=omega0)
sim.set_body(1, omega=omega1)
sim.set_body(2, omega=omega2)
def _default_steps(nx: int, u_lb: float, step_multiplier: float) -> int:
base = int(round(2.0 * float(nx) / (3.0 * float(u_lb))))
return int(round(base * float(step_multiplier)))
def run_case(
case_id: str,
slug: str,
features: dict,
*,
out_dir: Path,
steps: int,
report_every: int,
batch: int = 1,
device_id: int = 0,
) -> dict:
compat = _ensure_compat_config(CONFIG_PATH)
sim = Simulation(compat, device_id=device_id)
layout = _add_triangle_cylinders(sim)
sim.initialize()
u_lb = float(sim.lbm_cfg.velocity)
dx_phys = CYLINDER_DIAMETER_M / DIAMETER_CELLS
dt_phys = dx_phys * (u_lb / INLET_U_PHYS_M_S)
cylinders = cylinders_from_triangle_layout(layout)
print(f"--- {case_id} {slug} steps={steps} u_lb={u_lb} dt_phys={dt_phys} batch={batch} ---")
stream = sim.stream
batch_size = max(1, int(batch))
# Main loop: precompute actions, batch-step, read forces/sensors at intervals
for batch_start in range(0, steps, batch_size):
batch_end = min(batch_start + batch_size, steps)
for j in range(batch_start, batch_end):
t_phys = j * dt_phys
a1, a2, a3 = _actions_at_time(t_phys, features)
w1 = _action_to_omega_lb(a1, u_lb)
w2 = _action_to_omega_lb(a2, u_lb)
w3 = _action_to_omega_lb(a3, u_lb)
if SWAP_ACTION23_BODIES:
_set_body_omegas(sim, w1, w3, w2)
else:
_set_body_omegas(sim, w1, w2, w3)
n = batch_end - batch_start
sim.stepper.step(
n,
action_gpu=sim.bodies.action_gpu,
obs_gpu=sim.bodies.obs_gpu,
stream=stream,
)
if report_every > 0 and (batch_end % report_every == 0 or batch_end == steps):
stream.synchronize()
for bid in range(sim.bodies.count):
fx = sim.bodies.read_force(bid)
print(
f" step={batch_end} body={bid}"
f" fx={float(fx[0]):+.6f} fy={float(fx[1]):+.6f}",
flush=True,
)
stream.synchronize()
macro = sim.get_macroscopic()
vort = compute_vorticity(macro["ux"], macro["uy"])
png = out_dir / f"vorticity_{case_id}_{slug}.png"
ckpt = out_dir / f"state_{case_id}_{slug}.h5"
sim.save_checkpoint(str(ckpt))
render_info = render_vorticity_field(
vort,
nx=int(sim.lbm_cfg.nx),
ny=int(sim.lbm_cfg.ny),
out_path=str(png),
cylinders=cylinders,
vmin=VORT_VMIN,
vmax=VORT_VMAX,
minimal_axes=True,
)
sim.close()
summary = {
"case_id": case_id,
"slug": slug,
"steps": int(steps),
"batch": int(batch),
"u_lb": u_lb,
"dt_phys": dt_phys,
"vort_png": str(png),
"checkpoint": str(ckpt),
"vort_range_data": [float(vort.min()), float(vort.max())],
"vort_plot_range": [VORT_VMIN, VORT_VMAX],
"swap_action23_bodies": bool(SWAP_ACTION23_BODIES),
"render": render_info,
}
with (out_dir / f"summary_{case_id}_{slug}.json").open("w", encoding="utf-8") as f:
json.dump(summary, f, indent=2)
print(f" saved {png}")
return summary
def main() -> int:
ap = argparse.ArgumentParser(description="exp_ctrl_matrix vorticity batch")
ap.add_argument("--out-dir", type=str, default=str(DEFAULT_OUT))
ap.add_argument(
"--step-multiplier",
type=float,
default=2.0,
help=(
"Steps = multiplier * round(2*nx/(3*u_lb)); "
"use 2.0 after halving nx/ny."
),
)
ap.add_argument(
"--steps",
type=int,
default=FIXED_STEPS,
help=f"Total LBM steps (default {FIXED_STEPS}).",
)
ap.add_argument("--report-every", type=int, default=20000)
ap.add_argument("--cases", type=str, default="",
help="Comma list e.g. C0,C1 or empty=all.")
ap.add_argument(
"--batch", type=int, default=1,
help=(
"Batch N steps between action uploads. Default 1 (each step). "
"With --batch 10, actions are computed and uploaded every 10 steps. "
"Saves kernel launch overhead at the cost of control-signal interpolation."
),
)
ap.add_argument("--device-id", type=int, default=0, help="GPU device id.")
args = ap.parse_args()
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
selected = {c.strip() for c in args.cases.split(",") if c.strip()} \
if args.cases else None
summaries = []
with CONFIG_PATH.open("r", encoding="utf-8") as f:
grid_cfg = json.load(f)["grid"]
nx = int(grid_cfg["nx"])
ny = int(grid_cfg["ny"])
u_lb = 0.04
base_steps = int(round(2.0 * nx / (3.0 * u_lb)))
steps = int(args.steps) if int(args.steps) > 0 \
else _default_steps(nx, u_lb, args.step_multiplier)
print(
f"Output: {out_dir} | grid={nx}x{ny} | base_steps={base_steps} "
f"x{args.step_multiplier} -> {steps} | batch={args.batch} | "
f"device={args.device_id} | vort [{VORT_VMIN}, {VORT_VMAX}]"
)
for case_id, slug, features in CONTROL_CASES:
if selected and case_id not in selected:
continue
summaries.append(
run_case(
case_id,
slug,
features,
out_dir=out_dir,
steps=steps,
report_every=int(args.report_every),
batch=int(args.batch),
device_id=int(args.device_id),
)
)
manifest = {
"grid": {"nx": nx, "ny": ny},
"base_steps": base_steps,
"step_multiplier": float(args.step_multiplier),
"steps": steps,
"batch": int(args.batch),
"device_id": int(args.device_id),
"vort_vmin": VORT_VMIN,
"vort_vmax": VORT_VMAX,
"swap_action23_bodies": bool(SWAP_ACTION23_BODIES),
"cases": summaries,
}
manifest_path = out_dir / "manifest.json"
with manifest_path.open("w", encoding="utf-8") as f:
json.dump(manifest, f, indent=2)
print(f"Manifest: {manifest_path}")
return 0
if __name__ == "__main__":
raise SystemExit(main())