{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "e9ca5704",
      "metadata": {},
      "source": [
        "## CLEAN CFD workflow (use this section only)\n",
        "\n",
        "Ignore previous CFD cells above.\n",
        "\n",
        "Run cells in order:\n",
        "\n",
        "1. `CLEAN 1` imports\n",
        "2. `CLEAN 2` parameters\n",
        "3. `CLEAN 3` helpers\n",
        "4. `CLEAN 4` init/restore\n",
        "5. `CLEAN 5` compute chunk\n",
        "6. `CLEAN 6` plot-only (repeat this one as needed)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "id": "7e0fd1a4",
      "metadata": {},
      "outputs": [],
      "source": [
        "# ==================== CLEAN 1. Imports ====================\n",
        "from pathlib import Path\n",
        "import math\n",
        "import json\n",
        "import tempfile\n",
        "\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import pycuda.driver as cuda\n",
        "\n",
        "from CelerisLab import Simulation\n",
        "from CelerisLab.common.streakline import Streakline, ReleaseConfig, IntegratorConfig\n",
        "from CelerisLab.common.render import compute_vorticity, render_vorticity_field\n",
        "from CelerisLab.common.preprocess import cylinders_from_triangle_layout\n",
        "\n",
        "INLET_U_PHYS_M_S = 0.009028\n",
        "CYLINDER_DIAMETER_M = 0.010\n",
        "CENTER_SPACING_M = 0.015\n",
        "DIAMETER_CELLS = 20.0\n",
        "RAMP_TIME_S = 5.0\n",
        "INITIAL_ACTIONS_M_S = (0.0, 0.0, 0.0)\n",
        "OMEGA_SIGN_FROM_ACTION = -1.0  # experiment positive(clockwise) -> kernel omega(CCW positive)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "id": "40d27562",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "WORKDIR: /home/frank14f/CelerisLab/tests\n",
            "CONFIG: ../src/CelerisLab/configs/config_lbm_three_cylinder_triangle.json\n"
          ]
        }
      ],
      "source": [
        "# ==================== CLEAN 2. Parameters (edit here) ====================\n",
        "cwd = Path.cwd()\n",
        "WORKDIR = (cwd / \"tests\") if (cwd / \"src\").exists() else cwd\n",
        "WORKDIR.mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "CFD_PARAMS = {\n",
        "    \"config_path\": str((cwd / \"src/CelerisLab/configs/config_lbm_three_cylinder_triangle.json\") if (cwd / \"src\").exists() else Path(\"../src/CelerisLab/configs/config_lbm_three_cylinder_triangle.json\")),\n",
        "    \"inlet_scheme\": \"regularized\",\n",
        "    \"checkpoint_path\": str(WORKDIR / \"triangle_state_clean.h5\"),\n",
        "    \"output_streak_png\": str(WORKDIR / \"streakline_triangle_clean.png\"),\n",
        "    \"output_streak_meta\": str(WORKDIR / \"streakline_triangle_clean.json\"),\n",
        "    \"output_vort_png\": str(WORKDIR / \"vorticity_triangle_clean.png\"),\n",
        "    \"report_every\": 10000,\n",
        "}\n",
        "\n",
        "RUN_PARAMS = {\n",
        "    \"steady\": False,\n",
        "    \"steps_this_chunk\": 0,  # 0 => 2 * nx / (3 * U_lb)\n",
        "    \"load_checkpoint\": False,\n",
        "}\n",
        "\n",
        "STREAK_PARAMS = {\n",
        "    \"sample_every\": 200,\n",
        "    \"inject_per_seed\": 1,\n",
        "    # Per seed: 5 particles along +x with small spacing (strip mode).\n",
        "    \"release_mode\": \"strip\",\n",
        "    \"downstream_count\": 5,\n",
        "    \"downstream_spacing\": 1.0,\n",
        "    \"max_particle_age\": None,\n",
        "    \"age_decay_steps\": 80000.0,\n",
        "    \"blur_sigma\": 0.8,\n",
        "    \"num_threads\": 0,\n",
        "}\n",
        "\n",
        "ACTION_FEATURES = {\n",
        "    \"action1\": {\n",
        "        \"mean\": 0.0,\n",
        "        \"components\": [\n",
        "            # (0.2708, 0.0044, -2.6190),\n",
        "            # (0.8124, 0.0012, 2.8560),\n",
        "            (0.2708, 0.000, -2.6190),\n",
        "        ],\n",
        "    },\n",
        "    \"action2\": {\n",
        "        # \"mean\": -0.0140,\n",
        "        \"mean\": -0.0,\n",
        "        \"components\": [\n",
        "            # (0.2708, 0.0075, -0.9330),\n",
        "            # (0.8124, 0.0015, 0.8010),\n",
        "            # (0.5418, 0.0005, 1.8540),\n",
        "            (0.2708, 0.000, -0.9330),\n",
        "        ],\n",
        "    },\n",
        "    \"action3\": {\n",
        "        # \"mean\": 0.0140,\n",
        "        \"mean\": 0.0,\n",
        "        \"components\": [\n",
        "            # (0.2708, 0.0075, -1.3980),\n",
        "            # (0.8124, 0.0015, 2.2080),\n",
        "            # (0.5418, 0.0005, 1.8100),\n",
        "            (0.2708, 0.000, -1.3980),\n",
        "        ],\n",
        "    },\n",
        "}\n",
        "\n",
        "print(\"WORKDIR:\", WORKDIR)\n",
        "print(\"CONFIG:\", CFD_PARAMS[\"config_path\"])\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "f055f7dd",
      "metadata": {},
      "outputs": [],
      "source": [
        "# ==================== CLEAN 3. Helpers ====================\n",
        "def generate_signal(t_phys, feature):\n",
        "    value = float(feature[\"mean\"])\n",
        "    for freq_hz, amp, phase in feature[\"components\"]:\n",
        "        value += amp * math.cos(2.0 * math.pi * freq_hz * t_phys + phase)\n",
        "    return value\n",
        "\n",
        "\n",
        "def ramp_factor(elapsed_s):\n",
        "    if elapsed_s <= 0.0:\n",
        "        return 0.0\n",
        "    if elapsed_s >= RAMP_TIME_S:\n",
        "        return 1.0\n",
        "    x = elapsed_s / RAMP_TIME_S\n",
        "    return 0.5 * (1.0 - math.cos(math.pi * x))\n",
        "\n",
        "\n",
        "def actions_m_s_at_time(t_phys, steady=False):\n",
        "    if steady:\n",
        "        return (\n",
        "            float(ACTION_FEATURES[\"action1\"][\"mean\"]),\n",
        "            float(ACTION_FEATURES[\"action2\"][\"mean\"]),\n",
        "            float(ACTION_FEATURES[\"action3\"][\"mean\"]),\n",
        "        )\n",
        "    s1 = generate_signal(t_phys, ACTION_FEATURES[\"action1\"])\n",
        "    s2 = generate_signal(t_phys, ACTION_FEATURES[\"action2\"])\n",
        "    s3 = generate_signal(t_phys, ACTION_FEATURES[\"action3\"])\n",
        "    r = ramp_factor(t_phys)\n",
        "    a1 = INITIAL_ACTIONS_M_S[0] * (1.0 - r) + s1 * r\n",
        "    a2 = INITIAL_ACTIONS_M_S[1] * (1.0 - r) + s2 * r\n",
        "    a3 = INITIAL_ACTIONS_M_S[2] * (1.0 - r) + s3 * r\n",
        "    return a1, a2, a3\n",
        "\n",
        "\n",
        "def action_to_omega_lb(action_m_s, u_lb):\n",
        "    u_surf_lb = action_m_s * (u_lb / INLET_U_PHYS_M_S)\n",
        "    r_lb = DIAMETER_CELLS / 2.0\n",
        "    return OMEGA_SIGN_FROM_ACTION * (u_surf_lb / r_lb)\n",
        "\n",
        "\n",
        "def set_body_omegas(sim, omega0, omega1, omega2):\n",
        "    bodies = sim.bodies\n",
        "    dim = sim.lbm_cfg.dim\n",
        "    slot = 3 * dim\n",
        "    bodies.action.fill(0.0)\n",
        "    bodies.action[(0 * slot) + slot - 1] = np.float32(omega0)\n",
        "    bodies.action[(1 * slot) + slot - 1] = np.float32(omega1)\n",
        "    bodies.action[(2 * slot) + slot - 1] = np.float32(omega2)\n",
        "    cuda.memcpy_htod(bodies.action_gpu, bodies.action)\n",
        "\n",
        "\n",
        "def ensure_compat_config(config_path, preferred_scheme=\"regularized\"):\n",
        "    config_file = Path(config_path)\n",
        "    with config_file.open(\"r\", encoding=\"utf-8\") as f:\n",
        "        cfg = json.load(f)\n",
        "\n",
        "    method = cfg.setdefault(\"method\", {})\n",
        "    inlet = method.setdefault(\"inlet\", {})\n",
        "    outlet = method.setdefault(\"outlet\", {})\n",
        "\n",
        "    changed = False\n",
        "    if \"scheme\" not in inlet:\n",
        "        inlet[\"scheme\"] = preferred_scheme\n",
        "        changed = True\n",
        "    if \"regularized_neq_damp\" not in inlet:\n",
        "        inlet[\"regularized_neq_damp\"] = 0.5\n",
        "        changed = True\n",
        "    if \"blend_alpha\" not in outlet:\n",
        "        outlet[\"blend_alpha\"] = 0.7\n",
        "        changed = True\n",
        "    if \"backflow_clamp\" not in outlet:\n",
        "        outlet[\"backflow_clamp\"] = True\n",
        "        changed = True\n",
        "\n",
        "    if not changed:\n",
        "        return str(config_file)\n",
        "\n",
        "    tmp = tempfile.NamedTemporaryFile(\n",
        "        mode=\"w\",\n",
        "        suffix=\"_compat_config_lbm.json\",\n",
        "        delete=False,\n",
        "        encoding=\"utf-8\",\n",
        "    )\n",
        "    with tmp:\n",
        "        json.dump(cfg, tmp, indent=4)\n",
        "    return tmp.name\n",
        "\n",
        "\n",
        "def triangle_layout(cfg):\n",
        "    dx_phys = CYLINDER_DIAMETER_M / DIAMETER_CELLS\n",
        "    spacing_lb = CENTER_SPACING_M / dx_phys\n",
        "    radius_lb = DIAMETER_CELLS / 2.0\n",
        "\n",
        "    y_center = 0.5 * (cfg.ny - 1)\n",
        "    x_cluster_center = cfg.nx / 3.0\n",
        "    x_apex = x_cluster_center - (math.sqrt(3.0) / 3.0) * spacing_lb\n",
        "    x_rear = x_apex + (math.sqrt(3.0) / 2.0) * spacing_lb\n",
        "    y_upper = y_center + 0.5 * spacing_lb\n",
        "    y_lower = y_center - 0.5 * spacing_lb\n",
        "\n",
        "    return {\n",
        "        \"x_apex\": x_apex,\n",
        "        \"x_rear\": x_rear,\n",
        "        \"y_center\": y_center,\n",
        "        \"y_upper\": y_upper,\n",
        "        \"y_lower\": y_lower,\n",
        "        \"radius_lb\": radius_lb,\n",
        "    }\n",
        "\n",
        "\n",
        "def add_triangle_cylinders(sim):\n",
        "    layout = triangle_layout(sim.lbm_cfg)\n",
        "    sim.add_cylinder(center=(layout[\"x_apex\"], layout[\"y_center\"]), radius=layout[\"radius_lb\"])\n",
        "    sim.add_cylinder(center=(layout[\"x_rear\"], layout[\"y_lower\"]), radius=layout[\"radius_lb\"])\n",
        "    sim.add_cylinder(center=(layout[\"x_rear\"], layout[\"y_upper\"]), radius=layout[\"radius_lb\"])\n",
        "    return layout\n",
        "\n",
        "\n",
        "def build_release_points_for_triangle(layout):\n",
        "    release_x = min(layout[\"x_apex\"], layout[\"x_rear\"]) - 4.0 * DIAMETER_CELLS\n",
        "    y_low_edge = layout[\"y_lower\"] - layout[\"radius_lb\"]\n",
        "    y_high_edge = layout[\"y_upper\"] + layout[\"radius_lb\"]\n",
        "    ys = np.linspace(y_low_edge, y_high_edge, 4, dtype=np.float32)\n",
        "    return np.column_stack([np.full(4, release_x, dtype=np.float32), ys]).astype(np.float32)\n",
        "\n",
        "\n",
        "def filter_particles_outside_cylinders(points, layout):\n",
        "    centers = np.array(\n",
        "        [\n",
        "            [layout[\"x_apex\"], layout[\"y_center\"]],\n",
        "            [layout[\"x_rear\"], layout[\"y_lower\"]],\n",
        "            [layout[\"x_rear\"], layout[\"y_upper\"]],\n",
        "        ],\n",
        "        dtype=np.float32,\n",
        "    )\n",
        "    dist2_min = np.full(points.shape[0], np.inf, dtype=np.float32)\n",
        "    for center in centers:\n",
        "        d2 = (points[:, 0] - center[0]) ** 2 + (points[:, 1] - center[1]) ** 2\n",
        "        dist2_min = np.minimum(dist2_min, d2)\n",
        "    return dist2_min > (layout[\"radius_lb\"] ** 2)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "69904ae5",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Initialized from scratch.\n",
            "u_lb=0.04, dt_phys=0.0022153300841825435\n"
          ]
        }
      ],
      "source": [
        "# ==================== CLEAN 4. Init/restore session ====================\n",
        "if \"CLEAN_STATE\" not in globals():\n",
        "    CLEAN_STATE = {\"sim\": None, \"steps_done\": 0, \"u_lb\": None, \"dt_phys\": None}\n",
        "\n",
        "if CLEAN_STATE[\"sim\"] is None:\n",
        "    compat_config_path = ensure_compat_config(\n",
        "        CFD_PARAMS[\"config_path\"], preferred_scheme=CFD_PARAMS[\"inlet_scheme\"]\n",
        "    )\n",
        "    sim = Simulation(compat_config_path)\n",
        "    layout = add_triangle_cylinders(sim)\n",
        "    sim.initialize()\n",
        "\n",
        "    u_lb = float(sim.lbm_cfg.velocity)\n",
        "    dx_phys = CYLINDER_DIAMETER_M / DIAMETER_CELLS\n",
        "    dt_phys = dx_phys * (u_lb / INLET_U_PHYS_M_S)\n",
        "\n",
        "    CLEAN_STATE[\"sim\"] = sim\n",
        "    CLEAN_STATE[\"layout\"] = layout\n",
        "    CLEAN_STATE[\"u_lb\"] = u_lb\n",
        "    CLEAN_STATE[\"dt_phys\"] = dt_phys\n",
        "    CLEAN_STATE[\"steps_done\"] = 0\n",
        "\n",
        "    ckpt = Path(CFD_PARAMS[\"checkpoint_path\"])\n",
        "    if RUN_PARAMS[\"load_checkpoint\"] and ckpt.exists():\n",
        "        sim.load_checkpoint(str(ckpt))\n",
        "        CLEAN_STATE[\"steps_done\"] = int(sim.stepper.step_count)\n",
        "        print(f\"Restored checkpoint: {ckpt}, step={CLEAN_STATE['steps_done']}\")\n",
        "    else:\n",
        "        print(\"Initialized from scratch.\")\n",
        "\n",
        "    print(f\"Using config: {compat_config_path}\")\n",
        "    print(f\"Inlet scheme: {sim.lbm_cfg.inlet_scheme}\")\n",
        "else:\n",
        "    print(\"Session already initialized.\")\n",
        "\n",
        "print(f\"u_lb={CLEAN_STATE['u_lb']}, dt_phys={CLEAN_STATE['dt_phys']}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "id": "2e462edd",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "chunk_step=10000/100000 global=10000 a=(+0.00000,+0.00000,+0.00000)\n",
            "chunk_step=20000/100000 global=20000 a=(+0.00000,+0.00000,+0.00000)\n",
            "chunk_step=30000/100000 global=30000 a=(+0.00000,+0.00000,+0.00000)\n",
            "chunk_step=40000/100000 global=40000 a=(+0.00000,+0.00000,+0.00000)\n",
            "chunk_step=50000/100000 global=50000 a=(+0.00000,+0.00000,+0.00000)\n",
            "chunk_step=60000/100000 global=60000 a=(+0.00000,+0.00000,+0.00000)\n",
            "chunk_step=70000/100000 global=70000 a=(+0.00000,+0.00000,+0.00000)\n",
            "chunk_step=80000/100000 global=80000 a=(+0.00000,+0.00000,+0.00000)\n",
            "chunk_step=90000/100000 global=90000 a=(+0.00000,+0.00000,+0.00000)\n",
            "chunk_step=100000/100000 global=100000 a=(+0.00000,+0.00000,+0.00000)\n",
            "Compute done. step=100000\n",
            "Vorticity stats: min=-2.888157e-02, max=+2.866180e-02\n",
            "Saved cache: /home/frank14f/CelerisLab/tests/triangle_latest_fields_clean.npz\n",
            "Saved checkpoint: /home/frank14f/CelerisLab/tests/triangle_state_clean.h5\n"
          ]
        }
      ],
      "source": [
        "# ==================== CLEAN 5. Compute chunk only ====================\n",
        "sim = CLEAN_STATE[\"sim\"]\n",
        "if sim is None:\n",
        "    raise RuntimeError(\"Run CLEAN 4 first.\")\n",
        "\n",
        "layout = CLEAN_STATE[\"layout\"]\n",
        "steps = int(RUN_PARAMS[\"steps_this_chunk\"])\n",
        "if steps <= 0:\n",
        "    steps = int(round(2.0 * sim.lbm_cfg.nx / (3.0 * CLEAN_STATE[\"u_lb\"])))\n",
        "\n",
        "steady = bool(RUN_PARAMS[\"steady\"])\n",
        "report_every = max(1, int(CFD_PARAMS[\"report_every\"]))\n",
        "sample_every = max(1, int(STREAK_PARAMS[\"sample_every\"]))\n",
        "n_threads = configure_compute_threads(int(STREAK_PARAMS.get(\"num_threads\", 0)))\n",
        "print(f\"Host compute threads (BLAS/blur): {n_threads}\")\n",
        "\n",
        "release_cfg = ReleaseConfig(\n",
        "    mode=str(STREAK_PARAMS.get(\"release_mode\", \"strip\")),\n",
        "    line_count=1,\n",
        "    line_span=0.0,\n",
        "    downstream_count=max(1, int(STREAK_PARAMS.get(\"downstream_count\", 5))),\n",
        "    downstream_spacing=float(STREAK_PARAMS.get(\"downstream_spacing\", 1.0)),\n",
        "    inject_per_seed=max(1, int(STREAK_PARAMS[\"inject_per_seed\"])),\n",
        ")\n",
        "max_age = STREAK_PARAMS.get(\"max_particle_age\")\n",
        "integrator_cfg = IntegratorConfig(\n",
        "    alpha_t=0.25,\n",
        "    alpha_x=0.40,\n",
        "    max_particle_age=None if max_age is None else float(max_age),\n",
        "    num_threads=n_threads,\n",
        ")\n",
        "\n",
        "base_release = build_release_points_for_triangle(layout)\n",
        "release_points = build_release_points(base_release, release_cfg).astype(np.float64)\n",
        "cylinders = cylinders_from_triangle_layout(layout)\n",
        "print(\n",
        "    f\"Release seeds: {base_release.shape[0]} base -> {release_points.shape[0]} \"\n",
        "    f\"(+x x{release_cfg.downstream_count}, spacing={release_cfg.downstream_spacing})\"\n",
        ")\n",
        "\n",
        "# Streakline = continuous release at fixed seeds + current particle cloud at end (not path history).\n",
        "particles = np.empty((0, 2), dtype=np.float64)\n",
        "ages = np.empty((0,), dtype=np.float64)\n",
        "prev_frame = None\n",
        "chunk_start_step = int(CLEAN_STATE[\"steps_done\"])\n",
        "\n",
        "for i in range(steps):\n",
        "    t_phys = CLEAN_STATE[\"steps_done\"] * CLEAN_STATE[\"dt_phys\"]\n",
        "    a1, a2, a3 = actions_m_s_at_time(t_phys, steady=steady)\n",
        "    w0 = action_to_omega_lb(a1, CLEAN_STATE[\"u_lb\"])\n",
        "    w2 = action_to_omega_lb(a2, CLEAN_STATE[\"u_lb\"])\n",
        "    w3 = action_to_omega_lb(a3, CLEAN_STATE[\"u_lb\"])\n",
        "    # action2 -> rear-upper, action3 -> rear-lower (swap vs body add order)\n",
        "    set_body_omegas(sim, w0, w3, w2)\n",
        "    sim.run(1)\n",
        "    CLEAN_STATE[\"steps_done\"] += 1\n",
        "\n",
        "    if CLEAN_STATE[\"steps_done\"] % sample_every == 0:\n",
        "        macro = sim.get_macroscopic()\n",
        "        curr_frame = FlowFrame(\n",
        "            step=int(CLEAN_STATE[\"steps_done\"]),\n",
        "            ux=macro[\"ux\"].astype(np.float64, copy=False),\n",
        "            uy=macro[\"uy\"].astype(np.float64, copy=False),\n",
        "        )\n",
        "\n",
        "        # Continuous dye injection: new particles always born at upstream seeds.\n",
        "        new_particles = np.repeat(release_points, release_cfg.inject_per_seed, axis=0).astype(np.float64)\n",
        "        if particles.size:\n",
        "            particles = np.vstack([particles, new_particles])\n",
        "            ages = np.concatenate([ages, np.zeros(new_particles.shape[0], dtype=np.float64)])\n",
        "        else:\n",
        "            particles = new_particles.copy()\n",
        "            ages = np.zeros(new_particles.shape[0], dtype=np.float64)\n",
        "\n",
        "        if prev_frame is not None and particles.size:\n",
        "            particles, ages, _, _ = advance_particles_between_frames(\n",
        "                particles,\n",
        "                ages,\n",
        "                prev_frame,\n",
        "                curr_frame,\n",
        "                nx=sim.lbm_cfg.nx,\n",
        "                ny=sim.lbm_cfg.ny,\n",
        "                cfg=integrator_cfg,\n",
        "                cylinders=cylinders,\n",
        "            )\n",
        "\n",
        "        prev_frame = curr_frame\n",
        "\n",
        "    if (i + 1) % report_every == 0:\n",
        "        print(\n",
        "            f\"chunk_step={i+1}/{steps} global={CLEAN_STATE['steps_done']} \"\n",
        "            f\"a=({a1:+.5f},{a2:+.5f},{a3:+.5f}) particles={particles.shape[0]}\"\n",
        "        )\n",
        "\n",
        "sim.save_checkpoint(CFD_PARAMS[\"checkpoint_path\"])\n",
        "\n",
        "if particles.size == 0:\n",
        "    raise RuntimeError(\"No alive particles at end. Increase steps or lower sample_every.\")\n",
        "\n",
        "streak_png = Path(CFD_PARAMS[\"output_streak_png\"])\n",
        "render_info = render_streakline_density(\n",
        "    particles,\n",
        "    ages,\n",
        "    nx=sim.lbm_cfg.nx,\n",
        "    ny=sim.lbm_cfg.ny,\n",
        "    out_path=str(streak_png),\n",
        "    cylinders=cylinders,\n",
        "    age_decay_steps=float(STREAK_PARAMS[\"age_decay_steps\"]),\n",
        "    blur_sigma=float(STREAK_PARAMS[\"blur_sigma\"]),\n",
        "    minimal_axes=True,\n",
        "    background_color=(1.0, 1.0, 1.0),\n",
        "    streak_color=(1.0, 0.0, 0.0),\n",
        "    show_release_points=False,\n",
        ")\n",
        "\n",
        "streak_meta = {\n",
        "    \"chunk_start_step\": chunk_start_step,\n",
        "    \"chunk_end_step\": int(CLEAN_STATE[\"steps_done\"]),\n",
        "    \"sample_every\": sample_every,\n",
        "    \"release_points_base\": base_release.tolist(),\n",
        "    \"release_points_dense\": release_points.tolist(),\n",
        "    \"release_x_4D_upstream\": float(base_release[0, 0]),\n",
        "    \"particle_count\": int(particles.shape[0]),\n",
        "    \"max_particle_age\": max_age,\n",
        "    \"age_decay_steps\": float(STREAK_PARAMS[\"age_decay_steps\"]),\n",
        "    \"render_mode\": \"particle_cloud_streakline\",\n",
        "    \"render\": render_info,\n",
        "}\n",
        "meta_path = Path(CFD_PARAMS[\"output_streak_meta\"])\n",
        "meta_path.parent.mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "macro_end = sim.get_macroscopic()\n",
        "ux_end = macro_end[\"ux\"].astype(np.float32, copy=False)\n",
        "uy_end = macro_end[\"uy\"].astype(np.float32, copy=False)\n",
        "vort_end = compute_vorticity(ux_end, uy_end)\n",
        "vort_png = Path(CFD_PARAMS[\"output_vort_png\"])\n",
        "vort_info = render_vorticity_field(\n",
        "    vort_end,\n",
        "    nx=sim.lbm_cfg.nx,\n",
        "    ny=sim.lbm_cfg.ny,\n",
        "    out_path=str(vort_png),\n",
        "    cylinders=cylinders,\n",
        "    minimal_axes=True,\n",
        ")\n",
        "streak_meta[\"vorticity\"] = vort_info\n",
        "\n",
        "CLEAN_STATE[\"streak_particles\"] = particles\n",
        "CLEAN_STATE[\"streak_age\"] = ages\n",
        "CLEAN_STATE[\"streak_image\"] = plt.imread(streak_png)\n",
        "CLEAN_STATE[\"vort\"] = vort_end\n",
        "CLEAN_STATE[\"vort_image\"] = plt.imread(vort_png)\n",
        "\n",
        "with meta_path.open(\"w\", encoding=\"utf-8\") as f:\n",
        "    json.dump(streak_meta, f, indent=2)\n",
        "\n",
        "print(f\"Compute done. step={CLEAN_STATE['steps_done']}\")\n",
        "print(f\"Saved streakline image: {streak_png}\")\n",
        "print(f\"Saved vorticity image:  {vort_png}\")\n",
        "print(\n",
        "    f\"Vorticity stats: min={vort_end.min():+.6e}, max={vort_end.max():+.6e} \"\n",
        "    f\"(plot range [{vort_info['vmin']:+.3e}, {vort_info['vmax']:+.3e}])\"\n",
        ")\n",
        "print(f\"Saved streakline meta:  {meta_path}\")\n",
        "print(f\"Saved checkpoint:       {CFD_PARAMS['checkpoint_path']}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "618af266",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Plotting step=100000, vort min=-2.888157e-02, max=+2.866180e-02\n",
            "Saved figure: /home/frank14f/CelerisLab/tests/vorticity_triangle_clean.png\n"
          ]
        },
        {
          "data": {
            "image/png": "[encoded data, 1816648 characters]",
            "text/plain": [
              "<Figure size 6000x1200 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# ==================== CLEAN 6. Plot only (repeat) ====================\n",
        "# Streakline\n",
        "img_path = Path(CFD_PARAMS[\"output_streak_png\"])\n",
        "if \"streak_image\" in CLEAN_STATE:\n",
        "    img = CLEAN_STATE[\"streak_image\"]\n",
        "else:\n",
        "    if not img_path.exists():\n",
        "        raise RuntimeError(\"No streakline image. Run CLEAN 5 first.\")\n",
        "    img = plt.imread(img_path)\n",
        "\n",
        "step = int(CLEAN_STATE.get(\"steps_done\", -1))\n",
        "sim = CLEAN_STATE.get(\"sim\")\n",
        "if sim is None:\n",
        "    raise RuntimeError(\"Simulation state not found. Run CLEAN 4 first.\")\n",
        "print(f\"Plotting streakline at step={step}, image={img_path}\")\n",
        "\n",
        "fig = plt.figure(figsize=(sim.lbm_cfg.nx / 100.0, sim.lbm_cfg.ny / 100.0), dpi=100)\n",
        "ax = fig.add_axes([0.0, 0.0, 1.0, 1.0])\n",
        "ax.axis(\"off\")\n",
        "ax.imshow(img, origin=\"lower\", interpolation=\"nearest\", aspect=\"auto\")\n",
        "plt.show()\n",
        "\n",
        "vort_path = Path(CFD_PARAMS[\"output_vort_png\"])\n",
        "if \"vort_image\" in CLEAN_STATE:\n",
        "    vort_img = CLEAN_STATE[\"vort_image\"]\n",
        "elif vort_path.exists():\n",
        "    vort_img = plt.imread(vort_path)\n",
        "else:\n",
        "    raise RuntimeError(\"No vorticity image. Run CLEAN 5 first.\")\n",
        "print(f\"Vorticity image: {vort_path}\")\n",
        "fig2 = plt.figure(figsize=(sim.lbm_cfg.nx / 100.0, sim.lbm_cfg.ny / 100.0), dpi=100)\n",
        "ax2 = fig2.add_axes([0.0, 0.0, 1.0, 1.0])\n",
        "ax2.axis(\"off\")\n",
        "ax2.imshow(vort_img, origin=\"lower\", interpolation=\"nearest\", aspect=\"auto\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1a4f8cf2",
      "metadata": {},
      "outputs": [],
      "source": [
        "# ==================== CLEAN 7. Optional close ====================\n",
        "# if CLEAN_STATE.get(\"sim\") is not None:\n",
        "#     CLEAN_STATE[\"sim\"].close()\n",
        "#     CLEAN_STATE[\"sim\"] = None\n",
        "#     print(\"Simulation closed.\")"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "pycuda_3_10",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.13"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
