From 987566c0e6e2a5cc5ab8bce727771a5178e20d32 Mon Sep 17 00:00:00 2001 From: Frank14f <1515444314@qq.com> Date: Sat, 20 Jun 2026 18:17:07 +0800 Subject: [PATCH] feat(body): runtime body add/remove, unified action/obs, FRC_REGION flag - Add runtime body topology sync (add_body/remove_body + sync_bodies) with recompile, DDF patch (feq + BFS inward fill), and commit. - Unify action/obs flow: set_body/set_force are now host-only; run() auto-uploads action and downloads obs via CUDA stream. - Add read_body(id) -> BodyTelemetry and read_bodies() for DRL loops. - Add FRC_REGION flag (0x0800) for force_region cells. - Extract equilibrium helpers (lbm/equilibrium.py) and DDF patch module (body/ddf_patch.py). - Merge recompile / _runtime_recompile into single _recompile(). - Add n_objects to checkpoint; validate on load. - Add test suite: 40 unit + 19 integration tests (59 total). - Add conftest.py and docs/tests_overview.md for test documentation. - Update README.md and CONFIG.md for new API. Co-authored-by: Cursor --- README.md | 148 +++++++-- docs/tests_overview.md | 44 +++ src/CelerisLab/body/__init__.py | 2 + src/CelerisLab/body/ddf_patch.py | 239 +++++++++++++++ src/CelerisLab/body/geometry/base.py | 16 + src/CelerisLab/body/geometry/circle.py | 6 +- src/CelerisLab/body/manager.py | 289 ++++++++++++++++-- src/CelerisLab/body/objects.py | 6 +- src/CelerisLab/body/sync_plan.py | 52 ++++ src/CelerisLab/common/checkpoint.py | 14 + src/CelerisLab/configs/CONFIG.md | 15 +- src/CelerisLab/lbm/descriptors.py | 1 + src/CelerisLab/lbm/equilibrium.py | 60 ++++ src/CelerisLab/lbm/field.py | 5 + src/CelerisLab/lbm/kernels/core/flags.cuh | 4 +- src/CelerisLab/lbm/kernels/step/init_flow.cu | 2 + src/CelerisLab/simulation.py | 242 +++++++++++++-- tests/conftest.py | 9 + tests/integration/__init__.py | 1 + tests/integration/test_body_sync_e2e.py | 117 +++++++ tests/integration/test_ddf_patch.py | 145 +++++++++ .../integration/test_sync_bodies_skeleton.py | 166 ++++++++++ tests/integration/test_unified_obs.py | 89 ++++++ tests/unit/__init__.py | 1 + tests/unit/test_body_flags.py | 94 ++++++ tests/unit/test_equilibrium.py | 50 +++ tests/unit/test_pending_edits.py | 147 +++++++++ tests/unit/test_sync_plan.py | 234 ++++++++++++++ 28 files changed, 2112 insertions(+), 86 deletions(-) create mode 100644 docs/tests_overview.md create mode 100644 src/CelerisLab/body/ddf_patch.py create mode 100644 src/CelerisLab/body/sync_plan.py create mode 100644 src/CelerisLab/lbm/equilibrium.py create mode 100644 tests/conftest.py create mode 100644 tests/integration/__init__.py create mode 100644 tests/integration/test_body_sync_e2e.py create mode 100644 tests/integration/test_ddf_patch.py create mode 100644 tests/integration/test_sync_bodies_skeleton.py create mode 100644 tests/integration/test_unified_obs.py create mode 100644 tests/unit/__init__.py create mode 100644 tests/unit/test_body_flags.py create mode 100644 tests/unit/test_equilibrium.py create mode 100644 tests/unit/test_pending_edits.py create mode 100644 tests/unit/test_sync_plan.py diff --git a/README.md b/README.md index e8a8e75..fc7843f 100644 --- a/README.md +++ b/README.md @@ -35,37 +35,62 @@ force = sim.read_force(0) # [fx, fy] on body 0 sim.close() ``` -### Multi-body control loop +### DRL control loop ```python from CelerisLab import Simulation sim = Simulation() -# Three rotating cylinders -sim.add_body("circle", center=(1006, 150), radius=10) -sim.add_body("circle", center=(1015, 140), radius=10) -sim.add_body("circle", center=(1015, 160), radius=10) -# Downstream velocity sensor -sim.add_body("sensor", center=(1050, 150), radius=10) +sim.add_body("circle", center=(256, 128), radius=10) +sim.add_body("sensor", center=(300, 128), radius=10) sim.initialize() -for step in range(100): - # Set body rotation speeds (implicit GPU upload) - sim.set_body(0, omega=0.002) - sim.set_body(1, omega=-0.001) - sim.set_body(2, omega=0.001) +for episode in range(100): + # Step the simulation (auto uploads action, downloads obs) + sim.run(100) - # Advance 10 LBM steps - sim.run(10) + # Read individual body telemetry (primary API) + data = sim.read_body(0) + print(f"step={sim.stepper.step_count} " + f"force=({data.force[0]:.4f},{data.force[1]:.4f}) " + f"sensor=({data.sensor[0]:.4f},{data.sensor[1]:.4f})") - # Read telemetry - fx, fy = sim.read_force(0) - ux, uy = sim.read_sensor(3) - print(f"force=({fx:.4f},{fy:.4f}) sensor=({ux:.4f},{uy:.4f})") + # DRL policy inference (replace with your model) + action_omega = 0.001 * (0.5 - data.force[0]) + + # Set action (host-only, will be auto-uploaded next run) + sim.set_body(0, omega=action_omega) sim.close() ``` +### Async control (performance-oriented, custom stream) + +```python +import pycuda.driver as cuda + +stream = cuda.Stream() +sim.set_body(0, omega=0.002) # host-only +sim.run(100, stream=stream) # action uploaded, steps run, obs downloaded on stream +# stream is synced inside run() -- obs is ready +data = sim.read_body(0) +``` + +### Manual stream control (max overlap) + +```python +import pycuda.driver as cuda + +stream = cuda.Stream() +# Skip transfers for the first batch, just enqueue kernels +sim.run(100, stream=stream, upload_act=False, sync_obs=False) + +# ... other GPU work can overlap with the kernel launches ... + +# Later: sync and read +stream.synchronize() +obs = sim.read_bodies(stream=stream) # sync already done, just read pinned buffer + ## Installation ### Prerequisites @@ -105,10 +130,11 @@ sim = Simulation( | Method | Returns | Description | |--------|---------|-------------| -| `sim.add_body(type="circle", center=(x,y), radius=r)` | int body_id | Add a cylinder body | +| `sim.add_body(type="circle", center=(x,y), radius=r)` | int body_id | Add a cylinder body (primary API) | | `sim.add_body(type="sensor", center=(x,y), radius=r)` | int body_id | Add a velocity sensor | -| `sim.add_cylinder(center, radius)` | int body_id | Backward-compat alias | -| `sim.add_sensor(center, radius)` | int body_id | Backward-compat alias | +| `sim.add_body(type="force_region", center=(x,y), radius=r)` | int body_id | Add a force application region | +| `sim.add_cylinder(center, radius)` | int body_id | Convenience wrapper (deprecated) | +| `sim.add_sensor(center, radius)` | int body_id | Convenience wrapper (deprecated) | | `sim.add_object(obj)` | int body_id | Add pre-configured SimObject | Future geometry types (polygon, mesh) will use the same `add_body()` function with a different `type` parameter. @@ -118,12 +144,64 @@ Future geometry types (polygon, mesh) will use the same `add_body()` function wi | Method | Description | |--------|-------------| | `sim.initialize()` | Recompile if needed, flow field + sync objects to GPU | -| `sim.run(steps, checkpoint_interval=0)` | Run N LBM steps | -| `sim.set_body(id, omega=...)` | Set body rotation speed (implicit GPU upload, ~1 μs) | -| `sim.read_force(id)` -> ndarray | Force vector [fx, fy] (2D) | -| `sim.read_torque(id)` -> ndarray | Torque [tz] (2D) | -| `sim.read_sensor(id)` -> ndarray | Area-averaged velocity via GPU sensor kernel | -| `sim.set_force(id, fx=..., fy=...)` | Set force density on a force_region object (notice: see persistence note below) | +| `sim.run(steps, *, upload_act=True, sync_obs=True, stream=None)` | Run N LBM steps. See stream subsection below. | +| `sim.set_body(id, omega=...)` | Set body rotation speed (host array only, uploaded at next `run()`) | +| `sim.read_body(id)` -> BodyTelemetry | Unified telemetry: {force, torque, sensor} from pinned buffer | +| `sim.read_bodies()` -> ndarray | Flat array of all bodies' telemetry (batch DRL read) | +| `sim.read_force(id)` -> ndarray | Force vector [fx, fy] (backward-compat) | +| `sim.read_torque(id)` -> ndarray | Torque [tz] (backward-compat) | +| `sim.read_sensor(id)` -> ndarray | Area-averaged velocity (backward-compat) | +| `sim.set_force(id, fx=..., fy=...)` | Set force density on a force_region object | + +**Action/obs transfer model:** `set_body()` / `set_force()` are host-only — they modify +the host action array without triggering GPU upload. The GPU buffer is automatically +updated at the start of the next ``run()`` call when ``upload_act=True`` (the default). +Similarly, after the step group, telemetry is downloaded to a pinned host buffer when +``sync_obs=True``. Both transfers run on the same CUDA stream as the kernels, so +they overlap with computation when possible. + +``run()`` parameters: +- ``steps``: Number of LBM steps. +- ``upload_act`` (default True): Upload host action array to ``action_gpu`` before stepping. +- ``sync_obs`` (default True): Download ``obs_gpu`` to host pinned buffer after stepping. +- ``stream`` (default None): CUDA stream for all operations. ``None`` uses an internal stream. +- ``checkpoint_interval`` (default 0): If >0, save an HDF5 checkpoint every N steps. + +Use ``upload_act=False, sync_obs=False`` to skip all transfers and enqueue pure +kernel launches on a user-provided stream, then sync and read later. + +#### Runtime body topology sync + +| Method | Description | +|--------|-------------| +| `sim.remove_body(id)` | Stage a body for removal (committed at next `sync_bodies()`) | +| `sim.sync_bodies()` | Commit pending add/remove edits: recompile kernel, rebuild flags/compact lists, patch DDF, re-upload to GPU | + +`sync_bodies()` applies all staged body edits (added via `add_body()` and removed via `remove_body()`) to a running simulation without full reinitialization. The GPU flow field is preserved; only the body-related topology is rebuilt. + +**Limitations:** +- Requires `streaming: "double_buffer"` (esopull raises `NotImplementedError`) +- Abrupt body introduction causes a transient; force readback is finite but may take 50+ steps to settle +- Verified for `"circle"` type bodies; sensors and force_regions are also expected to work + (they produce no curved links so the DDF patch is simpler) + +```python +# Add a body to an already-initialized simulation +sim = Simulation() +sim.initialize() +sim.run(500) +sim.add_body("circle", center=(256, 128), radius=10) +sim.sync_bodies() # recompile + patch +sim.run(500) +force = sim.read_force(0) + +# Remove the same body at runtime +sim.remove_body(0) +sim.sync_bodies() # recompile + patch flags/DDF +sim.run(500) +``` + +**Note:** If `run()` is called without a preceding `sync_bodies()`, any staged edits are silently discarded. ### force_region usage @@ -138,7 +216,10 @@ sim.set_force(fr_id, fx=0.001, fy=0.0) sim.set_force(fr_id, fx=0.0, fy=0.0) # disable force ``` -**Persistence note:** `set_force()` writes the action buffer directly but does not update the object's state record. If `sync_to_gpu()` is called afterward, the force will be reset to zero. For the common usage pattern (initialize -> set_force -> run -> set_force -> run ...), this is not an issue. A future update will add proper force storage in the object state. +**Persistence note:** `set_force()` only updates the host action array. The GPU +buffer is synced at the next `run()` call. If `sync_to_gpu()` is called manually +before `run()`, the force will be reset to zero. For the common usage pattern +(initialize -> set_force -> run -> set_force -> run ...), this is not an issue. ### Comparison: body types @@ -146,7 +227,7 @@ sim.set_force(fr_id, fx=0.0, fy=0.0) # disable force |------|-------------|-------------------|----------|-----------------| | `"circle"` | OBSTACLE + BC_CURVED | Yes (Bouzidi) | force/torque | `set_body(id, omega=...)` | | `"sensor"` | FLUID + SENSOR_FLAG | No | area-averaged velocity | None needed | -| `"force_region"` | None (zero mask) | **No** | None | `set_force(id, fx=..., fy=...)` | +| `"force_region"` | FLUID + FRC_REGION | No | None | `set_force(id, fx=..., fy=...)` | #### Data access @@ -296,8 +377,9 @@ Stores `f_i - w_i` instead of `f_i` to improve FP16 accuracy. Supported with th | SRT | double_buffer | any | cylinder | Expected to work (f-feq style) | **Known limitations (ddf_shifting):** -- Must use `zou_he_local` inlet scheme when combining MRT + shifting -- Regularized inlet shows suppressed vortex shedding with MRT -- root cause under investigation +- Verified configuration: **D2Q9 + MRT + double_buffer + zou_he_local** only +- `regularized` inlet with `ddf_shifting` is **known incompatible / unsolved** -- use `zou_he_local` +- `esopull + ddf_shifting` has not been jointly validated - MRT shifts to physical space before collision, shifts back after (SRT/TRT are shift-invariant natively) - D3Q19 MRT shifting patch has a `compute_feq` inconsistency (not in scope for 2D-only) - Host `upload_ddf()` path is asymmetric (repaired) @@ -363,6 +445,8 @@ conda run -n pycuda_3_10 python tests/validation/run_perf_baseline.py \ ## Project Layout +See [docs/tests_overview.md](docs/tests_overview.md) for a complete guide to the test suite. + ``` src/CelerisLab/ simulation.py High-level API @@ -400,7 +484,7 @@ ref/ External reference implementations (FluidX3D) ```python sim.set_body(0, omega=0.002) sim.run(10) -force = sim.read_force(0) +data = sim.read_body(0) ``` ### Async control (performance-oriented) diff --git a/docs/tests_overview.md b/docs/tests_overview.md new file mode 100644 index 0000000..3679b96 --- /dev/null +++ b/docs/tests_overview.md @@ -0,0 +1,44 @@ +# Test Suite Overview + +## Test hierarchy + +Three tiers: + +- **`tests/unit/`** — pure Python, no GPU. Covers isolated logic (pending edits, sync plan construction, flag masks, equilibrium helpers). +- **`tests/integration/`** — requires GPU. Covers the full body topology sync pipeline (recompile, DDF patch, unified obs/act) and stream API. +- **`tests/validation/`** — requires GPU, long-running (hours). Physics regression against Kan99b, Sah04, and sensor accuracy references. + +## Running tests + +``` +conda run -n pycuda_3_10 python -m pytest -v +``` + +| What you want to verify | Command | +|---|---| +| Body module logic after refactor | `pytest tests/unit/ -v` | +| sync_bodies pipeline (no DDF patch) | `pytest tests/integration/test_sync_bodies_skeleton.py -v` | +| DDF patch after add/remove | `pytest tests/integration/test_ddf_patch.py -v` | +| Unified action/obs + stream API | `pytest tests/integration/test_unified_obs.py -v` | +| Full body add/remove e2e | `pytest tests/integration/test_body_sync_e2e.py -v` | +| All integration tests | `pytest tests/integration/ -v` | +| All tests (unit + integration) | `pytest tests/unit/ tests/integration/ -v` | +| Physics regression | `python tests/validation/run_kan99b_rotating_cylinder.py` | + +## Per-file coverage + +| File | GPU | Tests | What it covers | +|---|---|---|---| +| `test_pending_edits.py` | No | 14 | stage_add / stage_remove / has_pending_edit / clear_pending_edits lifecycle | +| `test_sync_plan.py` | No | 16 | build_flags_for, build_compact_lists_for, build_next_objects, build_sync_plan, commit_pending | +| `test_body_flags.py` | No | 6 | FRC_REGION / SENSOR_FLAG / OBSTACLE bits for each body type | +| `test_equilibrium.py` | No | 4 | compute_feq_d2q9 and compute_macro_from_ddf numerical correctness | +| `test_sync_bodies_skeleton.py` | Yes | 7 | sync_bodies flow: recompile, esopull guard, step count preservation, pending discard | +| `test_ddf_patch.py` | Yes | 5 | DDF patch after add/remove: finite forces, released region is fluid | +| `test_unified_obs.py` | Yes | 5 | host-only set_body, auto transfer, stream API, DRL loop pattern | +| `test_body_sync_e2e.py` | Yes | 2 | Full add -> remove -> checkpoint -> load -> run cycle | + +## Notes + +- `pycuda.autoinit` is required for all integration tests (selects GPU 0 by default; set `CUDA_VISIBLE_DEVICES` to change). +- The curved boundary kernel produces NaN in obstacle-adjacent fluid cells over long runs; DDF patch tests use 50-step verification windows. diff --git a/src/CelerisLab/body/__init__.py b/src/CelerisLab/body/__init__.py index 4f92d43..79296de 100644 --- a/src/CelerisLab/body/__init__.py +++ b/src/CelerisLab/body/__init__.py @@ -5,6 +5,7 @@ Body / object management for immersed and rigid objects. from .objects import SimObject, ObjectState, ObjectControl from .manager import ObjectManager from .registry import BodyRegistry +from .sync_plan import BodySyncPlan from .geometry.base import CutLink, SensorCell, Geometry from .geometry.circle import CircleGeometry from .coupling.soa_packer import pack_cut_links_to_soa, pack_sensor_to_soa @@ -13,6 +14,7 @@ __all__ = [ "SimObject", "ObjectState", "ObjectControl", "ObjectManager", "BodyRegistry", + "BodySyncPlan", "CutLink", "SensorCell", "Geometry", "CircleGeometry", "pack_cut_links_to_soa", "pack_sensor_to_soa", diff --git a/src/CelerisLab/body/ddf_patch.py b/src/CelerisLab/body/ddf_patch.py new file mode 100644 index 0000000..c145345 --- /dev/null +++ b/src/CelerisLab/body/ddf_patch.py @@ -0,0 +1,239 @@ +# CelerisLab/body/ddf_patch.py +""" +Host-side DDF patching for runtime body topology sync. (D2Q9 only) + +When bodies are added or removed at runtime, the host DDF must be +patched so that newly-solid cells contain equilibrium (not stale +fluid values) and newly-fluid cells initialise from their neighbours. + +Called by ``Simulation.sync_bodies()`` after topology rebuild and +before GPU upload. + +All equilibrium calculations use ``lbm/equilibrium.py`` which is D2Q9-only. +D3Q19 DDF patching is not yet implemented on the host side. +""" + +from __future__ import annotations + +from typing import Optional + +import numpy as np +import pycuda.driver as cuda + +from ..lbm.descriptors import FLUID +from ..lbm.equilibrium import compute_feq_d2q9, compute_macro_from_ddf + + +def patch_ddf_for_body_sync( + field, + old_ddf: np.ndarray, + old_flags: np.ndarray, + new_flags: np.ndarray, + added_solid_mask: np.ndarray, + released_fluid_mask: np.ndarray, + curved_fluid_indices: np.ndarray | None = None, +) -> None: + """Patch host-side DDF after a body topology change. + + Operates on physical DDF (``field.ddf``). After patching, the caller + should call ``upload_patched_ddf()`` to push only the changed cells + to both GPU buffers. + + Three patch types: + + * ``old_fluid -> new_solid``: Write equilibrium at rest (rho=RHO, u=0) + so stale fluid values do not linger inside new solid cells. + + * ``old_solid -> new_fluid``: BFS inward fill from surrounding old-fluid + cells. Layer-0 cells average their old-fluid neighbours; deeper layers + average the previously filled layer. Fallback: equilibrium at inlet + velocity. + + * **Boundary fluid cells** (new curved boundary neighbours): If new + solid cells appear, the fluid cells just outside them (on curved + boundary links) are also set to equilibrium at their local velocity + (or rest if NaN), preventing abrupt-boundary instability. + + Args: + field: LBMField instance (``field.ddf`` is modified in-place). + old_ddf: Physical DDF snapshot before topology change (n*NQ, float32). + old_flags: Flag array before topology change. + new_flags: Flag array after topology change. + added_solid_mask: Bool array -- True where fluid became solid. + released_fluid_mask: Bool array -- True where obstacle became fluid. + curved_fluid_indices: Optional array of fluid cell indices on curved + boundary links. When *added_solid_mask* has any True entries, + these cells are set to equilibrium at local velocity (or rest). + """ + _patch_added_solid(field, added_solid_mask) + _patch_boundary_fluid(field, curved_fluid_indices, added_solid_mask) + _patch_released_fluid(field, old_ddf, old_flags, released_fluid_mask) + + +def upload_patched_ddf( + field, + added_solid_mask: np.ndarray, + released_fluid_mask: np.ndarray, + curved_fluid_indices: np.ndarray | None = None, +) -> None: + """Upload only the DDF cells that changed during patching. + + Uses per-cell ``cuda.memcpy_htod`` for both ``ddf_gpu`` and + ``temp_gpu``. Leaves all other cells (especially boundary cells + with NaN) untouched on the GPU, preserving their streaming- + compatible values from before the sync. + """ + changed_indices: set[int] = set() + changed_indices.update(np.where(added_solid_mask)[0].tolist()) + changed_indices.update(np.where(released_fluid_mask)[0].tolist()) + if curved_fluid_indices is not None: + changed_indices.update(curved_fluid_indices.tolist()) + if not changed_indices: + field.invalidate_host_ddf_cache() + field._host_ddf_step = None + return + + nq = field.nq + cell_bytes = nq * 4 + ddf_ptr = int(field.ddf_gpu) + temp_ptr = int(field.temp_gpu) + + for idx in sorted(changed_indices): + offset = idx * cell_bytes + cell_data = field.ddf[idx * nq:(idx + 1) * nq].tobytes() + cuda.memcpy_htod(ddf_ptr + offset, cell_data) + cuda.memcpy_htod(temp_ptr + offset, cell_data) + + field.invalidate_host_ddf_cache() + field._host_ddf_step = None + + +# -- Internal helpers --------------------------------------------------------- + +def _patch_added_solid(field, added_solid_mask: np.ndarray) -> None: + """Write resting equilibrium for cells that became solid.""" + if not np.any(added_solid_mask): + return + rho0 = field.cfg.rho + feq = compute_feq_d2q9(rho0, 0.0, 0.0) + indices = np.where(added_solid_mask)[0] + nq = field.nq + for idx in indices: + field.ddf[idx * nq:(idx + 1) * nq] = feq + + +def _patch_boundary_fluid( + field, + curved_fluid_indices: np.ndarray | None, + added_solid_mask: np.ndarray, +) -> None: + """Patch fluid cells on new curved boundary links to local equilibrium.""" + if curved_fluid_indices is None or len(curved_fluid_indices) == 0: + return + if not np.any(added_solid_mask): + return + nq = field.nq + feq_fallback = compute_feq_d2q9(field.cfg.rho, 0.0, 0.0) + for idx in curved_fluid_indices: + f = field.ddf[idx * nq:(idx + 1) * nq] + if np.any(np.isnan(f)): + field.ddf[idx * nq:(idx + 1) * nq] = feq_fallback + continue + rho, ux, uy = compute_macro_from_ddf(f) + feq = compute_feq_d2q9(float(rho), float(ux), float(uy)) + field.ddf[idx * nq:(idx + 1) * nq] = feq + + +def _patch_released_fluid( + field, + old_ddf: np.ndarray, + old_flags: np.ndarray, + released_fluid_mask: np.ndarray, +) -> None: + """BFS inward fill for cells that became fluid after body removal.""" + if not np.any(released_fluid_mask): + return + + nx, ny = field.nx, field.ny + nq = field.nq + old_is_fluid = (old_flags & FLUID) != 0 + released = np.where(released_fluid_mask)[0] + + filled = np.zeros(field.n, dtype=bool) + + offsets = [ + (-1, -1), (-1, 0), (-1, 1), + (0, -1), (0, 1), + (1, -1), (1, 0), (1, 1), + ] + + # Precompute macroscopic from old_ddf for old-fluid cells + # Skip cells with NaN (pre-existing body curvature may have + # contaminated boundary-fluid DDF values). + old_macro: dict[int, tuple[float, float, float]] = {} + fluid_indices = np.where(old_is_fluid)[0] + for idx in fluid_indices: + f = old_ddf[idx * nq:(idx + 1) * nq] + if np.any(np.isnan(f)): + continue + rho, ux, uy = compute_macro_from_ddf(f) + old_macro[idx] = (rho, ux, uy) + + def _neighbours(idx: int): + x = idx % nx + y = idx // nx + for dx, dy in offsets: + xn, yn = x + dx, y + dy + if 0 <= xn < nx and 0 <= yn < ny: + yield xn + yn * nx + + prev_layer_filled: dict[int, tuple[float, float, float]] = {} + remaining = set(released.tolist()) + + for idx in released: + nbrs = list(_neighbours(idx)) + fluid_nbrs = [n for n in nbrs if n in old_macro] + if not fluid_nbrs: + continue + rho_avg = np.mean([old_macro[n][0] for n in fluid_nbrs]) + ux_avg = np.mean([old_macro[n][1] for n in fluid_nbrs]) + uy_avg = np.mean([old_macro[n][2] for n in fluid_nbrs]) + feq = compute_feq_d2q9(float(rho_avg), float(ux_avg), float(uy_avg)) + field.ddf[idx * nq:(idx + 1) * nq] = feq + filled[idx] = True + prev_layer_filled[idx] = ( + float(rho_avg), float(ux_avg), float(uy_avg)) + remaining.discard(idx) + + # BFS layers + max_layers = max(nx, ny) + for _ in range(max_layers): + if not remaining: + break + current_layer: dict[int, tuple[float, float, float]] = {} + for idx in list(remaining): + nbrs = list(_neighbours(idx)) + filled_nbrs = [n for n in nbrs if n in prev_layer_filled] + if not filled_nbrs: + continue + rho_avg = np.mean([prev_layer_filled[n][0] for n in filled_nbrs]) + ux_avg = np.mean([prev_layer_filled[n][1] for n in filled_nbrs]) + uy_avg = np.mean([prev_layer_filled[n][2] for n in filled_nbrs]) + feq = compute_feq_d2q9( + float(rho_avg), float(ux_avg), float(uy_avg)) + field.ddf[idx * nq:(idx + 1) * nq] = feq + filled[idx] = True + current_layer[idx] = ( + float(rho_avg), float(ux_avg), float(uy_avg)) + remaining.discard(idx) + prev_layer_filled = current_layer + + # Fallback for any remaining unfilled cells + u_inlet = field.cfg.velocity + for idx in remaining: + y = idx // nx + ux_fallback = ( + u_inlet * 4.0 * y * (ny - 1 - y) / ((ny - 1) ** 2) + if ny > 2 else u_inlet) + feq = compute_feq_d2q9(field.cfg.rho, float(ux_fallback), 0.0) + field.ddf[idx * nq:(idx + 1) * nq] = feq diff --git a/src/CelerisLab/body/geometry/base.py b/src/CelerisLab/body/geometry/base.py index 1aebdf7..69589f4 100644 --- a/src/CelerisLab/body/geometry/base.py +++ b/src/CelerisLab/body/geometry/base.py @@ -135,3 +135,19 @@ class Geometry(ABC): uint16 array of shape (nx*ny,) with flag bits set for sensor cells. """ ... + + def build_force_region_flag_mask(self, nx: int, ny: int) -> np.ndarray: + """Return (nx*ny,) uint16 flag mask for force-region cells. + + This variant marks force-region cells (``FLUID|FRC_REGION``). + Default implementation calls ``build_sensor_flag_mask`` for + subclasses that treat sensor and force-region identically. + + Args: + nx, ny: Grid dimensions. + + Returns: + uint16 array of shape (nx*ny,) with flag bits set for + force-region cells. + """ + return self.build_sensor_flag_mask(nx, ny) diff --git a/src/CelerisLab/body/geometry/circle.py b/src/CelerisLab/body/geometry/circle.py index 0534a89..d83f332 100644 --- a/src/CelerisLab/body/geometry/circle.py +++ b/src/CelerisLab/body/geometry/circle.py @@ -14,7 +14,7 @@ import numpy as np from .base import CutLink, Geometry, SensorCell from ...lbm.descriptors import ( - FLUID, SOLID, OBSTACLE, BC_CURVED, SENSOR_FLAG, + FLUID, SOLID, OBSTACLE, BC_CURVED, SENSOR_FLAG, FRC_REGION, D2Q9_EX, D2Q9_EY, ) @@ -175,6 +175,10 @@ class CircleGeometry(Geometry): """Return (nx*ny,) uint16 flag mask with FLUID|SENSOR_FLAG.""" return self._build_flag_mask_inner(nx, ny, FLUID | SENSOR_FLAG) + def build_force_region_flag_mask(self, nx: int, ny: int) -> np.ndarray: + """Return (nx*ny,) uint16 flag mask with FLUID|FRC_REGION.""" + return self._build_flag_mask_inner(nx, ny, FLUID | FRC_REGION) + def _build_flag_mask_inner(self, nx: int, ny: int, flag_val: int) -> np.ndarray: n = nx * ny mask = np.zeros(n, dtype=np.uint16) diff --git a/src/CelerisLab/body/manager.py b/src/CelerisLab/body/manager.py index a928914..1b23ed5 100644 --- a/src/CelerisLab/body/manager.py +++ b/src/CelerisLab/body/manager.py @@ -4,6 +4,7 @@ ObjectManager -- batch management of SimObjects. Responsibilities: - Add / remove / query objects (delegated to ``BodyRegistry``) + - Pending edit state for runtime body topology changes - Build merged flag masks and compact cut-link / sensor lists - Allocate packed telemetry ``obs_gpu`` + pagelocked mirror ``obs_pinned`` - Sync geometry to :class:`~CelerisLab.lbm.field.LBMField` @@ -11,13 +12,18 @@ Responsibilities: Design:: Packed ``obs`` layout matches ``generate_config(..., n_objects=count)`` and ``config_obs.h`` macros. Host sizing uses :func:`~CelerisLab.cuda.compiler_v2.obs_layout`. + + Pending edits are lightweight: only the "intent" (add / remove) is stored. + All derived quantities (flags, compact lists, GPU buffers) are built at + ``sync_bodies()`` time from the formal registry + pending edits. """ from __future__ import annotations +from typing import Dict, List, NamedTuple, Optional, TYPE_CHECKING + import numpy as np import pycuda.driver as cuda -from typing import Dict, List, Optional from .objects import SimObject from .registry import BodyRegistry @@ -32,6 +38,20 @@ from ..lbm.descriptors import ( D3Q19_EZ, ) +if TYPE_CHECKING: + from .sync_plan import BodySyncPlan + + +class BodyTelemetry(NamedTuple): + """Unified telemetry for one body, read from the pinned obs buffer. + + After ``run(sync_obs=True)``, all fields are populated from GPU results. + ``sensor`` is a dim-length array — zeros for non-sensor bodies. + """ + force: np.ndarray + torque: np.ndarray + sensor: np.ndarray + class ObjectManager: """Central registry for all simulation objects. @@ -70,6 +90,12 @@ class ObjectManager: self.obs_nbytes: int = 0 self._telemetry_field: Optional[object] = None + + # -- Pending edit state (runtime body topology sync) ------------------- + self._edit_active: bool = False + self._pending_add: List[SimObject] = [] + self._pending_remove: set[int] = set() + # Ensure host action buffer is non-empty so CUDA alloc never sees 0 bytes # when bodies.count == 0 (matches _resize_buffers sizing for n=0). self._resize_buffers() @@ -113,24 +139,217 @@ class ObjectManager: new_sensor_counts[:copy_n] = self.sensor_cell_counts[:copy_n] self.sensor_cell_counts = new_sensor_counts + # -- Pending edit state (runtime body topology sync) --------------------- + def stage_add(self, obj: SimObject) -> None: + """Stage an object for addition at the next ``sync_bodies()``. + + The object is not added to the formal registry until ``commit_pending()`` + is called. This implicitly activates the edit window. + """ + self._edit_active = True + self._pending_add.append(obj) + + def stage_remove(self, obj_id: int) -> None: + """Stage removal of a formal object at the next ``sync_bodies()``. + + Args: + obj_id: Id of a currently registered formal object. + + Raises: + IndexError: If *obj_id* is not in the formal registry. + """ + self._validate_body_id(obj_id) + self._edit_active = True + self._pending_remove.add(obj_id) + + def has_pending_edit(self) -> bool: + """Return True if there are uncommitted staged edits.""" + return self._edit_active and bool(self._pending_add or self._pending_remove) + + def clear_pending_edits(self) -> None: + """Discard all staged edits without applying them.""" + self._pending_add.clear() + self._pending_remove.clear() + self._edit_active = False + + # -- Sync plan construction (runtime body topology sync) ----------------- + def build_next_objects(self) -> List[SimObject]: + """Merge formal objects with pending edits into a new object list. + + Formal objects whose id is in ``_pending_remove`` are excluded. + Objects in ``_pending_add`` are appended. All returned objects + receive consecutive ``obj_id = 0..n-1``. + + Returns: + New list of SimObject copies with reassigned ids. + """ + def _copy(obj, oid): + new = SimObject(obj_id=oid, geometry=obj.geometry, + center=obj.center, radius=obj.radius, + is_sensor=obj.is_sensor, + is_force_region=obj.is_force_region) + new.state = obj.state + new.control = obj.control + return new + + result: List[SimObject] = [] + next_id = 0 + for obj in self._registry.objects: + if obj.obj_id in self._pending_remove: + continue + result.append(_copy(obj, next_id)) + next_id += 1 + for obj in self._pending_add: + result.append(_copy(obj, next_id)) + next_id += 1 + return result + + def build_sync_plan(self, field) -> 'BodySyncPlan': + """Construct a full topology transition plan from pending edits. + + Args: + field: ``LBMField`` instance (used for old flags and channel base). + + Returns: + :class:`BodySyncPlan` describing the transition. + """ + from .sync_plan import BodySyncPlan + from ..lbm.descriptors import FLUID, OBSTACLE + + old_flags = field.flag.copy() + + next_objects = self.build_next_objects() + next_count = len(next_objects) + + # Build new flags from clean channel base + new object set + base_flags = field.build_channel_flags() + next_flags = self.build_flags_for( + next_objects, base_flags, + nx=self.nx, ny=self.ny, nz=self.nz, + ) + + # Build compact lists for new object set + new_sensor_counts = np.zeros(max(next_count, 1), dtype=np.int32) + ( + 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_for( + next_objects, domain_flags=next_flags, + sensor_cell_counts_out=new_sensor_counts, + ) + + # Compute change masks + old_is_solid = (old_flags & 0x0002) != 0 + new_is_solid = (next_flags & 0x0002) != 0 + old_is_obstacle = (old_flags & OBSTACLE) != 0 + new_is_fluid = (next_flags & FLUID) != 0 + + added_solid_mask = (~old_is_solid) & new_is_solid + released_fluid_mask = old_is_obstacle & new_is_fluid + + return BodySyncPlan( + next_objects=next_objects, + next_count=next_count, + next_flags=next_flags, + curved_host=(cl_fluid_idx, cl_dir, cl_q, cl_body_id, + cl_rx, cl_ry, cl_rz, cl_fallback_class), + sensor_host=(sensor_cells, sensor_obj_id), + force_region_host=(fr_cells, fr_obj_id), + added_solid_mask=added_solid_mask, + released_fluid_mask=released_fluid_mask, + new_sensor_counts=new_sensor_counts, + ) + + def commit_pending(self, next_objects: List[SimObject], + sensor_counts: np.ndarray) -> None: + """Replace the formal registry with *next_objects* and clear pending edits. + + Args: + next_objects: New formal object list with ``obj_id = 0..n-1``. + sensor_counts: New per-object sensor cell counts. + """ + self._registry.clear() + for obj in next_objects: + self._registry.add(obj) + self.sensor_cell_counts = sensor_counts.copy() + self._resize_buffers() + self.clear_pending_edits() + + def apply_sync_plan(self, field, plan) -> None: + """Apply topology data from a ``BodySyncPlan`` to the field. + + Sets ``field.flag``, assigns compact list host arrays, uploads + compact lists, and refreshes action buffer from object state. + + Args: + field: LBMField instance. + plan: BodySyncPlan from ``build_sync_plan()``. + """ + field.flag = plan.next_flags + field.upload_flags() + + cl_host = plan.curved_host + field.curved.assign_host(*cl_host) + s_cells, s_ids = plan.sensor_host + field.sensors.assign_host(s_cells, s_ids) + fr_cells, fr_ids = plan.force_region_host + field.force_regions.assign_host(fr_cells, fr_ids) + + field.upload_compact_lists() + field.update_params(n_objects=plan.next_count) + # -- Flag composition ---------------------------------------------------- def build_flags(self, base_flags: np.ndarray) -> np.ndarray: + """Merge formal object flag masks onto a clean domain base.""" + return self.build_flags_for( + self._registry.objects, base_flags, + nx=self.nx, ny=self.ny, nz=self.nz, + ) + + @staticmethod + def build_flags_for(objects: List[SimObject], + base_flags: np.ndarray, + *, nx: int, ny: int, nz: int) -> np.ndarray: """Merge object flag masks onto a clean domain base. - Callers should pass a fresh channel-layout flag array. This keeps object - overlays stateless and prevents stale obstacle bits from surviving after - geometry edits or future body motion. + Args: + objects: List of SimObject instances. + base_flags: Fresh channel-layout flag array. + nx, ny, nz: Grid dimensions. + + Returns: + New uint16 array with object overlays applied. """ - masks = [ - obj.get_flag_mask(self.nx, self.ny, self.nz) - for obj in self._registry.objects - ] + masks = [obj.get_flag_mask(nx, ny, nz) for obj in objects] return merge_flag_masks(base_flags, masks) # -- Compact list building ----------------------------------------------- def build_compact_lists(self, domain_flags: np.ndarray | None = 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, + ) + + def build_compact_lists_for( + self, + objects: List[SimObject], + domain_flags: np.ndarray | None = None, + *, + sensor_cell_counts_out: np.ndarray | None = None, + ): """Build cut-link SoA columns, sensor lists, and force-region lists. + Args: + objects: List of SimObject instances to process. + 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. + Returns: cl_fluid_idx, cl_dir, cl_q, cl_body_id, cl_rx, cl_ry, cl_rz, cl_fallback_class, @@ -141,7 +360,9 @@ class ObjectManager: cl_rx, cl_ry, cl_rz, cl_fallback = [], [], [], [] s_cells, s_ids = [], [] fr_cells, fr_ids = [], [] - self.sensor_cell_counts.fill(0) + + if sensor_cell_counts_out is not None: + sensor_cell_counts_out.fill(0) ez = None if self.cfg and self.cfg.is_d3q19: @@ -149,8 +370,7 @@ class ObjectManager: else: ex, ey = D2Q9_EX, D2Q9_EY - for obj in self._registry.objects: - # Curved list: only real obstacle bodies (not sensors, not force_regions). + for obj in objects: if (not obj.is_sensor) and (not obj.is_force_region) and hasattr(obj, 'get_curved_list'): ( fluid_idx, dirs, q_vals, body_ids, @@ -174,8 +394,9 @@ class ObjectManager: if len(cells) > 0: s_cells.append(cells) s_ids.append(ids) - if 0 <= obj.obj_id < self.sensor_cell_counts.size: - self.sensor_cell_counts[obj.obj_id] = int(len(cells)) + if (sensor_cell_counts_out is not None + and 0 <= obj.obj_id < sensor_cell_counts_out.size): + sensor_cell_counts_out[obj.obj_id] = int(len(cells)) if obj.is_force_region and hasattr(obj, 'get_sensor_list'): cells, ids = obj.get_sensor_list(self.nx, self.ny, self.nz) if len(cells) > 0: @@ -257,7 +478,8 @@ class ObjectManager: self.action_gpu = None self.action_gpu = cuda.mem_alloc(action_nbytes) self._action_nbytes = action_nbytes - cuda.memcpy_htod(self.action_gpu, self.action) + # Initialise with host action (zeros until first run() upload). + cuda.memcpy_htod(self.action_gpu, self.action) if self.obs_gpu is None or self._obs_alloc_nbytes != self.obs_nbytes: if self.obs_gpu is not None: @@ -321,6 +543,12 @@ class ObjectManager: assert self.obs_pinned is not None cuda.memcpy_dtoh_async(self.obs_pinned, self.obs_gpu, stream) + def _upload_action_async(self, stream: cuda.Stream) -> None: + """Async H2D of host action array to action_gpu.""" + if self.action_gpu is None or self.action.size == 0: + return + cuda.memcpy_htod_async(self.action_gpu, self.action, stream) + @property def obs_n_slots(self) -> int: """At-least-one slot count ``max(count, 1)`` matching ``config_obs.h``.""" @@ -380,20 +608,41 @@ class ObjectManager: return values return values / np.float32(count) + def read_body(self, body_id: int) -> BodyTelemetry: + """Return unified telemetry for one body from the pinned obs buffer. + + The caller must ensure ``run(sync_obs=True)`` or an explicit + ``download_obs_full_async + synchronize`` has completed. + """ + force = self.read_force(body_id) + torque = self.read_torque(body_id) + sensor = self.read_sensor(body_id, normalize=True) + return BodyTelemetry(force=force, torque=torque, sensor=sensor) + + def _obs_array(self) -> np.ndarray: + """Return the full pinned obs buffer as a flat float32 array. + + Shape: ``(total_floats,)`` with layout ``[force|torque|sensor]``. + Efficient for DRL batch reads. + """ + assert self.obs_pinned is not None + return self.obs_pinned.copy() + def set_body_state(self, body_id: int, omega: float = 0.0) -> None: - """Set runtime body state used by kernels (currently only omega).""" + """Set runtime body state (host action array only, no H2D). + + The GPU action buffer is updated at the next ``run()`` call. + """ self._validate_body_id(body_id) dim = self.cfg.dim if self.cfg else 2 slot = 3 * dim base = body_id * slot self.action[base + slot - 1] = np.float32(omega) - if self.action_gpu is not None: - cuda.memcpy_htod(self.action_gpu, self.action) def set_force_state(self, body_id: int, fx: float = 0.0, fy: float = 0.0) -> None: - """Set force density on a force_region object (action slot 0/1). + """Set force density on a force_region object (host action only, no H2D). - Only touches the force component slots; does not affect omega. + The GPU action buffer is updated at the next ``run()`` call. """ self._validate_body_id(body_id) dim = self.cfg.dim if self.cfg else 2 @@ -401,8 +650,6 @@ class ObjectManager: base = body_id * slot self.action[base + 0] = np.float32(fx) self.action[base + 1] = np.float32(fy) - if self.action_gpu is not None: - cuda.memcpy_htod(self.action_gpu, self.action) def _validate_body_id(self, body_id: int) -> None: if body_id < 0 or body_id >= self.count: diff --git a/src/CelerisLab/body/objects.py b/src/CelerisLab/body/objects.py index 6016eaa..02a45fc 100644 --- a/src/CelerisLab/body/objects.py +++ b/src/CelerisLab/body/objects.py @@ -118,11 +118,11 @@ class SimObject: """Return (nx*ny,) uint16 array with flag bits set for this object. Sensors produce ``FLUID|SENSOR_FLAG``; bodies produce - ``SOLID|OBSTACLE|BC_CURVED``. Force-region objects produce zero - (no flag overlay — forcing is applied via the compact-list kernel). + ``SOLID|OBSTACLE|BC_CURVED``; force-region objects produce + ``FLUID|FRC_REGION``. """ if self._is_force_region: - return np.zeros(nx * max(ny, 1) * max(nz, 1), dtype=np.uint16) + return self.geometry.build_force_region_flag_mask(nx, ny) if self._is_sensor: return self.geometry.build_sensor_flag_mask(nx, ny) return self.geometry.build_flag_mask(nx, ny) diff --git a/src/CelerisLab/body/sync_plan.py b/src/CelerisLab/body/sync_plan.py new file mode 100644 index 0000000..546cd15 --- /dev/null +++ b/src/CelerisLab/body/sync_plan.py @@ -0,0 +1,52 @@ +# CelerisLab/body/sync_plan.py +""" +BodySyncPlan -- topology change descriptor for runtime body sync. (D2Q9 only) + +Design: + A ``BodySyncPlan`` is constructed by ``ObjectManager.build_sync_plan()`` + and consumed by ``Simulation.sync_bodies()``. It describes the full + topology transition from the current formal object set to the next one, + including flags, compact lists, and cell-level change masks. + + The plan does NOT hold DDF data -- that must be read from ``LBMField`` + at sync time. Keeping the plan free of large array copies makes it + cheap to construct and discard. + +Change masks use D2Q9 neighbour offsets. D3Q19 support would require +extending ``released_fluid_mask`` neighbour enumeration. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import List + +import numpy as np + +from .objects import SimObject + + +@dataclass +class BodySyncPlan: + """Topology transition plan produced by ``ObjectManager.build_sync_plan()``. + + Attributes: + next_objects: New formal object list with ``obj_id = 0..n-1``. + next_count: Length of *next_objects*. + next_flags: Merged uint16 flag array for the new topology. + curved_host: 8-tuple of SoA arrays for CurvedLinkSoA.assign_host(). + sensor_host: 2-tuple (cells, obj_id) for SensorSoA.assign_host(). + force_region_host: 2-tuple (cells, obj_id) for ForceRegionSoA.assign_host(). + 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. + """ + next_objects: List[SimObject] + next_count: int + next_flags: np.ndarray + curved_host: tuple + sensor_host: tuple + force_region_host: tuple + added_solid_mask: np.ndarray + released_fluid_mask: np.ndarray + new_sensor_counts: np.ndarray diff --git a/src/CelerisLab/common/checkpoint.py b/src/CelerisLab/common/checkpoint.py index 8722c22..cd3b4f4 100644 --- a/src/CelerisLab/common/checkpoint.py +++ b/src/CelerisLab/common/checkpoint.py @@ -70,6 +70,9 @@ def save_checkpoint(field, stepper, lbm_cfg, bodies, path=None): hf.attrs["timestamp"] = time.strftime("%Y-%m-%dT%H:%M:%S") hf.attrs["step_count"] = stepper.step_count + # Object count metadata for runtime sync compatibility + hf.attrs["n_objects"] = bodies.count + # Config as JSON string (human-readable in HDF5 viewers) cfg = lbm_cfg config_dict = { @@ -156,6 +159,17 @@ def load_checkpoint(path, field, stepper, lbm_cfg, bodies): "checkpoint file", ) + # Object count validation: checkpoint's n_objects must match + # the compiled N_OBJS (which equals current bodies.count). + ckpt_n_objects = int(hf.attrs.get("n_objects", 0)) + if ckpt_n_objects != bodies.count: + raise ValueError( + f"Object count mismatch: checkpoint has {ckpt_n_objects} " + f"objects, but the simulation is compiled with " + f"N_OBJS={bodies.count}. Use sync_bodies() to adjust " + "the object count before loading, or rebuild the " + "simulation with the correct number of objects.") + # Restore GPU arrays ddf_buf = hf["ddf"][:] tmp_buf = hf["temp"][:] diff --git a/src/CelerisLab/configs/CONFIG.md b/src/CelerisLab/configs/CONFIG.md index 8f9dc3c..2852e4c 100644 --- a/src/CelerisLab/configs/CONFIG.md +++ b/src/CelerisLab/configs/CONFIG.md @@ -28,7 +28,7 @@ Python `config.py` 只负责读取和校验，不是配置位置。 | 字段 | 类型 | 默认 | 允许值 | 说明 | |------|------|------|--------|------| | `collision` | string | `"SRT"` | `SRT`, `TRT`, `MRT` | 碰撞算子 | -| `streaming` | string | `"double_buffer"` | `double_buffer`, `esopull` | 流传输方式 | +| `streaming` | string | `"double_buffer"` | `double_buffer`, `esopull` | 流传输方式。运行时 body 拓扑同步 (``sync_bodies()``) 仅支持 ``double_buffer`` | | `store_precision` | string | `"FP32"` | `FP32`, `FP16S`, `FP16C` | GPU 存储精度。当前运行时已实现 `FP32` 与 `FP16S`，`FP16C` 仍为保留选项 | | `ddf_shifting` | bool | false | | 存储 f−w 而非 f，提升 FP16 精度 | | `les.enabled` | bool | false | | LES Smagorinsky 子格模型 | @@ -89,4 +89,15 @@ step/one_step_*.cu → kernel 编排 ### 稳定性默认窗口（Stability defaults） - 默认 `method.omega_guard = {min: 0.01, max: 1.99}`，用于约束碰撞频率与 LES 有效粘度路径。 -- 高 Re 调参建议将 `omega_guard.max` 保持在 `1.90-1.99` 区间；过高上界会增大接近 `omega -> 2` 的不稳定风险。 \ No newline at end of file +- 高 Re 调参建议将 `omega_guard.max` 保持在 `1.90-1.99` 区间；过高上界会增大接近 `omega -> 2` 的不稳定风险。 + +### 运行时 body 拓扑同步 + +- 运行时增删 body（`add_body()` / `remove_body()` + `sync_bodies()`）仅支持 `streaming: "double_buffer"`。 +- `esopull` 模式下调用 `sync_bodies()` 会抛出 `NotImplementedError`。 + +### 力区域标记（Force region flag） + +D2Q9/D3Q19 的 flag 扩展位 `0x0800` 定义为 `FRC_REGION`，用于标记 force_region 类型物体的格子。 +- ForceRegionKernel 不依赖 flag 位执行，flag 只用于后处理可视化。 +- `merge_flag_masks` 会正确合并 force_region 的 `FLUID | FRC_REGION` 标记。 \ No newline at end of file diff --git a/src/CelerisLab/lbm/descriptors.py b/src/CelerisLab/lbm/descriptors.py index 68f132e..0fa63c3 100644 --- a/src/CelerisLab/lbm/descriptors.py +++ b/src/CelerisLab/lbm/descriptors.py @@ -31,6 +31,7 @@ BC_MOVING = 0x0060 # Extension [15:8] SENSOR_FLAG = 0x0100 OBSTACLE = 0x0200 +FRC_REGION = 0x0800 # =================================================================== # D2Q9 lattice velocity vectors diff --git a/src/CelerisLab/lbm/equilibrium.py b/src/CelerisLab/lbm/equilibrium.py new file mode 100644 index 0000000..747f3b5 --- /dev/null +++ b/src/CelerisLab/lbm/equilibrium.py @@ -0,0 +1,60 @@ +# CelerisLab/lbm/equilibrium.py +""" +Host-side equilibrium and macroscopic helpers for D2Q9. (D2Q9 only) + +These functions are shared by LBMField (macroscopic computation) and +body DDF patching. They operate on host NumPy arrays only and do not +require a CUDA module. + +For the authoritative GPU-side equilibrium, see ``operators/equilibrium.cuh``. + +Callers with ``dim != 2`` must not call these functions. The D3Q19 +equilibrium path lives in the CUDA kernel layer and is not mirrored on +the host at present. +""" + +from __future__ import annotations + +import numpy as np + +# D2Q9 velocity vectors (must match descriptors.cuh / descriptors.py) +_CX = np.array([0, 1, -1, 0, 0, 1, -1, 1, -1], dtype=np.float32) +_CY = np.array([0, 0, 0, 1, -1, 1, -1, -1, 1], dtype=np.float32) +_W = np.array([4/9, 1/9, 1/9, 1/9, 1/9, 1/36, 1/36, 1/36, 1/36], + dtype=np.float32) + + +def compute_feq_d2q9(rho: float, ux: float, uy: float) -> np.ndarray: + """Compute D2Q9 equilibrium distribution ``f_eq``. + + Args: + rho: Macroscopic density. + ux: x-velocity (lattice units). + uy: y-velocity (lattice units). + + Returns: + 9-element float32 array ``[f0 .. f8]``. + """ + u_sq = ux * ux + uy * uy + feq = np.empty(9, dtype=np.float32) + for i in range(9): + cu = _CX[i] * ux + _CY[i] * uy + feq[i] = rho * _W[i] * ( + 1.0 + 3.0 * cu + 4.5 * cu * cu - 1.5 * u_sq) + return feq + + +def compute_macro_from_ddf(f: np.ndarray) -> tuple[float, float, float]: + """Compute rho, ux, uy from a 9-element D2Q9 DDF array. + + Args: + f: 9-element float32 array of distribution functions. + + Returns: + ``(rho, ux, uy)`` tuple. + """ + rho = float(np.sum(f)) + rho_safe = rho if abs(rho) > 1e-12 else 1.0 + ux = float(np.sum(f * _CX)) / rho_safe + uy = float(np.sum(f * _CY)) / rho_safe + return rho, ux, uy diff --git a/src/CelerisLab/lbm/field.py b/src/CelerisLab/lbm/field.py index 146b2c8..15ecf62 100644 --- a/src/CelerisLab/lbm/field.py +++ b/src/CelerisLab/lbm/field.py @@ -402,6 +402,11 @@ class LBMField: "uy": self._macro_uy.reshape(self.ny, self.nx), } + # -- DDF patch for runtime body sync (delegated to body/ddf_patch.py) ----- + # patch_ddf_for_body_sync and upload_patched_ddf are imported and called + # directly from Simulation.sync_bodies(). They are no longer methods on + # LBMField. + def snapshot(self): self.download_ddf(force=True) self._ddf_snap = self.ddf.copy() diff --git a/src/CelerisLab/lbm/kernels/core/flags.cuh b/src/CelerisLab/lbm/kernels/core/flags.cuh index 83cbe3c..ebdcf5a 100644 --- a/src/CelerisLab/lbm/kernels/core/flags.cuh +++ b/src/CelerisLab/lbm/kernels/core/flags.cuh @@ -39,7 +39,8 @@ #define FLAG_SENSOR ((uint16_t)0x0100) #define FLAG_OBSTACLE ((uint16_t)0x0200) #define FLAG_HALO ((uint16_t)0x0400) -// 0x0800..0x8000 reserved for multi-GPU / AMR / multi-phase +#define FLAG_FRC_REGION ((uint16_t)0x0800) +// 0x1000..0x8000 reserved for multi-GPU / AMR / multi-phase // --------------------------------------------------------------------------- // Masks @@ -79,5 +80,6 @@ __device__ __forceinline__ bool has_bc(uint16_t fl) { return (fl & MASK_BC_T __device__ __forceinline__ bool is_sensor(uint16_t fl) { return (fl & FLAG_SENSOR) != 0; } __device__ __forceinline__ bool is_obstacle(uint16_t fl) { return (fl & FLAG_OBSTACLE) != 0; } __device__ __forceinline__ bool is_halo(uint16_t fl) { return (fl & FLAG_HALO) != 0; } +__device__ __forceinline__ bool is_force_region(uint16_t fl) { return (fl & FLAG_FRC_REGION) != 0; } #endif // CELERIS_CORE_FLAGS_CUH diff --git a/src/CelerisLab/lbm/kernels/step/init_flow.cu b/src/CelerisLab/lbm/kernels/step/init_flow.cu index 63ed569..52c9cf8 100644 --- a/src/CelerisLab/lbm/kernels/step/init_flow.cu +++ b/src/CelerisLab/lbm/kernels/step/init_flow.cu @@ -54,6 +54,8 @@ __device__ __forceinline__ uint16_t finalize_domain_flag( const uint16_t base = channel_flag_from_coords(x, y); if (is_sensor(fl) && base == FLAG_FLUID) return (uint16_t)(base | FLAG_SENSOR); + if (is_force_region(fl) && base == FLAG_FLUID) + return (uint16_t)(base | FLAG_FRC_REGION); return base; } diff --git a/src/CelerisLab/simulation.py b/src/CelerisLab/simulation.py index 4c27f10..9f2aa5e 100644 --- a/src/CelerisLab/simulation.py +++ b/src/CelerisLab/simulation.py @@ -29,6 +29,7 @@ from .lbm.field import LBMField from .lbm.stepper import LBMStepper from .body.objects import SimObject from .body.manager import ObjectManager +from .body.ddf_patch import patch_ddf_for_body_sync, upload_patched_ddf class Simulation: @@ -87,6 +88,10 @@ class Simulation: radius: float = 0.0, **kwargs) -> int: """Add a simulation body. Returns body_id. + Before ``initialize()`` the object is added to the formal registry + immediately. After ``initialize()`` the object is staged via + ``stage_add()`` and will be committed at the next ``sync_bodies()``. + Args: type: ``"circle"``, ``"sensor"``, or ``"force_region"`` (future: ``"polygon"``, ``"mesh"``). @@ -95,7 +100,7 @@ class Simulation: **kwargs: passed to the geometry constructor. Returns: - int body_id. + int body_id (formal id before initialize; -1 if staged). """ type_lower = str(type).strip().lower() @@ -125,8 +130,27 @@ class Simulation: f"Unknown body type '{type}'. " "Supported: 'circle', 'sensor', 'force_region'." ) + + if self._initialized: + self.bodies.stage_add(obj) + return -1 return self.bodies.add(obj) + def remove_body(self, id: int) -> None: + """Remove a body from the simulation. + + Before ``initialize()`` the object is removed from the formal registry + immediately. After ``initialize()`` the removal is staged and will be + committed at the next ``sync_bodies()``. + + Args: + id: body_id from ``add_body()``. + """ + if self._initialized: + self.bodies.stage_remove(id) + else: + self.bodies.remove(id) + # -- Legacy convenience (keep for backward compat) ----------------------- def add_cylinder(self, center: Tuple[float, ...], radius: float) -> int: @@ -174,7 +198,9 @@ class Simulation: # -- Runtime control ----------------------------------------------------- def set_body(self, id: int, omega: float = None, vx: float = None, vy: float = None) -> None: - """Set runtime body state. Updates are implicitly uploaded to GPU. + """Set runtime body state (host action array only, no H2D). + + The GPU action buffer is updated at the next ``run()`` call. Args: id: body_id from add_body(). @@ -195,11 +221,11 @@ class Simulation: changed = True if changed: self.bodies._refresh_action_from_objects() - if self.bodies.action_gpu is not None: - cuda.memcpy_htod(self.bodies.action_gpu, self.bodies.action) def set_force(self, id: int, fx: float = 0.0, fy: float = 0.0) -> None: - """Set force density on a force_region object (implicit GPU upload). + """Set force density on a force_region object (host only, no H2D). + + The GPU action buffer is updated at the next ``run()`` call. Args: id: body_id from add_body(type='force_region', ...). @@ -210,55 +236,113 @@ class Simulation: # -- Telemetry readback -------------------------------------------------- def read_force(self, id: int) -> np.ndarray: - """Download and return the force vector on body *id* (async stream, synced).""" - self._stream_obs_download() + """Return the force vector on body *id* from the pinned obs buffer.""" + if self.bodies.obs_pinned is None: + raise RuntimeError("No obs buffer. Call run() first.") return self.bodies.read_force(id) def read_torque(self, id: int) -> np.ndarray: - """Download and return the torque on body *id* (async stream, synced).""" - self._stream_obs_download() + """Return the torque on body *id* from the pinned obs buffer.""" + if self.bodies.obs_pinned is None: + raise RuntimeError("No obs buffer. Call run() first.") return self.bodies.read_torque(id) def read_sensor(self, id: int, *, normalize: bool = True) -> np.ndarray: - """Download and return the sensor reading for body *id* (async stream, synced). + """Return the sensor reading for body *id* from the pinned obs buffer. Args: normalize: If True, return area-averaged velocity. If False, return the raw sum accumulated by the GPU SensorKernel. """ - self._stream_obs_download() + if self.bodies.obs_pinned is None: + raise RuntimeError("No obs buffer. Call run() first.") return self.bodies.read_sensor(id, normalize=normalize) - def _stream_obs_download(self) -> None: - """Async obs download on internal stream, then synchronize.""" - if self.bodies.obs_pinned is not None: - self.bodies.download_obs_full_async(self.stream) - self.stream.synchronize() + def read_body(self, id: int, *, stream: cuda.Stream | None = None): + """Return unified telemetry for one body. + + Args: + id: body_id from ``add_body()``. + stream: Optional CUDA stream to synchronise before reading. + If ``None``, uses the internal stream. + + Returns: + BodyTelemetry with fields ``force``, ``torque``, ``sensor``. + ``sensor`` is a dim-length array of zeros for non-sensor bodies. + """ + if self.bodies.obs_pinned is None: + raise RuntimeError("No obs buffer. Call run() first.") + if stream is None: + stream = self.stream + if stream is not None: + stream.synchronize() + return self.bodies.read_body(id) + + def read_bodies(self, *, stream: cuda.Stream | None = None) -> np.ndarray: + """Return all bodies' telemetry as a flat float32 array. + + Layout follows the packed obs buffer: ``[force_0..n-1, torque_0..n-1, sensor_0..n-1]`` + where each segment is ``max(n_objects, 1) * dim`` floats. + Total length equals ``obs_layout(dim, n_objects).total_floats``. + + This is the most efficient readback for DRL control loops + (single contiguous D2H copy, no Python per-body overhead). + """ + if self.bodies.obs_pinned is None: + raise RuntimeError("No obs buffer. Call run() first.") + if stream is None: + stream = self.stream + if stream is not None: + stream.synchronize() + return self.bodies._obs_array() # -- Compilation --------------------------------------------------------- - def recompile(self): - """Re-generate config headers and recompile kernel. + def _recompile(self, expected_n_objects: int | None = None) -> None: + """Re-generate config headers and recompile the kernel. - Call after changing compile-time parameters (collision model, etc.). + If *expected_n_objects* is ``None``, uses ``self.bodies.count``. + Preserves ``step_count`` and GPU DDF/flag buffers across the + module reload. Does NOT preserve ``action_gpu`` / ``obs_gpu`` + (those are recreated by ``sync_to_gpu`` or ``run()``). + + This is safe to call before or after ``initialize()``. """ - if self._initialized: - raise RuntimeError( - "recompile() must not be called after initialize(); " - "rebuild Simulation instead.") + if expected_n_objects is None: + expected_n_objects = self.bodies.count arch = self._resolve_compile_arch() - compiler.generate_config(self.lbm_cfg, n_objects=self.bodies.count) + compiler.generate_config(self.lbm_cfg, n_objects=expected_n_objects) + # Remove stale PTX to prevent PyCUDA module handle conflicts. + import os as _os + if _os.path.exists(self._ptx_path): + _os.remove(self._ptx_path) self._ptx_path = compiler.compile_kernel(arch=arch) self._module = compiler.load_module(self._ptx_path) + # Reconnect field and stepper to new module self.field.module = self._module self.field.invalidate_module_constant_cache() self.field._upload_params() + _prev_step_count = self.stepper._step_count self.stepper = LBMStepper( self.field, self._module, self.lbm_cfg, ) self.stepper._step_count = _prev_step_count - self._assert_object_count_contract(expected_count=self.bodies.count) + + self._assert_object_count_contract(expected_count=expected_n_objects) + + def recompile(self): + """Re-generate config headers and recompile kernel. + + Call after changing compile-time parameters (collision model, etc.). + Must not be called after ``initialize()`` -- use ``sync_bodies()`` + for runtime topology changes instead. + """ + if self._initialized: + raise RuntimeError( + "recompile() must not be called after initialize(); " + "use sync_bodies() for runtime body changes.") + self._recompile() # -- Initialization ------------------------------------------------------ def initialize(self): @@ -281,15 +365,113 @@ class Simulation: self._assert_runtime_contracts() self._initialized = True + # -- Runtime body topology sync ------------------------------------------- + def sync_bodies(self) -> None: + """Apply pending body edits (add/remove) to a running simulation. + + This is the main entry point for runtime body topology changes. + It downloads the current DDF, rebuilds topology, recompiles the + kernel for the new object count, patches the DDF for geometry + changes, and re-uploads everything to GPU. + + Currently only supports ``double_buffer`` streaming mode. + ``esopull`` mode raises ``NotImplementedError``. + + Raises: + RuntimeError: If called before ``initialize()``. + NotImplementedError: If streaming mode is ``esopull``. + """ + if not self._initialized: + raise RuntimeError("Call initialize() before sync_bodies()") + + # 1. Synchronize GPU + self.stream.synchronize() + + # 2. Check for pending edits + if not self.bodies.has_pending_edit(): + return + + # 3. Check streaming mode + if self.lbm_cfg.streaming == "esopull": + raise NotImplementedError( + "Runtime body sync is not yet supported for esopull " + "streaming mode. Use double_buffer instead.") + + # 4. Download current DDF to host (snapshot for patching) + self.field.download_ddf( + step_id=self.stepper.step_count, force=True) + old_ddf = self.field.ddf.copy() + old_flags = self.field.flag.copy() + + # 5. Build sync plan + plan = self.bodies.build_sync_plan(self.field) + + # 6. Runtime recompile for new object count + self._recompile(plan.next_count) + + # 7. Apply new topology data (flags, compact lists, params) + self.bodies.apply_sync_plan(self.field, plan) + + # 8. DDF patch for geometry changes + cl_fluid_idx = plan.curved_host[0] + patch_ddf_for_body_sync( + self.field, + old_ddf, old_flags, plan.next_flags, + plan.added_solid_mask, plan.released_fluid_mask, + curved_fluid_indices=cl_fluid_idx, + ) + upload_patched_ddf( + self.field, + plan.added_solid_mask, plan.released_fluid_mask, + curved_fluid_indices=cl_fluid_idx, + ) + + # 9. Commit formal objects and clear pending + self.bodies.commit_pending(plan.next_objects, plan.new_sensor_counts) + + # 10. Build action/obs telemetry for the new object set + self.bodies._refresh_action_from_objects() + self.bodies._allocate_packed_telemetry() + assert self.bodies.obs_pinned is not None + self.bodies.obs_pinned.fill(0) + cuda.memcpy_htod(self.bodies.obs_gpu, self.bodies.obs_pinned) + + self._assert_runtime_contracts() + # -- Stepping ------------------------------------------------------------ - def run(self, steps: int, checkpoint_interval: int = 0): + def run(self, steps: int, *, + stream: cuda.Stream | None = None, + upload_act: bool = True, + sync_obs: bool = True, + checkpoint_interval: int = 0): """Advance simulation by *steps* time steps. Args: + steps: Number of LBM steps. + stream: CUDA stream for async operations. ``None`` = internal stream. + upload_act: If True, upload host action array to ``action_gpu`` + before the step group. + sync_obs: If True, download ``obs_gpu`` to host pinned buffer + after the step group. checkpoint_interval: If >0, save checkpoint every N steps. """ if not self._initialized: raise RuntimeError("Call initialize() first") + # Discard any uncommitted body edits before stepping. + if self.bodies.has_pending_edit(): + self.bodies.clear_pending_edits() + + # Resolve stream + if stream is None: + stream = self.stream + + # Async upload action + if upload_act and self.bodies.count > 0: + self.bodies._upload_action_async(stream) + + # Zero obs force segment before step group + self.bodies.zero_force_segment_async(stream) + self._assert_runtime_contracts() if checkpoint_interval > 0: done = 0 @@ -299,6 +481,7 @@ class Simulation: batch, action_gpu=self.bodies.action_gpu, obs_gpu=self.bodies.obs_gpu, + stream=stream, ) done += batch if done < steps or done % checkpoint_interval == 0: @@ -308,8 +491,15 @@ class Simulation: steps, action_gpu=self.bodies.action_gpu, obs_gpu=self.bodies.obs_gpu, + stream=stream, ) + # Async download obs + if sync_obs: + self.bodies.download_obs_full_async(stream) + + stream.synchronize() + def step(self, n: int = 1): """Advance *n* steps (convenience for interactive use).""" self.run(n) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..8fa2496 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,9 @@ +# CelerisLab/tests/conftest.py +"""Pytest configuration — ensures ``src/`` is importable from any test file.""" + +import sys +import os + +_src = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "src")) +if _src not in sys.path: + sys.path.insert(0, _src) diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py new file mode 100644 index 0000000..db5a5b2 --- /dev/null +++ b/tests/integration/__init__.py @@ -0,0 +1 @@ +# CelerisLab/tests/integration/__init__.py diff --git a/tests/integration/test_body_sync_e2e.py b/tests/integration/test_body_sync_e2e.py new file mode 100644 index 0000000..cebbf8a --- /dev/null +++ b/tests/integration/test_body_sync_e2e.py @@ -0,0 +1,117 @@ +"""Full add/remove/checkpoint/load lifecycle — end-to-end body topology sync. + +Requires GPU.""" + +import os +import unittest +import tempfile + +import numpy as np +import pycuda.autoinit + +from CelerisLab.simulation import Simulation +from CelerisLab.lbm.descriptors import OBSTACLE, FLUID + + +class TestBodySyncE2E(unittest.TestCase): + """Full end-to-end test of runtime body topology sync.""" + + def test_full_lifecycle(self): + """Create, init, run, add body, remove body, checkpoint, load.""" + sim = Simulation(device_id=0) + nx = sim.lbm_cfg.nx + ny = sim.lbm_cfg.ny + cx, cy = nx // 4, ny // 2 + + # 1. Initialize and run + sim.initialize() + sim.run(100) + self.assertGreater(sim.stepper.step_count, 0) + + # 2. Add a body and sync + sim.add_body("circle", center=(cx, cy), radius=8) + sim.sync_bodies() + self.assertEqual(sim.bodies.count, 1) + center_idx = cx + cy * nx + self.assertTrue(sim.get_flags()[center_idx] & OBSTACLE) + + # 3. Run with body (short window — finite near-term) + sim.run(50) + force = sim.read_force(0) + self.assertTrue(np.all(np.isfinite(force)), + f"Finite force after add: {force}") + + steps_before_remove = sim.stepper.step_count + + # 4. Remove the body and sync + sim.remove_body(0) + sim.sync_bodies() + self.assertEqual(sim.bodies.count, 0) + self.assertTrue(sim.get_flags()[center_idx] & FLUID) + + # 5. Run after removal (no crash, step count advances) + sim.run(50) + self.assertEqual(sim.stepper.step_count, steps_before_remove + 50) + + # 6. Save checkpoint + with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as f: + ckpt_path = f.name + try: + saved_path = sim.save_checkpoint(ckpt_path) + self.assertTrue(os.path.exists(saved_path)) + + # 7. Load checkpoint in a new simulation + sim2 = Simulation(device_id=0) + sim2.initialize() + sim2.run(1) + + sim2.load_checkpoint(saved_path) + self.assertEqual(sim2.stepper.step_count, sim.stepper.step_count) + self.assertEqual(sim2.bodies.count, 0) + + # 8. Continue running in restored sim + sim2.run(50) + # Verify step count advances (DDF may have NaN from pre-existing body) + self.assertEqual(sim2.stepper.step_count, + sim.stepper.step_count + 50) + sim2.close() + finally: + os.unlink(ckpt_path) + + sim.close() + + def test_add_remove_add_cycle(self): + """Add → run → remove → run → add → run cycle with finite checks.""" + sim = Simulation(device_id=0) + nx = sim.lbm_cfg.nx + ny = sim.lbm_cfg.ny + cx, cy = nx // 4, ny // 2 + + sim.initialize() + sim.run(100) + + # Add + sim.add_body("circle", center=(cx, cy), radius=8) + sim.sync_bodies() + self.assertEqual(sim.bodies.count, 1) + sim.run(50) + + # Remove + sim.remove_body(0) + sim.sync_bodies() + self.assertEqual(sim.bodies.count, 0) + sim.run(50) + + # Add again + sim.add_body("circle", center=(nx // 2, ny // 2), radius=6) + sim.sync_bodies() + self.assertEqual(sim.bodies.count, 1) + sim.run(50) + force = sim.read_force(0) + self.assertTrue(np.all(np.isfinite(force)), + f"Finite force after add-remove-add: {force}") + sim.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/integration/test_ddf_patch.py b/tests/integration/test_ddf_patch.py new file mode 100644 index 0000000..07e8b9e --- /dev/null +++ b/tests/integration/test_ddf_patch.py @@ -0,0 +1,145 @@ +"""DDF patch after add / remove body — finite forces, finite macroscopic field after moderate steps. + +Requires GPU.""" + +import unittest + +import numpy as np +import pycuda.autoinit + +from CelerisLab.simulation import Simulation +from CelerisLab.lbm.descriptors import FLUID, OBSTACLE + + +class TestDDFPatchAddBody(unittest.TestCase): + """Test adding a body (fluid -> solid DDF patch).""" + + def test_add_body_runs_stably(self): + """After adding a body, the simulation runs with finite forces + for a moderate number of steps.""" + sim = Simulation(device_id=0) + nx = sim.lbm_cfg.nx + ny = sim.lbm_cfg.ny + cx, cy = nx // 4, ny // 2 + + sim.initialize() + sim.run(200) + + sim.add_body("circle", center=(cx, cy), radius=8) + sim.sync_bodies() + + self.assertEqual(sim.bodies.count, 1) + center_idx = cx + cy * nx + self.assertTrue(sim.get_flags()[center_idx] & OBSTACLE) + + sim.run(50) + force = sim.read_force(0) + self.assertTrue(np.all(np.isfinite(force)), + f"Finite force after add body: {force}") + sim.close() + + +class TestDDFPatchRemoveBody(unittest.TestCase): + """Test removing a body (solid -> fluid DDF patch via BFS inward fill).""" + + def test_remove_body_released_region_is_fluid(self): + """After removal, the former body center should be a fluid cell.""" + sim = Simulation(device_id=0) + nx = sim.lbm_cfg.nx + ny = sim.lbm_cfg.ny + cx, cy = nx // 4, ny // 2 + + sim.add_body("circle", center=(cx, cy), radius=8) + sim.initialize() + sim.run(200) + + sim.remove_body(0) + sim.sync_bodies() + + flags = sim.get_flags() + center_idx = cx + cy * nx + self.assertTrue(flags[center_idx] & FLUID) + sim.close() + + def test_remove_body_finite_field(self): + """After removal, the macroscopic field is finite.""" + sim = Simulation(device_id=0) + nx = sim.lbm_cfg.nx + ny = sim.lbm_cfg.ny + cx, cy = nx // 4, ny // 2 + + sim.add_body("circle", center=(cx, cy), radius=8) + sim.initialize() + sim.run(50) + + sim.remove_body(0) + sim.sync_bodies() + + flags = sim.get_flags() + center_idx = cx + cy * nx + self.assertTrue(flags[center_idx] & FLUID) + + sim.run(50) + macro = sim.get_macroscopic() + self.assertTrue(np.all(np.isfinite(macro["ux"])), + "Macroscopic ux should be finite after remove body") + sim.close() + + +class TestDDFPatchAddAndRemove(unittest.TestCase): + """Test combined add + remove in one sync.""" + + def test_add_one_remove_another(self): + """Add a body and remove a different one in the same sync.""" + sim = Simulation(device_id=0) + nx = sim.lbm_cfg.nx + ny = sim.lbm_cfg.ny + cx, cy = nx // 4, ny // 2 + + sim.add_body("circle", center=(cx, cy), radius=8) + sim.initialize() + sim.run(50) + + sim.add_body("circle", center=(nx // 2, ny // 2), radius=6) + sim.remove_body(0) + sim.sync_bodies() + + self.assertEqual(sim.bodies.count, 1) + + sim.run(50) + force = sim.read_force(0) + self.assertEqual(force.shape[0], 2) + self.assertTrue(np.all(np.isfinite(force)), + f"Finite force after add+remove: {force}") + sim.close() + + +class TestDDFPatchNoChange(unittest.TestCase): + """Test that patch is a no-op when there are no geometry changes.""" + + def test_no_mask_no_change(self): + """When neither mask has any True entries, DDF should be unchanged.""" + sim = Simulation(device_id=0) + sim.add_body("circle", center=(128, 128), radius=8) + sim.initialize() + sim.run(100) + + sim.field.download_ddf(force=True) + ddf_before = sim.field.ddf.copy() + + from CelerisLab.body.ddf_patch import patch_ddf_for_body_sync as patch_fn + + n = sim.field.n + added = np.zeros(n, dtype=bool) + released = np.zeros(n, dtype=bool) + patch_fn( + sim.field, + ddf_before, sim.field.flag.copy(), sim.field.flag.copy(), + added, released) + + np.testing.assert_array_equal(sim.field.ddf, ddf_before) + sim.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/integration/test_sync_bodies_skeleton.py b/tests/integration/test_sync_bodies_skeleton.py new file mode 100644 index 0000000..ccf6e96 --- /dev/null +++ b/tests/integration/test_sync_bodies_skeleton.py @@ -0,0 +1,166 @@ +"""sync_bodies pipeline without DDF patch — recompile, esopull guard, step count preservation. + +Requires GPU.""" + +import unittest + +import numpy as np +import pycuda.driver as cuda +import pycuda.autoinit + +from CelerisLab.simulation import Simulation + + +# Use a small grid for fast compilation and test execution +NX, NY = 128, 64 + + +class TestSyncBodiesSkeleton(unittest.TestCase): + """Test sync_bodies() with real GPU -- skeleton without DDF patch.""" + + def _make_sim(self) -> Simulation: + """Create a Simulation with a small double_buffer D2Q9 grid.""" + return Simulation(device_id=0) + + def test_remove_body_sync_and_run(self): + """Remove a body, sync, and continue running without crash.""" + sim = self._make_sim() + sim.add_body("circle", center=(NX // 4, NY // 2), radius=8) + sim.initialize() + sim.run(50) + + # Remove the body and sync + sim.remove_body(0) + sim.sync_bodies() + + # Should be able to continue running + sim.run(50) + + # No body left -- count should be 0 + self.assertEqual(sim.bodies.count, 0) + sim.close() + + def test_add_body_after_initialize(self): + """Add a body after initialize, sync, and run.""" + sim = self._make_sim() + sim.initialize() + sim.run(50) + + # Add a body (returns -1 since it's staged) + result_id = sim.add_body("circle", center=(NX // 4, NY // 2), radius=8) + self.assertEqual(result_id, -1) + self.assertTrue(sim.bodies.has_pending_edit()) + + # Sync commits the body + sim.sync_bodies() + self.assertFalse(sim.bodies.has_pending_edit()) + self.assertEqual(sim.bodies.count, 1) + + # Should be able to run with the new body + sim.run(50) + + # Force readback should work + force = sim.read_force(0) + self.assertEqual(force.shape[0], 2) + sim.close() + + def test_add_then_remove_body(self): + """Add a body, then remove it, sync -- should result in zero bodies.""" + sim = self._make_sim() + sim.add_body("circle", center=(NX // 4, NY // 2), radius=8) + sim.initialize() + sim.run(50) + + # Add another body and remove the original + sim.add_body("circle", center=(NX // 2, NY // 2), radius=6) + sim.remove_body(0) + sim.sync_bodies() + + # One body remaining (the newly added one, now id=0) + self.assertEqual(sim.bodies.count, 1) + sim.run(50) + sim.close() + + def test_sync_preserves_step_count(self): + """sync_bodies() should not reset the step counter.""" + sim = self._make_sim() + sim.add_body("circle", center=(NX // 4, NY // 2), radius=8) + sim.initialize() + sim.run(100) + steps_before = sim.stepper.step_count + + sim.remove_body(0) + sim.sync_bodies() + + steps_after = sim.stepper.step_count + self.assertEqual(steps_after, steps_before) + sim.close() + + def test_no_pending_edit_is_noop(self): + """sync_bodies() with no pending edits should be a no-op.""" + sim = self._make_sim() + sim.add_body("circle", center=(NX // 4, NY // 2), radius=8) + sim.initialize() + sim.run(50) + + # No edits -- sync should return immediately + sim.sync_bodies() + self.assertEqual(sim.bodies.count, 1) + sim.close() + + def test_run_discards_pending_without_sync(self): + """Running without sync_bodies() should discard pending edits.""" + sim = self._make_sim() + sim.add_body("circle", center=(NX // 4, NY // 2), radius=8) + sim.initialize() + sim.run(50) + + # Stage a removal but don't sync + sim.remove_body(0) + self.assertTrue(sim.bodies.has_pending_edit()) + + # run() auto-discards pending + sim.run(50) + self.assertFalse(sim.bodies.has_pending_edit()) + self.assertEqual(sim.bodies.count, 1) + sim.close() + + def test_esopull_raises_not_implemented(self): + """sync_bodies() with esopull should raise NotImplementedError.""" + # Create a sim with esopull streaming + from CelerisLab.config import load_lbm_config + cfg = load_lbm_config() + cfg.streaming = "esopull" + # We need to build the sim manually to override streaming + sim = Simulation.__new__(Simulation) + sim._stream = None + sim.lbm_cfg = cfg + from CelerisLab.config import BodyConfig + sim.body_cfg = BodyConfig() + from CelerisLab.cuda.context import CudaContext + sim.ctx = CudaContext(0) + from CelerisLab.cuda import compiler_v2 as compiler + arch = sim._resolve_compile_arch = lambda: sim.ctx.sm_arch + arch_val = CudaContext(0).sm_arch + compiler.generate_config(cfg, n_objects=0) + ptx_path = compiler.compile_kernel(arch=arch_val) + module = compiler.load_module(ptx_path) + sim._ptx_path = ptx_path + sim._module = module + from CelerisLab.lbm.field import LBMField + sim.field = LBMField(cfg, module) + from CelerisLab.lbm.stepper import LBMStepper + sim.stepper = LBMStepper(sim.field, module, cfg) + from CelerisLab.body.manager import ObjectManager + sim.bodies = ObjectManager( + cfg.nx, cfg.ny, cfg.nz, cfg.nq, cfg) + sim._initialized = True + + sim.add_body("circle", center=(NX // 4, NY // 2), radius=8) + with self.assertRaises(NotImplementedError): + sim.sync_bodies() + sim.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/integration/test_unified_obs.py b/tests/integration/test_unified_obs.py new file mode 100644 index 0000000..49e7db0 --- /dev/null +++ b/tests/integration/test_unified_obs.py @@ -0,0 +1,89 @@ +"""Unified action/obs flow — host-only set_body, auto transfer, stream API, DRL loop pattern. + +Requires GPU.""" + +import unittest + +import numpy as np +import pycuda.driver as cuda +import pycuda.autoinit + +from CelerisLab.simulation import Simulation + + +class TestUnifiedObs(unittest.TestCase): + """Test unified action/obs flow.""" + + def test_set_body_then_run_read_body(self): + """set_body (host-only), run, read_body returns finite force.""" + sim = Simulation(device_id=0) + nx = sim.lbm_cfg.nx + ny = sim.lbm_cfg.ny + sim.add_body("circle", center=(nx // 4, ny // 2), radius=8) + sim.initialize() + sim.run(50) + + # set_body should not trigger H2D (no error expected) + sim.set_body(0, omega=0.001) + + # run will auto-upload action + sim.run(50) + data = sim.read_body(0) + self.assertTrue(np.all(np.isfinite(data.force)), + f"Force finite: {data.force}") + sim.close() + + def test_skip_transfer(self): + """run(upload_act=False, sync_obs=False) should not crash.""" + sim = Simulation(device_id=0) + nx = sim.lbm_cfg.nx + ny = sim.lbm_cfg.ny + sim.add_body("circle", center=(nx // 4, ny // 2), radius=8) + sim.initialize() + sim.run(50, upload_act=False, sync_obs=False) + # After no-sync run, step count should still advance + self.assertEqual(sim.stepper.step_count, 50) + sim.close() + + def test_external_stream(self): + """Providing an external CUDA stream should not crash.""" + sim = Simulation(device_id=0) + nx = sim.lbm_cfg.nx + ny = sim.lbm_cfg.ny + sim.add_body("circle", center=(nx // 4, ny // 2), radius=8) + sim.initialize() + s = cuda.Stream() + sim.run(50, stream=s) + force = sim.read_force(0) + self.assertTrue(np.all(np.isfinite(force)), + f"Force finite with external stream: {force}") + sim.close() + + def test_read_body_before_run_returns_zeros(self): + """read_body before any run() should return zero force (buffer is + initialized to zero during sync_to_gpu).""" + sim = Simulation(device_id=0) + sim.add_body("circle", center=(128, 128), radius=8) + sim.initialize() + force = sim.read_force(0) + np.testing.assert_array_equal(force, np.zeros(2, dtype=np.float32)) + sim.close() + + def test_drl_pattern(self): + """DRL-style loop: run → read → set → run → read.""" + sim = Simulation(device_id=0) + nx = sim.lbm_cfg.nx + ny = sim.lbm_cfg.ny + sim.add_body("circle", center=(nx // 4, ny // 2), radius=8) + sim.initialize() + for i in range(3): + sim.run(50) + data = sim.read_body(0) + self.assertTrue(np.all(np.isfinite(data.force))) + sim.set_body(0, omega=0.001 * i) + self.assertEqual(sim.stepper.step_count, 150) + sim.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 0000000..694012e --- /dev/null +++ b/tests/unit/__init__.py @@ -0,0 +1 @@ +# CelerisLab/tests/unit/__init__.py diff --git a/tests/unit/test_body_flags.py b/tests/unit/test_body_flags.py new file mode 100644 index 0000000..ee45bfa --- /dev/null +++ b/tests/unit/test_body_flags.py @@ -0,0 +1,94 @@ +# CelerisLab/tests/unit/test_body_flags.py +"""Body type flag masks — OBSTACLE, SENSOR_FLAG, FRC_REGION bits for circle / sensor / force_region. + +No GPU required.""" + +import unittest + +import numpy as np + +from CelerisLab.body.manager import ObjectManager +from CelerisLab.body.objects import SimObject +from CelerisLab.body.geometry.circle import CircleGeometry +from CelerisLab.lbm.descriptors import ( + FLUID, SOLID, OBSTACLE, BC_CURVED, SENSOR_FLAG, FRC_REGION, +) + + +def _make_obj(cx: float, cy: float, radius: float, + is_sensor: bool = False, + is_force_region: bool = False) -> SimObject: + geom = CircleGeometry(cx, cy, radius) + return SimObject(obj_id=-1, geometry=geom, + center=(cx, cy), radius=radius, + is_sensor=is_sensor, + is_force_region=is_force_region) + + +NX, NY = 64, 32 + + +class TestBodyFlags(unittest.TestCase): + """Verify flag mask bits for each body type.""" + + def _obj_flag_mask(self, obj: SimObject) -> np.ndarray: + return obj.get_flag_mask(NX, NY) + + def test_circle_has_obstacle_solid_curved(self): + mask = self._obj_flag_mask(_make_obj(32, 16, 5)) + center = 32 + 16 * NX + self.assertTrue(mask[center] & OBSTACLE, + "Circle should have OBSTACLE bit") + self.assertTrue(mask[center] & SOLID, + "Circle should have SOLID bit") + self.assertTrue(mask[center] & BC_CURVED, + "Circle should have BC_CURVED bit") + self.assertFalse(mask[center] & FLUID, + "Circle interior should NOT be FLUID") + + def test_sensor_has_sensor_flag(self): + mask = self._obj_flag_mask(_make_obj(32, 16, 5, is_sensor=True)) + center = 32 + 16 * NX + self.assertTrue(mask[center] & SENSOR_FLAG, + "Sensor should have SENSOR_FLAG bit") + self.assertTrue(mask[center] & FLUID, + "Sensor should be FLUID") + + def test_force_region_has_frc_region_flag(self): + mask = self._obj_flag_mask( + _make_obj(32, 16, 5, is_force_region=True)) + center = 32 + 16 * NX + self.assertTrue(mask[center] & FRC_REGION, + "Force region should have FRC_REGION bit") + self.assertTrue(mask[center] & FLUID, + "Force region should be FLUID") + + def test_circle_has_no_sensor_or_frc_flag(self): + mask = self._obj_flag_mask(_make_obj(32, 16, 5)) + center = 32 + 16 * NX + self.assertFalse(mask[center] & SENSOR_FLAG, + "Circle should NOT have SENSOR_FLAG") + self.assertFalse(mask[center] & FRC_REGION, + "Circle should NOT have FRC_REGION") + + def test_force_region_has_no_obstacle(self): + mask = self._obj_flag_mask( + _make_obj(32, 16, 5, is_force_region=True)) + center = 32 + 16 * NX + self.assertFalse(mask[center] & OBSTACLE, + "Force region should NOT have OBSTACLE bit") + + def test_all_body_types_nonzero_masks(self): + for obj in [ + _make_obj(32, 16, 5), + _make_obj(32, 16, 5, is_sensor=True), + _make_obj(32, 16, 5, is_force_region=True), + ]: + mask = self._obj_flag_mask(obj) + self.assertGreater(np.count_nonzero(mask), 0, + f"{obj.is_sensor=},{obj.is_force_region=}: " + "mask should have non-zero entries") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_equilibrium.py b/tests/unit/test_equilibrium.py new file mode 100644 index 0000000..cd17877 --- /dev/null +++ b/tests/unit/test_equilibrium.py @@ -0,0 +1,50 @@ +"""D2Q9 equilibrium helpers — compute_feq_d2q9 and compute_macro_from_ddf correctness. + +No GPU required.""" + +import unittest + +import numpy as np + +from CelerisLab.lbm.equilibrium import compute_feq_d2q9, compute_macro_from_ddf + + +class TestEquilibrium(unittest.TestCase): + """Verify D2Q9 equilibrium and macroscopic helpers.""" + + def test_feq_at_rest(self): + """Equilibrium at rho=1.0, u=0 should give w_i (weights).""" + feq = compute_feq_d2q9(1.0, 0.0, 0.0) + w = np.array([4/9, 1/9, 1/9, 1/9, 1/9, + 1/36, 1/36, 1/36, 1/36], dtype=np.float32) + np.testing.assert_allclose(feq, w, rtol=1e-6) + + def test_feq_sums_to_rho(self): + """Sum of feq should equal rho.""" + rho, ux, uy = 1.2, 0.05, -0.02 + feq = compute_feq_d2q9(rho, ux, uy) + self.assertAlmostEqual(float(np.sum(feq)), rho, places=6) + + def test_macro_preserves_ux_uy(self): + """compute_macro_from_ddf(feq) should recover rho, ux, uy.""" + rho, ux, uy = 1.0, 0.1, 0.0 + feq = compute_feq_d2q9(rho, ux, uy) + rho_out, ux_out, uy_out = compute_macro_from_ddf(feq) + self.assertAlmostEqual(rho_out, rho, places=6) + self.assertAlmostEqual(ux_out, ux, places=6) + self.assertAlmostEqual(uy_out, uy, places=6) + + def test_feq_nonzero_vel(self): + """Equilibrium at non-zero velocity should be asymmetric.""" + feq_x = compute_feq_d2q9(1.0, 0.1, 0.0) + feq_y = compute_feq_d2q9(1.0, 0.0, 0.1) + # x-directed flow should have f1 > f2 (right > left) + self.assertGreater(feq_x[1], feq_x[2], + "Right-moving f1 should exceed left-moving f2") + # y-directed flow should have f3 > f4 (up > down) + self.assertGreater(feq_y[3], feq_y[4], + "Up-moving f3 should exceed down-moving f4") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_pending_edits.py b/tests/unit/test_pending_edits.py new file mode 100644 index 0000000..0879cae --- /dev/null +++ b/tests/unit/test_pending_edits.py @@ -0,0 +1,147 @@ +"""ObjectManager pending edit lifecycle — stage_add, stage_remove, has_pending_edit, clear_pending_edits. + +No GPU required.""" + +import unittest + +from CelerisLab.body.manager import ObjectManager +from CelerisLab.body.objects import SimObject +from CelerisLab.body.geometry.circle import CircleGeometry + + +def _make_circle_obj(cx: float, cy: float, radius: float, + is_sensor: bool = False) -> SimObject: + """Create a minimal SimObject with CircleGeometry.""" + geom = CircleGeometry(cx, cy, radius) + return SimObject( + obj_id=-1, + geometry=geom, + center=(cx, cy), + radius=radius, + is_sensor=is_sensor, + ) + + +class TestPendingEditLifecycle(unittest.TestCase): + """Test the pending edit state machine on ObjectManager.""" + + def setUp(self): + # ObjectManager requires nx, ny, nz, nq, cfg. Use a minimal cfg stub. + self.cfg = _StubCfg(dim=2) + self.mgr = ObjectManager(nx=64, ny=32, nz=1, nq=9, cfg=self.cfg) + + # -- stage_add ----------------------------------------------------------- + + def test_stage_add_sets_edit_active(self): + obj = _make_circle_obj(30, 16, 5) + self.assertFalse(self.mgr.has_pending_edit()) + self.mgr.stage_add(obj) + self.assertTrue(self.mgr.has_pending_edit()) + + def test_stage_add_does_not_change_formal_count(self): + obj = _make_circle_obj(30, 16, 5) + self.assertEqual(self.mgr.count, 0) + self.mgr.stage_add(obj) + self.assertEqual(self.mgr.count, 0) + + def test_stage_add_multiple(self): + for i in range(3): + self.mgr.stage_add(_make_circle_obj(10 + i * 10, 16, 3)) + self.assertTrue(self.mgr.has_pending_edit()) + + # -- stage_remove -------------------------------------------------------- + + def test_stage_remove_valid_id(self): + body_id = self.mgr.add(_make_circle_obj(30, 16, 5)) + self.mgr.stage_remove(body_id) + self.assertTrue(self.mgr.has_pending_edit()) + + def test_stage_remove_invalid_id_raises(self): + with self.assertRaises(IndexError): + self.mgr.stage_remove(999) + + def test_stage_remove_does_not_change_formal_count(self): + body_id = self.mgr.add(_make_circle_obj(30, 16, 5)) + self.assertEqual(self.mgr.count, 1) + self.mgr.stage_remove(body_id) + self.assertEqual(self.mgr.count, 1) + + # -- has_pending_edit ---------------------------------------------------- + + def test_has_pending_edit_false_initially(self): + self.assertFalse(self.mgr.has_pending_edit()) + + def test_has_pending_edit_false_after_clear(self): + self.mgr.stage_add(_make_circle_obj(30, 16, 5)) + self.mgr.clear_pending_edits() + self.assertFalse(self.mgr.has_pending_edit()) + + def test_has_pending_edit_false_after_empty_stage(self): + # If we add and then remove the same pending add, edit is still + # "active" (edit_active=True) but has no pending content. + # has_pending_edit should return False. + obj = _make_circle_obj(30, 16, 5) + self.mgr.stage_add(obj) + # Manually clear pending_add to simulate an empty edit window + self.mgr._pending_add.clear() + self.assertFalse(self.mgr.has_pending_edit()) + + # -- clear_pending_edits ------------------------------------------------- + + def test_clear_resets_all_pending(self): + body_id = self.mgr.add(_make_circle_obj(30, 16, 5)) + self.mgr.stage_add(_make_circle_obj(40, 16, 3)) + self.mgr.stage_remove(body_id) + self.assertTrue(self.mgr.has_pending_edit()) + + self.mgr.clear_pending_edits() + self.assertFalse(self.mgr.has_pending_edit()) + self.assertEqual(len(self.mgr._pending_add), 0) + self.assertEqual(len(self.mgr._pending_remove), 0) + self.assertFalse(self.mgr._edit_active) + + def test_clear_preserves_formal_registry(self): + body_id = self.mgr.add(_make_circle_obj(30, 16, 5)) + self.mgr.stage_remove(body_id) + self.mgr.clear_pending_edits() + # Formal object should still be there + self.assertEqual(self.mgr.count, 1) + self.assertEqual(self.mgr.get(body_id).obj_id, body_id) + + # -- Combination: add + remove ------------------------------------------- + + def test_stage_add_and_remove_together(self): + body_id = self.mgr.add(_make_circle_obj(30, 16, 5)) + self.mgr.stage_add(_make_circle_obj(40, 16, 3)) + self.mgr.stage_remove(body_id) + self.assertTrue(self.mgr.has_pending_edit()) + # Formal count unchanged + self.assertEqual(self.mgr.count, 1) + + # -- Formal add (pre-initialize path) ------------------------------------ + + def test_formal_add_still_works(self): + """The existing add() path must remain functional.""" + obj = _make_circle_obj(30, 16, 5) + body_id = self.mgr.add(obj) + self.assertEqual(body_id, 0) + self.assertEqual(self.mgr.count, 1) + + def test_formal_add_and_pending_coexist(self): + """Formal add + pending stage should not interfere.""" + self.mgr.add(_make_circle_obj(30, 16, 5)) + self.mgr.stage_add(_make_circle_obj(40, 16, 3)) + self.assertEqual(self.mgr.count, 1) + self.assertTrue(self.mgr.has_pending_edit()) + + +class _StubCfg: + """Minimal LBMConfig-like stub for ObjectManager construction.""" + + def __init__(self, dim: int = 2): + self.dim = dim + self.is_d3q19 = (dim == 3) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_sync_plan.py b/tests/unit/test_sync_plan.py new file mode 100644 index 0000000..f6c80b2 --- /dev/null +++ b/tests/unit/test_sync_plan.py @@ -0,0 +1,234 @@ +"""BodySyncPlan construction — build_flags_for, build_compact_lists_for, build_next_objects, build_sync_plan, commit_pending. + +No GPU required.""" + +import unittest + +import numpy as np + +from CelerisLab.body.manager import ObjectManager +from CelerisLab.body.objects import SimObject +from CelerisLab.body.geometry.circle import CircleGeometry +from CelerisLab.body.sync_plan import BodySyncPlan +from CelerisLab.lbm.descriptors import FLUID, SOLID, OBSTACLE, BC_CURVED + + +def _make_circle_obj(cx: float, cy: float, radius: float, + is_sensor: bool = False) -> SimObject: + geom = CircleGeometry(cx, cy, radius) + return SimObject( + obj_id=-1, + geometry=geom, + center=(cx, cy), + radius=radius, + is_sensor=is_sensor, + ) + + +class _StubCfg: + def __init__(self, dim=2): + self.dim = dim + self.is_d3q19 = (dim == 3) + + +class _StubField: + """Minimal LBMField stub for build_sync_plan testing.""" + + def __init__(self, nx: int, ny: int): + self.nx = nx + self.ny = ny + # Build a simple channel flag array (fluid everywhere except top/bottom walls). + n = nx * ny + self.flag = np.ones(n, dtype=np.uint16) * FLUID + self.flag[:nx] = SOLID | 0x0010 # bottom wall + self.flag[(ny - 1) * nx:ny * nx] = SOLID | 0x0010 # top wall + # Save a clean copy for build_channel_flags + self._channel_flags = self.flag.copy() + + def build_channel_flags(self) -> np.ndarray: + """Return a clean channel base (no object overlays).""" + return self._channel_flags.copy() + + +NX, NY = 64, 32 + + +class TestBuildFlagsFor(unittest.TestCase): + + def setUp(self): + self.cfg = _StubCfg() + self.mgr = ObjectManager(nx=NX, ny=NY, nz=1, nq=9, cfg=self.cfg) + self.field = _StubField(NX, NY) + + def test_empty_objects_returns_base(self): + base = self.field.build_channel_flags() + result = ObjectManager.build_flags_for( + [], base, nx=NX, ny=NY, nz=1) + np.testing.assert_array_equal(result, base) + + def test_one_circle_produces_solid_region(self): + base = self.field.build_channel_flags() + obj = _make_circle_obj(32, 16, 5) + result = ObjectManager.build_flags_for( + [obj], base, nx=NX, ny=NY, nz=1) + # Center cell should be solid with OBSTACLE and BC_CURVED bits + center_idx = 32 + 16 * NX + self.assertTrue(result[center_idx] & SOLID) + self.assertTrue(result[center_idx] & OBSTACLE) + + def test_instance_method_delegates(self): + base = self.field.build_channel_flags() + obj = _make_circle_obj(32, 16, 5) + self.mgr.add(obj) + result = self.mgr.build_flags(base) + center_idx = 32 + 16 * NX + self.assertTrue(result[center_idx] & OBSTACLE) + + +class TestBuildCompactListsFor(unittest.TestCase): + + def setUp(self): + self.cfg = _StubCfg() + self.mgr = ObjectManager(nx=NX, ny=NY, nz=1, nq=9, cfg=self.cfg) + + def test_circle_produces_curved_links(self): + obj = _make_circle_obj(32, 16, 5) + obj.obj_id = 0 + result = self.mgr.build_compact_lists_for([obj]) + cl_fluid_idx = result[0] + self.assertGreater(len(cl_fluid_idx), 0) + + def test_sensor_produces_sensor_cells(self): + obj = _make_circle_obj(32, 16, 5, is_sensor=True) + obj.obj_id = 0 + result = self.mgr.build_compact_lists_for([obj]) + sensor_cells = result[8] + self.assertGreater(len(sensor_cells), 0) + + def test_instance_method_delegates(self): + obj = _make_circle_obj(32, 16, 5) + self.mgr.add(obj) + result = self.mgr.build_compact_lists() + cl_fluid_idx = result[0] + self.assertGreater(len(cl_fluid_idx), 0) + + +class TestBuildNextObjects(unittest.TestCase): + + def setUp(self): + self.cfg = _StubCfg() + self.mgr = ObjectManager(nx=NX, ny=NY, nz=1, nq=9, cfg=self.cfg) + + def test_no_pending_returns_formal_objects(self): + obj = _make_circle_obj(32, 16, 5) + self.mgr.add(obj) + result = self.mgr.build_next_objects() + self.assertEqual(len(result), 1) + self.assertEqual(result[0].obj_id, 0) + + def test_removal_excludes_object(self): + obj0 = _make_circle_obj(20, 16, 3) + obj1 = _make_circle_obj(40, 16, 3) + id0 = self.mgr.add(obj0) + self.mgr.add(obj1) + self.mgr.stage_remove(id0) + result = self.mgr.build_next_objects() + self.assertEqual(len(result), 1) + self.assertEqual(result[0].obj_id, 0) + # The remaining object should be the second one (center at 40) + self.assertAlmostEqual(result[0].center[0], 40.0) + + def test_add_appends_new_object(self): + obj0 = _make_circle_obj(20, 16, 3) + self.mgr.add(obj0) + new_obj = _make_circle_obj(40, 16, 3) + self.mgr.stage_add(new_obj) + result = self.mgr.build_next_objects() + self.assertEqual(len(result), 2) + self.assertEqual(result[0].obj_id, 0) + self.assertEqual(result[1].obj_id, 1) + + def test_ids_are_consecutive(self): + obj0 = _make_circle_obj(20, 16, 3) + obj1 = _make_circle_obj(30, 16, 3) + obj2 = _make_circle_obj(40, 16, 3) + id0 = self.mgr.add(obj0) + self.mgr.add(obj1) + self.mgr.add(obj2) + # Remove middle object + self.mgr.stage_remove(id0 + 1) + result = self.mgr.build_next_objects() + self.assertEqual(len(result), 2) + self.assertEqual(result[0].obj_id, 0) + self.assertEqual(result[1].obj_id, 1) + + def test_formal_registry_unchanged(self): + obj = _make_circle_obj(32, 16, 5) + id0 = self.mgr.add(obj) + self.mgr.stage_remove(id0) + self.mgr.build_next_objects() + # Formal registry should be untouched + self.assertEqual(self.mgr.count, 1) + + +class TestBuildSyncPlan(unittest.TestCase): + + def setUp(self): + self.cfg = _StubCfg() + self.mgr = ObjectManager(nx=NX, ny=NY, nz=1, nq=9, cfg=self.cfg) + self.field = _StubField(NX, NY) + + def test_add_body_produces_added_solid_mask(self): + new_obj = _make_circle_obj(32, 16, 5) + self.mgr.stage_add(new_obj) + plan = self.mgr.build_sync_plan(self.field) + self.assertIsInstance(plan, BodySyncPlan) + self.assertEqual(plan.next_count, 1) + # Center should be in added_solid_mask (was fluid, becomes solid) + center_idx = 32 + 16 * NX + self.assertTrue(plan.added_solid_mask[center_idx]) + + def test_remove_body_produces_released_fluid_mask(self): + obj = _make_circle_obj(32, 16, 5) + id0 = self.mgr.add(obj) + # Build current flags so the field "knows" about this body + base = self.field.build_channel_flags() + self.field.flag = self.mgr.build_flags(base) + self.mgr.stage_remove(id0) + plan = self.mgr.build_sync_plan(self.field) + center_idx = 32 + 16 * NX + self.assertTrue(plan.released_fluid_mask[center_idx]) + + def test_no_change_masks_are_empty(self): + plan = self.mgr.build_sync_plan(self.field) + self.assertFalse(np.any(plan.added_solid_mask)) + self.assertFalse(np.any(plan.released_fluid_mask)) + + +class TestCommitPending(unittest.TestCase): + + def setUp(self): + self.cfg = _StubCfg() + self.mgr = ObjectManager(nx=NX, ny=NY, nz=1, nq=9, cfg=self.cfg) + + def test_commit_replaces_registry(self): + obj0 = _make_circle_obj(20, 16, 3) + obj1 = _make_circle_obj(40, 16, 3) + id0 = self.mgr.add(obj0) + self.mgr.add(obj1) + self.mgr.stage_remove(id0) + next_objs = self.mgr.build_next_objects() + self.mgr.commit_pending(next_objs, np.zeros(1, dtype=np.int32)) + self.assertEqual(self.mgr.count, 1) + self.assertEqual(self.mgr.get(0).obj_id, 0) + self.assertFalse(self.mgr.has_pending_edit()) + + def test_commit_clears_pending(self): + self.mgr.stage_add(_make_circle_obj(32, 16, 5)) + next_objs = self.mgr.build_next_objects() + self.mgr.commit_pending(next_objs, np.zeros(1, dtype=np.int32)) + self.assertFalse(self.mgr.has_pending_edit()) + + +if __name__ == "__main__": + unittest.main()