diff --git a/LegacyCelerisLab/driver.py b/LegacyCelerisLab/driver.py index 153770c..4c1b921 100644 --- a/LegacyCelerisLab/driver.py +++ b/LegacyCelerisLab/driver.py @@ -111,6 +111,7 @@ class FlowField: self.objects = {} self.action = np.zeros(0, dtype=self.DATA_TYPE) self.obs = np.zeros(0, dtype=self.DATA_TYPE) + self._control_interval = None initflow( self.flag_gpu, @@ -337,26 +338,55 @@ class FlowField: # if type == "taylor": # self.vortex_config[6] = - def run(self, num_steps: int, action_target: np.ndarray): - if ( - action_target.size != len(self.objects) - or action_target.dtype != self.DATA_TYPE - ): + def _validate_run(self, num_steps: int, action_target: np.ndarray): + if type(num_steps) is not int or num_steps < 1: + raise ValueError("num_steps must be a positive integer") + if action_target.size != len(self.objects) or action_target.dtype != self.DATA_TYPE: raise ValueError("action data type or size does not match the objects.") - elif len(self.objects) == 0: + if len(self.objects) == 0: raise ValueError("No objects have been added to the flow field.") - weight = 0.1 - stream = cuda.Stream() - action_pinned = cuda.pagelocked_empty_like(self.action) - action_pinned[:] = self.action - obs_pinned = cuda.pagelocked_empty_like(self.obs) + def run(self, num_steps: int, action_target: np.ndarray): + """Advance one complete legacy interval (original public behavior).""" + if self._control_interval is not None: + raise RuntimeError("run is unavailable while a control interval is active") + self.begin_control_interval(num_steps, action_target) + self.run_control_segment(num_steps) + self.end_control_interval() + + def begin_control_interval(self, total_steps: int, action_target: np.ndarray): + """Start one policy interval that may be split only to read/save fields.""" + if self._control_interval is not None: + raise RuntimeError("a control interval is already active") + self._validate_run(total_steps, action_target) self.error_flag[0] = 0 cuda.memcpy_htod(self.error_flag_gpu, self.error_flag) self.obs[:] = 0 - for i in range(num_steps): - action_pinned = (1 - weight) * action_pinned + weight * action_target - cuda.memcpy_htod_async(self.action_gpu, action_pinned, stream) + action = cuda.pagelocked_empty_like(self.action) + action[:] = self.action + self._control_interval = { + "total_steps": total_steps, + "completed_steps": 0, + "target": action_target.copy(), + "action": action, + "obs_steps": cuda.pagelocked_empty((total_steps, self.obs.size), dtype=self.DATA_TYPE), + "stream": cuda.Stream(), + } + + def run_control_segment(self, num_steps: int): + """Advance part of the active interval without resetting smoothing or obs.""" + state = self._control_interval + if state is None: + raise RuntimeError("no control interval is active") + if type(num_steps) is not int or num_steps < 1: + raise ValueError("num_steps must be a positive integer") + if state["completed_steps"] + num_steps > state["total_steps"]: + raise ValueError("segment exceeds the active control interval") + stream = state["stream"] + start = state["completed_steps"] + for local_step in range(num_steps): + state["action"] = 0.9 * state["action"] + 0.1 * state["target"] + cuda.memcpy_htod_async(self.action_gpu, state["action"], stream) self.step( self.flag_gpu, self.ddf_gpu, @@ -375,13 +405,26 @@ class FlowField: stream=stream, ) self.ddf_gpu, self.temp_gpu = self.temp_gpu, self.ddf_gpu - cuda.memcpy_dtoh_async(obs_pinned, self.obs_gpu, stream) + cuda.memcpy_dtoh_async(state["obs_steps"][start + local_step], self.obs_gpu, stream) cuda.memset_d32_async(self.obs_gpu, 0, self.obs.size, stream) - self.obs += obs_pinned stream.synchronize() - self.obs = (self.obs / num_steps).astype(self.DATA_TYPE) + state["completed_steps"] += num_steps + + def end_control_interval(self): + """Publish obs once, only at the original policy-control boundary.""" + state = self._control_interval + if state is None: + raise RuntimeError("no control interval is active") + if state["completed_steps"] != state["total_steps"]: + raise RuntimeError("cannot end a control interval before its boundary") + self.obs[:] = 0 + for step_obs in state["obs_steps"]: + self.obs += step_obs + self.obs = (self.obs / state["total_steps"]).astype(self.DATA_TYPE) cuda.memcpy_dtoh(self.error_flag, self.error_flag_gpu) self.last_error_flag = int(self.error_flag[0]) + self._control_interval = None + return self.obs def has_numeric_error(self) -> bool: return bool(self.last_error_flag != 0) diff --git a/src/CCD_analysis/DUAL_CLOCK_SAMPLING.md b/src/CCD_analysis/DUAL_CLOCK_SAMPLING.md new file mode 100644 index 0000000..d94d87b --- /dev/null +++ b/src/CCD_analysis/DUAL_CLOCK_SAMPLING.md @@ -0,0 +1,173 @@ +# Legacy 双时钟流场采样 + +## 目的 + +Legacy DRL 模型只应在固定控制周期边界接收 observation 并预测下一动作。例如： + +- Illusion：每 600 lattice steps 决策一次； +- Kármán：每 800 lattice steps 决策一次。 + +流场分析通常需要另一套保存时钟，例如每 257 或 317 lattice steps 保存一次 DDF。这两个时钟不必整除，也不应相互改变。 + +错误做法是在每个流场保存点调用普通 `ff.run(segment_steps, action)`。Legacy `run()` 每次都会： + +1. 从 `ff.action` 重新初始化 action smoothing； +2. 清空 observation 累计； +3. 将该段 observation 除以段长度并发布。 + +这样会把流场保存点错误地变成控制边界，改变闭环控制。 + +## 正确语义 + +一个 DRL 控制周期只执行一次以下生命周期： + +```python +ff.begin_control_interval(sample_interval, action_target) +ff.run_control_segment(first_length) +ff.get_ddf() # 只读流场 +ff.run_control_segment(second_length) +ff.get_ddf() +# ...累计恰好 sample_interval steps... +obs = ff.end_control_interval() +``` + +其中： + +- `begin_control_interval()` 只在真实控制边界调用一次； +- `run_control_segment()` 可以按任意正整数长度调用多次； +- segment 之间延续同一个 interval-local action smoothing； +- segment 之间延续同一个 raw observation 累计； +- `get_ddf()` 只将当前 GPU DDF 复制到 host，不更新 observation、FIFO、reward、action 或 policy； +- 只有累计推进完整控制周期后才能调用 `end_control_interval()`； +- `end_control_interval()` 发布完整控制周期的平均 observation； +- FIFO、observation normalization 和下一次 `model.predict()` 仍然只在控制边界执行一次。 + +普通 `ff.run(total_steps, action_target)` 保持兼容，它等价于 begin、一次完整 segment、end。 + +## 推荐复用接口 + +通用实现位于： + +- `LegacyCelerisLab/driver.py`：control-interval session； +- `src/CCD_analysis/utils/dual_clock.py`：绝对 field-step 调度和 `ux/uy` 保存； +- `src/CCD_analysis/scripts/collect_illusion.py`、`collect_karman.py`：实际接入示例。 + +OID 若使用同一个 Legacy `FlowField`，应直接复用 `CCD_analysis.utils.dual_clock.DualClockCollector`，不要复制一套控制循环。 + +### 固定流场间隔 + +```python +from CCD_analysis.utils.dual_clock import ( + DualClockCollector, + field_steps_from_interval, +) + +control_interval = 600 +control_count = 200 +field_interval = 257 + +collector = DualClockCollector( + ff, + control_interval=control_interval, + control_count=control_count, + u0=u0, + field_steps=field_steps_from_interval( + control_interval * control_count, + field_interval, + ), +) + +obs = initial_obs +for control_index in range(control_count): + action, _ = model.predict(obs, deterministic=True) + action_target = decode_action(action) + + # 此调用内部可能在一个控制周期内读取多次或零次 DDF， + # 但始终只结束一个真实控制周期。 + raw_obs = collector.run_interval(control_index, action_target) + + fifo.append(select_observation(raw_obs)) + obs = normalize(select_observation(raw_obs)) + +collector.save("fields.npz") +``` + +### 任意、不均匀保存位置 + +`field_steps` 使用从 rollout 开始计数的绝对 lattice step，并且必须严格递增： + +```python +collector = DualClockCollector( + ff, + control_interval=800, + control_count=200, + u0=u0, + field_steps=[137, 800, 913, 1721, 2400], +) +``` + +位于控制边界的 field step 只保存一次，同时仍只产生一次完整控制 observation。 + +## 输出合同 + +`collector.save("fields.npz")` 保存： + +- `ux`, `uy`：读取时刻的无量纲速度场； +- `lattice_steps`：每个场对应的绝对 lattice step； +- `control_indices`：所属的零基控制周期； +- `control_offsets`：在该控制周期内的位置，范围为 `1..control_interval`； +- `control_interval`：模型固定控制周期。 + +后续 CCD/OID 对齐必须使用 `lattice_steps`，不能通过数组下标猜测时间。 + +## 命令示例 + +所有 Legacy GPU 采集使用 `pycuda_3_10` 和 GPU 2： + +```bash +PYTHONPATH=src:. conda run -n pycuda_3_10 python \ + src/CCD_analysis/scripts/collect_illusion.py \ + --scene illusion_1.0L --device 2 --steps 200 \ + --field-interval 257 --output-dir /path/to/run/illusion_1.0L +``` + +```bash +PYTHONPATH=src:. conda run -n pycuda_3_10 python \ + src/CCD_analysis/scripts/collect_karman.py \ + --scene karman_re100 --device 2 --steps 200 \ + --field-interval 317 --output-dir /path/to/run/karman_re100 +``` + +不提供 `--field-interval` 时使用原来的单时钟 `ff.run()` 路径，不保存 `fields.npz`。 + +## 实际闭环验证 + +2026-07-22 在 `pycuda_3_10`、物理 GPU 2 上进行了 200 个控制周期的 A/B：baseline 不读取中间 DDF，实验组使用与控制周期不整除的流场间隔，因此读取位置持续在控制周期内漂移。 + +### Illusion 1.0L + +- 控制周期：600；流场间隔：257； +- baseline 末 100 周期 DTW similarity：0.973265； +- DDF 读取组：0.976044； +- 差值：+0.002779； +- 平均 reward 差值：-0.001633。 + +### Kármán Re=100 + +- 控制周期：800；流场间隔：317； +- baseline DTW similarity：0.954562； +- DDF 读取组：0.953442； +- 差值：-0.001119（约 -0.12%）； +- 两组末 100 周期传感器序列经相位补偿后的相互 DTW similarity：0.965680。 + +最终受控流场保持相同类型的尾流结构，没有出现控制失效。因此对同一条连续在线闭环轨迹，可以认为控制期间独立读取 DDF 不会造成有意义的控制效果下降。 + +## 边界和注意事项 + +- 该方法解决的是同一条连续闭环轨迹中的独立流场采样。 +- 它不保证不同 checkpoint restore 分支 bitwise 相同；Legacy CUDA kernel 本身存在数值非确定性。 +- 不要把两个独立 restore 分支当成逐点 bitwise 配对数据。 +- 不要在 active control interval 中调用普通 `ff.run()`。 +- 不要提前调用 `end_control_interval()`。 +- 不要在同一 control interval 内切换 action target。 +- DDF 读取后必须继续同一个 session，直至原固定控制边界。 diff --git a/src/CCD_analysis/PIPELINE.md b/src/CCD_analysis/PIPELINE.md index a5a59a8..faeccc9 100644 --- a/src/CCD_analysis/PIPELINE.md +++ b/src/CCD_analysis/PIPELINE.md @@ -28,6 +28,10 @@ scripts/collect_*.py → detect_period.py → compute_correction_ (CPU) ``` +## Dual-clock field sampling + +Legacy DRL 的控制时钟与 DDF 保存时钟必须独立：模型仍只在固定控制周期边界读取完整 observation 和预测动作，DDF 可以在周期内任意 lattice step 读取。通用 API、OID 复用方式、输出时间轴合同和 GPU A/B 结果见 [`DUAL_CLOCK_SAMPLING.md`](DUAL_CLOCK_SAMPLING.md)。 + ## Key Conventions ### Geometry (unified 2026-06-25) diff --git a/src/CCD_analysis/README.md b/src/CCD_analysis/README.md index ac424cc..f434163 100644 --- a/src/CCD_analysis/README.md +++ b/src/CCD_analysis/README.md @@ -6,6 +6,8 @@ Analyzes DRL-controlled fluidic pinball using **correction-field decomposition** **→ [`PIPELINE.md`](PIPELINE.md)** — pipeline overview, results index, conventions, new training integration guide. +**→ [`DUAL_CLOCK_SAMPLING.md`](DUAL_CLOCK_SAMPLING.md)** — 在固定 DRL 控制周期内按独立时钟读取/保存 DDF；CCD 与 OID 共用。 + ## Quick Commands ```bash diff --git a/src/CCD_analysis/scripts/collect_illusion.py b/src/CCD_analysis/scripts/collect_illusion.py index 38990bf..c213cfd 100644 --- a/src/CCD_analysis/scripts/collect_illusion.py +++ b/src/CCD_analysis/scripts/collect_illusion.py @@ -17,6 +17,7 @@ import os import sys import time from collections import deque +from pathlib import Path import numpy as np @@ -36,6 +37,7 @@ from CCD_analysis.utils.cfd_interface import ( calc_lag, calc_dtw_sim, ) from CCD_analysis.utils.resampling import analyze_harmonics, gen_target_states_at +from CCD_analysis.utils.dual_clock import DualClockCollector, field_steps_from_interval DATA_TYPE = np.float32 L0 = 20.0 @@ -44,9 +46,17 @@ FIFO_LEN = 150 CONV_LEN = 36 -def run_single(scene_name: str, device_id: int, n_steps: int) -> dict: +def run_single( + scene_name: str, + device_id: int, + n_steps: int, + *, + field_interval: int | None = None, + output_dir: str | None = None, +) -> dict: cfg = get_scene(scene_name) - out_dir = data_dir_for_scene(scene_name) + out_dir = output_dir or data_dir_for_scene(scene_name) + Path(out_dir).mkdir(parents=True, exist_ok=True) u0 = cfg["u0"] si = cfg["sample_interval"] ac_scale = cfg["action_scale"] @@ -152,6 +162,15 @@ def run_single(scene_name: str, device_id: int, n_steps: int) -> dict: obs = np.zeros(s_dim, dtype=np.float32) sens_c, forc_c, act_c, rew_c, sim_c = [], [], [], [], [] + collector = None + if field_interval is not None: + collector = DualClockCollector( + ff, + control_interval=si, + control_count=n_steps, + u0=u0, + field_steps=field_steps_from_interval(si * n_steps, field_interval), + ) for step in range(n_steps): action, _ = model.predict(obs, deterministic=True) @@ -163,7 +182,10 @@ def run_single(scene_name: str, device_id: int, n_steps: int) -> dict: temp_a[3:6] = omega ff.context.push() - ff.run(si, temp_a) + if collector is None: + ff.run(si, temp_a) + else: + collector.run_interval(step, temp_a) ff.context.pop() obs_slice = ff.obs.copy()[0:12] @@ -219,6 +241,8 @@ def run_single(scene_name: str, device_id: int, n_steps: int) -> dict: np.savez(os.path.join(out_dir, "controlled.npz"), sensors=sens_arr, forces=forc_arr, actions=act_arr, rewards=np.array(rew_c, dtype=np.float32)) + if collector is not None: + collector.save(Path(out_dir) / "fields.npz") save_vorticity_png(os.path.join(out_dir, "vorticity_controlled.png"), vorticity_from_ddf(ff, u0=u0), @@ -229,7 +253,13 @@ def run_single(scene_name: str, device_id: int, n_steps: int) -> dict: avg_sim = float(np.mean(sim_c[-tail:])) if sim_c else 0.0 print(f" reward={avg_reward:.4f} similarity={avg_sim:.4f}") - result = {"scene": scene_name, "similarity": avg_sim, "avg_reward": avg_reward} + result = { + "scene": scene_name, + "similarity": avg_sim, + "avg_reward": avg_reward, + "control_interval": si, + "field_interval": field_interval, + } with open(os.path.join(out_dir, "result.json"), "w") as f: json.dump(result, f, indent=2) @@ -245,6 +275,10 @@ def main(): help="Diameter shortcut (0.75, 1.0, 1.5)") ap.add_argument("--device", type=int, default=2) ap.add_argument("--steps", type=int, default=200) + ap.add_argument("--field-interval", type=int, default=None, + help="Save fields every N absolute lattice steps") + ap.add_argument("--output-dir", type=str, default=None, + help="Optional isolated output directory") args = ap.parse_args() if args.diameter is not None: @@ -257,7 +291,10 @@ def main(): return 1 t0 = time.time() - r = run_single(scene_name, args.device, args.steps) + r = run_single( + scene_name, args.device, args.steps, + field_interval=args.field_interval, output_dir=args.output_dir, + ) print(f"Done in {time.time()-t0:.1f}s: sim={r['similarity']:.4f}") diff --git a/src/CCD_analysis/scripts/collect_karman.py b/src/CCD_analysis/scripts/collect_karman.py index 963aaa2..9b42601 100644 --- a/src/CCD_analysis/scripts/collect_karman.py +++ b/src/CCD_analysis/scripts/collect_karman.py @@ -15,6 +15,7 @@ import os import sys import time from collections import deque +from pathlib import Path import numpy as np @@ -28,6 +29,7 @@ if _SRC not in sys.path: from LegacyCelerisLab import FlowField from CCD_analysis.configs import get_scene, get_scene_list, data_dir_for_scene, model_path_for_scene, LEGACY_CFG_DIR +from CCD_analysis.utils.dual_clock import DualClockCollector, field_steps_from_interval from CCD_analysis.utils.cfd_interface import ( load_legacy_configs, build_karman_cloak_env, add_pinball, build_observation, @@ -39,9 +41,17 @@ DATA_TYPE = np.float32 L0 = 20.0 -def run_single(scene_name: str, device_id: int, n_steps: int) -> dict: +def run_single( + scene_name: str, + device_id: int, + n_steps: int, + *, + field_interval: int | None = None, + output_dir: str | None = None, +) -> dict: cfg = get_scene(scene_name) - out_dir = data_dir_for_scene(scene_name) + out_dir = output_dir or data_dir_for_scene(scene_name) + Path(out_dir).mkdir(parents=True, exist_ok=True) u0 = cfg["u0"] si = cfg["sample_interval"] ac_scale = cfg["action_scale"] @@ -120,6 +130,15 @@ def run_single(scene_name: str, device_id: int, n_steps: int) -> dict: sens_c, forc_c, act_c, rew_c = [], [], [], [] obs = np.zeros(s_dim, dtype=np.float32) + collector = None + if field_interval is not None: + collector = DualClockCollector( + ff, + control_interval=si, + control_count=n_steps, + u0=u0, + field_steps=field_steps_from_interval(si * n_steps, field_interval), + ) for step in range(n_steps): action, _ = model.predict(obs, deterministic=True) @@ -129,7 +148,10 @@ def run_single(scene_name: str, device_id: int, n_steps: int) -> dict: action_arr = scale_action(action, scale=ac_scale, bias=ac_bias, u0=u0, n_total_bodies=n_obj) ff.context.push() - ff.run(si, action_arr) + if collector is None: + ff.run(si, action_arr) + else: + collector.run_interval(step, action_arr) ff.context.pop() obs_slice = ff.obs.copy()[2:14] @@ -156,6 +178,8 @@ def run_single(scene_name: str, device_id: int, n_steps: int) -> dict: np.savez(os.path.join(out_dir, "controlled.npz"), sensors=sens_arr, forces=forc_arr, actions=act_arr, rewards=rew_arr) + if collector is not None: + collector.save(Path(out_dir) / "fields.npz") save_vorticity_png(os.path.join(out_dir, "vorticity_controlled.png"), vorticity_from_ddf(ff, u0=u0), @@ -165,7 +189,13 @@ def run_single(scene_name: str, device_id: int, n_steps: int) -> dict: sim_score = compute_similarity(target_states, sens_arr, 30) print(f" reward={avg_reward:.4f} similarity={sim_score:.4f}") - result = {"scene": scene_name, "similarity": sim_score, "avg_reward": avg_reward} + result = { + "scene": scene_name, + "similarity": sim_score, + "avg_reward": avg_reward, + "control_interval": si, + "field_interval": field_interval, + } with open(os.path.join(out_dir, "result.json"), "w") as f: json.dump(result, f, indent=2) @@ -181,6 +211,10 @@ def main(): help="Re number shortcut (50, 100, 200, 400)") ap.add_argument("--device", type=int, default=2) ap.add_argument("--steps", type=int, default=200) + ap.add_argument("--field-interval", type=int, default=None, + help="Save fields every N absolute lattice steps") + ap.add_argument("--output-dir", type=str, default=None, + help="Optional isolated output directory") args = ap.parse_args() if args.re is not None: @@ -193,7 +227,10 @@ def main(): return 1 t0 = time.time() - r = run_single(scene_name, args.device, args.steps) + r = run_single( + scene_name, args.device, args.steps, + field_interval=args.field_interval, output_dir=args.output_dir, + ) print(f"Done in {time.time()-t0:.1f}s: sim={r['similarity']:.4f}") diff --git a/src/CCD_analysis/tests/test_dual_clock.py b/src/CCD_analysis/tests/test_dual_clock.py new file mode 100644 index 0000000..e28c70f --- /dev/null +++ b/src/CCD_analysis/tests/test_dual_clock.py @@ -0,0 +1,65 @@ +import numpy as np +import pytest + +from CCD_analysis.utils.dual_clock import DualClockCollector, field_steps_from_interval + + +class FakeFlowField: + def __init__(self): + self.absolute_step = 0 + self.ddf = np.zeros(9, dtype=np.float32) + self.begin_calls = 0 + self.end_calls = 0 + self._completed = 0 + self._total = 0 + + def begin_control_interval(self, total_steps, action_target): + self.begin_calls += 1 + self._total = total_steps + self._completed = 0 + + def run_control_segment(self, steps): + self.absolute_step += steps + self._completed += steps + + def end_control_interval(self): + assert self._completed == self._total + self.end_calls += 1 + return np.array([self.absolute_step], dtype=np.float32) + + def get_ddf(self): + self.ddf[:] = self.absolute_step + + +def test_field_steps_from_nondivisible_interval(): + assert field_steps_from_interval(16, 3) == (3, 6, 9, 12, 15) + + +def test_irregular_field_reads_do_not_add_control_boundaries(monkeypatch): + def fake_velocity(ff, u0): + ff.get_ddf() + return ff.ddf[:3].copy(), ff.ddf[3:6].copy() + + monkeypatch.setattr("CCD_analysis.utils.dual_clock.get_velocity_field", fake_velocity) + ff = FakeFlowField() + collector = DualClockCollector( + ff, + control_interval=8, + control_count=2, + u0=0.01, + field_steps=[3, 6, 9, 12, 15, 16], + ) + collector.run_interval(0, np.zeros(1, dtype=np.float32)) + collector.run_interval(1, np.zeros(1, dtype=np.float32)) + assert collector.saved_steps == [3, 6, 9, 12, 15, 16] + assert ff.begin_calls == 2 + assert ff.end_calls == 2 + assert ff.absolute_step == 16 + + +def test_collector_rejects_invalid_field_timeline(): + with pytest.raises(ValueError): + DualClockCollector( + FakeFlowField(), control_interval=8, control_count=2, + u0=0.01, field_steps=[3, 3], + ) diff --git a/src/CCD_analysis/utils/dual_clock.py b/src/CCD_analysis/utils/dual_clock.py new file mode 100644 index 0000000..74d35c4 --- /dev/null +++ b/src/CCD_analysis/utils/dual_clock.py @@ -0,0 +1,101 @@ +"""Minimal dual-clock field sampling for legacy CelerisLab rollouts.""" +from __future__ import annotations + +from collections.abc import Iterable +from pathlib import Path + +import numpy as np + +from CCD_analysis.utils.cfd_interface import get_velocity_field + + +def field_steps_from_interval(total_steps: int, field_interval: int) -> tuple[int, ...]: + """Return absolute field-save steps within a finite rollout.""" + if type(total_steps) is not int or total_steps < 1: + raise ValueError("total_steps must be a positive integer") + if type(field_interval) is not int or field_interval < 1: + raise ValueError("field_interval must be a positive integer") + return tuple(range(field_interval, total_steps + 1, field_interval)) + + +class DualClockCollector: + """Keep policy decisions fixed while reading fields at independent steps.""" + + def __init__( + self, + flow_field, + *, + control_interval: int, + control_count: int, + u0: float, + field_steps: Iterable[int], + ): + if type(control_interval) is not int or control_interval < 1: + raise ValueError("control_interval must be a positive integer") + if type(control_count) is not int or control_count < 1: + raise ValueError("control_count must be a positive integer") + self.flow_field = flow_field + self.control_interval = control_interval + self.control_count = control_count + self.u0 = u0 + self.field_steps = tuple(field_steps) + horizon = control_interval * control_count + if any(type(step) is not int for step in self.field_steps): + raise ValueError("field steps must be integers") + if tuple(sorted(set(self.field_steps))) != self.field_steps: + raise ValueError("field steps must be strictly increasing and unique") + if any(step < 1 or step > horizon for step in self.field_steps): + raise ValueError("field steps must lie within the rollout") + self._next_field = 0 + self._next_control = 0 + self.ux: list[np.ndarray] = [] + self.uy: list[np.ndarray] = [] + self.saved_steps: list[int] = [] + + def run_interval(self, control_index: int, action_target: np.ndarray) -> np.ndarray: + """Advance exactly one policy interval, pausing only at due field steps.""" + if control_index != self._next_control: + raise ValueError("control intervals must run sequentially") + start = control_index * self.control_interval + stop = start + self.control_interval + due = [] + while self._next_field < len(self.field_steps): + step = self.field_steps[self._next_field] + if step > stop: + break + if step > start: + due.append(step) + self._next_field += 1 + + ff = self.flow_field + ff.begin_control_interval(self.control_interval, action_target) + current = start + boundaries = due if due and due[-1] == stop else [*due, stop] + for step in boundaries: + if step > current: + ff.run_control_segment(step - current) + current = step + if step in due: + ux, uy = get_velocity_field(ff, u0=self.u0) + self.ux.append(ux) + self.uy.append(uy) + self.saved_steps.append(step) + obs = ff.end_control_interval().copy() + self._next_control += 1 + return obs + + def save(self, path: str | Path) -> None: + """Save fields and their absolute lattice-step timeline.""" + if self._next_control != self.control_count: + raise RuntimeError("cannot save an incomplete rollout") + destination = Path(path) + destination.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed( + destination, + ux=np.asarray(self.ux, dtype=np.float32), + uy=np.asarray(self.uy, dtype=np.float32), + lattice_steps=np.asarray(self.saved_steps, dtype=np.int64), + control_indices=(np.asarray(self.saved_steps, dtype=np.int64) - 1) // self.control_interval, + control_offsets=(np.asarray(self.saved_steps, dtype=np.int64) - 1) % self.control_interval + 1, + control_interval=np.asarray(self.control_interval, dtype=np.int64), + )