Wire Jets through ObjectManager and Simulation

ObjectManager.jets holds scene-level Jets. build_compact_lists_for
collects true FP64 hits (SimObject.get_curved_list hits_out) row-aligned
with the cut-link columns; BodySyncPlan carries them as curved_hits, and
sync_to_gpu / apply_sync_plan / the pre-sync check pass jets+hits to
assign_surface. Jets require D2Q9/FP32/unshifted/double_buffer and may
only drive links of stationary bodies. Simulation.add_jet (before
initialize) and set_jet_amplitude (velocity column re-upload only);
save_checkpoint refuses scenes with jets. Fine-path wiring not included.

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

src/CelerisLab/body/manager.py

@@ -70,6 +70,7 @@ class ObjectManager:
         self.nq = nq
         self.cfg = cfg
         self._registry = BodyRegistry()
+        self.jets: list = []  # scene-level Jets; selection is spatial, not per body
 
         self.action_gpu: Optional[cuda.DeviceAllocation] = None
         self.obs_gpu: Optional[cuda.DeviceAllocation] = None
@@ -247,6 +248,7 @@ class ObjectManager:
 
         # Build compact lists for new object set
         new_sensor_counts = np.zeros(max(next_count, 1), dtype=np.int32)
+        hits = []
         (
             cl_fluid_idx, cl_dir, cl_q, cl_body_id,
             cl_rx, cl_ry, cl_rz, cl_fallback_class,
@@ -254,7 +256,7 @@ class ObjectManager:
             fr_cells, fr_obj_id,
         ) = self.build_compact_lists_for(
             next_objects, domain_flags=next_flags,
-            sensor_cell_counts_out=new_sensor_counts,
+            sensor_cell_counts_out=new_sensor_counts, hits_out=hits,
         )
 
         # Compute change masks
@@ -277,6 +279,7 @@ class ObjectManager:
             added_solid_mask=added_solid_mask,
             released_fluid_mask=released_fluid_mask,
             new_sensor_counts=new_sensor_counts,
+            curved_hits=hits[0],
         )
 
     def commit_pending(self, next_objects: List[SimObject],
@@ -309,7 +312,8 @@ class ObjectManager:
 
         cl_host = plan.curved_host
         field.curved.assign_host(*cl_host)
-        field.curved.assign_surface(plan.next_objects, self.cfg, plan.next_flags)
+        field.curved.assign_surface(plan.next_objects, self.cfg, plan.next_flags,
+                                    jets=self.jets, hits=plan.curved_hits)
         s_cells, s_ids = plan.sensor_host
         field.sensors.assign_host(s_cells, s_ids)
         fr_cells, fr_ids = plan.force_region_host
@@ -344,13 +348,13 @@ class ObjectManager:
         return merge_flag_masks(base_flags, masks)
 
     # -- Compact list building -----------------------------------------------
-    def build_compact_lists(self, domain_flags: np.ndarray | None = None):
+    def build_compact_lists(self, domain_flags: np.ndarray | None = None, *, hits_out=None):
         """Build cut-link SoA columns, sensor lists, and force-region lists
         from the formal registry objects.
         """
         return self.build_compact_lists_for(
             self._registry.objects, domain_flags,
-            sensor_cell_counts_out=self.sensor_cell_counts,
+            sensor_cell_counts_out=self.sensor_cell_counts, hits_out=hits_out,
         )
 
     def build_compact_lists_for(
@@ -359,6 +363,7 @@ class ObjectManager:
         domain_flags: np.ndarray | None = None,
         *,
         sensor_cell_counts_out: np.ndarray | None = None,
+        hits_out: list | None = None,
     ):
         """Build cut-link SoA columns, sensor lists, and force-region lists.
 
@@ -367,6 +372,8 @@ class ObjectManager:
             domain_flags: Optional uint16 flags for donor-cell classification.
             sensor_cell_counts_out: Int array to fill with per-object sensor
                 cell counts.  If provided, ``fill(0)`` is called first.
+            hits_out: List that receives one FP64 (n, 2) array of true hit
+                points row-aligned with the cut-link columns (for Jets).
 
         Returns:
             cl_fluid_idx, cl_dir, cl_q, cl_body_id, cl_rx, cl_ry, cl_rz,
@@ -376,6 +383,7 @@ class ObjectManager:
         """
         cl_idx, cl_dir, cl_q, cl_bid = [], [], [], []
         cl_rx, cl_ry, cl_rz, cl_fallback = [], [], [], []
+        cl_hit = []
         s_cells, s_ids = [], []
         fr_cells, fr_ids = [], []
         if domain_flags is None:
@@ -404,7 +412,7 @@ class ObjectManager:
                 ) = obj.get_curved_list(
                     self.nx, self.ny, self.nq,
                     nz=self.nz, ex=ex, ey=ey, ez=ez,
-                    domain_flags=domain_flags,
+                    domain_flags=domain_flags, hits_out=cl_hit,
                 )
                 if len(fluid_idx) > 0:
                     cl_idx.append(fluid_idx)
@@ -452,12 +460,15 @@ class ObjectManager:
         r_fr_cells = np.concatenate(fr_cells) if fr_cells else np.zeros(0, dtype=np.uint32)
         r_fr_ids = np.concatenate(fr_ids) if fr_ids else np.zeros(0, dtype=np.int32)
 
-        if any(obj.surface_actuation is not None for obj in objects):
+        r_hit = np.concatenate(cl_hit) if cl_hit else np.zeros((0, 2))
+        if hits_out is not None:
+            hits_out.append(r_hit)
+        if self.jets or any(obj.surface_actuation is not None for obj in objects):
             # Reject unresolved apertures/ownership before sync or restart mutates state.
             from ..lbm.curved_links import CurvedLinkSoA
             links = CurvedLinkSoA()
             links.assign_host(r_idx, r_dir, r_q, r_bid, r_rx, r_ry, r_rz, r_fallback)
-            links.assign_surface(objects, self.cfg, domain_flags)
+            links.assign_surface(objects, self.cfg, domain_flags, jets=self.jets, hits=r_hit)
 
         return (
             r_idx, r_dir, r_q, r_bid,
@@ -479,13 +490,14 @@ class ObjectManager:
         self._telemetry_field = field
 
         next_flags = self.build_flags(field.build_channel_flags()) if rebuild_flags else field.flag
+        hits = []
 
         (
             cl_fluid_idx, cl_dir, cl_q, cl_body_id,
             cl_rx, cl_ry, cl_rz, cl_fallback_class,
             sensor_cells, sensor_obj_id,
             fr_cells, fr_obj_id,
-        ) = self.build_compact_lists(domain_flags=next_flags)
+        ) = self.build_compact_lists(domain_flags=next_flags, hits_out=hits)
 
         if rebuild_flags:
             field.flag = next_flags
@@ -495,7 +507,8 @@ class ObjectManager:
             cl_fluid_idx, cl_dir, cl_q, cl_body_id,
             cl_rx, cl_ry, cl_rz, cl_fallback_class,
         )
-        field.curved.assign_surface(self.objects, self.cfg, field.flag)
+        field.curved.assign_surface(self.objects, self.cfg, field.flag,
+                                    jets=self.jets, hits=hits[0])
         field.sensors.assign_host(sensor_cells, sensor_obj_id)
         field.force_regions.assign_host(fr_cells, fr_obj_id)
 

src/CelerisLab/body/objects.py

@@ -195,12 +195,13 @@ class SimObject:
     def get_curved_list(
         self, nx: int, ny: int, nq: int,
         nz: int = 1, ex=None, ey=None, ez=None,
-        domain_flags=None,
+        domain_flags=None, hits_out=None,
     ):
         """Return 8-column SoA arrays for curved-boundary kernel.
 
         Delegates to the geometry object for cut-link enumeration and
-        to ``pack_cut_links_to_soa`` for SoA packing.
+        to ``pack_cut_links_to_soa`` for SoA packing. If ``hits_out`` is a
+        list, the true FP64 hit points (n, 2) are appended row-aligned.
 
         Returns:
             8-tuple matching CurvedLinkSoA.assign_host().
@@ -210,6 +211,9 @@ class SimObject:
         links = self.geometry.build_cut_links(nx, ny, domain_flags=domain_flags)
         for cl in links:
             cl.body_id = self.obj_id
+        if hits_out is not None:
+            hits_out.append(np.array([(cl.hit_x, cl.hit_y) for cl in links],
+                                     np.float64).reshape(-1, 2))
         return pack_cut_links_to_soa(links)
 
     def get_sensor_list(self, nx: int, ny: int, nz: int = 1):

src/CelerisLab/body/sync_plan.py

@@ -40,6 +40,7 @@ class BodySyncPlan:
         added_solid_mask: Bool array -- True where ``old_fluid -> new_solid``.
         released_fluid_mask: Bool array -- True where ``old_obstacle -> new_fluid``.
         new_sensor_counts: Int array of sensor cell counts per object.
+        curved_hits: FP64 (n, 2) true hit points row-aligned with curved_host.
     """
     next_objects: List[SimObject]
     next_count: int
@@ -50,3 +51,4 @@ class BodySyncPlan:
     added_solid_mask: np.ndarray
     released_fluid_mask: np.ndarray
     new_sensor_counts: np.ndarray
+    curved_hits: np.ndarray | None = None

src/CelerisLab/lbm/curved_links.py

@@ -73,8 +73,16 @@ def surface_link_velocity(objects, cfg, flags, fluid_idx, dir, body_id, rx, ry,
         from ..body.jets import jet_link_shapes, jet_link_velocity
         if hits is None or len(hits) != len(dir):
             raise ValueError("Jets require true hit points row-aligned with the cut links")
+        if (cfg is None or not cfg.is_d2q9 or cfg.store_precision != "FP32"
+                or cfg.streaming != "double_buffer" or cfg.ddf_shifting):
+            raise ValueError("Jets require D2Q9/FP32/unshifted/double_buffer")
         shape, jet_id = jet_link_shapes(jets, hits)
         driven = jet_id >= 0
+        owners = {obj.obj_id: obj for obj in objects}
+        for bid in np.unique(body_id[driven]):
+            state = vars(owners[int(bid)].state)
+            if any(state.get(k, 0.) != 0 for k in ("vx", "vy", "omega")):
+                raise ValueError("Jets may only drive links of stationary bodies")
         if np.any(surface_aperture & driven):
             raise ValueError("Jet overlaps a SurfaceActuation aperture")
         surface_velocity[driven] = jet_link_velocity(jets, shape, jet_id)[driven]

src/CelerisLab/simulation.py

@@ -157,6 +157,31 @@ class Simulation:
             return -1
         return self.bodies.add(obj)
 
+    def add_jet(self, jet) -> int:
+        """Register a body-independent Jet before ``initialize()``; returns its index.
+
+        The jet drives every scene cut link whose true hit point lies in its
+        region (overlaps rejected at initialize). Not checkpointed.
+        """
+        from .body.jets import Jet
+        if not isinstance(jet, Jet):
+            raise TypeError("add_jet expects a Jet")
+        if self._initialized:
+            raise RuntimeError("Jets must be added before initialize()")
+        self.bodies.jets.append(jet)
+        return len(self.bodies.jets) - 1
+
+    @context_bound
+    def set_jet_amplitude(self, index: int, amplitude: float) -> None:
+        """Change one jet's signed amplitude; re-uploads per-link velocity only."""
+        if self._initialized:
+            self.field.curved.set_jet_amplitude(index, amplitude, self.stream)
+            return
+        from .body.jets import Jet
+        jet = self.bodies.jets[index]
+        Jet(jet.region, jet.field, jet.profile, amplitude)  # validate before mutation
+        jet.amplitude = float(amplitude)
+
     def remove_body(self, id: int) -> None:
         """Remove a body from the simulation.
 
@@ -433,6 +458,8 @@ class Simulation:
             self.field.flag = self.bodies.build_flags(self.field.flag)
         self.field.upload_flags()
 
+        if self.bodies.jets and self.bodies.count == 0:
+            raise ValueError("Jets need bodies with cut links")
         self.stepper.initialize()
         if self.bodies.count > 0:
             self.bodies.sync_to_gpu(self.field, rebuild_flags=False)
@@ -644,6 +671,8 @@ class Simulation:
     def save_checkpoint(self, path: str = None):
         """Save full simulation state to an HDF5 file."""
         from .common.checkpoint import save_checkpoint
+        if self.bodies.jets:
+            raise NotImplementedError("Jets are not checkpointed yet")
         return save_checkpoint(self.field, self.stepper, self.lbm_cfg,
                                self.bodies, path)
 

tests/unit/test_jets.py

@@ -1,5 +1,6 @@
 """Body-independent Jet: true-hit spatial selection, fields, profiles, overlap, amplitude."""
 import math
+from types import SimpleNamespace
 
 import numpy as np
 import pytest
@@ -8,6 +9,7 @@ from CelerisLab.body.jets import (Jet, Parallel, Radial, Sector, Tube,
                                   fine_hits_to_physical, jet_link_shapes)
 from CelerisLab.body.objects import SurfaceActuation
 from CelerisLab.lbm.curved_links import CurvedLinkSoA, surface_link_velocity
+from .test_state_regressions import device, make_field, manager
 from .test_surface_actuation import circle, projected
 
 CENTER, RADIUS = (6., 5.), 2.999
@@ -132,3 +134,46 @@ def test_legacy_adapter_measured_equivalence(spec):
     diff = np.abs(cl.surface_velocity - legacy.surface_velocity).max()
     print(f"legacy adapter max |dU| = {diff:.3e} (exact={diff == 0})")
     assert diff <= 4*np.finfo(np.float32).eps*abs(spec.normal_velocity)
+
+
+def jet_case(device, jets):
+    field = make_field(device)
+    bodies = manager(field, [circle(radius=2.)])
+    bodies.jets.extend(jets)
+    bodies.sync_to_gpu(field)
+    return field, bodies
+
+
+def test_manager_passes_true_hits_and_reuploads_only_velocity(device):
+    # CPU memory double only: checks host/device bookkeeping, not GPU numerics.
+    jet = Jet(Sector(CENTER, 0., .6, 3.), Radial(CENTER, inward=False), "quartic", 2e-3)
+    field, bodies = jet_case(device, [jet])
+    cl = field.curved
+    assert cl.jet_id is not None and (cl.jet_id == 0).sum() >= 2
+    velocity_gpu = cl.surface_velocity_gpu
+    np.testing.assert_array_equal(np.frombuffer(velocity_gpu.data, np.float32).reshape(-1, 2),
+                                  cl.surface_velocity)
+    geometry = {name: bytes(getattr(cl, name + "_gpu").data)
+                for name in ("fluid_idx", "dir", "q", "body_id", "rx", "ry")}
+    writes, before = device.writes, cl.surface_velocity.copy()
+    cl.set_jet_amplitude(0, -1e-3, SimpleNamespace(synchronize=lambda: None))
+    assert device.writes == writes + 1 and cl.surface_velocity_gpu is velocity_gpu
+    np.testing.assert_array_equal(np.frombuffer(velocity_gpu.data, np.float32).reshape(-1, 2),
+                                  cl.surface_velocity)
+    assert all(bytes(getattr(cl, name + "_gpu").data) == data for name, data in geometry.items())
+    np.testing.assert_allclose(cl.surface_velocity, -.5*before, rtol=1e-6, atol=0)
+    np.testing.assert_array_equal(cl.surface_velocity[cl.jet_id < 0], 0.)
+
+
+def test_manager_rejects_overlapping_jets_and_moving_bodies(device):
+    a = Jet(Sector(CENTER, 0., .6, 3.), Radial(CENTER))
+    b = Jet(Tube(CENTER, 0., 3., .5), Parallel(0.))
+    with pytest.raises(ValueError, match="overlaps"):
+        jet_case(device, [a, b])
+    field = make_field(device)
+    moving = circle(radius=2.)
+    moving.state.omega = 1e-3
+    bodies = manager(field, [moving])
+    bodies.jets.append(a)
+    with pytest.raises(ValueError, match="stationary"):
+        bodies.sync_to_gpu(field)