Multipatch: FixedPatchSet2D host metadata and opt-in FixedPatchStepper path

Several separated (>= 4 coarse cells) body-free 2:1 patches share one
level-wide fine grid stored by a sparse BlockLayout2D over the patch
supports. FixedPatchSet2D re-addresses each FixedPatch2D's tables to level
coordinates (same rows/child order, canonical patch order). The stepper
accepts a sequence of rectangles as bounds (requires fine_block and
base_block; rejects fine_objects/jets), reusing the existing blocked step
sequence and kernels. Helpers use a covered mask instead of one rectangle
and skip unstored fine holes. No kernel changes.

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

src/CelerisLab/lbm/refinement.py

@@ -114,6 +114,19 @@ class FixedPatch2D:
     def fine_origin(self) -> tuple[float, float]:
         return self.bounds[0] - HALO / 2, self.bounds[1] - HALO / 2
 
+    @property
+    def fine_level_origin(self) -> tuple[int, int]:
+        """(y, x) of this window's first cell on the level-wide 2:1 fine grid."""
+        return 2 * self.bounds[1] - HALO, 2 * self.bounds[0] - HALO
+
+    @property
+    def covered(self) -> np.ndarray:
+        """Dense coarse mask of cells replaced by the patch (not leaves)."""
+        x0, y0, x1, y1 = self.bounds
+        mask = np.zeros(self.coarse_shape, bool)
+        mask[y0:y1, x0:x1] = True
+        return mask
+
     @property
     def table_bytes(self) -> int:
         return sum(a.nbytes for a in (self.fine_kind, self.ghost_fine,
@@ -138,6 +151,104 @@ class FixedPatch2D:
             raise ValueError("levels are not synchronized: fine_steps must equal 2*coarse_steps")
 
 
+MIN_PATCH_GAP = 4  # coarse cells: the full 4-fine-cell supports stay disjoint
+
+
+@dataclass(frozen=True, eq=False)
+class FixedPatchSet2D:
+    """Several separated FixedPatch2D rectangles sharing one 2:1 fine level.
+
+    The fine level is the whole (2*ny, 2*nx) grid in level coordinates; only
+    the patch supports (active region + HALO ghost band) are stored, by a
+    sparse BlockLayout2D from fine_layout(). Tables have FixedPatch2D meaning
+    and row/child order, re-addressed to level cells: ghost_fine are level
+    fine cells, coalesce_fine are q*N+cell level slots; coarse ids unchanged.
+    Patches are kept in canonical (y0, x0) order, so input order is irrelevant.
+    Active regions must be >= MIN_PATCH_GAP coarse cells apart (Chebyshev).
+    """
+    coarse_shape: tuple[int, int]
+    patches: tuple
+    tau_coarse: float
+    tau_fine: float
+    fine_kind: np.ndarray
+    ghost_fine: np.ndarray
+    ghost_coarse: np.ndarray
+    coalesce_coarse: np.ndarray
+    coalesce_fine: np.ndarray
+
+    @classmethod
+    def from_config(cls, cfg: LBMConfig, rectangles) -> "FixedPatchSet2D":
+        rectangles = tuple(rectangles)
+        if not rectangles:
+            raise ValueError("at least one patch rectangle is required")
+        patches = sorted((FixedPatch2D.from_config(cfg, r) for r in rectangles),
+                         key=lambda p: (p.bounds[1], p.bounds[0]))
+        for i, a in enumerate(patches):
+            for b in patches[i + 1:]:
+                (ax0, ay0, ax1, ay1), (bx0, by0, bx1, by1) = a.bounds, b.bounds
+                if max(ay0 - by1, by0 - ay1, ax0 - bx1, bx0 - ax1) < MIN_PATCH_GAP:
+                    raise ValueError(f"patches {a.bounds} and {b.bounds} are closer "
+                                     f"than {MIN_PATCH_GAP} coarse cells")
+        ny, nx = patches[0].coarse_shape
+        fy, fx = 2 * ny, 2 * nx
+        if fy * fx * 9 > np.iinfo(np.uint64).max // 2:
+            raise ValueError("fine level exceeds 64-bit indexing")
+        kind = np.full((fy, fx), OUTER_GHOST, np.uint8)
+        ghost_fine, coalesce_fine = [], []
+        for p in patches:
+            y0, x0 = p.fine_level_origin
+            ly, lx = p.fine_shape
+            kind[y0:y0 + ly, x0:x0 + lx] = p.fine_kind
+            gy, gx = np.divmod(p.ghost_fine.astype(np.int64), lx)
+            ghost_fine.append((gy + y0) * fx + gx + x0)
+            q, cell = np.divmod(p.coalesce_fine.astype(np.int64), ly * lx)
+            cy, cx = np.divmod(cell, lx)
+            coalesce_fine.append(q * fy * fx + (cy + y0) * fx + cx + x0)
+        ghost_coarse = np.concatenate([p.ghost_coarse for p in patches])
+        coalesce_coarse = np.concatenate([p.coalesce_coarse for p in patches])
+        covered = np.zeros((ny, nx), bool)
+        for p in patches:
+            covered |= p.covered
+        cells = ny * nx
+        if (covered.ravel()[ghost_coarse].any()
+                or covered.ravel()[coalesce_coarse % np.uint64(cells)].any()
+                or len(np.unique(coalesce_coarse)) != len(coalesce_coarse)):
+            raise ValueError("patch transfers overlap another patch")
+        return cls((ny, nx), tuple(patches), patches[0].tau_coarse, patches[0].tau_fine,
+                   _readonly(kind, np.uint8), _readonly(np.concatenate(ghost_fine), np.uint64),
+                   _readonly(ghost_coarse, np.uint32), _readonly(coalesce_coarse, np.uint64),
+                   _readonly(np.concatenate(coalesce_fine), np.uint64).reshape(-1, 4))
+
+    @property
+    def fine_shape(self) -> tuple[int, int]:
+        return self.fine_kind.shape
+
+    @property
+    def bounds(self) -> tuple:
+        return tuple(p.bounds for p in self.patches)
+
+    @property
+    def covered(self) -> np.ndarray:
+        mask = np.zeros(self.coarse_shape, bool)
+        for p in self.patches:
+            mask |= p.covered
+        return mask
+
+    def fine_layout(self, block):
+        """Sparse non-periodic layout storing exactly the patch supports."""
+        from .blocks import BlockLayout2D
+        by, bx = block
+        origins = set()
+        for p in self.patches:
+            y0, x0 = p.fine_level_origin
+            ly, lx = p.fine_shape
+            if y0 % by or x0 % bx or ly % by or lx % bx:
+                raise ValueError(f"patch {p.bounds} support is not aligned to fine blocks {block}")
+            origins.update((y, x) for y in range(y0, y0 + ly, by) for x in range(x0, x0 + lx, bx))
+        return BlockLayout2D.build(*self.fine_shape, by, bx, periodic=(False, False),
+                                   origins=sorted(origins))
+
+
 @dataclass(frozen=True, eq=False)
 class FineBoundary2D:
     """Fixed local projection plus auditable, finite transfer dependency sets.

src/CelerisLab/lbm/refinement_stepper.py

@@ -18,7 +18,8 @@ from .blocks import BlockLayout2D, blocked_link_tables
 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, prepare_fine_boundary
+from .refinement import (ACTIVE, HALO, OUTER_GHOST, FixedPatch2D, FixedPatchSet2D,
+                         prepare_fine_boundary)
 from .descriptors import SOLID, D2Q9_EX, D2Q9_EY
 
 
@@ -30,10 +31,11 @@ def blocked_fine_tables(patch, layout, order):
     unused halo entries OUTER). Coarse donors/targets are unchanged.
     """
     owner = layout.owner.ravel().astype(np.intp)
+    owned = layout.owned.ravel()  # sparse multipatch levels leave holes unstored
     cells = patch.fine_shape[0] * patch.fine_shape[1]
     q, cell = np.divmod(patch.coalesce_fine[order], np.uint64(cells))
     kind = np.full(layout.nblocks * int(np.prod(layout.local_shape)), OUTER_GHOST, np.uint8)
-    kind[owner] = patch.fine_kind.ravel()
+    kind[owner[owned]] = patch.fine_kind.ravel()[owned]
     return (owner[patch.ghost_fine.astype(np.intp)].astype(np.uint64),
             layout.slot(owner[cell.astype(np.intp)], q), kind)
 
@@ -50,9 +52,7 @@ def blocked_base_tables(patch, layout, order, wall_aware=False):
               + 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 = np.ones(patch.coarse_shape, bool)
-    x0, y0, x1, y1 = patch.bounds
-    leaves[y0:y1, x0:x1] = False
+    leaves = ~patch.covered
     if wall_aware:
         leaves[[0, -1]] = False
     kind[owner[leaves.ravel()].astype(np.intp)] = ACTIVE
@@ -110,6 +110,13 @@ class FixedPatchStepper:
     ``jets`` (body.jets.Jet) use the same edge frame as ``fine_objects`` and
     drive fine cut links selected at their true hits; ``set_jet_amplitude``
     re-uploads the per-link velocity column only.
+
+    Several patches: pass ``bounds`` as a sequence of (x0, y0, x1, y1)
+    rectangles (FixedPatchSet2D, >= 4 coarse cells apart). This requires
+    fine_block and base_block, is body-free (no fine_objects/jets yet), and
+    uses one level-wide fine grid (2*ny, 2*nx): public fine arrays have that
+    shape, NaN outside the patches' ACTIVE cells. The step sequence is the
+    blocked one above, with every patch's transfers in one launch each.
     """
     def __init__(self, cfg, bounds, *, context, coarse_boundary="periodic", fine_objects=(),
                  fine_block=None, base_block=None, jets=()):
@@ -117,14 +124,24 @@ class FixedPatchStepper:
             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")
-        self.patch = FixedPatch2D.from_config(cfg, bounds)
+        bounds = tuple(bounds)
+        multipatch = len(bounds) > 0 and not isinstance(bounds[0], (Integral, np.integer))
+        if multipatch:
+            if fine_block is None or base_block is None:
+                raise ValueError("several patches require fine_block and base_block")
+            if tuple(fine_objects) or tuple(jets):
+                raise ValueError("several patches do not support fine_objects or jets yet")
+            self.patch = FixedPatchSet2D.from_config(cfg, bounds)
+        else:
+            self.patch = FixedPatch2D.from_config(cfg, bounds)
         self._wall_aware = coarse_boundary == "slot_reservoir"
         if self._wall_aware:
             if (cfg.y_wall_bc != "on_site_velocity" or cfg.top_wall_bc != "reservoir"
                     or cfg.velocity != 0 or not 2 <= cfg.wall_slot_width < cfg.nx):
                 raise ValueError("slot_reservoir requires on_site_velocity, reservoir, "
                                  "velocity=0 and 2 <= wall_slot_width < nx")
-            if self.patch.bounds[1] < 4 or self.patch.bounds[3] > cfg.ny - 4:
+            rects = self.patch.bounds if multipatch else (self.patch.bounds,)
+            if any(r[1] < 4 or r[3] > cfg.ny - 4 for r in rects):
                 raise ValueError("slot_reservoir requires y0 >= 4 and y1 <= ny-4")
         fine_objects = tuple(fine_objects)
         if base_block is not None and fine_block is None:
@@ -143,7 +160,8 @@ class FixedPatchStepper:
             if len(fine_block) != 2:
                 raise ValueError("fine_block must be (by, bx)")
             fy, fx = self.patch.fine_shape
-            self.fine_layout = BlockLayout2D.build(fy, fx, *fine_block, periodic=(False, False))
+            self.fine_layout = (self.patch.fine_layout(fine_block) if multipatch else
+                                BlockLayout2D.build(fy, fx, *fine_block, periodic=(False, False)))
             if (self.fine_layout.nblocks > MAX_GRID_DIM
                     or self.fine_layout.local_shape[0] > MAX_GRID_DIM):
                 raise ValueError("fine blocks exceed launch limits")
@@ -297,8 +315,7 @@ class FixedPatchStepper:
         return self._step_count
 
     def _mask_nonleaves(self, coarse, fine):
-        x0, y0, x1, y1 = self.patch.bounds
-        coarse[:, y0:y1, x0:x1] = np.nan
+        coarse[:, self.patch.covered] = np.nan
         if self._wall_aware:
             coarse[:, (0, -1), :] = np.nan
         fine[:, ~self._fine_leaf] = np.nan
@@ -311,9 +328,7 @@ class FixedPatchStepper:
                       for a in (coarse_g, fine_g)]
             if any(a.shape != level.shape for a, level in zip(arrays, self._levels)):
                 raise ValueError("Expected Q-major (9,ny,nx) arrays for both levels")
-            x0, y0, x1, y1 = self.patch.bounds
-            leaf = np.ones(self.patch.coarse_shape, dtype=bool)
-            leaf[y0:y1, x0:x1] = False
+            leaf = ~self.patch.covered
             if self._wall_aware:
                 leaf[[0, -1]] = False
             for array, mask in zip(arrays, (leaf, self._fine_leaf)):

tests/unit/test_refinement.py

@@ -163,5 +163,72 @@ class FixedPatchTests(unittest.TestCase):
                 p.check_synchronization(c, f, phase=phase)
 
 
+class FixedPatchSetTests(unittest.TestCase):
+    """Several separated patches: level-wide re-addressing of FixedPatch2D tables."""
+    FACE = ((4, 4, 8, 12), (12, 4, 16, 12))          # 16x24 coarse, face gap 4
+    S1 = ((4, 4, 12, 12), (20, 4, 28, 12))           # 16x32 coarse, translation copy
+
+    def build(self, rects, ny=16, nx=24):
+        return refinement.FixedPatchSet2D.from_config(LBMConfig(nx=nx, ny=ny, viscosity=.1), rects)
+
+    def test_tables_are_shifted_single_patch_tables(self):
+        s = self.build(self.FACE)
+        fy, fx = s.fine_shape
+        self.assertEqual((fy, fx), (32, 48))
+        rows = cells = 0
+        for p in s.patches:
+            y0, x0 = p.fine_level_origin
+            ly, lx = p.fine_shape
+            np.testing.assert_array_equal(s.fine_kind[y0:y0 + ly, x0:x0 + lx], p.fine_kind)
+            gy, gx = np.divmod(p.ghost_fine.astype(np.int64), lx)
+            np.testing.assert_array_equal(s.ghost_fine[cells:cells + len(gy)],
+                                          (gy + y0) * fx + gx + x0)
+            np.testing.assert_array_equal(s.ghost_coarse[cells:cells + len(gy)], p.ghost_coarse)
+            q, c = np.divmod(p.coalesce_fine.astype(np.int64), ly * lx)
+            cy, cx = np.divmod(c, lx)
+            np.testing.assert_array_equal(s.coalesce_fine[rows:rows + len(q)],
+                                          q * fy * fx + (cy + y0) * fx + cx + x0)
+            np.testing.assert_array_equal(s.coalesce_coarse[rows:rows + len(q)], p.coalesce_coarse)
+            rows, cells = rows + len(q), cells + len(gy)
+        self.assertEqual((rows, cells), (len(s.coalesce_coarse), len(s.ghost_fine)))
+        self.assertEqual(int((s.fine_kind == ACTIVE).sum()), 4 * int(s.covered.sum()))
+
+    def test_input_order_is_irrelevant(self):
+        a, b = self.build(self.FACE), self.build(self.FACE[::-1])
+        for name in ("fine_kind", "ghost_fine", "ghost_coarse", "coalesce_coarse", "coalesce_fine"):
+            np.testing.assert_array_equal(getattr(a, name), getattr(b, name))
+        self.assertEqual(a.bounds, b.bounds)
+
+    def test_one_patch_set_matches_fixed_patch(self):
+        p = FixedPatch2D.from_config(config(), (4, 4, 12, 12))
+        s = self.build(((4, 4, 12, 12),), nx=16)
+        np.testing.assert_array_equal(s.covered, p.covered)
+        self.assertEqual((s.tau_coarse, s.tau_fine), (p.tau_coarse, p.tau_fine))
+        np.testing.assert_array_equal(s.coalesce_coarse, p.coalesce_coarse)
+        np.testing.assert_array_equal(s.ghost_coarse, p.ghost_coarse)
+
+    def test_gap_and_overlap_rejected(self):
+        for gap in (0, 2, 3):
+            with self.subTest(gap=gap), self.assertRaisesRegex(ValueError, "closer"):
+                self.build(((4, 4, 8, 12), (8 + gap, 4, 12 + gap, 12)))
+        with self.assertRaises(ValueError):
+            self.build(())
+        with self.assertRaises(ValueError):  # FixedPatch2D margin still applies
+            self.build(((1, 4, 8, 12), (12, 4, 16, 12)))
+
+    def test_sparse_fine_layout_stores_exactly_the_supports(self):
+        for rects, nx in ((self.FACE, 24), (self.S1, 32)):
+            s = self.build(rects, nx=nx)
+            layout = s.fine_layout((4, 4))
+            support = np.zeros(s.fine_shape, bool)
+            for p in s.patches:
+                y0, x0 = p.fine_level_origin
+                support[y0:y0 + p.fine_shape[0], x0:x0 + p.fine_shape[1]] = True
+            np.testing.assert_array_equal(layout.owned, support)
+            self.assertTrue(layout.owned.ravel()[s.ghost_fine.astype(np.intp)].all())
+        with self.assertRaisesRegex(ValueError, "aligned"):
+            self.build(self.FACE).fine_layout((5, 4))
+
+
 if __name__ == "__main__":
     unittest.main(verbosity=2)