Support bounded three-level separated patch unions

src/CelerisLab/lbm/refinement.py

@@ -318,11 +318,35 @@ def nest_patch(cfg, parent, child_bounds):
     would read a covered or OUTER cell without a child-supplied arrival; and
     child coverage of L1's own two-hop restriction footprint toward L0.
     """
+    child = FixedPatch2D.from_config(_middle_config(cfg, parent), child_bounds)
+    return child, _nested_kind(parent, child)
+
+
+def nest_patch_set(cfg, parent, child_rectangles_in_parent):
+    """Separated child union in the parent's allocated L1 array coordinates.
+
+    The root edge coordinate of L1 index (x, y) is
+    ``parent.fine_origin + (x, y)/2``. For a set parent its origin is zero;
+    a scalar parent's window includes its ghost band and has a shifted origin.
+    The child set allocates a dense public extent twice that L1 array, while
+    its blocked storage omits holes. Both adjacent pairs retain gap >= 4 in
+    their respective parent-cell units. Validate support over the whole union.
+    """
+    child = FixedPatchSet2D.from_config(
+        _middle_config(cfg, parent), child_rectangles_in_parent)
+    return child, _nested_kind(parent, child)
+
+
+def _middle_config(cfg, parent):
+    fy, fx = parent.fine_shape
+    return replace(cfg, nx=fx, ny=fy, viscosity=2 * cfg.viscosity, y_wall_bc="bounce_back",
+                   top_wall_bc="velocity", wall_slot_width=0, wall_normal_velocity=0.0,
+                   velocity=0.0)
+
+
+def _nested_kind(parent, child):
+    """Validate composed transfer support and preserve inherited ghost kinds."""
     fy, fx = parent.fine_shape
-    middle = replace(cfg, nx=fx, ny=fy, viscosity=2 * cfg.viscosity, y_wall_bc="bounce_back",
-                     top_wall_bc="velocity", wall_slot_width=0, wall_normal_velocity=0.0,
-                     velocity=0.0)
-    child = FixedPatch2D.from_config(middle, child_bounds)
     kind = parent.fine_kind
     covered = child.covered
     active = (kind == ACTIVE) & ~covered
@@ -346,7 +370,7 @@ def nest_patch(cfg, parent, child_bounds):
         raise ValueError("child covers the parent's restriction footprint")
     composed = kind.copy()
     composed[covered] = OUTER_GHOST
-    return child, _readonly(composed, np.uint8)
+    return _readonly(composed, np.uint8)
 
 
 def prepare_fine_boundary(cfg, patch, objects, *, wall_aware=False, jets=(),

src/CelerisLab/lbm/refinement_stepper.py

@@ -19,7 +19,7 @@ from .curved_links import surface_link_velocity
 from ..body.jets import jet_link_shapes, jet_link_velocity
 from .block_stepper import MAX_GRID_DIM
 from .refinement import (ACTIVE, HALO, OUTER_GHOST, FixedPatch2D, FixedPatchSet2D,
-                         nest_patch, prepare_fine_boundary)
+                         nest_patch, nest_patch_set, prepare_fine_boundary)
 from .descriptors import SOLID, D2Q9_EX, D2Q9_EY
 
 
@@ -61,23 +61,34 @@ class _Pair:
     parent_cells: int    # parent block window cells (explode donor stride)
 
 
-def blocked_base_tables(patch, layout, order, wall_aware=False):
+def blocked_base_tables(patch, layout, order, wall_aware=False, *, composed_kind=None):
     """Parent donor cells and receiver-local arrival slots, in transfer row order.
 
     Stage arrivals only after child reads and parent halo gather, immediately
     before parent pull. Source *owner* slots are wrong at block seams.
+    Nested pairs must pass composed_kind to retain inherited INNER/OUTER and
+    sparse holes; donors and receivers must be owned, uncovered ACTIVE cells.
     """
     owner = layout.owner.ravel()
     q, receiver = np.divmod(patch.coalesce_coarse[order], np.uint64(owner.size))
+    if composed_kind is None:
+        dense_kind = np.where(patch.covered, OUTER_GHOST, ACTIVE).astype(np.uint8)
+        if wall_aware:
+            dense_kind[[0, -1]] = OUTER_GHOST
+    else:
+        dense_kind = np.asarray(composed_kind)
+        if dense_kind.shape != patch.coarse_shape:
+            raise ValueError("composed parent kind has the wrong shape")
+    donors = patch.ghost_coarse.astype(np.intp)
+    receivers = receiver.astype(np.intp)
+    legal = layout.owned.ravel() & (dense_kind.ravel() == ACTIVE) & ~patch.covered.ravel()
+    if not legal[donors].all() or not legal[receivers].all():
+        raise ValueError("child donors/receivers need owned uncovered parent ACTIVE cells")
     offset = (np.asarray(D2Q9_EY)[q.astype(np.intp)] * layout.local_shape[1]
               + np.asarray(D2Q9_EX)[q.astype(np.intp)])
-    source = owner[receiver.astype(np.intp)].astype(np.int64) - offset
-    kind = np.full(layout.nblocks * int(np.prod(layout.local_shape)), OUTER_GHOST, np.uint8)
-    leaves = ~patch.covered
-    if wall_aware:
-        leaves[[0, -1]] = False
-    kind[owner[leaves.ravel()].astype(np.intp)] = ACTIVE
-    return (owner[patch.ghost_coarse].astype(np.uint64), layout.slot(source, q), kind)
+    source = owner[receivers].astype(np.int64) - offset
+    return (owner[donors].astype(np.uint64), layout.slot(source, q),
+            blocked_kind(layout, dense_kind))
 
 
 @dataclass
@@ -171,6 +182,16 @@ class FixedPatchStepper:
     shape, NaN outside the patches' ACTIVE cells. The step sequence is the
     blocked one above, with every patch's transfers in one launch each.
 
+    Three-level separated unions: ``child_rectangles_in_parent=[...]`` is
+    mutually exclusive with scalar ``child`` and accepts either parent form.
+    Coordinates index the allocated L1 array, including ghost bands and holes:
+    root edge = patch.fine_origin + L1 index / 2. A parent set has origin zero;
+    a scalar parent retains its shifted window origin. The child set's public
+    extent is twice the allocated L1 shape, even though storage is sparse.
+    Each pair requires gap >= 4 in its own parent-cell units. All regions share
+    one buffer pair per level and the fixed 1:2:4 schedule; L2-only bodies/Jets
+    below apply unchanged, with object IDs assigned over the whole level.
+
     Three-level single chain: ``child`` bounds use L1 window cells; all three
     levels must be blocked. ``child_objects`` and ``child_jets`` use the L0
     edge frame, with L2 dx=dt=.25 and unchanged lattice velocity. L1 bodies
@@ -181,11 +202,14 @@ class FixedPatchStepper:
     """
     def __init__(self, cfg, bounds, *, context, coarse_boundary="periodic", fine_objects=(),
                  fine_block=None, base_block=None, jets=(), child=None, child_block=None,
-                 child_objects=(), child_jets=()):
+                 child_objects=(), child_jets=(), child_rectangles_in_parent=None):
         if HALO != 4:
             raise RuntimeError("refinement.cu requires a four-cell fine halo")
         if coarse_boundary not in ("periodic", "slot_reservoir"):
             raise ValueError("coarse_boundary must be periodic or slot_reservoir")
+        if child is not None and child_rectangles_in_parent is not None:
+            raise ValueError("child and child_rectangles_in_parent are mutually exclusive")
+        nested = child is not None or child_rectangles_in_parent is not None
         bounds = tuple(bounds)
         multipatch = len(bounds) > 0 and not isinstance(bounds[0], (Integral, np.integer))
         if multipatch:
@@ -205,7 +229,7 @@ class FixedPatchStepper:
                 raise ValueError("slot_reservoir requires y0 >= 4 and y1 <= ny-4")
         fine_objects, jets = tuple(fine_objects), tuple(jets)
         child_objects, child_jets = tuple(child_objects), tuple(child_jets)
-        if child is None and (child_objects or child_jets):
+        if not nested and (child_objects or child_jets):
             raise ValueError("child_objects/child_jets require a nested child")
         if base_block is not None and fine_block is None:
             raise ValueError("base_block requires fine_block")
@@ -230,17 +254,23 @@ class FixedPatchStepper:
                 raise ValueError("fine blocks exceed launch limits")
         self.child_patch = self.child_layout = None
         self._middle_kind = None
-        if child is not None:
-            if multipatch or fine_block is None or base_block is None or child_block is None:
-                raise ValueError("a nested child needs one patch plus fine_block, base_block "
-                                 "and child_block")
+        if nested:
+            if child is not None and multipatch:
+                raise ValueError("scalar child requires one parent patch; use child_rectangles_in_parent")
+            if fine_block is None or base_block is None or child_block is None:
+                raise ValueError("a nested child needs fine_block, base_block and child_block")
             if tuple(fine_objects) or tuple(jets):
                 raise ValueError("a nested child supports bodies/Jets only on L2; use child_objects/child_jets")
             if len(child_block) != 2:
                 raise ValueError("child_block must be (by, bx)")
-            self.child_patch, self._middle_kind = nest_patch(cfg, self.patch, child)
-            self.child_layout = BlockLayout2D.build(*self.child_patch.fine_shape, *child_block,
-                                                    periodic=(False, False))
+            if child_rectangles_in_parent is not None:
+                self.child_patch, self._middle_kind = nest_patch_set(
+                    cfg, self.patch, child_rectangles_in_parent)
+                self.child_layout = self.child_patch.fine_layout(child_block)
+            else:
+                self.child_patch, self._middle_kind = nest_patch(cfg, self.patch, child)
+                self.child_layout = BlockLayout2D.build(*self.child_patch.fine_shape, *child_block,
+                                                       periodic=(False, False))
             if (self.child_layout.nblocks > MAX_GRID_DIM
                     or self.child_layout.local_shape[0] > MAX_GRID_DIM):
                 raise ValueError("child blocks exceed launch limits")
@@ -365,7 +395,8 @@ class FixedPatchStepper:
                     g, cp = self._levels[2], self.child_patch
                     order = np.argsort(cp.coalesce_coarse)
                     ghosts, children, g_kind = blocked_fine_tables(cp, g.layout, order)
-                    donors, targets, _ = blocked_base_tables(cp, f.layout, order)
+                    donors, targets, _ = blocked_base_tables(
+                        cp, f.layout, order, composed_kind=self._middle_kind)
                     g.kind = self._upload(g_kind)
                     g.halo_copy = g.module.get_function("BlockHaloCopy")
                     self._child = _Pair(g.module.get_function("BlockExplode"),