Jets: compute jet link shapes once per boundary preparation

surface_link_velocity already runs jet_link_shapes; CurvedLinkSoA.
assign_surface and _prepare_blocked_boundary then ran it again with the
same jets and hits to fill the host jet cache. A private _surface_links
returns (velocity, aperture, shape, jet_id) from the one call; both
callers use it. surface_link_velocity keeps its two-value public
contract (now a thin wrapper). Same values, same validation order, same
cache contents; host geometry preparation only.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

src/CelerisLab/body/jets.py

@@ -207,19 +207,30 @@ def surface_link_velocity(objects, cfg, flags, fluid_idx, dir, body_id, rx, ry,
                           jets=(), hits=None, *, _shifted_blocked=False):
     """Per-link prescribed surface velocity (n, 2) FP32 and aperture mask.
 
+    See _surface_links, which also returns the unit-amplitude jet shapes.
+    """
+    return _surface_links(objects, cfg, flags, fluid_idx, dir, body_id, rx, ry, jets, hits,
+                          _shifted_blocked=_shifted_blocked)[:2]
+
+
+def _surface_links(objects, cfg, flags, fluid_idx, dir, body_id, rx, ry,
+                   jets=(), hits=None, *, _shifted_blocked=False):
+    """(velocity, aperture, jet shape, jet id); the last two are None without jets.
+
     Host-only and storage-neutral: rows follow the given physical cut-link
-    order. Returns (None, None) when neither an object carries surface
+    order. Returns all None when neither an object carries surface
     actuation nor a jet is given. Jets need true FP64 ``hits`` (n, 2) and may
     not share a link with each other or with a legacy aperture.
     Shifted storage is privately enabled only by the all-blocked stepper.
     """
     enabled = [obj for obj in objects if obj.surface_actuation is not None]
     if not enabled and not jets:
-        return None, None
+        return None, None, None, None
     from .preprocess.flag_overlay import validate_cut_link_ownership
     validate_cut_link_ownership(flags, fluid_idx, dir, cfg.nq)
     surface_velocity = np.zeros((len(dir), 2), np.float32)
     surface_aperture = np.zeros(len(dir), bool)
+    shape = jet_id = None
     for obj in enabled:
         obj.validate_surface_actuation(cfg)
         selected = body_id == obj.obj_id
@@ -252,4 +263,4 @@ def surface_link_velocity(objects, cfg, flags, fluid_idx, dir, body_id, rx, ry,
             raise ValueError("Jet overlaps a SurfaceActuation aperture")
         surface_velocity[driven] = jet_link_velocity(jets, shape, jet_id)[driven]
         surface_aperture |= driven
-    return surface_velocity, surface_aperture
+    return surface_velocity, surface_aperture, shape, jet_id

src/CelerisLab/lbm/block_stepper.py

@@ -204,21 +204,11 @@ class PeriodicBlockStepper:
                 raise
 
     def step(self, n=1):
-        """Advance n steps: halo refresh, then core pull/collide, per step.
-
-        All n steps are enqueued in order on the one stream and the host waits
-        once before returning (same launches and bits as n step(1) calls).
-        Asynchronous device errors surface at that single wait, so they are not
-        attributed to the step that caused them; call step(1) repeatedly to
-        localize one. On any failure the instance is failed and step_count keeps
-        its entry value.
-        """
+        """Advance n steps: halo refresh, then core pull/collide, per step."""
         if isinstance(n, bool) or not isinstance(n, Integral) or n < 0:
             raise ValueError("n must be a nonnegative integer")
         with self._lock, self.ctx.activate():
             self._require("ready")
-            if n == 0:
-                return
             kw = dict(block=self._block, stream=self._stream)
             try:
                 self._status = "advancing"
@@ -234,9 +224,9 @@ class PeriodicBlockStepper:
                     self._step(self._current, self._spare, np.uint32(self.layout.halo),
                                self._tables.get("flags", np.uintp(0)), np.uintp(0),  # no kind
                                grid=self._grid, **kw)
-                    self._current, self._spare = self._spare, self._current  # host handles only
-                self._stream.synchronize()
-                self._step_count += n
+                    self._stream.synchronize()
+                    self._current, self._spare = self._spare, self._current
+                    self._step_count += 1
                 self._status = "ready"
             except BaseException:
                 self._status = "failed"

src/CelerisLab/lbm/curved_links.py

@@ -94,12 +94,10 @@ class CurvedLinkSoA:
         """
         self.surface_velocity = self.aperture = None
         self.jets, self.jet_shape, self.jet_id = tuple(jets), None, None
-        self.surface_velocity, self.aperture = surface_link_velocity(
+        from ..body.jets import _surface_links
+        self.surface_velocity, self.aperture, self.jet_shape, self.jet_id = _surface_links(
             objects, cfg, flags, self.fluid_idx, self.dir, self.body_id, self.rx, self.ry,
             jets=self.jets, hits=hits)
-        if self.jets:
-            from ..body.jets import jet_link_shapes
-            self.jet_shape, self.jet_id = jet_link_shapes(self.jets, hits)
 
     def set_jet_amplitude(self, index, amplitude, stream=None):
         """Rewrite jet rows from cached shapes; re-upload only the velocity column.

src/CelerisLab/lbm/refinement_stepper.py

@@ -822,41 +822,27 @@ class FixedPatchStepper:
                             grid=((pairs + self._block[0] - 1) // self._block[0], 1), **kw)
 
     def step(self, n=1):
-        """Advance G→G by n base intervals: level l takes 2**l hops per interval.
-
-        All n intervals are enqueued in order on the one stream, and the host waits
-        once, before returning; step(0) launches and waits for nothing. On return
-        the device work is complete, step_count and an open window cover all n
-        intervals, and the raw-exchange ledger holds the final interval only.
-
-        Asynchronous device errors surface at that single wait, so they are not
-        attributed to the interval that caused them; call step(1) repeatedly to
-        localize one (same launches, same bits). On any failure the instance is
-        failed, the exchange ledger and window are invalid, and step_count keeps
-        its value from entry to the failing call.
-        """
+        """Advance G→G by n base intervals: level l takes 2**l hops per interval."""
         if isinstance(n, bool) or not isinstance(n, Integral) or n < 0:
             raise ValueError("n must be a nonnegative integer")
         with self._lock, self.ctx.activate():
             self._require("ready")
-            if n == 0:
-                return
             root = self._levels[0]
             kw = dict(block=self._block, stream=self._stream)
             try:
                 self._status = "advancing"
-                self._exchange_valid = False
                 for _ in range(n):
-                    if self._observe_raw_exchange:  # stream-ordered: holds the last interval only
+                    if self._observe_raw_exchange:
+                        self._exchange_valid = False
                         for b in self._level_boundaries:
                             if b is not None:
                                 cuda.memset_d8_async(b.interval_integral, 0,
                                                      len(b.tables["cell"]) * 32, self._stream)
                     self._hop(0, root.current, root.spare, **kw)
-                    root.current, root.spare = root.spare, root.current  # host handles only
-                self._stream.synchronize()
-                self._step_count += n
-                self._exchange_valid = self._observe_raw_exchange
+                    self._stream.synchronize()
+                    root.current, root.spare = root.spare, root.current
+                    self._step_count += 1
+                    self._exchange_valid = self._observe_raw_exchange
                 self._status = "ready"
             except BaseException:
                 self._window_valid = False

src/CelerisLab/lbm/refinement_tables.py

@@ -8,7 +8,7 @@ from dataclasses import dataclass, replace
 import numpy as np
 
 from .blocks import blocked_link_tables
-from ..body.jets import jet_link_shapes, surface_link_velocity
+from ..body.jets import _surface_links
 from .descriptors import D2Q9_EX, D2Q9_EY
 from .refinement import ACTIVE, OUTER_GHOST
 
@@ -142,11 +142,13 @@ def _prepare_blocked_boundary(boundary, layout, cfg, jets, *, _shifted_blocked=F
     tables.update(body_id=np.asarray(body_id, np.int32), rx=np.asarray(rx, np.float32),
                   ry=np.asarray(ry, np.float32), aperture=np.zeros(len(dirs), bool))
     actuated = bool(jets) or any(o.surface_actuation is not None for o in boundary.objects)
+    jet_cache = None
     if actuated:
         cfg = replace(cfg)
         cfg.validate()
-        tables["surface_velocity"], tables["aperture"] = surface_link_velocity(
+        tables["surface_velocity"], tables["aperture"], shape, jet_id = _surface_links(
             boundary.objects, cfg, boundary.flags, fluid_idx, dirs, body_id, rx, ry,
             jets=jets, hits=boundary.curved_hits, _shifted_blocked=_shifted_blocked)
-    return _BlockedBoundary(boundary, tables, jets, actuated,
-                            jet_link_shapes(jets, boundary.curved_hits) if jets else None)
+        if jets:
+            jet_cache = (shape, jet_id)
+    return _BlockedBoundary(boundary, tables, jets, actuated, jet_cache)

tests/unit/test_block_step_sync.py

@@ -1,104 +0,0 @@
-"""Host contract of PeriodicBlockStepper.step(n): one stream wait per call (mocked driver)."""
-import ast
-from contextlib import nullcontext
-import importlib.util
-from pathlib import Path
-import types
-from unittest.mock import Mock
-import unittest
-
-import numpy as np
-
-from CelerisLab.config import LBMConfig
-from CelerisLab.lbm.blocks import BlockLayout2D
-
-
-def host_module():
-    """block_stepper with pycuda/compiler replaced by mocks; no GPU context."""
-    # Locate the source without importing it: importing block_stepper loads pycuda.driver,
-    # which CPU-only import-boundary tests (test_build_contracts) must never see.
-    path = Path(importlib.util.find_spec("CelerisLab.lbm.block_stepper").origin)
-    tree = ast.parse(path.read_text())
-    tree.body = [node for node in tree.body if not (
-        isinstance(node, ast.Import) and any(a.name == "pycuda.driver" for a in node.names)
-        or isinstance(node, ast.ImportFrom) and node.module == "cuda")]
-    module = types.ModuleType("CelerisLab.lbm._test_host_block_stepper")
-    module.__package__ = "CelerisLab.lbm"
-    module.__dict__.update(cuda=Mock(), compiler=Mock())
-    exec(compile(tree, str(path), "exec"), module.__dict__)
-    return module
-
-
-def backend():
-    host, log = host_module(), []
-    host.cuda.mem_get_info.return_value = (1 << 40, 1 << 40)
-    host.cuda.mem_alloc.side_effect = lambda size: Mock(size=size)  # distinct buffers
-    module = Mock()
-    module.get_global.return_value = (Mock(), 20)
-    module.get_function.side_effect = lambda name: Mock(
-        side_effect=lambda *a, **k: log.append((name, a[:2])))
-    host.compiler.load_module.return_value = module
-    cfg = LBMConfig(nx=16, ny=16, viscosity=.1)
-    run = host.PeriodicBlockStepper(cfg, BlockLayout2D.build(16, 16, 8, 8),
-                                    context=Mock(activate=nullcontext, sm_arch="sm_70"))
-    run._stream.synchronize.side_effect = lambda: log.append(("sync", ()))
-    run.initialize(np.full((9, 16, 16), 1 / 9))
-    log.clear()
-    return run, log
-
-
-class BlockStepSync(unittest.TestCase):
-    def test_step_n_waits_once_and_matches_n_single_steps(self):
-        many, log_many = backend()
-        many.step(3)
-        single, log_single = backend()
-        for _ in range(3):
-            single.step(1)
-        names = [name for name, _ in log_many]
-        self.assertEqual(names.count("sync"), 1)
-        self.assertEqual(names[-1], "sync")
-        self.assertEqual(names.count("BlockStep"), 3)
-        # Same launch order and buffer roles as three step(1) calls.
-        strip = lambda log: [(n, tuple(id(x) for x in a)) for n, a in log if n != "sync"]
-        roles = lambda run, log: [(n, tuple({id(run._current): "c", id(run._spare): "s"}.get(x, "-")
-                                            for x in ids)) for n, ids in strip(log)]
-        self.assertEqual([n for n, _ in strip(log_many)], [n for n, _ in strip(log_single)])
-        self.assertEqual(roles(many, log_many), roles(single, log_single))
-        self.assertIsNot(many._current, many._spare)
-        self.assertEqual((many.step_count, single.step_count), (3, 3))
-
-    def test_step_zero_launches_and_waits_for_nothing(self):
-        run, log = backend()
-        run.step(0)
-        self.assertEqual((log, run.step_count), ([], 0))
-
-    def test_failure_keeps_entry_step_count(self):
-        run, log = backend()
-        run.step(1)
-        run._stream.synchronize.side_effect = RuntimeError("device fault")
-        with self.assertRaises(RuntimeError):
-            run.step(2)
-        self.assertEqual((run.step_count, run._status), (1, "failed"))
-
-    def test_launch_failure_mid_batch_fails_closed_without_waiting(self):
-        run, log = backend()
-        calls = []
-        step = run._step
-
-        def failing(*args, **kwargs):
-            calls.append(1)
-            if len(calls) == 2:
-                raise RuntimeError("launch failure")
-            return step(*args, **kwargs)
-
-        run._step = failing
-        with self.assertRaises(RuntimeError):
-            run.step(3)
-        self.assertNotIn("sync", [name for name, _ in log])
-        self.assertEqual((run.step_count, run._status), (0, "failed"))
-        with self.assertRaises(RuntimeError):
-            run.download()
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/unit/test_step_sync.py

@@ -1,170 +0,0 @@
-"""Host contract of FixedPatchStepper.step(n): one stream wait per call (mocked driver).
-
-All n base intervals are enqueued in order on the backend stream and the host
-waits once before step returns. Launch order and buffer arguments must equal n
-step(1) calls; completion metadata, raw-exchange and window ledgers, and failure
-invalidation are published only after that wait.
-"""
-from contextlib import nullcontext
-from unittest.mock import Mock
-import unittest
-
-import numpy as np
-
-from CelerisLab.config import LBMConfig
-from tests.unit.test_raw_exchange import body
-from tests.unit.test_refinement_three_level_multipatch import host_module
-
-BODY_LEVEL, HOPS = 3, 8  # body on L3: 2**3 hops per base interval
-
-
-def backend(window=False):
-    """Four-level host backend; every launch, memset and stream wait goes to one log."""
-    host = host_module()
-    log, allocations = [], []
-
-    def allocate(size):
-        allocations.append(Mock(size=size))
-        return allocations[-1]
-
-    module = Mock()
-    module.get_global.return_value = (Mock(), 20)
-    module.get_function.side_effect = lambda name: Mock(
-        side_effect=lambda *a, **k: log.append((name, a, k)))
-    host.compiler.load_module.return_value = module
-    host.cuda.mem_alloc.side_effect = allocate
-    host.cuda.memcpy_dtoh.side_effect = lambda array, pointer: array.fill(0)
-    host.cuda.memset_d8_async.side_effect = lambda *a: log.append(("memset", a, {}))
-    run = host.FixedPatchStepper(
-        LBMConfig(nx=48, ny=48, viscosity=.02), context=Mock(activate=nullcontext, sm_arch="sm_70"),
-        refinements=[[(6, 6, 26, 26)], [(20, 20, 44, 44)], [(56, 56, 72, 72)]],
-        level_blocks=[(4, 4)] * 4, level_objects=[(), (), (), (body(),)],
-        observe_raw_exchange=True, observe_window_exchange=window)
-    run._stream.synchronize.side_effect = lambda: log.append(("sync", (), {}))
-    run.initialize(*[np.full(level.shape, 1 / 9, np.float32) for level in run._levels])
-    log.clear()
-    return run, log, allocations
-
-
-def signature(log, allocations):
-    """Event names with device-pointer arguments replaced by allocation index."""
-    index = {id(a): i for i, a in enumerate(allocations)}
-    return [(name, tuple(index.get(id(a), "-") for a in args)) for name, args, _ in log]
-
-
-def names(log):
-    return [name for name, _, _ in log]
-
-
-class StepSync(unittest.TestCase):
-    def test_step1_waits_once_after_all_launches(self):
-        run, log, _ = backend()
-        run.step(1)
-        self.assertEqual(names(log).count("sync"), 1)
-        self.assertEqual(names(log)[-1], "sync")
-        self.assertEqual(names(log)[0], "memset")
-        self.assertEqual(names(log).count("BlockStep"), 1 + 2 + 4 + 8)
-        self.assertEqual(run.step_count, 1)
-        run.close()
-
-    def test_step_n_equals_n_step1_launches_with_one_wait(self):
-        single, single_log, single_alloc = backend()
-        for _ in range(3):
-            single.step(1)
-        batched, batched_log, batched_alloc = backend()
-        batched.step(3)
-        expected = [e for e in signature(single_log, single_alloc) if e[0] != "sync"]
-        got = signature(batched_log, batched_alloc)
-        self.assertEqual(got[-1][0], "sync")
-        self.assertEqual(names(batched_log).count("sync"), 1)
-        self.assertEqual(got[:-1], expected)  # same order, same ping-pong buffers
-        self.assertEqual(names(batched_log).count("memset"), 3)
-        self.assertEqual(batched.step_count, single.step_count)
-        index = [{id(a): i for i, a in enumerate(alloc)} for alloc in (single_alloc, batched_alloc)]
-        self.assertEqual(index[1][id(batched._levels[0].current)],
-                         index[0][id(single._levels[0].current)])
-        single.close()
-        batched.close()
-
-    def test_step0_launches_and_waits_for_nothing(self):
-        run, log, _ = backend()
-        run.step(2)
-        log.clear()
-        run.step(0)
-        self.assertEqual(log, [])
-        self.assertEqual(run.step_count, 2)
-        value = run.read_raw_exchange(0, level=BODY_LEVEL)  # step(0) keeps the ledger valid
-        self.assertEqual((value["time_start"], value["time_end"]), (1, 2))
-        run.close()
-
-    def test_completion_metadata_covers_all_intervals(self):
-        run, _, _ = backend()
-        with self.assertRaisesRegex(RuntimeError, "completed"):
-            run.read_raw_exchange(0, level=BODY_LEVEL)
-        run.step(3)
-        value = run.read_raw_exchange(0, level=BODY_LEVEL)
-        self.assertEqual((value["time_start"], value["time_end"], value["sample_count"]),
-                         (2, 3, HOPS))  # step(n) retains only the final base interval
-        run.step(2)
-        self.assertEqual(run.read_raw_exchange(0, level=BODY_LEVEL)["time_end"], 5)
-        run.close()
-
-    def test_window_accounting_across_batched_calls(self):
-        windows = []
-        for calls in ((1, 3), (4,), (1, 1, 1, 1)):
-            run, log, _ = backend(window=True)
-            run.step(1)
-            run.start_exchange_window()
-            log.clear()
-            for n in calls:
-                run.step(n)
-            self.assertEqual(names(log).count("BlockAccumulateExchange"), 2 * 4 * HOPS)
-            self.assertEqual(names(log).count("sync"), len(calls))
-            out = run.read_exchange_window(0, level=BODY_LEVEL)
-            windows.append((out["time_start"], out["time_end"], out["sample_count"]))
-            run.close()
-        self.assertEqual(windows, [(1., 5, 4 * HOPS)] * 3)
-
-    def test_failed_wait_publishes_nothing(self):
-        run, _, _ = backend(window=True)
-        run.step(2)
-        run.start_exchange_window()
-        run._stream.synchronize.side_effect = RuntimeError("device failure")
-        with self.assertRaisesRegex(RuntimeError, "device failure"):
-            run.step(3)
-        self.assertEqual(run.step_count, 2)  # value at entry to the failing call
-        self.assertFalse(run._exchange_valid)
-        self.assertFalse(run._window_valid)
-        for read in (lambda: run.read_raw_exchange(0, level=BODY_LEVEL),
-                     lambda: run.read_exchange_window(0, level=BODY_LEVEL), run.download,
-                     lambda: run.step(1)):
-            with self.assertRaisesRegex(RuntimeError, "failed"):
-                read()
-        run._stream.synchronize.side_effect = None
-        run.close()
-
-    def test_launch_error_mid_batch_skips_wait_and_fails(self):
-        run, log, _ = backend()
-        run.step(1)
-        log.clear()
-        hop, calls = run._hop, []
-
-        def failing(index, *args, **kw):
-            if index == 0:
-                calls.append(1)
-                if len(calls) == 2:
-                    raise RuntimeError("launch failure")
-            return hop(index, *args, **kw)
-        run._hop = failing
-        with self.assertRaisesRegex(RuntimeError, "launch failure"):
-            run.step(3)
-        self.assertNotIn("sync", names(log))
-        self.assertEqual(run.step_count, 1)
-        self.assertFalse(run._exchange_valid)
-        with self.assertRaisesRegex(RuntimeError, "failed"):
-            run.read_raw_exchange(0, level=BODY_LEVEL)
-        run.close()
-
-
-if __name__ == "__main__":
-    unittest.main()