DynamisLab/archive/analysis_crossre_scripts/phase3_validate.py
2026-06-09 18:46:59 +08:00

307 lines
9.9 KiB
Python

# analysis_crossre/scripts/phase3_validate.py
"""Phase 3: closed-loop validation using cross-Re SINDy control law.
Usage::
conda run -n pycuda_3_10 python phase3_validate.py \\
--device 2 --out-dir output/analysis_crossre/sindy_val
conda run -n pycuda_3_10 python phase3_validate.py \\
--validate-re 35,70,150 --device 2
conda run -n pycuda_3_10 python phase3_validate.py \\
--baseline-only --validate-re 35 --device 2
"""
from __future__ import annotations
import argparse
import json
import os
import sys
import time
from collections import deque
import numpy as np
_PROJ = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
if _PROJ not in sys.path:
sys.path.insert(0, _PROJ)
from LegacyCelerisLab import FlowField # noqa: E402
from utils import (
nu_from_re,
load_legacy_configs,
build_karman_cloak_env,
add_pinball,
build_observation,
scale_action,
action_to_physical,
compute_dimensionless,
compute_v3_symbols,
save_vorticity_png,
vorticity_from_ddf,
compute_similarity,
)
from cfg import (
CONFIG_DIR,
OUTPUT_DIR,
MODEL_DIR,
SAMPLE_INTERVAL,
FIFO_LEN,
CONV_LEN,
S_DIM,
A_DIM,
ACTION_SCALE,
ACTION_BIAS,
U0,
RE_CASES_TRAIN,
RE_CASES_VALIDATION,
RE_LABEL_MAP,
)
DATA_TYPE = np.float32
def load_cross_re_coef(sindy_results_path: str, threshold: float) -> dict:
"""Load v3 cross-Re coefficients.
Returns dict ``{cylinder_name: {"coef": np.ndarray, "feat_names": list, "has_bias": bool}}``
"""
with open(sindy_results_path) as f:
data = json.load(f)
cross = data["cross_re"]
coefs = {}
for ch_entry in cross["channels"]:
name = ch_entry["cylinder"]
feat_names = ch_entry["feature_names"]
coef_full = np.array(ch_entry["best_coef"], dtype=np.float64)
has_bias = ch_entry["has_bias"]
scale = np.max(np.abs(coef_full))
if scale > 0 and threshold > 0:
mask = np.abs(coef_full) / scale >= threshold
else:
mask = np.ones_like(coef_full, dtype=bool)
coef = coef_full * mask
nz = int(np.sum(mask))
print(f" {name}: total={len(coef_full)} nz={nz} threshold={threshold} "
f"R2={ch_entry['best']['r2']:.4f}")
coefs[name] = {"coef": coef, "feat_names": feat_names,
"has_bias": has_bias, "nz": nz, "r2": ch_entry["best"]["r2"]}
return coefs
def predict_omega_v3(
obs_slice: np.ndarray,
actions_prev: np.ndarray,
actions_prev2: np.ndarray,
coefs: dict,
mu: float,
u0: float = 0.01,
) -> np.ndarray:
"""Predict physical omega using v3 dimensionless features.
Front: no bias term.
Bottom/Top: with bias term.
All 3 independently (no exchange symmetry constraint).
Parameters
----------
obs_slice : (12,) raw [sensor(6), force(6)] in lattice units
actions_prev : (3,) omega(t-1)
actions_prev2 : (3,) omega(t-2)
"""
sensors = obs_slice[0:6].astype(np.float64).reshape(1, 6)
forces = obs_slice[6:12].astype(np.float64).reshape(1, 6)
a_prev = actions_prev.astype(np.float64).reshape(1, 3)
a_prev2 = actions_prev2.astype(np.float64).reshape(1, 3)
# Dimensionless
dim = compute_dimensionless(sensors, forces, u0=u0, d=20.0)
# Build v3 features
Theta_f, Theta_top, names = compute_v3_symbols(
dim, a_prev, a_prev2, mu=mu, include_mu=(mu > 0))
# Predict
omega = np.zeros(3, dtype=np.float64)
omega[0] = float(Theta_f[0] @ coefs["front"]["coef"]) # front (no bias)
omega[1] = float(Theta_top[0] @ coefs["bottom"]["coef"]) # bottom
omega[2] = float(Theta_top[0] @ coefs["top"]["coef"]) # top
return omega
def run_sindy_controlled(
re_code: int,
coefs: dict,
device_id: int,
output_root: str,
*,
n_steps: int = 150,
) -> dict:
"""Run closed-loop validation with SINDy control law."""
os.makedirs(output_root, exist_ok=True)
nu = nu_from_re(re_code, u0=U0)
mu = 2.0 / re_code # 1 / Re_D
label = RE_LABEL_MAP.get(re_code, f"Re{re_code}")
print(f"\n{'='*60}")
print(f"SINDy Validation: {label} nu={nu:.6f} mu={mu:.6f}")
print(f"{'='*60}")
# Build environment (same as Phase 1)
cuda_cfg, field_cfg = load_legacy_configs(CONFIG_DIR)
field_cfg = field_cfg._replace(viscosity=float(nu))
# Phase 1: dist + sensors + target
ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
target_states, _ = build_karman_cloak_env(
ff, u0=U0, l0=20.0, sample_interval=SAMPLE_INTERVAL,
fifo_len=FIFO_LEN, data_type=DATA_TYPE,
)
# Phase 2: pinball + norm
norm = add_pinball(
ff, l0=20.0, u0=U0, sample_interval=SAMPLE_INTERVAL,
fifo_len=FIFO_LEN, data_type=DATA_TYPE,
action_bias=ACTION_BIAS,
)
np.savez(os.path.join(output_root, "target.npz"), target_states=target_states)
# --- Uncontrolled rollout ---
print(" uncontrolled rollout ...")
ff.restore_ddf()
ff.apply_ddf()
sens_unc, forc_unc = [], []
for _ in range(n_steps):
ff.run(SAMPLE_INTERVAL, np.zeros(7, dtype=DATA_TYPE))
obs_slice = ff.obs.copy()[2:14]
sens_unc.append(obs_slice[0:6])
forc_unc.append(obs_slice[6:12])
np.savez(os.path.join(output_root, "uncontrolled.npz"),
sensors=np.array(sens_unc, dtype=np.float32),
forces=np.array(forc_unc, dtype=np.float32))
# Uncontrolled vorticity
omega_unc = vorticity_from_ddf(ff, u0=U0)
save_vorticity_png(os.path.join(output_root, "vorticity_uncontrolled.png"),
omega_unc, title=f"{label} uncontrolled")
# --- SINDy controlled rollout ---
print(f" SINDy controlled rollout ({n_steps} steps) ...")
ff.restore_ddf()
ff.apply_ddf()
# Bias FIFO
fifo = deque(maxlen=FIFO_LEN)
bias_action = scale_action(
np.zeros(3, dtype=np.float32),
scale=ACTION_SCALE, bias=ACTION_BIAS, u0=U0, n_total_bodies=7,
)
for _ in range(FIFO_LEN):
ff.run(SAMPLE_INTERVAL, bias_action)
fifo.append(ff.obs.copy()[2:14])
sens_sc, forc_sc, omega_sc = [], [], []
omega_bias = action_to_physical(
np.zeros((1, 3), dtype=np.float32),
scale=ACTION_SCALE, bias=ACTION_BIAS, u0=U0,
).flatten()
actions_prev = omega_bias.copy()
actions_prev2 = omega_bias.copy()
for step in range(n_steps):
obs_slice = fifo[-1] if len(fifo) > 0 else np.zeros(12, dtype=np.float32)
omega_pred = predict_omega_v3(obs_slice, actions_prev, actions_prev2, coefs, mu, u0=U0)
omega_sc.append(omega_pred.copy())
# Convert action to legacy array and apply
norm_action = (omega_pred / U0 - np.array(ACTION_BIAS, dtype=np.float64)) / ACTION_SCALE
norm_action = np.clip(norm_action, -1.0, 1.0).astype(np.float32)
action_arr = scale_action(
norm_action,
scale=ACTION_SCALE, bias=ACTION_BIAS, u0=U0, n_total_bodies=7,
)
ff.run(SAMPLE_INTERVAL, action_arr)
obs_slice_new = ff.obs.copy()[2:14]
fifo.append(obs_slice_new)
sens_sc.append(obs_slice_new[0:6])
forc_sc.append(obs_slice_new[6:12])
actions_prev = omega_pred
sens_sc_arr = np.array(sens_sc, dtype=np.float32)
forc_sc_arr = np.array(forc_sc, dtype=np.float32)
omega_sc_arr = np.array(omega_sc, dtype=np.float32)
np.savez(os.path.join(output_root, "sindy_controlled.npz"),
sensors=sens_sc_arr, forces=forc_sc_arr,
omegas=omega_sc_arr)
# Vorticity
omega_vort = vorticity_from_ddf(ff, u0=U0)
save_vorticity_png(os.path.join(output_root, "vorticity_sindy_controlled.png"),
omega_vort, title=f"{label} SINDy-controlled")
# Similarity
sim = compute_similarity(target_states, sens_sc_arr, CONV_LEN)
print(f" SINDy similarity: {sim:.4f}")
del ff
result = {"re_code": re_code, "mu": mu,
"sindy_similarity": sim,
"n_steps": n_steps}
with open(os.path.join(output_root, "result.json"), "w") as f:
json.dump(result, f, indent=2)
return result
def main():
ap = argparse.ArgumentParser(description="Phase 3: cross-Re SINDy validation")
ap.add_argument("--sindy-results", type=str,
default=os.path.join(OUTPUT_DIR, "sindy", "sindy_results_v3.json"),
help="Path to Phase 2 SINDy results JSON (v3 dimensionless)")
ap.add_argument("--validate-re", type=str, default="35,70,150",
help="Comma-separated validation Re codes")
ap.add_argument("--device", type=int, default=0, help="GPU device ID")
ap.add_argument("--steps", type=int, default=150,
help="Number of inference steps")
ap.add_argument("--threshold", type=float, default=0.002,
help="SINDy sparsity threshold (default: 0.002)")
ap.add_argument("--out-dir", type=str,
default=os.path.join(OUTPUT_DIR, "sindy_val"),
help="Output root for validation results")
args = ap.parse_args()
validate_re = [int(r) for r in args.validate_re.split(",")]
os.makedirs(args.out_dir, exist_ok=True)
# Load cross-Re coefficients
print(f"\nLoading cross-Re coefficients from {args.sindy_results}")
coefs = load_cross_re_coef(args.sindy_results, args.threshold)
for name in ["front", "bottom", "top"]:
print(f" {name}: nz={coefs[name]['nz']}, R2={coefs[name]['r2']:.4f}, "
f"threshold={args.threshold}")
t_start = time.time()
# Run for each validation Re
for re_code in validate_re:
out_dir = os.path.join(args.out_dir, f"re{re_code}")
result = run_sindy_controlled(
re_code, coefs, args.device, out_dir, n_steps=args.steps,
)
print(f" Done: Re{re_code} -> {out_dir}")
elapsed = time.time() - t_start
print(f"\nTotal time: {elapsed:.1f}s")
return 0
if __name__ == "__main__":
sys.exit(main())