From 61e82ec90aa861f1c9416c834a3238a3ada62a58 Mon Sep 17 00:00:00 2001 From: Frank14f <1515444314@qq.com> Date: Sat, 8 Aug 2026 15:50:49 +0800 Subject: [PATCH] feat(eval): publish cycle-mean wake acquisition Add deterministic phase-filtered V5 and Legacy acquisition with complete-cycle mean fields, then evaluate controlled wakes against target and zero baselines offline. Co-authored-by: Cursor --- src/drl_pinball/acquisition.py | 145 ++- src/drl_pinball/case_registry.py | 89 ++ src/drl_pinball/eval/acquire_v5.py | 560 ++++++++++ src/drl_pinball/eval/infer_train.py | 959 ++++++++++-------- src/drl_pinball/eval/tests/test_wake_l2.py | 87 ++ src/drl_pinball/eval/wake_l2.py | 146 +++ src/drl_pinball/legacy_test/__main__.py | 4 + src/drl_pinball/legacy_test/acquire.py | 819 +++++++++++++++ src/drl_pinball/legacy_test/cases.py | 36 + .../legacy_test/core/legacy_env_builder.py | 302 +++--- .../legacy_test/core/model_loader.py | 98 +- src/drl_pinball/legacy_test/metrics.py | 144 +++ src/drl_pinball/legacy_test/runtime.py | 64 ++ .../legacy_test/tests/test_acquire.py | 202 ++++ .../legacy_test/tests/test_acquire_matrix.py | 217 ++++ .../legacy_test/tests/test_acquire_roles.py | 232 +++++ .../legacy_test/tests/test_contracts.py | 287 ++++++ .../tests/test_nonstandard_acquire.py | 312 ++++++ tests/test_drl_pinball_acquisition.py | 57 ++ tests/test_drl_pinball_infer_train.py | 271 +++++ tests/test_drl_pinball_v5_acquire.py | 565 +++++++++++ 21 files changed, 5022 insertions(+), 574 deletions(-) create mode 100644 src/drl_pinball/case_registry.py create mode 100644 src/drl_pinball/eval/acquire_v5.py create mode 100644 src/drl_pinball/eval/tests/test_wake_l2.py create mode 100644 src/drl_pinball/eval/wake_l2.py create mode 100644 src/drl_pinball/legacy_test/__main__.py create mode 100644 src/drl_pinball/legacy_test/acquire.py create mode 100644 src/drl_pinball/legacy_test/cases.py create mode 100644 src/drl_pinball/legacy_test/metrics.py create mode 100644 src/drl_pinball/legacy_test/runtime.py create mode 100644 src/drl_pinball/legacy_test/tests/test_acquire.py create mode 100644 src/drl_pinball/legacy_test/tests/test_acquire_matrix.py create mode 100644 src/drl_pinball/legacy_test/tests/test_acquire_roles.py create mode 100644 src/drl_pinball/legacy_test/tests/test_contracts.py create mode 100644 src/drl_pinball/legacy_test/tests/test_nonstandard_acquire.py create mode 100644 tests/test_drl_pinball_infer_train.py create mode 100644 tests/test_drl_pinball_v5_acquire.py diff --git a/src/drl_pinball/acquisition.py b/src/drl_pinball/acquisition.py index 5476f4b..9054620 100644 --- a/src/drl_pinball/acquisition.py +++ b/src/drl_pinball/acquisition.py @@ -198,6 +198,127 @@ def assign_phase(times, sensors): return {"phase": phase, "cycle_id": ids, "crossing_times": crossings} +def smooth_center_uy(uy): + """One-pass three-point binomial smoother ``[1, 2, 1]/4`` with copied endpoints.""" + signal = np.asarray(uy, dtype=np.float64) + if signal.ndim != 1 or len(signal) < 1 or not np.all(np.isfinite(signal)): + raise ValueError("center uy must be a finite one-dimensional array") + if len(signal) == 1: + return signal.copy() + if len(signal) == 2: + return signal.copy() + out = np.empty_like(signal) + out[0] = signal[0] + out[-1] = signal[-1] + out[1:-1] = 0.25 * signal[:-2] + 0.5 * signal[1:-1] + 0.25 * signal[2:] + return out + + +def _filter_crossings_by_minimum_interval(crossings, minimum_interval): + crossings = np.asarray(crossings, dtype=np.float64) + if crossings.ndim != 1 or len(crossings) < 1 or not np.all(np.isfinite(crossings)): + raise ValueError("crossing times must be finite and one-dimensional") + if type(minimum_interval) not in (int, float) or not np.isfinite(minimum_interval) or minimum_interval < 0: + raise ValueError("minimum_interval must be a non-negative finite number") + if minimum_interval == 0 or len(crossings) == 1: + return crossings.copy(), np.asarray([], dtype=np.float64) + accepted = [float(crossings[0])] + rejected = [] + for value in crossings[1:]: + if float(value) - accepted[-1] < float(minimum_interval): + rejected.append(float(value)) + else: + accepted.append(float(value)) + return np.asarray(accepted, dtype=np.float64), np.asarray(rejected, dtype=np.float64) + + +def assign_periodic_phase(times, sensors, *, minimum_crossing_interval): + """Standard periodic phase from smoothed center uy with minimum crossing gap. + + Leaves :func:`assign_phase` unchanged for Erase/special callers. Uses one-pass + ``[1,2,1]/4`` smoothing, then interpolated rising zero crossings, then drops + crossings closer than ``minimum_crossing_interval`` to the previous accepted one. + """ + time, sensor = _times(times, "times"), np.asarray(sensors) + if sensor.ndim != 2 or sensor.shape[0] != len(time) or sensor.shape[1] < 4: + raise ValueError("sensors must have shape (time, >=4)") + raw = np.asarray(sensor[:, 3], dtype=np.float64) + if not np.all(np.isfinite(raw)): + raise ValueError("center uy must be finite") + smoothed = smooth_center_uy(raw) + left = np.flatnonzero((smoothed[:-1] <= 0.0) & (smoothed[1:] > 0.0)) + if len(left) < 2: + raise ValueError("at least two rising crossings are required") + weight = -smoothed[left] / (smoothed[left + 1] - smoothed[left]) + raw_crossings = time[left] + weight * (time[left + 1] - time[left]) + crossings, rejected = _filter_crossings_by_minimum_interval( + raw_crossings, minimum_crossing_interval + ) + if len(crossings) < 2: + raise ValueError("at least two accepted rising crossings are required") + cycle = np.searchsorted(crossings, time, side="right") - 1 + valid = (cycle >= 0) & (cycle < len(crossings) - 1) & (time < crossings[-1]) + ids = np.full(len(time), -1, dtype=np.int64) + phase = np.full(len(time), np.nan, dtype=np.float64) + ids[valid] = cycle[valid] + start = crossings[cycle[valid]] + phase[valid] = TAU * (time[valid] - start) / (crossings[cycle[valid] + 1] - start) + return { + "phase": phase, + "cycle_id": ids, + "crossing_times": crossings, + "raw_crossing_times": np.asarray(raw_crossings, dtype=np.float64), + "rejected_crossing_times": rejected, + "smoothed_center_uy": smoothed, + "smoothing_kernel": [0.25, 0.5, 0.25], + "minimum_crossing_interval": float(minimum_crossing_interval), + "accepted_crossing_count": int(len(crossings)), + "rejected_crossing_count": int(len(rejected)), + "complete_cycle_count": int(len(crossings) - 1), + } + + +def complete_cycle_field_indices(field_times, crossing_times): + """Return indices with ``t_first <= field_time < t_last`` for accepted crossings.""" + times = _times(field_times, "field_times") + crossings = _times(crossing_times, "crossing_times") + if len(crossings) < 2: + raise ValueError("at least two accepted rising crossings are required") + first, last = float(crossings[0]), float(crossings[-1]) + indices = np.flatnonzero((times >= first) & (times < last)).astype(np.int64) + if len(indices) < 1: + raise ValueError("no field boundaries fall within complete cycles") + return { + "field_indices": indices, + "first_crossing_time": first, + "last_crossing_time": last, + "complete_cycle_count": int(len(crossings) - 1), + "mean_field_count": int(len(indices)), + } + + +def accumulate_mean_fields(ux, uy, field_indices): + """Arithmetic mean over complete-cycle boundary fields; returns 2-D arrays.""" + ux = np.asarray(ux) + uy = np.asarray(uy) + indices = np.asarray(field_indices, dtype=np.int64) + if ux.ndim != 3 or uy.shape != ux.shape or ux.dtype.kind not in "iuf" or uy.dtype.kind not in "iuf": + raise ValueError("ux/uy must be numeric arrays with shape (n, ny, nx)") + if indices.ndim != 1 or len(indices) < 1: + raise ValueError("field_indices must be a non-empty one-dimensional array") + if np.any(indices < 0) or np.any(indices >= ux.shape[0]): + raise ValueError("field_indices out of range") + selected_ux = np.asarray(ux[indices], dtype=np.float64) + selected_uy = np.asarray(uy[indices], dtype=np.float64) + if not np.all(np.isfinite(selected_ux)) or not np.all(np.isfinite(selected_uy)): + raise ValueError("mean fields require finite selected samples") + return { + "mean_ux": selected_ux.mean(axis=0).astype(np.float32, copy=False), + "mean_uy": selected_uy.mean(axis=0).astype(np.float32, copy=False), + "mean_field_count": int(len(indices)), + } + + def select_phase_fields(field_times, crossing_times, phases=EIGHT_PHASES): """Select one boundary globally for each target phase across complete cycles.""" times = _times(field_times, "field_times") @@ -461,8 +582,28 @@ def write_folded_artifacts(npz_path, csv_path, folded): def publish_selected_fields(destination, arrays): - """Atomic no-replace NPZ publication at a caller-supplied path.""" - data, _ = _arrays(arrays) + """Atomic no-replace NPZ publication at a caller-supplied path. + + Stacked phase/event snapshots must share a common first dimension. Optional + complete-cycle ``mean_ux``/``mean_uy`` are stored as plain 2-D fields whose + shape matches the stacked spatial trailing dimensions. + """ + raw = {name: np.asarray(value) for name, value in arrays.items()} + means = {} + for name in ("mean_ux", "mean_uy"): + if name in raw: + means[name] = raw.pop(name) + data, _ = _arrays(raw) + if means: + spatial = data["ux"].shape[1:] if "ux" in data else data["uy"].shape[1:] + for name, value in means.items(): + if value.ndim != 2 or value.dtype.kind not in "iuf": + raise ValueError(f"{name} must be a numeric two-dimensional array") + if value.shape != spatial: + raise ValueError(f"{name} spatial shape must match stacked fields") + if not np.all(np.isfinite(value)): + raise ValueError(f"{name} must be finite") + data.update(means) target = Path(destination) target.parent.mkdir(parents=True, exist_ok=True) if target.exists(): diff --git a/src/drl_pinball/case_registry.py b/src/drl_pinball/case_registry.py new file mode 100644 index 0000000..71338e7 --- /dev/null +++ b/src/drl_pinball/case_registry.py @@ -0,0 +1,89 @@ +"""Canonical registry for active V5 scratch-training cases. + +Canonical case IDs do not encode the training method. Historical transfer +(``_tr``) names are archive-only and are intentionally absent here. +""" +from __future__ import annotations + +import argparse +from dataclasses import dataclass +from pathlib import Path +from typing import Mapping + +REPO_ROOT = Path(__file__).resolve().parents[2] +TRAIN_DIR = REPO_ROOT / "src" / "drl_pinball" / "train" +CONFIG_DIR = REPO_ROOT / "configs" + + +@dataclass(frozen=True) +class CaseSpec: + case_id: str + config: str + calibration: str + trainer: str + si: int + seeds: tuple[int, ...] + scene_type: str + target_diam: float | None = None + + @property + def config_path(self) -> Path: + return CONFIG_DIR / self.config + + @property + def calibration_path(self) -> Path: + return TRAIN_DIR / "calibrations" / self.calibration / "calibration.json" + + @property + def trainer_path(self) -> Path: + return TRAIN_DIR / self.trainer + + def model_dir(self, seed: int) -> Path: + return TRAIN_DIR / "output" / f"{self.case_id}_seed{seed}" / "models" + + +_CASES = ( + CaseSpec("kar_re100", "config_lbm_karman_2000x600.json", "kar_re100", "train_karman.py", 800, (41, 42, 43, 44, 45), "karman"), + CaseSpec("kar_d075", "config_lbm_karman_2000x600.json", "kar_d075", "train_karman.py", 800, (44,), "karman"), + CaseSpec("kar_d15", "config_lbm_karman_2000x600.json", "kar_d15", "train_karman.py", 800, (45,), "karman"), + CaseSpec("kar_d2", "config_lbm_karman_2000x600.json", "kar_d2", "train_karman.py", 800, (45,), "karman"), + CaseSpec("ill_075L", "config_lbm_karman_2000x600.json", "ill_075L", "train_illusion.py", 1100, (43,), "illusion", 0.75), + CaseSpec("ill_1L", "config_lbm_karman_2000x600.json", "ill_1L", "train_illusion.py", 1200, (43,), "illusion", 1.0), + CaseSpec("ill_15L", "config_lbm_karman_2000x600.json", "ill_15L", "train_illusion.py", 1200, (43,), "illusion", 1.5), + CaseSpec("ill_2L", "config_lbm_karman_2000x600.json", "ill_2L", "train_illusion.py", 1200, (43,), "illusion", 2.0), + CaseSpec("kar_re60", "config_lbm_karman_2000x600_re60.json", "kar_re60", "train_karman.py", 800, (43,), "karman"), + CaseSpec("kar_re200", "config_lbm_karman_2000x600_re200.json", "kar_re200", "train_karman.py", 500, (43,), "karman"), + CaseSpec("kar_re400", "config_lbm_karman_2000x600_re400.json", "kar_re400", "train_karman.py", 400, (43,), "karman"), +) + +CASE_REGISTRY: Mapping[str, CaseSpec] = {case.case_id: case for case in _CASES} +CASE_IDS = tuple(CASE_REGISTRY) + + +def get_case(case_id: str) -> CaseSpec: + try: + return CASE_REGISTRY[case_id] + except KeyError as exc: + raise KeyError(f"Unknown canonical V5 case {case_id!r}; expected one of {CASE_IDS}") from exc + + +def main() -> int: + parser = argparse.ArgumentParser(description="Query the canonical V5 case registry") + parser.add_argument("case_id", choices=CASE_IDS) + parser.add_argument("field", choices=("config", "calibration", "trainer", "si", "seeds", "scene_type")) + args = parser.parse_args() + case = get_case(args.case_id) + values = { + "config": case.config_path, + "calibration": case.calibration_path, + "trainer": case.trainer_path, + "si": case.si, + "seeds": ",".join(map(str, case.seeds)), + "scene_type": case.scene_type, + } + print(values[args.field]) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/drl_pinball/eval/acquire_v5.py b/src/drl_pinball/eval/acquire_v5.py new file mode 100644 index 0000000..971cdb4 --- /dev/null +++ b/src/drl_pinball/eval/acquire_v5.py @@ -0,0 +1,560 @@ +"""V5 controlled-role phase acquisition; CFD imports remain lazy.""" +from __future__ import annotations + +import argparse +import json +import os +import shutil +from pathlib import Path +from typing import Any, Callable + +import numpy as np + +from drl_pinball.acquisition import ( + accumulate_mean_fields, assign_periodic_phase, cleanup_scratch, + complete_cycle_field_indices, create_scratch, default_reproduction_mapping, + dual_cycle_dtw, pooled_phase_bins, prepare_role_output, publish_role_output, + publish_selected_fields, + select_phase_fields, validate_modern_fields, validate_output_storage, + write_boundary_artifacts, write_phase_cycle_artifacts, +) +from drl_pinball.case_registry import CASE_IDS, get_case +from drl_pinball.eval import infer_train + +ROLES = ("controlled", "target", "zero") +WARMUP_CONTROL_STEPS, COLLECT_BOUNDARIES = 750, 250 +# Reviewed from run-local target center-sensor uy crossings; this controls only +# the rolling DTW window. Phase itself always comes from measured crossings. +CYCLE_WINDOWS = { + "kar_re100": 30, "kar_d075": 24, "kar_d15": 39, "kar_d2": 47, + "ill_075L": 18, "ill_1L": 19, "ill_15L": 26, "ill_2L": 31, + "kar_re60": 39, "kar_re200": 40, "kar_re400": 45, +} +CYCLE_WINDOW_REL_TOL = 0.15 +MIN_FREE_BYTES = 4 * 1024**3 + + +def _atomic_json(path: Path, payload: dict[str, Any]) -> None: + temporary = path.with_name(f".{path.name}.tmp") + with temporary.open("w", encoding="utf-8") as stream: + json.dump(payload, stream, indent=2, allow_nan=False) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + + +def _unwrap_step(step_result): + obs, reward, done, info = step_result + item = info[0] if isinstance(info, (list, tuple)) else info + ended = bool(np.asarray(done).reshape(-1)[0]) + return obs, float(np.asarray(reward).reshape(-1)[0]), ended, item + + +def _raw_sample(raw_env, scene_type): + raw = np.asarray(raw_env._read_obs(), dtype=np.float32) + sample = raw[2:14] if scene_type == "karman" else raw[0:12] + if raw.ndim != 1 or sample.shape != (12,) or not np.all(np.isfinite(sample)): + raise ValueError(f"{scene_type} raw observation must expose finite 6-sensor/6-force data") + return sample + + +def _raw_boundary(raw_env, scene_type, action, reward, info, control_index): + sample = _raw_sample(raw_env, scene_type) + target_omega = np.asarray(raw_env._action_to_omega(action), dtype=np.float32).reshape(3) + effective = np.asarray(raw_env.smoother._state, dtype=np.float32).reshape(3).copy() + lattice = int(raw_env.sim.stepper.step_count) + return { + "physical_time": lattice * float(raw_env._cal["U0"]) / int(raw_env._cal["grid"]["nx"]), + "lattice_step": lattice, "control_index": int(control_index), + "sensors": sample[:6].copy(), "forces": sample[6:12].copy(), + "action_normalized": np.asarray(action, dtype=np.float32).reshape(3).copy(), + "commanded_target_omega": target_omega, "effective_smoothed_omega": effective, + "reward_raw": float(reward), "cd": float(info.get("cd", np.nan)), + "cl": float(info.get("cl", np.nan)), "r_cd": float(info.get("r_cd", np.nan)), + "r_cl": float(info.get("r_cl", np.nan)), "r_sim": float(info.get("r_sim", np.nan)), + "floor_pen": float(info.get("floor_pen", np.nan)), + "native_reward_dtw": float(info["sim"]), + } + + +def _capture_fields(raw_env): + holder = {} + raw_env._gpu_block(lambda: holder.update(raw_env.sim.get_macroscopic())) + nx, ny = int(raw_env.sim.lbm_cfg.nx), int(raw_env.sim.lbm_cfg.ny) + return validate_modern_fields( + {name: np.asarray(holder[name], dtype=np.float32) for name in ("rho", "ux", "uy")}, + nx=nx, ny=ny, + ) + + +def _collect_controlled(model, vec_env, raw_env, scratch: Path, + capture_fields: Callable = _capture_fields, *, scene_type="karman"): + obs = vec_env.reset() + for step in range(WARMUP_CONTROL_STEPS): + action, _ = model.predict(obs, deterministic=True) + obs, _, ended, _ = _unwrap_step(vec_env.step(action)) + if ended: + raise RuntimeError(f"controlled acquisition terminated during warmup at step {step + 1}") + rows = [] + field_buffer = None + for boundary in range(COLLECT_BOUNDARIES): + action, _ = model.predict(obs, deterministic=True) + obs, reward, ended, info = _unwrap_step(vec_env.step(action)) + if ended: + raise RuntimeError(f"controlled acquisition terminated at boundary {boundary}") + rows.append(_raw_boundary(raw_env, scene_type, action, reward, info, WARMUP_CONTROL_STEPS + boundary + 1)) + fields = capture_fields(raw_env) + if field_buffer is None: + field_buffer = _allocate_field_buffer(fields["ux"].shape) + _store_field_candidate(field_buffer, boundary, fields) + return rows, field_buffer + + +def _allocate_field_buffer(shape_yx): + ny, nx = map(int, shape_yx) + return { + "ux": np.empty((COLLECT_BOUNDARIES, ny, nx), dtype=np.float32), + "uy": np.empty((COLLECT_BOUNDARIES, ny, nx), dtype=np.float32), + } + + +def _store_field_candidate(field_buffer, boundary: int, fields): + field_buffer["ux"][boundary] = np.asarray(fields["ux"], dtype=np.float32) + field_buffer["uy"][boundary] = np.asarray(fields["uy"], dtype=np.float32) + + +def _physical_zero_action(raw_env): + bias = np.asarray(raw_env._cal["ACTION_BIAS"], dtype=np.float32) + scale = float(raw_env._cal["ACTION_SCALE"]) + if bias.shape != (3,) or not np.all(np.isfinite(bias)) or not np.isfinite(scale) or scale <= 0: + raise ValueError("ACTION_BIAS/ACTION_SCALE must define a finite three-action mapping") + action = (-bias / np.float32(scale)).reshape(1, 3) + if np.any(action < -1.0) or np.any(action > 1.0): + raise ValueError("physical-zero counterbias falls outside normalized action bounds") + omega = np.asarray(raw_env._action_to_omega(action), dtype=np.float64).reshape(3) + if not np.all(np.isfinite(omega)) or not np.allclose(omega, 0.0, rtol=0.0, atol=1e-12): + raise ValueError("normalized counterbias does not produce physical zero omega") + return action + + +def _collect_zero(vec_env, raw_env, scratch: Path, + capture_fields: Callable = _capture_fields, *, scene_type="karman"): + obs = vec_env.reset() + zero = _physical_zero_action(raw_env) + for step in range(WARMUP_CONTROL_STEPS): + obs, _, ended, _ = _unwrap_step(vec_env.step(zero)) + if ended: + raise RuntimeError(f"zero acquisition terminated during warmup at step {step + 1}") + rows = [] + field_buffer = None + for boundary in range(COLLECT_BOUNDARIES): + obs, reward, ended, info = _unwrap_step(vec_env.step(zero)) + if ended: + raise RuntimeError(f"zero acquisition terminated at boundary {boundary}") + rows.append(_raw_boundary(raw_env, scene_type, zero, reward, info, + WARMUP_CONTROL_STEPS + boundary + 1)) + fields = capture_fields(raw_env) + if field_buffer is None: + field_buffer = _allocate_field_buffer(fields["ux"].shape) + _store_field_candidate(field_buffer, boundary, fields) + return rows, field_buffer + + +class _TargetRuntime: + def __init__(self, sim, sensor_ids, calibration): + self.sim, self.sensor_ids, self.calibration = sim, sensor_ids, calibration + + def close(self): + self.sim.close() + + +def _create_target_runtime(case, bundle, device_id, simulation_factory=None, + spinup_runner=None): + if simulation_factory is None: + from CelerisLab import Simulation + simulation_factory = Simulation + cal = bundle["calibration"] + nx, ny = int(cal["grid"]["nx"]), int(cal["grid"]["ny"]) + center_y = (ny - 1) / 2.0 + l0 = float(cal.get("L0", 20.0)) + sim = simulation_factory(lbm_config_path=str(bundle["config_path"]), device_id=device_id) + sim._assert_object_count_contract = lambda *args, **kwargs: None + if case.scene_type == "karman": + body_x, sensor_x = 600.0, 1200.0 + radius = float(cal.get("dist_radius", 1.0)) * l0 + else: + body_x, sensor_x = 400.0, 600.0 + radius = float(case.target_diam) * l0 + if not np.isfinite(radius) or radius <= 0: + raise ValueError("target radius must be finite and positive") + cylinder_id = sim.add_body("circle", center=(body_x, center_y, 0.0), radius=radius) + sensors = [sim.add_body("sensor", center=(sensor_x, center_y + dy, 0.0), radius=5.0) + for dy in (40.0, 0.0, -40.0)] + if cylinder_id != 0 or sensors != [1, 2, 3] or sim.bodies.count != 4: + raise ValueError("target body order must be cylinder id0 then sensor ids1..3") + cylinder = sim.bodies.get(0) + sensor_bodies = [sim.bodies.get(index) for index in sensors] + if bool(getattr(cylinder, "_is_sensor", True)) or any( + not bool(getattr(body, "_is_sensor", False)) for body in sensor_bodies): + raise ValueError("target body types do not match cylinder/sensor order") + sim.initialize() + if sim.bodies.count != 4 or [sim.bodies.get(i).obj_id for i in range(4)] != list(range(4)): + raise ValueError("initialized target body registry changed id/order/count") + spinup_steps = int(4 * nx / float(cal["U0"])) + if spinup_steps <= 0: + raise ValueError("target physical spin-up must be positive") + if spinup_runner is None: + sim.run(spinup_steps, zero_obs=True) + else: + spinup_runner(sim, spinup_steps) + if sim.bodies.count != 4: + raise ValueError("target body count changed during physical spin-up") + return _TargetRuntime(sim, sensors, cal) + + +def _capture_target_fields(runtime): + holder = {} + runtime.sim.ctx._ctx.push() + try: + holder.update(runtime.sim.get_macroscopic()) + finally: + runtime.sim.ctx._ctx.pop() + nx, ny = int(runtime.sim.lbm_cfg.nx), int(runtime.sim.lbm_cfg.ny) + return validate_modern_fields( + {name: np.asarray(holder[name], dtype=np.float32) for name in ("rho", "ux", "uy")}, + nx=nx, ny=ny, + ) + + +def _target_boundary(runtime, control_index): + sensor = [] + for sensor_id in runtime.sensor_ids: + value = runtime.sim.read_sensor(sensor_id, normalize=True) + sensor.extend((float(value[0]), float(value[1]))) + lattice = int(runtime.sim.stepper.step_count) + nan3, nan6 = np.full(3, np.nan, dtype=np.float32), np.full(6, np.nan, dtype=np.float32) + return { + "physical_time": lattice * float(runtime.calibration["U0"]) / int(runtime.calibration["grid"]["nx"]), + "lattice_step": lattice, "control_index": int(control_index), + "sensors": np.asarray(sensor, dtype=np.float32), "forces": nan6, + "action_normalized": nan3, "commanded_target_omega": nan3.copy(), + "effective_smoothed_omega": nan3.copy(), "reward_raw": np.nan, "cd": np.nan, + "cl": np.nan, "r_cd": np.nan, "r_cl": np.nan, "r_sim": np.nan, + "floor_pen": np.nan, "native_reward_dtw": np.nan, + } + + +def _collect_target(runtime, scratch: Path, si: int, + capture_fields: Callable = _capture_target_fields): + for _ in range(WARMUP_CONTROL_STEPS): + runtime.sim.run(si, zero_obs=True, sync_obs=True) + rows = [] + field_buffer = None + for boundary in range(COLLECT_BOUNDARIES): + runtime.sim.run(si, zero_obs=True, sync_obs=True) + rows.append(_target_boundary(runtime, WARMUP_CONTROL_STEPS + boundary + 1)) + fields = capture_fields(runtime) + if field_buffer is None: + field_buffer = _allocate_field_buffer(fields["ux"].shape) + _store_field_candidate(field_buffer, boundary, fields) + return rows, field_buffer + + +def _columns(rows): + return {name: np.asarray([row[name] for row in rows]) for name in rows[0]} + + +def _phase_values(columns): + excluded = {"physical_time", "lattice_step", "control_index", "phase", "cycle_id"} + return {name: value for name, value in columns.items() if name not in excluded} + + +def _finite_mean_or_none(values): + values = np.asarray(values, dtype=np.float64) + finite = values[np.isfinite(values)] + return None if not len(finite) else float(np.mean(finite)) + + +def _minimum_crossing_interval(physical_time, cycle_length): + times = np.asarray(physical_time, dtype=np.float64) + if len(times) < 2: + raise ValueError("physical_time must contain at least two boundaries") + dt = float(np.median(np.diff(times))) + if not np.isfinite(dt) or dt <= 0: + raise ValueError("physical_time spacing must be positive and finite") + return 0.5 * float(cycle_length) * dt + + +PHASE_FIELD_KEYS = { + "ux", "uy", "mean_ux", "mean_uy", "field_indices", "cycle_id", + "target_phase", "actual_phase", "phase_error", +} + + +def _finalize(role_dir: Path, scratch: Path, rows, field_buffer, bundle, storage, + sensor_dtw_factor: float, role: str, unavailable_fields, *, case, seed, + cycle_length): + columns = _columns(rows) + dtw = dual_cycle_dtw( + bundle["target_states"], columns["sensors"] * sensor_dtw_factor, + columns["native_reward_dtw"], cycle_length=cycle_length, lag_channel=3, + ) + columns.update(target_normalized_dtw=dtw["target_normalized_dtw"], + target_normalized_dtw_lag=dtw["target_normalized_dtw_lag"]) + min_gap = _minimum_crossing_interval(columns["physical_time"], cycle_length) + phase = assign_periodic_phase( + columns["physical_time"], columns["sensors"], minimum_crossing_interval=min_gap, + ) + if len(phase["crossing_times"]) < 4: + raise ValueError("at least three complete center-uy cycles are required") + columns.update(phase=phase["phase"], cycle_id=phase["cycle_id"]) + write_boundary_artifacts(role_dir / "timeseries.npz", role_dir / "timeseries.csv", columns) + phase_values = _phase_values(columns) + write_phase_cycle_artifacts( + role_dir / "phase_cycle.npz", role_dir / "phase_cycle.csv", + pooled_phase_bins(phase["phase"], phase_values, bins=32), + ) + selected = select_phase_fields(columns["physical_time"], phase["crossing_times"]) + mean_span = complete_cycle_field_indices(columns["physical_time"], phase["crossing_times"]) + mean = accumulate_mean_fields(field_buffer["ux"], field_buffer["uy"], mean_span["field_indices"]) + published = { + "ux": np.asarray(field_buffer["ux"][selected["field_indices"]], dtype=np.float32), + "uy": np.asarray(field_buffer["uy"][selected["field_indices"]], dtype=np.float32), + "mean_ux": mean["mean_ux"], + "mean_uy": mean["mean_uy"], + } + published.update(selected) + field_path = publish_selected_fields(role_dir / "phase_fields.npz", published) + with np.load(field_path, allow_pickle=False) as saved: + if (saved["ux"].shape[0] != 8 or saved["mean_ux"].shape != saved["ux"].shape[1:] + or saved["mean_uy"].shape != saved["uy"].shape[1:] + or set(saved.files) != PHASE_FIELD_KEYS): + raise ValueError("published phase fields failed exact-key validation") + _atomic_json(role_dir / "dtw_summary.json", { + "native_pipeline": ("unavailable for target role" if role == "target" else + "info[sim] copied exactly at each post-step boundary"), + "target_normalized_dtw": dtw["metadata"]["definition"], + "cycle_length": cycle_length, "lag_channel": 3, + "target_scale": dtw["metadata"]["scale"].tolist(), + "native_mean": _finite_mean_or_none(columns["native_reward_dtw"]), + "target_normalized_dtw_finite_mean": _finite_mean_or_none(columns["target_normalized_dtw"]), + }) + cleanup_scratch(scratch, root=role_dir / "scratch") + scratch.parent.rmdir() + _atomic_json(role_dir / "metadata.json", { + "schema": "drl-pinball-v5-acquisition-v2", "case_id": case.case_id, "role": role, + "seed": (seed if role == "controlled" else None), "warmup_control_steps": WARMUP_CONTROL_STEPS, + "collected_post_step_boundaries": COLLECT_BOUNDARIES, "si": case.si, + "cycle_window": cycle_length, "cycle_window_source": "reviewed explicit CYCLE_WINDOWS mapping", + "units": {"physical_time": "tU0/NX", + "lattice_step": "completed solver lattice steps", + "control_index": "one-based completed policy boundary index", + "sensors": "raw Celeris normalized sensor velocity", "forces": "raw Celeris normalized force", + "action_normalized": "dimensionless policy action [-1,1]", + "commanded_target_omega": "lattice radians per lattice time", + "effective_smoothed_omega": "lattice radians per lattice time", + "reward_raw": "native environment scalar", "dtw": "dimensionless similarity", + "ux": "lattice velocity", "uy": "lattice velocity"}, + "phase_contract": "smoothed rising crossings of sensors[:,3] with min-gap filter; complete half-open cycles", + "phase_smoothing_kernel": phase["smoothing_kernel"], + "minimum_crossing_interval": phase["minimum_crossing_interval"], + "accepted_crossing_count": phase["accepted_crossing_count"], + "rejected_crossing_count": phase["rejected_crossing_count"], + "complete_cycle_count": phase["complete_cycle_count"], + "mean_field_count": mean_span["mean_field_count"], + "mean_first_crossing_time": mean_span["first_crossing_time"], + "mean_last_crossing_time": mean_span["last_crossing_time"], + "dtw_contract": "native info[sim] preserved; separate target-max-abs-normalized rolling six-channel DTW", + "dtw_lag_channel": 3, + "phase_variables": list(phase_values), "sensor_dtw_factor": sensor_dtw_factor, + "sensor_dtw_conversion_contract": "persist raw native sensors; multiply only the in-memory DTW input", + "model": (infer_train._file_identity(bundle["model_path"]) if role == "controlled" else None), + "normalizer": (infer_train._file_identity(bundle["vecnormalize_path"]) if role == "controlled" else None), + "artifact_resolution": infer_train._bundle_metadata(bundle), + "config": infer_train._file_identity(bundle["config_path"]), + "resolved_optane_path": str(storage["resolved_output_root"]), + "scratch_cleanup": {"complete": not scratch.exists(), "path": str(scratch)}, + "role_semantics": ({ + "kind": "physical-zero/uncontrolled", + "normalized_action": "counterbias -ACTION_BIAS/ACTION_SCALE", + "commanded_omega": "validated numerically zero before rollout", + } if role == "zero" else ({ + "kind": "target disturbance-cylinder trajectory", + } if role == "target" else {"kind": "frozen-policy controlled trajectory"})), + "unavailable_fields": list(unavailable_fields), + "deliberately_not_persisted_fields": ["rho (returned by get_macroscopic and used only for in-memory field validation)"], + "candidate_field_storage": "single-role in-memory FP32 ux/uy buffer; no boundary_*.npz scratch", + }) + + +def _validate_staged_role(role_dir: Path) -> None: + expected = {"timeseries.npz", "timeseries.csv", "phase_cycle.npz", "phase_cycle.csv", + "phase_fields.npz", "dtw_summary.json", "metadata.json"} + actual = {path.name for path in role_dir.iterdir()} + if actual != expected: + raise ValueError(f"staged role files mismatch: expected {sorted(expected)}, got {sorted(actual)}") + with np.load(role_dir / "timeseries.npz", allow_pickle=False) as saved: + expected_series = { + "physical_time", "lattice_step", "control_index", "sensors", "forces", + "action_normalized", "commanded_target_omega", "effective_smoothed_omega", + "reward_raw", "cd", "cl", "r_cd", "r_cl", "r_sim", "floor_pen", + "native_reward_dtw", "target_normalized_dtw", + "target_normalized_dtw_lag", "phase", "cycle_id", + } + if set(saved.files) != expected_series or saved["sensors"].shape != (COLLECT_BOUNDARIES, 6): + raise ValueError("timeseries exact schema failed") + with np.load(role_dir / "phase_fields.npz", allow_pickle=False) as saved: + if (set(saved.files) != PHASE_FIELD_KEYS or saved["ux"].shape[0] != 8 + or saved["uy"].shape != saved["ux"].shape + or saved["mean_ux"].shape != saved["ux"].shape[1:] + or saved["mean_uy"].shape != saved["uy"].shape[1:]): + raise ValueError("phase fields exact schema failed") + for name in ("dtw_summary.json", "metadata.json"): + with (role_dir / name).open(encoding="utf-8") as stream: + json.load(stream) + + +def _validate_acquisition_bundle(case, seed, bundle): + if set(CYCLE_WINDOWS) != set(CASE_IDS): + raise ValueError("reviewed cycle-window mapping must exactly cover the case registry") + if case.scene_type not in {"karman", "illusion"}: + raise ValueError(f"unsupported V5 scene type {case.scene_type!r}") + if seed not in case.seeds or int(bundle["seed"]) != seed: + raise ValueError("resolved bundle seed does not match the registered request") + if Path(bundle["config_path"]).resolve() != case.config_path.resolve(): + raise ValueError("resolved config is not the registry config") + calibration = bundle["calibration"] + source_calibration = infer_train._load_json(bundle["calibration_path"]) + if int(calibration.get("SI", -1)) != case.si or int(source_calibration.get("SI", -1)) != case.si: + raise ValueError("run-local and resolved SI must match the case registry") + grid = calibration.get("grid") + if not isinstance(grid, dict) or int(grid.get("nx", 0)) <= 0 or int(grid.get("ny", 0)) <= 0: + raise ValueError("calibration must declare a positive two-dimensional grid") + config = infer_train._load_json(bundle["config_path"]) + config_grid = config.get("grid", {}) + if (int(config_grid.get("nx", 0)), int(config_grid.get("ny", 0))) != (int(grid["nx"]), int(grid["ny"])): + raise ValueError("config and calibration grids do not match") + u0 = float(calibration.get("U0", np.nan)) + if not np.isfinite(u0) or u0 <= 0 or not np.isclose( + u0, float(config.get("physics", {}).get("velocity", np.nan)), rtol=0.0, atol=1e-12): + raise ValueError("config and calibration U0 do not match") + target = np.asarray(bundle["target_states"]) + if target.shape != (150, 6) or target.dtype.kind != "f" or not np.all(np.isfinite(target)): + raise ValueError("resolved target must be finite floating shape (150, 6)") + center_uy = np.asarray(target[:, 3], dtype=np.float64) + left = np.flatnonzero((center_uy[:-1] <= 0.0) & (center_uy[1:] > 0.0)) + if len(left) < 3: + raise ValueError("target center-sensor uy must expose at least two cycle periods") + crossings = left - center_uy[left] / (center_uy[left + 1] - center_uy[left]) + measured = float(np.median(np.diff(crossings))) + reviewed = CYCLE_WINDOWS[case.case_id] + if not np.isclose(measured, reviewed, rtol=CYCLE_WINDOW_REL_TOL, atol=0.0): + raise ValueError( + f"reviewed cycle window {reviewed} disagrees with target center-uy period {measured:.3f}" + ) + if case.scene_type == "illusion" and case.target_diam is None: + raise ValueError("Illusion case must declare target geometry") + return reviewed + + +def _validate_shared_role_identity(case, requested_seed, requested_bundle): + """Prove target/zero physical inputs are seed-invariant before one case output.""" + requested_calibration = dict(requested_bundle["calibration"]) + requested_calibration.pop("config_path", None) + requested_target = np.asarray(requested_bundle["target_states"]) + for seed in case.seeds: + bundle = requested_bundle if seed == requested_seed else infer_train._resolve_seed_artifacts(case, seed) + _validate_acquisition_bundle(case, seed, bundle) + calibration = dict(bundle["calibration"]) + calibration.pop("config_path", None) + if calibration != requested_calibration: + raise ValueError(f"{case.case_id} seed {seed} has different physical calibration") + if not np.allclose( + np.asarray(bundle["target_states"]), requested_target, rtol=0.0, atol=1e-6): + raise ValueError(f"{case.case_id} seed {seed} has a different physical target") + if Path(bundle["config_path"]).resolve() != Path(requested_bundle["config_path"]).resolve(): + raise ValueError(f"{case.case_id} seed {seed} has a different registry config") + + +def acquire_role(role, *, case_id="kar_re100", seed=None, output_root=None, + overwrite=False, device_id=0, + storage_validator=validate_output_storage, runtime_factory=None, + target_runtime_factory=None, finalizer=_finalize): + if role not in ROLES: + raise ValueError(f"unknown V5 acquisition role {role!r}") + case = get_case(case_id) + seed = case.seeds[0] if seed is None else seed + # Artifact, calibration, target, config, and reviewed-window checks precede + # storage mutation and all solver construction for every role. + bundle = infer_train._resolve_seed_artifacts(case, seed) + cycle_length = _validate_acquisition_bundle(case, seed, bundle) + if role != "controlled": + _validate_shared_role_identity(case, seed, bundle) + storage = storage_validator(repo_mapping=default_reproduction_mapping(infer_train.REPO_ROOT), + output_root=output_root, min_free_bytes=MIN_FREE_BYTES) + output_case = f"{case.case_id}_seed{seed}" if role == "controlled" else case.case_id + prepared = prepare_role_output(storage, "v5", output_case, role, overwrite=overwrite) + role_dir, scratch_root = prepared["role_dir"], prepared["scratch_root"] + scratch = create_scratch(scratch_root) + runtime = None + try: + if role == "target": + runtime = (_create_target_runtime(case, bundle, device_id) + if target_runtime_factory is None + else target_runtime_factory(case, bundle, device_id)) + rows, fields = _collect_target(runtime, scratch, case.si) + sensor_factor = float(bundle["calibration"].get("SENSOR_CC", 1.0)) + unavailable = ["forces", "action_normalized", "commanded_target_omega", + "effective_smoothed_omega", "reward_raw", "cd", "cl", "r_cd", + "r_cl", "r_sim", "floor_pen", "native_reward_dtw"] + else: + if runtime_factory is None: + env = infer_train._create_env(case, bundle, device_id) + model, vec_env = infer_train._wrap_and_load(env, bundle) + raw_env = infer_train._get_raw_env(vec_env) + runtime = env + else: + runtime, model, vec_env, raw_env = runtime_factory(case, bundle, device_id) + if role == "controlled": + rows, fields = _collect_controlled( + model, vec_env, raw_env, scratch, scene_type=case.scene_type) + else: + rows, fields = _collect_zero(vec_env, raw_env, scratch, scene_type=case.scene_type) + sensor_factor = float(getattr(raw_env, "_dtw_sensor_factor", 1.0)) + unavailable = [] + finalizer(role_dir, scratch, rows, fields, bundle, storage, sensor_factor, + role, unavailable, case=case, seed=seed, cycle_length=cycle_length) + _validate_staged_role(role_dir) + return publish_role_output(prepared) + except Exception: + if prepared["staging_dir"].exists(): + shutil.rmtree(prepared["staging_dir"]) + raise + finally: + if runtime is not None: + runtime.close() + + +def acquire_controlled(**kwargs): + return acquire_role("controlled", **kwargs) + + +def main() -> int: + parser = argparse.ArgumentParser(description="Acquire canonical V5 phase evidence") + parser.add_argument("--case", choices=CASE_IDS, default="kar_re100") + parser.add_argument("--seed", type=int) + parser.add_argument("--role", choices=ROLES, required=True) + parser.add_argument("--device-id", type=int, default=0) + parser.add_argument("--output-root", type=Path) + parser.add_argument("--overwrite", action="store_true") + args = parser.parse_args() + path = acquire_role(args.role, case_id=args.case, seed=args.seed, + output_root=args.output_root, overwrite=args.overwrite, + device_id=args.device_id) + print(path) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/drl_pinball/eval/infer_train.py b/src/drl_pinball/eval/infer_train.py index 1bc0ea6..4397c26 100644 --- a/src/drl_pinball/eval/infer_train.py +++ b/src/drl_pinball/eval/infer_train.py @@ -1,26 +1,17 @@ #!/usr/bin/env python3 -"""V5 Train model inference engine — skeleton injection with per-seed VecNormalize. +"""Canonical V5 evaluator for retained scratch policies. -For each scene in TRAIN_SCENES: - 1. Create V5 CFD env (KarmanCloakEnv / IllusionCloakEnv) - 2. For every seed: fresh skeleton + VecNormalize.load(seed.pkl) + weight injection - 3. 360-step deterministic rollout, pick best by tail-180 reward - 4. Save signals.npz + vorticity PNGs (controlled, target, zero) - -Uses skeleton injection per-seed (one skeleton per seed, one VecNormalize load per seed) -to avoid cloudpickle/numpy PPO.load() deserialization issues while still using -trained VecNormalize stats for proper observation normalization. - -Usage: - conda run -n pycuda_3_10 python -u infer_train.py --device-id 0 - conda run -n pycuda_3_10 python -u infer_train.py --device-id 0 --scene re100 +The case registry is the only case authority. CFD dependencies are imported lazily +so bundle and reference validation remain CPU-only. """ from __future__ import annotations import argparse +import csv +import hashlib import io import json -import os +import shutil import sys import time import zipfile @@ -28,443 +19,581 @@ from pathlib import Path from typing import Any, Dict, List, Tuple import numpy as np -import pycuda.driver as cuda; cuda.init() -_REPO = str(Path(__file__).resolve().parents[3]) -_SRC = Path(_REPO) / "src" -for p in [_REPO, str(_SRC)]: - if p not in sys.path: - sys.path.insert(0, p) +REPO_ROOT = Path(__file__).resolve().parents[3] +SRC_ROOT = REPO_ROOT / "src" +TRAIN_DIR = SRC_ROOT / "drl_pinball" / "train" +OUT_BASE = Path(__file__).resolve().parent / "output" / "train" +REFERENCE_CSV = TRAIN_DIR / "results/latest/tables/latest_eval_summary.csv" +REFERENCE_JSON = TRAIN_DIR / "results/latest/tables/latest_results_summary.json" +NUM_STEPS = 360 +TAIL_WINDOW = 180 +_POLICY_SENSOR_UNITS = "legacy-policy-v1" +REFERENCE_TABLE_ATOL = 1e-6 +REPRO_REWARD_ATOL = 0.02 +REPRO_COMPONENT_ATOL = 0.02 +REPRO_DTW_ATOL = 0.02 +REPRO_ACTION_MEAN_ATOL = 0.03 -import torch -from torch.nn import Module as TorchModule -from stable_baselines3 import PPO -from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize +for path in (REPO_ROOT, SRC_ROOT, TRAIN_DIR): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) -from CelerisLab.common.render import compute_vorticity, render_vorticity_field - -_TRAIN_DIR = _SRC / "drl_pinball" / "train" -sys.path.insert(0, str(_TRAIN_DIR)) - -from symmetry_wrapper import SymmetryAugmentWrapper - -from drl_pinball.eval.scene_manifest import TRAIN_SCENES - -_OUT_BASE = Path(__file__).resolve().parent / "output" / "train" +from drl_pinball.case_registry import CASE_IDS, CaseSpec, get_case -class Sin(TorchModule): - def __init__(self): super().__init__() - def forward(self, x): return torch.sin(x) +def log(message: str) -> None: + print(f"[{time.strftime('%H:%M:%S')}] {message}", flush=True) -_device = None # set at startup +def _load_json(path: Path) -> Any: + with path.open() as stream: + return json.load(stream) -def log(msg: str) -> None: - print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True) +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() -# --------------------------------------------------------------------------- -# Env creation -# --------------------------------------------------------------------------- -def _create_env(scene: Dict[str, Any], device_id: int): - scene_type = scene["scene_type"] - config_path = Path(scene["config_path"]).resolve() - cal_path = Path(scene["calibration_path"]).resolve() - si = scene["si"] - - with open(cal_path) as f: - cal = json.load(f) - cal["SI"] = si - - cal_dir = cal_path.parent - target_path = cal_dir / "target.npy" - if not target_path.exists(): - raise FileNotFoundError(f"target.npy not found at {target_path}") - target_states = np.load(str(target_path)) - - if scene_type == "karman": - from env_karman import KarmanCloakEnv - log(f" Creating KarmanCloakEnv (config={config_path.name}, SI={si})...") - t0 = time.perf_counter() - env = KarmanCloakEnv( - device_id=device_id, seed=41, - calibration=cal, config_path=str(config_path), - target_states=target_states, - ) - log(f" Env ready in {time.perf_counter() - t0:.0f}s") - return env - elif scene_type == "illusion": - from env_illusion import IllusionCloakEnv - target_harmonics_path = cal_dir / "target_harmonics.json" - if not target_harmonics_path.exists(): - raise FileNotFoundError(f"target_harmonics.json not found at {target_harmonics_path}") - with open(target_harmonics_path) as f: - target_harmonics = json.load(f) - target_diam = float(scene.get("target_diam", 1.0)) - cal["target_diam"] = target_diam - log(f" Creating IllusionCloakEnv (config={config_path.name}, diam={target_diam}L, SI={si})...") - t0 = time.perf_counter() - env = IllusionCloakEnv( - device_id=device_id, seed=41, - calibration=cal, config_path=str(config_path), - target_states=target_states, - target_harmonics=target_harmonics, - target_diam=target_diam, - ) - log(f" Env ready in {time.perf_counter() - t0:.0f}s") - return env - else: - raise ValueError(f"Unknown scene_type: {scene_type}") +def _file_identity(path: Path) -> Dict[str, Any]: + path = path.resolve() + stat = path.stat() + return {"path": str(path), "size_bytes": stat.st_size, "sha256": _sha256(path)} -# --------------------------------------------------------------------------- -# Extract policy weights from SB3 .zip -# --------------------------------------------------------------------------- -def _extract_state_dict(zip_path: str) -> dict: - with zipfile.ZipFile(zip_path, "r") as zf: - with zf.open("policy.pth") as f: - return torch.load(io.BytesIO(f.read()), map_location="cpu", - weights_only=False) - - -# --------------------------------------------------------------------------- -# Raw env extraction (walk wrapper chain) -# --------------------------------------------------------------------------- -def _get_raw_env(vec_env): - dummy = vec_env.venv - inner = dummy.envs[0] - if hasattr(inner, 'prob'): - return inner.env - return inner - - -# --------------------------------------------------------------------------- -# Saving -# --------------------------------------------------------------------------- -def _save_vorticity(env, out_path: Path, scene: Dict[str, Any]) -> None: - sim = env.sim - macro = sim.get_macroscopic() - vort = compute_vorticity(macro["ux"], macro["uy"]) - nx = int(sim.lbm_cfg.nx) - ny = int(sim.lbm_cfg.ny) - cylinders = [] - for body_id in range(sim.bodies.count): - body = sim.bodies.get(body_id) - if hasattr(body, 'center'): - c = body.center - r = body.radius if hasattr(body, 'radius') else 10.0 - cylinders.append(((c[0], c[1]), r)) - render_vorticity_field(vort, nx=nx, ny=ny, out_path=str(out_path), - cylinders=cylinders, vmin=-0.03, vmax=0.03) - - -# --------------------------------------------------------------------------- -# Per-seed eval — correct pattern: PPO.load(no env) + VecNormalize(env, training=False) -# --------------------------------------------------------------------------- -def _eval_one_seed(env, seed_label: str, seed_dir_str: str, - num_steps: int) -> Dict[str, Any]: - """Load model WITHOUT env, wrap env with VecNormalize.load (single layer, frozen).""" - seed_dir = Path(seed_dir_str) - model_path = seed_dir / "best_model.zip" - norm_path = seed_dir.parent / "vec_normalize.pkl" - - if not model_path.exists(): - log(f" [seed {seed_label}] SKIP: no best_model.zip") - return {"seed": seed_label, "avg_reward": -float("inf"), "skip": True} - if not norm_path.exists(): - log(f" [seed {seed_label}] SKIP: no vec_normalize.pkl") - return {"seed": seed_label, "avg_reward": -float("inf"), "skip": True} - - log(f" [seed {seed_label}] loading...") - - # All seeds: PPO.load fails (numpy._core.numeric cloudpickle issue). - # Use skeleton injection directly — avoids GPU pollution from failed load. - wrapped = SymmetryAugmentWrapper(env, prob=0.0, seed=41, rollout_len=2048) - vec_env = DummyVecEnv([lambda: wrapped]) - vec_env = VecNormalize.load(str(norm_path), vec_env) +def _freeze_vecnormalize(vec_env): + """Freeze saved observation moments and disable reward normalization.""" + vec_env.training = False vec_env.norm_reward = False + return vec_env - skeleton = PPO( - "MlpPolicy", - policy_kwargs={"activation_fn": Sin, "net_arch": [64, 64]}, - env=vec_env, device=_device, - n_steps=2048, batch_size=64, n_epochs=10, - learning_rate=3e-4, gamma=0.995, verbose=0, + +def _save_seed_signals(out_dir: Path, result: Dict[str, Any]) -> Path: + path = out_dir / f"seed_{result['seed']}_signals.npz" + np.savez_compressed( + path, sensors=result["sensors"], forces=result["forces"], + actions=result["actions"], rewards=result["rewards"], + r_cd=result["r_cd_series"], r_cl=result["r_cl_series"], + r_sim=result["r_sim_series"], sim_raw=result["sim_raw_series"], ) - sd = _extract_state_dict(str(model_path)) - skeleton.policy.load_state_dict(sd, strict=False) - - raw_env = _get_raw_env(vec_env) - sens_cc = float(raw_env._cal.get("SENSOR_CC", 78.0)) - return _rollout(skeleton, vec_env, raw_env, sens_cc, seed_label, num_steps) + return path -def _rollout(model, vec_env, raw_env, sens_cc, seed_label, num_steps): - """Deterministic rollout, model.predict receives normalized obs from vec_env.""" - ep_rew, ep_r_cd, ep_r_cl, ep_r_sim, ep_sim_raw = [], [], [], [], [] - sig_s, sig_f, sig_a = [], [], [] +def _case_dict(case: CaseSpec, seeds: tuple[int, ...]) -> Dict[str, Any]: + return { + "case_id": case.case_id, + "scene_type": case.scene_type, + "config_path": str(case.config_path), + "si": case.si, + "target_diam": case.target_diam, + "seeds": tuple(seeds), + "num_steps": NUM_STEPS, + } + + + +def _normalizer_lineage(case_id: str, run_dir: Path) -> tuple[Path, str]: + """Select the matrix-proven standalone-eval normalizer, never by fallback.""" + if case_id not in CASE_IDS: + raise ValueError(f"No retained-artifact normalizer lineage for {case_id!r}") + if case_id == "kar_d075": + return ( + run_dir / "best_vecnormalize.pkl", + "kar_d075 retained standalone-eval A/B requires best checkpoint normalizer", + ) + return ( + run_dir / "vec_normalize.pkl", + "canonical retained standalone-eval final compatibility alias", + ) + + +def _calibration_schema(calibration: Dict[str, Any]) -> str: + schema = calibration.get("schema_version") + units = calibration.get("sensor_units") + if schema == "drl-pinball-calibration-v2" and units == "celeris-area-time-average": + return "native-v2" + if schema is None and units is None: + return "legacy-schema-less" + return "ambiguous" + + +def _adapt_native_illusion_products(target: np.ndarray, harmonics: Any, + factor: float) -> tuple[np.ndarray, Any]: + """Adapt native target products in memory to legacy calibration units.""" + target = np.asarray(target).copy() + if target.shape != (150, 6): + raise ValueError(f"Expected Illusion target shape (150, 6), got {target.shape}") + target *= factor + wrapped = isinstance(harmonics, dict) + items = harmonics.get("harmonics") if wrapped else harmonics + if not isinstance(items, list) or len(items) != 8: + raise ValueError("Expected exactly 8 Illusion harmonic channels") + adapted = [] + for index, source in enumerate(items): + item = dict(source) + if index < 6: + item["dc"] = float(item["dc"]) * factor + item["amps"] = [float(value) * factor for value in item["amps"]] + adapted.append(item) + if not wrapped: + return target, adapted + document = dict(harmonics) + document["harmonics"] = adapted + return target, document + +def _resolve_seed_artifacts(case: CaseSpec, seed: int) -> Dict[str, Any]: + """Resolve retained artifacts and explicit matrix-backed lineage exceptions.""" + if seed not in case.seeds: + raise ValueError(f"Seed {seed} is not registered for {case.case_id}: {case.seeds}") + run_dir = case.model_dir(seed).parent.resolve() + normalizer_path, normalizer_rationale = _normalizer_lineage(case.case_id, run_dir) + paths = { + "model_path": run_dir / "models/best_model.zip", + "vecnormalize_path": normalizer_path, + "calibration_path": run_dir / "calibration.json", + "target_path": run_dir / "target.npy", + } + missing = [path for path in paths.values() if not path.exists()] + if missing: + raise FileNotFoundError( + f"Incomplete run-local bundle for {case.case_id} seed {seed}; missing: " + + ", ".join(str(path) for path in missing) + ) + calibration = dict(_load_json(paths["calibration_path"])) + run_schema = _calibration_schema(calibration) + target_states = np.load(paths["target_path"], allow_pickle=False) + target_harmonics = None + adaptation = { + "applied": False, "factor": 1.0, "run_calibration_schema": run_schema, + "registry_calibration_schema": None, "source_target_sha256": _sha256(paths["target_path"]), + "registry_target_sha256": None, + } + if case.scene_type == "illusion": + calibration_dir = case.calibration_path.parent.resolve() + paths["registry_calibration_path"] = case.calibration_path.resolve() + paths["calibration_target_path"] = calibration_dir / "target.npy" + paths["harmonics_path"] = calibration_dir / "target_harmonics.json" + missing = [path for key, path in paths.items() if key in ( + "registry_calibration_path", "calibration_target_path", "harmonics_path" + ) and not path.exists()] + if missing: + raise FileNotFoundError( + f"Incomplete registry calibration bundle for {case.case_id}; missing: " + + ", ".join(str(path) for path in missing) + ) + registry_calibration = dict(_load_json(paths["registry_calibration_path"])) + registry_schema = _calibration_schema(registry_calibration) + run_target_hash = adaptation["source_target_sha256"] + registry_target_hash = _sha256(paths["calibration_target_path"]) + if run_target_hash != registry_target_hash: + raise ValueError( + f"Illusion target mismatch for {case.case_id} seed {seed}: run-local " + f"target.npy SHA256 {run_target_hash} != registry calibration target.npy " + f"SHA256 {registry_target_hash}" + ) + if run_schema != "legacy-schema-less" or registry_schema != "native-v2": + raise ValueError( + f"Ambiguous Illusion calibration schemas for {case.case_id}: " + f"run={run_schema}, registry={registry_schema}; expected legacy-schema-less/native-v2" + ) + factor = float(calibration.get("SENSOR_CC", 78.0)) + target_states, target_harmonics = _adapt_native_illusion_products( + target_states, _load_json(paths["harmonics_path"]), factor + ) + adaptation.update( + applied=True, factor=factor, registry_calibration_schema=registry_schema, + registry_target_sha256=registry_target_hash, + harmonics_source_sha256=_sha256(paths["harmonics_path"]), + source_units="registry native-v2", destination_units="run-local legacy calibration", + ) + elif run_schema == "ambiguous": + raise ValueError(f"Ambiguous run-local calibration schema for {case.case_id}") + calibration["SI"] = case.si + return { + "seed": str(seed), "run_dir": run_dir, "calibration": calibration, + "target_states": target_states, "target_harmonics": target_harmonics, + "policy_sensor_units": _POLICY_SENSOR_UNITS, + "normalizer_lineage_rationale": normalizer_rationale, + "illusion_native_to_legacy_adaptation": adaptation, + "config_path": case.config_path.resolve(), **paths, + } + + +def _bundle_metadata(bundle: Dict[str, Any]) -> Dict[str, Any]: + artifact_keys = ( + "model_path", "vecnormalize_path", "calibration_path", "target_path", + "registry_calibration_path", "calibration_target_path", "harmonics_path", + ) + artifacts = [_file_identity(bundle[key]) for key in artifact_keys if bundle.get(key)] + return { + "run_dir": str(bundle["run_dir"]), + "policy_sensor_units": bundle["policy_sensor_units"], + "normalizer_lineage_rationale": bundle["normalizer_lineage_rationale"], + "illusion_native_to_legacy_adaptation": bundle["illusion_native_to_legacy_adaptation"], + "artifacts": artifacts, + } + + +def _prepare_output(case_id: str, overwrite: bool, output_root: Path | None = None) -> Path: + out_dir = (OUT_BASE if output_root is None else output_root) / case_id + if out_dir.exists() and any(out_dir.iterdir()): + if not overwrite: + raise FileExistsError( + f"Baseline output already exists at {out_dir}; pass --overwrite explicitly to replace it" + ) + shutil.rmtree(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + return out_dir + + +def _extract_state_dict(zip_path: Path): + import torch + with zipfile.ZipFile(zip_path) as archive: + with archive.open("policy.pth") as stream: + return torch.load(io.BytesIO(stream.read()), map_location="cpu", weights_only=False) + + +def _runtime_imports(): + import pycuda.driver as cuda + cuda.init() + import torch + from stable_baselines3 import PPO + from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize + from symmetry_wrapper import SymmetryAugmentWrapper + return torch, PPO, DummyVecEnv, VecNormalize, SymmetryAugmentWrapper + + +def _create_env(case: CaseSpec, bundle: Dict[str, Any], device_id: int): + cal = dict(bundle["calibration"]) + common = dict( + device_id=device_id, seed=int(bundle["seed"]), calibration=cal, + config_path=str(bundle["config_path"]), target_states=bundle["target_states"], + policy_sensor_units=bundle["policy_sensor_units"], + ) + if case.scene_type == "karman": + from env_karman import KarmanCloakEnv + return KarmanCloakEnv(**common) + from env_illusion import IllusionCloakEnv + return IllusionCloakEnv( + **common, target_harmonics=bundle["target_harmonics"], + target_diam=float(case.target_diam), + ) + + +def _wrap_and_load(env, bundle: Dict[str, Any]): + torch, PPO, DummyVecEnv, VecNormalize, SymmetryAugmentWrapper = _runtime_imports() + wrapped = SymmetryAugmentWrapper(env, prob=0.0, seed=41, rollout_len=2048) + vec_env = _freeze_vecnormalize(VecNormalize.load( + str(bundle["vecnormalize_path"]), DummyVecEnv([lambda: wrapped]) + )) + class Sin(torch.nn.Module): + def forward(self, value): + return torch.sin(value) + model = PPO( + "MlpPolicy", policy_kwargs={"activation_fn": Sin, "net_arch": [64, 64]}, + env=vec_env, device=torch.device("cpu"), n_steps=2048, batch_size=64, + n_epochs=10, learning_rate=3e-4, gamma=0.995, verbose=0, + ) + try: + model.policy.load_state_dict(_extract_state_dict(bundle["model_path"]), strict=False) + except Exception: + model = PPO.load(str(bundle["model_path"]), env=vec_env, + device=torch.device("cpu"), custom_objects={"activation_fn": Sin}) + return model, vec_env + + +def _get_raw_env(vec_env): + inner = vec_env.venv.envs[0] + return inner.env if hasattr(inner, "prob") else inner + + +def _rollout(model, vec_env, raw_env, seed: str) -> Dict[str, Any]: + series = {name: [] for name in ("rewards", "r_cd", "r_cl", "r_sim", "sim_raw")} + sensors, forces, actions = [], [], [] obs = vec_env.reset() - - for step in range(num_steps): + sensor_factor = float(getattr(raw_env, "_dtw_sensor_factor", raw_env._cal.get("SENSOR_CC", 78.0))) + for _ in range(NUM_STEPS): action, _ = model.predict(obs, deterministic=True) obs, reward, done, info = vec_env.step(action) - - inf = info[0] if isinstance(info, list) else info - ep_rew.append(float(reward[0])) - ep_r_cd.append(float(inf.get("r_cd", 0))) - ep_r_cl.append(float(inf.get("r_cl", 0))) - ep_r_sim.append(float(inf.get("r_sim", 0))) - ep_sim_raw.append(float(inf.get("sim", 0))) - - try: - raw = raw_env._read_obs() - if "Karman" in type(raw_env).__name__: - sl = raw[2:14] - else: - sl = raw[:] - sig_s.append(sl[0:6] * sens_cc) - sig_f.append(sl[6:12]) - except Exception: - sig_s.append(np.zeros(6, dtype=np.float32)) - sig_f.append(np.zeros(6, dtype=np.float32)) - sig_a.append(np.asarray(action, dtype=np.float32).flatten()) + item = info[0] if isinstance(info, list) else info + series["rewards"].append(float(reward[0])) + for key in ("r_cd", "r_cl", "r_sim"): + series[key].append(float(item.get(key, 0.0))) + series["sim_raw"].append(float(item.get("sim", 0.0))) + raw = raw_env._read_obs() + sample = raw[2:14] if "Karman" in type(raw_env).__name__ else raw + sensors.append(sample[:6] * sensor_factor) + forces.append(sample[6:12]) + actions.append(np.asarray(action, dtype=np.float32).ravel()) if done[0]: - break - - tail = 180 - sl = slice(-tail, None) if len(ep_rew) >= tail else slice(None) - avg_r = float(np.mean(ep_rew[sl])) - - log(f" reward={avg_r:.4f} r_cd={float(np.mean(ep_r_cd[sl])):.3f} " - f"r_cl={float(np.mean(ep_r_cl[sl])):.3f} " - f"r_sim={float(np.mean(ep_r_sim[sl])):.3f} " - f"sim_raw={float(np.mean(ep_sim_raw[sl])):.3f}") - - return { - "seed": seed_label, - "avg_reward": avg_r, - "r_cd": float(np.mean(ep_r_cd[sl])), "r_cl": float(np.mean(ep_r_cl[sl])), - "r_sim": float(np.mean(ep_r_sim[sl])), - "sim_raw": float(np.mean(ep_sim_raw[sl])), - "sensors": np.array(sig_s, dtype=np.float32), - "forces": np.array(sig_f, dtype=np.float32), - "actions": np.array(sig_a, dtype=np.float32), - "rewards": np.array(ep_rew, dtype=np.float32), - "skip": False, + raise RuntimeError(f"Rollout terminated before the required {NUM_STEPS} steps") + tail = slice(-TAIL_WINDOW, None) + result = { + "seed": seed, "skip": False, + "sensors": np.asarray(sensors, dtype=np.float32), + "forces": np.asarray(forces, dtype=np.float32), + "actions": np.asarray(actions, dtype=np.float32), } - - -# --------------------------------------------------------------------------- -# Per-scene pipeline -# --------------------------------------------------------------------------- -def evaluate_train_scene(scene: Dict[str, Any], device_id: int) -> None: - scene_id = scene["scene_id"] - num_steps = scene["num_steps"] - out_dir = _OUT_BASE / scene_id - out_dir.mkdir(parents=True, exist_ok=True) - - log(f"\n{'='*70}") - log(f" Train eval: {scene_id} (SI={scene['si']}, steps={num_steps})") - log(f"{'='*70}") - - # Create CFD env once, reused across seeds - env = _create_env(scene, device_id) - - # Evaluate each seed (independent skeleton + VecNormalize per seed) - all_results = [] - for seed_label, seed_dir in scene["seeds"]: - t0 = time.perf_counter() - result = _eval_one_seed(env, seed_label, seed_dir, num_steps) - result["_dt"] = time.perf_counter() - t0 - all_results.append(result) - - # Pick best seed - valid = [r for r in all_results if not r.get("skip", False)] - if not valid: - log(" WARNING: No valid seeds found, skipping scene.") - env.close() - return - best = max(valid, key=lambda r: r["avg_reward"]) - log(f" Best: seed={best['seed']}, reward={best['avg_reward']:.4f}") - - # Save all-seeds summary - summary = [{"seed": r["seed"], "reward": r["avg_reward"], - "r_cd": r.get("r_cd", 0), "r_cl": r.get("r_cl", 0), - "r_sim": r.get("r_sim", 0), "sim_raw": r.get("sim_raw", 0), - "dt_sec": r.get("_dt", 0)} for r in all_results] - with open(out_dir / "all_seeds.json", "w") as f: - json.dump(summary, f, indent=2) - - # Save best seed's full signals - np.savez_compressed(out_dir / "signals.npz", - sensors=best["sensors"], forces=best["forces"], - actions=best["actions"], rewards=best["rewards"]) - - # Vorticity: re-create env, load best seed skeleton, run, capture - log(" Generating vorticity for best seed...") - env2 = _create_env(scene, device_id) - seed_dir = Path(next(s[1] for s in scene["seeds"] if s[0] == best["seed"])) - norm_path = seed_dir.parent / "vec_normalize.pkl" - model_path = seed_dir / "best_model.zip" - - wrapped2 = SymmetryAugmentWrapper(env2, prob=0.0, seed=41, rollout_len=2048) - vec_env2 = DummyVecEnv([lambda: wrapped2]) - vec_env2 = VecNormalize.load(str(norm_path), vec_env2) - vec_env2.norm_reward = False - skeleton2 = PPO( - "MlpPolicy", - policy_kwargs={"activation_fn": Sin, "net_arch": [64, 64]}, - env=vec_env2, device=_device, - n_steps=2048, batch_size=64, n_epochs=10, - learning_rate=3e-4, gamma=0.995, verbose=0, + for key, values in series.items(): + result[key if key == "rewards" else f"{key}_series"] = np.asarray(values, dtype=np.float32) + result.update( + avg_reward=float(np.mean(result["rewards"][tail])), + r_cd=float(np.mean(result["r_cd_series"][tail])), + r_cl=float(np.mean(result["r_cl_series"][tail])), + r_sim=float(np.mean(result["r_sim_series"][tail])), + sim_raw=float(np.mean(result["sim_raw_series"][tail])), + tail_window=TAIL_WINDOW, ) - sd = _extract_state_dict(str(model_path)) - skeleton2.policy.load_state_dict(sd, strict=False) - - obs = vec_env2.reset() - for _ in range(num_steps): - action, _ = skeleton2.predict(obs, deterministic=True) - obs, reward, done, info = vec_env2.step(action) - if done[0]: - break - _save_vorticity(env2, out_dir / "vorticity_controlled.png", scene) - env2.close() - - # Target vorticity - log(" Generating target vorticity...") - _generate_target_vorticity(scene, device_id, out_dir) - - # Zero-action baseline - log(" Generating zero-action baseline...") - env3 = _create_env(scene, device_id) - zero_omega = env3._action_to_omega(np.zeros(3, dtype=np.float32)) - for _ in range(num_steps): - smoothed = env3.smoother(zero_omega) - env3._set_omega(smoothed) - env3._gpu_block(lambda: env3.sim.run(scene["si"], zero_obs=True)) - _save_vorticity(env3, out_dir / "vorticity_zero.png", scene) - env3.close() - - # Metrics - metrics = { - "best_seed": best["seed"], - "dtw_sim_v5": float(best.get("sim_raw", 0)), - "reward_mean": float(best["avg_reward"]), - "r_cd_mean": float(best.get("r_cd", 0)), - "r_cl_mean": float(best.get("r_cl", 0)), - "r_sim_mean": float(best.get("r_sim", 0)), - "sim_raw_mean": float(best.get("sim_raw", 0)), - "aF_mean": float(np.mean(best["actions"][:, 0])), - "aB_mean": float(np.mean(best["actions"][:, 1])), - "aT_mean": float(np.mean(best["actions"][:, 2])), - } - with open(out_dir / "metrics.json", "w") as f: - json.dump(metrics, f, indent=2) - - env.close() - log(f" {scene_id} complete.") + return result -def _generate_target_vorticity(scene: Dict[str, Any], device_id: int, - out_dir: Path) -> None: +def _save_vorticity(env, path: Path) -> None: + from CelerisLab.common.render import compute_vorticity, render_vorticity_field + macro = env.sim.get_macroscopic() + vorticity = compute_vorticity(macro["ux"], macro["uy"]) + cylinders = [] + for index in range(env.sim.bodies.count): + body = env.sim.bodies.get(index) + if hasattr(body, "center"): + cylinders.append(((body.center[0], body.center[1]), getattr(body, "radius", 10.0))) + render_vorticity_field( + vorticity, nx=int(env.sim.lbm_cfg.nx), ny=int(env.sim.lbm_cfg.ny), + out_path=str(path), cylinders=cylinders, vmin=-0.001, vmax=0.001, + ) + + +def _generate_target_vorticity(case: CaseSpec, bundle: Dict[str, Any], device_id: int, path: Path) -> None: from CelerisLab import Simulation - - scene_type = scene["scene_type"] - config_path = Path(scene["config_path"]).resolve() - si = scene["si"] - warmup = int(4.0 * 2000 / 0.01) - L0 = 20.0 - U0 = 0.01 - - sim = Simulation(lbm_config_path=str(config_path), device_id=device_id) - - if scene_type == "karman": - with open(scene["calibration_path"]) as f: - cal = json.load(f) - dist_radius = float(cal.get("dist_radius", 1.0)) * L0 - nx_cfg = int(cal.get("grid", {}).get("nx", 2000)) - ny_cfg = int(cal.get("grid", {}).get("ny", 600)) - CENTER_Y = float(ny_cfg - 1) / 2.0 - DIST_X = 600.0 - SENS_X = 1200.0 - - sim.add_body("circle", center=(DIST_X, CENTER_Y, 0.0), radius=dist_radius) - sim.add_body("sensor", center=(SENS_X, CENTER_Y + 40.0, 0.0), radius=5.0) - sim.add_body("sensor", center=(SENS_X, CENTER_Y, 0.0), radius=5.0) - sim.add_body("sensor", center=(SENS_X, CENTER_Y - 40.0, 0.0), radius=5.0) - elif scene_type == "illusion": - with open(scene["calibration_path"]) as f: - cal = json.load(f) - target_diam = float(scene.get("target_diam", 1.0)) - nx_cfg = int(cal.get("grid", {}).get("nx", 2000)) - ny_cfg = int(cal.get("grid", {}).get("ny", 600)) - CENTER_Y = float(ny_cfg - 1) / 2.0 - TARGET_X = 400.0 - SENS_X = 600.0 - - sim.add_body("circle", center=(TARGET_X, CENTER_Y, 0.0), radius=target_diam * L0) - sim.add_body("sensor", center=(SENS_X, CENTER_Y + 40.0, 0.0), radius=5.0) - sim.add_body("sensor", center=(SENS_X, CENTER_Y, 0.0), radius=5.0) - sim.add_body("sensor", center=(SENS_X, CENTER_Y - 40.0, 0.0), radius=5.0) + from CelerisLab.common.render import compute_vorticity, render_vorticity_field + cal = bundle["calibration"] + nx, ny = int(cal["grid"]["nx"]), int(cal["grid"]["ny"]) + sim = Simulation(lbm_config_path=str(bundle["config_path"]), device_id=device_id) + center_y = (ny - 1) / 2.0 + if case.scene_type == "karman": + sim.add_body("circle", center=(600.0, center_y, 0.0), radius=float(cal.get("dist_radius", 1.0)) * 20.0) + sensor_x = 1200.0 else: + sim.add_body("circle", center=(400.0, center_y, 0.0), radius=float(case.target_diam) * 20.0) + sensor_x = 600.0 + for dy in (40.0, 0.0, -40.0): + sim.add_body("sensor", center=(sensor_x, center_y + dy, 0.0), radius=5.0) + try: + sim.initialize() + sim.run(int(4.0 * 2000 / 0.01), zero_obs=True) + for _ in range(30): + sim.run(case.si, zero_obs=True) + macro = sim.get_macroscopic() + vort = compute_vorticity(macro["ux"], macro["uy"]) + cylinders = [((sim.bodies.get(i).center[0], sim.bodies.get(i).center[1]), + getattr(sim.bodies.get(i), "radius", 10.0)) + for i in range(sim.bodies.count) if hasattr(sim.bodies.get(i), "center")] + render_vorticity_field(vort, nx=nx, ny=ny, out_path=str(path), + cylinders=cylinders, vmin=-0.001, vmax=0.001) + finally: sim.close() + + +def _compact_seed_summary(item: Dict[str, Any], metrics_only: bool) -> Dict[str, Any]: + return { + "seed": item["seed"], "reward": item["avg_reward"], "tail_window": TAIL_WINDOW, + "steps_recorded": NUM_STEPS, "r_cd": item["r_cd"], "r_cl": item["r_cl"], + "r_sim": item["r_sim"], "sim_raw": item["sim_raw"], + "signals_path": None if metrics_only else f"seed_{item['seed']}_signals.npz", + "model_path": item["model_path"], "model_sha256": item["model_sha256"], + "vecnormalize_path": item["vecnormalize_path"], + "vecnormalize_training": item["vecnormalize_training"], "norm_reward": item["norm_reward"], + "artifact_resolution": item["artifact_resolution"], + } + + +def _build_metrics(case_id: str, best: Dict[str, Any]) -> Dict[str, Any]: + action_tail = best["actions"][-TAIL_WINDOW:] + return { + "case_id": case_id, "best_seed": best["seed"], "eval_steps": NUM_STEPS, + "tail_window": TAIL_WINDOW, "reward_mean": best["avg_reward"], + "r_cd_mean": best["r_cd"], "r_cl_mean": best["r_cl"], + "r_sim_mean": best["r_sim"], "sim_raw_mean": best["sim_raw"], + "dtw_sim_v5": best["sim_raw"], + "aF_mean": float(np.mean(action_tail[:, 0])), + "aT_mean": float(np.mean(action_tail[:, 1])), + "aB_mean": float(np.mean(action_tail[:, 2])), + "action_mean_window": TAIL_WINDOW, + "policy_sensor_units": _POLICY_SENSOR_UNITS, + } + + +def evaluate_case(case: CaseSpec, seeds: tuple[int, ...], device_id: int, + overwrite: bool, metrics_only: bool = False, + output_root: Path = OUT_BASE) -> None: + # Validate every strict bundle before an explicit overwrite can remove output. + bundles = {str(seed): _resolve_seed_artifacts(case, seed) for seed in seeds} + out_dir = _prepare_output(case.case_id, overwrite, output_root) + results = [] + for seed in seeds: + bundle = bundles[str(seed)] + env = _create_env(case, bundle, device_id) + try: + model, vec_env = _wrap_and_load(env, bundle) + raw_env = _get_raw_env(vec_env) + result = _rollout(model, vec_env, raw_env, str(seed)) + result.update( + model_path=str(bundle["model_path"]), model_sha256=_sha256(bundle["model_path"]), + vecnormalize_path=str(bundle["vecnormalize_path"]), + vecnormalize_training=vec_env.training, norm_reward=vec_env.norm_reward, + artifact_resolution=_bundle_metadata(bundle), + ) + if not metrics_only: + _save_seed_signals(out_dir, result) + results.append(result) + finally: + env.close() + best = max(results, key=lambda item: item["avg_reward"]) + summary = [_compact_seed_summary(item, metrics_only) for item in results] + metrics = _build_metrics(case.case_id, best) + (out_dir / "all_seeds.json").write_text(json.dumps(summary, indent=2) + "\n") + (out_dir / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n") + if metrics_only: + print(json.dumps({"metrics": metrics, "all_seeds": summary}, indent=2)) return - sim.initialize() - sim.run(warmup, zero_obs=True) - for _ in range(30): - sim.run(si, zero_obs=True) + np.savez_compressed(out_dir / "signals.npz", sensors=best["sensors"], forces=best["forces"], + actions=best["actions"], rewards=best["rewards"]) + best_bundle = bundles[best["seed"]] + env = _create_env(case, best_bundle, device_id) + try: + model, vec_env = _wrap_and_load(env, best_bundle) + raw_env = _get_raw_env(vec_env) + obs = vec_env.reset() + for _ in range(NUM_STEPS): + action, _ = model.predict(obs, deterministic=True) + obs, _, done, _ = vec_env.step(action) + if done[0]: + raise RuntimeError("Controlled vorticity replay terminated early") + _save_vorticity(raw_env, out_dir / "vorticity_controlled.png") + finally: + env.close() + _generate_target_vorticity(case, best_bundle, device_id, out_dir / "vorticity_target.png") + env = _create_env(case, best_bundle, device_id) + try: + zero = env._action_to_omega(np.zeros(3, dtype=np.float32)) + for _ in range(NUM_STEPS): + env._set_omega(env.smoother(zero)) + env._gpu_block(lambda: env.sim.run(case.si, zero_obs=True)) + _save_vorticity(env, out_dir / "vorticity_zero.png") + finally: + env.close() - macro = sim.get_macroscopic() - vort = compute_vorticity(macro["ux"], macro["uy"]) - cylinders = [ - ((ctr[0], ctr[1]), rad) - for body in [sim.bodies.get(i) for i in range(sim.bodies.count)] - if hasattr(body, 'center') - for ctr, rad in [((body.center[0], body.center[1]), - body.radius if hasattr(body, 'radius') else 10.0)] - ] - render_vorticity_field(vort, nx=nx_cfg, ny=ny_cfg, - out_path=str(out_dir / "vorticity_target.png"), - cylinders=cylinders, vmin=-0.03, vmax=0.03) - sim.close() + +def validate_outputs(case_ids: tuple[str, ...], output_root: Path | None = None) -> tuple[list[str], list[str]]: + """Check strict reference consistency and tolerant fresh-run reproduction.""" + output_root = OUT_BASE if output_root is None else output_root + with REFERENCE_CSV.open(newline="") as stream: + csv_rows = {row["case_id"]: row for row in csv.DictReader(stream)} + payload = _load_json(REFERENCE_JSON) + json_rows = {row["case_id"]: row for row in payload["evaluation"]} + if payload["scope"]["eval_steps"] != NUM_STEPS or payload["scope"]["action_tail_steps"] != TAIL_WINDOW: + raise ValueError("Reference summary does not use the canonical 360/tail-180 contract") + + reports, failures = [], [] + checks = { + "reward_mean": ("eval_reward_mean", REPRO_REWARD_ATOL), + "r_cd_mean": ("eval_r_cd_mean", REPRO_COMPONENT_ATOL), + "r_cl_mean": ("eval_r_cl_mean", REPRO_COMPONENT_ATOL), + "r_sim_mean": ("eval_r_sim_mean", REPRO_COMPONENT_ATOL), + "dtw_sim_v5": ("dtw_similarity", REPRO_DTW_ATOL), + "aF_mean": ("action_front_mean", REPRO_ACTION_MEAN_ATOL), + "aT_mean": ("action_top_mean", REPRO_ACTION_MEAN_ATOL), + "aB_mean": ("action_bottom_mean", REPRO_ACTION_MEAN_ATOL), + } + for case_id in case_ids: + path = output_root / case_id / "metrics.json" + if not path.exists(): + failures.append(f"{case_id}: missing {path}") + continue + metrics = _load_json(path) + if case_id not in csv_rows or case_id not in json_rows: + failures.append(f"{case_id}: absent from latest reference summaries") + continue + csv_row, json_row = csv_rows[case_id], json_rows[case_id] + + table_fields = ("selected_seed", *(reference_key for reference_key, _ in checks.values())) + tables_consistent = True + for field in table_fields: + if field == "selected_seed": + consistent = int(csv_row[field]) == int(json_row[field]) + else: + consistent = np.isclose( + float(csv_row[field]), float(json_row[field]), + rtol=0.0, atol=REFERENCE_TABLE_ATOL, + ) + if not consistent: + tables_consistent = False + failures.append( + f"{case_id}: CSV/JSON disagree for {field} " + f"(strict atol={REFERENCE_TABLE_ATOL})" + ) + if not tables_consistent: + continue + + selected_seed = int(json_row["selected_seed"]) + try: + actual_seed = int(metrics["best_seed"]) + except (KeyError, TypeError, ValueError): + failures.append(f"{case_id}: metrics best_seed is missing or invalid") + continue + if actual_seed != selected_seed: + failures.append( + f"{case_id}: metrics best_seed={actual_seed} does not match " + f"reference selected_seed={selected_seed}" + ) + + for metric_key, (reference_key, tolerance) in checks.items(): + actual = float(metrics[metric_key]) + expected = float(json_row[reference_key]) + delta = actual - expected + passed = abs(delta) <= tolerance + reports.append( + f"{case_id}: {metric_key} actual={actual:.9f} " + f"reference={expected:.9f} delta={delta:+.9f} " + f"atol={tolerance:.3f} {'PASS' if passed else 'FAIL'}" + ) + if not passed: + failures.append( + f"{case_id}: {metric_key} delta={delta:+.9f} exceeds " + f"reproduction atol={tolerance:.3f}" + ) + return reports, failures -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- def main() -> int: - global _device - parser = argparse.ArgumentParser(description="Evaluate V5 train models") + parser = argparse.ArgumentParser(description="Evaluate one canonical V5 case/seed by default") + parser.add_argument("--case", choices=CASE_IDS, default=CASE_IDS[0]) + parser.add_argument("--seed", type=int) + parser.add_argument("--all-seeds", action="store_true") parser.add_argument("--device-id", type=int, default=0) - parser.add_argument("--scene", type=str, default=None) + parser.add_argument("--overwrite", action="store_true") + parser.add_argument("--metrics-only", action="store_true") + parser.add_argument("--output-root", type=Path, default=OUT_BASE) + parser.add_argument("--validate", action="store_true", help="CPU-only validation of retained outputs") + parser.add_argument("--validate-all", action="store_true") args = parser.parse_args() - - _device = torch.device(f"cuda:{args.device_id}") - log(f"GPU: {args.device_id}, torch device: {_device}") - - scenes = TRAIN_SCENES - if args.scene: - key = args.scene - scenes = [s for s in TRAIN_SCENES if s["scene_id"] == key - or s["scene_id"].startswith(key)] - if not scenes: - log(f"ERROR: No scene matching '{args.scene}'") - return 1 - - for i, scene in enumerate(scenes): - if i > 0: - prev_cfg = scenes[i - 1]["config_path"] - curr_cfg = scene["config_path"] - if prev_cfg != curr_cfg: - log(f"Waiting 120s before config switch...") - time.sleep(120) - - evaluate_train_scene(scene, args.device_id) - - log("All train scenes complete.") + if args.validate or args.validate_all: + case_ids = CASE_IDS if args.validate_all else (args.case,) + reports, failures = validate_outputs(case_ids, args.output_root) + for report in reports: + print(report) + for failure in failures: + print(f"ERROR: {failure}") + return 1 if failures else 0 + case = get_case(args.case) + if args.seed is not None and args.all_seeds: + parser.error("--seed and --all-seeds are mutually exclusive") + seeds = case.seeds if args.all_seeds else (args.seed if args.seed is not None else case.seeds[0],) + evaluate_case(case, tuple(seeds), args.device_id, args.overwrite, args.metrics_only, args.output_root) return 0 diff --git a/src/drl_pinball/eval/tests/test_wake_l2.py b/src/drl_pinball/eval/tests/test_wake_l2.py new file mode 100644 index 0000000..24b2318 --- /dev/null +++ b/src/drl_pinball/eval/tests/test_wake_l2.py @@ -0,0 +1,87 @@ +import csv,json +import numpy as np +import pytest +from drl_pinball.eval import wake_l2 + +def test_formulas_and_wrap(): + z=np.zeros((8,4,5)); ux=np.r_[np.ones((4,4,5)),-np.ones((4,4,5))]; uy=np.full_like(z,2) + g=wake_l2.field_metrics(ux,uy,z,z,2,(slice(None),slice(None)),periodic=True) + assert g["E_phase8"]==pytest.approx(np.sqrt(5)/2); assert g["E_mean"] is None; assert g["E_snapshot"]==pytest.approx(np.sqrt(5)/2) + assert wake_l2.mean_field_metric(np.ones((4,5)),np.full((4,5),2.),np.zeros((4,5)),np.zeros((4,5)),2,(slice(None),slice(None)))==pytest.approx(np.sqrt(5)/2) + assert wake_l2.wrapped_phase_delta([2*np.pi-.1,.1],[0,2*np.pi-.1])==pytest.approx([-.1,.2]) +def test_steady_and_geometry_guards(): + z=np.zeros((1,3,4)); g=wake_l2.field_metrics(np.ones_like(z),z,z,z,2,(slice(None),slice(None)),periodic=False) + assert g=={"E_mean":.5,"E_snapshot":.5,"E_slot":None,"E_phase8":None} + with pytest.raises(ValueError,match="outside"): wake_l2.roi_slices((10,10),D=2,xs=4,downstream_solid_extent=-1) + with pytest.raises(ValueError,match="not downstream"): wake_l2.roi_slices((30,50),D=1,xs=20,downstream_solid_extent=14) +def test_frequency_and_fail_closed(): + t=np.arange(41.); s=np.zeros((41,6)); s[:,3]=np.sin(2*np.pi*(t-1)/10); g=wake_l2.frequency_diagnostics(t,s) + assert g["median_period"]==pytest.approx(10); assert g["dominant_frequency"]==pytest.approx(.1) + assert g["crossing_count"]==4; assert g["period_cv"]==pytest.approx(0) + with pytest.raises(ValueError,match="three rising"): wake_l2.frequency_diagnostics(t[:10],s[:10]) + s[2,3]=np.nan + with pytest.raises(ValueError,match="finite"): wake_l2.frequency_diagnostics(t,s) +def metadata(rd,role,u=.5): + rd.mkdir(parents=True); (rd/"metadata.json").write_text(json.dumps({"schema":"drl-pinball-legacy-acquisition-v2","case_id":"steady","role":role,"source_provenance":{"builder_config":{"values":{"u0":u}}}})) +def test_synthetic_steady_and_atomic_output(tmp_path): + b=tmp_path/"legacy"/"steady" + for role,val in (("target",0),("constant",1),("zero",2)): + rd=b/role; metadata(rd,role); a=np.full((1,9,30),val,dtype=np.float32); np.savez(rd/"late_field.npz",ux=a,uy=np.zeros_like(a),field_indices=[0]) + row=wake_l2._evaluate(**wake_l2.discover(tmp_path,allow_partial=True)[0],geometry_override={"D":.5,"xs":10,"y0":4,"downstream_solid_extent":6}) + assert row["E_mean_ctl"]==pytest.approx(2); assert row["E_mean_zero"]==pytest.approx(4); assert row["eta_mean"]==pytest.approx(.5); assert row["E_phase8_ctl"] is None + assert row["E_mean_ctl"]==row["E_snapshot_ctl"] + assert row["nonperiodic_metric_kind"]=="late_snapshot_error" + payload={"schema":wake_l2.SCHEMA,"rows":[row]}; jp,cp=tmp_path/"a.json",tmp_path/"a.csv"; wake_l2.write_outputs(payload,jp,cp) + assert json.loads(jp.read_text())["rows"][0]["eta_mean"]==pytest.approx(.5) + with cp.open(newline="") as f: + csv_row=next(csv.DictReader(f)); assert csv_row["case_id"]=="steady"; assert csv_row["nonperiodic_metric_kind"]=="late_snapshot_error" + with pytest.raises(FileExistsError): wake_l2.write_outputs(payload,jp,cp) +def test_schema_and_matrix_fail_closed(tmp_path): + p=tmp_path/"bad.npz"; a=np.zeros((8,2,2)); np.savez(p,ux=a,uy=a,mean_ux=np.zeros((2,2)),mean_uy=np.zeros((2,2)),field_indices=range(8),cycle_id=range(8),target_phase=wake_l2.PHASES,actual_phase=wake_l2.PHASES,phase_error=np.zeros(8),extra=1) + with pytest.raises(ValueError,match="exact key"): wake_l2._field(p,True) + p2=tmp_path/"old.npz"; np.savez(p2,ux=a,uy=a,field_indices=range(8),cycle_id=range(8),target_phase=wake_l2.PHASES,actual_phase=wake_l2.PHASES,phase_error=np.zeros(8)) + with pytest.raises(ValueError,match="exact key"): wake_l2._field(p2,True) + (tmp_path/"v5").mkdir(); (tmp_path/"legacy").mkdir() + with pytest.raises(ValueError,match="full matrix"): wake_l2.discover(tmp_path) + +def test_v2_interpretation_and_periodic_selection_diagnostics(tmp_path,monkeypatch): + monkeypatch.setattr(wake_l2,"discover",lambda root,allow_partial=False: []) + payload=wake_l2.evaluate_root(tmp_path,allow_partial=True) + assert payload["schema"]=="drl-pinball-wake-l2-v2"; assert payload["version"]==2 + assert payload["roi"]["x_over_D"]==[-6,14] + assert "not the mean of the eight phase snapshots" in payload["metric_interpretation"]["E_mean"] + assert "complete-cycle" in payload["metric_interpretation"]["E_mean"] + assert "not a repeated-crossing" in payload["metric_interpretation"]["E_phase8"] + assert "smoothed rising crossings" in payload["metric_interpretation"]["phase_reference"] + assert payload["mask_provenance"]["solver_exact_mask_persisted"] is False + assert payload["mask_provenance"]["velocity_zero_mask_inference"] is False + assert payload["frequency_interpretation"]["dominant_frequency_unit"]=="NX/(lattice_time*U0)" + assert payload["formula"]["E_mean"].startswith("RMS difference of complete-cycle mean fields") + b=tmp_path/"periodic" + t=np.arange(41.); sensors=np.zeros((41,6)); sensors[:,3]=np.sin(2*np.pi*(t-1)/10) + for role,val,indices,cycles,mean_val in ( + ("target",0,range(8),range(8),0.0), + ("controlled",1,[0,1,1,3,4,5,6,7],[2,2,3,3,4,4,5,5],3.0), + ("zero",2,range(8),np.zeros(8,int),4.0), + ): + rd=b/role; metadata(rd,role); a=np.full((8,9,30),val,dtype=np.float32) + mean=np.full((9,30),mean_val,dtype=np.float32) + np.savez(rd/"phase_fields.npz",ux=a,uy=np.zeros_like(a),mean_ux=mean,mean_uy=np.zeros_like(mean), + field_indices=indices,cycle_id=cycles,target_phase=wake_l2.PHASES,actual_phase=wake_l2.PHASES,phase_error=np.zeros(8)) + np.savez(rd/"timeseries.npz",physical_time=t,sensors=sensors) + row=wake_l2._evaluate(b/"controlled",b/"target",b/"zero",generation="legacy",case_id="steady",role="controlled",seed=None,periodic=True,geometry_override={"D":.5,"xs":10,"y0":4,"downstream_solid_extent":6}) + assert row["nonperiodic_metric_kind"]=="eight_phase_nearest_snapshot_diagnostic" + assert row["role_field_indices"]==[0,1,1,3,4,5,6,7]; assert row["role_duplicate_field_index"] is True + assert row["role_unique_cycle_count"]==4; assert row["target_unique_cycle_count"]==8; assert row["zero_unique_cycle_count"]==1 + assert row["role_cycle_id"]==[2,2,3,3,4,4,5,5] + assert row["role_phase_delta_slots"]==pytest.approx(np.zeros(8)) + assert row["role_crossing_count"]==4; assert row["role_period_cv"]==pytest.approx(0) + assert all(np.isfinite(row[k]) for k in ("role_period_cv","target_period_cv","zero_period_cv","role_dominant_frequency")) + # Counterexample: eight-phase slot mean differs from complete-cycle mean fields. + # Controlled phase snapshots are all 1 => slot-mean error vs target 0 would be 2, + # but complete-cycle mean_ux=3 => E_mean_ctl = |3-0|/u0 = 6. + assert row["E_mean_ctl"]==pytest.approx(6) + assert row["E_mean_zero"]==pytest.approx(8) + assert row["eta_mean"]==pytest.approx(0.25) + assert row["E_phase8_ctl"]==pytest.approx(2) + assert row["E_mean_ctl"] != pytest.approx(row["E_phase8_ctl"]) diff --git a/src/drl_pinball/eval/wake_l2.py b/src/drl_pinball/eval/wake_l2.py new file mode 100644 index 0000000..17442f0 --- /dev/null +++ b/src/drl_pinball/eval/wake_l2.py @@ -0,0 +1,146 @@ +"""Offline wake-field L2 evaluation; never imports CFD/GPU modules.""" +from __future__ import annotations +import argparse,csv,json,math,os,tempfile +from pathlib import Path +import numpy as np +SCHEMA="drl-pinball-wake-l2-v2" +PHASES=np.arange(8)*np.pi/4 +V5={"kar_re100","kar_d075","kar_d15","kar_d2","ill_075L","ill_1L","ill_15L","ill_2L","kar_re60","kar_re200","kar_re400"} +LEGACY={"karman_re50","karman_re100","karman_re200","karman_re400","illusion_075L","illusion_1L","illusion_15L"} +EXPECTED=(15,11,7,1) +PHASE_KEYS={"ux","uy","mean_ux","mean_uy","field_indices","cycle_id","target_phase","actual_phase","phase_error"}; STEADY_KEYS={"ux","uy","field_indices"} +CSV_FIELDS=("generation","case_id","seed","role","u0","roi_x0","roi_x1","roi_y0","roi_y1","nonperiodic_metric_kind","E_mean_ctl","E_mean_zero","eta_mean","E_phase8_ctl","E_phase8_zero","eta_phase8","E_snapshot_ctl","E_snapshot_zero","eta_snapshot","role_phase_delta_max_abs","role_phase_delta_rms","zero_phase_delta_max_abs","zero_phase_delta_rms","role_phase_error_max_abs","target_phase_error_max_abs","zero_phase_error_max_abs","role_unique_cycle_count","target_unique_cycle_count","zero_unique_cycle_count","role_duplicate_field_index","target_duplicate_field_index","zero_duplicate_field_index","role_median_period","target_median_period","zero_median_period","role_period_cv","target_period_cv","zero_period_cv","role_crossing_count","target_crossing_count","zero_crossing_count","role_dominant_frequency","target_dominant_frequency","zero_dominant_frequency","role_relative_frequency_mismatch","zero_relative_frequency_mismatch") +def wrapped_phase_delta(a,b): + a,b=np.asarray(a,float),np.asarray(b,float) + if a.shape!=b.shape or not np.all(np.isfinite(a)) or not np.all(np.isfinite(b)): raise ValueError("phase arrays must be equal-shape finite arrays") + return (a-b+np.pi)%(2*np.pi)-np.pi +def roi_slices(shape,*,D,xs,y0=None,downstream_solid_extent): + ny,nx=map(int,shape); y0=(ny-1)/2 if y0 is None else float(y0) + if D<=0 or not all(np.isfinite(v) for v in (D,xs,y0,downstream_solid_extent)): raise ValueError("invalid geometry") + x0,x1=math.ceil(xs-6*D),math.floor(xs+14*D); ylo,yhi=math.ceil(y0-5*D),math.floor(y0+5*D) + if x0<0 or ylo<0 or x1>=nx or yhi>=ny: raise ValueError("ROI outside grid") + if x0<=downstream_solid_extent: raise ValueError("ROI not downstream of solid extent") + return slice(ylo,yhi+1),slice(x0,x1+1) +def mean_field_metric(rux,ruy,tux,tuy,u0,roi): + a=[np.asarray(v,float) for v in (rux,ruy,tux,tuy)] + if len({v.shape for v in a})!=1 or a[0].ndim!=2: raise ValueError("mean fields must be equal-shape 2-D arrays") + if u0<=0 or not np.isfinite(u0) or any(not np.all(np.isfinite(v)) for v in a): raise ValueError("non-finite mean field/U0") + ys,xs=roi; du=(a[0][ys,xs]-a[2][ys,xs])/u0; dv=(a[1][ys,xs]-a[3][ys,xs])/u0 + return float(np.sqrt(np.mean(du*du+dv*dv))) +def field_metrics(rux,ruy,tux,tuy,u0,roi,*,periodic): + a=[np.asarray(v) for v in (rux,ruy,tux,tuy)]; k=8 if periodic else 1 + if len({v.shape for v in a})!=1 or a[0].ndim!=3 or a[0].shape[0]!=k: raise ValueError("invalid field shapes") + if u0<=0 or not np.isfinite(u0) or any(not np.all(np.isfinite(v)) for v in a): raise ValueError("non-finite field/U0") + ys,xs=roi; du=(a[0][:,ys,xs].astype(float)-a[2][:,ys,xs])/u0; dv=(a[1][:,ys,xs].astype(float)-a[3][:,ys,xs])/u0 + slots=np.sqrt(np.mean(du*du+dv*dv,axis=(1,2))) + # Periodic E_mean is computed separately from complete-cycle mean_ux/mean_uy. + if periodic: + return {"E_mean":None,"E_snapshot":float(slots[0]),"E_slot":slots.tolist(),"E_phase8":float(np.sqrt(np.mean(du*du+dv*dv)))} + return {"E_mean":float(slots[0]),"E_snapshot":float(slots[0]),"E_slot":None,"E_phase8":None} +def frequency_diagnostics(times,sensors): + t,s=np.asarray(times,float),np.asarray(sensors,float) + if t.ndim!=1 or s.ndim!=2 or s.shape[0]!=len(t) or s.shape[1]<4 or not np.all(np.isfinite(t)) or not np.all(np.isfinite(s)) or np.any(np.diff(t)<=0): raise ValueError("timeseries must be finite and ordered") + y=s[:,3]; i=np.flatnonzero((y[:-1]<=0)&(y[1:]>0)) + if len(i)<3: raise ValueError("at least three rising crossings required") + c=t[i]-y[i]*(t[i+1]-t[i])/(y[i+1]-y[i]); periods=np.diff(c); p=float(np.median(periods)); cv=float(np.std(periods)/np.mean(periods)) + if p<=0 or not np.isfinite(p) or not np.isfinite(cv): raise ValueError("invalid period") + return {"median_period":p,"dominant_frequency":1/p,"crossing_count":int(len(c)),"period_cv":cv} +def _json(p): + with Path(p).open(encoding="utf8") as f: d=json.load(f) + if not isinstance(d,dict): raise ValueError(f"{p}: JSON object required") + return d +def _meta(rd,generation,case,role): + d=_json(rd/"metadata.json"); schema=f"drl-pinball-{'v5' if generation=='v5' else 'legacy'}-acquisition-v2" + if (d.get("schema"),d.get("case_id"),d.get("role"))!=(schema,case,role): raise ValueError(f"{rd}: metadata mismatch") + if generation=="v5": + cp=Path(d.get("config",{}).get("path","")); cp=cp if cp.is_absolute() else (rd/cp).resolve(); u=_json(cp).get("physics",{}).get("velocity") + else: u=d.get("source_provenance",{}).get("builder_config",{}).get("values",{}).get("u0") + try: u=float(u) + except (TypeError,ValueError) as e: raise ValueError(f"{rd}: invalid U0") from e + if u<=0 or not np.isfinite(u): raise ValueError(f"{rd}: invalid U0") + return d,u +def _field(path,periodic): + keys=PHASE_KEYS if periodic else STEADY_KEYS + with np.load(path,allow_pickle=False) as z: + if set(z.files)!=keys: raise ValueError(f"{path}: exact key schema mismatch") + d={k:np.array(z[k]) for k in z.files} + stacked=("ux","uy","field_indices")+ (("cycle_id","target_phase","actual_phase","phase_error") if periodic else ()) + if d["ux"].shape!=d["uy"].shape or d["ux"].ndim!=3 or d["ux"].shape[0]!=(8 if periodic else 1) or any(not np.all(np.isfinite(d[k])) for k in stacked): raise ValueError(f"{path}: shape/finite validation failed") + if periodic: + if d["mean_ux"].shape!=d["ux"].shape[1:] or d["mean_uy"].shape!=d["uy"].shape[1:] or d["mean_ux"].ndim!=2: raise ValueError(f"{path}: mean field shape mismatch") + if not np.all(np.isfinite(d["mean_ux"])) or not np.all(np.isfinite(d["mean_uy"])): raise ValueError(f"{path}: non-finite mean fields") + if d["target_phase"].shape!=(8,) or not np.allclose(d["target_phase"],PHASES,rtol=0,atol=1e-7): raise ValueError(f"{path}: noncanonical phases") + return d +def _freq(rd): + with np.load(rd/"timeseries.npz",allow_pickle=False) as z: + if not {"physical_time","sensors"}<=set(z.files): raise ValueError(f"{rd}: timeseries keys missing") + return frequency_diagnostics(z["physical_time"],z["sensors"]) +def _pdiag(f,reference): + d=wrapped_phase_delta(f["actual_phase"],reference["actual_phase"]); return d.tolist(),float(np.max(abs(d))),float(np.sqrt(np.mean(d*d))),float(np.max(abs(f["phase_error"]))) +def _selection_diag(f): + indices=np.asarray(f["field_indices"]); cycles=np.asarray(f["cycle_id"]) + return {"field_indices":indices.tolist(),"cycle_id":cycles.tolist(),"unique_cycle_count":int(len(np.unique(cycles))),"duplicate_field_index":bool(len(np.unique(indices))!=len(indices))} +def _geometry(g,c): + return ((600.,440.) if c.startswith("ill_") else (1200.,1036.)) if g=="v5" else ((600.,440.) if c.startswith("illusion_") else (800.,636.)) +def _evaluate(role_dir,target_dir,zero_dir,*,generation,case_id,role,seed,periodic,geometry_override=None): + _,u=_meta(role_dir,generation,case_id,role); _,ut=_meta(target_dir,generation,case_id,"target"); _,uz=_meta(zero_dir,generation,case_id,"zero") + if not u==ut==uz: raise ValueError("role/target/zero U0 mismatch") + fn="phase_fields.npz" if periodic else "late_field.npz"; rf,tf,zf=(_field(p/fn,periodic) for p in (role_dir,target_dir,zero_dir)) + if len({rf["ux"].shape,tf["ux"].shape,zf["ux"].shape})!=1: raise ValueError("role field shapes differ") + if periodic and (not np.array_equal(rf["target_phase"],tf["target_phase"]) or not np.array_equal(zf["target_phase"],tf["target_phase"])): raise ValueError("target_phase mismatch") + xs,extent=_geometry(generation,case_id); o=geometry_override or {}; roi=roi_slices(rf["ux"].shape[1:],D=float(o.get("D",20)),xs=float(o.get("xs",xs)),y0=o.get("y0"),downstream_solid_extent=float(o.get("downstream_solid_extent",extent))) + ctl=field_metrics(rf["ux"],rf["uy"],tf["ux"],tf["uy"],u,roi,periodic=periodic); zero=field_metrics(zf["ux"],zf["uy"],tf["ux"],tf["uy"],u,roi,periodic=periodic) + if periodic: + ctl["E_mean"]=mean_field_metric(rf["mean_ux"],rf["mean_uy"],tf["mean_ux"],tf["mean_uy"],u,roi) + zero["E_mean"]=mean_field_metric(zf["mean_ux"],zf["mean_uy"],tf["mean_ux"],tf["mean_uy"],u,roi) + def eta(k): + if zero[k] is None or zero[k]<=0: raise ValueError(f"zero {k} must be positive") + return 1-ctl[k]/zero[k] + row={"generation":generation,"case_id":case_id,"seed":seed,"role":role,"u0":u,"roi_x0":roi[1].start,"roi_x1":roi[1].stop-1,"roi_y0":roi[0].start,"roi_y1":roi[0].stop-1,"nonperiodic_metric_kind":"eight_phase_nearest_snapshot_diagnostic" if periodic else "late_snapshot_error","E_mean_ctl":ctl["E_mean"],"E_mean_zero":zero["E_mean"],"eta_mean":eta("E_mean"),"E_phase8_ctl":ctl["E_phase8"],"E_phase8_zero":zero["E_phase8"],"eta_phase8":eta("E_phase8") if periodic else None,"E_snapshot_ctl":ctl["E_snapshot"],"E_snapshot_zero":zero["E_snapshot"],"eta_snapshot":eta("E_snapshot"),"E_slot_ctl":ctl["E_slot"],"E_slot_zero":zero["E_slot"]} + for p in ("role","target","zero"): + for k in ("phase_delta_slots","phase_delta_max_abs","phase_delta_rms","phase_error_max_abs","field_indices","cycle_id","unique_cycle_count","duplicate_field_index","median_period","period_cv","crossing_count","dominant_frequency","relative_frequency_mismatch"): row[f"{p}_{k}"]=None + if periodic: + fs={"role":rf,"target":tf,"zero":zf}; fq={p:_freq(d) for p,d in (("role",role_dir),("target",target_dir),("zero",zero_dir))}; targetf=fq["target"]["dominant_frequency"] + for p,f in fs.items(): + row[f"{p}_phase_delta_slots"],row[f"{p}_phase_delta_max_abs"],row[f"{p}_phase_delta_rms"],row[f"{p}_phase_error_max_abs"]=_pdiag(f,tf) + for k,v in _selection_diag(f).items(): row[f"{p}_{k}"]=v + for k in ("median_period","period_cv","crossing_count","dominant_frequency"): row[f"{p}_{k}"]=fq[p][k] + row[f"{p}_relative_frequency_mismatch"]=abs(fq[p]["dominant_frequency"]-targetf)/targetf + return row +def discover(root,*,allow_partial=False): + root=Path(root); v5,legacy=root/"v5",root/"legacy"; ctl=[]; seen=set() + for rd in sorted(v5.glob("*_seed*/controlled")): + m=_json(rd/"metadata.json"); key=(m.get("case_id"),m.get("seed")) + if key[0] not in V5 or not isinstance(key[1],int) or key in seen: raise ValueError(f"{rd}: invalid/duplicate mapping") + seen.add(key); ctl.append((rd,*key)) + vc={c for _,c,_ in ctl}; lp={p.name for p in legacy.iterdir() if p.is_dir() and p.name in LEGACY}; counts=(len(ctl),len(vc),len(lp),int((legacy/"steady").is_dir())) + if not allow_partial and (counts!=EXPECTED or vc!=V5 or lp!=LEGACY): raise ValueError(f"full matrix discovery failed: {counts} != {EXPECTED}") + jobs=[dict(role_dir=r,target_dir=v5/c/"target",zero_dir=v5/c/"zero",generation="v5",case_id=c,role="controlled",seed=s,periodic=True) for r,c,s in ctl] + for c in sorted(lp): + b=legacy/c; jobs.append(dict(role_dir=b/"controlled",target_dir=b/"target",zero_dir=b/"zero",generation="legacy",case_id=c,role="controlled",seed=None,periodic=True)) + if (legacy/"steady").is_dir(): + b=legacy/"steady"; jobs.append(dict(role_dir=b/"constant",target_dir=b/"target",zero_dir=b/"zero",generation="legacy",case_id="steady",role="constant",seed=None,periodic=False)) + if not jobs: raise ValueError("no evaluable jobs") + return jobs +def evaluate_root(root,*,allow_partial=False): + return {"schema":SCHEMA,"version":2,"roi":{"D":20,"x_over_D":[-6,14],"abs_y_over_D_max":5,"coordinates":"inclusive lattice cells; y0=(ny-1)/2","geometry_guard":"Every included scenario has ROI x lower bound strictly downstream of its conservative solid extent; the registered -6..14D ROI is therefore expected to contain fluid only."},"metric_interpretation":{"E_mean":"For periodic cases, U0-normalized ROI RMS difference of complete-cycle arithmetic mean fields (mean_ux/mean_uy) over accepted rising crossings [t_first,t_last). This is not the mean of the eight phase snapshots. For steady, E_mean and E_snapshot are the same late-snapshot error.","E_phase8":"Discrete phase-resolved diagnostic from eight independently phase-referenced nearest-boundary snapshots; not a repeated-crossing phase-conditioned ensemble metric.","phase_reference":"Target, controlled, and zero are independently phased by smoothed rising crossings of each role's center-sensor uy with a minimum-gap filter. A shared target_phase slot is a common mathematical phase coordinate, not independent evidence that physical events are exactly equivalent across roles.","phase_alignment":"No global phase minimization or interpolation is performed.","steady":"nonperiodic_metric_kind=late_snapshot_error.","periodic":"nonperiodic_metric_kind=eight_phase_nearest_snapshot_diagnostic."},"mask_provenance":{"status":"geometry-guarded ROI diagnostic","solver_exact_mask_persisted":False,"known_solid_cells_expected_in_roi":False,"basis":"All included ROI x lower bounds are strictly downstream of the conservative solid extent.","claim_limit":"Without a persisted solver-exact mask, this is not a solver-mask-proven exact common-fluid metric.","velocity_zero_mask_inference":False,"prohibition":"velocity==0 is not used or permitted as a fluid/solid mask inference."},"frequency_interpretation":{"physical_time_input_unit":"lattice_time*U0/NX (tU0/NX)","dominant_frequency_unit":"NX/(lattice_time*U0)","use":"Relative mismatch only among target, controlled, and zero within the same case.","cross_generation_limit":"Not for Legacy/V5 physical-frequency comparison."},"formula":{"slot":"sqrt(mean_ROI((du_x^2+du_y^2)/U0^2))","E_phase8":"RMS over eight slots and cells","E_mean":"RMS difference of complete-cycle mean fields mean_ux/mean_uy","eta":"1-E_role/E_zero","period_cv":"std(rising-crossing periods)/mean(rising-crossing periods)"},"limitations":["Each phase slot is one snapshot, not a repeated-crossing average.","phi0 is the canonical slot-0 snapshot.","No global phase minimization or interpolation.","E_mean uses only complete cycles between first and last accepted rising crossings."],"exclusions":["legacy erase","legacy vortex","other directories"],"rows":[_evaluate(**j) for j in discover(root,allow_partial=allow_partial)]} +def _atomic(path,writer): + path=Path(path); path.parent.mkdir(parents=True,exist_ok=True); fd,tmp=tempfile.mkstemp(prefix="."+path.name+".",suffix=".tmp",dir=path.parent) + try: + with os.fdopen(fd,"w",encoding="utf8",newline="") as f: writer(f); f.flush(); os.fsync(f.fileno()) + os.replace(tmp,path) + except BaseException: + try: os.unlink(tmp) + except FileNotFoundError: pass + raise +def write_outputs(payload,json_path,csv_path,*,overwrite=False): + jp,cp=Path(json_path),Path(csv_path) + if jp.resolve()==cp.resolve(): raise ValueError("outputs must differ") + if not overwrite and any(p.exists() for p in (jp,cp)): raise FileExistsError("refusing overwrite") + def wj(f): json.dump(payload,f,indent=2,allow_nan=False); f.write("\n") + def wc(f): + w=csv.DictWriter(f,fieldnames=CSV_FIELDS,extrasaction="ignore"); w.writeheader(); w.writerows(payload["rows"]) + _atomic(jp,wj); _atomic(cp,wc) +def main(argv=None): + default=Path(__file__).resolve().parents[1]/"data"/"reproduction"; p=argparse.ArgumentParser(); p.add_argument("--root",type=Path,default=default); p.add_argument("--output-json",type=Path); p.add_argument("--output-csv",type=Path); p.add_argument("--overwrite",action="store_true"); p.add_argument("--allow-partial",action="store_true"); a=p.parse_args(argv); payload=evaluate_root(a.root,allow_partial=a.allow_partial); write_outputs(payload,a.output_json or a.root/"wake_l2_summary.json",a.output_csv or a.root/"wake_l2_summary.csv",overwrite=a.overwrite); return 0 +if __name__=="__main__": raise SystemExit(main()) diff --git a/src/drl_pinball/legacy_test/__main__.py b/src/drl_pinball/legacy_test/__main__.py new file mode 100644 index 0000000..4dc025f --- /dev/null +++ b/src/drl_pinball/legacy_test/__main__.py @@ -0,0 +1,4 @@ +from .run import main + +if __name__ == "__main__": + main() diff --git a/src/drl_pinball/legacy_test/acquire.py b/src/drl_pinball/legacy_test/acquire.py new file mode 100644 index 0000000..aa586a2 --- /dev/null +++ b/src/drl_pinball/legacy_test/acquire.py @@ -0,0 +1,819 @@ +"""Legacy Karman Re100 role acquisition; solver imports stay lazy.""" +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import shutil +from collections import deque +from pathlib import Path +from typing import Any, Callable + +import numpy as np + +from drl_pinball.acquisition import ( + accumulate_mean_fields, assign_phase, assign_periodic_phase, + cleanup_scratch, complete_cycle_field_indices, create_scratch, + default_reproduction_mapping, dual_cycle_dtw, pooled_phase_bins, + prepare_role_output, publish_role_output, publish_selected_fields, + select_phase_fields, validate_output_storage, write_boundary_artifacts, + write_phase_cycle_artifacts, +) +from .cases import get_case +from .metrics import erase_reward_terms, frozen_reference_comparison, reward_terms, sha256_file +from .core.dtw_metrics import gen_target_states_at +from .runtime import erase_policy_observation, load_policy_norm, policy_observation, reset_runtime, run_historical_interval + +CASE_NAME, ROLE = "karman_re100", "controlled" # compatibility defaults +PERIODIC_CASES = ("karman_re50", "karman_re100", "karman_re200", "karman_re400", + "illusion_075L", "illusion_1L", "illusion_15L") +SUPPORTED_CASES = PERIODIC_CASES + ("steady", "vortex_lamb", "vortex_taylor", "erase") +PERIODIC_ROLES = ("controlled", "target", "zero") +STEADY_ROLES = ("target", "constant", "zero") +ROLES = PERIODIC_ROLES +WARMUP_INTERVALS, COLLECT_BOUNDARIES = 480, 160 +NX, U0, SI, CYCLE_LENGTH = 1280, 0.01, 800, 30 # compatibility defaults +MIN_FREE_BYTES = 4 * 1024**3 + + +def _atomic_json(path: Path, payload: dict[str, Any]) -> None: + temporary = path.with_name(f".{path.name}.tmp") + with temporary.open("w", encoding="utf-8") as stream: + json.dump(payload, stream, indent=2) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + + +def _identity(path: Path) -> dict[str, Any]: + return {"path": str(path.resolve()), "sha256": sha256_file(str(path)), "bytes": path.stat().st_size} + + +def _canonical_json_value(value): + """Convert builder contract values to deterministic JSON-native values.""" + if isinstance(value, dict): + return {str(key): _canonical_json_value(current) for key, current in value.items()} + if isinstance(value, (list, tuple)): + return [_canonical_json_value(current) for current in value] + if isinstance(value, np.ndarray): + return _canonical_json_value(value.tolist()) + if isinstance(value, np.generic): + return value.item() + if value is None or isinstance(value, (str, int, float, bool)): + return value + raise TypeError(f"builder config contains non-JSON value {type(value).__name__}") + + +def _optional_file_provenance(path: Path | None): + if path is None: + return None + path = Path(path) + return {"path": str(path.resolve()), "exists": path.is_file(), + "file_identity": _identity(path) if path.is_file() else None} + + +def _generated_array_identity(value) -> dict[str, Any]: + """Identify a run-generated numeric array by dtype, shape, and C-order content.""" + array = np.asarray(value) + if array.dtype.hasobject: + raise TypeError("run-generated identity does not support object arrays") + descriptor = {"dtype": array.dtype.str, "shape": list(array.shape), "order": "C"} + header = json.dumps(descriptor, sort_keys=True, separators=(",", ":"), + allow_nan=False).encode("ascii") + content = np.ascontiguousarray(array).tobytes(order="C") + digest = hashlib.sha256(header + b"\0" + content).hexdigest() + return {"identity_kind": "run-generated-array-content", **descriptor, + "content_sha256": digest, "bytes": len(content)} + + +def _generated_config_identity(config) -> dict[str, Any]: + """Identify a run-generated builder config by canonical JSON content.""" + values = _canonical_json_value(config) + canonical = json.dumps(values, sort_keys=True, separators=(",", ":"), + ensure_ascii=True, allow_nan=False) + return {"identity_kind": "run-generated-builder-config", "values": values, + "canonical_json_sha256": hashlib.sha256(canonical.encode("ascii")).hexdigest()} + + +def _role_semantics(role: str) -> str: + return { + "controlled": "frozen-policy controlled trajectory", + "target": "live builder-generated target trajectory; actions, rewards, and policy normalization unavailable", + "zero": "physical-zero/uncontrolled trajectory using counter-bias normalized action; frozen norm used only for native reward evaluation", + "constant": "established steady constant control at case.action_bias*U0 from normalized action zero", + }[role] + + +def _controlled_reference(role, bundle, columns, conv_len): + if role != "controlled": + return None + return frozen_reference_comparison( + str(bundle["reference_path"]), columns["sensors"], columns["forces"], + columns["action_normalized"], conv_len, + ) + + +def _case_roles(case): + if case.name not in SUPPORTED_CASES: + raise ValueError(f"unsupported Legacy acquisition case {case.name!r}") + return STEADY_ROLES if case.scene == "steady" else PERIODIC_ROLES + + +def _resolve_bundle( + repo_root: Path, case_name: str = CASE_NAME, role: str = "controlled", +) -> dict[str, Any]: + """Resolve immutable CPU artifacts before solver construction.""" + case = get_case(case_name) + _case_roles(case) + frozen = repo_root / "src" / "SR_analysis" / "data" / case.scene / (case.reference_case or case.name) + needs_model = case.model is not None and role == "controlled" + needs_frozen_policy = (case.model is not None and role != "target" and + case.scene != "erase" and + not (case.scene == "vortex" and role == "zero")) + norm_path = frozen / "norm.json" if needs_frozen_policy else None + norm = load_policy_norm(str(norm_path)) if norm_path is not None else None + model = None + if needs_model: + from .core.model_loader import model_path + model = Path(model_path(case.model)) + return {"case": case, "model_path": model, "norm_path": norm_path, "norm": norm, + "reference_path": frozen / "controlled.npz"} + + +def _default_runtime(case, bundle, device_id, role="controlled"): + from .core.legacy_env_builder import ( + build_illusion, build_illusion_target, build_karman_cloak, build_karman_target, + build_erase, build_steady_cloak, build_steady_target, build_vortex, + ) + if case.scene == "vortex": + data = build_vortex(device_id, case.re_code, vortex_type=case.vortex_type, action_scale=case.action_scale, target_only=(role == "target")) + elif case.scene == "erase": + data = build_erase(device_id, case.re_code, target_only=(role == "target")) + elif case.scene == "karman": + data = (build_karman_target(device_id, case.re_code, sample_interval=case.sample_interval) + if role == "target" else + build_karman_cloak(device_id, case.re_code, sample_interval=case.sample_interval)) + elif case.scene == "illusion": + kwargs = dict(device_id=device_id, re_code=case.re_code, + target_diameter_L=case.target_radius_l, sample_interval=case.sample_interval) + data = build_illusion_target(**kwargs) if role == "target" else build_illusion(**kwargs) + elif case.scene == "steady": + data = (build_steady_target(device_id, case.re_code, sample_interval=case.sample_interval) + if role == "target" else build_steady_cloak(device_id, case.re_code)) + else: + raise ValueError(f"unsupported Legacy acquisition scene {case.scene!r}") + model = None + if role == "controlled": + from .core.model_loader import load_model + model = load_model(case.model) + return data, model, data["flow_field"] + + +def _run_interval(ff, steps: int, command: np.ndarray) -> None: + context = getattr(ff, "context", None) + if context is None: + run_historical_interval(ff, steps, command) + return + context.push() + try: + run_historical_interval(ff, steps, command) + finally: + context.pop() + + +def _command(case, action, n_objects, dtype): + action = np.asarray(action, dtype=np.float32).reshape(3) + command = np.zeros(n_objects, dtype=dtype) + command[-3:] = (action * case.action_scale + np.asarray(case.action_bias)) * U0 + return action, command + + +def _clock(ff, completed_intervals: int, sample_interval: int = SI) -> tuple[int, int, int | None]: + """Return relative policy clocks plus the solver's absolute lattice clock.""" + lattice_step = completed_intervals * sample_interval + absolute = None + if hasattr(ff, "solver_clock_state"): + state = ff.solver_clock_state() + absolute = int(state["solver_absolute_lattice_clock"]) + if absolute < lattice_step: + raise ValueError("solver lattice clock regressed below acquisition timeline") + return lattice_step, completed_intervals, absolute + + +def _capture_boundary_field(ff, nx: int, ny: int) -> dict[str, np.ndarray]: + """Capture public Legacy q/U0, convert to physical q/RHO_ref, and canonicalize.""" + velocity = ff.current_step_velocity_field() + if not isinstance(velocity, tuple) or len(velocity) != 2: + raise ValueError("public Legacy velocity accessor must return (ux, uy)") + flags = np.asarray(ff.completed_flags_xy()) + if flags.dtype != np.uint8 or flags.shape != (nx, ny): + raise ValueError("completed Legacy flags must be exact uint8 (nx,ny)") + fluid_xy = (flags & np.uint8(0b00000001)) != 0 + u0 = float(ff.field_config.velocity) + if not np.isfinite(u0) or u0 <= 0: + raise ValueError("Legacy field_config.velocity U0 must be positive and finite") + result = {} + for name, field in zip(("ux", "uy"), velocity): + value = np.asarray(field) + if value.shape != (nx, ny) or value.dtype.kind != "f": + raise ValueError("public Legacy velocity must be floating (nx,ny) q/U0") + if not np.isfinite(value[fluid_xy]).all(): + raise ValueError("public Legacy fluid velocity contains non-finite values") + if not np.array_equal(value[~fluid_xy], np.zeros(np.count_nonzero(~fluid_xy), value.dtype)): + raise ValueError("public Legacy velocity must be exact zero on nonfluid cells") + result[name] = np.asarray((value * u0).T, dtype=np.float32) + return result + + +def _target_sensors(target_states): + target = np.asarray(target_states) + if target.ndim != 2: + raise ValueError("target_states must be two-dimensional") + if target.shape[1] == 6: + return target + if target.shape[1] == 8: + return target[:, 2:8] + raise ValueError("target_states must expose exactly six sensor channels") + + +def _observation_slices(case, role, boundary_obs): + if role == "target": + return ((boundary_obs[2:8], np.full(6, np.nan, np.float32)) + if case.scene in ("karman", "illusion") else + (boundary_obs[:6], np.full(6, np.nan, np.float32))) + raw = boundary_obs[:12] if case.scene in ("illusion", "steady") else boundary_obs[2:14] + return raw[:6], raw[6:12] + + +def _collect_role(role, data, model, ff, policy_norm, scratch: Path, + capture_field: Callable = _capture_boundary_field, case_name: str = CASE_NAME): + case = get_case(case_name) + if role not in _case_roles(case): + raise ValueError(f"unknown acquisition role {role!r} for {case.name}") + config = data["config"]; si = int(case.sample_interval) + expected_objects = 4 if role == "target" and case.scene in ("karman", "illusion") else (3 if role == "target" else (7 if case.scene == "karman" else 6)) + if (int(config["nx"]), int(config["sample_interval"]), int(config["n_obj_total"])) != (NX, si, expected_objects): + raise ValueError(f"built Legacy {case.name} {role} lattice/control contract changed") + if role == "target": + fifo = obs = None + else: + saved = np.asarray(data["norm"]["save_states"], dtype=np.float32) + obs_dim = 14 if case.scene == "illusion" else 12 + fifo, obs = reset_runtime(ff, saved, obs_dim) + if not isinstance(fifo, deque) or len(fifo) != len(saved): + raise ValueError("canonical reset FIFO contract failed") + rows = [] + field_buffer = None + total = WARMUP_INTERVALS + COLLECT_BOUNDARIES + harmonics = data.get("target_harmonics") + steady = case.scene == "steady" + for interval in range(total): + if role == "controlled": + action, _ = model.predict(obs, deterministic=True) + action, command = _command(case, action, expected_objects, ff.DATA_TYPE) + elif role == "zero": + action = -np.asarray(case.action_bias, np.float32) / np.float32(case.action_scale) + action, command = _command(case, action, expected_objects, ff.DATA_TYPE) + if not np.allclose(command[-3:], 0.0, rtol=0.0, atol=np.finfo(ff.DATA_TYPE).eps): + raise ValueError("Legacy zero role counter-bias did not cancel the affine command") + command[-3:] = np.zeros(3, dtype=ff.DATA_TYPE) + elif role == "constant": + action = np.zeros(3, np.float32) + action, command = _command(case, action, expected_objects, ff.DATA_TYPE) + else: + action = np.full(3, np.nan, np.float32); command = np.zeros(expected_objects, ff.DATA_TYPE) + _run_interval(ff, si, command) + boundary_obs = np.asarray(ff.obs, np.float32).copy() + sensors, forces = _observation_slices(case, role, boundary_obs) + if role == "target": + effective = np.full(3, np.nan, np.float32) + terms = {name: np.nan for name in ("reward", "reward_cd", "reward_cl", "native_legacy_dtw")} + else: + raw = np.concatenate((sensors, forces)).astype(np.float32) + fifo.append(raw) + if steady: + terms = {name: np.nan for name in ("reward", "reward_cd", "reward_cl", "native_legacy_dtw")} + else: + terms = reward_terms(case, data["target_states"], harmonics, np.asarray(fifo), policy_norm["force_norm_fact"], interval) + if not steady: + target_force = gen_target_states_at(interval + 1, harmonics)[:2] if harmonics is not None else None + obs = policy_observation(raw, policy_norm, target_force=target_force) + effective = np.asarray(ff.current_effective_action(), np.float32)[-3:].copy() + if interval < WARMUP_INTERVALS: + continue + boundary = interval - WARMUP_INTERVALS + lattice_step, control_clock, absolute = _clock(ff, interval + 1, si) + native = float(terms["native_legacy_dtw"]) + row = {"physical_time": lattice_step * U0 / NX, "lattice_step": lattice_step, + "control_index": control_clock, "sensors": sensors.copy(), "forces": forces.copy(), + "action_normalized": action.copy(), + "commanded_target_omega": (np.full(3, np.nan, np.float32) if role == "target" else command[-3:].astype(np.float32, copy=True)), + "effective_smoothed_omega": effective, "reward_raw": float(terms["reward"]), + "reward_cd": float(terms["reward_cd"]), "reward_cl": float(terms["reward_cl"]), + "reward_sim": float(np.exp(-10 * abs(native - 1))) if np.isfinite(native) else np.nan, + "native_reward_dtw": native} + if absolute is not None: row["solver_absolute_lattice_step"] = absolute + rows.append(row) + if not steady or boundary == COLLECT_BOUNDARIES - 1: + fields = capture_field(ff, int(config["nx"]), int(config["ny"])) + if field_buffer is None: + count = 1 if steady else COLLECT_BOUNDARIES + field_buffer = _allocate_field_buffer(fields["ux"].shape, count) + store_at = 0 if steady else boundary + _store_field_candidate(field_buffer, store_at, fields) + return rows, field_buffer + + +def _collect(data, model, ff, policy_norm, scratch: Path, capture_field: Callable = _capture_boundary_field): + return _collect_role("controlled", data, model, ff, policy_norm, scratch, capture_field=capture_field) + + +def _allocate_field_buffer(shape_yx, count: int = COLLECT_BOUNDARIES): + ny, nx = map(int, shape_yx) + return { + "ux": np.empty((int(count), ny, nx), dtype=np.float32), + "uy": np.empty((int(count), ny, nx), dtype=np.float32), + } + + +def _store_field_candidate(field_buffer, boundary: int, fields): + field_buffer["ux"][boundary] = np.asarray(fields["ux"], dtype=np.float32) + field_buffer["uy"][boundary] = np.asarray(fields["uy"], dtype=np.float32) + + +def _columns(rows): + return {name: np.asarray([row[name] for row in rows]) for name in rows[0]} + + +def _require_periodic_crossings(phase, minimum: int = 4): + count = len(phase["crossing_times"]) + if count < minimum: + raise ValueError(f"periodic publication requires at least {minimum} center-uy rising crossings") + return count + + +def _phase_values(columns, role="controlled"): + excluded = {"physical_time", "lattice_step", "solver_absolute_lattice_step", + "control_index", "phase", "cycle_id"} + if role == "target": + return {"sensors": columns["sensors"]} + return {name: value for name, value in columns.items() if name not in excluded} + + +def _minimum_crossing_interval(physical_time, cycle_length): + times = np.asarray(physical_time, dtype=np.float64) + if len(times) < 2: + raise ValueError("physical_time must contain at least two boundaries") + dt = float(np.median(np.diff(times))) + if not np.isfinite(dt) or dt <= 0: + raise ValueError("physical_time spacing must be positive and finite") + return 0.5 * float(cycle_length) * dt + + +PHASE_FIELD_KEYS = { + "ux", "uy", "mean_ux", "mean_uy", "field_indices", "cycle_id", + "target_phase", "actual_phase", "phase_error", +} + + +def _finalize(role_dir: Path, scratch: Path, rows, field_buffer, data, bundle, storage, role="controlled"): + case = bundle["case"]; columns = _columns(rows) + phase_meta = {} + if case.scene == "steady": + write_boundary_artifacts(role_dir / "timeseries.npz", role_dir / "timeseries.csv", columns) + publish_selected_fields(role_dir / "late_field.npz", { + "ux": np.asarray(field_buffer["ux"][:1], dtype=np.float32), + "uy": np.asarray(field_buffer["uy"][:1], dtype=np.float32), + "field_indices": np.asarray([len(rows) - 1], dtype=np.int64), + }) + phase_values = {} + else: + dtw = dual_cycle_dtw(_target_sensors(data["target_states"]), columns["sensors"], + columns["native_reward_dtw"], cycle_length=case.conv_len, lag_channel=3) + columns.update(target_normalized_dtw=dtw["target_normalized_dtw"], + target_normalized_dtw_lag=dtw["target_normalized_dtw_lag"]) + min_gap = _minimum_crossing_interval(columns["physical_time"], case.conv_len) + phase = assign_periodic_phase( + columns["physical_time"], columns["sensors"], minimum_crossing_interval=min_gap, + ) + _require_periodic_crossings(phase) + columns.update(phase=phase["phase"], cycle_id=phase["cycle_id"]) + write_boundary_artifacts(role_dir / "timeseries.npz", role_dir / "timeseries.csv", columns) + phase_values = _phase_values(columns, role) + write_phase_cycle_artifacts(role_dir / "phase_cycle.npz", role_dir / "phase_cycle.csv", pooled_phase_bins(phase["phase"], phase_values, bins=32)) + selected = select_phase_fields(columns["physical_time"], phase["crossing_times"]) + mean_span = complete_cycle_field_indices(columns["physical_time"], phase["crossing_times"]) + mean = accumulate_mean_fields(field_buffer["ux"], field_buffer["uy"], mean_span["field_indices"]) + published = { + "ux": np.asarray(field_buffer["ux"][selected["field_indices"]], dtype=np.float32), + "uy": np.asarray(field_buffer["uy"][selected["field_indices"]], dtype=np.float32), + "mean_ux": mean["mean_ux"], + "mean_uy": mean["mean_uy"], + } + published.update(selected) + field_path = publish_selected_fields(role_dir / "phase_fields.npz", published) + with np.load(field_path, allow_pickle=False) as saved: + if (saved["ux"].shape[0] != 8 or set(saved.files) != PHASE_FIELD_KEYS + or saved["mean_ux"].shape != saved["ux"].shape[1:] + or saved["mean_uy"].shape != saved["uy"].shape[1:]): + raise ValueError("published phase fields failed exact-key validation") + reference = _controlled_reference(role, bundle, columns, case.conv_len) + _atomic_json(role_dir / "dtw_summary.json", {"native_pipeline": "unavailable for target role" if role == "target" else "original Legacy reward DTW", + "target_normalized_dtw": dtw["metadata"]["definition"], "historical_window": case.conv_len, + "lag_channel": 3, "target_scale": dtw["metadata"]["scale"].tolist(), + "native_mean": (None if role == "target" else float(np.nanmean(columns["native_reward_dtw"]))), + "target_normalized_dtw_finite_mean": float(np.nanmean(columns["target_normalized_dtw"])), + "crossing_count": int(len(phase["crossing_times"])), + "frozen_reference_comparison": reference}) + phase_meta = { + "phase_contract": "smoothed rising crossings of sensors[:,3] with min-gap filter; complete half-open cycles", + "phase_smoothing_kernel": phase["smoothing_kernel"], + "minimum_crossing_interval": phase["minimum_crossing_interval"], + "accepted_crossing_count": phase["accepted_crossing_count"], + "rejected_crossing_count": phase["rejected_crossing_count"], + "complete_cycle_count": phase["complete_cycle_count"], + "mean_field_count": mean_span["mean_field_count"], + "mean_first_crossing_time": mean_span["first_crossing_time"], + "mean_last_crossing_time": mean_span["last_crossing_time"], + "candidate_field_storage": "single-role in-memory FP32 ux/uy buffer; no boundary_*.npz scratch", + } + cleanup_scratch(scratch, root=role_dir / "scratch"); scratch.parent.rmdir() + _atomic_json(role_dir / "metadata.json", {"schema": "drl-pinball-legacy-acquisition-v2", "case_id": case.name, "scene": case.scene, "role": role, + "warmup_control_steps": WARMUP_INTERVALS, "collected_post_step_boundaries": COLLECT_BOUNDARIES, + "si": case.sample_interval, "historical_window": case.conv_len, "phase_sensor": "fixed center sensor uy (sensors[:,3])" if case.scene != "steady" else None, + "field_contract": {"quantity": "q/RHO_ref", "RHO_ref": 1, "source": "public current_step_velocity_field q/U0 multiplied by field_config.velocity U0"}, + "field_count": 1 if case.scene == "steady" else 8, "nonperiodic": case.scene == "steady", + "role_semantics": _role_semantics(role), + "timeline_contract": "V5 lattice_step is an absolute solver count; Legacy physical_time is derived only from relative Legacy lattice_step", + "source_provenance": { + "model_file": _identity(bundle["model_path"]) if role == "controlled" else None, + "normalizer_file": (_identity(bundle["norm_path"]) if bundle["norm_path"] is not None else None), + "frozen_reference_file": (_optional_file_provenance(bundle["reference_path"]) + if role == "controlled" else None), + "target_states": _generated_array_identity(data["target_states"]), + "builder_config": _generated_config_identity(data["config"]), + }, + "model": _identity(bundle["model_path"]) if role == "controlled" else None, + "normalizer": (_identity(bundle["norm_path"]) if bundle["norm_path"] is not None else None), + "normalization_source": ("frozen_policy" if bundle["norm_path"] is not None else + ("builder_recomputed" if role != "target" else "not_applicable")), + "scratch_cleanup": {"complete": not scratch.exists(), "path": str(scratch)}, + "resolved_optane_path": str(storage["resolved_output_root"]), "phase_variables": list(phase_values), + **phase_meta}) + + +def _validate_staged_role(role_dir: Path, case_name: str = CASE_NAME) -> None: + steady = get_case(case_name).scene == "steady" + expected = ({"timeseries.npz", "timeseries.csv", "late_field.npz", "metadata.json"} if steady else + {"timeseries.npz", "timeseries.csv", "phase_cycle.npz", "phase_cycle.csv", "phase_fields.npz", "dtw_summary.json", "metadata.json"}) + actual = {path.name for path in role_dir.iterdir()} + if actual != expected: raise ValueError(f"staged role files mismatch: expected {sorted(expected)}, got {sorted(actual)}") + field_name = "late_field.npz" if steady else "phase_fields.npz" + with np.load(role_dir / field_name, allow_pickle=False) as saved: + count = 1 if steady else 8 + if saved["ux"].shape[0] != count or saved["uy"].shape[0] != count: raise ValueError("field count contract failed") + if not steady: + if (set(saved.files) != PHASE_FIELD_KEYS + or saved["mean_ux"].shape != saved["ux"].shape[1:] + or saved["mean_uy"].shape != saved["uy"].shape[1:]): + raise ValueError("phase fields exact schema failed") + with (role_dir / "metadata.json").open(encoding="utf-8") as stream: json.load(stream) + + +def _acquire_standard(role, *, case_name=CASE_NAME, output_root=None, overwrite=False, device_id=0, + repo_root=None, storage_validator=validate_output_storage, + runtime_factory=None, finalizer=_finalize, **unused): + case = get_case(case_name) + if role not in _case_roles(case): raise ValueError(f"unknown acquisition role {role!r} for {case_name}") + repo = Path(repo_root) if repo_root is not None else Path(__file__).resolve().parents[3] + storage = storage_validator(repo_mapping=default_reproduction_mapping(repo), + output_root=output_root, min_free_bytes=MIN_FREE_BYTES) + if case_name == CASE_NAME and role == "controlled": + bundle = _resolve_bundle(repo) + else: + bundle = _resolve_bundle(repo, case_name, role) + prepared = prepare_role_output(storage, "legacy", case_name, role, overwrite=overwrite) + scratch = create_scratch(prepared["scratch_root"]) + ff = None + try: + if runtime_factory is None: + data, model, ff = _default_runtime(bundle["case"], bundle, device_id, role) + else: + data, model, ff = runtime_factory(bundle["case"], bundle, device_id, role) + policy_norm = (None if role == "target" else + (data["norm"] if bundle["case"].scene == "steady" else bundle["norm"])) + rows, fields = _collect_role(role, data, model, ff, policy_norm, scratch, case_name=case_name) + finalizer(prepared["role_dir"], scratch, rows, fields, data, bundle, storage, role) + _validate_staged_role(prepared["role_dir"], case_name) + return publish_role_output(prepared) + except Exception: + if prepared["staging_dir"].exists(): + shutil.rmtree(prepared["staging_dir"]) + raise + finally: + if ff is not None and hasattr(ff, "close"): + ff.close() + + + +def _vortex_offset(case_name, value): + if type(value) is not int: + raise ValueError("vortex y offset must be an integer L0 multiple") + allowed = {0} if case_name == "vortex_lamb" else {-2, -1, 0, 1, 2} + if value not in allowed: + raise ValueError(f"{case_name} vortex y offset must be one of {sorted(allowed)}") + return value + +def _vortex_scenario(case_name, offset): + _vortex_offset(case_name, offset) + label = {0:"y000", -2:"ym2L", -1:"ym1L", 1:"yp1L", 2:"yp2L"}[offset] + return f"{case_name}_{label}" + +def _vortex_event_channel(case_name): + if case_name == "vortex_lamb": + return 1, "upper_sensor_uy" + if case_name == "vortex_taylor": + return 3, "center_sensor_uy" + raise ValueError(f"unknown Vortex case {case_name!r}") + + +def _select_vortex_event(sensors, case_name): + sensors = np.asarray(sensors, np.float64) + if sensors.ndim != 2 or sensors.shape[1:] != (6,) or len(sensors) != 150 or not np.all(np.isfinite(sensors)): + raise ValueError("Vortex event selection requires finite sensors shape (150,6)") + channel, channel_name = _vortex_event_channel(case_name) + uy = sensors[:, channel] + baseline = float(np.median(uy[:20])) + score = np.abs(uy - baseline) + lower, upper = 38, 87 + peak = lower + int(np.argmax(score[lower:upper + 1])) + indices = peak + np.asarray([-10, -5, 0, 5, 10], np.int64) + return {"baseline_uy": baseline, "event_score": score, + "selected_event_score": score[indices], "peak_index": np.int64(peak), + "field_indices": indices, + "relative_offsets": np.asarray([-10, -5, 0, 5, 10], np.int64), + "search_lower_boundary": lower, "search_upper_boundary": upper, + "event_channel_index": channel, "event_channel": channel_name} + + +def _canonical_vortex_source(case_name): + return "vortex_lamb_y000" if case_name == "vortex_lamb" else "vortex_taylor_y000" + + +def _load_vortex_event_source(storage, case_name): + source = Path(storage["resolved_output_root"]) / "legacy" / _canonical_vortex_source(case_name) / "controlled" / "event_summary.json" + if not source.is_file(): + raise FileNotFoundError(f"canonical Vortex event source is required before this run: {source}") + try: + with source.open(encoding="utf-8") as stream: + summary = json.load(stream) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"malformed canonical Vortex event source {source}") from exc + channel, channel_name = _vortex_event_channel(case_name) + expected_source = {"scenario": _canonical_vortex_source(case_name), "role": "controlled", + "channel_index": channel, "channel": channel_name} + required = {"schema", "event_source", "common_peak_index", "field_indices", "relative_offsets"} + if not isinstance(summary, dict) or not required.issubset(summary): + raise ValueError(f"malformed canonical Vortex event source {source}") + try: + peak = int(summary["common_peak_index"]) + indices = np.asarray(summary["field_indices"], np.int64) + offsets = np.asarray(summary["relative_offsets"], np.int64) + except (TypeError, ValueError) as exc: + raise ValueError(f"malformed canonical Vortex event source {source}") from exc + if (summary["schema"] != "drl-pinball-legacy-vortex-event-v2" or + summary["event_source"] != expected_source or indices.shape != (5,) or + not np.array_equal(offsets, [-10, -5, 0, 5, 10]) or + not np.array_equal(indices, peak + offsets) or peak < 38 or peak > 87): + raise ValueError(f"malformed canonical Vortex event source {source}") + return {"source_path": source, "event_source": expected_source, + "peak_index": peak, "field_indices": indices, "relative_offsets": offsets} + +def _erase_phase(times, center_uy, disturbance_force_y): + try: + phase = _assign_scalar_phase(times, center_uy) + if len(phase["crossing_times"]) >= 4: + return phase, "center_sensor_uy", False + except ValueError as exc: + if "rising crossings" not in str(exc): raise + phase = _assign_scalar_phase(times, disturbance_force_y) + if len(phase["crossing_times"]) < 4: + raise ValueError("Erase publication requires at least 4 rising crossings in fixed disturbance-force fallback") + return phase, "disturbance_force_y", True + +def _assign_scalar_phase(times, probe): + probe=np.asarray(probe,np.float64); times=np.asarray(times,np.float64) + if probe.shape != times.shape or probe.ndim != 1 or not np.all(np.isfinite(probe)): + raise ValueError("phase probe must be one-dimensional, aligned, and finite") + synthetic=np.zeros((len(times),4),np.float64); synthetic[:,3]=probe + return assign_phase(times,synthetic) + +def _special_observation(case, role, raw): + raw=np.asarray(raw,np.float32) + if case.scene == "erase": + if role == "target": return raw[:6], np.full(2,np.nan,np.float32), np.full(6,np.nan,np.float32), raw[:6] + if raw.shape != (14,): raise ValueError("Erase observation must contain exactly 14 channels") + return raw[:6], raw[6:8], raw[8:14], raw + if raw.shape != (6 if role == "target" else 12,): raise ValueError("Vortex observation inventory changed") + return raw[:6], np.full(2,np.nan,np.float32), (np.full(6,np.nan,np.float32) if role=="target" else raw[6:12]), raw + +def _special_schedule(case, target): + if case.scene == "erase": + return (0 if target else 480), 160 + return 0, 150 + + +def _special_timeline_metadata(case, target): + warmup, collected = _special_schedule(case, target) + return {"warmup_control_steps": warmup, + "collected_post_step_boundaries": collected, + "historical_episode_steps": 150 if case.scene == "vortex" else None} + + +def _collect_special(role,data,model,ff,norm,scratch,case,capture_field=_capture_boundary_field): + target = role == "target"; expected = 3 if target else (7 if case.scene=="erase" else 6) + if int(data["config"]["n_obj_total"]) != expected: raise ValueError("special scene object inventory changed") + if target: fifo=None; obs=None + else: + saved=np.asarray(data["norm"]["save_states"],np.float32) + fifo,obs=reset_runtime(ff,saved,14 if case.scene=="erase" else 12) + if case.scene == "erase": + if saved.ndim != 2 or saved.shape[1] != 14: + raise ValueError("Erase reset FIFO must retain exactly 14 raw channels") + obs=np.zeros(12,dtype=np.float32) + warmup,collect=_special_schedule(case,target) + rows=[]; fields=[] + for interval in range(warmup+collect): + if role=="controlled": action,_=model.predict(obs,deterministic=True); action,command=_command(case,action,expected,ff.DATA_TYPE) + elif role=="zero": + action=-np.asarray(case.action_bias,np.float32)/np.float32(case.action_scale); action,command=_command(case,action,expected,ff.DATA_TYPE); command[-3:]=0 + else: action=np.full(3,np.nan,np.float32); command=np.zeros(expected,ff.DATA_TYPE) + _run_interval(ff,case.sample_interval,command); sensors,disturbance,pinball,raw=_special_observation(case,role,np.asarray(ff.obs,np.float32).copy()) + if target: terms={}; effective=np.full(3,np.nan,np.float32) + else: + fifo.append(raw.copy()) + if case.scene=="erase": terms=erase_reward_terms(np.asarray(fifo),np.mean(data["target_states"],axis=0),norm,case.conv_len); obs=erase_policy_observation(raw,norm) + else: terms=reward_terms(case,data["target_states"],None,np.asarray(fifo),norm["force_norm_fact"],interval); obs=policy_observation(raw,norm) + effective=np.asarray(ff.current_effective_action(),np.float32)[-3:].copy() + if interval < warmup: continue + lattice,clock,absolute=_clock(ff,interval+1,case.sample_interval) + row={"physical_time":lattice*U0/NX,"lattice_step":lattice,"control_index":clock,"sensors":sensors.copy(), + "disturbance_force":disturbance.copy(),"pinball_forces":pinball.copy(),"action_normalized":action.copy(), + "commanded_target_omega":np.full(3,np.nan,np.float32) if target else command[-3:].astype(np.float32), + "effective_smoothed_omega":effective,"reward_raw":float(terms.get("reward",np.nan)), + "reward_u":float(terms.get("reward_u",np.nan)),"reward_v":float(terms.get("reward_v",np.nan)), + "native_reward_dtw":float(terms.get("native_legacy_dtw",np.nan)),"native_lag":float(terms.get("native_lag",np.nan)), + "native_component_similarity":np.asarray(terms.get("native_component_similarity",[np.nan,np.nan]))} + if absolute is not None: row["solver_absolute_lattice_step"]=absolute + rows.append(row); value=capture_field(ff,int(data["config"]["nx"]),int(data["config"]["ny"])); path=scratch/f"boundary_{len(rows)-1:04d}.npz"; np.savez(path,**value); fields.append(path) + return rows,fields + +def _publish_fields(path,field_paths,selection): + result={"ux":[],"uy":[]}; indices=np.asarray(selection["field_indices"],np.int64) + for index in indices: + with np.load(field_paths[int(index)],allow_pickle=False) as item: + result["ux"].append(item["ux"]); result["uy"].append(item["uy"]) + result={key:np.stack(value) for key,value in result.items()}; result.update(selection); publish_selected_fields(path,result) + +def _special_normalization_source(case, role): + if role == "target": + return "not_applicable" + if case.scene == "erase": + return "builder_recomputed_historical" + if case.scene == "vortex" and role == "zero": + return "builder_recomputed_zero_runtime" + return "frozen_policy" + + +def _special_normalizer_file(bundle): + return _identity(bundle["norm_path"]) if bundle["norm_path"] is not None else None + + +def _finalize_special(role_dir,scratch,rows,field_paths,data,bundle,storage,role,event_source=None): + case=bundle["case"]; columns=_columns(rows); target=role=="target"; event_peak_index=None; event_diagnostics=None + if case.scene=="vortex": + channel, channel_name = _vortex_event_channel(case.name) + own_uy = np.asarray(columns["sensors"][:, channel], np.float64) + own_baseline = float(np.median(own_uy[:20])) + own_score = np.abs(own_uy - own_baseline) + columns["event_score"] = own_score + if event_source is None: + local = _select_vortex_event(columns["sensors"], case.name) + event_source = {"event_source": {"scenario": _canonical_vortex_source(case.name), + "role": "controlled", "channel_index": channel, "channel": channel_name}, + "peak_index": int(local["peak_index"]), + "field_indices": local["field_indices"], + "relative_offsets": local["relative_offsets"]} + indices = np.asarray(event_source["field_indices"], np.int64) + offsets = np.asarray(event_source["relative_offsets"], np.int64) + event_selection = {"field_indices": indices, "relative_offsets": offsets, + "selected_event_score": own_score[indices]} + write_boundary_artifacts(role_dir/"timeseries.npz",role_dir/"timeseries.csv",columns) + _publish_fields(role_dir/"event_fields.npz",field_paths,event_selection) + event_peak_index = int(event_source["peak_index"]) + event_diagnostics = {"search_lower_boundary": 38, "search_upper_boundary": 87, + "diagnostic_channel_index": channel, "diagnostic_channel": channel_name} + _atomic_json(role_dir/"event_summary.json", {"schema":"drl-pinball-legacy-vortex-event-v2", + "definition":"canonical controlled fixed-channel absolute deviation from median first 20; earliest argmax in inclusive boundaries 38..87", + "event_source":event_source["event_source"], "common_peak_index":event_peak_index, + "field_indices":indices.tolist(), "relative_offsets":offsets.tolist(), + "selected_event_score":own_score[indices].tolist(), "diagnostic_baseline_uy":own_baseline, + "diagnostic_peak_score":float(own_score[event_peak_index]), **event_diagnostics}) + field_count=5; nonperiodic=True; probe=None + elif target: + write_boundary_artifacts(role_dir/"timeseries.npz",role_dir/"timeseries.csv",columns); _publish_fields(role_dir/"late_field.npz",field_paths,{"field_indices":np.asarray([len(rows)-1],np.int64)}); field_count=1; nonperiodic=True; probe=None + err=columns["sensors"]-np.mean(data["target_states"],axis=0); _atomic_json(role_dir/"erase_summary.json",{"native_metrics":"unavailable for clean target","clean_sensor_rmse":float(np.sqrt(np.mean(err*err))),"phase_probe":None,"nonperiodic_clean_reference":True}) + else: + phase,probe,fallback=_erase_phase(columns["physical_time"],columns["sensors"][:,3],columns["disturbance_force"][:,1]) + columns.update(phase=phase["phase"],cycle_id=phase["cycle_id"]); write_boundary_artifacts(role_dir/"timeseries.npz",role_dir/"timeseries.csv",columns) + values=_phase_values(columns,role); write_phase_cycle_artifacts(role_dir/"phase_cycle.npz",role_dir/"phase_cycle.csv",pooled_phase_bins(phase["phase"],values,bins=32)); _publish_fields(role_dir/"phase_fields.npz",field_paths,select_phase_fields(columns["physical_time"],phase["crossing_times"])); field_count=8; nonperiodic=False + err=columns["sensors"]-np.mean(data["target_states"],axis=0); _atomic_json(role_dir/"erase_summary.json",{"native_reward_mean":float(np.mean(columns["reward_raw"])),"reward_u_mean":float(np.mean(columns["reward_u"])),"reward_v_mean":float(np.mean(columns["reward_v"])),"native_reward_dtw_mean":float(np.mean(columns["native_reward_dtw"])),"clean_sensor_rmse":float(np.sqrt(np.mean(err*err))),"phase_probe":probe,"phase_fallback_used":fallback,"force_balance_diagnostic":"native enhanced formula is scale invariant; normalized diagnostic equals native by definition"}) + cleanup_scratch(scratch,root=role_dir/"scratch"); scratch.parent.rmdir() + timeline=_special_timeline_metadata(case,target) + _atomic_json(role_dir/"metadata.json",{"schema":"drl-pinball-legacy-nonstandard-acquisition-v1","case_id":case.name,"scene":case.scene,"role":role,**timeline,"field_count":field_count,"nonperiodic":nonperiodic,"phase_probe":probe,"event_peak_index":event_peak_index,"event_source":(event_source["event_source"] if event_source else None),"event_diagnostics":event_diagnostics,"normalization_source":_special_normalization_source(case,role),"role_semantics":(("physical-zero/uncontrolled trajectory using counter-bias normalized action; builder-recomputed historical norm" if role=="zero" else "frozen-policy controlled trajectory using builder-recomputed historical norm") if case.scene=="erase" and not target else ("physical-zero/uncontrolled trajectory using counter-bias normalized action; builder-recomputed zero-runtime norm" if case.scene=="vortex" and role=="zero" else _role_semantics(role))),"field_contract":{"quantity":"q/RHO_ref","RHO_ref":1},"source_provenance":{"model_file":_identity(bundle["model_path"]) if role=="controlled" else None,"normalizer_file":_special_normalizer_file(bundle),"target_states":_generated_array_identity(data["target_states"]),"builder_config":_generated_config_identity(data["config"])},"resolved_optane_path":str(storage["resolved_output_root"])}) + +def _validate_special(role_dir,case,role): + if case.scene=="vortex": expected={"timeseries.npz","timeseries.csv","event_fields.npz","event_summary.json","metadata.json"}; field="event_fields.npz"; count=5 + elif role=="target": expected={"timeseries.npz","timeseries.csv","late_field.npz","erase_summary.json","metadata.json"}; field="late_field.npz"; count=1 + else: expected={"timeseries.npz","timeseries.csv","phase_cycle.npz","phase_cycle.csv","phase_fields.npz","erase_summary.json","metadata.json"}; field="phase_fields.npz"; count=8 + actual={x.name for x in role_dir.iterdir()} + if actual!=expected: raise ValueError(f"staged role files mismatch: expected {sorted(expected)}, got {sorted(actual)}") + with np.load(role_dir/field,allow_pickle=False) as saved: + expected_keys = ({"ux", "uy", "field_indices", "relative_offsets", "selected_event_score"} + if case.scene == "vortex" else None) + if expected_keys is not None and set(saved.files) != expected_keys: + raise ValueError(f"Vortex event field keys mismatch: expected {sorted(expected_keys)}, got {sorted(saved.files)}") + if any(np.asarray(saved[name]).ndim < 1 or np.asarray(saved[name]).shape[0] != count + for name in saved.files): + raise ValueError("special field arrays must be non-scalar with exact field count") + if case.scene == "vortex": + field_indices = np.asarray(saved["field_indices"], np.int64) + relative_offsets = np.asarray(saved["relative_offsets"], np.int64) + if case.scene == "vortex": + try: + with (role_dir/"event_summary.json").open(encoding="utf-8") as stream: + summary = json.load(stream) + channel, channel_name = _vortex_event_channel(case.name) + expected_source = {"scenario": _canonical_vortex_source(case.name), "role": "controlled", + "channel_index": channel, "channel": channel_name} + peak = int(summary["common_peak_index"]) + if (summary["schema"] != "drl-pinball-legacy-vortex-event-v2" or + summary["event_source"] != expected_source or + not np.array_equal(field_indices, peak + relative_offsets) or + not np.array_equal(field_indices, np.asarray(summary["field_indices"], np.int64)) or + not np.array_equal(relative_offsets, [-10, -5, 0, 5, 10])): + raise ValueError + except (OSError, KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: + raise ValueError("Vortex event summary/source contract failed") from exc + +def _acquire_special(role,case_name,output_root,overwrite,device_id,repo_root,storage_validator,runtime_factory,vortex_y_offset_l0): + case=get_case(case_name); repo=Path(repo_root) if repo_root else Path(__file__).resolve().parents[3]; offset=_vortex_offset(case_name,vortex_y_offset_l0) if case.scene=="vortex" else 0 + storage=storage_validator(repo_mapping=default_reproduction_mapping(repo),output_root=output_root,min_free_bytes=MIN_FREE_BYTES); bundle=_resolve_bundle(repo,case_name,role) + scenario=_vortex_scenario(case_name,offset) if case.scene=="vortex" else case_name + source_run = case.scene == "vortex" and role == "controlled" and offset == 0 + event_source = None if not case.scene == "vortex" or source_run else _load_vortex_event_source(storage, case_name) + prepared=prepare_role_output(storage,"legacy",scenario,role,overwrite=overwrite); scratch=create_scratch(prepared["scratch_root"]); ff=None + try: + if runtime_factory is not None: data,model,ff=runtime_factory(case,bundle,device_id,role) + else: + from .core.legacy_env_builder import build_erase,build_vortex + data=build_vortex(device_id,case.re_code,vortex_type=case.vortex_type,action_scale=case.action_scale,y_offset_l0=offset,target_only=role=="target",role=role) if case.scene=="vortex" else build_erase(device_id,case.re_code,target_only=role=="target") + model=None + if role=="controlled": + from .core.model_loader import load_model + model=load_model(case.model) + ff=data["flow_field"] + norm=None if role=="target" else (data["norm"] if case.scene=="erase" or (case.scene=="vortex" and role=="zero") else bundle["norm"]); rows,fields=_collect_special(role,data,model,ff,norm,scratch,case); _finalize_special(prepared["role_dir"],scratch,rows,fields,data,bundle,storage,role,event_source); _validate_special(prepared["role_dir"],case,role); return publish_role_output(prepared) + except Exception: + if prepared["staging_dir"].exists(): shutil.rmtree(prepared["staging_dir"]) + raise + finally: + if ff is not None and hasattr(ff,"close"): ff.close() + +def acquire_role(role, *, case_name=CASE_NAME, output_root=None, overwrite=False, device_id=0, repo_root=None, storage_validator=validate_output_storage, runtime_factory=None, finalizer=_finalize, vortex_y_offset_l0=0): + case=get_case(case_name) + if role not in _case_roles(case): raise ValueError(f"unknown acquisition role {role!r} for {case_name}") + if case.scene in ("vortex","erase"): + return _acquire_special(role,case_name,output_root,overwrite,device_id,repo_root,storage_validator,runtime_factory,vortex_y_offset_l0) + return _acquire_standard(role,case_name=case_name,output_root=output_root,overwrite=overwrite,device_id=device_id,repo_root=repo_root,storage_validator=storage_validator,runtime_factory=runtime_factory,finalizer=finalizer) + +def acquire_controlled(**kwargs): + return acquire_role("controlled", **kwargs) + + +def main() -> int: + parser = argparse.ArgumentParser(description="Acquire Legacy full-matrix evidence") + parser.add_argument("case", choices=SUPPORTED_CASES) + parser.add_argument("--role", choices=ROLES + ("constant",), required=True) + parser.add_argument("--device-id", type=int, default=0) + parser.add_argument("--vortex-y-offset-l0", type=int, default=0) + parser.add_argument("--output-root", type=Path); parser.add_argument("--overwrite", action="store_true") + args = parser.parse_args() + path = acquire_role(args.role, case_name=args.case, output_root=args.output_root, overwrite=args.overwrite, device_id=args.device_id, vortex_y_offset_l0=args.vortex_y_offset_l0) + print(path); return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/drl_pinball/legacy_test/cases.py b/src/drl_pinball/legacy_test/cases.py new file mode 100644 index 0000000..29eb89c --- /dev/null +++ b/src/drl_pinball/legacy_test/cases.py @@ -0,0 +1,36 @@ +"""Canonical case matrix for LegacyCelerisLab reproduction.""" +from __future__ import annotations +from dataclasses import dataclass + +@dataclass(frozen=True) +class Case: + name: str + scene: str + model: str | None + re_code: float = 100.0 + sample_interval: int = 800 + steps: int = 200 + action_scale: float = 8.0 + action_bias: tuple[float, float, float] = (0.0, -4.0, 4.0) + conv_len: int = 30 + target_radius_l: float | None = None + vortex_type: str | None = None + reference_case: str | None = None + +CASES = { + **{f"karman_re{re}": Case(f"karman_re{re}", "karman", f"d1a3o12_re{re}", re_code=float(re)) for re in (50, 100, 200, 400)}, + "illusion_075L": Case("illusion_075L", "illusion", "d1a3o14_250525_imit_075L_2U_400S", sample_interval=400, action_bias=(0.0, -2.0, 2.0), conv_len=36, target_radius_l=0.75, reference_case="illusion_0.75L"), + "illusion_1L": Case("illusion_1L", "illusion", "d1a3o14_250525_imit_1L_2U_600S", sample_interval=600, action_bias=(0.0, -2.0, 2.0), conv_len=36, target_radius_l=1.0), + "illusion_15L": Case("illusion_15L", "illusion", "d1a3o14_250525_imit_15L_2U", action_bias=(0.0, -2.0, 2.0), conv_len=36, target_radius_l=1.5, reference_case="illusion_1.5L"), + "vortex_lamb": Case("vortex_lamb", "vortex", "vortex_lamb", steps=150, action_scale=4.0, vortex_type="lamb"), + "vortex_taylor": Case("vortex_taylor", "vortex", "vortex_taylor", steps=150, action_scale=4.0, vortex_type="taylor"), + "erase": Case("erase", "erase", "d1a3o12_250729_250326_erase", sample_interval=600, action_scale=8.0, action_bias=(0.0, -8.0, 8.0), conv_len=36), + "steady": Case("steady", "steady", None, action_bias=(0.0, -5.1, 5.1)), +} +ERASE_UNSUPPORTED = "" # compatibility name; Erase is acquisition-only and supported. + +def get_case(name: str) -> Case: + try: + return CASES[name] + except KeyError as exc: + raise KeyError(f"unknown case {name!r}; choose from {sorted(CASES)} or erase") from exc diff --git a/src/drl_pinball/legacy_test/core/legacy_env_builder.py b/src/drl_pinball/legacy_test/core/legacy_env_builder.py index 77725b7..a6c3ff2 100644 --- a/src/drl_pinball/legacy_test/core/legacy_env_builder.py +++ b/src/drl_pinball/legacy_test/core/legacy_env_builder.py @@ -35,6 +35,7 @@ if _REPO not in sys.path: from LegacyCelerisLab import FlowField # noqa: E402 from LegacyCelerisLab import utils as legacy_utils # noqa: E402 +from ..runtime import run_historical_interval # noqa: E402 # --------------------------------------------------------------------------- # Constants @@ -58,9 +59,20 @@ def _center_y(ff: FlowField) -> float: return (ff.FIELD_SHAPE[1] - 1) / 2.0 +def _field_xy(ff: FlowField) -> Tuple[int, int]: + """Return Legacy's 2-D lattice dimensions from its 3-D shape contract.""" + shape = tuple(ff.FIELD_SHAPE) + if len(shape) != 3: + raise ValueError(f"Legacy FlowField.FIELD_SHAPE must be (nx, ny, nz), got {shape}") + nx, ny, nz = shape + if nz != 1: + raise ValueError(f"Legacy builder requires nz=1, got FIELD_SHAPE={shape}") + return int(nx), int(ny) + + def _stabilize(ff: FlowField, n_obj: int) -> None: steps = int(4 * ff.FIELD_SHAPE[0] / U0) - ff.run(steps, np.zeros(n_obj, dtype=DATA_TYPE)) + run_historical_interval(ff, steps, np.zeros(n_obj, dtype=DATA_TYPE)) def _compute_karman_norm(fifo: np.ndarray) -> Dict[str, Any]: @@ -82,6 +94,37 @@ def _compute_karman_norm(fifo: np.ndarray) -> Dict[str, Any]: # Karman Cloak (dist-cyl + 3 sensors + 3 pinball = 7 objects) # --------------------------------------------------------------------------- +def build_karman_target( + device_id: int = 0, + re_code: float = 100.0, + *, + sample_interval: int = 800, +) -> Dict[str, Any]: + """Build only the canonical disturbance-cylinder Karman target stage.""" + viscosity = _nu_from_re(re_code) + cuda_cfg = legacy_utils.load_cuda_config(os.path.join(CONFIG_DIR, "config_cuda.json")) + field_cfg = legacy_utils.load_flow_field_config(os.path.join(CONFIG_DIR, "config_flowfield.json")) + field_cfg = field_cfg._replace(viscosity=float(viscosity)) + ff = FlowField(field_cfg, cuda_cfg, device_id=device_id) + cy = _center_y(ff) + nx, ny = _field_xy(ff) + ff.add_cylinder((10.0 * L0, cy, 0.0), 1.0 * L0) + for y_off in (2.0, 0.0, -2.0): + ff.add_sensor((40.0 * L0, cy + y_off * L0, 0.0), SENSOR_RADIUS) + if ff.obs.size // 2 != 4: + raise ValueError("canonical Karman target requires disturbance plus three sensors") + _stabilize(ff, 4) + target = np.empty((FIFO_LEN, 6), dtype=DATA_TYPE) + for index in range(FIFO_LEN): + run_historical_interval(ff, sample_interval, np.zeros(4, dtype=DATA_TYPE)) + target[index] = ff.obs.copy()[2:8] + config = {"device_id": device_id, "viscosity": viscosity, "re_code": re_code, + "u0": U0, "sample_interval": sample_interval, "fifo_len": FIFO_LEN, + "conv_len": CONV_LEN, "nx": nx, "ny": ny, "n_obj_total": 4, + "obs_slice": (2, 8), "s_dim": 6} + return {"flow_field": ff, "target_states": target, "config": config} + + def build_karman_cloak( device_id: int = 0, re_code: float = 100.0, @@ -119,7 +162,7 @@ def build_karman_cloak( # Record target (sensor signals only, no pinball) target = np.empty((0, 6), dtype=DATA_TYPE) for _ in range(FIFO_LEN): - ff.run(sample_interval, np.zeros(4, dtype=DATA_TYPE)) + run_historical_interval(ff, sample_interval, np.zeros(4, dtype=DATA_TYPE)) target = np.vstack((target, ff.obs.copy()[2:8])) # Phase 2: Add pinball @@ -139,7 +182,7 @@ def build_karman_cloak( # Zero-action norm collection fifo = deque(maxlen=FIFO_LEN) for _ in range(FIFO_LEN): - ff.run(sample_interval, np.zeros(7, dtype=DATA_TYPE)) + run_historical_interval(ff, sample_interval, np.zeros(7, dtype=DATA_TYPE)) fifo.append(ff.obs.copy()[2:14]) norm = _compute_karman_norm(np.array(fifo, dtype=DATA_TYPE)) @@ -152,7 +195,7 @@ def build_karman_cloak( fifo.clear() for _ in range(FIFO_LEN): - ff.run(sample_interval, bias_arr) + run_historical_interval(ff, sample_interval, bias_arr) fifo.append(ff.obs.copy()[2:14]) save_states = np.array(list(fifo), dtype=DATA_TYPE) ff.apply_ddf() @@ -181,6 +224,29 @@ def build_karman_cloak( return {"flow_field": ff, "target_states": target, "norm": norm, "config": config} +def build_steady_target( + device_id: int = 0, re_code: float = 100.0, *, sample_interval: int = 800, +) -> Dict[str, Any]: + """Build the clean-channel three-sensor steady target.""" + viscosity = _nu_from_re(re_code) + cuda_cfg = legacy_utils.load_cuda_config(os.path.join(CONFIG_DIR, "config_cuda.json")) + field_cfg = legacy_utils.load_flow_field_config(os.path.join(CONFIG_DIR, "config_flowfield.json"))._replace(viscosity=float(viscosity)) + ff = FlowField(field_cfg, cuda_cfg, device_id=device_id) + cy = _center_y(ff); nx, ny = _field_xy(ff) + for y_off in (2.0, 0.0, -2.0): + ff.add_sensor((40.0 * L0, cy + y_off * L0, 0.0), SENSOR_RADIUS) + _stabilize(ff, 3) + target = np.empty((FIFO_LEN, 6), dtype=DATA_TYPE) + for index in range(FIFO_LEN): + run_historical_interval(ff, sample_interval, np.zeros(3, dtype=DATA_TYPE)) + target[index] = ff.obs.copy()[:6] + config = {"device_id": device_id, "viscosity": viscosity, "re_code": re_code, + "u0": U0, "sample_interval": sample_interval, "fifo_len": FIFO_LEN, + "conv_len": CONV_LEN, "nx": nx, "ny": ny, "n_obj_total": 3, + "obs_slice": (0, 6), "s_dim": 6} + return {"flow_field": ff, "target_states": target, "config": config} + + # --------------------------------------------------------------------------- # Steady Cloak (3 sensors + 3 pinball = 6 objects, no dist-cyl) # --------------------------------------------------------------------------- @@ -230,7 +296,7 @@ def build_steady_cloak( target = np.empty((0, 6), dtype=DATA_TYPE) for _ in range(FIFO_LEN): - ff2.run(800, np.zeros(3, dtype=DATA_TYPE)) + run_historical_interval(ff2, 800, np.zeros(3, dtype=DATA_TYPE)) target = np.vstack((target, ff2.obs.copy()[0:6])) del ff2 @@ -241,7 +307,7 @@ def build_steady_cloak( # Norm fifo = deque(maxlen=FIFO_LEN) for _ in range(FIFO_LEN): - ff.run(800, np.zeros(6, dtype=DATA_TYPE)) + run_historical_interval(ff, 800, np.zeros(6, dtype=DATA_TYPE)) fifo.append(ff.obs.copy()[0:12]) temp = np.array(fifo, dtype=DATA_TYPE) force_norm_fact = 6.0 * float(np.max(np.abs(temp[:, 6:12]))) @@ -259,7 +325,7 @@ def build_steady_cloak( fifo.clear() for _ in range(FIFO_LEN): - ff.run(800, bias_arr) + run_historical_interval(ff, 800, bias_arr) fifo.append(ff.obs.copy()[0:12]) save_states = np.array(list(fifo), dtype=DATA_TYPE) ff.apply_ddf() @@ -293,6 +359,33 @@ def build_steady_cloak( return {"flow_field": ff, "target_states": target, "norm": norm, "config": config} +def build_illusion_target( + device_id: int = 0, re_code: float = 100.0, *, target_diameter_L: float = 1.0, + sample_interval: int = 600, +) -> Dict[str, Any]: + """Build the target-cylinder stage and retain its public flow field.""" + viscosity = _nu_from_re(re_code) + cuda_cfg = legacy_utils.load_cuda_config(os.path.join(CONFIG_DIR, "config_cuda.json")) + field_cfg = legacy_utils.load_flow_field_config(os.path.join(CONFIG_DIR, "config_flowfield.json"))._replace(viscosity=float(viscosity)) + ff = FlowField(field_cfg, cuda_cfg, device_id=device_id) + cy = _center_y(ff); nx, ny = _field_xy(ff) + ff.add_cylinder((20.0 * L0, cy, 0.0), target_diameter_L * L0) + for y_off in (2.0, 0.0, -2.0): + ff.add_sensor((30.0 * L0, cy + y_off * L0, 0.0), SENSOR_RADIUS) + _stabilize(ff, 4) + target = np.empty((FIFO_LEN, 8), dtype=DATA_TYPE) + for index in range(FIFO_LEN): + run_historical_interval(ff, sample_interval, np.zeros(4, dtype=DATA_TYPE)) + target[index] = ff.obs.copy()[:8] + from .dtw_metrics import analyze_harmonics + harmonics = analyze_harmonics(target, n_harmonics=5) + config = {"device_id": device_id, "viscosity": viscosity, "re_code": re_code, + "u0": U0, "sample_interval": sample_interval, "fifo_len": FIFO_LEN, + "conv_len": 36, "nx": nx, "ny": ny, "n_obj_total": 4, + "obs_slice": (0, 8), "s_dim": 8, "target_diameter_L": target_diameter_L} + return {"flow_field": ff, "target_states": target, "target_harmonics": harmonics, "config": config} + + # --------------------------------------------------------------------------- # Illusion (target cylinder + 3 sensors at illusion positions, then pinball + sensors) # --------------------------------------------------------------------------- @@ -333,7 +426,7 @@ def build_illusion( # Record target (8 channels: cyl_force[2] + sensors[6]) target_states = np.empty((0, 8), dtype=DATA_TYPE) for _ in range(FIFO_LEN): - ff_target.run(sample_interval, np.zeros(4, dtype=DATA_TYPE)) + run_historical_interval(ff_target, sample_interval, np.zeros(4, dtype=DATA_TYPE)) target_states = np.vstack((target_states, ff_target.obs.copy()[0:8])) # Harmonics analysis (FFT) @@ -364,7 +457,7 @@ def build_illusion( # Norm fifo = deque(maxlen=FIFO_LEN) for _ in range(FIFO_LEN): - ff.run(sample_interval, np.zeros(6, dtype=DATA_TYPE)) + run_historical_interval(ff, sample_interval, np.zeros(6, dtype=DATA_TYPE)) fifo.append(ff.obs.copy()[0:12]) norm = _compute_karman_norm(np.array(fifo, dtype=DATA_TYPE)) @@ -378,10 +471,12 @@ def build_illusion( fifo.clear() for _ in range(FIFO_LEN): - ff.run(sample_interval, bias_arr) + run_historical_interval(ff, sample_interval, bias_arr) fifo.append(ff.obs.copy()[0:12]) save_states = np.array(list(fifo), dtype=DATA_TYPE) - ff.apply_ddf() + # Original Illusion CustomEnv overwrites the checkpoint after bias FIFO. + ff.get_ddf() + ff.save_ddf() norm["save_states"] = save_states norm["action_bias"] = list(action_bias) @@ -419,128 +514,73 @@ def build_illusion( # --------------------------------------------------------------------------- def build_vortex( - device_id: int = 0, - re_code: float = 100.0, - *, - vortex_type: str = "lamb", - action_scale: float = 4.0, - action_bias: Tuple[float, float, float] = (0.0, -4.0, 4.0), + device_id: int = 0, re_code: float = 100.0, *, vortex_type: str = "lamb", + action_scale: float = 4.0, action_bias: Tuple[float, float, float] = (0.0, -4.0, 4.0), + y_offset_l0: int = 0, target_only: bool = False, role: Optional[str] = None, ) -> Dict[str, Any]: - """Build vortex cloaking environment. - - Target phase: vortex at x=10*L0 + 3 sensors. - Pinball phase: vortex at x=15*L0 + pinball + 3 sensors. - - Object order (pinball phase): sensor0(0), sensor1(1), sensor2(2), - front(3), top(4), bottom(5). - - MAX_STEPS = 150 (transient event). - """ + """Build Vortex with an explicit role-bound runtime handoff.""" + if role is None: + role = "target" if target_only else "controlled" + if role not in ("controlled", "target", "zero"): + raise ValueError(f"invalid Vortex builder role {role!r}") + if target_only != (role == "target"): + raise ValueError("Vortex target_only and role semantics disagree") + allowed = {0} if vortex_type == "lamb" else {-2, -1, 0, 1, 2} if vortex_type == "taylor" else set() + if type(y_offset_l0) is not int or y_offset_l0 not in allowed: + raise ValueError(f"invalid {vortex_type} vortex y offset {y_offset_l0}; allowed {sorted(allowed)}") viscosity = _nu_from_re(re_code) cuda_cfg = legacy_utils.load_cuda_config(os.path.join(CONFIG_DIR, "config_cuda.json")) - field_cfg = legacy_utils.load_flow_field_config(os.path.join(CONFIG_DIR, "config_flowfield.json")) - field_cfg = field_cfg._replace(viscosity=float(viscosity)) - - vortex_strength = 0.5 * U0 if vortex_type == "lamb" else 0.03 * U0 - - # Phase 1: Sensors-only env -> record clean channel -> add vortex -> record target - ff_sensors = FlowField(field_cfg, cuda_cfg, device_id=device_id) - cy_s = _center_y(ff_sensors) - for y_off in [2.0, 0.0, -2.0]: - ff_sensors.add_sensor((40.0 * L0, cy_s + y_off * L0, 0.0), SENSOR_RADIUS) - - n_sens = ff_sensors.obs.size // 2 - assert n_sens == 3 - - _stabilize(ff_sensors, 3) - - # Record clean channel baseline - ff_sensors.get_ddf() - ff_sensors.save_ddf() - - # Add vortex at x=10*L0 and record target - ff_sensors.add_vortex( - (10.0 * L0, cy_s, 0.0), - 2.0 * L0, - vortex_strength, - 0.0, - vortex_type, - ) - - target_states = np.empty((0, 6), dtype=DATA_TYPE) - for _ in range(FIFO_LEN): - ff_sensors.run(800, np.zeros(3, dtype=DATA_TYPE)) - target_states = np.vstack((target_states, ff_sensors.obs.copy()[0:6])) - - del ff_sensors - - # Phase 2: Pinball + sensors + vortex + field_cfg = legacy_utils.load_flow_field_config(os.path.join(CONFIG_DIR, "config_flowfield.json"))._replace(viscosity=float(viscosity)) ff = FlowField(field_cfg, cuda_cfg, device_id=device_id) - cy = _center_y(ff) - NX, NY = ff.FIELD_SHAPE[0], ff.FIELD_SHAPE[1] - - for y_off in [2.0, 0.0, -2.0]: - ff.add_sensor((40.0 * L0, cy + y_off * L0, 0.0), SENSOR_RADIUS) # id=0,1,2 - ff.add_cylinder((30.0 * L0, cy, 0.0), PINBALL_RADIUS) # front, id=3 - ff.add_cylinder((31.3 * L0, cy + 0.75 * L0, 0.0), PINBALL_RADIUS) # top, id=4 - ff.add_cylinder((31.3 * L0, cy - 0.75 * L0, 0.0), PINBALL_RADIUS) # bottom, id=5 - - n_total = ff.obs.size // 2 - assert n_total == 6, f"Expected 6 objects, got {n_total}" - - _stabilize(ff, 6) - - ff.get_ddf() + cy = _center_y(ff); nx, ny = _field_xy(ff) + for y_off in (2.0, 0.0, -2.0): + ff.add_sensor((40.0 * L0, cy + y_off * L0, 0.0), SENSOR_RADIUS) + run_historical_interval(ff, int(nx / U0), np.zeros(3, dtype=DATA_TYPE)) + ff.get_ddf(); ff.save_ddf() + strength = (0.5 if vortex_type == "lamb" else 0.03) * U0 + target_vortex = (10.0 * L0, cy + y_offset_l0 * L0, 0.0) + ff.add_vortex(target_vortex, 2.0 * L0, strength, 0.0, vortex_type) + target = np.empty((0, 6), dtype=DATA_TYPE) + for _ in range(FIFO_LEN): + run_historical_interval(ff, 800, np.zeros(3, dtype=DATA_TYPE)); target = np.vstack((target, ff.obs.copy())) + target_config = {"device_id": device_id, "viscosity": viscosity, "re_code": re_code, "u0": U0, "sample_interval": 800, "fifo_len": FIFO_LEN, "conv_len": CONV_LEN, "nx": nx, "ny": ny, "n_obj_total": 3, "obs_slice": (0, 6), "s_dim": 6, "max_steps": 150, "vortex_type": vortex_type, "vortex_y_offset_l0": y_offset_l0, "vortex_role": role, "runtime_handoff": "fresh_second_injection_from_saved_pre_injection_sensors_state"} + if target_only: + ff.restore_ddf(); ff.apply_ddf() + ff.add_vortex(target_vortex, 2.0 * L0, strength, 0.0, vortex_type) + return {"flow_field": ff, "target_states": target, "config": target_config} + ff.restore_ddf(); ff.apply_ddf() + ff.add_cylinder((30.0 * L0, cy, 0.0), PINBALL_RADIUS) + ff.add_cylinder((31.3 * L0, cy + 0.75 * L0, 0.0), PINBALL_RADIUS) + ff.add_cylinder((31.3 * L0, cy - 0.75 * L0, 0.0), PINBALL_RADIUS) + run_historical_interval(ff, int(nx / U0), np.zeros(6, dtype=DATA_TYPE)) + bias = np.zeros(6, dtype=DATA_TYPE); bias[3:] = np.asarray(action_bias, dtype=DATA_TYPE) * U0 + if role == "controlled": + run_historical_interval(ff, int(nx / U0), bias) + acquisition_vortex = (15.0 * L0, cy + y_offset_l0 * L0, 0.0) + pretrigger_development_steps = int(2 * nx / U0) + if role == "zero": + ff.add_vortex(acquisition_vortex, 2.0 * L0, strength, 0.0, vortex_type) + run_historical_interval( + ff, pretrigger_development_steps, np.zeros(6, dtype=DATA_TYPE) + ) + ff.add_vortex(acquisition_vortex, 2.0 * L0, strength, 0.0, vortex_type) + if role == "zero": + # Zero needs a true zero-base + fresh-vortex checkpoint. Controlled + # deliberately retains the historical stale-host save semantics. + ff.get_ddf() ff.save_ddf() - - # Norm fifo = deque(maxlen=FIFO_LEN) for _ in range(FIFO_LEN): - ff.run(800, np.zeros(6, dtype=DATA_TYPE)) - fifo.append(ff.obs.copy()[0:12]) - norm = _compute_karman_norm(np.array(fifo, dtype=DATA_TYPE)) - - # Bias FIFO - ff.apply_ddf() - bias_arr = np.zeros(6, dtype=DATA_TYPE) - bias_arr[3] = float(action_bias[0] * U0) - bias_arr[4] = float(action_bias[1] * U0) - bias_arr[5] = float(action_bias[2] * U0) - - fifo.clear() + run_historical_interval(ff, 800, np.zeros(6, dtype=DATA_TYPE)); fifo.append(ff.obs.copy()) + norm = _compute_karman_norm(np.asarray(fifo, dtype=DATA_TYPE)) + ff.apply_ddf(); fifo.clear() + runtime_command = bias if role == "controlled" else np.zeros(6, dtype=DATA_TYPE) for _ in range(FIFO_LEN): - ff.run(800, bias_arr) - fifo.append(ff.obs.copy()[0:12]) - save_states = np.array(list(fifo), dtype=DATA_TYPE) - ff.apply_ddf() - - norm["save_states"] = save_states - norm["action_bias"] = list(action_bias) - norm["n_obj_total"] = 6 - norm["vortex_type"] = vortex_type - norm["vortex_strength"] = vortex_strength - - config = { - "device_id": device_id, - "viscosity": viscosity, - "re_code": re_code, - "u0": U0, - "sample_interval": 800, - "fifo_len": FIFO_LEN, - "conv_len": CONV_LEN, - "nx": NX, - "ny": NY, - "n_obj_total": 6, - "action_scale": action_scale, - "action_bias": list(action_bias), - "obs_slice": (0, 12), - "s_dim": 12, - "max_steps": 150, - "vortex_type": vortex_type, - } - - return {"flow_field": ff, "target_states": target_states, "norm": norm, "config": config} - + run_historical_interval(ff, 800, runtime_command); fifo.append(ff.obs.copy()) + saved = np.asarray(fifo, dtype=DATA_TYPE); ff.apply_ddf() + norm.update(save_states=saved, action_bias=list(action_bias), n_obj_total=6, vortex_type=vortex_type, vortex_strength=strength) + config = {"device_id": device_id, "viscosity": viscosity, "re_code": re_code, "u0": U0, "sample_interval": 800, "fifo_len": FIFO_LEN, "conv_len": CONV_LEN, "nx": nx, "ny": ny, "n_obj_total": 6, "action_scale": action_scale, "action_bias": list(action_bias), "obs_slice": (0, 12), "s_dim": 12, "max_steps": 150, "vortex_type": vortex_type, "vortex_y_offset_l0": y_offset_l0, "vortex_role": role, "pre_injection_stabilization": "zero_then_bias" if role == "controlled" else "exact_physical_zero_with_vortex_pretrigger_2NX_over_U0", "pretrigger_vortex": role == "zero", "pretrigger_development_steps": pretrigger_development_steps if role == "zero" else 0, "runtime_handoff": "saved_fresh_injection_checkpoint"} + return {"flow_field": ff, "target_states": target, "norm": norm, "config": config} # --------------------------------------------------------------------------- # Erase (sensors(0-2) + dist-cyl(r=0.75L, id=3) + pinball(4-6) = 7 objects) @@ -551,6 +591,7 @@ def build_erase( re_code: float = 100.0, *, action_bias: Tuple[float, float, float] = (0.0, -8.0, 8.0), + target_only: bool = False, ) -> Dict[str, Any]: """Build erase environment (cancel upstream disturbance to clean flow). @@ -572,14 +613,17 @@ def build_erase( n_clean = ff_clean.obs.size // 2 assert n_clean == 3 - _stabilize(ff_clean, 3) + run_historical_interval(ff_clean, int(2 * ff_clean.FIELD_SHAPE[0] / U0), np.zeros(3, dtype=DATA_TYPE)) target = np.empty((0, 6), dtype=DATA_TYPE) for _ in range(FIFO_LEN): - ff_clean.run(600, np.zeros(3, dtype=DATA_TYPE)) + run_historical_interval(ff_clean, 600, np.zeros(3, dtype=DATA_TYPE)) target = np.vstack((target, ff_clean.obs.copy()[0:6])) # Target = mean (steady, not periodic) target_mean = np.mean(target, axis=0, dtype=DATA_TYPE) + if target_only: + nx, ny = _field_xy(ff_clean) + return {"flow_field": ff_clean, "target_states": target, "config": {"device_id": device_id, "viscosity": viscosity, "re_code": re_code, "u0": U0, "sample_interval": 600, "fifo_len": FIFO_LEN, "conv_len": 36, "nx": nx, "ny": ny, "n_obj_total": 3, "obs_slice": (0, 6), "s_dim": 6}} del ff_clean # Phase 2: Full erase env @@ -605,7 +649,7 @@ def build_erase( # Norm (erase-specific: full obs[0:14], force_norm uses pinball only) fifo = deque(maxlen=FIFO_LEN) for _ in range(FIFO_LEN): - ff.run(600, np.zeros(7, dtype=DATA_TYPE)) + run_historical_interval(ff, 600, np.zeros(7, dtype=DATA_TYPE)) fifo.append(ff.obs.copy()[0:14]) temp = np.array(fifo, dtype=DATA_TYPE) @@ -624,7 +668,7 @@ def build_erase( fifo.clear() for _ in range(FIFO_LEN): - ff.run(600, bias_arr) + run_historical_interval(ff, 600, bias_arr) fifo.append(ff.obs.copy()[0:14]) save_states = np.array(list(fifo), dtype=DATA_TYPE) ff.apply_ddf() diff --git a/src/drl_pinball/legacy_test/core/model_loader.py b/src/drl_pinball/legacy_test/core/model_loader.py index 6be09a3..7b3af3d 100644 --- a/src/drl_pinball/legacy_test/core/model_loader.py +++ b/src/drl_pinball/legacy_test/core/model_loader.py @@ -1,40 +1,82 @@ -# legacy_test/core/model_loader.py -"""PPO model loader for legacy test scripts. +"""Canonical, non-archive PPO loader for legacy evaluation scripts. -Wraps the reproduce ModelInventory to provide a simpler interface for -Track A test scripts that only need to load models onto CPU. +Only the models explicitly used by ``legacy_test/test_*.py`` are accepted. +Paths resolve directly under repository-root ``models/``; no discovery scan and +no archived Python inventory participates in evaluation. """ - from __future__ import annotations import os -import sys +from dataclasses import dataclass from typing import Optional -import numpy as np - _REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "..")) -_SRC = os.path.join(_REPO, "src") -for p in [_REPO, _SRC]: - if p not in sys.path: - sys.path.insert(0, p) - -from drl_pinball.reproduce.configs.model_inventory import ModelInventory # noqa: E402 - -_inventory = ModelInventory() +_MODEL_ROOT = os.path.join(_REPO, "models") -def load_model(name: str) -> "PPO": - """Load a PPO model onto CPU for inference. - - Returns a PPO model loaded from ``models/{subdir}/{name}.zip`` - with Sin activation and correct observation/action spaces. - - Delegates to ModelInventory.load(name, device="cpu"). - """ - return _inventory.load(name, device="cpu") +@dataclass(frozen=True) +class ModelSpec: + subdir: str + s_dim: int + scene: str -def list_models(scene: Optional[str] = None) -> list: - """List available model names, optionally filtered by scene.""" - return _inventory.list_models(scene) +# S_DIM and names are the canonical legacy-test README/script contracts. +MODEL_SPECS = { + "d1a3o12_re50": ModelSpec("old", 12, "karman"), + "d1a3o12_re100": ModelSpec("old", 12, "karman"), + "d1a3o12_re200": ModelSpec("old", 12, "karman"), + "d1a3o12_re400": ModelSpec("old", 12, "karman"), + "vortex_lamb": ModelSpec("old", 12, "vortex"), + "vortex_taylor": ModelSpec("old", 12, "vortex"), + "d1a3o14_250525_imit_075L_2U_400S": ModelSpec("250525", 14, "illusion"), + "d1a3o14_250525_imit_1L_2U_600S": ModelSpec("250525", 14, "illusion"), + "d1a3o14_250525_imit_15L_2U": ModelSpec("250525", 14, "illusion"), + "d1a3o12_250729_250326_erase": ModelSpec("250729", 12, "erase"), +} + + +def model_path(name: str) -> str: + """Resolve one allowlisted model to an explicit root-model zip path.""" + try: + spec = MODEL_SPECS[name] + except KeyError as exc: + raise KeyError(f"unknown canonical legacy model {name!r}") from exc + path = os.path.join(_MODEL_ROOT, spec.subdir, f"{name}.zip") + if not os.path.isfile(path): + raise FileNotFoundError(f"canonical legacy model not found: {path}") + return path + + +def _dummy_env(s_dim: int): + import gymnasium as gym + import numpy as np + from gymnasium import spaces + + class LegacyInferenceEnv(gym.Env): + observation_space = spaces.Box(-1.0, 1.0, shape=(s_dim,), dtype=np.float32) + action_space = spaces.Box(-1.0, 1.0, shape=(3,), dtype=np.float32) + + def reset(self, *, seed=None, options=None): + super().reset(seed=seed) + return np.zeros(s_dim, np.float32), {} + + def step(self, action): + return np.zeros(s_dim, np.float32), 0.0, False, False, {} + + return LegacyInferenceEnv() + + +def load_model(name: str): + """Load one canonical legacy PPO zip on CPU with its correct S_DIM.""" + from stable_baselines3 import PPO + spec = MODEL_SPECS.get(name) + if spec is None: + raise KeyError(f"unknown canonical legacy model {name!r}") + return PPO.load(model_path(name), env=_dummy_env(spec.s_dim), device="cpu") + + +def list_models(scene: Optional[str] = None) -> list[str]: + """List allowlisted model names, optionally filtered by canonical scene.""" + return sorted(name for name, spec in MODEL_SPECS.items() + if scene is None or spec.scene == scene) diff --git a/src/drl_pinball/legacy_test/metrics.py b/src/drl_pinball/legacy_test/metrics.py new file mode 100644 index 0000000..eea6efe --- /dev/null +++ b/src/drl_pinball/legacy_test/metrics.py @@ -0,0 +1,144 @@ +"""Legacy reward DTW and normalized offline target-scale metrics.""" +from __future__ import annotations +import hashlib +import os +import numpy as np +from .core.dtw_metrics import calc_dtw_sim, calc_lag, compute_similarity_vortex, gen_target_states_at + +def target_scale_cycle_similarity(target, state, conv_len, *, target_start=0): + target = np.asarray(target, dtype=np.float64)[:, target_start:target_start + 6] + state = np.asarray(state, dtype=np.float64)[:, :6] + scale = np.maximum(np.max(np.abs(target), axis=0), 1e-12) + result = cycle_similarity(target / scale, state / scale, conv_len) + result["scale"] = scale.tolist() + result["definition"] = "offline target-channel max-abs normalized legacy cycle DTW" + return result + +def cycle_similarity(target, state, conv_len): + target = np.asarray(target, dtype=np.float64); state = np.asarray(state, dtype=np.float64) + lag = calc_lag(target[conv_len:2*conv_len, 1], state[-conv_len:, 1]) + scores = [calc_dtw_sim(np.roll(target[:, i], -lag)[conv_len:2*conv_len], state[-conv_len:, i]) for i in range(6)] + return {"similarity": float(np.mean(scores)), "per_channel": {str(i): float(v) for i,v in enumerate(scores)}, "lag": int(lag)} + +def native_similarity(case, target, fifo, step): + if case.scene == "vortex": + return compute_similarity_vortex(target, fifo[:, :6], step, case.conv_len) + start = 2 if case.scene == "illusion" else 0 + return cycle_similarity(np.asarray(target)[:, start:start + 6], np.asarray(fifo)[:, :6], case.conv_len)["similarity"] + +def reward_terms(case, target, harmonics, fifo, force_norm, step): + fifo = np.asarray(fifo, dtype=np.float64) + forces = fifo[-1, 6:12] / force_norm + sim = native_similarity(case, target, fifo, step) + if case.scene == "illusion": + desired = gen_target_states_at(step, harmonics)[:2] / force_norm + cd, cl = forces[0::2].sum(), forces[1::2].sum() + r_cd, r_cl, r_sim = np.exp(-abs((cd-desired[0])*10)), np.exp(-abs((cl-desired[1])*10)), np.exp(-10*abs(sim-1)) + reward = .3*r_cd + .3*r_cl + .4*r_sim + else: + cd, cl = forces[0::2].mean(), forces[1::2].mean() + r_cd, r_cl, r_sim = np.exp(-abs(cd*20)), np.exp(-abs(cl*80)), np.exp(-10*abs(sim-1)) + weights = (.2, .3, .5) if case.scene == "vortex" else (.3, .4, .3) + reward = weights[0]*r_cd + weights[1]*r_cl + weights[2]*r_sim + return {"reward": float(min(reward, 1.0)), "reward_cd": float(r_cd), "reward_cl": float(r_cl), "native_legacy_dtw": float(sim)} + + +def sha256_file(path): + digest = hashlib.sha256() + with open(path, "rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + +def channel_correlations(reference, current): + reference = np.asarray(reference, dtype=np.float64) + current = np.asarray(current, dtype=np.float64) + n = min(len(reference), len(current)) + reference, current = reference[:n], current[:n] + result = [] + for index in range(reference.shape[1]): + x = reference[:, index] - reference[:, index].mean() + y = current[:, index] - current[:, index].mean() + denominator = np.sqrt(np.sum(x*x) * np.sum(y*y)) + result.append(float(np.sum(x*y) / denominator) if denominator >= 1e-12 else 0.0) + return result + +def rms_error(reference, current): + reference = np.asarray(reference, dtype=np.float64) + current = np.asarray(current, dtype=np.float64) + n = min(len(reference), len(current)) + return float(np.sqrt(np.mean((reference[:n] - current[:n]) ** 2))) + +def aggregate_rewards(reward_rows, conv_len): + if not reward_rows: + return None + keys = ("reward", "reward_cd", "reward_cl", "native_legacy_dtw") + values = {key: np.asarray([row[key] for row in reward_rows], dtype=np.float64) for key in keys} + tail = min(int(conv_len), len(reward_rows)) + return { + "mean": float(values["reward"].mean()), + "tail": float(values["reward"][-tail:].mean()), + "final": float(values["reward"][-1]), + "tail_reward_cd": float(values["reward_cd"][-tail:].mean()), + "tail_reward_cl": float(values["reward_cl"][-tail:].mean()), + "tail_native_legacy_dtw": float(values["native_legacy_dtw"][-tail:].mean()), + "tail_length": tail, + } + +def frozen_reference_comparison(path, sensors, forces, actions, conv_len, threshold=0.95): + if not os.path.isfile(path): + return None + with np.load(path) as frozen: + reference = {name: np.asarray(frozen[name], dtype=np.float64) for name in ("sensors", "forces", "actions")} + sensor_count = min(len(reference["sensors"]), len(sensors)) + sensor_dtw = cycle_similarity( + reference["sensors"][:sensor_count], np.asarray(sensors)[:sensor_count], conv_len + ) + return { + "reference_path": os.path.abspath(path), + "reference_sha256": sha256_file(path), + "sensor_dtw": sensor_dtw, + "sensor_correlations": channel_correlations(reference["sensors"], sensors), + "force_correlations": channel_correlations(reference["forces"], forces), + "action_correlations": channel_correlations(reference["actions"], actions), + "sensor_rms": rms_error(reference["sensors"], sensors), + "force_rms": rms_error(reference["forces"], forces), + "action_rms": rms_error(reference["actions"], actions), + "threshold": float(threshold), + "passed": bool(sensor_dtw["similarity"] > threshold), + } + + +def _erase_enhanced_similarity(target, state): + target = np.asarray(target, np.float64); state = np.asarray(state, np.float64) + target_std, state_std = max(float(target.std()), 1e-8), max(float(state.std()), 1e-8) + amplitude = min(target_std, state_std) / max(target_std, state_std) + mean_diff = abs(float(target.mean() - state.mean())) + mean_scale = max(abs(float(target.mean())), abs(float(state.mean())), 1e-8) + n, m = len(target), len(state); matrix = np.full((n + 1, m + 1), np.inf); matrix[0, 0] = 0.0 + for i in range(1, n + 1): + for j in range(1, m + 1): + matrix[i, j] = abs(target[i-1]-state[j-1]) + min(matrix[i-1,j], matrix[i,j-1], matrix[i-1,j-1]) + scale = max(float(np.max(np.abs(target))), float(np.max(np.abs(state))), 1e-8) + dtw = max(0.0, 1.0 - matrix[n,m] / (n * scale)) + return .8*dtw + .1*amplitude + .1/(1.0 + mean_diff/mean_scale*10.0) + +def erase_reward_terms(fifo, target_sensors, norm, conv_len=36): + """Exact rolling same-rollout reward from legacy_env_erase.py.""" + states = np.asarray(fifo, np.float64) + if states.ndim != 2 or states.shape[1] != 14 or len(states) < 2*conv_len: + raise ValueError("Erase reward requires a (time,14) FIFO of at least 2*conv_len") + sens_norm=np.asarray(norm["sens_norm_fact"],np.float64); sens_dev=np.asarray(norm["sens_deviation"],np.float64) + target_sensors=np.asarray(target_sensors,np.float64); sensors=(states[-1,:6]-sens_dev)/sens_norm; target=(target_sensors-sens_dev)/sens_norm + reference,current=-states[conv_len:2*conv_len,7],states[-conv_len:,9] + corr=np.correlate(reference-reference.mean(),current-current.mean(),"full"); lag=int(np.arange(-len(reference)+1,len(reference))[np.argmax(corr)]) + components=[] + for component in range(2): + wanted=-np.roll(states[:,component+6],-lag)[conv_len:2*conv_len] + actual=states[-conv_len:,component+8]+states[-conv_len:,component+10]+states[-conv_len:,component+12] + components.append(float(_erase_enhanced_similarity(wanted,actual))) + sim=float(np.mean(components)); diff_u=float(np.mean(np.abs(sensors[[0,2,4]]-target[[0,2,4]]))) + diff_v=float(np.mean(np.abs(sensors[[1,3,5]]-target[[1,3,5]]))); amp_v=float(np.std(states[:,1]+states[:,3]+states[:,5])/sens_norm[3]) + reward_u=float(np.exp(-abs(diff_u*40.0))); reward_v=float(.7*np.exp(-abs(amp_v*20.0))+.3*np.exp(-abs(diff_v*20.0))) + reward=float(min(.4*reward_u+.4*reward_v+.2*sim,1.0)) + return {"reward":reward,"reward_u":reward_u,"reward_v":reward_v,"native_legacy_dtw":sim,"native_lag":lag,"native_component_similarity":np.asarray(components)} diff --git a/src/drl_pinball/legacy_test/runtime.py b/src/drl_pinball/legacy_test/runtime.py new file mode 100644 index 0000000..b75fdc7 --- /dev/null +++ b/src/drl_pinball/legacy_test/runtime.py @@ -0,0 +1,64 @@ +"""Runtime contracts shared by every Legacy reproduction case.""" +from __future__ import annotations +from collections import deque +import json +import os +import numpy as np + +def run_historical_interval(flow_field, num_steps, action_target): + """Run one interval with the pre-persistence Legacy EMA seed contract.""" + shape = np.asarray(flow_field.action).shape + flow_field.action = np.zeros(shape, dtype=flow_field.DATA_TYPE) + flow_field.run(num_steps, action_target) + +def load_policy_norm(path): + """Load and validate one explicitly mapped frozen training normalization.""" + if not os.path.isfile(path): + raise FileNotFoundError(f"Frozen policy normalization not found: {path}") + with open(path, encoding="utf-8") as handle: + document = json.load(handle) + required = {"force_norm_fact", "sens_deviation", "sens_norm_fact"} + missing = required - set(document) + if missing: + raise ValueError(f"Frozen policy normalization missing keys: {sorted(missing)}") + force = float(document["force_norm_fact"]) + deviation = np.asarray(document["sens_deviation"], dtype=np.float32) + scale = np.asarray(document["sens_norm_fact"], dtype=np.float32) + if not np.isfinite(force) or force <= 0: + raise ValueError("Frozen force_norm_fact must be finite and positive") + if deviation.shape != (6,) or scale.shape != (6,): + raise ValueError("Frozen sensor normalization arrays must each have shape (6,)") + if not np.all(np.isfinite(deviation)) or not np.all(np.isfinite(scale)) or np.any(scale <= 0): + raise ValueError("Frozen sensor normalization values must be finite and scales positive") + return {"force_norm_fact": force, "sens_deviation": deviation, "sens_norm_fact": scale} + +def compute_norm(rows): + rows = np.asarray(rows, dtype=np.float32) + return {"force_norm_fact": float(6*np.max(np.abs(rows[:, 6:12]))), + "sens_deviation": rows[:, :6].mean(axis=0), + "sens_norm_fact": 5*np.max(np.abs(rows[:, :6]-rows[:, :6].mean(axis=0)), axis=0)} + +def reset_runtime(flow_field, saved_fifo, s_dim): + flow_field.restore_ddf() + flow_field.apply_ddf() + fifo = deque((np.asarray(x, dtype=np.float32).copy() for x in saved_fifo), maxlen=len(saved_fifo)) + return fifo, np.zeros(s_dim, dtype=np.float32) + +def policy_observation(raw, norm, *, target_force=None): + raw = np.asarray(raw, dtype=np.float32) + force = raw[6:12] / norm["force_norm_fact"] + sensors = (raw[:6] - norm["sens_deviation"]) / norm["sens_norm_fact"] + result = np.hstack((force, sensors)) + if target_force is not None: + result = np.hstack((result, np.asarray(target_force)/norm["force_norm_fact"])) + return np.clip(result, -1, 1).astype(np.float32) + + +def erase_policy_observation(raw, norm): + """Original Erase ordering: pinball forces then normalized sensors.""" + raw = np.asarray(raw, dtype=np.float32) + if raw.shape != (14,) or not np.all(np.isfinite(raw)): + raise ValueError("Erase raw observation must be 14 finite channels") + result = np.hstack((raw[8:14] / norm["force_norm_fact"], + (raw[:6] - norm["sens_deviation"]) / norm["sens_norm_fact"])) + return np.clip(result, -1, 1).astype(np.float32) diff --git a/src/drl_pinball/legacy_test/tests/test_acquire.py b/src/drl_pinball/legacy_test/tests/test_acquire.py new file mode 100644 index 0000000..bdec36d --- /dev/null +++ b/src/drl_pinball/legacy_test/tests/test_acquire.py @@ -0,0 +1,202 @@ +from pathlib import Path + +import numpy as np +import pytest + +import src.drl_pinball.legacy_test.acquire as acquire + + +class FakeModel: + def __init__(self): + self.calls = 0 + + def predict(self, obs, deterministic): + assert deterministic is True + self.calls += 1 + return np.array([0.25, -0.5, 0.75], dtype=np.float32), None + + +class FakeFlow: + DATA_TYPE = np.float32 + + def __init__(self): + self.action = np.ones(7, dtype=np.float32) + self.obs = np.zeros(14, dtype=np.float32) + self.intervals = 0 + self.capture_intervals = [] + + def restore_ddf(self): + pass + + def apply_ddf(self): + pass + + def run(self, steps, command): + assert steps == acquire.SI + assert np.array_equal(self.action, np.zeros(7, dtype=np.float32)) + self.intervals += 1 + self.action = command.copy() + phase = 2 * np.pi * self.intervals / 30 + self.obs[2:8] = [np.sin(phase + index) for index in range(6)] + self.obs[8:14] = np.arange(6) + self.intervals + + def current_effective_action(self): + return self.action.copy() + + def solver_clock_state(self): + return {"solver_absolute_lattice_clock": 1_000_000 + self.intervals * acquire.SI, + "solver_absolute_control_clock": self.intervals} + + +def _data(): + target = np.column_stack([ + np.sin(2 * np.pi * np.arange(150) / 30 + index) for index in range(6) + ]).astype(np.float32) + return { + "config": {"nx": 1280, "ny": 2, "sample_interval": 800, "n_obj_total": 7}, + "norm": {"save_states": np.zeros((150, 12), dtype=np.float32)}, + "target_states": target, + } + + +def test_controlled_schedule_zero_ema_and_post_interval_field_capture(monkeypatch, tmp_path): + ff, model = FakeFlow(), FakeModel() + monkeypatch.setattr(acquire, "reward_terms", lambda *args: { + "reward": 0.5, "reward_cd": 0.6, "reward_cl": 0.7, "native_legacy_dtw": 0.8, + }) + + def capture(flow, nx, ny): + assert flow.intervals > acquire.WARMUP_INTERVALS + flow.capture_intervals.append(flow.intervals) + return {"ux": np.full((ny, nx), flow.intervals, np.float32), + "uy": np.full((ny, nx), -flow.intervals, np.float32)} + + rows, fields = acquire._collect(_data(), model, ff, { + "force_norm_fact": 1.0, "sens_deviation": np.zeros(6), "sens_norm_fact": np.ones(6), + }, tmp_path, capture_field=capture) + + assert (acquire.WARMUP_INTERVALS, acquire.COLLECT_BOUNDARIES) == (480, 160) + assert ff.intervals == model.calls == 640 + assert len(rows) == 160 + assert fields["ux"].shape == fields["uy"].shape == (160, 2, 1280) + assert ff.capture_intervals == list(range(481, 641)) + assert rows[0]["control_index"] == 481 and rows[-1]["control_index"] == 640 + assert rows[0]["lattice_step"] == 481 * 800 + assert rows[-1]["lattice_step"] == 640 * 800 + assert rows[0]["solver_absolute_lattice_step"] == 1_000_000 + 481 * 800 + assert rows[0]["solver_absolute_lattice_step"] != rows[0]["lattice_step"] + assert rows[0]["physical_time"] == rows[0]["lattice_step"] * acquire.U0 / acquire.NX + assert np.allclose(rows[0]["commanded_target_omega"], [0.02, -0.08, 0.1]) + assert np.array_equal(rows[0]["commanded_target_omega"], rows[0]["effective_smoothed_omega"]) + assert np.allclose(fields["ux"][0], 481) and np.allclose(fields["uy"][-1], -640) + assert not list(tmp_path.glob("boundary_*.npz")) + assert set(rows[0]) == { + "physical_time", "lattice_step", "solver_absolute_lattice_step", "control_index", + "sensors", "forces", + "action_normalized", "commanded_target_omega", "effective_smoothed_omega", + "reward_raw", "reward_cd", "reward_cl", "reward_sim", "native_reward_dtw", + } + + +def test_legacy_metadata_schema_and_timeline_wording_are_explicit(): + source = Path(acquire.__file__).read_text(encoding="utf-8") + assert '"schema": "drl-pinball-legacy-acquisition-v2"' in source + assert '"schema": "drl-pinball-v5-acquisition-v1"' not in source + assert '"solver_absolute_lattice_step"' in source + assert 'V5 lattice_step is an absolute solver count' in source + assert 'derived only from relative Legacy lattice_step' in source + + +def test_capture_uses_public_q_over_u0_and_transposes_to_physical_velocity(): + ny, nx = 2, 3 + + class PublicFlow: + class Config: + velocity = 0.01 + field_config = Config() + + def current_step_velocity_field(self): + return (np.array([[1, 2], [3, 4], [0, 6]], dtype=np.float32), + np.array([[-1, -2], [-3, -4], [0, -6]], dtype=np.float32)) + + def completed_flags_xy(self): + return np.array([[1, 1], [0b00010001, 1], [2, 1]], dtype=np.uint8) + + fields = acquire._capture_boundary_field(PublicFlow(), nx, ny) + assert fields["ux"].shape == fields["uy"].shape == (ny, nx) + assert np.allclose(fields["ux"], [[0.01, 0.03, 0.0], [0.02, 0.04, 0.06]]) + assert np.allclose(fields["uy"], -fields["ux"]) + assert fields["ux"].dtype == fields["uy"].dtype == np.float32 + + +def test_capture_rejects_nonfluid_public_velocity(): + class BadFlow: + class Config: + velocity = 0.01 + field_config = Config() + def current_step_velocity_field(self): + return np.ones((1, 1), np.float32), np.zeros((1, 1), np.float32) + def completed_flags_xy(self): + return np.zeros((1, 1), np.uint8) + + with pytest.raises(ValueError, match="exact zero on nonfluid"): + acquire._capture_boundary_field(BadFlow(), 1, 1) + +def test_staged_exact_keys(tmp_path): + role = tmp_path / "role" + role.mkdir() + for name in ("timeseries.csv", "phase_cycle.csv"): + (role / name).write_text("x\n", encoding="ascii") + required = { + "physical_time": np.arange(2), "lattice_step": np.arange(2), + "solver_absolute_lattice_step": np.arange(2) + 100, + "control_index": np.arange(2), "sensors": np.zeros((2, 6)), + "forces": np.zeros((2, 6)), "action_normalized": np.zeros((2, 3)), + "commanded_target_omega": np.zeros((2, 3)), + "effective_smoothed_omega": np.zeros((2, 3)), "reward_raw": np.zeros(2), + "reward_cd": np.zeros(2), "reward_cl": np.zeros(2), "reward_sim": np.zeros(2), + "native_reward_dtw": np.zeros(2), "target_normalized_dtw": np.zeros(2), + "target_normalized_dtw_lag": np.zeros(2), "phase": np.zeros(2), + "cycle_id": np.zeros(2), + } + np.savez(role / "timeseries.npz", **required) + np.savez(role / "phase_cycle.npz", x=np.zeros(2)) + np.savez(role / "phase_fields.npz", ux=np.zeros((8, 2, 3)), uy=np.zeros((8, 2, 3)), + mean_ux=np.zeros((2, 3)), mean_uy=np.zeros((2, 3)), + field_indices=np.arange(8), cycle_id=np.arange(8), target_phase=np.arange(8), + actual_phase=np.arange(8), phase_error=np.zeros(8)) + (role / "dtw_summary.json").write_text("{}", encoding="utf-8") + (role / "metadata.json").write_text("{}", encoding="utf-8") + acquire._validate_staged_role(role) + + +def test_failure_removes_transaction_and_preserves_existing(monkeypatch, tmp_path): + root = tmp_path / "out" + final = root / "legacy" / acquire.CASE_NAME / acquire.ROLE + final.mkdir(parents=True) + marker = final / "keep.txt" + marker.write_text("old", encoding="utf-8") + model = tmp_path / "model.zip"; model.write_bytes(b"model") + norm = tmp_path / "norm.json"; norm.write_text("{}", encoding="utf-8") + bundle = {"case": acquire.get_case(acquire.CASE_NAME), "model_path": model, + "norm_path": norm, "norm": {}, "reference_path": tmp_path / "missing.npz"} + monkeypatch.setattr(acquire, "_resolve_bundle", lambda repo: bundle) + monkeypatch.setattr(acquire, "_collect_role", lambda *args, **kwargs: ([], [])) + + def storage_validator(**kwargs): + root.mkdir(exist_ok=True) + return {"resolved_output_root": root, "device": root.stat().st_dev} + + def runtime_factory(case, current_bundle, device, role): + assert role == "controlled" + return _data(), FakeModel(), FakeFlow() + + def fail(*args): + raise RuntimeError("finalize failed") + + with pytest.raises(RuntimeError, match="finalize failed"): + acquire.acquire_controlled(output_root=root, overwrite=True, repo_root=tmp_path, + storage_validator=storage_validator, + runtime_factory=runtime_factory, finalizer=fail) + assert marker.read_text(encoding="utf-8") == "old" + assert not list(final.parent.glob(".controlled.staging-*")) diff --git a/src/drl_pinball/legacy_test/tests/test_acquire_matrix.py b/src/drl_pinball/legacy_test/tests/test_acquire_matrix.py new file mode 100644 index 0000000..91917a0 --- /dev/null +++ b/src/drl_pinball/legacy_test/tests/test_acquire_matrix.py @@ -0,0 +1,217 @@ +import numpy as np +import pytest + +import src.drl_pinball.legacy_test.acquire as acquire + + +def test_allowlist_roles_and_case_sample_intervals(): + assert acquire.SUPPORTED_CASES == ( + "karman_re50", "karman_re100", "karman_re200", "karman_re400", + "illusion_075L", "illusion_1L", "illusion_15L", "steady", "vortex_lamb", "vortex_taylor", "erase", + ) + assert [acquire.get_case(name).sample_interval for name in + ("karman_re50", "illusion_075L", "illusion_1L", "illusion_15L", "steady")] == [800, 400, 600, 800, 800] + assert acquire._case_roles(acquire.get_case("steady")) == ("target", "constant", "zero") + assert acquire._case_roles(acquire.get_case("illusion_075L")) == ("controlled", "target", "zero") + + +def test_scene_observation_slices_are_exact(): + obs14 = np.arange(14, dtype=np.float32) + sensors, forces = acquire._observation_slices(acquire.get_case("karman_re50"), "controlled", obs14) + assert np.array_equal(sensors, obs14[2:8]) and np.array_equal(forces, obs14[8:14]) + obs12 = np.arange(12, dtype=np.float32) + for name in ("illusion_075L", "steady"): + sensors, forces = acquire._observation_slices(acquire.get_case(name), "zero", obs12) + assert np.array_equal(sensors, obs12[:6]) and np.array_equal(forces, obs12[6:12]) + target8 = np.arange(8, dtype=np.float32) + sensors, forces = acquire._observation_slices(acquire.get_case("illusion_075L"), "target", target8) + assert np.array_equal(sensors, target8[2:8]) and np.isnan(forces).all() + + +def test_illusion_policy_observation_is_exact_14d_next_harmonic_step(monkeypatch, tmp_path): + case = acquire.get_case("illusion_075L") + class Model: + def predict(self, obs, deterministic=True): + assert obs.shape == (14,) + return np.zeros(3, np.float32), None + class Flow: + DATA_TYPE = np.float32 + def __init__(self): self.obs=np.zeros(12,np.float32); self.action=np.zeros(6,np.float32); self.i=0 + def restore_ddf(self): pass + def apply_ddf(self): pass + def run(self, steps, command): self.i += 1; self.action=command.copy(); self.obs=np.arange(12,dtype=np.float32) + def current_effective_action(self): return self.action + data={"config":{"nx":1280,"ny":1,"sample_interval":400,"n_obj_total":6}, + "norm":{"save_states":np.zeros((2,12),np.float32)}, "target_states":np.zeros((80,8)), + "target_harmonics":[{}]*8} + seen=[] + monkeypatch.setattr(acquire, "gen_target_states_at", lambda step, harmonics: seen.append(step) or np.arange(8,dtype=float)+step) + monkeypatch.setattr(acquire, "reward_terms", lambda *args: {"reward":1.,"reward_cd":1.,"reward_cl":1.,"native_legacy_dtw":1.}) + monkeypatch.setattr(acquire, "policy_observation", lambda raw,norm,target_force=None: np.r_[raw,target_force].astype(np.float32)) + monkeypatch.setattr(acquire, "WARMUP_INTERVALS", 1); monkeypatch.setattr(acquire, "COLLECT_BOUNDARIES", 1) + acquire._collect_role("controlled", data, Model(), Flow(), {"force_norm_fact":1.}, tmp_path, + capture_field=lambda *args: {"ux":np.zeros((1,1280)),"uy":np.zeros((1,1280))}, + case_name=case.name) + assert seen == [1, 2] + + +def test_steady_staging_has_no_phase_or_dtw_artifacts(tmp_path): + role=tmp_path/"role"; role.mkdir() + np.savez(role/"timeseries.npz", x=np.zeros(1)); (role/"timeseries.csv").write_text("x\n0\n") + np.savez(role/"late_field.npz", ux=np.zeros((1,2,3)), uy=np.zeros((1,2,3)), field_indices=np.array([159])) + (role/"metadata.json").write_text("{}") + acquire._validate_staged_role(role, "steady") + assert not any("phase" in path.name or "dtw" in path.name for path in role.iterdir()) + + +def test_nonstandard_cases_are_allowlisted(): + for name in ("vortex_lamb", "vortex_taylor", "erase"): + assert acquire.get_case(name).name == name + assert acquire._case_roles(acquire.get_case(name)) == ("controlled", "target", "zero") + + +def test_target_and_steady_bundle_resolution_do_not_require_frozen_norm(monkeypatch, tmp_path): + monkeypatch.setattr(acquire, "load_policy_norm", lambda path: pytest.fail("frozen norm loaded")) + target = acquire._resolve_bundle(tmp_path, "illusion_075L", "target") + steady = acquire._resolve_bundle(tmp_path, "steady", "constant") + assert target["norm"] is target["norm_path"] is None + assert steady["norm"] is steady["norm_path"] is None + + +def test_vortex_zero_bundle_uses_builder_norm_when_frozen_norm_is_absent(monkeypatch, tmp_path): + monkeypatch.setattr(acquire, "load_policy_norm", lambda path: pytest.fail("frozen norm loaded")) + bundle = acquire._resolve_bundle(tmp_path, "vortex_lamb", "zero") + assert bundle["norm"] is bundle["norm_path"] is None + assert acquire._special_normalization_source(bundle["case"], "zero") == ( + "builder_recomputed_zero_runtime" + ) + assert acquire._special_normalizer_file(bundle) is None + + +def test_vortex_controlled_bundle_still_requires_frozen_norm(tmp_path): + with pytest.raises(FileNotFoundError, match="Frozen policy normalization not found"): + acquire._resolve_bundle(tmp_path, "vortex_lamb", "controlled") + + +def test_steady_constant_and_zero_match_active_physical_commands(monkeypatch, tmp_path): + case = acquire.get_case("steady") + class Flow: + DATA_TYPE = np.float32 + def __init__(self): + self.obs = np.zeros(12, np.float32); self.action = np.zeros(6, np.float32); self.commands = [] + def restore_ddf(self): pass + def apply_ddf(self): pass + def run(self, steps, command): self.commands.append(command.copy()); self.action = command.copy() + def current_effective_action(self): return self.action + data = {"config":{"nx":1280,"ny":1,"sample_interval":800,"n_obj_total":6}, + "norm":{"save_states":np.zeros((2,12),np.float32)}} + norm = {"force_norm_fact":1., "sens_deviation":np.zeros(6), "sens_norm_fact":np.ones(6)} + monkeypatch.setattr(acquire, "WARMUP_INTERVALS", 1); monkeypatch.setattr(acquire, "COLLECT_BOUNDARIES", 1) + for role, expected in (("constant", np.asarray(case.action_bias) * acquire.U0), + ("zero", np.zeros(3))): + flow = Flow() + rows, _ = acquire._collect_role(role, data, None, flow, norm, tmp_path, + capture_field=lambda *args: {"ux":np.zeros((1,1280)), "uy":np.zeros((1,1280))}, + case_name="steady") + assert len(flow.commands) == 2 + assert all(np.allclose(command[-3:], expected) for command in flow.commands) + assert np.allclose(rows[0]["commanded_target_omega"], expected) + + +def test_periodic_publication_requires_four_crossings(): + with pytest.raises(ValueError, match="at least 4 center-uy rising crossings"): + acquire._require_periodic_crossings({"crossing_times": np.arange(3)}) + assert acquire._require_periodic_crossings({"crossing_times": np.arange(4)}) == 4 + + +def test_illusion_collector_actions_exactly_match_active_run_convention(monkeypatch, tmp_path): + case = acquire.get_case("illusion_075L") + predicted = [np.array([.25, -.5, .75], np.float32), + np.array([-.125, .375, -.625], np.float32), + np.array([.5, 0., -.25], np.float32)] + class Model: + def __init__(self): self.index = 0 + def predict(self, obs, deterministic=True): + result = predicted[self.index]; self.index += 1; return result, None + class Flow: + DATA_TYPE = np.float32 + def __init__(self): + self.obs=np.zeros(12,np.float32); self.action=np.zeros(6,np.float32); self.commands=[] + def restore_ddf(self): pass + def apply_ddf(self): pass + def run(self, steps, command): self.commands.append(command.copy()); self.action=command.copy(); self.obs += 1 + def current_effective_action(self): return self.action + data={"config":{"nx":1280,"ny":1,"sample_interval":400,"n_obj_total":6}, + "norm":{"save_states":np.zeros((2,12),np.float32)}, "target_states":np.zeros((80,8)), + "target_harmonics":[{}]*8} + norm={"force_norm_fact":1., "sens_deviation":np.zeros(6), "sens_norm_fact":np.ones(6)} + monkeypatch.setattr(acquire, "WARMUP_INTERVALS", 1); monkeypatch.setattr(acquire, "COLLECT_BOUNDARIES", 2) + monkeypatch.setattr(acquire, "reward_terms", lambda *args: {"reward":1.,"reward_cd":1.,"reward_cl":1.,"native_legacy_dtw":1.}) + monkeypatch.setattr(acquire, "gen_target_states_at", lambda step, harmonics: np.arange(8,dtype=float)+step) + flow=Flow() + acquire._collect_role("controlled", data, Model(), flow, norm, tmp_path, + capture_field=lambda *args: {"ux":np.zeros((1,1280)),"uy":np.zeros((1,1280))}, + case_name=case.name) + # Exact convention in active run.py: (normalized_action * action_scale + action_bias) * U0. + expected=[(action * case.action_scale + np.asarray(case.action_bias)) * acquire.U0 for action in predicted] + assert all(np.allclose(command[-3:], wanted) for command, wanted in zip(flow.commands, expected)) + + +def test_generated_array_identity_binds_dtype_shape_and_content(): + value = np.arange(12, dtype=np.float32).reshape(3, 4) + first = acquire._generated_array_identity(value) + assert first == acquire._generated_array_identity(value.copy()) + assert first["identity_kind"] == "run-generated-array-content" + assert first["dtype"] == "= 14: + self.obs[8:14] = np.arange(6) + self.intervals + + def current_effective_action(self): + return self.action.copy() + + def solver_clock_state(self): + return {"solver_absolute_lattice_clock": self.absolute_base + self.intervals * acquire.SI, + "solver_absolute_control_clock": self.intervals} + + +def data(n_objects=7): + phase = 2 * np.pi * np.arange(150) / acquire.CYCLE_LENGTH + target = np.column_stack([np.sin(phase + index) for index in range(6)]).astype(np.float32) + result = { + "config": {"nx": 1280, "ny": 2, "sample_interval": 800, + "n_obj_total": n_objects}, + "target_states": target, + } + if n_objects == 7: + result["norm"] = {"save_states": np.zeros((150, 12), dtype=np.float32)} + return result + + +def norm(): + return {"force_norm_fact": 1.0, "sens_deviation": np.zeros(6), + "sens_norm_fact": np.ones(6)} + + +def capture(flow, nx, ny): + assert flow.intervals > acquire.WARMUP_INTERVALS + flow.capture_intervals.append(flow.intervals) + return {"ux": np.zeros((ny, nx), np.float32), "uy": np.zeros((ny, nx), np.float32)} + + +def test_target_builder_uses_disturbance_then_three_sensors_without_pinball(monkeypatch): + class TargetFlow: + DATA_TYPE = np.float32 + FIELD_SHAPE = (1280, 2, 1) + + def __init__(self, field_cfg, cuda_cfg, device_id): + self.obs = np.zeros(0, np.float32) + self.objects = [] + self.interval = 0 + + def add_cylinder(self, center, radius): + self.objects.append(("cylinder", center, radius)) + self.obs = np.zeros(2 * len(self.objects), np.float32) + + def add_sensor(self, center, radius): + self.objects.append(("sensor", center, radius)) + self.obs = np.zeros(2 * len(self.objects), np.float32) + + class Config: + def _replace(self, **kwargs): + return self + + monkeypatch.setattr(builder, "FlowField", TargetFlow) + monkeypatch.setattr(builder.legacy_utils, "load_cuda_config", lambda path: object()) + monkeypatch.setattr(builder.legacy_utils, "load_flow_field_config", lambda path: Config()) + monkeypatch.setattr(builder, "_stabilize", lambda flow, count: None) + monkeypatch.setattr(builder, "FIFO_LEN", 2) + + def interval(flow, steps, command): + flow.interval += 1 + flow.obs[0:2] = [100 + flow.interval, 200 + flow.interval] + flow.obs[2:8] = np.arange(6) + 10 * flow.interval + + monkeypatch.setattr(builder, "run_historical_interval", interval) + result = builder.build_karman_target(0, 100.0, sample_interval=800) + + flow = result["flow_field"] + assert [item[0] for item in flow.objects] == ["cylinder", "sensor", "sensor", "sensor"] + assert flow.objects[0][1][0] == 10.0 * builder.L0 + assert [item[1][1] for item in flow.objects[1:]] == [40.5, 0.5, -39.5] + assert result["config"]["n_obj_total"] == 4 + assert np.array_equal(result["target_states"], + np.asarray([np.arange(6) + 10, np.arange(6) + 20], np.float32)) + + +def test_zero_role_schedule_actions_rewards_and_fields(monkeypatch, tmp_path): + flow = RoleFlow(7, 14) + reward_calls = [] + + def rewards(*args): + reward_calls.append(args) + return {"reward": 0.4, "reward_cd": 0.5, "reward_cl": 0.6, + "native_legacy_dtw": 0.7} + + monkeypatch.setattr(acquire, "reward_terms", rewards) + rows, fields = acquire._collect_role( + "zero", data(), None, flow, norm(), tmp_path, capture_field=capture, + ) + + assert flow.intervals == len(reward_calls) == 640 + assert flow.capture_intervals == list(range(481, 641)) + assert len(rows) == 160 + assert fields["ux"].shape == fields["uy"].shape == (160, 2, 1280) + assert not list(tmp_path.glob("boundary_*.npz")) + counter_bias = np.array([0.0, 0.5, -0.5], np.float32) + assert np.array_equal(rows[0]["action_normalized"], counter_bias) + assert all(np.array_equal(command, np.zeros(7, np.float32)) for command in flow.commands) + assert np.array_equal(rows[0]["commanded_target_omega"], np.zeros(3, np.float32)) + assert np.allclose(rows[0]["effective_smoothed_omega"], 0.0, atol=1e-7) + assert rows[0]["reward_raw"] == 0.4 + assert rows[0]["native_reward_dtw"] == 0.7 + assert rows[0]["control_index"] == 481 and rows[-1]["control_index"] == 640 + + +def test_zero_role_metadata_names_physical_zero_counter_bias(): + source = Path(acquire.__file__).read_text(encoding="utf-8") + assert "physical-zero/uncontrolled" in source + assert "counter-bias normalized action" in source + + +def test_target_role_ordering_nan_contract_and_schedule(monkeypatch, tmp_path): + flow = RoleFlow(4, 8) + monkeypatch.setattr(acquire, "reward_terms", lambda *args: pytest.fail("target has no reward")) + rows, fields = acquire._collect_role( + "target", data(n_objects=4), None, flow, None, tmp_path, capture_field=capture, + ) + + assert flow.intervals == 640 + assert all(np.array_equal(command, np.zeros(4, np.float32)) for command in flow.commands) + assert len(rows) == 160 + assert fields["ux"].shape == (160, 2, 1280) + first = rows[0] + phase = 2 * np.pi * 481 / acquire.CYCLE_LENGTH + expected = np.asarray([np.sin(phase + index) for index in range(6)], np.float32) + assert np.allclose(first["sensors"], expected) + assert first["sensors"][0] != flow.obs[0] + for name in ("forces", "action_normalized", "commanded_target_omega", + "effective_smoothed_omega"): + assert np.isnan(first[name]).all() + for name in ("reward_raw", "reward_cd", "reward_cl", "reward_sim", + "native_reward_dtw"): + assert np.isnan(first[name]) + assert first["solver_absolute_lattice_step"] == 2_000_000 + 481 * 800 + + +def test_target_sensor_reference_selects_legacy_target_channels(): + six = np.arange(60, dtype=float).reshape(10, 6) + eight = np.column_stack((np.full((10, 2), -1.0), six)) + assert np.array_equal(acquire._target_sensors(six), six) + assert np.array_equal(acquire._target_sensors(eight), six) + with pytest.raises(ValueError, match="six sensor channels"): + acquire._target_sensors(np.zeros((10, 7))) + + +def test_target_phase_artifacts_exclude_unavailable_nan_variables(): + columns = {"sensors": np.ones((4, 6)), "forces": np.full((4, 6), np.nan), + "action_normalized": np.full((4, 3), np.nan), + "physical_time": np.arange(4), "lattice_step": np.arange(4), + "solver_absolute_lattice_step": np.arange(4), + "control_index": np.arange(4), "phase": np.arange(4), + "cycle_id": np.arange(4)} + values = acquire._phase_values(columns, "target") + assert set(values) == {"sensors"} + assert np.isfinite(values["sensors"]).all() + + +@pytest.mark.parametrize("role", ["target", "zero"]) +def test_role_failure_transaction_preserves_existing(monkeypatch, tmp_path, role): + root = tmp_path / "out" + final = root / "legacy" / acquire.CASE_NAME / role + final.mkdir(parents=True) + marker = final / "keep.txt" + marker.write_text("old", encoding="utf-8") + model = tmp_path / "model.zip"; model.write_bytes(b"model") + norm_path = tmp_path / "norm.json"; norm_path.write_text("{}", encoding="utf-8") + bundle = {"case": acquire.get_case(acquire.CASE_NAME), "model_path": model, + "norm_path": norm_path, "norm": norm(), + "reference_path": tmp_path / "missing.npz"} + monkeypatch.setattr(acquire, "_resolve_bundle", lambda repo, *args: bundle) + monkeypatch.setattr(acquire, "_collect_role", lambda *args, **kwargs: ([], [])) + + def storage_validator(**kwargs): + root.mkdir(exist_ok=True) + return {"resolved_output_root": root, "device": root.stat().st_dev} + + def runtime_factory(case, current_bundle, device, current_role): + n_objects = 4 if current_role == "target" else 7 + return data(n_objects), None, RoleFlow(n_objects, 8 if current_role == "target" else 14) + + with pytest.raises(RuntimeError, match="finalize failed"): + acquire.acquire_role( + role, output_root=root, overwrite=True, repo_root=tmp_path, + storage_validator=storage_validator, runtime_factory=runtime_factory, + finalizer=lambda *args: (_ for _ in ()).throw(RuntimeError("finalize failed")), + ) + assert marker.read_text(encoding="utf-8") == "old" + assert not list(final.parent.glob(f".{role}.staging-*")) + + +def test_cli_exposes_all_roles(): + source = Path(acquire.__file__).read_text(encoding="utf-8") + assert 'choices=ROLES' in source diff --git a/src/drl_pinball/legacy_test/tests/test_contracts.py b/src/drl_pinball/legacy_test/tests/test_contracts.py new file mode 100644 index 0000000..8e3717a --- /dev/null +++ b/src/drl_pinball/legacy_test/tests/test_contracts.py @@ -0,0 +1,287 @@ +import ast +from pathlib import Path + +import numpy as np +import src.drl_pinball.legacy_test.metrics as legacy_metrics +import src.drl_pinball.legacy_test.render as legacy_render +import src.drl_pinball.legacy_test.run as legacy_run +from src.drl_pinball.legacy_test.cases import CASES, ERASE_UNSUPPORTED, get_case +from src.drl_pinball.legacy_test.metrics import ( + aggregate_rewards, frozen_reference_comparison, reward_terms, sha256_file, + target_scale_cycle_similarity, +) +from src.drl_pinball.legacy_test.render import ddf_to_velocity +from src.drl_pinball.legacy_test.core.legacy_env_builder import _field_xy +from src.drl_pinball.legacy_test.runtime import ( + compute_norm, load_policy_norm, policy_observation, reset_runtime, + run_historical_interval, +) + +def test_config_matrix(): + assert set(CASES)=={"karman_re50","karman_re100","karman_re200","karman_re400","illusion_075L","illusion_1L","illusion_15L","vortex_lamb","vortex_taylor","erase","steady"} + assert [CASES[f"karman_re{x}"].re_code for x in (50,100,200,400)] == [50,100,200,400] + assert [CASES[x].sample_interval for x in ("illusion_075L","illusion_1L","illusion_15L")] == [400,600,800] + assert get_case("erase").model == "d1a3o12_250729_250326_erase" + +def test_target_scale_metric_is_scale_invariant(): + t=np.column_stack([np.sin(np.arange(90)/7+i) * (i+1) for i in range(6)]) + result=target_scale_cycle_similarity(t,t,30) + factors=np.arange(1,7,dtype=np.float64) + scaled=target_scale_cycle_similarity(t*factors,t*factors,30) + assert np.isclose(result["similarity"],scaled["similarity"]) + assert len(result["scale"])==6 + +def test_norm_recomputed_from_rows(): + rows=np.arange(120,dtype=np.float32).reshape(10,12) + norm=compute_norm(rows) + assert norm["force_norm_fact"] == 6*np.max(np.abs(rows[:,6:12])) + assert not np.allclose(norm["sens_deviation"],0) + +class FakeFlow: + def __init__(self): self.calls=[] + def restore_ddf(self): self.calls.append("restore") + def apply_ddf(self): self.calls.append("apply") + +def test_reset_restores_ddf_fifo_and_zero_observation(): + ff=FakeFlow(); saved=np.arange(36,dtype=np.float32).reshape(3,12) + fifo,obs=reset_runtime(ff,saved,12) + assert ff.calls==["restore","apply"] + assert np.array_equal(np.asarray(fifo),saved) + assert np.array_equal(obs,np.zeros(12,dtype=np.float32)) + +def test_ddf_conversion_uses_reference_density_not_population_sum(): + ny,nx=3,4; f=np.zeros((ny,nx,9),dtype=np.float32); f[:,:,0]=7; f[:,:,1]=.01 + flags=np.ones((nx,ny),dtype=np.uint8) + ux,uy=ddf_to_velocity(f.transpose(2,0,1).reshape(-1),flags,nx,ny) + assert np.allclose(ux,.01) + assert np.allclose(uy,0) + + +def test_render_final_accepts_real_legacy_3d_shape(monkeypatch, tmp_path): + class RenderFlow: + FIELD_SHAPE = (4, 3, 1) + + def __init__(self): + f = np.zeros((3, 4, 9), dtype=np.float32) + f[:, :, 0] = 1.0 + f[:, :, 1] = 0.25 + self.ddf = f.transpose(2, 0, 1).reshape(-1) + self.flag = np.ones(12, dtype=np.uint8) + self.downloaded = False + + def get_ddf(self): + self.downloaded = True + + def completed_flags_xy(self): + return self.flag.reshape((3, 4)).T.copy() + + captured = {} + + def fake_compute_vorticity(ux, uy): + captured["velocity_shape"] = (ux.shape, uy.shape) + return np.zeros_like(ux) + + def fake_render(vort, **kwargs): + captured["vorticity_shape"] = vort.shape + captured["render_kwargs"] = kwargs + return {"image_path": kwargs["out_path"]} + + monkeypatch.setattr(legacy_render, "compute_vorticity", fake_compute_vorticity) + monkeypatch.setattr(legacy_render, "render_vorticity_field", fake_render) + + flow = RenderFlow() + out_path = tmp_path / "final_vorticity.png" + result = legacy_render.render_final(flow, str(out_path)) + + assert flow.downloaded + assert captured["velocity_shape"] == ((3, 4), (3, 4)) + assert captured["vorticity_shape"] == (3, 4) + assert captured["render_kwargs"] == { + "nx": 4, + "ny": 3, + "out_path": str(out_path), + "vmin": -0.001, + "vmax": 0.001, + } + assert result == {"image_path": str(out_path)} + + +def test_reward_aggregation_uses_case_tail(): + rows = [ + {"reward": float(i), "reward_cd": float(i + 1), "reward_cl": float(i + 2), + "native_legacy_dtw": float(i + 3)} + for i in range(5) + ] + summary = aggregate_rewards(rows, 3) + assert summary == { + "mean": 2.0, + "tail": 3.0, + "final": 4.0, + "tail_reward_cd": 4.0, + "tail_reward_cl": 5.0, + "tail_native_legacy_dtw": 6.0, + "tail_length": 3, + } + +def test_frozen_reference_comparison_reports_exact_metrics(tmp_path): + time = np.arange(90, dtype=np.float64) + phase = 2.0 * np.pi * time / 30.0 + sensors = np.column_stack([np.sin(phase + index) for index in range(6)]) + forces = np.column_stack([np.cos(phase + index) for index in range(6)]) + actions = np.column_stack([np.sin(phase + index) for index in range(3)]) + path = tmp_path / "controlled.npz" + np.savez_compressed(path, sensors=sensors, forces=forces, actions=actions, rewards=time) + + result = frozen_reference_comparison(str(path), sensors, forces, actions, 30) + + assert result["reference_path"] == str(path.resolve()) + assert len(result["reference_sha256"]) == 64 + assert np.isclose(result["sensor_dtw"]["similarity"], 1.0) + assert np.allclose(result["sensor_correlations"], 1.0) + assert np.allclose(result["force_correlations"], 1.0) + assert np.allclose(result["action_correlations"], 1.0) + assert result["sensor_rms"] == result["force_rms"] == result["action_rms"] == 0.0 + assert result["threshold"] == 0.95 + assert result["passed"] is True + + +def test_historical_interval_zeros_each_ema_seed_and_forwards_target(): + class IntervalFlow: + DATA_TYPE = np.float32 + + def __init__(self): + self.action = np.array([4.0, -3.0], dtype=np.float32) + self.calls = [] + + def run(self, steps, target): + self.calls.append((self.action.copy(), steps, target.copy())) + self.action[:] = 9.0 + + flow = IntervalFlow() + first = np.array([1.0, 2.0], dtype=np.float32) + second = np.array([-1.0, 3.0], dtype=np.float32) + + run_historical_interval(flow, 8, first) + run_historical_interval(flow, 12, second) + + assert np.array_equal(flow.calls[0][0], np.zeros(2, dtype=np.float32)) + assert np.array_equal(flow.calls[1][0], np.zeros(2, dtype=np.float32)) + assert flow.calls[0][1] == 8 + assert flow.calls[1][1] == 12 + assert np.array_equal(flow.calls[0][2], first) + assert np.array_equal(flow.calls[1][2], second) + assert flow.calls[0][0].dtype == flow.calls[1][0].dtype == np.float32 + +def test_active_builder_and_runner_use_historical_interval_adapter_only(): + root = Path(__file__).resolve().parents[1] + for relative in ("run.py", "core/legacy_env_builder.py"): + tree = ast.parse((root / relative).read_text()) + direct_calls = [ + node for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "run" + ] + assert direct_calls == [], f"{relative} must route all FlowField intervals through adapter" + + +def test_frozen_policy_norm_drives_observation_reward_and_provenance(monkeypatch, tmp_path): + path = tmp_path / "norm.json" + path.write_text( + '{"force_norm_fact": 2.0, "sens_deviation": [1, 1, 1, 1, 1, 1], ' + '"sens_norm_fact": [2, 2, 2, 2, 2, 2]}', + encoding="utf-8", + ) + frozen = load_policy_norm(str(path)) + recomputed = { + "force_norm_fact": 8.0, + "sens_deviation": np.zeros(6, dtype=np.float32), + "sens_norm_fact": np.ones(6, dtype=np.float32), + } + raw = np.array([3.0] * 6 + [1.0] * 6, dtype=np.float32) + + observation = policy_observation(raw, frozen) + recomputed_observation = policy_observation(raw, recomputed) + + assert np.allclose(observation[:6], 0.5) + assert np.allclose(observation[6:], 1.0) + assert not np.array_equal(observation, recomputed_observation) + + monkeypatch.setattr(legacy_metrics, "native_similarity", lambda *args: 1.0) + reward = reward_terms( + CASES["karman_re50"], np.zeros((60, 6)), None, + np.tile(raw, (30, 1)), frozen["force_norm_fact"], 0, + ) + expected_cd = np.exp(-abs(0.5 * 20.0)) + expected_cl = np.exp(-abs(0.5 * 80.0)) + assert np.isclose(reward["reward_cd"], expected_cd) + assert np.isclose(reward["reward_cl"], expected_cl) + assert len(sha256_file(str(path))) == 64 + +def test_missing_frozen_policy_norm_fails_closed(tmp_path): + missing = tmp_path / "missing" / "norm.json" + try: + load_policy_norm(str(missing)) + except FileNotFoundError as exc: + assert str(missing) in str(exc) + else: + raise AssertionError("missing mapped policy norm must fail closed") + + +def test_metrics_only_output_helpers_have_no_filesystem_or_render(monkeypatch, tmp_path): + calls = [] + + monkeypatch.setattr(legacy_run.os.path, "isdir", lambda path: calls.append(("isdir", path))) + monkeypatch.setattr(legacy_run.shutil, "rmtree", lambda path: calls.append(("rmtree", path))) + monkeypatch.setattr(legacy_run.os, "makedirs", lambda path: calls.append(("makedirs", path))) + monkeypatch.setattr(legacy_run.np, "savez_compressed", lambda *args, **kwargs: calls.append(("save", args[0]))) + monkeypatch.setattr(legacy_run, "_json", lambda path, value: calls.append(("json", path))) + monkeypatch.setattr(legacy_run, "render_final", lambda ff, path: calls.append(("render", path))) + + out = str(tmp_path / "case") + legacy_run._prepare_case_output(out, metrics_only=True) + legacy_run._write_case_output( + out, True, object(), np.zeros((1, 6)), np.zeros((1, 6)), + np.zeros((1, 3)), [], np.zeros((1, 6)), np.zeros((1, 12)), {}, {}, + ) + + assert calls == [] + +def test_normal_output_helpers_write_defined_artifacts(monkeypatch, tmp_path): + calls = [] + out = str(tmp_path / "case") + + monkeypatch.setattr(legacy_run.np, "savez_compressed", lambda *args, **kwargs: calls.append(Path(args[0]).name)) + monkeypatch.setattr(legacy_run, "_json", lambda path, value: calls.append(Path(path).name)) + monkeypatch.setattr(legacy_run, "render_final", lambda ff, path: calls.append(Path(path).name)) + + legacy_run._write_case_output( + out, False, object(), np.zeros((1, 6)), np.zeros((1, 6)), + np.zeros((1, 3)), [], np.zeros((1, 6)), np.zeros((1, 12)), {}, {}, + ) + + assert calls == [ + "signals.npz", "reset_contract.npz", "metrics.json", "norm.json", + "final_vorticity.png", + ] + + +def test_builder_uses_first_two_legacy_shape_dimensions_and_validates_z(): + class ShapeFlow: + def __init__(self, shape): + self.FIELD_SHAPE = shape + + assert _field_xy(ShapeFlow((1280, 512, 1))) == (1280, 512) + try: + _field_xy(ShapeFlow((1280, 512, 2))) + except ValueError as exc: + assert "nz=1" in str(exc) + else: + raise AssertionError("builder must reject non-singleton z shape") + + try: + _field_xy(ShapeFlow((1280, 512))) + except ValueError as exc: + assert "(nx, ny, nz)" in str(exc) + else: + raise AssertionError("builder must reject two-dimensional shape") diff --git a/src/drl_pinball/legacy_test/tests/test_nonstandard_acquire.py b/src/drl_pinball/legacy_test/tests/test_nonstandard_acquire.py new file mode 100644 index 0000000..e903062 --- /dev/null +++ b/src/drl_pinball/legacy_test/tests/test_nonstandard_acquire.py @@ -0,0 +1,312 @@ +import numpy as np +import pytest +from src.drl_pinball.legacy_test import acquire +from src.drl_pinball.legacy_test.core import legacy_env_builder as builder +from src.drl_pinball.legacy_test.metrics import erase_reward_terms +from src.drl_pinball.legacy_test.runtime import erase_policy_observation + +def test_vortex_labels_and_offsets_are_exact(): + assert acquire._vortex_scenario("vortex_lamb",0)=="vortex_lamb_y000" + assert [acquire._vortex_scenario("vortex_taylor",x) for x in (-2,-1,0,1,2)] == ["vortex_taylor_ym2L","vortex_taylor_ym1L","vortex_taylor_y000","vortex_taylor_yp1L","vortex_taylor_yp2L"] + for name,value in (("vortex_lamb",1),("vortex_taylor",3)): + with pytest.raises(ValueError): acquire._vortex_offset(name,value) + +def test_lamb_event_uses_only_upper_uy_and_earliest_exact_five(): + sensors=np.zeros((150,6)); sensors[62,1]=4; sensors[60,3]=100; sensors[63,5]=200 + event=acquire._select_vortex_event(sensors,"vortex_lamb") + assert int(event["peak_index"])==62 + assert event["event_channel_index"]==1 + assert np.array_equal(event["field_indices"],[52,57,62,67,72]) + assert (event["search_lower_boundary"],event["search_upper_boundary"])==(38,87) + + +def test_taylor_event_uses_only_center_uy_y0_source_channel(): + sensors=np.zeros((150,6)); sensors[55,1]=100; sensors[70,3]=4; sensors[65,5]=200 + event=acquire._select_vortex_event(sensors,"vortex_taylor") + assert int(event["peak_index"])==70 + assert event["event_channel_index"]==3 + assert np.array_equal(event["field_indices"],[60,65,70,75,80]) + + +def test_event_baseline_window_and_earliest_tie_are_exact(): + sensors=np.zeros((150,6)); sensors[:20,1]=2; sensors[50,1]=5; sensors[51,1]=-1 + event=acquire._select_vortex_event(sensors,"vortex_lamb") + assert event["baseline_uy"]==2 + assert int(event["peak_index"])==50 + sensors[149,1]=100 + assert int(acquire._select_vortex_event(sensors,"vortex_lamb")["peak_index"])==50 + + +def test_vortex_schedule_matches_historical_episode(): + lamb=acquire.get_case("vortex_lamb"); taylor=acquire.get_case("vortex_taylor") + for target in (False,True): + assert acquire._special_schedule(lamb,target)==(0,150) + assert acquire._special_schedule(taylor,target)==(0,150) + assert lamb.steps==taylor.steps==150 + assert acquire._special_timeline_metadata(taylor,False)=={ + "warmup_control_steps":0,"collected_post_step_boundaries":150, + "historical_episode_steps":150} + +def test_erase_observation_inventory_and_policy_order(): + raw=np.arange(14,dtype=np.float32); norm={"force_norm_fact":2.,"sens_deviation":np.zeros(6),"sens_norm_fact":np.ones(6)*100} + obs=erase_policy_observation(raw,norm) + assert np.array_equal(obs[:4],np.clip(raw[8:12]/2,-1,1)) + sensors,dist,pinball,full=acquire._special_observation(acquire.get_case("erase"),"controlled",raw) + assert np.array_equal(sensors,raw[:6]) and np.array_equal(dist,raw[6:8]) and np.array_equal(pinball,raw[8:14]) and np.array_equal(full,raw) + +def test_erase_controlled_reset_policy_obs_is_exact_12d_zero(monkeypatch,tmp_path): + case=acquire.get_case("erase") + norm={"force_norm_fact":2.,"sens_deviation":np.zeros(6,dtype=np.float32), + "sens_norm_fact":np.ones(6,dtype=np.float32), + "save_states":np.full((150,14),7,dtype=np.float32)} + data={"config":{"n_obj_total":7,"nx":2,"ny":1},"norm":norm, + "target_states":np.zeros((150,6),dtype=np.float32)} + class Model: + def __init__(self): self.observations=[] + def predict(self,obs,deterministic=True): + self.observations.append(np.asarray(obs).copy()) + return np.zeros(3,dtype=np.float32),None + class Flow: + DATA_TYPE=np.float32 + def __init__(self): + self.action=np.zeros(7,dtype=np.float32); self.obs=np.zeros(14,dtype=np.float32); self.step=0 + def restore_ddf(self): pass + def apply_ddf(self): pass + def run(self,steps,command): + self.step+=1; self.action=command.copy() + self.obs=np.arange(14,dtype=np.float32)+self.step + self.obs[6:8]=[1000+self.step,2000+self.step] + def current_effective_action(self): return self.action + monkeypatch.setattr(acquire,"_special_schedule",lambda case,target:(0,2)) + monkeypatch.setattr(acquire,"erase_reward_terms",lambda *args,**kwargs:{"reward":0.,"reward_u":0.,"reward_v":0.,"native_legacy_dtw":0.}) + model=Model(); flow=Flow() + acquire._collect_special("controlled",data,model,flow,norm,tmp_path,case, + capture_field=lambda *args:{"ux":np.zeros((1,2),np.float32),"uy":np.zeros((1,2),np.float32)}) + assert len(model.observations)==2 + assert model.observations[0].shape==(12,) + assert model.observations[0].dtype==np.float32 + assert np.array_equal(model.observations[0],np.zeros(12,dtype=np.float32)) + raw=np.arange(14,dtype=np.float32)+1; raw[6:8]=[1001,2001] + expected=erase_policy_observation(raw,norm) + assert model.observations[1].shape==(12,) + assert np.array_equal(model.observations[1],expected) + assert np.array_equal(model.observations[1][:6],np.clip(raw[8:14]/2,-1,1)) + assert not np.any(np.isin(model.observations[1],raw[6:8])) + + +def test_erase_reward_formula_exact_simple_fixture(): + t=np.arange(150,dtype=float); states=np.zeros((150,14)); states[:,6]=np.sin(t/5); states[:,7]=np.cos(t/5) + states[:,8]=-.5*states[:,6]; states[:,10]=-.25*states[:,6]; states[:,12]=-.25*states[:,6] + states[:,9]=-.5*states[:,7]; states[:,11]=-.25*states[:,7]; states[:,13]=-.25*states[:,7] + norm={"force_norm_fact":1.,"sens_deviation":np.zeros(6),"sens_norm_fact":np.ones(6)} + result=erase_reward_terms(states,np.zeros(6),norm) + assert 0 <= result["native_legacy_dtw"] <= 1 + assert np.isclose(result["reward"],.4*result["reward_u"]+.4*result["reward_v"]+.2*result["native_legacy_dtw"]) + assert result["native_component_similarity"].shape==(2,) + +def test_erase_phase_fixed_probe_then_fallback(): + times=np.arange(160,dtype=float); center=np.ones(160); disturbance=np.sin(2*np.pi*times/20) + phase,probe,fallback=acquire._erase_phase(times,center,disturbance) + assert len(phase["crossing_times"])>=4 and probe=="disturbance_force_y" and fallback + +def _field_candidates(tmp_path, count=5): + paths=[] + for index in range(count): + path=tmp_path/f"candidate_{index}.npz" + np.savez(path,ux=np.full((2,3),index),uy=np.full((2,3),-index)) + paths.append(path) + return paths + + +def test_publish_fields_reproduces_scalar_selection_rejection(tmp_path): + destination=tmp_path/"event_fields.npz" + selection={"field_indices":np.arange(5),"relative_offsets":np.arange(5), + "selected_event_score":np.arange(5,dtype=float),"peak_index":np.int64(2)} + with pytest.raises(ValueError,match="named numeric non-scalar arrays required"): + acquire._publish_fields(destination,_field_candidates(tmp_path),selection) + assert not destination.exists() + + +def test_publish_fields_writes_only_aligned_event_arrays(tmp_path): + destination=tmp_path/"event_fields.npz" + selection={"field_indices":np.arange(5),"relative_offsets":np.array([-10,-5,0,5,10]), + "selected_event_score":np.linspace(0,1,5)} + acquire._publish_fields(destination,_field_candidates(tmp_path),selection) + with np.load(destination,allow_pickle=False) as saved: + assert set(saved.files)=={"ux","uy","field_indices","relative_offsets","selected_event_score"} + assert all(saved[name].ndim>=1 and saved[name].shape[0]==5 for name in saved.files) + assert "peak_index" not in saved.files and "event_score" not in saved.files + + +def test_special_exact_inventory_validation(tmp_path): + role=tmp_path/"role"; role.mkdir() + np.savez(role/"timeseries.npz",x=np.zeros(5)) + np.savez(role/"event_fields.npz",ux=np.zeros((5,1,1)),uy=np.zeros((5,1,1)), + field_indices=np.array([-8,-3,2,7,12]),relative_offsets=np.array([-10,-5,0,5,10]), + selected_event_score=np.zeros(5)) + (role/"timeseries.csv").write_text("x\n") + (role/"event_summary.json").write_text(__import__("json").dumps({ + "schema":"drl-pinball-legacy-vortex-event-v2", + "event_source":{"scenario":"vortex_lamb_y000","role":"controlled","channel_index":1,"channel":"upper_sensor_uy"}, + "common_peak_index":2,"field_indices":[-8,-3,2,7,12],"relative_offsets":[-10,-5,0,5,10]})) + (role/"metadata.json").write_text("{}") + acquire._validate_special(role,acquire.get_case("vortex_lamb"),"controlled") + with np.load(role/"event_fields.npz") as saved: + assert set(saved.files)=={"ux","uy","field_indices","relative_offsets","selected_event_score"} + np.savez(role/"event_fields.npz",ux=np.zeros((5,1,1)),uy=np.zeros((5,1,1)), + field_indices=np.arange(5),relative_offsets=np.arange(5), + selected_event_score=np.zeros(5),peak_index=np.asarray(2)) + with pytest.raises(ValueError,match="Vortex event field keys mismatch"): + acquire._validate_special(role,acquire.get_case("vortex_lamb"),"controlled") + np.savez(role/"event_fields.npz",ux=np.zeros((5,1,1)),uy=np.zeros((5,1,1)), + field_indices=np.arange(5),relative_offsets=np.arange(5), + selected_event_score=np.zeros(5)) + (role/"extra").write_text("") + with pytest.raises(ValueError,match="staged role files mismatch"): acquire._validate_special(role,acquire.get_case("vortex_lamb"),"controlled") + + +def test_special_failure_transaction_preserves_existing(monkeypatch,tmp_path): + root=tmp_path/"out"; final=root/"legacy"/"vortex_lamb_y000"/"controlled" + final.mkdir(parents=True); marker=final/"keep.txt"; marker.write_text("old") + model=tmp_path/"model.zip"; model.write_bytes(b"model") + norm_path=tmp_path/"norm.json"; norm_path.write_text("{}") + bundle={"case":acquire.get_case("vortex_lamb"),"model_path":model,"norm_path":norm_path, + "norm":{"force_norm_fact":1.},"reference_path":tmp_path/"missing.npz"} + monkeypatch.setattr(acquire,"_resolve_bundle",lambda *args:bundle) + monkeypatch.setattr(acquire,"_collect_special",lambda *args,**kwargs:([],[])) + monkeypatch.setattr(acquire,"_finalize_special",lambda *args:(_ for _ in ()).throw(RuntimeError("scalar publication failed"))) + class Flow: + def close(self): pass + def storage_validator(**kwargs): + root.mkdir(exist_ok=True); return {"resolved_output_root":root,"device":root.stat().st_dev} + def runtime_factory(*args): return {"norm":{}},object(),Flow() + with pytest.raises(RuntimeError,match="scalar publication failed"): + acquire.acquire_role("controlled",case_name="vortex_lamb",output_root=root,overwrite=True, + repo_root=tmp_path,storage_validator=storage_validator,runtime_factory=runtime_factory) + assert marker.read_text()=="old" + assert not list(final.parent.glob(".controlled.staging-*")) + + +def _write_event_source(root, case_name="vortex_lamb", peak=62): + scenario=acquire._canonical_vortex_source(case_name) + channel,name=acquire._vortex_event_channel(case_name) + path=root/"legacy"/scenario/"controlled"/"event_summary.json" + path.parent.mkdir(parents=True) + path.write_text(__import__("json").dumps({"schema":"drl-pinball-legacy-vortex-event-v2", + "event_source":{"scenario":scenario,"role":"controlled","channel_index":channel,"channel":name}, + "common_peak_index":peak,"field_indices":[peak-10,peak-5,peak,peak+5,peak+10], + "relative_offsets":[-10,-5,0,5,10]})) + return path + + +def test_common_source_indices_reused_across_roles_and_taylor_offsets(tmp_path): + storage={"resolved_output_root":tmp_path} + _write_event_source(tmp_path,"vortex_lamb",61) + _write_event_source(tmp_path,"vortex_taylor",73) + assert np.array_equal(acquire._load_vortex_event_source(storage,"vortex_lamb")["field_indices"],[51,56,61,66,71]) + for offset in (-2,-1,0,1,2): + source=acquire._load_vortex_event_source(storage,"vortex_taylor") + assert np.array_equal(source["field_indices"],[63,68,73,78,83]) + assert source["event_source"]["scenario"]=="vortex_taylor_y000" + + +def test_missing_or_malformed_common_source_fails_closed(tmp_path): + storage={"resolved_output_root":tmp_path} + with pytest.raises(FileNotFoundError,match="required before this run"): + acquire._load_vortex_event_source(storage,"vortex_lamb") + path=_write_event_source(tmp_path,"vortex_lamb") + path.write_text("{}") + with pytest.raises(ValueError,match="malformed canonical"): + acquire._load_vortex_event_source(storage,"vortex_lamb") + + +class _VortexBuilderFlow: + DATA_TYPE=np.float32 + FIELD_SHAPE=(1280,2,1) + def __init__(self,*args,**kwargs): + self.obs=np.zeros(0,np.float32); self.action=np.zeros(0,np.float32) + self.runs=[]; self.vortices=[]; self.actions=[]; self.restore_count=0; self.active_injection=0 + def add_sensor(self,*args): self.obs=np.zeros(self.obs.size+2,np.float32); self.action=np.zeros(self.obs.size//2,np.float32) + def add_cylinder(self,*args): self.obs=np.zeros(self.obs.size+2,np.float32); self.action=np.zeros(self.obs.size//2,np.float32) + def add_vortex(self,*args): + self.vortices.append(args); self.actions.append(("add_vortex",args)); self.active_injection+=1 + def get_ddf(self): pass + def save_ddf(self): pass + def restore_ddf(self): self.restore_count+=1; self.active_injection=0 + def apply_ddf(self): pass + + +class _BuilderConfig: + def _replace(self,**kwargs): return self + + +def _patch_vortex_builder(monkeypatch): + monkeypatch.setattr(builder,"FlowField",_VortexBuilderFlow) + monkeypatch.setattr(builder.legacy_utils,"load_cuda_config",lambda path:object()) + monkeypatch.setattr(builder.legacy_utils,"load_flow_field_config",lambda path:_BuilderConfig()) + monkeypatch.setattr(builder,"FIFO_LEN",2) + def interval(flow,steps,command): + copied=np.asarray(command).copy() + flow.runs.append((steps,copied)); flow.actions.append(("run",steps,copied)) + flow.action=copied + flow.obs[:]=flow.active_injection + monkeypatch.setattr(builder,"run_historical_interval",interval) + + +def test_target_builder_handoff_is_fresh_second_injection(monkeypatch): + _patch_vortex_builder(monkeypatch) + result=builder.build_vortex(target_only=True,role="target") + flow=result["flow_field"] + long_step=int(1280/builder.U0) + assert np.array_equal(result["target_states"],np.ones((2,6),np.float32)) + assert len(flow.vortices)==2 and flow.restore_count==1 and flow.active_injection==1 + assert [action[0] for action in flow.actions]==["run","add_vortex","run","run","add_vortex"] + assert flow.actions[0][1]==long_step and len(flow.actions[0][2])==3 + assert all(action[1]==800 and len(action[2])==3 for action in flow.actions[2:4]) + assert flow.vortices[0]==flow.vortices[1] + assert result["config"]["runtime_handoff"]=="fresh_second_injection_from_saved_pre_injection_sensors_state" + + +def test_zero_builder_pretriggers_then_formally_injects_at_exact_zero(monkeypatch): + _patch_vortex_builder(monkeypatch) + zero=builder.build_vortex(role="zero",vortex_type="taylor",y_offset_l0=2) + flow=zero["flow_field"] + long_step=int(1280/builder.U0); development_steps=int(2*1280/builder.U0) + six_object_actions=[action for action in flow.actions + if action[0]=="add_vortex" and action[1][0][0]==15*builder.L0 + or action[0]=="run" and len(action[2])==6 and action[1] in (long_step,development_steps)] + assert [action[0] for action in six_object_actions]==["run","add_vortex","run","add_vortex"] + assert six_object_actions[0][1]==long_step + assert np.array_equal(six_object_actions[0][2],np.zeros(6,np.float32)) + assert six_object_actions[2][1]==development_steps + assert np.array_equal(six_object_actions[2][2],np.zeros(6,np.float32)) + assert six_object_actions[1][1]==six_object_actions[3][1] + vortex=six_object_actions[1][1] + assert vortex[0]==(15*builder.L0,1/2+2*builder.L0,0.0) + assert vortex[1]==2*builder.L0 and vortex[2]==0.03*builder.U0 and vortex[4]=="taylor" + assert zero["config"]["pre_injection_stabilization"]=="exact_physical_zero_with_vortex_pretrigger_2NX_over_U0" + assert zero["config"]["pretrigger_vortex"] is True + assert zero["config"]["pretrigger_development_steps"]==development_steps + + +def test_controlled_builder_keeps_zero_bias_then_one_formal_vortex(monkeypatch): + _patch_vortex_builder(monkeypatch) + controlled=builder.build_vortex(role="controlled") + flow=controlled["flow_field"] + long_step=int(1280/builder.U0) + acquisition_actions=[action for action in flow.actions + if action[0]=="add_vortex" and action[1][0][0]==15*builder.L0 + or action[0]=="run" and len(action[2])==6 and action[1]==long_step] + assert [action[0] for action in acquisition_actions]==["run","run","add_vortex"] + assert np.array_equal(acquisition_actions[0][2],np.zeros(6,np.float32)) + assert np.array_equal(acquisition_actions[1][2][3:],np.asarray([0,-4,4],np.float32)*builder.U0) + assert sum(vortex[0][0]==15*builder.L0 for vortex in flow.vortices)==1 + assert controlled["config"]["pre_injection_stabilization"]=="zero_then_bias" + assert controlled["config"]["pretrigger_vortex"] is False + assert controlled["config"]["pretrigger_development_steps"]==0 + + +def test_vortex_builder_role_and_target_only_are_bound(monkeypatch): + _patch_vortex_builder(monkeypatch) + with pytest.raises(ValueError,match="semantics disagree"): + builder.build_vortex(target_only=True,role="controlled") diff --git a/tests/test_drl_pinball_acquisition.py b/tests/test_drl_pinball_acquisition.py index d036eca..d1d1d57 100644 --- a/tests/test_drl_pinball_acquisition.py +++ b/tests/test_drl_pinball_acquisition.py @@ -9,8 +9,11 @@ import numpy as np import pytest from drl_pinball.acquisition import ( + accumulate_mean_fields, assign_phase, + assign_periodic_phase, cleanup_scratch, + complete_cycle_field_indices, create_scratch, default_reproduction_mapping, decode_legacy_physical_velocity, @@ -20,6 +23,7 @@ from drl_pinball.acquisition import ( publish_role_output, publish_selected_fields, select_phase_fields, + smooth_center_uy, validate_output_storage, validate_modern_fields, write_boundary_artifacts, @@ -184,6 +188,59 @@ def test_selection_is_global_eight_with_metadata_and_index_tie(): assert np.allclose(result["phase_error"], [0.0] * 8, atol=1e-15) +def test_smooth_center_uy_binomial_and_endpoint_copy(): + assert np.allclose(smooth_center_uy([1.0, 5.0, 1.0]), [1.0, 3.0, 1.0]) + assert np.allclose(smooth_center_uy([2.0, 4.0]), [2.0, 4.0]) + + +def test_assign_phase_default_behavior_unchanged_with_jitter(): + times = np.arange(12, dtype=float) + sensors = np.zeros((12, 6)) + sensors[:, 3] = [-1, 1, -0.1, 0.1, -1, 1, -1, 1, -1, 1, -1, 1] + result = assign_phase(times, sensors) + assert len(result["crossing_times"]) == 6 + assert set(result) == {"phase", "cycle_id", "crossing_times"} + + +def test_assign_periodic_phase_filters_jitter_with_minimum_interval(): + times = np.arange(20, dtype=float) + sensors = np.zeros((20, 6)) + # Deep brief negative between rises so smoothing keeps a close false crossing. + sensors[:, 3] = [-2, -2, -1, 1, 2, -3, 1, 2, 2, -2, -1, 1, 2, 2, -2, -1, 1, 2, 2, -2] + raw = assign_phase(times, sensors) + filtered = assign_periodic_phase(times, sensors, minimum_crossing_interval=4.0) + assert len(raw["crossing_times"]) >= 4 + assert filtered["rejected_crossing_count"] >= 1 + assert 5.75 in filtered["rejected_crossing_times"] + assert filtered["accepted_crossing_count"] == len(filtered["crossing_times"]) + assert filtered["complete_cycle_count"] == filtered["accepted_crossing_count"] - 1 + assert np.all(np.diff(filtered["crossing_times"]) >= 4.0 - 1e-12) + assert filtered["smoothing_kernel"] == [0.25, 0.5, 0.25] + first, last = filtered["crossing_times"][0], filtered["crossing_times"][-1] + valid = (times >= first) & (times < last) + assert np.all(filtered["cycle_id"][valid] >= 0) + assert np.all(filtered["cycle_id"][~valid] < 0) + + +def test_complete_cycle_mean_excludes_outside_first_last_crossing(): + times = np.arange(10, dtype=float) + crossings = np.array([2.0, 6.0, 9.0]) + selected = complete_cycle_field_indices(times, crossings) + assert selected["field_indices"].tolist() == [2, 3, 4, 5, 6, 7, 8] + assert selected["mean_field_count"] == 7 + assert selected["complete_cycle_count"] == 2 + ux = np.arange(10, dtype=np.float32).reshape(10, 1, 1) + uy = (2 * np.arange(10)).astype(np.float32).reshape(10, 1, 1) + mean = accumulate_mean_fields(ux, uy, selected["field_indices"]) + assert mean["mean_ux"].shape == (1, 1) + assert mean["mean_uy"].shape == (1, 1) + assert float(mean["mean_ux"][0, 0]) == pytest.approx(5.0) + assert float(mean["mean_uy"][0, 0]) == pytest.approx(10.0) + # Eight-slot average of selected phase snapshots must remain a different concept. + eight = accumulate_mean_fields(ux, uy, np.array([2, 3, 4, 5, 6, 7, 8, 8])) + assert float(eight["mean_ux"][0, 0]) != pytest.approx(float(mean["mean_ux"][0, 0])) + + def test_pooled_32_bins_population_std_and_original_indices(): result = pooled_phase_bins( np.array([np.pi / 32, 3 * np.pi / 32, np.pi / 32, np.nan]), diff --git a/tests/test_drl_pinball_infer_train.py b/tests/test_drl_pinball_infer_train.py new file mode 100644 index 0000000..8254567 --- /dev/null +++ b/tests/test_drl_pinball_infer_train.py @@ -0,0 +1,271 @@ +from __future__ import annotations + +import ast +import csv +import json +from pathlib import Path + +import numpy as np + + +ROOT = Path(__file__).resolve().parents[1] +SOURCE = ROOT / "src/drl_pinball/eval/infer_train.py" + + +def _load_helpers(*names: str): + tree = ast.parse(SOURCE.read_text()) + selected = [node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name in names] + namespace = { + "Path": Path, "Dict": dict, "Any": object, "List": list, "Tuple": tuple, + "np": np, "json": json, "csv": csv, "hashlib": __import__("hashlib"), "CaseSpec": object, "shutil": __import__("shutil"), + "NUM_STEPS": 360, "TAIL_WINDOW": 180, "_POLICY_SENSOR_UNITS": "legacy-policy-v1", + "REFERENCE_TABLE_ATOL": 1e-6, "REPRO_REWARD_ATOL": 0.02, + "REPRO_COMPONENT_ATOL": 0.02, "REPRO_DTW_ATOL": 0.02, + "REPRO_ACTION_MEAN_ATOL": 0.03, "OUT_BASE": ROOT / "src/drl_pinball/eval/output/train", + } + exec(compile(ast.Module(body=selected, type_ignores=[]), str(SOURCE), "exec"), namespace) + return [namespace[name] for name in names] + + +def test_vecnormalize_is_frozen_for_canonical_inference(): + (freeze,) = _load_helpers("_freeze_vecnormalize") + class FakeVecNormalize: + training = True + norm_reward = True + vec = freeze(FakeVecNormalize()) + assert vec.training is False + assert vec.norm_reward is False + + +def test_seed_signal_serialization_contains_only_compact_series(tmp_path): + (save_signals,) = _load_helpers("_save_seed_signals") + values = np.ones(4, dtype=np.float32) + result = { + "seed": "43", "sensors": np.ones((4, 6)), "forces": np.ones((4, 6)), + "actions": np.ones((4, 3)), "rewards": values, "r_cd_series": values, + "r_cl_series": values, "r_sim_series": values, "sim_raw_series": values, + } + path = save_signals(tmp_path, result) + with np.load(path) as artifact: + assert set(artifact.files) == {"sensors", "forces", "actions", "rewards", "r_cd", "r_cl", "r_sim", "sim_raw"} + assert not ({"ux", "uy", "vorticity"} & set(artifact.files)) + + +def test_evaluator_contract_is_canonical_and_gpu_imports_are_lazy(): + source = SOURCE.read_text() + assert "scene_manifest" not in source + assert "provenance" not in source + assert "from drl_pinball.case_registry import" in source + assert "NUM_STEPS = 360" in source + assert "TAIL_WINDOW = 180" in source + assert "deterministic=True" in source + assert "vmin=-0.001, vmax=0.001" in source + tree = ast.parse(source) + top_imports = [node for node in tree.body if isinstance(node, (ast.Import, ast.ImportFrom))] + assert all("pycuda" not in ast.unparse(node) for node in top_imports) + + +def test_prepare_output_preserves_existing_baseline(tmp_path): + (prepare,) = _load_helpers("_prepare_output") + existing = tmp_path / "kar_re100" + existing.mkdir() + sentinel = existing / "metrics.json" + sentinel.write_text("baseline") + try: + prepare("kar_re100", False, tmp_path) + except FileExistsError as exc: + assert "--overwrite" in str(exc) + else: + raise AssertionError("existing baseline must be protected") + assert prepare("kar_re100", True, tmp_path) == existing + assert not sentinel.exists() + + +def test_illusion_bundle_uses_verified_calibration_harmonics(tmp_path): + load_json, sha256, schema, adapt, lineage, resolve = _load_helpers( + "_load_json", "_sha256", "_calibration_schema", + "_adapt_native_illusion_products", "_normalizer_lineage", "_resolve_seed_artifacts") + lineage.__globals__["CASE_IDS"] = ("ill_1L",) + resolve.__globals__.update( + _sha256=sha256, _load_json=load_json, _calibration_schema=schema, + _adapt_native_illusion_products=adapt, _normalizer_lineage=lineage) + class Case: + case_id = "ill_1L" + seeds = (43,) + scene_type = "illusion" + si = 1200 + config_path = tmp_path / "config.json" + calibration_path = tmp_path / "calibrations/ill_1L/calibration.json" + def model_dir(self, seed): + return tmp_path / f"ill_1L_seed{seed}/models" + case = Case() + run = case.model_dir(43).parent + (run / "models").mkdir(parents=True) + case.calibration_path.parent.mkdir(parents=True) + case.calibration_path.write_text(json.dumps({ + "schema_version": "drl-pinball-calibration-v2", + "sensor_units": "celeris-area-time-average", + })) + (run / "models/best_model.zip").write_bytes(b"model") + (run / "vec_normalize.pkl").write_bytes(b"normalizer") + np.save(run / "target.npy", np.ones((150, 6), dtype=np.float32)) + np.save(case.calibration_path.parent / "target.npy", np.ones((150, 6), dtype=np.float32)) + (run / "calibration.json").write_text(json.dumps({"SENSOR_CC": 78.0})) + + try: + resolve(case, 43) + except FileNotFoundError as exc: + assert "target_harmonics.json" in str(exc) + assert "registry calibration bundle" in str(exc) + else: + raise AssertionError("missing calibration harmonics must fail") + + harmonics_path = case.calibration_path.parent / "target_harmonics.json" + harmonics = [ + {"dc": float(i + 1), "amps": [float(i + 2)], "freqs": [0.1], "phases": [0.2]} + for i in range(8) + ] + harmonics_path.write_text(json.dumps(harmonics)) + np.save(case.calibration_path.parent / "target.npy", np.zeros((150, 6), dtype=np.float32)) + try: + resolve(case, 43) + except ValueError as exc: + assert "target mismatch" in str(exc) + assert "SHA256" in str(exc) + else: + raise AssertionError("mismatched calibration target must fail") + + np.save(case.calibration_path.parent / "target.npy", np.ones((150, 6), dtype=np.float32)) + bundle = resolve(case, 43) + assert bundle["target_path"] == (run / "target.npy").resolve() + assert bundle["harmonics_path"] == harmonics_path.resolve() + assert np.array_equal(bundle["target_states"], np.full((150, 6), 78.0, dtype=np.float32)) + adapted = bundle["target_harmonics"] + assert adapted[0]["dc"] == 78.0 and adapted[0]["amps"] == [156.0] + assert adapted[0]["freqs"] == [0.1] and adapted[0]["phases"] == [0.2] + assert adapted[5]["dc"] == 468.0 + assert adapted[6] == harmonics[6] and adapted[7] == harmonics[7] + metadata = bundle["illusion_native_to_legacy_adaptation"] + assert metadata["applied"] is True and metadata["factor"] == 78.0 + assert metadata["source_target_sha256"] == metadata["registry_target_sha256"] + + + +def test_normalizer_lineage_selects_d075_best_and_others_final(tmp_path): + (lineage,) = _load_helpers("_normalizer_lineage") + lineage.__globals__["CASE_IDS"] = ("kar_d075", "kar_re60") + d075_path, d075_reason = lineage("kar_d075", tmp_path) + other_path, other_reason = lineage("kar_re60", tmp_path) + assert d075_path == tmp_path / "best_vecnormalize.pkl" + assert "A/B" in d075_reason + assert other_path == tmp_path / "vec_normalize.pkl" + assert "final compatibility alias" in other_reason + try: + lineage("unknown", tmp_path) + except ValueError as exc: + assert "No retained-artifact normalizer lineage" in str(exc) + else: + raise AssertionError("unregistered lineage must fail closed") + + +def test_action_means_use_tail_window_only(): + (build_metrics,) = _load_helpers("_build_metrics") + actions = np.vstack((np.full((180, 3), 10.0), np.array([[1.0, 2.0, 3.0]] * 180))) + best = { + "seed": "45", "avg_reward": 0.9, "r_cd": 0.8, "r_cl": 0.7, + "r_sim": 0.6, "sim_raw": 0.5, "actions": actions, + } + metrics = build_metrics("kar_re100", best) + assert (metrics["aF_mean"], metrics["aT_mean"], metrics["aB_mean"]) == (1.0, 2.0, 3.0) + assert metrics["action_mean_window"] == 180 + + +def test_metrics_only_returns_before_signals_and_vorticity_paths(): + source = SOURCE.read_text() + tree = ast.parse(source) + evaluate = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "evaluate_case") + metrics_guard = next( + node for node in evaluate.body + if isinstance(node, ast.If) and ast.unparse(node.test) == "metrics_only" + ) + assert isinstance(metrics_guard.body[-1], ast.Return) + guard_line = metrics_guard.lineno + save_seed_line = next(node.lineno for node in ast.walk(evaluate) if isinstance(node, ast.Call) and ast.unparse(node.func) == "_save_seed_signals") + replay_lines = [node.lineno for node in ast.walk(evaluate) if isinstance(node, ast.Call) and ast.unparse(node.func) in {"_save_vorticity", "_generate_target_vorticity"}] + assert save_seed_line < guard_line + save_parent = next(node for node in ast.walk(evaluate) if isinstance(node, ast.If) and any(getattr(child, "lineno", -1) == save_seed_line for child in ast.walk(node))) + assert ast.unparse(save_parent.test) == "not metrics_only" + assert replay_lines and all(line > guard_line for line in replay_lines) + +def _configure_validation(validate, load_json, tmp_path, metrics, csv_values, json_values): + metrics_dir = tmp_path / "out/kar_re100" + metrics_dir.mkdir(parents=True) + (metrics_dir / "metrics.json").write_text(json.dumps(metrics)) + fields = [ + "case_id", "selected_seed", "eval_reward_mean", "eval_r_cd_mean", + "eval_r_cl_mean", "eval_r_sim_mean", "dtw_similarity", + "action_front_mean", "action_top_mean", "action_bottom_mean", + ] + csv_path = tmp_path / "latest.csv" + csv_path.write_text( + ",".join(fields) + "\n" + + ",".join(str(csv_values[field]) for field in fields) + "\n" + ) + json_path = tmp_path / "latest.json" + json_path.write_text(json.dumps({ + "scope": {"eval_steps": 360, "action_tail_steps": 180}, + "evaluation": [json_values], + })) + validate.__globals__.update( + OUT_BASE=tmp_path / "out", REFERENCE_CSV=csv_path, + REFERENCE_JSON=json_path, _load_json=load_json, + ) + + +def test_validation_accepts_documented_fresh_re100_deltas_and_reports_them(tmp_path): + load_json, validate = _load_helpers("_load_json", "validate_outputs") + reference = { + "case_id": "kar_re100", "selected_seed": 45, + "eval_reward_mean": .931258, "eval_r_cd_mean": .981787, + "eval_r_cl_mean": .980641, "eval_r_sim_mean": .856324, + "dtw_similarity": .918458, "action_front_mean": .015262, + "action_top_mean": -.239834, "action_bottom_mean": .226212, + } + fresh = { + "best_seed": "45", "reward_mean": .937275, "r_cd_mean": .980295, + "r_cl_mean": .988632, "r_sim_mean": .866492, "dtw_sim_v5": .923246, + "aF_mean": .033458, "aT_mean": -.243172, "aB_mean": .214681, + } + _configure_validation(validate, load_json, tmp_path, fresh, reference, reference) + reports, failures = validate(("kar_re100",)) + assert failures == [] + assert len(reports) == 8 + assert all("delta=" in report and report.endswith("PASS") for report in reports) + + +def test_validation_keeps_strict_tables_and_fails_reproduction_gate(tmp_path): + load_json, validate = _load_helpers("_load_json", "validate_outputs") + reference = { + "case_id": "kar_re100", "selected_seed": 45, + "eval_reward_mean": 1.0, "eval_r_cd_mean": 2.0, + "eval_r_cl_mean": 3.0, "eval_r_sim_mean": 4.0, + "dtw_similarity": 5.0, "action_front_mean": 0.1, + "action_top_mean": 0.2, "action_bottom_mean": 0.3, + } + csv_reference = dict(reference, eval_reward_mean=1.000002) + fresh = { + "best_seed": 44, "reward_mean": 1.021, "r_cd_mean": 2.0, + "r_cl_mean": 3.0, "r_sim_mean": 4.0, "dtw_sim_v5": 5.0, + "aF_mean": 0.1, "aT_mean": 0.2, "aB_mean": 0.3, + } + _configure_validation(validate, load_json, tmp_path, fresh, csv_reference, reference) + reports, failures = validate(("kar_re100",)) + assert reports == [] + assert any("CSV/JSON disagree" in failure and "strict atol=1e-06" in failure for failure in failures) + + csv_reference["eval_reward_mean"] = reference["eval_reward_mean"] + _configure_validation(validate, load_json, tmp_path / "second", fresh, csv_reference, reference) + reports, failures = validate(("kar_re100",)) + assert any(report.endswith("FAIL") and "reward_mean" in report for report in reports) + assert any("best_seed=44" in failure for failure in failures) + assert any("exceeds reproduction atol=0.020" in failure for failure in failures) diff --git a/tests/test_drl_pinball_v5_acquire.py b/tests/test_drl_pinball_v5_acquire.py new file mode 100644 index 0000000..8f543f7 --- /dev/null +++ b/tests/test_drl_pinball_v5_acquire.py @@ -0,0 +1,565 @@ +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest + +from drl_pinball.eval import acquire_v5 + + +class FakeModel: + def __init__(self): + self.calls = 0 + + def predict(self, obs, deterministic): + assert deterministic is True + self.calls += 1 + return np.array([[0.1, -0.2, 0.3]], dtype=np.float32), None + + +class FakeVecEnv: + def __init__(self, raw): + self.raw = raw + self.steps = 0 + self.events = [] + + def reset(self): + self.events.append("reset") + return np.zeros((1, 12), dtype=np.float32) + + def step(self, action): + self.steps += 1 + self.raw.control_step = self.steps + self.raw.sim.stepper.step_count = self.steps * 800 + self.raw.smoother._state = self.raw._action_to_omega(action) * 0.5 + self.events.append(("step", self.steps)) + info = {"sim": self.steps / 1000, "r_cd": 1, "r_cl": 2, "r_sim": 3, "floor_pen": 4} + return np.zeros((1, 12)), np.array([5.0]), np.array([False]), [info] + + +class FakeRaw: + def __init__(self): + self.control_step = 0 + self._cal = {"U0": 0.01, "grid": {"nx": 2000, "ny": 600}, + "ACTION_BIAS": [0.0, 0.0, 0.0], "ACTION_SCALE": 12.0} + self.smoother = type("Smoother", (), {"_state": np.zeros(3)})() + self.sim = type("Sim", (), {})() + self.sim.stepper = type("Stepper", (), {"step_count": 0})() + + def _action_to_omega(self, action): + return np.asarray(action).reshape(3) * 2 + + def _read_obs(self): + return np.arange(14, dtype=np.float32) + self.control_step + + +def test_controlled_schedule_is_reset_750_warmup_then_250_post_step_fields(tmp_path): + raw, model = FakeRaw(), FakeModel() + vec = FakeVecEnv(raw) + captures = [] + + def capture(env): + captures.append((env.control_step, len(vec.events))) + value = np.full((2, 3), env.control_step, dtype=np.float32) + return {"rho": value, "ux": value, "uy": value} + + rows, buffer = acquire_v5._collect_controlled(model, vec, raw, tmp_path, capture) + assert vec.events[0] == "reset" + assert vec.steps == model.calls == 1000 + assert len(rows) == len(captures) == 250 + assert buffer["ux"].shape == buffer["uy"].shape == (250, 2, 3) + assert captures[0][0] == rows[0]["control_index"] == 751 + assert captures[-1][0] == rows[-1]["control_index"] == 1000 + assert all(event_count == step + 1 for step, event_count in captures) + assert rows[0]["native_reward_dtw"] == pytest.approx(0.751) + assert np.allclose(rows[0]["commanded_target_omega"], [0.2, -0.4, 0.6]) + assert np.allclose(rows[0]["effective_smoothed_omega"], [0.1, -0.2, 0.3]) + assert np.allclose(buffer["ux"][0], 751) and np.allclose(buffer["uy"][-1], 1000) + assert not list(tmp_path.glob("boundary_*.npz")) + + +def test_field_capture_runs_inside_env_cuda_context_and_validates_shape(): + events = [] + raw = type("Raw", (), {})() + raw.sim = type("Sim", (), {})() + raw.sim.lbm_cfg = type("Cfg", (), {"nx": 3, "ny": 2})() + + def macro(): + events.append("macro") + value = np.ones((2, 3), dtype=np.float32) + return {"rho": value, "ux": value, "uy": value} + + raw.sim.get_macroscopic = macro + raw._gpu_block = lambda fn: (events.append("push"), fn(), events.append("pop")) + result = acquire_v5._capture_fields(raw) + assert events == ["push", "macro", "pop"] + assert result["ux"].shape == (2, 3) + assert set(result) == {"rho", "ux", "uy"} + + +def test_zero_uses_full_vec_step_schedule_and_zero_action(tmp_path): + raw = FakeRaw() + vec = FakeVecEnv(raw) + rows, buffer = acquire_v5._collect_zero(vec, raw, tmp_path, lambda env: { + name: np.ones((2, 3), dtype=np.float32) for name in ("rho", "ux", "uy") + }) + assert vec.events[0] == "reset" + assert vec.steps == 1000 and len(rows) == 250 + assert buffer["ux"].shape == (250, 2, 3) + assert np.array_equal(rows[0]["action_normalized"], np.zeros(3, dtype=np.float32)) + assert rows[0]["native_reward_dtw"] == pytest.approx(0.751) + assert rows[-1]["control_index"] == 1000 + assert set(rows[0]) == set(acquire_v5._target_boundary(type("Target", (), { + "sensor_ids": (0, 1, 2), "calibration": {"U0": 0.01, "grid": {"nx": 2000}}, + "sim": type("Sim", (), {"stepper": type("Stepper", (), {"step_count": 800})(), + "read_sensor": lambda self, sid, normalize: (0.0, 0.0)})() + })(), 1)) + + +def test_target_geometry_schedule_order_and_nan_contract(tmp_path): + class Sim: + def __init__(self): + self.added, self.runs, self.closed = [], [], False + self._objects = [] + self.bodies = type("Bodies", (), { + "get": lambda owner, index: self._objects[index], + "count": property(lambda owner: len(self._objects)), + })() + self.stepper = type("Stepper", (), {"step_count": 0})() + self.lbm_cfg = type("Cfg", (), {"nx": 3, "ny": 2})() + context = type("Context", (), {"push": lambda self: None, "pop": lambda self: None})() + self.ctx = type("Cuda", (), {"_ctx": context})() + + def add_body(self, kind, **kwargs): + self.added.append((kind, kwargs)) + body_id = len(self.added) - 1 + self._objects.append(type("Body", (), { + "obj_id": body_id, "_is_sensor": kind == "sensor", + })()) + return body_id + + def initialize(self): + self.initialized = True + + def run(self, steps, **kwargs): + self.runs.append((steps, kwargs)) + self.stepper.step_count += steps + + def read_sensor(self, sensor_id, normalize=True): + assert normalize is True + return np.array([sensor_id + 0.1, sensor_id + 0.2]) + + def get_macroscopic(self): + value = np.ones((2, 3), dtype=np.float32) + return {"rho": value, "ux": value, "uy": value} + + def close(self): + self.closed = True + + sim = Sim() + bundle = {"calibration": {"grid": {"nx": 3, "ny": 2}, "dist_radius": 1.25, + "L0": 20.0, "U0": 0.01}, + "config_path": Path("config.json")} + case = type("Case", (), {"scene_type": "karman", "target_diam": None})() + spinups = [] + runtime = acquire_v5._create_target_runtime( + case, bundle, 2, simulation_factory=lambda **_: sim, + spinup_runner=lambda target_sim, steps: spinups.append((target_sim, steps)), + ) + assert sim.added == [ + ("circle", {"center": (600.0, 0.5, 0.0), "radius": 25.0}), + ("sensor", {"center": (1200.0, 40.5, 0.0), "radius": 5.0}), + ("sensor", {"center": (1200.0, 0.5, 0.0), "radius": 5.0}), + ("sensor", {"center": (1200.0, -39.5, 0.0), "radius": 5.0}), + ] + assert spinups == [(sim, 1200)] + rows, buffer = acquire_v5._collect_target(runtime, tmp_path, 800) + assert sim.runs == [(800, {"zero_obs": True, "sync_obs": True})] * 1000 + assert len(rows) == 250 + assert buffer["ux"].shape == buffer["uy"].shape == (250, 2, 3) + assert np.allclose(buffer["ux"][0], 1.0) + assert not list(tmp_path.glob("boundary_*.npz")) + assert np.allclose(rows[0]["sensors"], [1.1, 1.2, 2.1, 2.2, 3.1, 3.2]) + for name in ("forces", "action_normalized", "commanded_target_omega", + "effective_smoothed_omega"): + assert np.all(np.isnan(rows[0][name])) + for name in ("reward_raw", "cd", "cl", "r_cd", "r_cl", "r_sim", "floor_pen", + "native_reward_dtw"): + assert np.isnan(rows[0][name]) + runtime.close() + assert sim.closed + + +def test_finalize_converts_sensors_only_for_dtw(tmp_path, monkeypatch): + captured = {} + n = 150 + times = np.arange(n, dtype=float) + sensors = np.column_stack([np.sin(2 * np.pi * times / 30 + i) for i in range(6)]) + rows = [] + for i in range(n): + rows.append({"physical_time": float(i), "lattice_step": i * 800, + "control_index": i + 1, "sensors": sensors[i], "forces": np.ones(6), + "action_normalized": np.zeros(3), "commanded_target_omega": np.zeros(3), + "effective_smoothed_omega": np.zeros(3), "reward_raw": 1.0, + "cd": 1.0, "cl": 1.0, "r_cd": 1.0, "r_cl": 1.0, "r_sim": 1.0, + "floor_pen": 0.0, "native_reward_dtw": 1.0}) + scratch_root = tmp_path / "scratch" + scratch = scratch_root / "candidate" + scratch.mkdir(parents=True) + fields = { + "ux": np.ones((n, 2, 3), dtype=np.float32), + "uy": np.ones((n, 2, 3), dtype=np.float32), + } + identity = tmp_path / "identity" + identity.write_bytes(b"read-only") + bundle = {"target_states": sensors * 7.0, "model_path": identity, + "vecnormalize_path": identity, "config_path": identity} + original = acquire_v5.dual_cycle_dtw + + def observe(target, state, native, **kwargs): + captured["state"] = state.copy() + captured["lag_channel"] = kwargs["lag_channel"] + return original(target, state, native, **kwargs) + + monkeypatch.setattr(acquire_v5, "dual_cycle_dtw", observe) + monkeypatch.setattr(acquire_v5.infer_train, "_file_identity", lambda path: {"path": str(path)}) + monkeypatch.setattr(acquire_v5.infer_train, "_bundle_metadata", lambda bundle: {}) + case = type("Case", (), {"case_id": "kar_re100", "si": 800})() + acquire_v5._finalize(tmp_path, scratch, rows, fields, bundle, + {"resolved_output_root": tmp_path}, 7.0, "zero", [], + case=case, seed=45, cycle_length=30) + assert np.allclose(captured["state"], sensors * 7.0) + assert captured["lag_channel"] == 3 + import json + assert json.loads((tmp_path / "dtw_summary.json").read_text())["lag_channel"] == 3 + assert json.loads((tmp_path / "metadata.json").read_text())["dtw_lag_channel"] == 3 + with np.load(tmp_path / "timeseries.npz", allow_pickle=False) as saved: + assert np.allclose(saved["sensors"], sensors) + assert identity.read_bytes() == b"read-only" + + +def test_collection_failure_cleans_only_transaction_scratch(tmp_path): + scratch_root = tmp_path / "scratch" + scratch_root.mkdir() + scratch = acquire_v5.create_scratch(scratch_root) + sibling = tmp_path / "immutable-model.zip" + sibling.write_bytes(b"model") + (scratch / "partial.npz").write_bytes(b"partial") + acquire_v5.cleanup_scratch(scratch, root=scratch_root) + assert not scratch.exists() + assert sibling.read_bytes() == b"model" + + + +def test_acquire_finalize_failure_leaves_no_partial_role(tmp_path, monkeypatch): + final_role = tmp_path / "v5" / "karman_re100" / "controlled" + sentinel = tmp_path / "immutable-model.zip" + sentinel.write_bytes(b"model") + storage = {"resolved_output_root": tmp_path, "device": tmp_path.stat().st_dev} + bundle = {"model_path": sentinel, "vecnormalize_path": sentinel} + monkeypatch.setattr(acquire_v5, "get_case", lambda _: type( + "Case", (), {"case_id": "kar_re100", "scene_type": "karman", "si": 800, + "seeds": (45,)})()) + monkeypatch.setattr(acquire_v5.infer_train, "_resolve_seed_artifacts", lambda *_: bundle) + monkeypatch.setattr(acquire_v5, "_validate_acquisition_bundle", lambda *_: 30) + monkeypatch.setattr(acquire_v5, "_validate_shared_role_identity", lambda *_: None) + monkeypatch.setattr(acquire_v5, "_collect_controlled", lambda *_ , **__: ([], [])) + + class Env: + def close(self): + pass + + runtime = lambda *_: (Env(), object(), object(), type("Raw", (), {})()) + + def fail(staging, *_, **__): + (staging / "timeseries.npz").write_bytes(b"partial") + raise RuntimeError("injected finalize failure") + + with pytest.raises(RuntimeError, match="injected finalize failure"): + acquire_v5.acquire_controlled( + output_root=tmp_path, overwrite=True, + storage_validator=lambda **_: storage, runtime_factory=runtime, finalizer=fail, + ) + assert not final_role.exists() + case_dir = final_role.parent + assert not case_dir.exists() or list(case_dir.iterdir()) == [] + assert sentinel.read_bytes() == b"model" + + + +def test_target_full_finalize_exact_products_and_unavailable_metadata(tmp_path, monkeypatch): + n = 150 + times = np.arange(n, dtype=float) + sensors = np.column_stack([np.sin(2 * np.pi * times / 30 + i) for i in range(6)]) + rows = [] + for i in range(n): + nan3, nan6 = np.full(3, np.nan), np.full(6, np.nan) + rows.append({"physical_time": float(i), "lattice_step": i * 800, + "control_index": i + 1, "sensors": sensors[i], "forces": nan6, + "action_normalized": nan3, "commanded_target_omega": nan3, + "effective_smoothed_omega": nan3, "reward_raw": np.nan, + "cd": np.nan, "cl": np.nan, "r_cd": np.nan, "r_cl": np.nan, + "r_sim": np.nan, "floor_pen": np.nan, "native_reward_dtw": np.nan}) + role_dir = tmp_path / "role" + role_dir.mkdir() + scratch = role_dir / "scratch" / "candidate" + scratch.mkdir(parents=True) + fields = { + "ux": np.ones((n, 2, 3), dtype=np.float32), + "uy": np.ones((n, 2, 3), dtype=np.float32), + } + identity = tmp_path / "identity" + identity.write_bytes(b"read-only") + unavailable = ["forces", "action_normalized", "reward_raw", "native_reward_dtw"] + monkeypatch.setattr(acquire_v5.infer_train, "_file_identity", lambda path: {"path": str(path)}) + monkeypatch.setattr(acquire_v5.infer_train, "_bundle_metadata", lambda bundle: {}) + acquire_v5._finalize( + role_dir, scratch, rows, fields, + {"target_states": sensors, "model_path": identity, "vecnormalize_path": identity, + "config_path": identity}, + {"resolved_output_root": tmp_path}, 1.0, "target", unavailable, + case=type("Case", (), {"case_id": "kar_re100", "si": 800})(), + seed=45, cycle_length=30, + ) + monkeypatch.setattr(acquire_v5, "COLLECT_BOUNDARIES", n) + acquire_v5._validate_staged_role(role_dir) + expected = {"timeseries.npz", "timeseries.csv", "phase_cycle.npz", "phase_cycle.csv", + "phase_fields.npz", "dtw_summary.json", "metadata.json", "identity"} + expected.remove("identity") + assert {path.name for path in role_dir.iterdir()} == expected + import json + metadata = json.loads((role_dir / "metadata.json").read_text()) + summary = json.loads((role_dir / "dtw_summary.json").read_text()) + assert metadata["role"] == "target" and metadata["seed"] is None + assert metadata["unavailable_fields"] == unavailable + assert metadata["candidate_field_storage"].startswith("single-role in-memory") + assert metadata["phase_smoothing_kernel"] == [0.25, 0.5, 0.25] + assert metadata["mean_field_count"] > 0 + assert summary["native_mean"] is None + with np.load(role_dir / "timeseries.npz", allow_pickle=False) as saved: + assert np.all(np.isnan(saved["native_reward_dtw"])) + assert np.allclose(saved["sensors"], sensors) + with np.load(role_dir / "phase_cycle.npz", allow_pickle=False) as saved: + assert "sensors_pooled" in saved.files and "reward_raw_mean" in saved.files + assert len(saved["sensors_pooled"]) > 0 + with np.load(role_dir / "phase_fields.npz", allow_pickle=False) as saved: + assert set(saved.files) == acquire_v5.PHASE_FIELD_KEYS + assert saved["mean_ux"].shape == saved["ux"].shape[1:] + assert saved["mean_uy"].shape == saved["uy"].shape[1:] + +def test_cli_enables_all_roles_without_replay_rejection(): + source = Path(acquire_v5.__file__).read_text() + assert 'parser.add_argument("--case", choices=CASE_IDS' in source + assert 'parser.add_argument("--seed", type=int)' in source + assert 'parser.add_argument("--role", choices=ROLES' in source + assert "replay is not implemented" not in source + assert "acquire_role(args.role" in source + + +def test_target_unavailable_summary_is_strict_json(tmp_path): + path = tmp_path / "summary.json" + acquire_v5._atomic_json(path, {"native_mean": None, "unavailable_fields": ["reward_raw"]}) + text = path.read_text() + assert "NaN" not in text and '"native_mean": null' in text + + + +def test_scene_aware_raw_sample_layouts(): + karman = type("Raw", (), {"_read_obs": lambda self: np.arange(14, dtype=np.float32)})() + illusion = type("Raw", (), {"_read_obs": lambda self: np.arange(12, dtype=np.float32)})() + assert np.array_equal(acquire_v5._raw_sample(karman, "karman"), np.arange(2, 14)) + assert np.array_equal(acquire_v5._raw_sample(illusion, "illusion"), np.arange(12)) + with pytest.raises(ValueError, match="6-sensor/6-force"): + acquire_v5._raw_sample(illusion, "karman") + + +def test_illusion_target_geometry_and_case_si(tmp_path): + class Sim: + def __init__(self): + self.added, self.runs = [], [] + self._objects = [] + self.bodies = type("Bodies", (), { + "get": lambda owner, index: self._objects[index], + "count": property(lambda owner: len(self._objects)), + })() + self.stepper = type("Stepper", (), {"step_count": 0})() + self.lbm_cfg = type("Cfg", (), {"nx": 3, "ny": 2})() + def add_body(self, kind, **kwargs): + self.added.append((kind, kwargs)) + body_id = len(self.added) - 1 + self._objects.append(type("Body", (), { + "obj_id": body_id, "_is_sensor": kind == "sensor", + })()) + return body_id + def initialize(self): pass + def run(self, steps, **kwargs): + self.runs.append(steps); self.stepper.step_count += steps + def read_sensor(self, sensor_id, normalize=True): return (0.1, 0.2) + def get_macroscopic(self): + value = np.ones((2, 3), dtype=np.float32) + return {name: value for name in ("rho", "ux", "uy")} + def close(self): pass + sim = Sim() + context = type("Context", (), {"push": lambda self: None, "pop": lambda self: None})() + sim.ctx = type("Cuda", (), {"_ctx": context})() + case = type("Case", (), {"scene_type": "illusion", "target_diam": 1.5})() + bundle = {"calibration": {"grid": {"nx": 3, "ny": 2}, "L0": 20.0, "U0": 0.01}, + "config_path": Path("config.json")} + spinups = [] + runtime = acquire_v5._create_target_runtime( + case, bundle, 0, lambda **_: sim, + spinup_runner=lambda target_sim, steps: spinups.append((target_sim, steps)), + ) + assert sim.added[0] == ("circle", {"center": (400.0, 0.5, 0.0), "radius": 30.0}) + assert [item[1]["center"][0] for item in sim.added[1:]] == [600.0] * 3 + assert spinups == [(sim, 1200)] + acquire_v5._collect_target(runtime, tmp_path, 1200) + assert sim.runs == [1200] * 1000 + + +def test_bundle_validation_covers_registry_and_fails_before_storage(tmp_path, monkeypatch): + target = np.zeros((150, 6), dtype=np.float32) + phase = 2 * np.pi * np.arange(150) / 30 + target[:, 3] = np.sin(phase) + config = tmp_path / "config.json" + calibration = tmp_path / "calibration.json" + config.write_text('{"grid":{"nx":2000,"ny":600},"physics":{"velocity":0.01}}') + calibration.write_text('{"SI":800}') + case = type("Case", (), { + "case_id": "kar_re100", "scene_type": "karman", "si": 800, + "seeds": (45,), "target_diam": None, "config_path": config, + })() + bundle = {"seed": "45", "config_path": config, "calibration_path": calibration, + "calibration": {"SI": 800, "U0": 0.01, "grid": {"nx": 2000, "ny": 600}}, + "target_states": target} + assert set(acquire_v5.CYCLE_WINDOWS) == set(acquire_v5.CASE_IDS) + assert acquire_v5._validate_acquisition_bundle(case, 45, bundle) == 30 + calibration.write_text('{"SI":500}') + with pytest.raises(ValueError, match="SI"): + acquire_v5._validate_acquisition_bundle(case, 45, bundle) + + +def test_output_paths_seed_qualify_controlled_only(tmp_path, monkeypatch): + case = type("Case", (), {"case_id": "kar_re100", "scene_type": "karman", + "si": 800, "seeds": (45,), "target_diam": None})() + bundle = {"seed": "45"} + storage = {"resolved_output_root": tmp_path, "device": tmp_path.stat().st_dev} + monkeypatch.setattr(acquire_v5, "get_case", lambda _: case) + monkeypatch.setattr(acquire_v5.infer_train, "_resolve_seed_artifacts", lambda *_: bundle) + monkeypatch.setattr(acquire_v5, "_validate_acquisition_bundle", lambda *_: 30) + monkeypatch.setattr(acquire_v5, "_validate_staged_role", lambda *_: None) + monkeypatch.setattr(acquire_v5, "publish_role_output", lambda prepared: prepared["final_role_dir"]) + class Env: + def close(self): pass + monkeypatch.setattr(acquire_v5, "_collect_controlled", lambda *_, **__: ([], [])) + def finalize(role_dir, scratch, *args, **kwargs): + acquire_v5.cleanup_scratch(scratch, root=role_dir / "scratch") + scratch.parent.rmdir() + result = acquire_v5.acquire_role( + "controlled", case_id="kar_re100", seed=45, output_root=tmp_path, + storage_validator=lambda **_: storage, + runtime_factory=lambda *_: (Env(), object(), object(), object()), finalizer=finalize, + ) + assert result == tmp_path / "v5/kar_re100_seed45/controlled" + + + +def test_shared_roles_require_seed_invariant_physical_identity(monkeypatch): + case = type("Case", (), {"case_id": "kar_re100", "seeds": (41, 42)})() + target = np.ones((150, 6), dtype=np.float32) + base = {"seed": "41", "calibration": {"SI": 800, "config_path": "old"}, + "target_states": target, "config_path": Path("config.json")} + other = {"seed": "42", "calibration": {"SI": 800, "config_path": "new"}, + "target_states": target.copy(), "config_path": Path("config.json")} + monkeypatch.setattr(acquire_v5, "_validate_acquisition_bundle", lambda *_: 30) + monkeypatch.setattr(acquire_v5.infer_train, "_resolve_seed_artifacts", lambda *_: other) + acquire_v5._validate_shared_role_identity(case, 41, base) + other["target_states"] = target + np.float32(1e-3) + with pytest.raises(ValueError, match="different physical target"): + acquire_v5._validate_shared_role_identity(case, 41, base) + + + +def test_target_runtime_fails_closed_on_body_id_order_and_count(): + class Bodies: + def __init__(self, sim): self.sim = sim + @property + def count(self): return len(self.sim.objects) + def get(self, index): return self.sim.objects[index] + class BadSim: + def __init__(self): + self.objects = [] + self.bodies = Bodies(self) + def add_body(self, kind, **kwargs): + body_id = len(self.objects) + 1 + self.objects.append(type("Body", (), { + "obj_id": body_id, "_is_sensor": kind == "sensor", + })()) + return body_id + def initialize(self): raise AssertionError("must fail before initialize") + case = type("Case", (), {"scene_type": "karman", "target_diam": None})() + bundle = {"calibration": {"grid": {"nx": 2000, "ny": 600}, "U0": 0.01, + "L0": 20.0}, "config_path": Path("config.json")} + with pytest.raises(ValueError, match="body order"): + acquire_v5._create_target_runtime(case, bundle, 0, lambda **_: BadSim()) + + +def test_physical_zero_counterbias_shape_and_rollout(tmp_path): + class BiasedRaw(FakeRaw): + def __init__(self): + super().__init__() + self._cal.update(ACTION_BIAS=[1.5, -3.0, 0.75], ACTION_SCALE=6.0) + def _action_to_omega(self, action): + action = np.asarray(action, dtype=np.float32).reshape(3) + return action * self._cal["ACTION_SCALE"] + np.asarray( + self._cal["ACTION_BIAS"], dtype=np.float32) + raw = BiasedRaw() + vec = FakeVecEnv(raw) + actions = [] + original_step = vec.step + def step(action): + actions.append(np.asarray(action).copy()) + return original_step(action) + vec.step = step + acquire_v5._collect_zero(vec, raw, tmp_path, lambda env: { + name: np.ones((2, 3), dtype=np.float32) for name in ("rho", "ux", "uy") + }) + expected = np.array([[-0.25, 0.5, -0.125]], dtype=np.float32) + assert actions and all(action.shape == (1, 3) for action in actions) + assert all(np.array_equal(action, expected) for action in actions) + + +def test_acquire_role_propagates_case_seed_si_and_scene(tmp_path, monkeypatch): + case = type("Case", (), {"case_id": "ill_1L", "scene_type": "illusion", + "si": 1200, "seeds": (43,), "target_diam": 1.0})() + bundle = {"seed": "43"} + storage = {"resolved_output_root": tmp_path, "device": tmp_path.stat().st_dev} + observed = {} + monkeypatch.setattr(acquire_v5, "get_case", lambda case_id: case) + monkeypatch.setattr(acquire_v5.infer_train, "_resolve_seed_artifacts", + lambda selected_case, seed: bundle) + monkeypatch.setattr(acquire_v5, "_validate_acquisition_bundle", lambda *args: 19) + monkeypatch.setattr(acquire_v5, "_validate_staged_role", lambda *_: None) + monkeypatch.setattr(acquire_v5, "publish_role_output", lambda prepared: prepared["final_role_dir"]) + class Env: + def close(self): pass + def collect(model, vec, raw, scratch, **kwargs): + observed["scene_type"] = kwargs["scene_type"] + return [], [] + monkeypatch.setattr(acquire_v5, "_collect_controlled", collect) + def finalize(role_dir, scratch, *args, **kwargs): + observed.update(case=kwargs["case"], seed=kwargs["seed"], cycle=kwargs["cycle_length"]) + acquire_v5.cleanup_scratch(scratch, root=role_dir / "scratch") + scratch.parent.rmdir() + result = acquire_v5.acquire_role( + "controlled", case_id="ill_1L", seed=43, output_root=tmp_path, + storage_validator=lambda **_: storage, + runtime_factory=lambda selected_case, selected_bundle, device: ( + Env(), object(), object(), type("Raw", (), {"_dtw_sensor_factor": 78.0})()), + finalizer=finalize, + ) + assert observed == {"scene_type": "illusion", "case": case, "seed": 43, "cycle": 19} + assert result == tmp_path / "v5/ill_1L_seed43/controlled"