Frank14fandCursor
61e82ec90a
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 <cursoragent@cursor.com>
403 lines
17 KiB
Python
403 lines
17 KiB
Python
import csv
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from collections import namedtuple
|
|
from pathlib import Path
|
|
|
|
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,
|
|
dual_cycle_dtw,
|
|
pooled_phase_bins,
|
|
prepare_role_output,
|
|
publish_role_output,
|
|
publish_selected_fields,
|
|
select_phase_fields,
|
|
smooth_center_uy,
|
|
validate_output_storage,
|
|
validate_modern_fields,
|
|
write_boundary_artifacts,
|
|
write_phase_cycle_artifacts,
|
|
)
|
|
|
|
Usage = namedtuple("Usage", "total used free")
|
|
|
|
|
|
def _storage_tree(tmp_path):
|
|
mount = tmp_path / "optane"
|
|
expected = mount / "DynamisLab" / "drl_pinball" / "reproduction"
|
|
expected.mkdir(parents=True)
|
|
mapping = tmp_path / "repo" / "src" / "drl_pinball" / "data" / "reproduction"
|
|
mapping.parent.mkdir(parents=True)
|
|
mapping.symlink_to(expected, target_is_directory=True)
|
|
return mount, expected, mapping
|
|
|
|
|
|
def _validate_tmp_storage(tmp_path, **overrides):
|
|
mount, expected, mapping = _storage_tree(tmp_path)
|
|
kwargs = {
|
|
"repo_mapping": mapping, "expected_root": expected, "optane_mount": mount,
|
|
"min_free_bytes": 100, "is_mount": lambda path: path == mount,
|
|
"disk_usage": lambda path: Usage(1000, 100, 900),
|
|
}
|
|
kwargs.update(overrides)
|
|
return validate_output_storage(**kwargs), mount, expected, mapping
|
|
|
|
|
|
def test_default_reproduction_mapping_is_derivable(tmp_path):
|
|
assert default_reproduction_mapping(tmp_path) == (
|
|
tmp_path / "src" / "drl_pinball" / "data" / "reproduction"
|
|
)
|
|
|
|
|
|
def test_storage_validation_defaults_output_to_repo_mapping(tmp_path):
|
|
storage, _, expected, mapping = _validate_tmp_storage(tmp_path)
|
|
assert storage["output_root"] == mapping
|
|
assert storage["resolved_output_root"] == expected
|
|
def test_storage_validation_accepts_exact_mapping_and_configurable_same_device(tmp_path):
|
|
|
|
mount, expected, mapping = _storage_tree(tmp_path)
|
|
output = mount / "alternate" / "campaign"
|
|
result = validate_output_storage(
|
|
repo_mapping=mapping, output_root=output, expected_root=expected,
|
|
optane_mount=mount, min_free_bytes=500, is_mount=lambda path: path == mount,
|
|
disk_usage=lambda path: Usage(1000, 100, 900),
|
|
)
|
|
assert result["resolved_output_root"] == output
|
|
assert result["free_bytes"] == 900
|
|
|
|
|
|
def test_storage_validation_rejects_mount_space_mapping_and_device_failures(tmp_path):
|
|
mount, expected, mapping = _storage_tree(tmp_path)
|
|
common = dict(repo_mapping=mapping, expected_root=expected, optane_mount=mount)
|
|
with pytest.raises(ValueError, match="not a mount"):
|
|
validate_output_storage(**common, is_mount=lambda path: False)
|
|
with pytest.raises(OSError, match="free space"):
|
|
validate_output_storage(
|
|
**common, min_free_bytes=901, is_mount=lambda path: True,
|
|
disk_usage=lambda path: Usage(1000, 100, 900),
|
|
)
|
|
wrong = tmp_path / "wrong"
|
|
wrong.mkdir()
|
|
mapping.unlink()
|
|
mapping.symlink_to(wrong, target_is_directory=True)
|
|
with pytest.raises(ValueError, match="resolve exactly"):
|
|
validate_output_storage(**common, is_mount=lambda path: True)
|
|
|
|
mapping.unlink()
|
|
mapping.symlink_to(expected, target_is_directory=True)
|
|
real_stat = os.stat
|
|
outside = tmp_path / "outside"
|
|
outside.mkdir()
|
|
|
|
def different_device(path):
|
|
value = real_stat(path)
|
|
if Path(path) == outside:
|
|
return type("Stat", (), {"st_dev": value.st_dev + 1})()
|
|
return value
|
|
|
|
with pytest.raises(ValueError, match="output root"):
|
|
validate_output_storage(
|
|
**common, output_root=outside / "new", is_mount=lambda path: True,
|
|
stat=different_device,
|
|
)
|
|
|
|
|
|
def test_prepare_role_output_conflict_and_scoped_overwrite(tmp_path):
|
|
storage, _, _, _ = _validate_tmp_storage(tmp_path)
|
|
prepared = prepare_role_output(storage, "v5", "karman_re100", "controlled")
|
|
(prepared["role_dir"] / "old.txt").write_text("old")
|
|
published = publish_role_output(prepared)
|
|
sentinel = published / "old.txt"
|
|
sibling = published.parent / "target"
|
|
sibling.mkdir()
|
|
(sibling / "keep.txt").write_text("keep")
|
|
with pytest.raises(FileExistsError):
|
|
prepare_role_output(storage, "v5", "karman_re100", "controlled")
|
|
replaced = prepare_role_output(storage, "v5", "karman_re100", "controlled", overwrite=True)
|
|
assert sentinel.exists()
|
|
(replaced["role_dir"] / "new.txt").write_text("new")
|
|
final = publish_role_output(replaced)
|
|
assert not sentinel.exists() and (final / "new.txt").exists()
|
|
assert (sibling / "keep.txt").read_text() == "keep"
|
|
with pytest.raises(ValueError, match="path component"):
|
|
prepare_role_output(storage, "v5", "karman_re100", "../target", overwrite=True)
|
|
|
|
|
|
def test_prepare_rejects_symlink_parent_without_following_it(tmp_path):
|
|
storage, _, _, _ = _validate_tmp_storage(tmp_path)
|
|
target = tmp_path / "external-pipeline"
|
|
target.mkdir()
|
|
pipeline = storage["resolved_output_root"] / "v5"
|
|
pipeline.symlink_to(target, target_is_directory=True)
|
|
with pytest.raises(ValueError, match="pipeline directory"):
|
|
prepare_role_output(storage, "v5", "karman_re100", "controlled", overwrite=True)
|
|
assert list(target.iterdir()) == []
|
|
def test_prepare_rejects_symlink_role_without_following_it(tmp_path):
|
|
|
|
storage, _, _, _ = _validate_tmp_storage(tmp_path)
|
|
target = tmp_path / "do-not-delete"
|
|
target.mkdir()
|
|
sentinel = target / "keep.txt"
|
|
sentinel.write_text("keep")
|
|
role = storage["resolved_output_root"] / "legacy" / "karman_re100" / "zero"
|
|
role.parent.mkdir(parents=True)
|
|
role.symlink_to(target, target_is_directory=True)
|
|
with pytest.raises(ValueError, match="must not be a symlink"):
|
|
prepare_role_output(storage, "legacy", "karman_re100", "zero", overwrite=True)
|
|
assert sentinel.read_text() == "keep"
|
|
|
|
|
|
def test_phase_complete_interpolated_half_open():
|
|
times = np.arange(13, dtype=float)
|
|
sensors = np.zeros((13, 6))
|
|
sensors[:, 3] = np.sin(2 * np.pi * (times - 0.25) / 4)
|
|
result = assign_phase(times, sensors)
|
|
assert np.allclose(result["crossing_times"], [0.29289322, 4.29289322, 8.29289322])
|
|
assert result["cycle_id"].tolist() == [-1] + [0] * 4 + [1] * 4 + [-1] * 4
|
|
assert np.isnan(result["phase"][9])
|
|
|
|
|
|
def test_crossing_exact_zero_and_plateau_are_single_crossings():
|
|
times = np.arange(10, dtype=float)
|
|
sensors = np.zeros((10, 6))
|
|
sensors[:, 3] = [-1, 0, 0, 1, 1, -1, 0, 0, 1, 1]
|
|
result = assign_phase(times, sensors)
|
|
assert np.array_equal(result["crossing_times"], [2.0, 7.0])
|
|
assert result["cycle_id"].tolist() == [-1, -1, 0, 0, 0, 0, 0, -1, -1, -1]
|
|
assert result["phase"][2] == 0.0
|
|
|
|
|
|
def test_selection_is_global_eight_with_metadata_and_index_tie():
|
|
field_times = np.concatenate((np.arange(0.5, 4.0, 0.5), np.arange(4.0, 8.0, 0.5)))
|
|
result = select_phase_fields(field_times, np.array([0.0, 4.0, 8.0]))
|
|
assert result["field_indices"].shape == (8,)
|
|
assert np.array_equal(result["field_indices"], [7, 0, 1, 2, 3, 4, 5, 6])
|
|
assert np.array_equal(result["cycle_id"], [1, 0, 0, 0, 0, 0, 0, 0])
|
|
assert np.allclose(result["actual_phase"], np.arange(8) * np.pi / 4, atol=1e-15)
|
|
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]),
|
|
{"sensors": np.array([[1.0, 10.0], [4.0, 40.0], [3.0, 30.0], [9.0, 90.0]])},
|
|
)
|
|
assert result["phase_bin_count"][:2].tolist() == [2, 1]
|
|
assert np.allclose(result["mean"]["sensors"][0], [2, 20])
|
|
assert np.allclose(result["std"]["sensors"][0], [1, 10])
|
|
assert result["sorted_index"].tolist() == [0, 2, 1]
|
|
assert np.array_equal(result["pooled"]["sensors"], [[1, 10], [3, 30], [4, 40]])
|
|
|
|
|
|
def test_named_phase_pooling_preserves_rows_with_unavailable_nan_variables():
|
|
phase = np.array([0.1, 0.2, 0.3])
|
|
result = pooled_phase_bins(phase, {
|
|
"sensors": np.array([[1.0], [2.0], [3.0]]),
|
|
"reward_raw": np.full(3, np.nan),
|
|
})
|
|
assert result["sorted_index"].tolist() == [0, 1, 2]
|
|
assert result["phase_bin_count"].sum() == 3
|
|
assert np.all(np.isnan(result["mean"]["reward_raw"]))
|
|
assert np.any(np.isfinite(result["mean"]["sensors"]))
|
|
|
|
|
|
def test_dual_dtw_preserves_native_series_and_separates_rolling_normalized():
|
|
step = np.arange(90)
|
|
target = np.column_stack(
|
|
[(channel + 1) * np.sin(2 * np.pi * step / 30 + channel) for channel in range(6)]
|
|
)
|
|
native = np.linspace(0.125, 0.875, 90, dtype=np.float32)
|
|
native_bytes = native.tobytes()
|
|
result = dual_cycle_dtw(target, target, native, cycle_length=30, epsilon=1e-9)
|
|
assert result["native_reward_dtw"] is native
|
|
assert result["native_reward_dtw"].tobytes() == native_bytes
|
|
assert np.all(np.isnan(result["target_normalized_dtw"][:29]))
|
|
assert np.all(np.isfinite(result["target_normalized_dtw"][29:]))
|
|
assert np.allclose(result["target_normalized_dtw"][[29, 59, 89]], 1.0)
|
|
assert result["metadata"]["lag_channel"] == 1
|
|
assert result["metadata"]["epsilon"] == 1e-9
|
|
assert np.array_equal(result["metadata"]["target_channels"], np.arange(6))
|
|
assert np.array_equal(result["metadata"]["state_channels"], np.arange(6))
|
|
|
|
|
|
def test_dual_dtw_unavailable_native_is_nan_series():
|
|
step = np.arange(60)
|
|
target = np.column_stack([np.sin(2 * np.pi * step / 20 + channel) for channel in range(6)])
|
|
result = dual_cycle_dtw(target, target, cycle_length=20)
|
|
assert result["native_reward_dtw"].shape == (60,)
|
|
assert np.all(np.isnan(result["native_reward_dtw"]))
|
|
assert result["target_normalized_dtw"].shape == (60,)
|
|
|
|
|
|
def test_boundary_and_complete_phase_cycle_artifact_schemas(tmp_path):
|
|
boundary = {
|
|
"time": np.array([0.0, 1.0]),
|
|
"sensor": np.array([[1.0, np.nan], [2.0, 3.0]]),
|
|
"field": np.array([4.0, np.nan]),
|
|
}
|
|
npz_path, csv_path = write_boundary_artifacts(
|
|
tmp_path / "raw.npz", tmp_path / "raw.csv", boundary
|
|
)
|
|
with np.load(npz_path, allow_pickle=False) as saved:
|
|
assert np.isnan(saved["sensor"][0, 1])
|
|
with csv_path.open(newline="") as stream:
|
|
rows = list(csv.DictReader(stream))
|
|
assert list(rows[0]) == ["time", "sensor_0", "sensor_1", "field"]
|
|
assert rows[0]["sensor_1"] == "nan"
|
|
|
|
folded = pooled_phase_bins(
|
|
np.array([0.1, 0.2]), {"sensors": np.array([[1.0, 10.0], [3.0, 30.0]])}
|
|
)
|
|
npz_path, csv_path = write_phase_cycle_artifacts(
|
|
tmp_path / "phase_cycle.npz", tmp_path / "phase_cycle.csv", folded
|
|
)
|
|
with np.load(npz_path, allow_pickle=False) as saved:
|
|
assert set(saved.files) == {"sorted_index", "sorted_phase", "phase_center",
|
|
"phase_bin_count", "sensors_pooled", "sensors_mean",
|
|
"sensors_std"}
|
|
assert np.array_equal(saved["sorted_index"], [0, 1])
|
|
assert np.array_equal(saved["sensors_pooled"], [[1.0, 10.0], [3.0, 30.0]])
|
|
assert len(saved["phase_bin_count"]) == 32
|
|
with csv_path.open(newline="") as stream:
|
|
rows = list(csv.DictReader(stream))
|
|
assert {"sensors_0_pooled", "sensors_1_mean", "sensors_1_std", "phase_bin_count"}.issubset(rows[0])
|
|
assert [row["representation"] for row in rows[:2]] == ["pooled", "pooled"]
|
|
assert all(row["representation"] == "bin" for row in rows[2:])
|
|
assert len(rows[2:]) == 32
|
|
assert [row["original_index"] for row in rows[:2]] == ["0", "1"]
|
|
|
|
|
|
def test_publication_and_cleanup(tmp_path):
|
|
root = tmp_path / "scratch"
|
|
scratch = create_scratch(root)
|
|
(scratch / "x").write_text("x")
|
|
destination = tmp_path / "pub" / "fields.npz"
|
|
publish_selected_fields(destination, {"ux": np.ones((2, 3, 4))})
|
|
with pytest.raises(FileExistsError):
|
|
publish_selected_fields(destination, {"ux": np.ones((1, 1, 1))})
|
|
cleanup_scratch(scratch, root=root)
|
|
assert not scratch.exists()
|
|
with pytest.raises(ValueError):
|
|
cleanup_scratch(root, root=root)
|
|
|
|
|
|
def test_legacy_physical_decoder_is_mask_aware_and_pressure_invariant():
|
|
ny, nx = 2, 3
|
|
populations = np.zeros((9, ny, nx), dtype=np.float32)
|
|
populations[1] = 0.01
|
|
populations[2] = -0.02
|
|
flags = np.ones((nx, ny), dtype=np.uint8)
|
|
flags[1, 0] = np.uint8(0b00010001) # sensor plus FLUID remains fluid
|
|
flags[2, 1] = np.uint8(0b00000010)
|
|
populations[:, 1, 2] = np.nan # ignored nonfluid garbage
|
|
|
|
first = decode_legacy_physical_velocity(
|
|
populations.reshape(-1), flags=flags, nx=nx, ny=ny,
|
|
include_pressure=True,
|
|
)
|
|
populations[0, :, :] = 7.0 # pressure/f0 changes, velocity does not
|
|
populations[0, 1, 2] = np.nan
|
|
second = decode_legacy_physical_velocity(
|
|
populations.reshape(-1), flags=flags, nx=nx, ny=ny,
|
|
include_pressure=True,
|
|
)
|
|
|
|
fluid = ((flags & 1) != 0).T
|
|
assert first["ux"].shape == (ny, nx)
|
|
assert np.allclose(first["ux"][fluid], 0.01)
|
|
assert np.allclose(first["uy"][fluid], -0.02)
|
|
assert np.array_equal(first["ux"], second["ux"])
|
|
assert np.array_equal(first["uy"], second["uy"])
|
|
assert first["ux"][1, 2] == first["uy"][1, 2] == first["pressure"][1, 2] == 0.0
|
|
assert not np.array_equal(first["pressure"][fluid], second["pressure"][fluid])
|
|
|
|
|
|
def test_legacy_physical_decoder_accepts_zero_sum_nonzero_momentum():
|
|
populations = np.zeros((9, 1, 1), dtype=np.float32)
|
|
populations[1, 0, 0] = 0.01
|
|
populations[0, 0, 0] = -0.01
|
|
decoded = decode_legacy_physical_velocity(
|
|
populations.reshape(-1), flags=np.ones((1, 1), np.uint8), nx=1, ny=1,
|
|
)
|
|
assert decoded["ux"][0, 0] == pytest.approx(0.01)
|
|
|
|
|
|
def test_legacy_physical_decoder_rejects_nonfinite_fluid_populations():
|
|
populations = np.zeros((9, 1, 1), dtype=np.float32)
|
|
populations[0, 0, 0] = np.nan
|
|
with pytest.raises(ValueError, match="fluid populations.*non-finite"):
|
|
decode_legacy_physical_velocity(
|
|
populations.reshape(-1), flags=np.ones((1, 1), np.uint8), nx=1, ny=1,
|
|
)
|
|
|
|
def test_module_import_does_not_load_gpu_dependencies():
|
|
code = (
|
|
"import sys; import drl_pinball.acquisition; "
|
|
"assert 'pycuda' not in sys.modules; assert 'torch' not in sys.modules"
|
|
)
|
|
subprocess.run([sys.executable, "-c", code], check=True)
|