CelerisLab/tests/postproc/run_exp_ctrl_matrix_streakline.py

# CelerisLab/tests/postproc/run_exp_ctrl_matrix_streakline.py
"""Streakline post-processing for exp_ctrl_matrix cases.

Runs full CFD, uses Streakline.observe() in the last STREAK_WINDOW steps,
renders streakline at the final step.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

_REPO = Path(__file__).resolve().parents[2]

import sys

sys.path.insert(0, str(_REPO / "src"))
sys.path.insert(0, str(_REPO / "tests" / "postproc"))

import run_exp_ctrl_matrix_vorticity as vort
from CelerisLab import Simulation
from CelerisLab.common.streakline import Streakline, ReleaseConfig, IntegratorConfig

DIAMETER_CELLS = vort.DIAMETER_CELLS
DEFAULT_OUT = _REPO / "tests" / "output" / "exp_ctrl_matrix_streak_ny300"
STREAK_WINDOW_STEPS = 20_000
STREAK_SAMPLE_EVERY = 50
STREAK_AGE_DECAY = 100_000.0
STREAK_BLUR_SIGMA = 0.8


def build_release_points_for_triangle(layout: dict) -> np.ndarray:
    release_x = min(layout["x_apex"], layout["x_rear"]) - 4.0 * DIAMETER_CELLS
    y_low_edge = layout["y_lower"] - layout["radius_lb"]
    y_high_edge = layout["y_upper"] + layout["radius_lb"]
    ys = np.linspace(y_low_edge, y_high_edge, 4, dtype=np.float64)
    return np.column_stack([np.full(4, release_x, dtype=np.float64), ys])


def _apply_body_actions(
    sim: Simulation, a1: float, a2: float, a3: float, u_lb: float
) -> None:
    w1 = vort._action_to_omega_lb(a1, u_lb)
    w2 = vort._action_to_omega_lb(a2, u_lb)
    w3 = vort._action_to_omega_lb(a3, u_lb)
    if vort.SWAP_ACTION23_BODIES:
        vort._set_body_omegas(sim, w1, w3, w2)
    else:
        vort._set_body_omegas(sim, w1, w2, w3)


def _cylinders_from_triangle_layout(layout: dict) -> list[tuple[tuple[float, float], float]]:
    radius = float(layout["radius_lb"])
    return [
        ((float(layout["x_apex"]), float(layout["y_center"])), radius),
        ((float(layout["x_rear"]), float(layout["y_lower"])), radius),
        ((float(layout["x_rear"]), float(layout["y_upper"])), radius),
    ]


def run_streak_case(
    case_id: str,
    slug: str,
    features: dict,
    *,
    out_dir: Path,
    total_steps: int,
    streak_window: int,
    sample_every: int,
    report_every: int,
) -> dict:
    streak_start = max(0, int(total_steps) - int(streak_window))
    compat = vort._ensure_compat_config(vort.CONFIG_PATH)
    sim = Simulation(compat)
    layout = vort._add_triangle_cylinders(sim)
    sim.initialize()
    u_lb = float(sim.lbm_cfg.velocity)
    dt_phys = (
        (vort.CYLINDER_DIAMETER_M / DIAMETER_CELLS)
        * (u_lb / vort.INLET_U_PHYS_M_S)
    )
    cylinders = _cylinders_from_triangle_layout(layout)

    release_cfg = ReleaseConfig(
        mode="strip",
        line_count=1,
        line_span=0.0,
        downstream_count=5,
        downstream_spacing=1.0,
        inject_per_seed=1,
    )
    integrator_cfg = IntegratorConfig(
        alpha_t=0.25, alpha_x=0.40, max_particle_age=None
    )
    base_release = build_release_points_for_triangle(layout)

    streak = Streakline(
        release_points=base_release,
        release_cfg=release_cfg,
        integrator_cfg=integrator_cfg,
        nx=int(sim.lbm_cfg.nx),
        ny=int(sim.lbm_cfg.ny),
        cylinders=cylinders,
    )

    print(
        f"--- {case_id} {slug} steps={total_steps} streak_window={streak_window} "
        f"(inject from step {streak_start}) ---"
    )

    for step in range(total_steps):
        t_phys = step * dt_phys
        a1, a2, a3 = vort._actions_at_time(t_phys, features)
        _apply_body_actions(sim, a1, a2, a3, u_lb)
        sim.run(1)

        # Feed streakline within the window
        if step >= streak_start and (step + 1) % sample_every == 0:
            macro = sim.get_macroscopic()
            streak.observe(ux=macro["ux"], uy=macro["uy"], step=int(step + 1))

        if report_every > 0 and (step + 1) % report_every == 0:
            print(
                f"  step {step+1}/{total_steps} a=({a1:+.5f},{a2:+.5f},{a3:+.5f}) "
                f"particles={streak.n_particles}"
            )

    if streak.n_particles == 0:
        raise RuntimeError(
            f"{case_id}: no particles in streak window; lower sample_every."
        )

    png = out_dir / f"streakline_{case_id}_{slug}.png"
    render_info = streak.render(
        str(png),
        age_decay_steps=STREAK_AGE_DECAY,
        blur_sigma=STREAK_BLUR_SIGMA,
        background_color=(1.0, 1.0, 1.0),
        streak_color=(1.0, 0.0, 0.0),
    )
    sim.close()

    summary = {
        "case_id": case_id,
        "slug": slug,
        "total_steps": int(total_steps),
        "streak_window_steps": int(streak_window),
        "streak_start_step": int(streak_start),
        "sample_every": int(sample_every),
        "particle_count_final": int(streak.n_particles),
        "release_points_dense": int(base_release.shape[0]),
        "streak_png": str(png),
        "swap_action23_bodies": bool(vort.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} particles={streak.n_particles}")
    return summary


def main() -> int:
    ap = argparse.ArgumentParser(description="exp_ctrl_matrix streakline batch")
    ap.add_argument("--out-dir", type=str, default=str(DEFAULT_OUT))
    ap.add_argument("--steps", type=int, default=vort.FIXED_STEPS)
    ap.add_argument("--streak-window", type=int, default=STREAK_WINDOW_STEPS)
    ap.add_argument("--sample-every", type=int, default=STREAK_SAMPLE_EVERY)
    ap.add_argument("--report-every", type=int, default=20000)
    ap.add_argument("--cases", type=str, default="")
    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
    )

    with vort.CONFIG_PATH.open("r", encoding="utf-8") as f:
        grid = json.load(f)["grid"]
    print(
        f"Output: {out_dir} | grid={grid['nx']}x{grid['ny']} | steps={args.steps} | "
        f"streak last {args.streak_window} steps, sample_every={args.sample_every}"
    )

    summaries = []
    for case_id, slug, features in vort.CONTROL_CASES:
        if selected and case_id not in selected:
            continue
        summaries.append(
            run_streak_case(
                case_id,
                slug,
                features,
                out_dir=out_dir,
                total_steps=int(args.steps),
                streak_window=int(args.streak_window),
                sample_every=int(args.sample_every),
                report_every=int(args.report_every),
            )
        )

    manifest = {
        "grid": grid,
        "steps": int(args.steps),
        "streak_window_steps": int(args.streak_window),
        "sample_every": int(args.sample_every),
        "swap_action23_bodies": bool(vort.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())