268 lines
9.7 KiB
Python
268 lines
9.7 KiB
Python
"""Phase 3: closed-loop for ablation modes.
|
|
|
|
Usage::
|
|
|
|
conda run -n pycuda_3_10 python phase3_ablation_val.py \\
|
|
--ablation-json output/analysis_crossre/sindy/ablation_results.json \\
|
|
--mode v21 --validate-re 70 --device 2
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
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
|
|
|
|
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_physical_symbols,
|
|
save_vorticity_png, vorticity_from_ddf, compute_similarity,
|
|
)
|
|
from cfg import (
|
|
CONFIG_DIR, OUTPUT_DIR, SAMPLE_INTERVAL, FIFO_LEN, CONV_LEN,
|
|
S_DIM, ACTION_SCALE, ACTION_BIAS, U0,
|
|
)
|
|
|
|
DATA_TYPE = np.float32
|
|
|
|
|
|
def load_ablation_coef(ablation_path, mode, channels_to_load=("front", "bottom", "top")):
|
|
"""Load coefficients for a specific ablation mode."""
|
|
with open(ablation_path) as f:
|
|
data = json.load(f)
|
|
|
|
mode_data = data[mode]
|
|
coefs = {}
|
|
for ch in mode_data["channels"]:
|
|
name = ch["cylinder"]
|
|
if name in channels_to_load:
|
|
coefs[name] = {
|
|
"coef": np.array(ch["best_coef"], dtype=np.float64),
|
|
"has_bias": ch["has_bias"],
|
|
}
|
|
return coefs
|
|
|
|
|
|
def predict_ablation(obs_slice, actions_prev, actions_prev2, coefs, mu, mode, u0=0.01):
|
|
"""Predict action using ablation mode coefficients.
|
|
|
|
obs_slice: (12,) raw lattice [sensor(6), force(6)]
|
|
"""
|
|
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)
|
|
|
|
is_dim = "v22" in mode or "v23" in mode or "v24" in mode or mode in ("v2_dimless",)
|
|
|
|
if is_dim:
|
|
dim = compute_dimensionless(sensors, forces, u0=u0, d=20.0)
|
|
# Build dimensionless features
|
|
sym = _build_dimensionless_features(dim, a_prev[0], a_prev2[0], mu)
|
|
Y_scale = 1.0 / u0 # predict nondim alpha = omega / U0
|
|
else:
|
|
# Lattice features (v2/v21)
|
|
a_prev_f = np.zeros((1, 3), dtype=np.float64)
|
|
a_prev2_f = np.zeros((1, 3), dtype=np.float64)
|
|
a_prev_f[0] = a_prev
|
|
a_prev2_f[0] = a_prev2
|
|
sym = _build_lattice_features(sensors, forces, a_prev_f, a_prev2_f, mu)
|
|
Y_scale = 1.0 # predict raw omega
|
|
|
|
# Build feature vector
|
|
feat_keys = [k for k in sym.keys() if k not in ("mu", "mu_u_a", "mu_v_a", "mu_Cd_tot", "mu_Cl_diff")]
|
|
feat_keys_mu = ["mu", "mu_u_a", "mu_v_a", "mu_Cd_tot", "mu_Cl_diff"]
|
|
|
|
feats = []
|
|
for k in feat_keys:
|
|
feats.append(float(sym[k][0]) if isinstance(sym[k], np.ndarray) else float(sym[k]))
|
|
if mu > 0:
|
|
for k in feat_keys_mu:
|
|
feats.append(float(sym[k][0]) if isinstance(sym[k], np.ndarray) else float(sym[k]))
|
|
|
|
omega = np.zeros(3, dtype=np.float64)
|
|
for ci, name in enumerate(["front", "bottom", "top"]):
|
|
c = coefs.get(name)
|
|
if c is None:
|
|
continue
|
|
coef_arr = c["coef"]
|
|
has_bias = c["has_bias"]
|
|
|
|
if has_bias:
|
|
feat_vec = np.array([1.0] + feats) if len(coef_arr) == len(feats) + 1 else np.array(feats)
|
|
else:
|
|
feat_vec = np.array(feats) if len(coef_arr) == len(feats) else np.array([1.0] + feats)
|
|
|
|
if len(feat_vec) != len(coef_arr):
|
|
feat_vec = np.array(feats) # fallback
|
|
|
|
pred = float(feat_vec @ coef_arr) * Y_scale
|
|
omega[ci] = pred
|
|
|
|
return omega
|
|
|
|
|
|
def _build_lattice_features(sensors, forces, a_prev, a_prev2, mu):
|
|
"""Build v2-style lattice features. All args 2D (1, N)."""
|
|
# Ensure 2D
|
|
if sensors.ndim == 1:
|
|
sensors = sensors.reshape(1, -1)
|
|
if forces.ndim == 1:
|
|
forces = forces.reshape(1, -1)
|
|
if a_prev.ndim == 1:
|
|
a_prev = a_prev.reshape(1, -1)
|
|
if a_prev2.ndim == 1:
|
|
a_prev2 = a_prev2.reshape(1, -1)
|
|
sym = compute_physical_symbols(sensors, forces, a_prev, a_prev2)
|
|
sym["mu"] = np.array([mu])
|
|
sym["mu_u_a"] = sym["u_a"] * mu
|
|
sym["mu_v_a"] = sym["v_a"] * mu
|
|
sym["mu_Cd_tot"] = sym["Fx_tot"] * mu
|
|
sym["mu_Cl_diff"] = sym["Fy_diff"] * mu
|
|
return sym
|
|
|
|
|
|
def _build_dimensionless_features(dim, a_prev, a_prev2, mu):
|
|
"""Build dimensionless features. a_prev/a_prev2 are 1D (3,) arrays."""
|
|
if a_prev.ndim > 1:
|
|
a_prev = a_prev.flatten()
|
|
if a_prev2.ndim > 1:
|
|
a_prev2 = a_prev2.flatten()
|
|
"""Build dimensionless features."""
|
|
T = 1
|
|
u_B, u_C, u_T = dim["u_hat_B"][0], dim["u_hat_C"][0], dim["u_hat_T"][0]
|
|
v_B, v_C, v_T = dim["v_hat_B"][0], dim["v_hat_C"][0], dim["v_hat_T"][0]
|
|
|
|
sym = {}
|
|
sym["u_m"] = np.array([(u_B + u_C + u_T) / 3.0])
|
|
sym["u_a"] = np.array([(u_T - u_B) / 2.0])
|
|
sym["u_c"] = np.array([u_C])
|
|
sym["u_curv"] = np.array([u_B - 2.0*u_C + u_T])
|
|
sym["v_a"] = np.array([(v_T - v_B) / 2.0])
|
|
sym["v_curv"] = np.array([v_B - 2.0*v_C + v_T])
|
|
sym["sin_ua"] = np.sin(np.pi * sym["u_a"])
|
|
sym["cos_ua"] = np.cos(np.pi * sym["u_a"])
|
|
|
|
sym["Cd_tot"] = np.array([dim["Cd_F"][0] + dim["Cd_T"][0] + dim["Cd_B"][0]])
|
|
sym["Cd_rear"] = np.array([dim["Cd_T"][0] + dim["Cd_B"][0]])
|
|
sym["Cd_diff"] = np.array([dim["Cd_T"][0] - dim["Cd_B"][0]])
|
|
sym["Cl_tot"] = np.array([dim["Cl_F"][0] + dim["Cl_T"][0] + dim["Cl_B"][0]])
|
|
sym["Cl_diff"] = np.array([dim["Cl_T"][0] - dim["Cl_B"][0]])
|
|
|
|
# Nondim actions
|
|
a_prev_n = a_prev / U0
|
|
a_prev2_n = a_prev2 / U0
|
|
sym["a0_lag1"] = np.array([a_prev_n[0]])
|
|
sym["a1_lag1"] = np.array([a_prev_n[1]])
|
|
sym["a2_lag1"] = np.array([a_prev_n[2]])
|
|
sym["da0"] = np.array([a_prev_n[0] - a_prev2_n[0]])
|
|
sym["da1"] = np.array([a_prev_n[1] - a_prev2_n[1]])
|
|
sym["da2"] = np.array([a_prev_n[2] - a_prev2_n[2]])
|
|
|
|
sym["mu"] = np.array([mu])
|
|
sym["mu_u_a"] = sym["u_a"] * mu
|
|
sym["mu_v_a"] = sym["v_a"] * mu
|
|
sym["mu_Cd_tot"] = sym["Cd_tot"] * mu
|
|
sym["mu_Cl_diff"] = sym["Cl_diff"] * mu
|
|
|
|
return sym
|
|
|
|
|
|
def run_closed_loop(re_code, coefs, mode, device_id, output_root, n_steps=100):
|
|
"""Run closed-loop for one ablation mode."""
|
|
os.makedirs(output_root, exist_ok=True)
|
|
nu = nu_from_re(re_code, u0=U0)
|
|
mu = 2.0 / re_code
|
|
|
|
cuda_cfg, field_cfg = load_legacy_configs(CONFIG_DIR)
|
|
field_cfg = field_cfg._replace(viscosity=float(nu))
|
|
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)
|
|
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)
|
|
|
|
# Controlled rollout
|
|
ff.restore_ddf()
|
|
ff.apply_ddf()
|
|
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 = []
|
|
actions_prev = action_to_physical(
|
|
np.zeros((1,3), dtype=np.float32), scale=ACTION_SCALE,
|
|
bias=ACTION_BIAS, u0=U0).flatten()
|
|
actions_prev2 = actions_prev.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_ablation(obs_slice, actions_prev, actions_prev2, coefs, mu, mode, u0=U0)
|
|
|
|
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])
|
|
actions_prev2 = actions_prev.copy()
|
|
actions_prev = omega_pred.copy()
|
|
|
|
sens_arr = np.array(sens_sc, dtype=np.float32)
|
|
sim = compute_similarity(target_states, sens_arr, CONV_LEN)
|
|
|
|
omega_vort = vorticity_from_ddf(ff, u0=U0)
|
|
save_vorticity_png(os.path.join(output_root, f"vorticity_{mode}.png"),
|
|
omega_vort, title=f"Re{re_code} {mode}")
|
|
|
|
del ff
|
|
result = {"re_code": re_code, "mode": mode, "similarity": sim}
|
|
with open(os.path.join(output_root, "result.json"), "w") as f:
|
|
json.dump(result, f, indent=2)
|
|
print(f" Re{re_code} {mode}: similarity={sim:.4f}")
|
|
return result
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--ablation-json", type=str, required=True)
|
|
ap.add_argument("--mode", type=str, default="v21")
|
|
ap.add_argument("--validate-re", type=str, default="70")
|
|
ap.add_argument("--device", type=int, default=2)
|
|
ap.add_argument("--steps", type=int, default=100)
|
|
ap.add_argument("--out-dir", type=str, default=os.path.join(OUTPUT_DIR, "sindy_val"))
|
|
args = ap.parse_args()
|
|
|
|
validate_re = [int(r) for r in args.validate_re.split(",")]
|
|
coefs = load_ablation_coef(args.ablation_json, args.mode)
|
|
os.makedirs(args.out_dir, exist_ok=True)
|
|
|
|
for rc in validate_re:
|
|
out_sub = os.path.join(args.out_dir, f"re{rc}")
|
|
run_closed_loop(rc, coefs, args.mode, args.device, out_sub, n_steps=args.steps)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|