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/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_annulus_ledger.py

@@ -1,101 +0,0 @@
-"""Synthetic checks of the annulus mass ledger. Host-only, stdlib; not solver evidence."""
-
-import random
-import sys
-from pathlib import Path
-import unittest
-
-sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "validation"))
-import annulus_ledger as al  # noqa: E402
-
-NX = NY = 24
-# Annulus 4 <= r <= 8 around (12, 12); dyadic values keep every sum exact.
-A = {(x, y) for x in range(NX) for y in range(NY) if 16 <= (x - 12) ** 2 + (y - 12) ** 2 <= 64}
-DELTA = 0.375
-
-
-def owner(seed=1):
-    rng = random.Random(seed)
-    return [[[rng.randint(1, 1024) / 1024 for _ in range(NX)] for _ in range(NY)] for _ in range(9)]
-
-
-def collide(P, source=None):
-    """Mass-conserving stand-in: rotate populations within each cell (+ optional source)."""
-    D = [[[P[(q + 1) % 9][y][x] for x in range(NX)] for y in range(NY)] for q in range(9)]
-    if source:
-        (x, y), s = source
-        D[0][y][x] += s
-    return D
-
-
-def link(interior):
-    """(x, y, q) with x in A and x - c_q inside (seam) or outside (edge) A."""
-    for (x, y) in sorted(A):
-        for q in range(1, 9):
-            src = ((x - al.CX[q]) % NX, (y - al.CY[q]) % NY)
-            if (src in A) == interior:
-                return x, y, q
-    raise AssertionError
-
-
-class Ledger(unittest.TestCase):
-    def setUp(self):
-        self.O = owner()
-        self.P = al.pull(self.O, NX, NY)
-
-    def run_ledger(self, P, D):
-        return al.ledger(self.O, P, D, A, NX, NY)
-
-    def test_ideal_pull_has_no_h_and_closes(self):
-        r = self.run_ledger(self.P, collide(self.P))
-        self.assertEqual((r["H_seam"], r["H_edge"], r["C"], r["residual"]), (0, 0, 0, 0))
-        self.assertEqual(r["dM"], -r["Q_exact"])
-        self.assertNotEqual(r["Q_exact"], 0)
-
-    def test_collision_source_is_c(self):
-        x, y = sorted(A)[3]
-        r = self.run_ledger(self.P, collide(self.P, ((x, y), DELTA)))
-        self.assertEqual((r["C"], r["H"], r["residual"]), (DELTA, 0, 0))
-
-    def test_interior_seam_perturbation(self):
-        x, y, q = link(interior=True)
-        self.P[q][y][x] += DELTA
-        r = self.run_ledger(self.P, collide(self.P))
-        self.assertEqual((r["H_seam"], r["H_edge"], r["residual"]), (DELTA, 0, 0))
-        # Invisible to any flux through the edge of A; only H explains the mass change.
-        self.assertEqual(al.staged_flux(self.O, self.P, A, NX, NY), r["Q_exact"])
-        self.assertEqual(r["dM"] + r["Q_exact"], DELTA)
-
-    def test_edge_perturbation_and_staged_flux_double_count(self):
-        x, y, q = link(interior=False)
-        self.P[q][y][x] += DELTA
-        r = self.run_ledger(self.P, collide(self.P))
-        self.assertEqual((r["H_seam"], r["H_edge"], r["residual"]), (0, DELTA, 0))
-        staged = al.staged_flux(self.O, self.P, A, NX, NY)
-        self.assertEqual(staged, r["Q_exact"] - DELTA)
-        # Staged flux + full H counts the staging change twice.
-        self.assertEqual(r["dM"] + staged - r["C"] - r["H"], -DELTA)
-        # Staged flux alone closes only without the edge part of H.
-        self.assertEqual(r["dM"] + staged - r["C"] - r["H_seam"], 0)
-
-    def test_owner_snapshot_after_staging_hides_the_change(self):
-        # Snapshot taken after staging wrote the same slot into the owner image.
-        x, y, q = link(interior=False)
-        sx, sy = (x - al.CX[q]) % NX, (y - al.CY[q]) % NY
-        self.P[q][y][x] += DELTA
-        late = [[row[:] for row in plane] for plane in self.O]
-        late[q][sy][sx] += DELTA
-        r = al.ledger(late, self.P, collide(self.P), A, NX, NY)
-        self.assertEqual((r["H"], r["residual"]), (0, 0))
-
-    def test_wrong_pull_direction_is_flagged(self):
-        # Indexing error (push instead of pull): H is nonzero with no physical cause.
-        bad = [[[self.O[q][(y + al.CY[q]) % NY][(x + al.CX[q]) % NX] for x in range(NX)]
-                for y in range(NY)] for q in range(9)]
-        r = self.run_ledger(bad, collide(bad))
-        self.assertNotEqual(r["H"], 0)
-        self.assertEqual(r["residual"], 0)
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/validation/annulus_ledger.py

@@ -1,75 +0,0 @@
-"""Per-hop mass ledger of a fluid cell set A (host-only, stdlib, D2Q9 pull scheme).
-
-Diagnostic definition only; no solver code, no GPU. Arrays are indexed [q][y][x] on a
-periodic nx x ny lattice; A is a set of (x, y) ACTIVE cells.
-
-  O  owner populations before any staging (canonical owner slots, start of the hop)
-  P  populations each cell actually pulls: P[q](x) is what x reads for direction q
-     (after halo copy, BFL/coalesce staging, ...)
-  D  populations written by the cell after collision
-
-With M0 = sum_{A,q} O and M1 = sum_{A,q} D (both owner snapshots of A):
-
-  M1 - M0 + Q_exact = C + H                                            (exact)
-
-  C       = sum_{x in A, q} (D_q(x) - P_q(x))              collision source
-  H       = sum_{x in A, q} (P_q(x) - O_q(x - c_q))        actual minus ideal owner input,
-            split into H_seam (x - c_q in A: internal pulls, e.g. tile seams) and
-            H_edge (x - c_q not in A: inputs crossing the edge of A)
-  Q_exact = sum_{y in A, y+c_q not in A} O_q(y) - sum_{y not in A, y+c_q in A} O_q(y)
-            owner-based net outward flux (outgoing minus incoming)
-
-Proof: sum_{A,q} O_q(x - c_q) - M0 = -Q_exact (relabel y = x - c_q per direction), and
-sum_{A,q} D = C + H + sum_{A,q} O_q(x - c_q). Q_exact must use O, not staged values: a
-flux built from staged inputs already contains H_edge, so adding the full H counts it
-twice (see staged_flux). The edge of A must not cross cut links, otherwise O at the solid
-end enters Q_exact.
-"""
-
-CX = (0, 1, -1, 0, 0, 1, -1, 1, -1)
-CY = (0, 0, 0, 1, -1, 1, -1, -1, 1)
-
-
-def _src(x, y, q, nx, ny):
-    return (x - CX[q]) % nx, (y - CY[q]) % ny
-
-
-def pull(O, nx, ny):
-    """Ideal pulled input: P[q](x) = O[q](x - c_q)."""
-    return [[[O[q][_src(x, y, q, nx, ny)[1]][_src(x, y, q, nx, ny)[0]]
-              for x in range(nx)] for y in range(ny)] for q in range(9)]
-
-
-def ledger(O, P, D, A, nx, ny):
-    A = set(A)
-    out = dict(dM=0.0, C=0.0, H_seam=0.0, H_edge=0.0, Q_exact=0.0)
-    for (x, y) in A:
-        for q in range(9):
-            out["dM"] += D[q][y][x] - O[q][y][x]
-            out["C"] += D[q][y][x] - P[q][y][x]
-            sx, sy = _src(x, y, q, nx, ny)
-            key = "H_seam" if (sx, sy) in A else "H_edge"
-            out[key] += P[q][y][x] - O[q][sy][sx]
-            if (sx, sy) not in A:
-                out["Q_exact"] -= O[q][sy][sx]          # incoming owner population
-            tx, ty = (x + CX[q]) % nx, (y + CY[q]) % ny
-            if (tx, ty) not in A:
-                out["Q_exact"] += O[q][y][x]            # outgoing owner population
-    out["H"] = out["H_seam"] + out["H_edge"]
-    out["residual"] = out["dM"] + out["Q_exact"] - out["C"] - out["H"]
-    return out
-
-
-def staged_flux(O, P, A, nx, ny):
-    """Outward flux with incoming terms from the staged inputs P (the wrong combination)."""
-    A = set(A)
-    q_out = 0.0
-    for (x, y) in A:
-        for q in range(9):
-            sx, sy = _src(x, y, q, nx, ny)
-            if (sx, sy) not in A:
-                q_out -= P[q][y][x]
-            tx, ty = (x + CX[q]) % nx, (y + CY[q]) % ny
-            if (tx, ty) not in A:
-                q_out += O[q][y][x]
-    return q_out