# OID_analysis/scripts/collect_controlled.py """ Collect DRL-controlled rollout (q_ctl) for Karman cloak and illusion scenes. Generates field time series for Delta-q_ctl computation. Usage: # Karman: conda run -n pycuda_3_10 python src/OID_analysis/scripts/collect_controlled.py \ --scene karman_re100 --device 1 --steps 500 # Illusion (3 diameters): conda run -n pycuda_3_10 python src/OID_analysis/scripts/collect_controlled.py \ --scene illusion_1.0L --device 3 --steps 500 conda run -n pycuda_3_10 python src/OID_analysis/scripts/collect_controlled.py \ --scene illusion_0.75L --device 3 --steps 500 conda run -n pycuda_3_10 python src/OID_analysis/scripts/collect_controlled.py \ --scene illusion_1.5L --device 3 --steps 500 """ from __future__ import annotations import argparse import json import os import sys import time from collections import deque 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 LegacyCelerisLab import FlowField # noqa: E402 from OID_analysis.utils.cfd_interface import ( # noqa: E402 load_legacy_configs, get_velocity_field, save_vorticity_png, vorticity_from_ddf, load_ppo_model, scale_action, build_observation, compute_similarity, calc_lag, calc_dtw_sim, analyze_harmonics, gen_target_states_at, ) from OID_analysis.configs import ( # noqa: E402 get_scene, get_scene_list, model_path_for_scene, LEGACY_CFG_DIR, ) DATA_TYPE = np.float32 L0 = 20.0 CENTER_Y = (512 - 1) / 2.0 FIFO_LEN = 150 CONV_LEN_DEFAULT = 30 CONV_LEN_ILLUSION = 36 def collect_single(scene_name: str, device_id: int, n_steps: int) -> dict: cfg = get_scene(scene_name) u0 = cfg["u0"] si = cfg["sample_interval"] ac_scale = cfg["action_scale"] ac_bias = cfg["action_bias"] n_obj = cfg["n_objects_env"] s_dim = cfg["s_dim"] source = cfg.get("source", "") out_dir = data_dir_for_scene(scene_name) os.makedirs(out_dir, exist_ok=True) # Check DDF+FIFO checkpoint exists ddf_ckpt_path = os.path.join(out_dir, "ddf_checkpoint.npy") fifo_ckpt_path = os.path.join(out_dir, "fifo_checkpoint.npy") has_blk = os.path.isfile(ddf_ckpt_path) and os.path.isfile(fifo_ckpt_path) # Load legacy configs cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR) field_cfg = field_cfg._replace(viscosity=float(cfg["nu"]), velocity=float(u0)) if not has_blk: print(f" No DDF checkpoint found, building env from scratch ...") ff = FlowField(field_cfg, cuda_cfg, device_id=device_id) if cfg["has_disturbance"]: # Karman layout: dist_cyl first ff.add_cylinder((10.0 * L0, CENTER_Y, 0.0), L0) for y_off in [2.0, 0.0, -2.0]: ff.add_sensor((cfg["sensor_x"] * L0, CENTER_Y + y_off * L0, 0.0), L0 / 4.0) n_phase1 = 4 ff.run(int(4 * 1280 / u0), np.zeros(n_phase1, dtype=DATA_TYPE)) else: # Illusion layout: sensors first for y_off in [2.0, 0.0, -2.0]: ff.add_sensor((cfg["sensor_x"] * L0, CENTER_Y + y_off * L0, 0.0), L0 / 4.0) # Add pinball ff.add_cylinder((cfg["pinball_front_x"] * L0, CENTER_Y, 0.0), L0 / 2.0) ff.add_cylinder((cfg["pinball_rear_x"] * L0, CENTER_Y + 0.75 * L0, 0.0), L0 / 2.0) ff.add_cylinder((cfg["pinball_rear_x"] * L0, CENTER_Y - 0.75 * L0, 0.0), L0 / 2.0) ff.run(int(4 * 1280 / u0), np.zeros(n_obj, dtype=DATA_TYPE)) ff.get_ddf() ff.save_ddf() # Norm obs_slice_start = cfg["obs_slice"][0] obs_slice_end = cfg["obs_slice"][1] fifo = deque(maxlen=FIFO_LEN) for _ in range(FIFO_LEN): ff.run(si, np.zeros(n_obj, dtype=DATA_TYPE)) fifo.append(ff.obs.copy()[obs_slice_start:obs_slice_end]) temp = np.array(fifo, dtype=DATA_TYPE) force_norm_fact = 6.0 * float(np.max(np.abs(temp[:, 6:12]))) sens_deviation = np.mean(temp[:, 0:6], axis=0).astype(DATA_TYPE) sens_norm_fact = np.zeros(6, dtype=DATA_TYPE) for i in range(6): sens_norm_fact[i] = 5.0 * float(np.max(np.abs(temp[:, i] - sens_deviation[i]))) norm = {"force_norm_fact": force_norm_fact, "sens_deviation": sens_deviation.tolist(), "sens_norm_fact": sens_norm_fact.tolist()} # Preset-action FIFO (matches legacy env) ff.apply_ddf() bias_arr = np.zeros(n_obj, dtype=DATA_TYPE) if cfg["has_disturbance"]: bias_arr[4] = -4.0 * u0 bias_arr[5] = 4.0 * u0 else: bias_arr[4] = -1.0 * u0 bias_arr[5] = 1.0 * u0 fifo.clear() for _ in range(FIFO_LEN): ff.run(si, bias_arr) fifo.append(ff.obs.copy()[obs_slice_start:obs_slice_end]) save_states_arr = np.array(fifo, dtype=DATA_TYPE) # Save checkpoint ff.get_ddf() np.save(ddf_ckpt_path, ff.ddf) np.save(fifo_ckpt_path, save_states_arr) with open(os.path.join(out_dir, "norm.json"), "w") as f: json.dump(norm, f, indent=2) print(f" Checkpoint saved to {out_dir}") else: print(f" Loading DDF+FIFO checkpoint from {out_dir}") # Load norm with open(os.path.join(out_dir, "norm.json")) as f: norm = json.load(f) # Rebuild env to get a fresh FlowField ff = FlowField(field_cfg, cuda_cfg, device_id=device_id) if cfg["has_disturbance"]: ff.add_cylinder((10.0 * L0, CENTER_Y, 0.0), L0) for y_off in [2.0, 0.0, -2.0]: ff.add_sensor((cfg["sensor_x"] * L0, CENTER_Y + y_off * L0, 0.0), L0 / 4.0) else: for y_off in [2.0, 0.0, -2.0]: ff.add_sensor((cfg["sensor_x"] * L0, CENTER_Y + y_off * L0, 0.0), L0 / 4.0) ff.add_cylinder((cfg["pinball_front_x"] * L0, CENTER_Y, 0.0), L0 / 2.0) ff.add_cylinder((cfg["pinball_rear_x"] * L0, CENTER_Y + 0.75 * L0, 0.0), L0 / 2.0) ff.add_cylinder((cfg["pinball_rear_x"] * L0, CENTER_Y - 0.75 * L0, 0.0), L0 / 2.0) # Restore DDF ff.ddf = np.load(ddf_ckpt_path) ff.apply_ddf() print(f" DDF checkpoint restored") # Save config with open(os.path.join(out_dir, "config.json"), "w") as f: json.dump({k: str(v) if not isinstance(v, (int, float, list, bool)) else v for k, v in cfg.items()}, f, indent=2) # ---- Target signals (needed for s_dim=14 illusion) ---- target_states = None target_harmonics = None if cfg_sid == "illusion": target_path = os.path.join(out_dir, "target.npz") harm_path = os.path.join(out_dir, "target_harmonics.json") if os.path.isfile(target_path) and os.path.isfile(harm_path): target_data = np.load(target_path) target_states = target_data["target_states"] with open(harm_path) as f: target_harmonics = json.load(f) print(f" Target loaded: {target_states.shape}") else: print(f" WARNING: no target found at {target_path}") # ---- PPO inference ---- obs_slice_start = cfg["obs_slice"][0] obs_slice_end = cfg["obs_slice"][1] # Load checkpoint FIFO state load_state = np.load(fifo_ckpt_path) fifo = deque(maxlen=FIFO_LEN) for s in load_state: fifo.append(s) model_path = model_path_for_scene(scene_name) if model_path is None: raise ValueError(f"No model path for {scene_name}") print(f" Loading model: {model_path} (s_dim={s_dim})") model = load_ppo_model(model_path, device=f"cuda:{device_id}", s_dim=s_dim, a_dim=3) model.set_random_seed(19) obs = np.zeros(s_dim, dtype=np.float32) sens_c, forc_c, act_c, ux_list, uy_list = [], [], [], [], [] for step in range(n_steps): action, _ = model.predict(obs, deterministic=True) action = action.astype(np.float32).flatten() act_c.append(action.copy()) # Build omega array temp = np.zeros(n_obj, dtype=DATA_TYPE) omega = (action * ac_scale + np.array(ac_bias, dtype=np.float32)) * u0 temp[n_obj - 3:] = omega ff.context.push() ff.run(si, temp) ff.context.pop() obs_slice = ff.obs.copy()[obs_slice_start:obs_slice_end] fifo.append(obs_slice) sens_c.append(obs_slice[0:6]) forc_c.append(obs_slice[6:12]) # Build next observation forces_norm = obs_slice[6:12] / norm["force_norm_fact"] sens_norm = (obs_slice[0:6] - norm["sens_deviation"]) / norm["sens_norm_fact"] if s_dim == 14 and target_harmonics is not None: target_vals = gen_target_states_at(step, target_harmonics) t_cd_n = float(target_vals[0]) / norm["force_norm_fact"] t_cl_n = float(target_vals[1]) / norm["force_norm_fact"] obs = np.clip(np.hstack([forces_norm, sens_norm, t_cd_n, t_cl_n]), -1.0, 1.0).astype(np.float32) else: obs = np.clip(np.hstack([forces_norm, sens_norm]), -1.0, 1.0).astype(np.float32) # Save field every step (for Delta-q_ctl POD) ux, uy = get_velocity_field(ff, u0=u0) ux_list.append(ux) uy_list.append(uy) # Save sens_arr = np.array(sens_c, dtype=np.float32) forc_arr = np.array(forc_c, dtype=np.float32) act_arr = np.array(act_c, dtype=np.float32) # Compute similarity conv_len = cfg.get("conv_len", CONV_LEN_DEFAULT) if target_states is not None: if cfg_sid == "karman": sim = compute_similarity(target_states, sens_arr, conv_len) elif cfg_sid == "illusion": # For illusion, target_states[:, 2:8] has the sensor references target_sensors = target_states[:, 2:8] if target_states.shape[1] >= 8 else target_states sim = compute_similarity(target_sensors, sens_arr, conv_len) else: sim = 0.0 print(f" similarity = {sim:.4f}") np.savez(os.path.join(out_dir, "controlled.npz"), sensors=sens_arr, forces=forc_arr, actions=act_arr) np.savez_compressed(os.path.join(out_dir, "fields.npz"), ux=np.stack(ux_list), uy=np.stack(uy_list)) omega_viz = vorticity_from_ddf(ff, u0=u0) save_vorticity_png(os.path.join(out_dir, "vorticity_controlled.png"), omega_viz, title=f"{scene_name} controlled") result = {"scene": scene_name, "n_steps": n_steps, "similarity": float(sim) if target_states is not None else 0.0} with open(os.path.join(out_dir, "result.json"), "w") as f: json.dump(result, f, indent=2) del ff, model print(f" Saved {n_steps} snapshots to {out_dir}") return result def main(): ap = argparse.ArgumentParser() ap.add_argument("--scene", type=str, required=True, help="Scene name: karman_re100, illusion_0.75L, illusion_1.0L, illusion_1.5L") ap.add_argument("--device", type=int, default=3) ap.add_argument("--steps", type=int, default=500) args = ap.parse_args() all_scenes = get_scene_list() if args.scene not in all_scenes: print(f"Unknown scene: {args.scene}. Available PPO scenes: " f"{[s for s in all_scenes if get_scene(s).get('source') == 'PPO_inference']}") return 1 t0 = time.time() r = collect_single(args.scene, args.device, args.steps) print(f"Done in {time.time() - t0:.1f}s: sim={r.get('similarity', 0):.4f}") if __name__ == "__main__": main()