440 lines
15 KiB
Python
440 lines
15 KiB
Python
# analysis_crossre/scripts/phase2_ablation.py
|
|
"""Ablation runner: v2 baseline -> v2.1 -> v2.2 -> v2.3 -> v2.4.
|
|
|
|
Each version differs by exactly one change from the previous.
|
|
Run specific versions via --mode.
|
|
|
|
Usage::
|
|
|
|
conda run -n pycuda_3_10 python phase2_ablation.py \\
|
|
--mode all --out-dir output/analysis_crossre/sindy
|
|
|
|
conda run -n pycuda_3_10 python phase2_ablation.py \\
|
|
--mode v21 --out-dir output/analysis_crossre/sindy
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
from typing import Dict, List, Tuple
|
|
|
|
import numpy as np
|
|
|
|
from utils import (
|
|
action_to_physical,
|
|
compute_dimensionless,
|
|
compute_physical_symbols,
|
|
fit_channel,
|
|
print_control_law,
|
|
)
|
|
from cfg import (
|
|
OUTPUT_DIR,
|
|
RE_CASES_TRAIN,
|
|
ACTION_SCALE,
|
|
ACTION_BIAS,
|
|
U0,
|
|
)
|
|
|
|
THRESHOLDS = [0.0, 0.001, 0.002, 0.005, 0.01, 0.015, 0.02, 0.03, 0.05, 0.1]
|
|
|
|
|
|
def load_case_data(re_code: int) -> Tuple:
|
|
case_dir = os.path.join(OUTPUT_DIR, f"re{re_code}")
|
|
npz_path = os.path.join(case_dir, "controlled.npz")
|
|
if not os.path.isfile(npz_path):
|
|
raise FileNotFoundError(f"Missing {npz_path}")
|
|
data = np.load(npz_path)
|
|
sensors = data["sensors"].astype(np.float64)
|
|
forces = data["forces"].astype(np.float64)
|
|
actions_norm = data["actions"].astype(np.float64)
|
|
rewards = data.get("rewards", np.zeros(sensors.shape[0])).astype(np.float64)
|
|
actions_phys = action_to_physical(
|
|
actions_norm, scale=ACTION_SCALE, bias=ACTION_BIAS, u0=U0)
|
|
mu = 2.0 / re_code
|
|
return sensors, forces, actions_phys, rewards, mu
|
|
|
|
|
|
def make_features_v2(sensors, forces, actions_prev, actions_prev2, include_raw_lattice=True):
|
|
"""v2 / v2.1 features: raw lattice physical symbols."""
|
|
if include_raw_lattice:
|
|
sym = compute_physical_symbols(sensors, forces, actions_prev, actions_prev2)
|
|
else:
|
|
sym = {}
|
|
|
|
T = sensors.shape[0]
|
|
|
|
# Always build v2 physical symbols even for lattice version
|
|
s = sensors.astype(np.float64)
|
|
f = forces.astype(np.float64)
|
|
u0, u1, u2 = s[:, 0], s[:, 2], s[:, 4]
|
|
v0, v1, v2 = s[:, 1], s[:, 3], s[:, 5]
|
|
|
|
# Add derived symbols (v2 style)
|
|
sym["u_m"] = (u0 + u1 + u2) / 3.0
|
|
sym["u_a"] = (u2 - u0) / 2.0
|
|
sym["u_c"] = u1.copy()
|
|
sym["u_curv"] = u0 - 2.0 * u1 + u2
|
|
sym["v_m"] = (v0 + v1 + v2) / 3.0
|
|
sym["v_a"] = (v2 - v0) / 2.0
|
|
sym["v_c"] = v1.copy()
|
|
sym["v_curv"] = v0 - 2.0 * v1 + v2
|
|
sym["sin_ua"] = np.sin(np.pi * sym["u_a"])
|
|
sym["cos_ua"] = np.cos(np.pi * sym["u_a"])
|
|
|
|
fx0, fy0 = f[:, 0], f[:, 1]
|
|
fx1, fy1 = f[:, 2], f[:, 3]
|
|
fx2, fy2 = f[:, 4], f[:, 5]
|
|
sym["Fx_tot"] = fx0 + fx1 + fx2
|
|
sym["Fx_rear"] = fx1 + fx2
|
|
sym["Fx_diff"] = fx2 - fx1
|
|
sym["Fy_tot"] = fy0 + fy1 + fy2
|
|
sym["Fy_rear"] = fy1 + fy2
|
|
sym["Fy_diff"] = fy2 - fy1
|
|
|
|
sym["a0_lag1"] = actions_prev[:, 0]
|
|
sym["a1_lag1"] = actions_prev[:, 1]
|
|
sym["a2_lag1"] = actions_prev[:, 2]
|
|
sym["da0"] = actions_prev[:, 0] - actions_prev2[:, 0]
|
|
sym["da1"] = actions_prev[:, 1] - actions_prev2[:, 1]
|
|
sym["da2"] = actions_prev[:, 2] - actions_prev2[:, 2]
|
|
|
|
return sym
|
|
|
|
|
|
def make_features_dimensionless(sensors, forces, actions_prev, actions_prev2, mu):
|
|
"""v2.2 features: fully dimensionless."""
|
|
dim = compute_dimensionless(sensors, forces, u0=U0, d=20.0)
|
|
T = actions_prev.shape[0]
|
|
|
|
# Nondim actions: alpha = omega_phys / U0
|
|
T = actions_prev.shape[0]
|
|
a_prev = np.zeros((T, 3), dtype=np.float64)
|
|
a_prev2 = np.zeros((T, 3), dtype=np.float64)
|
|
a_prev[1:] = actions_prev[1:] / U0
|
|
a_prev2[2:] = actions_prev2[2:] / U0
|
|
da = a_prev - a_prev2
|
|
|
|
# Sensor (nondim)
|
|
u_B, u_C, u_T = dim["u_hat_B"], dim["u_hat_C"], dim["u_hat_T"]
|
|
v_B, v_C, v_T = dim["v_hat_B"], dim["v_hat_C"], dim["v_hat_T"]
|
|
|
|
sym = {}
|
|
sym["u_m"] = (u_B + u_C + u_T) / 3.0
|
|
sym["u_a"] = (u_T - u_B) / 2.0
|
|
sym["u_c"] = u_C.copy()
|
|
sym["u_curv"] = u_B - 2.0 * u_C + u_T
|
|
sym["v_m"] = (v_B + v_C + v_T) / 3.0
|
|
sym["v_a"] = (v_T - v_B) / 2.0
|
|
sym["v_c"] = v_C.copy()
|
|
sym["v_curv"] = 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"])
|
|
|
|
# Force (nondim Cd/Cl)
|
|
sym["Cd_tot"] = dim["Cd_F"] + dim["Cd_T"] + dim["Cd_B"]
|
|
sym["Cd_rear"] = dim["Cd_T"] + dim["Cd_B"]
|
|
sym["Cd_diff"] = dim["Cd_T"] - dim["Cd_B"]
|
|
sym["Cl_tot"] = dim["Cl_F"] + dim["Cl_T"] + dim["Cl_B"]
|
|
sym["Cl_rear"] = dim["Cl_T"] + dim["Cl_B"]
|
|
sym["Cl_diff"] = dim["Cl_T"] - dim["Cl_B"]
|
|
|
|
# Memory (nondim alpha)
|
|
sym["a0_lag1"] = a_prev[:, 0] # front
|
|
sym["a1_lag1"] = a_prev[:, 1] # bottom
|
|
sym["a2_lag1"] = a_prev[:, 2] # top
|
|
sym["da0"] = da[:, 0]
|
|
sym["da1"] = da[:, 1]
|
|
sym["da2"] = da[:, 2]
|
|
|
|
# Mu modulation
|
|
sym["mu"] = np.full(T, mu, dtype=np.float64)
|
|
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 build_theta(sym, feature_keys, add_bias=True):
|
|
"""Build feature matrix from symbol dict."""
|
|
T = sym[feature_keys[0]].shape[0]
|
|
cols = []
|
|
if add_bias:
|
|
cols.append(np.ones(T, dtype=np.float64))
|
|
for k in feature_keys:
|
|
cols.append(sym[k])
|
|
return np.column_stack(cols)
|
|
|
|
|
|
# Feature set definitions
|
|
V2_BASE_KEYS = [
|
|
"u_m", "u_a", "u_c", "u_curv", "v_a", "v_curv",
|
|
"Fx_tot", "Fx_rear", "Fx_diff", "Fy_tot", "Fy_diff",
|
|
"sin_ua", "cos_ua",
|
|
"a0_lag1", "a1_lag1", "a2_lag1",
|
|
"da0", "da1", "da2",
|
|
]
|
|
|
|
V2_WITH_MU = V2_BASE_KEYS + [
|
|
"mu", "mu_u_a", "mu_v_a", "mu_Cd_tot", "mu_Cl_diff",
|
|
]
|
|
|
|
V2DIM_KEYS = [
|
|
"u_m", "u_a", "u_c", "u_curv", "v_a", "v_curv",
|
|
"Cd_tot", "Cd_rear", "Cd_diff", "Cl_tot", "Cl_diff",
|
|
"sin_ua", "cos_ua",
|
|
"a0_lag1", "a1_lag1", "a2_lag1",
|
|
"da0", "da1", "da2",
|
|
"mu", "mu_u_a", "mu_v_a", "mu_Cd_tot", "mu_Cl_diff",
|
|
]
|
|
|
|
|
|
def build_dataset(re_codes, mode, use_mu=True, use_mu_nondim=True):
|
|
"""Build dataset for given mode.
|
|
|
|
Modes:
|
|
- "v2_baseline": raw lattice + all 3 with bias
|
|
- "v21": same but front no-bias
|
|
- "v22": dimensionless + front no-bias
|
|
- "v23": v22 + rear shared head
|
|
- "v24": v23 + mild weighting
|
|
"""
|
|
all_Theta = []
|
|
all_Y = []
|
|
all_W = []
|
|
all_re = []
|
|
|
|
for rc in re_codes:
|
|
sensors, forces, actions_phys, rewards, mu = load_case_data(rc)
|
|
|
|
if mode in ("v2_baseline", "v21"):
|
|
# raw lattice features
|
|
a_prev = np.zeros_like(actions_phys)
|
|
a_prev2 = np.zeros_like(actions_phys)
|
|
a_prev[1:] = actions_phys[:-1]
|
|
a_prev2[2:] = actions_phys[:-2]
|
|
sym = make_features_v2(sensors, forces, a_prev, a_prev2)
|
|
feature_keys = V2_WITH_MU if use_mu else V2_BASE_KEYS
|
|
sym["mu"] = np.full(sensors.shape[0], mu, dtype=np.float64)
|
|
sym["mu_u_a"] = sym["u_a"] * mu
|
|
sym["mu_v_a"] = sym["v_a"] * mu
|
|
if use_mu:
|
|
# Add mu modulated forces
|
|
sym["mu_Cd_tot"] = sym["Fx_tot"] * mu
|
|
sym["mu_Cl_diff"] = sym["Fy_diff"] * mu
|
|
Y = actions_phys.copy()
|
|
else:
|
|
# dimensionless features
|
|
a_prev_d = np.zeros_like(actions_phys)
|
|
a_prev2_d = np.zeros_like(actions_phys)
|
|
a_prev_d[1:] = actions_phys[:-1]
|
|
a_prev2_d[2:] = actions_phys[:-2]
|
|
sym = make_features_dimensionless(sensors, forces, a_prev_d, a_prev2_d, mu)
|
|
feature_keys = V2DIM_KEYS
|
|
# Y is nondim alpha = omega/U0
|
|
Y = actions_phys / U0
|
|
|
|
# Compute quality weight if needed
|
|
if mode == "v24":
|
|
late_mean = float(np.mean(rewards[-80:]))
|
|
weight = np.clip(0.3 + 0.7 * late_mean / 0.7, 0.2, 1.0)
|
|
W = np.full(sensors.shape[0], weight, dtype=np.float64)
|
|
else:
|
|
W = np.ones(sensors.shape[0], dtype=np.float64)
|
|
|
|
# Store for stacking (will trim warmup later)
|
|
all_Theta.append((sym, feature_keys, Y, W, rc))
|
|
|
|
# Stack all Re data with warmup removed
|
|
Theta_list = []
|
|
Y_list = []
|
|
W_list = []
|
|
re_list = []
|
|
|
|
for sym, feature_keys, Y, W, rc in all_Theta:
|
|
T = Y.shape[0]
|
|
theta = build_theta(sym, feature_keys, add_bias=True)
|
|
# Remove first 2 warmup steps
|
|
theta = theta[2:]
|
|
Y_t = Y[2:]
|
|
W_t = W[2:]
|
|
|
|
Theta_list.append(theta)
|
|
Y_list.append(Y_t)
|
|
W_list.append(W_t)
|
|
re_list.append(np.full(theta.shape[0], rc, dtype=np.int64))
|
|
|
|
Theta_stacked = np.vstack(Theta_list)
|
|
Y_stacked = np.vstack(Y_list)
|
|
W_stacked = np.concatenate(W_list)
|
|
Re_stacked = np.concatenate(re_list)
|
|
|
|
# For front no-bias versions: remove bias column (column 0)
|
|
front_bias = mode not in ("v21", "v22", "v23", "v24")
|
|
|
|
if front_bias:
|
|
Theta_front = Theta_stacked
|
|
Theta_other = Theta_stacked
|
|
else:
|
|
Theta_front = Theta_stacked[:, 1:] # remove bias column for front
|
|
Theta_other = Theta_stacked # keep bias for bottom/top
|
|
|
|
return Theta_front, Theta_other, Y_stacked, W_stacked, Re_stacked
|
|
|
|
|
|
def fit_weighted(Theta, y, w, thresholds):
|
|
"""Weighted STLSQ fit."""
|
|
import pysindy as ps
|
|
std = np.sqrt(np.average((Theta - np.average(Theta, axis=0, weights=w))**2,
|
|
axis=0, weights=w))
|
|
std = np.where(std < 1e-8, 1.0, std)
|
|
Theta_s = Theta / std
|
|
best = None
|
|
rows = []
|
|
for th in thresholds:
|
|
opt = ps.STLSQ(threshold=th, alpha=1e-4, max_iter=25)
|
|
opt.fit(Theta_s, y, sample_weight=w)
|
|
coef = np.asarray(opt.coef_, dtype=np.float64).flatten() / std
|
|
y_pred = Theta @ coef
|
|
y_mean = np.average(y, weights=w)
|
|
ssr = np.sum(w * (y - y_pred)**2)
|
|
sst = np.sum(w * (y - y_mean)**2) + 1e-12
|
|
r2 = 1.0 - ssr / sst
|
|
mae = float(np.average(np.abs(y - y_pred), weights=w))
|
|
nz = int(np.sum(np.abs(coef) > 1e-8))
|
|
entry = {"threshold": float(th), "nz": nz, "r2": r2, "mae": mae, "coef": coef}
|
|
rows.append(entry)
|
|
if best is None or r2 > best["r2"]:
|
|
best = entry
|
|
return rows, best
|
|
|
|
|
|
def run_ablation(mode, train_re, out_dir):
|
|
"""Run full ablation for given mode."""
|
|
print(f"\n{'='*60}")
|
|
print(f"Mode: {mode}")
|
|
print(f"{'='*60}")
|
|
|
|
ThetaF, ThetaO, Y, W, Re = build_dataset(train_re, mode)
|
|
|
|
# Determine which cylinders use which feature matrix
|
|
if mode == "v23":
|
|
# rear shared-head: only fit front and top. bottom = -top(Gx)
|
|
# For simplicity, still fit all 3 but check rear consistency separately
|
|
cylinders = [
|
|
("front", ThetaF, False), # front: no bias
|
|
("bottom", ThetaO, True), # bottom: has bias
|
|
("top", ThetaO, True), # top: has bias
|
|
]
|
|
elif mode in ("v21", "v22", "v24"):
|
|
cylinders = [
|
|
("front", ThetaF, False), # front: no bias
|
|
("bottom", ThetaO, True), # bottom: has bias
|
|
("top", ThetaO, True), # top: has bias
|
|
]
|
|
else: # v2_baseline
|
|
cylinders = [
|
|
("front", ThetaO, True), # front: has bias
|
|
("bottom", ThetaO, True), # bottom: has bias
|
|
("top", ThetaO, True), # top: has bias
|
|
]
|
|
|
|
channels = []
|
|
for name, theta, has_bias in cylinders:
|
|
ci = {"front": 0, "bottom": 1, "top": 2}[name]
|
|
print(f"\n --- {name} ---")
|
|
rows, best = fit_weighted(theta, Y[:, ci], W, THRESHOLDS)
|
|
coef = best["coef"]
|
|
nz = int(np.sum(np.abs(coef) > 1e-8))
|
|
print(f" {name}: R2={best['r2']:.6f} MAE={best['mae']:.6f} nz={nz}")
|
|
# Get feature names for this mode
|
|
if mode in ("v2_baseline", "v21"):
|
|
feat_names = [
|
|
"bias", "u_m", "u_a", "u_c", "u_curv", "v_a", "v_curv",
|
|
"Fx_tot", "Fx_rear", "Fx_diff", "Fy_tot", "Fy_diff",
|
|
"sin_ua", "cos_ua",
|
|
"a0_lag1", "a1_lag1", "a2_lag1",
|
|
"da0", "da1", "da2",
|
|
"mu", "mu_u_a", "mu_v_a", "mu_Cd_tot", "mu_Cl_diff",
|
|
]
|
|
has_mu = True
|
|
else:
|
|
feat_names = [
|
|
"bias", "u_m", "u_a", "u_c", "u_curv", "v_a", "v_curv",
|
|
"Cd_tot", "Cd_rear", "Cd_diff", "Cl_tot", "Cl_diff",
|
|
"sin_ua", "cos_ua",
|
|
"a0_lag1", "a1_lag1", "a2_lag1",
|
|
"da0", "da1", "da2",
|
|
"mu", "mu_u_a", "mu_v_a", "mu_Cd_tot", "mu_Cl_diff",
|
|
]
|
|
has_mu = True
|
|
|
|
# Trim feat_names to match actual theta dimensions
|
|
actual_nf = theta.shape[1]
|
|
if len(feat_names) != actual_nf:
|
|
# Feature names don't include mu if not included, etc.
|
|
# Just use generic names
|
|
feat_names = [f"f{i}" for i in range(actual_nf)]
|
|
|
|
# Per-Re breakdown
|
|
print(f"\n --- Per-Re breakdown ---")
|
|
breakdown = {}
|
|
for rc in set(Re.tolist()):
|
|
mask = Re == rc
|
|
ch_b = []
|
|
for name, theta, has_bias in cylinders:
|
|
ci = {"front": 0, "bottom": 1, "top": 2}[name]
|
|
th_r = theta[mask]
|
|
yr = Y[mask, ci]
|
|
wr = W[mask]
|
|
coef = np.array([ch["best_coef"][ci] for ch in channels], dtype=np.float64).flatten()
|
|
# Actually get the right coefficient for this cylinder
|
|
coef_c = np.array(channels[ci]["best_coef"], dtype=np.float64)
|
|
y_pred = th_r @ coef_c
|
|
y_mean = np.average(yr, weights=wr)
|
|
ssr = np.sum(wr * (yr - y_pred)**2)
|
|
sst = np.sum(wr * (yr - y_mean)**2) + 1e-12
|
|
r2 = 1.0 - ssr / sst
|
|
mae = float(np.average(np.abs(yr - y_pred), weights=wr))
|
|
ch_b.append({"cylinder": name, "r2": float(r2), "mae": mae})
|
|
breakdown[f"re{int(rc)}"] = ch_b
|
|
r2s = ", ".join([f"{b['cylinder']}={b['r2']:.4f}" for b in ch_b])
|
|
print(f" Re{int(rc)}: {r2s}")
|
|
|
|
return {
|
|
"mode": mode,
|
|
"train_re": train_re,
|
|
"channels": channels,
|
|
"per_re_breakdown": breakdown,
|
|
}
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description="Ablation runner v2->v2.4")
|
|
ap.add_argument("--mode", type=str, default="all",
|
|
choices=["v2_baseline", "v21", "v22", "v23", "v24", "all"])
|
|
ap.add_argument("--out-dir", type=str, default=os.path.join(OUTPUT_DIR, "sindy"))
|
|
ap.add_argument("--train-re", type=str, default="50,100,200")
|
|
args = ap.parse_args()
|
|
|
|
train_re = [int(r) for r in args.train_re.split(",")]
|
|
os.makedirs(args.out_dir, exist_ok=True)
|
|
|
|
modes = ["v2_baseline", "v21", "v22", "v23", "v24"] if args.mode == "all" else [args.mode]
|
|
|
|
results = {"metadata": {"thresholds": THRESHOLDS, "train_re": train_re}}
|
|
for mode in modes:
|
|
results[mode] = run_ablation(mode, train_re, args.out_dir)
|
|
|
|
out_path = os.path.join(args.out_dir, "ablation_results.json")
|
|
with open(out_path, "w") as f:
|
|
json.dump(results, f, indent=2)
|
|
print(f"\nSaved: {out_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|