"""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())