test: opt in to window-ledger DFG observations with separate receipts

tests/validation/run_refined_dfg.py

@@ -45,7 +45,7 @@ def config(h, diameter, collision="SRT"):
     return cfg, body, p
 
 
-def plan(h, collision="SRT"):
+def plan(h, collision="SRT", observation_mode="last_base_step"):
     _, _, c = config(h, 40, collision)
     _, _, f = config(h, 80, collision)
     x0, y0, x1, y1 = BOUNDS
@@ -54,6 +54,7 @@ def plan(h, collision="SRT"):
     return dict(
         status="prepared_not_run",
         collision=collision,
+        observation_mode=observation_mode,
         root_gpu_gates_required=True,
         root=c,
         uniform=f,
@@ -156,10 +157,13 @@ def host_geometry(source, h, collision="SRT"):
     return cfg, pair, body, host
 
 
-def verify_source(source, collision="SRT"):
+def verify_source(source, collision="SRT", observation_mode="last_base_step"):
     if collision not in ("SRT", "MRT"):
         raise ValueError("Matched collision must be SRT or MRT")
-    name = "source_manifest.json" if collision == "SRT" else "source_manifest_mrt.json"
+    if observation_mode not in ("last_base_step", "window"):
+        raise ValueError("Unknown refined observation mode")
+    name = ("source_manifest_window.json" if observation_mode == "window" else
+            "source_manifest.json" if collision == "SRT" else "source_manifest_mrt.json")
     expected = json.loads((HERE / "data/refined_dfg" / name).read_text())
     for name, digest in expected["sha256"].items():
         if hashlib.sha256((source / name).read_bytes()).hexdigest() != digest:
@@ -167,6 +171,46 @@ def verify_source(source, collision="SRT"):
     return expected
 
 
+def advance_refined_interval(backend, n, deadline, observation_mode):
+    """Return the same root-unit interval impulse, never an already averaged force."""
+    impulse = np.zeros(2)
+    if observation_mode == "window":
+        begin = backend.step_count
+        check_deadline(deadline)
+        backend.start_exchange_window()
+        backend.step(n)
+        exchange = backend.read_exchange_window(0, level=1)
+        if (
+            exchange["units"] != "base_lattice_per_unit_depth"
+            or exchange["window"] != "window"
+            or exchange["sample_count"] != 2 * n
+            or exchange["time_start"] != begin
+            or exchange["time_end"] != begin + n
+            or exchange["dt"] != 0.5
+        ):
+            raise RuntimeError("Unexpected window force observation contract")
+        impulse += exchange["total"]["impulse"]
+        check_deadline(deadline)
+        return impulse
+    if observation_mode != "last_base_step":
+        raise ValueError("Unknown refined observation mode")
+    for _ in range(n):
+        check_deadline(deadline)
+        backend.step(1)
+        exchange = backend.read_raw_exchange(
+            0, level=1, window="last_base_step"
+        )
+        if (
+            exchange["units"] != "base_lattice_per_unit_depth"
+            or exchange["sample_count"] != 2
+            or exchange["time_end"] - exchange["time_start"] != 1
+        ):
+            raise RuntimeError("Unexpected force observation contract")
+        impulse += exchange["total"]["impulse"]
+        check_deadline(deadline)
+    return impulse
+
+
 def run_case(a, h, periodic, case):
     from CelerisLab.simulation import Simulation
     from CelerisLab.cuda.context import CudaContext
@@ -199,6 +243,7 @@ def run_case(a, h, periodic, case):
                     fine_block=FINE_BLOCK,
                     fine_objects=(body,),
                     observe_raw_exchange=True,
+                    **({"observe_window_exchange": True} if a.observation_mode == "window" else {}),
                 )
             )
             from CelerisLab.lbm.channel import channel_masks
@@ -248,20 +293,9 @@ def run_case(a, h, periodic, case):
                 impulse = np.zeros(2)
                 tick = time.monotonic()
                 if refined:
-                    for _ in range(n):
-                        check_deadline(a.deadline)
-                        backend.step(1)
-                        exchange = backend.read_raw_exchange(
-                            0, level=1, window="last_base_step"
-                        )
-                        if (
-                            exchange["units"] != "base_lattice_per_unit_depth"
-                            or exchange["sample_count"] != 2
-                            or exchange["time_end"] - exchange["time_start"] != 1
-                        ):
-                            raise RuntimeError("Unexpected force observation contract")
-                        impulse += exchange["total"]["impulse"]
-                        check_deadline(a.deadline)
+                    impulse = advance_refined_interval(
+                        backend, n, a.deadline, a.observation_mode
+                    )
                 else:
                     done = 0
                     while done < n:
@@ -345,6 +379,7 @@ def run_case(a, h, periodic, case):
                     ready = dict(
                         case=case,
                         collision=a.collision,
+                        observation_mode=a.observation_mode,
                         **a.provenance,
                         mode=a.mode,
                         step=begin,
@@ -359,6 +394,7 @@ def run_case(a, h, periodic, case):
     result = dict(
         case=case,
         collision=a.collision,
+        observation_mode=a.observation_mode,
         **a.provenance,
         mode=a.mode,
         history=rows,
@@ -428,12 +464,14 @@ def check_deadline(deadline):
         )
 
 
-def verify_pilot(path, manifest, collision="SRT"):
+def verify_pilot(path, manifest, collision="SRT", observation_mode="last_base_step"):
     if path is None:
         raise ValueError(
             "Full comparison requires --pilot-results from a completed matched pilot"
         )
     previous = read_finite_json(path / "plan.json")
+    if previous.get("observation_mode") != observation_mode:
+        raise ValueError("Pilot observation mode differs")
     if previous.get("collision") != collision:
         raise ValueError("Pilot collision differs")
     if previous["source_manifest"] != manifest or previous["bounds_root_edge"] != list(
@@ -446,7 +484,8 @@ def verify_pilot(path, manifest, collision="SRT"):
     for case, diameter in (("uniform", 80), ("refined", 40)):
         summary = read_finite_json(path / case / "summary.json")
         if (
-            summary.get("collision") != collision
+            summary.get("observation_mode") != observation_mode
+            or summary.get("collision") != collision
             or summary.get("case") != case
             or summary.get("mode") != "pilot"
             or summary.get("status") != "pilot_complete"
@@ -498,7 +537,8 @@ def verify_pilot(path, manifest, collision="SRT"):
                 raise ValueError("Unhealthy pilot observations")
         ready = read_finite_json(path / case / "ready.json")
         if (
-            ready.get("collision") != collision
+            ready.get("observation_mode") != observation_mode
+            or ready.get("collision") != collision
             or ready.get("case") != case
             or ready.get("mode") != "pilot"
             or ready.get("step") != interval
@@ -523,6 +563,8 @@ def verify_pilot(path, manifest, collision="SRT"):
 
 def main():
     parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--exchange-window", action="store_true",
+                        help="Opt in to refined window-ledger observations on the pinned candidate")
     parser.add_argument("--collision", choices=["SRT", "MRT"], default="SRT")
     parser.add_argument("--source", type=Path, default=DEFAULT_SOURCE)
     parser.add_argument("--output", type=Path, required=True)
@@ -536,13 +578,14 @@ def main():
     )
     a = parser.parse_args()
     a.source = a.source.resolve()
-    manifest = verify_source(a.source, a.collision)
+    a.observation_mode = "window" if a.exchange_window else "last_base_step"
+    manifest = verify_source(a.source, a.collision, a.observation_mode)
     h, periodic = helpers(a.source)
     from source_provenance import checkout_provenance
 
     a.provenance = checkout_provenance(a.source)
     calibration = (
-        verify_pilot(a.pilot_results, manifest, a.collision)
+        verify_pilot(a.pilot_results, manifest, a.collision, a.observation_mode)
         if a.mode == "full" else None
     )
     if a.mode != "plan" and (
@@ -559,7 +602,7 @@ def main():
     if a.output.exists() and any(a.output.iterdir()):
         raise ValueError("Output must be empty/new")
     a.output.mkdir(parents=True, exist_ok=True)
-    specification = plan(h, a.collision)
+    specification = plan(h, a.collision, a.observation_mode)
     specification.update(
         source_manifest=manifest,
         **a.provenance,
@@ -586,6 +629,7 @@ def main():
             failure = dict(
                 status="wall_time_limit",
                 collision=a.collision,
+                observation_mode=a.observation_mode,
                 **a.provenance,
                 case=case,
                 matched_case_pass=False,
@@ -601,6 +645,7 @@ def main():
             break
     comparison = dict(
         collision=a.collision,
+        observation_mode=a.observation_mode,
         **a.provenance,
         status="pilot_only" if a.mode == "pilot" else "incomplete",
         matched_case_pass=False,