CelerisLab/tests/run_exp_ctrl_matrix_vorticity.py

422 lines
13 KiB
Python

# CelerisLab/tests/run_exp_ctrl_matrix_vorticity.py
"""Batch vorticity images for three-cylinder control matrix (exp_ctrl_matrix.md).
Runs each control law to steady-like state, saves final vorticity PNG only (no streaklines).
"""
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
import pycuda.driver as cuda
_REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO / "src"))
from CelerisLab import Simulation
from CelerisLab.common.streakline import (
compute_vorticity,
cylinders_from_triangle_layout,
render_vorticity_field,
)
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_cylinder(center=(layout["x_apex"], layout["y_center"]), radius=layout["radius_lb"])
sim.add_cylinder(center=(layout["x_rear"], layout["y_lower"]), radius=layout["radius_lb"])
sim.add_cylinder(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:
bodies = sim.bodies
dim = sim.lbm_cfg.dim
slot = 3 * dim
bodies.action.fill(0.0)
bodies.action[(0 * slot) + slot - 1] = np.float32(omega0)
bodies.action[(1 * slot) + slot - 1] = np.float32(omega1)
bodies.action[(2 * slot) + slot - 1] = np.float32(omega2)
cuda.memcpy_htod(bodies.action_gpu, bodies.action)
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,
) -> dict:
compat = _ensure_compat_config(CONFIG_PATH)
sim = Simulation(compat)
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} ---")
for i in range(steps):
t_phys = i * 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)
sim.run(1)
if report_every > 0 and (i + 1) % report_every == 0:
print(f" step {i+1}/{steps} a=({a1:+.5f},{a2:+.5f},{a3:+.5f})")
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),
"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}, held fixed across grid tweaks).",
)
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.")
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} | 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),
)
)
manifest = {
"grid": {"nx": nx, "ny": ny},
"base_steps": base_steps,
"step_multiplier": float(args.step_multiplier),
"steps": steps,
"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())