test_runtime_child_remap: inject remap body rejects through the boundary record

The removed _boundaries list made the two body-reject patches no-ops.
Inject them via _device_boundaries (level-2 record with foreign host
geometry; non-None level-1 record). set_jet_amplitude docstring names
the public jets view.

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

src/CelerisLab/lbm/blocks.py

@@ -18,7 +18,9 @@ all NQ populations from an owner core cell into another block's halo cell.
 Periodicity is explicit per axis. Across a non-periodic exterior, halo cells
 have no pair and stay invalid (NaN after packing): no wrap, no implied BC.
 
-Preparation only: no CUDA, no boundary conditions, no refinement levels.
+Host preparation plus thin launch wrappers: no CUDA import or allocation
+(launch_* only call kernels and device tables the caller owns), no refinement
+levels; level_config and static_circle_tables prepare one blocked level.
 The one-cell same-level halo here is unrelated to the four-cell fine halo of
 the fixed refinement patch (refinement.py), whose semantics are unchanged.
 

src/CelerisLab/lbm/hierarchy.py

@@ -122,7 +122,7 @@ def leaf_partition(h):
     for level, kind in enumerate(h.kinds):
         if level:
             pair = h.pairs[level - 1]
-            py, px = (pair.fine_level_offset if isinstance(pair, FixedPatch2D) else (0, 0))
+            py, px = pair.fine_level_offset
             oy, ox = 2 * oy + py, 2 * ox + px
         y, x = np.nonzero(kind == ACTIVE)
         np.add.at(base, ((y + oy) >> level, (x + ox) >> level), 0.25 ** level)

src/CelerisLab/lbm/refinement.py

@@ -120,6 +120,11 @@ class FixedPatch2D:
         (an index offset, not an (x, y) origin)."""
         return 2 * self.bounds[1] - HALO, 2 * self.bounds[0] - HALO
 
+    @property
+    def rectangles(self) -> tuple:
+        """The one parent rectangle, as a sequence like FixedPatchSet2D.rectangles."""
+        return (self.bounds,)
+
     @property
     def covered(self) -> np.ndarray:
         """Dense coarse mask of cells replaced by the patch (not leaves)."""
@@ -225,11 +230,18 @@ class FixedPatchSet2D:
     def bounds(self) -> tuple:
         return tuple(p.bounds for p in self.patches)
 
+    rectangles = bounds
+
     @property
     def fine_origin(self) -> tuple[float, float]:
         """Lower edge of the level-wide fine grid in coarse physical units."""
         return 0.0, 0.0
 
+    @property
+    def fine_level_offset(self) -> tuple[int, int]:
+        """The level-wide grid itself: zero (y, x) index offset."""
+        return 0, 0
+
     @property
     def covered(self) -> np.ndarray:
         mask = np.zeros(self.coarse_shape, bool)

src/CelerisLab/lbm/refinement_stepper.py

@@ -111,6 +111,57 @@ def _check_levels(refinements, blocks, objects, jets):
             raise ValueError("bodies/Jets are supported only on the finest level of a nested hierarchy")
 
 
+# Host preparation (numpy only, before any device work); each returns one entry per level.
+
+def _level_layouts(cfg, pairs, level_blocks, wall_aware):
+    """Blocked storage layout of every level, checked against the launch limits."""
+    layouts = []
+    for level, block in enumerate(level_blocks):
+        if level == 0:
+            layout = BlockLayout2D.build(cfg.ny, cfg.nx, *block, periodic=(True, not wall_aware))
+        elif isinstance(pairs[level - 1], FixedPatchSet2D):
+            layout = pairs[level - 1].fine_layout(block)
+        else:
+            layout = BlockLayout2D.build(*pairs[level - 1].fine_shape, *block,
+                                         periodic=(False, False))
+        if layout.nblocks > MAX_GRID_DIM or layout.local_shape[0] > MAX_GRID_DIM:
+            raise ValueError(f"level {level} blocks exceed launch limits")
+        layouts.append(layout)
+    return layouts
+
+
+def _host_boundaries(hierarchy, level_cfgs, level_objects, level_jets, wall_aware, shifted):
+    """Host boundary geometry of every refined level in the root edge frame (None at L0)."""
+    hosts = [None]
+    for level in range(1, hierarchy.depth):
+        pair = hierarchy.pairs[level - 1]
+        if level == 1:  # the root config carries the wall-aware boundary fields
+            host = prepare_fine_boundary(level_cfgs[0], pair, level_objects[1],
+                                         wall_aware=wall_aware, jets=level_jets[1],
+                                         _shifted_blocked=shifted)
+        else:
+            origin, spacing = hierarchy.edge_frame(level - 1)
+            host = prepare_fine_boundary(level_cfgs[level - 1], pair, level_objects[level],
+                                         jets=level_jets[level], parent_origin=origin,
+                                         parent_spacing=spacing, _shifted_blocked=shifted)
+        hosts.append(host)
+    return hosts
+
+
+def _fluid_leaves(hierarchy, boundaries):
+    """ACTIVE cells of every level that no static circle or body makes solid."""
+    leaves = [kind == ACTIVE for kind in hierarchy.kinds]
+    for level, b in enumerate(boundaries):
+        if b is None:
+            continue
+        if b.host is None:  # single level: static root circles
+            leaves[level] &= ~b.tables["solid"]
+        else:
+            fine_shape = hierarchy.pairs[level - 1].fine_shape
+            leaves[level] &= (b.host.flags.reshape(fine_shape) & SOLID) == 0
+    return leaves
+
+
 class FixedPatchStepper:
     """D2Q9 SRT/MRT FP32 fixed 2:1 refinement hierarchy, built from level lists.
 
@@ -246,40 +297,15 @@ class FixedPatchStepper:
         self.hierarchy = build_hierarchy(cfg, refinements, wall_aware=self._wall_aware)
         pairs = self.hierarchy.pairs
         if self._wall_aware:
-            rects = (pairs[0].bounds,) if isinstance(pairs[0], FixedPatch2D) else pairs[0].bounds
-            if any(r[1] < 4 or r[3] > cfg.ny - 4 for r in rects):
+            if any(r[1] < 4 or r[3] > cfg.ny - 4 for r in pairs[0].rectangles):
                 raise ValueError("slot_reservoir requires y0 >= 4 and y1 <= ny-4")
-        self._layouts = []
-        for level, block in enumerate(level_blocks):
-            if level == 0:
-                layout = BlockLayout2D.build(cfg.ny, cfg.nx, *block,
-                                             periodic=(True, not self._wall_aware))
-            elif isinstance(pairs[level - 1], FixedPatchSet2D):
-                layout = pairs[level - 1].fine_layout(block)
-            else:
-                layout = BlockLayout2D.build(*pairs[level - 1].fine_shape, *block,
-                                             periodic=(False, False))
-            if layout.nblocks > MAX_GRID_DIM or layout.local_shape[0] > MAX_GRID_DIM:
-                raise ValueError(f"level {level} blocks exceed launch limits")
-            self._layouts.append(layout)
+        self._layouts = _level_layouts(cfg, pairs, level_blocks, self._wall_aware)
         if self._channel:
             base = blocked_coarse_tables(pairs[0], self._layouts[0], np.argsort(pairs[0].coalesce_coarse))
             validate_channel_transfers(pairs[0], self._layouts[0], base[1])
         self._level_cfgs = _kernel_configs(cfg, pairs, self._wall_aware)
-        # Host boundary geometry per level in the root edge frame (Hierarchy.edge_frame).
-        self._boundaries = [None]
-        for level in range(1, depth):
-            pair = pairs[level - 1]
-            if level == 1:  # the root config carries the wall-aware boundary fields
-                host = prepare_fine_boundary(cfg, pair, level_objects[1],
-                                             wall_aware=self._wall_aware, jets=level_jets[1],
-                                             _shifted_blocked=self._ddf_shifting)
-            else:
-                origin, spacing = self.hierarchy.edge_frame(level - 1)
-                host = prepare_fine_boundary(self._level_cfgs[level - 1], pair, level_objects[level],
-                                             jets=level_jets[level], parent_origin=origin,
-                                             parent_spacing=spacing, _shifted_blocked=self._ddf_shifting)
-            self._boundaries.append(host)
+        hosts = _host_boundaries(self.hierarchy, self._level_cfgs, level_objects, level_jets,
+                                 self._wall_aware, self._ddf_shifting)
         if not isinstance(observe_raw_exchange, bool):
             raise ValueError("observe_raw_exchange must be bool")
         self._observe_raw_exchange = observe_raw_exchange
@@ -294,23 +320,17 @@ class FixedPatchStepper:
         self._exchange_spacing = tuple(self.hierarchy.edge_frame(l)[1] for l in range(depth))
         self._exchange_centers = tuple(tuple(tuple(o.geometry.center) for o in objects)
                                        for objects in level_objects)
-        self._level_jets = level_jets
         root = None
         if depth == 1 and level_objects[0]:
             tables = static_circle_tables(self._layouts[0], level_objects[0])
             root = _BlockedBoundary(None, tables, (), False)
+        # The one boundary/Jet record per level; host geometry and Jets are views of it.
         self._device_boundaries = [root] + [
-            _prepare_blocked_boundary(self._boundaries[level], self._layouts[level],
+            _prepare_blocked_boundary(hosts[level], self._layouts[level],
                                       self._level_cfgs[level], level_jets[level],
                                       _shifted_blocked=self._ddf_shifting)
             for level in range(1, depth)]
-        self._leaf = [kind == ACTIVE for kind in self.hierarchy.kinds]
-        if root is not None:
-            self._leaf[0] &= ~root.tables["solid"]
-        for level in range(1, depth):
-            host = self._boundaries[level]
-            if host is not None:
-                self._leaf[level] &= (host.flags.reshape(pairs[level - 1].fine_shape) & SOLID) == 0
+        self._leaf = _fluid_leaves(self.hierarchy, self._device_boundaries)
         self._state = [mask.copy() for mask in self._leaf]
         if self._channel:
             self._leaf[0] &= channel_masks((cfg.ny, cfg.nx))[0]
@@ -556,7 +576,6 @@ class FixedPatchStepper:
             pairs = self.hierarchy.pairs
             if (self._wall_aware or self._channel or self.hierarchy.depth != 3
                     or not isinstance(pairs[1], FixedPatchSet2D)
-                    or self._boundaries[1] is not None
                     or any(b is not None for b in self._device_boundaries[:2])):
                 raise ValueError("remap_finest requires a periodic three-level hierarchy with a "
                                  "level-2 set and no level-0/1 bodies")
@@ -567,8 +586,7 @@ class FixedPatchStepper:
             old = RemapTopology.from_config(cfg, pairs[0], pairs[1].bounds,
                                            child_objects=objects, child_jets=jets,
                                            _shifted_blocked=self._ddf_shifting)
-            if (not same_physical_boundary(old.boundary, self._boundaries[2])
-                    or (self._boundaries[2] is not None and actual is None)):
+            if not same_physical_boundary(old.boundary, None if actual is None else actual.host):
                 raise ValueError("objects must reproduce the original root-edge-frame level-2 boundary")
             if actual is not None:
                 velocity = actual.tables.get("surface_velocity")
@@ -622,7 +640,7 @@ class FixedPatchStepper:
             return candidate
 
     def set_jet_amplitude(self, index, amplitude, *, level):
-        """Replace Jet ``index`` of ``level_jets[level]`` between steps; no geometry/table rebuild.
+        """Replace Jet ``index`` of level ``level`` (``jets[level]``) between steps; no table rebuild.
 
         Invalid input raises before any change. A device/sync error marks the
         instance failed (as step/initialize do); host state commits only after
@@ -634,26 +652,23 @@ class FixedPatchStepper:
                     or not 1 <= level < self.hierarchy.depth):
                 raise ValueError("Jet level must name an existing refined level (1 .. depth-1)")
             boundary = self._device_boundaries[level]
-            jets = list(self._level_jets[level])
+            # Host preparation: validate and build the new column without touching state.
+            jets = list(() if boundary is None else boundary.jets)
             jets[index] = jets[index].with_amplitude(amplitude)  # validates first
-            velocity = None
-            if boundary is not None and boundary.jet_cache is not None:
-                shape, jet_id = boundary.jet_cache
-                velocity = boundary.tables["surface_velocity"].copy()
-                driven = jet_id >= 0
-                velocity[driven] = jet_link_velocity(jets, shape, jet_id)[driven]
+            shape, jet_id = boundary.jet_cache
+            velocity = boundary.tables["surface_velocity"].copy()
+            driven = jet_id >= 0
+            velocity[driven] = jet_link_velocity(jets, shape, jet_id)[driven]
+            # Device action: an uncertain write fails the instance closed.
             try:
-                if velocity is not None:
-                    self._stream.synchronize()  # no in-flight step reads the column
-                    cuda.memcpy_htod(boundary.device["surface_velocity"], velocity)
+                self._stream.synchronize()  # no in-flight step reads the column
+                cuda.memcpy_htod(boundary.device["surface_velocity"], velocity)
             except BaseException:
                 self._status = "failed"
                 raise
-            if velocity is not None:
-                boundary.tables["surface_velocity"] = velocity
-            if boundary is not None:
-                boundary.jets = tuple(jets)
-            self._level_jets[level] = tuple(jets)
+            # Commit: the one boundary record holds both the column and the Jets.
+            boundary.tables["surface_velocity"] = velocity
+            boundary.jets = tuple(jets)
 
     def _blocked_boundary(self, level, buffer, **kw):
         b = level.boundary
@@ -901,9 +916,20 @@ class FixedPatchStepper:
                 raise
 
     # Read-only per-level state, indexed by level (0 = root). Pairs: hierarchy.pairs[l-1].
-    layouts = property(lambda self: tuple(self._layouts))
-    boundaries = property(lambda self: tuple(self._boundaries),
-                          doc="Host boundary geometry per level (None: no bodies/Jets).")
-    jets = property(lambda self: tuple(self._level_jets))
-    leaf = property(lambda self: tuple(m.copy() for m in self._leaf),
-                    doc="Per-level (ny, nx) masks of public (non-NaN) download cells.")
+    @property
+    def layouts(self):
+        return tuple(self._layouts)
+
+    @property
+    def boundaries(self):
+        """Host boundary geometry per level (None: no bodies/Jets on a refined level)."""
+        return tuple(None if b is None else b.host for b in self._device_boundaries)
+
+    @property
+    def jets(self):
+        return tuple(() if b is None else b.jets for b in self._device_boundaries)
+
+    @property
+    def leaf(self):
+        """Per-level (ny, nx) masks of public (non-NaN) download cells."""
+        return tuple(m.copy() for m in self._leaf)

tests/unit/test_runtime_child_remap.py

@@ -1,5 +1,6 @@
 """CPU mocks execute production remap lifecycle; no GPU numerical claim."""
 from dataclasses import replace
+from types import SimpleNamespace
 import unittest
 from unittest.mock import patch
 
@@ -69,10 +70,11 @@ class RuntimeChildRemap(unittest.TestCase):
         run._status = 'ready'
         for bad in ([], [(40,16,48,40)], [(16,16,32,24),(16,27,32,55)]):
             with self.assertRaises(ValueError): run.remap_finest(bad)
-        hosts, pairs = run._boundaries, run.hierarchy.pairs
+        records, pairs = run._device_boundaries, run.hierarchy.pairs
         for attribute, value in (('_wall_aware', True),
-                                 ('_boundaries', [*hosts[:2], object()]),        # child boundary
-                                 ('_boundaries', [hosts[0], object(), hosts[2]]),  # fine boundary
+                                 # level-2 record whose host geometry differs from objects=()
+                                 ('_device_boundaries', [*records[:2], SimpleNamespace(host=object(), jets=())]),
+                                 ('_device_boundaries', [records[0], object(), records[2]]),  # level-1 body
                                  ('hierarchy', replace(run.hierarchy, pairs=(pairs[0], None)))):
             with patch.object(run, attribute, value):
                 with self.assertRaises(ValueError): run.remap_finest(B)