diff --git a/src/CelerisLab/body/manager.py b/src/CelerisLab/body/manager.py index 46a5592..a97ec75 100644 --- a/src/CelerisLab/body/manager.py +++ b/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) diff --git a/src/CelerisLab/body/objects.py b/src/CelerisLab/body/objects.py index c3fb101..19bb74c 100644 --- a/src/CelerisLab/body/objects.py +++ b/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): diff --git a/src/CelerisLab/body/sync_plan.py b/src/CelerisLab/body/sync_plan.py index 546cd15..d6b0dc2 100644 --- a/src/CelerisLab/body/sync_plan.py +++ b/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 diff --git a/src/CelerisLab/lbm/curved_links.py b/src/CelerisLab/lbm/curved_links.py index 3fdf0e4..3ddaea2 100644 --- a/src/CelerisLab/lbm/curved_links.py +++ b/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] diff --git a/src/CelerisLab/simulation.py b/src/CelerisLab/simulation.py index 64cf484..fb92e10 100644 --- a/src/CelerisLab/simulation.py +++ b/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) diff --git a/tests/unit/test_jets.py b/tests/unit/test_jets.py index 4ddde23..1f66771 100644 --- a/tests/unit/test_jets.py +++ b/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)