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/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

@@ -178,7 +178,11 @@ extern "C" __global__ void BlockStep(const fpxx* in, fpxx* out, unsigned int hal
                 else if (gy == d_block_wall.ny - 2) reconstruct_top_reservoir(f, RHO);
             }
 #if COLLISION_MODEL == 2
-            collide_mrt_unforced(f);  // SRT path below unchanged.
+            // Uniform MRT sequence (collide_dispatch, zero forcing); 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/operators/collision_mrt.cuh

@@ -114,16 +114,6 @@ __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,7 +37,10 @@ __device__ unsigned int incoming_counts[NQ * TOTAL_CELLS];
 
 __device__ __forceinline__ void patch_collide(float* f, unsigned long k) {
 #if COLLISION_MODEL == 2
-    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
     float rho, ux, uy, feq[NQ];
     compute_rho_u(f, rho, ux, uy);

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

@@ -55,7 +55,11 @@ extern "C" __global__ void BlockChannelStep(
         }
         if (!wall && !outlet) {
 #if COLLISION_MODEL == 2
-            collide_mrt_unforced(f);
+            // Match uniform MRT arithmetic (collide_dispatch, zero forcing).
+            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_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_refinement_mrt.py

@@ -37,13 +37,8 @@ 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_unforced(f);"), 1)
-            self.assertNotIn("collide_mrt_no_force(", source)
-            self.assertNotIn("compute_rho_u(f, rho, ux, uy);\n            collide_mrt", source)
+            self.assertEqual(source.count("collide_mrt_no_force("), 1)
 
 
 if __name__ == "__main__": unittest.main()