Merge a454f023 (claude/jet-shapes-once) for integration review

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

README.md

@@ -125,8 +125,8 @@ with CudaContext(0) as ctx, FixedPatchStepper(cfg, context=ctx, refinements=[(8,
 
 Scope and limits:
 - The examples below use D2Q9, SRT, FP32, double buffer, unshifted, no LES or
-  forcing. Fixed 2:1 hierarchies can use level lists of arbitrary depth; the
+  forcing. Fixed 2:1 hierarchies use level lists of arbitrary depth; two- and
-  legacy two/three-level forms are shown below.
+  three-level examples below use that same constructor.
 - `coarse_boundary="periodic"` (default) or `"slot_reservoir"` (bottom slot, top
   reservoir; patch needs `y0 >= 4`, `y1 <= ny-4`). Periodic mode ignores the wall/slot config fields.
 - Object centers/radii use the coarse **cell-edge** frame (cell `I` spans `[I, I+1)`),
@@ -151,7 +151,7 @@ Scope and limits:
   inside the union's active region with
   its boundary stencil clear of transfer cells; a jet may not select links of two
   different bodies, even across regions. Wall-touching patches remain unsupported.
-  For nesting, use the level lists or legacy child forms below.
+  For nesting, add further `refinements` entries as shown below.
 
   ```python
   cfg = LBMConfig(nx=32, ny=16, viscosity=.1)
@@ -165,19 +165,20 @@ Scope and limits:
       coarse_g, fine_g = run.download()  # coarse: NaN on every covered cell;
                                          # fine: NaN on every non-ACTIVE cell
   ```
-- Legacy three-level scalar chain: add `child=(x0, y0, x1, y1)` in
+- Three-level nested windows: pass two scalar rectangles in `refinements` —
-  **L1 window cells** (the patch's fine grid including its 4-cell ghost band) and
+  L1 in L0 cells, L2 in **L1 window cells** (the parent's fine grid including its
-  all three `level_blocks` entries. L2 must sit inside L1's
+  4-cell ghost band) — plus all three `level_blocks` entries. L2 must sit inside
-  ACTIVE cells, away from L1's ghost band and transfer footprint (checked before
+  L1's ACTIVE cells, away from L1's ghost band and transfer footprint (checked
-  any GPU allocation). `initialize` and `download` take/return three arrays; each
+  before any GPU allocation). `initialize` and `download` take/return three
-  level is NaN on its non-ACTIVE (and child-covered) cells. One coarse step = 2 L1
+  arrays; each level is NaN on its non-ACTIVE (and child-covered) cells. One
-  steps = 4 L2 steps. One L1 patch only. Optional stationary circles and Jets
+  coarse step = 2 L1 steps = 4 L2 steps. One L1 window only in this scalar
-  live on L2 via `level_objects[2]` and `level_jets[2]`; L1 bodies/Jets retain
+  chain. Optional stationary circles and Jets live on L2 via `level_objects[2]`
-  their L1 meaning and must be empty in this mode. All object and Jet geometry
+  and `level_jets[2]`; L1 bodies/Jets retain their L1 meaning and must be empty
-  uses the **L0 edge frame**, not the child bounds' L1 window frame. L2 has
+  in this mode. All object and Jet geometry uses the **L0 edge frame**, not the
-  dx=dt=0.25, so prescribed lattice velocity is unchanged. Each whole body and
+  child bounds' L1 window frame. L2 has dx=dt=0.25, so prescribed lattice
-  its BFL stencil must be inside L2 ACTIVE cells, outside the transfer footprint.
+  velocity is unchanged. Each whole body and its BFL stencil must be inside L2
-  A Jet may select one physical body across multiple blocks, never two bodies.
+  ACTIVE cells, outside the transfer footprint. A Jet may select one physical
+  body across multiple blocks, never two bodies.
 
   ```python
   from CelerisLab.body.geometry.circle import CircleGeometry
@@ -199,16 +200,16 @@ Scope and limits:
       run.step(4)
       coarse_g, l1_g, l2_g = run.download()  # solids/covered/nonleaves are NaN
   ```
-- Level lists: `refinements[l]` describes level `l+1` in allocated level-`l`
+- Level lists (the only constructor): `refinements[l]` describes level `l+1`
-  integer cells; use a rectangle list for a union, or a scalar rectangle for a
+  in allocated level-`l` integer cells; use a rectangle list for a union, or a
-  single window (a scalar child requires a scalar parent). `level_blocks`,
+  scalar rectangle for a single window (a scalar child requires a scalar
-  `level_objects` and `level_jets` each have one entry per level, including L0.
+  parent). `level_blocks`, `level_objects` and `level_jets` each have one entry
-  These lists are mutually exclusive with the legacy patch keywords. Three or
+  per level, including L0. Three or more levels require all-blocked storage and
-  more levels require all-blocked storage and bodies/Jets only on the finest
+  bodies/Jets only on the finest level. Every adjacent pair must preserve
-  level. Every adjacent pair must preserve inherited ACTIVE coverage, valid
+  inherited ACTIVE coverage, valid donors/receivers, surviving pulls and the
-  donors/receivers, surviving pulls and the upstream transfer footprint.
+  upstream transfer footprint. Level `l` takes `2**l` steps per base interval.
-  Level `l` takes `2**l` steps per base interval. This describes implemented
+  This describes implemented topology and scheduling, not accuracy or
-  topology and scheduling, not accuracy or convergence qualification.
+  convergence qualification.
 - Blocked body levels can enable `observe_raw_exchange=True` and call
   `read_raw_exchange(body_id, level=...)` after a completed base interval.
   It reduces unique cut-link observations into wall/aperture/total force,

docs/fixed_three_level_multipatch.md

@@ -1,8 +1,8 @@
 # Fixed three-level rectangle unions
 
-This page describes the legacy three-level entry point of `FixedPatchStepper`,
+This page describes three-level `FixedPatchStepper` setups built from the
-with rectangle unions on either adjacent pair. The `refinements` level-list
+`refinements` level list, with rectangle unions on either adjacent pair. The
-entry point also supports arbitrary depth; three or more levels require
+same constructor supports arbitrary depth; three or more levels require
 all-blocked storage and bodies/Jets only on the finest level. This is not
 general AMR: there is no per-region timestep, moving geometry, automatic
 dynamic refinement or L1 body support in the three-level form.
@@ -12,15 +12,15 @@ No CUDA kernel changes are involved.
 
 ## Explicit coordinate frames
 
-Keep the existing scalar `child=(x0,y0,x1,y1)` for a single scalar parent patch.
+For a single scalar parent patch, pass a scalar rectangle as `refinements[1]`:
-Its coordinates retain their original meaning: half-open integer indices in the
+`(x0,y0,x1,y1)` in half-open integer indices of the allocated L1 window,
-allocated L1 window, including its ghost band.
+including its ghost band.
 
-A rectangle list as the second `refinements` entry (formerly the keyword
+A rectangle list as `refinements[1]` accepts a scalar or a set parent and
-`child_rectangles_in_parent=[...]`) accepts a scalar or a set parent and interprets every rectangle in
+interprets every rectangle in the actual allocated L1 array, including ghost
-the actual allocated L1 array, including ghost bands and holes. It never infers
+bands and holes. It never infers region-local coordinates. The physical root
-region-local coordinates. The physical root cell edge corresponding to an L1
+cell edge corresponding to an L1 index `(x,y)` is
-index `(x,y)` is `parent.fine_origin + (x,y)/2`:
+`parent.fine_origin + (x,y)/2`:
 
 - A scalar L1 parent has a shifted window origin.
 - A set L1 parent has origin `(0,0)` and public shape `(2*ny,2*nx)`.

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/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/kernels/blocks.cu

@@ -168,21 +168,15 @@ extern "C" __global__ void BlockStep(const fpxx* in, fpxx* out, unsigned int hal
         stream_pull_load(k, f, in + base, j);
         if (cell_kind == KIND_ACTIVE) {
             if (d_block_wall.nx != 0) {
-                const unsigned int bx = NX - 2 * halo, by = NY - 2 * halo;
+                unsigned int gx, gy;
-                const unsigned int nbx = d_block_wall.nx / bx;
+                block_global_xy(x, y, halo, d_block_wall.nx, gx, gy);
-                const unsigned int gx = (blockIdx.z % nbx) * bx + x - halo;
-                const unsigned int gy = (blockIdx.z / nbx) * by + y - halo;
                 if (gy == 1)
                     reconstruct_bottom_velocity(f, 0.0f, bottom_wall_normal_velocity(
                         gx, d_block_wall.nx, d_block_wall.slot_width, d_block_wall.amplitude));
                 else if (gy == d_block_wall.ny - 2) reconstruct_top_reservoir(f, RHO);
             }
 #if COLLISION_MODEL == 2
-            // Uniform MRT sequence (collide_dispatch, zero forcing); SRT path below unchanged.
+            collide_mrt_unforced(f);  // SRT path below unchanged.
-            float rho, ux, uy;
-            compute_rho_u(f, rho, ux, uy);
-            collide_mrt_no_force(f, rho, ux, uy,
-                fminf(OMEGA_COLLISION_MAX, fmaxf(OMEGA_COLLISION_MIN, d_params.omega)));
 #else
             float rho, ux, uy, feq[NQ];
             compute_rho_u(f, rho, ux, uy);

src/CelerisLab/lbm/kernels/core/layout.cuh

@@ -53,6 +53,18 @@ __device__ __forceinline__ void index_from_thread(unsigned int& x,
     k = (unsigned long)y * NX + x;
 }
 
+// Global cell of window cell (x, y) in block blockIdx.z of a row-major,
+// full-domain block layout domain_nx cells wide. NX/NY are the local strides.
+__device__ __forceinline__ void block_global_xy(unsigned int x, unsigned int y,
+                                                unsigned int halo, unsigned int domain_nx,
+                                                unsigned int& gx, unsigned int& gy)
+{
+    const unsigned int bx = NX - 2 * halo, by = NY - 2 * halo;
+    const unsigned int nbx = domain_nx / bx;
+    gx = (blockIdx.z % nbx) * bx + x - halo;
+    gy = (blockIdx.z / nbx) * by + y - halo;
+}
+
 // ---------------------------------------------------------------------------
 // Neighbor list (periodic BCs via modular arithmetic)
 // ---------------------------------------------------------------------------

src/CelerisLab/lbm/kernels/operators/collision_mrt.cuh

@@ -114,6 +114,16 @@ __device__ __forceinline__ void collide_mrt_no_force(float* __restrict__ g,
     collide_mrt(g, rho, ux, uy, Fin, omega);
 }
 
+// Unforced MRT with the clamped device omega, as in collide_dispatch with zero
+// forcing. Shared by the blocked, channel-root and dense-patch collisions.
+__device__ __forceinline__ void collide_mrt_unforced(float* __restrict__ g)
+{
+    float rho, ux, uy;
+    compute_rho_u(g, rho, ux, uy);
+    collide_mrt_no_force(g, rho, ux, uy,
+        fminf(OMEGA_COLLISION_MAX, fmaxf(OMEGA_COLLISION_MIN, d_params.omega)));
+}
+
 #elif LATTICE_MODEL == LATTICE_D3Q19
 
 __device__ __forceinline__ void collide_mrt(float* __restrict__ g,

src/CelerisLab/lbm/kernels/refinement.cu

@@ -37,10 +37,7 @@ __device__ unsigned int incoming_counts[NQ * TOTAL_CELLS];
 
 __device__ __forceinline__ void patch_collide(float* f, unsigned long k) {
 #if COLLISION_MODEL == 2
-    float rho, ux, uy;
+    collide_mrt_unforced(f);
-    compute_rho_u(f, rho, ux, uy);
-    collide_mrt_no_force(f, rho, ux, uy,
-        fminf(OMEGA_COLLISION_MAX, fmaxf(OMEGA_COLLISION_MIN, d_params.omega)));
 #else
     float rho, ux, uy, feq[NQ];
     compute_rho_u(f, rho, ux, uy);

src/CelerisLab/lbm/kernels/step/block_channel.cu

@@ -15,10 +15,8 @@ extern "C" __global__ void BlockChannelStep(
     unsigned int x, y; unsigned long k;
     index_from_thread(x, y, k);
     if (x < halo || y < halo || x >= NX - halo || y >= NY - halo) return;
-    const unsigned int bx = NX - 2 * halo, by = NY - 2 * halo;
+    unsigned int gx, gy;
-    const unsigned int nbx = d_block_channel.nx / bx;
+    block_global_xy(x, y, halo, d_block_channel.nx, gx, gy);
-    const unsigned int gx = (blockIdx.z % nbx) * bx + x - halo;
-    const unsigned int gy = (blockIdx.z / nbx) * by + y - halo;
     const unsigned long base = (unsigned long)blockIdx.z * NQ * TOTAL_CELLS;
     const unsigned long cell = (unsigned long)blockIdx.z * TOTAL_CELLS + k;
     float f[NQ];
@@ -55,11 +53,7 @@ extern "C" __global__ void BlockChannelStep(
         }
         if (!wall && !outlet) {
 #if COLLISION_MODEL == 2
-            // Match uniform MRT arithmetic (collide_dispatch, zero forcing).
+            collide_mrt_unforced(f);
-            float rho, ux, uy;
-            compute_rho_u(f, rho, ux, uy);
-            collide_mrt_no_force(f, rho, ux, uy,
-                fminf(OMEGA_COLLISION_MAX, fmaxf(OMEGA_COLLISION_MIN, d_params.omega)));
 #else
             // Match uniform SRT arithmetic, including its shifted equilibrium path.
             float rho, ux, uy, feq[NQ], forcing[NQ] = {};

src/CelerisLab/lbm/refinement_stepper.py

@@ -75,13 +75,13 @@ def _check_levels(refinements, blocks, objects, jets):
             raise ValueError(f"{name} needs one entry per level ({depth})")
     for level, block in enumerate(blocks):
         if block is not None and len(block) != 2:
-            raise ValueError(f"{_level_name(level)}_block must be (by, bx)")
+            raise ValueError(f"level_blocks[{level}] must be (by, bx)")
     if objects[0] or jets[0]:
         raise ValueError("base-level bodies/Jets are not supported")
     if any(blocks[level - 1] is not None and blocks[level] is None for level in range(1, depth)):
         raise ValueError("a blocked level requires blocked refined levels below it")
     if not is_rectangle(refinements[0]) and (blocks[0] is None or blocks[1] is None):
-        raise ValueError("several patches require fine_block and base_block")
+        raise ValueError("several L1 patches require level_blocks[0] and level_blocks[1]")
     if depth > 2:
         if any(block is None for block in blocks):
             raise ValueError("three or more levels must all be blocked")
@@ -149,9 +149,11 @@ class FixedPatchStepper:
     Nested levels: a child set's coordinates index the allocated parent array,
     including ghost bands and holes: root edge = parent fine_origin + index / 2
     (zero origin for a parent set; a scalar parent keeps its shifted window
-    origin). A scalar child requires a scalar parent. Each pair requires gap >= 4
+    origin). A scalar child requires a scalar parent. Rectangle unions impose no
-    in its own parent-cell units. All regions share one buffer pair per level
+    fixed gap rule: adjacent, overlapping and concave unions are accepted when
-    and the fixed schedule of 2**l hops per base step at level l. Bodies/Jets are supported only on the
+    their actual transfer dependencies close (FixedPatchSet2D /
+    nest_patch_set). All regions share one buffer pair per level and the fixed
+    schedule of 2**l hops per base step at level l. Bodies/Jets are supported only on the
     finest level of a nested hierarchy (L2 dx=dt=.25 in root units, unchanged
     lattice velocity); every whole circle and its BFL stencil must fit strictly
     away from that level's transfer footprint.

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,
+        if jets:
-                            jet_link_shapes(jets, boundary.curved_hits) if jets else None)
+            jet_cache = (shape, jet_id)
+    return _BlockedBoundary(boundary, tables, jets, actuated, jet_cache)

tests/unit/test_block_step_sync.py

@@ -1,6 +1,7 @@
 """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
@@ -8,14 +9,15 @@ import unittest
 
 import numpy as np
 
-import CelerisLab.lbm.block_stepper as b
 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."""
-    path = Path(b.__file__)
+    # 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)

tests/unit/test_refinement_mrt.py

@@ -37,8 +37,13 @@ class RefinementMRTHost(unittest.TestCase):
             self.assertIn("STORE_PRECISION != 0", source)
             self.assertIn('#if COLLISION_MODEL == 2\n#include "operators/collision_mrt.cuh"\n#endif', source)
         self.assertIn("USE_DDF_SHIFTING ||", refinement)  # dense/mixed shifted stays rejected
+        # One shared unforced MRT body; each blocked/patch path calls it exactly once.
+        mrt = (KERNELS / "operators/collision_mrt.cuh").read_text()
+        self.assertEqual(mrt.count("void collide_mrt_unforced("), 1)
         for source in (blocks, channel, refinement):
-            self.assertEqual(source.count("collide_mrt_no_force("), 1)
+            self.assertEqual(source.count("collide_mrt_unforced(f);"), 1)
+            self.assertNotIn("collide_mrt_no_force(", source)
+            self.assertNotIn("compute_rho_u(f, rho, ux, uy);\n            collide_mrt", source)
 
 
 if __name__ == "__main__": unittest.main()