DynamisLab/src/OID_analysis/analysis/steady_reanalysis.py

# OID_analysis/analysis/steady_reanalysis.py
"""
Steady cloak re-analysis with suppression/restoration metrics.
Replaces R^2 with physically meaningful metrics:
  - RMS reduction per zone
  - Recirculation length/area collapse
  - Enstrophy reduction
  - Force RMS

Usage:
    conda run -n sr_env python3 src/OID_analysis/analysis/steady_reanalysis.py
"""
from __future__ import annotations

import json
import os
import sys

import numpy as np

_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
if _REPO not in sys.path:
    sys.path.insert(0, _REPO)

from OID_analysis.configs import DATA_DIR, data_dir_for_scene  # noqa: E402


def compute_rms_reduction(ux_ctl, uy_ctl, ux_blk, uy_blk, zone_mask=None):
    """Compute RMS reduction ratio: 1 - RMS(ctl)/RMS(blk)"""
    if zone_mask is not None:
        ux_ctl = ux_ctl[:, zone_mask]
        uy_ctl = uy_ctl[:, zone_mask]
        ux_blk = ux_blk[:, zone_mask]
        uy_blk = uy_blk[:, zone_mask]

    rms_ctl_u = np.std(ux_ctl, axis=0)
    rms_ctl_v = np.std(uy_ctl, axis=0)
    rms_blk_u = np.std(ux_blk, axis=0)
    rms_blk_v = np.std(uy_blk, axis=0)

    rms_ctl = np.sqrt(np.mean(rms_ctl_u**2 + rms_ctl_v**2))
    rms_blk = np.sqrt(np.mean(rms_blk_u**2 + rms_blk_v**2))

    reduction = 1.0 - rms_ctl / (rms_blk + 1e-30)
    return float(reduction), float(rms_ctl), float(rms_blk)


def compute_enstrophy_reduction(ux_ctl, uy_ctl, ux_blk, uy_blk, zone_mask=None):
    """Compute enstrophy reduction ratio."""
    def zonal_enstrophy(ux, uy):
        omega = np.zeros((ux.shape[0], ux.shape[1], ux.shape[2]), dtype=np.float64)
        for t in range(min(ux.shape[0], 100)):  # subsample for speed
            omega[t] = np.gradient(uy[t].astype(np.float64), axis=1) - \
                       np.gradient(ux[t].astype(np.float64), axis=0)
        if zone_mask is not None:
            area = max(np.sum(zone_mask), 1)
            return np.mean(0.5 * omega[:, zone_mask]**2)
        return np.mean(0.5 * omega**2)

    ens_ctl = zonal_enstrophy(ux_ctl, uy_ctl)
    ens_blk = zonal_enstrophy(ux_blk, uy_blk)
    reduction = 1.0 - ens_ctl / (ens_blk + 1e-30)
    return float(reduction), float(ens_ctl), float(ens_blk)


def compute_recirculation_metrics(mean_u):
    """Compute Lr and Ar from mean u field."""
    ny, nx = mean_u.shape
    center_y = (ny - 1) / 2.0
    cl_y = int(center_y)

    # Lr: furthest downstream x on centerline where mean_u < 0
    u_cl = mean_u[cl_y, :]
    neg_idx = np.where(u_cl < 0)[0]
    Lr = float(neg_idx[-1]) if len(neg_idx) > 0 else 0.0

    # Ar: pixels where mean_u < 0
    Ar = float(np.sum(mean_u < 0))

    return Lr, Ar


def analyze_steady():
    print("=== Steady Cloak Re-Analysis ===")

    # Load fields
    blk_dir = data_dir_for_scene("pinball_baseline")
    ctl_dir = data_dir_for_scene("steady_cloak")

    f_blk = np.load(os.path.join(blk_dir, "fields.npz"))
    ux_blk, uy_blk = f_blk["ux"], f_blk["uy"]

    f_ctl = np.load(os.path.join(ctl_dir, "fields.npz"))
    ux_ctl, uy_ctl = f_ctl["ux"], f_ctl["uy"]

    # Equalize lengths
    N = min(ux_blk.shape[0], ux_ctl.shape[0])
    ux_blk, uy_blk = ux_blk[:N], uy_blk[:N]
    ux_ctl, uy_ctl = ux_ctl[:N], uy_ctl[:N]

    print(f"\nN snapshots: {N} (min of blk={ux_blk.shape[0]}, ctl={ux_ctl.shape[0]})")

    # Zone masks
    ny, nx = ux_blk.shape[1], ux_blk.shape[2]

    # Near-body: x=[580,660], y=[200,310] in lattice
    nb_mask = np.zeros((ny, nx), dtype=bool)
    nb_mask[200:310, 580:660] = True

    # Near-wake: x=[660,800], y=[180,330]
    nw_mask = np.zeros((ny, nx), dtype=bool)
    nw_mask[180:330, 660:800] = True

    # Downstream sensor zone: x=[790,810]
    ds_mask = np.zeros((ny, nx), dtype=bool)
    for sy in [215, 255, 295]:
        y0, y1 = max(0, sy-10), min(ny, sy+10)
        ds_mask[y0:y1, 790:810] = True

    zones = [
        ("near-body", nb_mask),
        ("near-wake", nw_mask),
        ("downstream", ds_mask),
        ("full-field", None),
    ]

    results = {}
    for zname, zmask in zones:
        print(f"\n  Zone: {zname}")
        rms_red, rms_c, rms_b = compute_rms_reduction(
            ux_ctl, uy_ctl, ux_blk, uy_blk, zmask)
        ens_red, ens_c, ens_b = compute_enstrophy_reduction(
            ux_ctl, uy_ctl, ux_blk, uy_blk, zmask)
        results[zname] = {
            "rms_reduction": rms_red,
            "rms_ctl": rms_c,
            "rms_blk": rms_b,
            "enstrophy_reduction": ens_red,
            "enstrophy_ctl": ens_c,
            "enstrophy_blk": ens_b,
        }
        print(f"    RMS reduction: {rms_red:.4f} (ctl={rms_c:.6f}, blk={rms_b:.6f})")
        print(f"    Enstrophy reduction: {ens_red:.4f} (ctl={ens_c:.6f}, blk={ens_b:.6f})")

    # Recirculation metrics for full field
    mean_u_ctl = np.mean(ux_ctl, axis=0)
    mean_u_blk = np.mean(ux_blk, axis=0)
    Lr_ctl, Ar_ctl = compute_recirculation_metrics(mean_u_ctl)
    Lr_blk, Ar_blk = compute_recirculation_metrics(mean_u_blk)
    results["recirculation"] = {
        "Lr_ctl_lattice": Lr_ctl,
        "Lr_blk_lattice": Lr_blk,
        "Lr_collapse": Lr_ctl / (Lr_blk + 1e-30),
        "Ar_ctl": Ar_ctl,
        "Ar_blk": Ar_blk,
        "Ar_collapse": Ar_ctl / (Ar_blk + 1e-30),
    }
    print(f"\n  Recirculation:")
    print(f"    Lr: ctl={Lr_ctl:.0f}, blk={Lr_blk:.0f}, collapse={Lr_ctl/(Lr_blk+1e-30):.4f}")
    print(f"    Ar: ctl={Ar_ctl:.0f}, blk={Ar_blk:.0f}, collapse={Ar_ctl/(Ar_blk+1e-30):.4f}")

    # Force metrics
    fp = os.path.join(blk_dir, "forces.npz")
    forces_blk = np.load(fp)["forces"][:N]
    fp_ctl = os.path.join(ctl_dir, "forces.npz")
    forces_ctl = np.load(fp_ctl)["forces"][:N]

    Fx_blk_rms = np.std(np.sum(forces_blk[:, 0::2], axis=1))
    Fy_blk_rms = np.std(np.sum(forces_blk[:, 1::2], axis=1))
    Fx_ctl_rms = np.std(np.sum(forces_ctl[:, 0::2], axis=1))
    Fy_ctl_rms = np.std(np.sum(forces_ctl[:, 1::2], axis=1))

    results["force"] = {
        "Fx_rms_blk": float(Fx_blk_rms),
        "Fx_rms_ctl": float(Fx_ctl_rms),
        "Fx_reduction": float(1.0 - Fx_ctl_rms / (Fx_blk_rms + 1e-30)),
        "Fy_rms_blk": float(Fy_blk_rms),
        "Fy_rms_ctl": float(Fy_ctl_rms),
        "Fy_reduction": float(1.0 - Fy_ctl_rms / (Fy_blk_rms + 1e-30)),
    }
    print(f"\n  Force:")
    print(f"    Fx RMS: blk={Fx_blk_rms:.6f}, ctl={Fx_ctl_rms:.6f}, reduction={results['force']['Fx_reduction']:.4f}")
    print(f"    Fy RMS: blk={Fy_blk_rms:.6f}, ctl={Fy_ctl_rms:.6f}, reduction={results['force']['Fy_reduction']:.4f}")

    # Save
    out_dir = os.path.join(DATA_DIR, "derived", "steady_metrics")
    os.makedirs(out_dir, exist_ok=True)
    with open(os.path.join(out_dir, "steady_reanalysis.json"), "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nSaved to {out_dir}/steady_reanalysis.json")

    # Summary table
    print(f"\n{'='*60}")
    print(f"STEADY CLOAK RE-ANALYSIS SUMMARY")
    print(f"{'='*60}")
    print(f"{'Metric':<30s} {'Uncontrolled':>12s} {'Controlled':>12s} {'Reduction':>12s}")
    print(f"{'-'*66}")
    for zname in ["near-body", "near-wake", "downstream"]:
        zr = results[zname]
        print(f"RMS ({zname}){'':<12s} {zr['rms_blk']:12.6f} {zr['rms_ctl']:12.6f} {zr['rms_reduction']:12.4f}")
    print(f"Lr (recirc length)     {Lr_blk:12.0f} {Lr_ctl:12.0f} {results['recirculation']['Lr_collapse']:12.4f}")
    print(f"Ar (recirc area)       {Ar_blk:12.0f} {Ar_ctl:12.0f} {results['recirculation']['Ar_collapse']:12.4f}")
    print(f"Fx RMS                 {Fx_blk_rms:12.6f} {Fx_ctl_rms:12.6f} {results['force']['Fx_reduction']:12.4f}")
    print(f"Fy RMS                 {Fy_blk_rms:12.6f} {Fy_ctl_rms:12.6f} {results['force']['Fy_reduction']:12.4f}")


if __name__ == "__main__":
    analyze_steady()