134 lines
4.8 KiB
Python
134 lines
4.8 KiB
Python
"""SINDy fitting for Illusion scenes.
|
|
|
|
Usage:
|
|
conda run -n pycuda_3_10 python sindy/run_illusion.py
|
|
conda run -n pycuda_3_10 python sindy/run_illusion.py --diameters 0.75,1.0
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
from typing import List, Optional
|
|
|
|
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)
|
|
_SRC = os.path.join(_REPO, "src")
|
|
if _SRC not in sys.path:
|
|
sys.path.insert(0, _SRC)
|
|
|
|
from SR_analysis.utils.sindy_fitter import fit_sindy, get_feature_matrix_from_data
|
|
from SR_analysis.configs import get_scene, get_scene_list
|
|
|
|
SINDY_DIR = os.path.join(os.path.dirname(__file__), "..", "sindy", "illusion")
|
|
THRESHOLDS = [0.0, 0.001, 0.002, 0.005, 0.01, 0.015, 0.02, 0.03, 0.05, 0.1]
|
|
|
|
|
|
def load_data(scene_name: str) -> tuple:
|
|
data_dir = os.path.join(os.path.dirname(__file__), "..", "data", "illusion", scene_name)
|
|
npz = np.load(os.path.join(data_dir, "controlled.npz"))
|
|
sensors = npz["sensors"].astype(np.float64)
|
|
forces = npz["forces"].astype(np.float64)
|
|
actions = npz["actions"].astype(np.float64)
|
|
return sensors, forces, actions
|
|
|
|
|
|
def run(scene_names: Optional[List[str]] = None):
|
|
if scene_names is None:
|
|
scene_names = get_scene_list("illusion")
|
|
|
|
per_scene = {}
|
|
|
|
for sn in scene_names:
|
|
print(f"\n{'='*60}")
|
|
print(f"Scene: {sn}")
|
|
print(f"{'='*60}")
|
|
|
|
cfg = get_scene(sn)
|
|
sensors, forces, actions_phys = load_data(sn)
|
|
mu = cfg["mu"]
|
|
print(f" T={sensors.shape[0]}, mu={mu:.6f}")
|
|
|
|
Theta_f, Theta_r, Y, fn_f, fn_r = get_feature_matrix_from_data(
|
|
sensors, forces, actions_phys, mu, u0=cfg["u0"],
|
|
alpha_mode=False, include_mu=True, n_warmup=2,
|
|
)
|
|
print(f" Front: {Theta_f.shape}, Rear: {Theta_r.shape}")
|
|
|
|
# Front channel
|
|
print(f"\n --- Front (no bias) ---")
|
|
front_results = fit_sindy(Theta_f, Y[:, 0], THRESHOLDS)
|
|
best_f = max(front_results, key=lambda r: r["r2"])
|
|
print(f" Best: th={best_f['threshold']:.4f} nz={best_f['nz']:2d} R2={best_f['r2']:.6f}")
|
|
|
|
# Top channel (rear shared-head)
|
|
print(f"\n --- Top (rear shared-head) ---")
|
|
top_results = fit_sindy(Theta_r, Y[:, 2], THRESHOLDS)
|
|
best_t = max(top_results, key=lambda r: r["r2"])
|
|
print(f" Best: th={best_t['threshold']:.4f} nz={best_t['nz']:2d} R2={best_t['r2']:.6f}")
|
|
|
|
# Bottom (independent)
|
|
print(f"\n --- Bottom (independent) ---")
|
|
bot_results = fit_sindy(Theta_r, Y[:, 1], THRESHOLDS)
|
|
best_b = max(bot_results, key=lambda r: r["r2"])
|
|
print(f" Best: th={best_b['threshold']:.4f} nz={best_b['nz']:2d} R2={best_b['r2']:.6f}")
|
|
|
|
per_scene[sn] = {
|
|
"scene": sn,
|
|
"re_code": cfg["re_code"],
|
|
"mu": mu,
|
|
"n_samples": Theta_f.shape[0],
|
|
"feature_names_front": fn_f,
|
|
"feature_names_rear": fn_r,
|
|
"front": {
|
|
"results": [{k: v for k, v in r.items() if k != "coef"} for r in front_results],
|
|
"best": {k: v for k, v in best_f.items() if k != "coef"},
|
|
"best_coef": best_f["coef"],
|
|
"sparsity_curve": [(r["threshold"], r["nz"], r["r2"]) for r in front_results],
|
|
},
|
|
"top": {
|
|
"results": [{k: v for k, v in r.items() if k != "coef"} for r in top_results],
|
|
"best": {k: v for k, v in best_t.items() if k != "coef"},
|
|
"best_coef": best_t["coef"],
|
|
"sparsity_curve": [(r["threshold"], r["nz"], r["r2"]) for r in top_results],
|
|
},
|
|
"bottom": {
|
|
"results": [{k: v for k, v in r.items() if k != "coef"} for r in bot_results],
|
|
"best": {k: v for k, v in best_b.items() if k != "coef"},
|
|
"best_coef": best_b["coef"],
|
|
"sparsity_curve": [(r["threshold"], r["nz"], r["r2"]) for r in bot_results],
|
|
},
|
|
}
|
|
|
|
os.makedirs(SINDY_DIR, exist_ok=True)
|
|
out_path = os.path.join(SINDY_DIR, "sindy_results.json")
|
|
result = {"thresholds": THRESHOLDS, "per_scene": per_scene}
|
|
with open(out_path, "w") as f:
|
|
json.dump(result, f, indent=2)
|
|
print(f"\nSaved: {out_path}")
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--diameters", type=str, default=None,
|
|
help="Comma-separated diameters (e.g. 0.75,1.0,1.5)")
|
|
ap.add_argument("--scene-names", type=str, default=None)
|
|
args = ap.parse_args()
|
|
|
|
if args.scene_names:
|
|
names = [s.strip() for s in args.scene_names.split(",")]
|
|
elif args.diameters:
|
|
names = [f"illusion_{d.strip()}L" for d in args.diameters.split(",")]
|
|
else:
|
|
names = None
|
|
|
|
run(names)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|