Jet review fixes: frozen Jet, validated amplitude replacement, test selector

Jet is a frozen dataclass; with_amplitude returns a validated copy, so
region/field/profile cannot drift from the cached per-link shapes and
amplitude cannot bypass validation. CurvedLinkSoA.set_jet_amplitude and
Simulation.set_jet_amplitude replace the Jet (Simulation also updates the
registered list so later geometry syncs keep the amplitude). Reject
non-finite hits; document edge rule (|s|<1 exclusive, radial/along
bounds inclusive). Fix the clipped-q test selector to receiver (9,5)
dir 8, as in test_surface_actuation.

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

src/CelerisLab/body/jets.py

@@ -7,12 +7,16 @@ Per-link velocity = amplitude * profile(s) * field(hit), where ``s`` in [-1, 1]
 is the region's lateral coordinate. Velocities are lattice units of the level
 whose links are evaluated; callers convert frames/units before evaluation.
 One physical link (receiver, direction) may be driven by at most one jet.
+
+Edge rule: the lateral bound |s| < 1 is exclusive (the profile vanishes
+there for parabolic/quartic); Sector r_min/r_max and Tube 0 <= along <= length
+are inclusive. Jets are immutable; an amplitude change replaces the Jet.
 """
 
 from __future__ import annotations
 
 import math
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 
 import numpy as np
 
@@ -100,11 +104,12 @@ class Parallel:
         return np.broadcast_to([math.cos(self.angle), math.sin(self.angle)], hits.shape)
 
 
-@dataclass
+@dataclass(frozen=True)
 class Jet:
     """Region (Sector|Tube), field (Radial|Parallel), profile name, signed amplitude.
 
-    Only ``amplitude`` may change after registration; it never rebuilds links.
+    Frozen: change amplitude with ``with_amplitude`` (validated, new object);
+    the cached per-link shapes stay valid because region/field/profile cannot change.
     """
     region: object
     field: object
@@ -118,6 +123,9 @@ class Jet:
         if abs(self.amplitude) >= 1:
             raise ValueError("Jet amplitude must satisfy abs(amplitude) < 1 (lattice units)")
 
+    def with_amplitude(self, amplitude):
+        return replace(self, amplitude=float(amplitude))
+
 
 def jet_link_shapes(jets, hits):
     """Unit-amplitude per-link vectors (n, 2) FP64 and jet index (n,) int32 (-1 = none).
@@ -126,6 +134,8 @@ def jet_link_shapes(jets, hits):
     links. Raises if a jet selects nothing or two jets select the same link.
     """
     hits = np.asarray(hits, np.float64).reshape(-1, 2)
+    if not np.isfinite(hits).all():
+        raise ValueError("Jet hit points must be finite")
     shape = np.zeros_like(hits)
     jet_id = np.full(len(hits), -1, np.int32)
     for j, jet in enumerate(jets):

src/CelerisLab/lbm/curved_links.py

@@ -149,10 +149,10 @@ class CurvedLinkSoA:
 
     def set_jet_amplitude(self, index, amplitude, stream=None):
         """Rewrite jet rows from cached shapes; re-upload only the velocity column."""
-        from ..body.jets import Jet, jet_link_velocity
-        jet = self.jets[index]
-        Jet(jet.region, jet.field, jet.profile, amplitude)  # validate before mutation
-        jet.amplitude = float(amplitude)
+        from ..body.jets import jet_link_velocity
+        jets = list(self.jets)
+        jets[index] = jets[index].with_amplitude(amplitude)  # validates before any change
+        self.jets = tuple(jets)
         driven = self.jet_id >= 0
         self.surface_velocity[driven] = jet_link_velocity(
             self.jets, self.jet_shape, self.jet_id)[driven]

src/CelerisLab/simulation.py

@@ -174,13 +174,10 @@ class Simulation:
     @context_bound
     def set_jet_amplitude(self, index: int, amplitude: float) -> None:
         """Change one jet's signed amplitude; re-uploads per-link velocity only."""
+        jet = self.bodies.jets[index].with_amplitude(amplitude)
         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)
+        self.bodies.jets[index] = jet  # later geometry syncs keep the new amplitude
 
     def remove_body(self, id: int) -> None:
         """Remove a body from the simulation.

tests/unit/test_jets.py

@@ -74,8 +74,8 @@ def test_selection_uses_true_hit_not_clipped_q():
     t = (6-math.sqrt(36-8*(9-2.999**2)))/4
     angle = math.atan2(t, 3-t)
     _, cl, _, _, hits = scene()
-    row = np.flatnonzero((cl.fluid_idx == 5*12 + 9) & np.isclose(cl.q, .01))
-    assert len(row) == 1
+    row = np.flatnonzero((cl.fluid_idx == 5*12 + 9) & (cl.dir == 8))
+    assert len(row) == 1 and np.isclose(cl.q[row[0]], .01)
     clipped = np.array([9., 5.]) + .01*np.array([-1., 1.])
     jet = Jet(Sector(CENTER, angle, 1e-4, RADIUS + .1), Radial(CENTER))
     _, jet_id = jet_link_shapes([jet], hits)
@@ -120,7 +120,22 @@ def test_amplitude_update_reuses_cached_shapes():
     assert np.array_equal(cl.aperture, ids >= 0)
     with pytest.raises(ValueError):
         cl.set_jet_amplitude(0, 1.5)
-    assert jet.amplitude == -2e-3
+    assert cl.jets[0].amplitude == -2e-3 and jet.amplitude == 1e-3  # registered Jet untouched
+    np.testing.assert_allclose(cl.surface_velocity, -2*before, rtol=1e-6, atol=0)
+
+
+def test_jet_definition_is_immutable_and_amplitude_validated():
+    jet = Jet(Sector(CENTER, 0., .5, 4.), Radial(CENTER), "parabolic", 1e-3)
+    for name, value in (("field", Parallel(0.)), ("region", Tube(CENTER, 0., 1., 1.)),
+                        ("profile", "uniform"), ("amplitude", float("nan"))):
+        with pytest.raises(AttributeError):
+            setattr(jet, name, value)
+    for bad in (float("nan"), float("inf"), 1.0):
+        with pytest.raises(ValueError):
+            jet.with_amplitude(bad)
+    assert jet.with_amplitude(-1e-3).field is jet.field and jet.amplitude == 1e-3
+    with pytest.raises(ValueError, match="finite"):
+        jet_link_shapes([jet], [[np.nan, 5.]])
 
 
 @pytest.mark.parametrize("spec", [SurfaceActuation(0., .6, .002), SurfaceActuation(2., 1.1, -.003)])