From da9044dc0c75b9b8bc81c3ff35a195ebc28bc25c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 14 Aug 2026 23:47:13 +0200 Subject: [PATCH 1/7] feat: add the MPM solver (2D and 3D) --- Cargo.toml | 12 +- crates/nexus_mpm2d/Cargo.toml | 53 ++ crates/nexus_mpm2d/build.rs | 19 + crates/nexus_mpm3d/Cargo.toml | 55 ++ crates/nexus_mpm3d/build.rs | 19 + crates/nexus_mpm_shaders2d/Cargo.toml | 49 ++ crates/nexus_mpm_shaders2d/build.rs | 3 + crates/nexus_mpm_shaders3d/Cargo.toml | 48 ++ crates/nexus_mpm_shaders3d/build.rs | 3 + src_mpm/grid/grid.rs | 547 +++++++++++++ src_mpm/grid/mod.rs | 4 + src_mpm/grid/sort.rs | 103 +++ src_mpm/lib.rs | 34 + src_mpm/models/drucker_prager.rs | 57 ++ src_mpm/models/mod.rs | 51 ++ src_mpm/pipeline.rs | 558 +++++++++++++ src_mpm/sampling/mod.rs | 15 + src_mpm/sampling/sample_polyline.rs | 57 ++ src_mpm/sampling/sample_trimesh.rs | 231 ++++++ src_mpm/solver/boundary_condition.rs | 39 + src_mpm/solver/g2p.rs | 70 ++ src_mpm/solver/g2p_cdf.rs | 41 + src_mpm/solver/grid_update.rs | 68 ++ src_mpm/solver/grid_update_cdf.rs | 41 + src_mpm/solver/mod.rs | 35 + src_mpm/solver/p2g.rs | 66 ++ src_mpm/solver/p2g_cdf.rs | 45 ++ src_mpm/solver/params.rs | 23 + src_mpm/solver/particle.rs | 648 +++++++++++++++ src_mpm/solver/particle_model.rs | 325 ++++++++ src_mpm/solver/particle_update.rs | 45 ++ src_mpm/solver/prep_readback.rs | 251 ++++++ src_mpm/solver/rigid_integrate.rs | 121 +++ src_mpm/solver/rigid_particle_update.rs | 60 ++ src_mpm/solver/timestep_bound.rs | 70 ++ src_mpm/trimesh.rs | 132 ++++ src_mpm_shaders/collision/collide.rs | 60 ++ src_mpm_shaders/collision/mod.rs | 1 + src_mpm_shaders/grid/grid.rs | 744 ++++++++++++++++++ src_mpm_shaders/grid/kernel.rs | 180 +++++ src_mpm_shaders/grid/mod.rs | 3 + src_mpm_shaders/grid/sort.rs | 667 ++++++++++++++++ src_mpm_shaders/lib.rs | 160 ++++ src_mpm_shaders/models/default.rs | 343 ++++++++ src_mpm_shaders/models/drucker_prager.rs | 268 +++++++ src_mpm_shaders/models/fluid.rs | 107 +++ src_mpm_shaders/models/interfaces.rs | 48 ++ src_mpm_shaders/models/linear_elasticity.rs | 99 +++ src_mpm_shaders/models/mod.rs | 9 + .../models/neo_hookean_elasticity.rs | 84 ++ src_mpm_shaders/models/snow.rs | 114 +++ src_mpm_shaders/models/specializations.rs | 1 + src_mpm_shaders/models/utils.rs | 176 +++++ src_mpm_shaders/solver/boundary_condition.rs | 114 +++ src_mpm_shaders/solver/g2p.rs | 474 +++++++++++ src_mpm_shaders/solver/g2p_cdf.rs | 407 ++++++++++ src_mpm_shaders/solver/grid_update.rs | 88 +++ src_mpm_shaders/solver/grid_update_cdf.rs | 132 ++++ src_mpm_shaders/solver/grid_update_collide.rs | 183 +++++ src_mpm_shaders/solver/mod.rs | 15 + src_mpm_shaders/solver/p2g.rs | 449 +++++++++++ src_mpm_shaders/solver/p2g_cdf.rs | 259 ++++++ src_mpm_shaders/solver/params.rs | 17 + src_mpm_shaders/solver/particle.rs | 238 ++++++ src_mpm_shaders/solver/particle_update.rs | 201 +++++ src_mpm_shaders/solver/prep_readback.rs | 469 +++++++++++ src_mpm_shaders/solver/rigid_impulses.rs | 185 +++++ .../solver/rigid_particle_update.rs | 54 ++ src_mpm_shaders/solver/timestep_bound.rs | 123 +++ 69 files changed, 10468 insertions(+), 2 deletions(-) create mode 100644 crates/nexus_mpm2d/Cargo.toml create mode 100644 crates/nexus_mpm2d/build.rs create mode 100644 crates/nexus_mpm3d/Cargo.toml create mode 100644 crates/nexus_mpm3d/build.rs create mode 100644 crates/nexus_mpm_shaders2d/Cargo.toml create mode 100644 crates/nexus_mpm_shaders2d/build.rs create mode 100644 crates/nexus_mpm_shaders3d/Cargo.toml create mode 100644 crates/nexus_mpm_shaders3d/build.rs create mode 100644 src_mpm/grid/grid.rs create mode 100644 src_mpm/grid/mod.rs create mode 100644 src_mpm/grid/sort.rs create mode 100644 src_mpm/lib.rs create mode 100644 src_mpm/models/drucker_prager.rs create mode 100644 src_mpm/models/mod.rs create mode 100644 src_mpm/pipeline.rs create mode 100644 src_mpm/sampling/mod.rs create mode 100644 src_mpm/sampling/sample_polyline.rs create mode 100644 src_mpm/sampling/sample_trimesh.rs create mode 100644 src_mpm/solver/boundary_condition.rs create mode 100644 src_mpm/solver/g2p.rs create mode 100644 src_mpm/solver/g2p_cdf.rs create mode 100644 src_mpm/solver/grid_update.rs create mode 100644 src_mpm/solver/grid_update_cdf.rs create mode 100644 src_mpm/solver/mod.rs create mode 100644 src_mpm/solver/p2g.rs create mode 100644 src_mpm/solver/p2g_cdf.rs create mode 100644 src_mpm/solver/params.rs create mode 100644 src_mpm/solver/particle.rs create mode 100644 src_mpm/solver/particle_model.rs create mode 100644 src_mpm/solver/particle_update.rs create mode 100644 src_mpm/solver/prep_readback.rs create mode 100644 src_mpm/solver/rigid_integrate.rs create mode 100644 src_mpm/solver/rigid_particle_update.rs create mode 100644 src_mpm/solver/timestep_bound.rs create mode 100644 src_mpm/trimesh.rs create mode 100644 src_mpm_shaders/collision/collide.rs create mode 100644 src_mpm_shaders/collision/mod.rs create mode 100644 src_mpm_shaders/grid/grid.rs create mode 100644 src_mpm_shaders/grid/kernel.rs create mode 100644 src_mpm_shaders/grid/mod.rs create mode 100644 src_mpm_shaders/grid/sort.rs create mode 100644 src_mpm_shaders/lib.rs create mode 100644 src_mpm_shaders/models/default.rs create mode 100644 src_mpm_shaders/models/drucker_prager.rs create mode 100644 src_mpm_shaders/models/fluid.rs create mode 100644 src_mpm_shaders/models/interfaces.rs create mode 100644 src_mpm_shaders/models/linear_elasticity.rs create mode 100644 src_mpm_shaders/models/mod.rs create mode 100644 src_mpm_shaders/models/neo_hookean_elasticity.rs create mode 100644 src_mpm_shaders/models/snow.rs create mode 100644 src_mpm_shaders/models/specializations.rs create mode 100644 src_mpm_shaders/models/utils.rs create mode 100644 src_mpm_shaders/solver/boundary_condition.rs create mode 100644 src_mpm_shaders/solver/g2p.rs create mode 100644 src_mpm_shaders/solver/g2p_cdf.rs create mode 100644 src_mpm_shaders/solver/grid_update.rs create mode 100644 src_mpm_shaders/solver/grid_update_cdf.rs create mode 100644 src_mpm_shaders/solver/grid_update_collide.rs create mode 100644 src_mpm_shaders/solver/mod.rs create mode 100644 src_mpm_shaders/solver/p2g.rs create mode 100644 src_mpm_shaders/solver/p2g_cdf.rs create mode 100644 src_mpm_shaders/solver/params.rs create mode 100644 src_mpm_shaders/solver/particle.rs create mode 100644 src_mpm_shaders/solver/particle_update.rs create mode 100644 src_mpm_shaders/solver/prep_readback.rs create mode 100644 src_mpm_shaders/solver/rigid_impulses.rs create mode 100644 src_mpm_shaders/solver/rigid_particle_update.rs create mode 100644 src_mpm_shaders/solver/timestep_bound.rs diff --git a/Cargo.toml b/Cargo.toml index c2043719..15d3db0f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,6 +10,10 @@ members = [ "crates/nexus_rbd3d", "crates/nexus_rbd_shaders2d", "crates/nexus_rbd_shaders3d", + "crates/nexus_mpm_shaders2d", + "crates/nexus_mpm_shaders3d", + "crates/nexus_mpm2d", + "crates/nexus_mpm3d", "crates/nexus_python3d", ] resolver = "2" @@ -68,13 +72,17 @@ crunchy = "0.2.4" # Shader crates nexus_rbd_shaders2d = { version = "0.4.0", path = "crates/nexus_rbd_shaders2d" } nexus_rbd_shaders3d = { version = "0.4.0", path = "crates/nexus_rbd_shaders3d" } +nexus_mpm_shaders2d = { version = "0.4.0", path = "crates/nexus_mpm_shaders2d" } +nexus_mpm_shaders3d = { version = "0.4.0", path = "crates/nexus_mpm_shaders3d" } -# Internal crates. rbd is pulled with default-features off so dependents can -# pick their own feature set, mirroring the rapier/parry pattern. +# Internal crates. rbd is pulled both with defaults and with default-features +# off (mpm wants just the dim feature), so it mirrors the rapier/parry pattern. nexus2d = { version = "0.4.0", path = "crates/nexus2d" } nexus3d = { version = "0.4.0", path = "crates/nexus3d" } nexus_rbd2d = { version = "0.4.0", path = "crates/nexus_rbd2d", default-features = false } nexus_rbd3d = { version = "0.4.0", path = "crates/nexus_rbd3d", default-features = false } +nexus_mpm2d = { version = "0.4.0", path = "crates/nexus_mpm2d" } +nexus_mpm3d = { version = "0.4.0", path = "crates/nexus_mpm3d" } nexus_viewer2d = { version = "0.4.0", path = "crates/nexus_viewer2d" } nexus_viewer3d = { version = "0.4.0", path = "crates/nexus_viewer3d" } diff --git a/crates/nexus_mpm2d/Cargo.toml b/crates/nexus_mpm2d/Cargo.toml new file mode 100644 index 00000000..7fe5f5dd --- /dev/null +++ b/crates/nexus_mpm2d/Cargo.toml @@ -0,0 +1,53 @@ +[package] +name = "nexus_mpm2d" +authors = { workspace = true } +description = "Cross-platform 2D GPU-accelerated rigid-body physics." +repository = { workspace = true } +version = { workspace = true } +edition = { workspace = true } +license = { workspace = true } + +[lib] +name = "nexus_mpm2d" +path = "../../src_mpm/lib.rs" +required-features = ["dim2"] + +[lints] +rust.unexpected_cfgs = { level = "warn", check-cfg = [ + 'cfg(feature, values("dim3"))', +] } + +[features] +default = ["dim2", "f32", "webgpu"] +dim2 = [] +f32 = [] +f64 = [] +webgpu = ["khal/webgpu"] +metal = ["khal/metal", "nexus_rbd2d/metal"] +cpu = ["nexus_mpm_shaders2d/cpu", "nexus_rbd2d/cpu", "vortx/cpu"] +cpu-parallel = ["cpu", "nexus_mpm_shaders2d/cpu-parallel", "nexus_rbd2d/cpu-parallel", "vortx/cpu-parallel"] +cuda = ["khal/cuda", "khal-builder/cuda", "nexus_mpm_shaders2d/cuda"] + +[dependencies] +nexus_mpm_shaders2d = { workspace = true } +nexus_rbd2d = { workspace = true, features = ["dim2"] } +glamx = { workspace = true } +include_dir = { workspace = true } +khal = { workspace = true } +vortx = { workspace = true } +bytemuck = { workspace = true } +web-time = { workspace = true } +static_assertions = { workspace = true } +bvh = { workspace = true } +rapier2d = { workspace = true, features = ["default"] } +parry2d = { workspace = true, features = ["default"] } + +[dev-dependencies] +futures-test = { workspace = true } +serial_test = { workspace = true } +approx = { workspace = true } +rand = { workspace = true } + +[build-dependencies] +khal-builder = {workspace = true} +nexus_mpm_shaders2d = { workspace = true } diff --git a/crates/nexus_mpm2d/build.rs b/crates/nexus_mpm2d/build.rs new file mode 100644 index 00000000..854e5bf8 --- /dev/null +++ b/crates/nexus_mpm2d/build.rs @@ -0,0 +1,19 @@ +use khal_builder::KhalBuilder; +use std::path::PathBuf; + +fn main() { + let output_dir = PathBuf::from(std::env::var_os("OUT_DIR").expect("OUT_DIR not set by cargo")) + .join("shaders-spirv"); + let target_arch = std::env::var("CARGO_CFG_TARGET_ARCH").unwrap(); + + let mut builder = KhalBuilder::from_dependency("nexus_mpm_shaders2d", true).feature("dim2"); + + // NOTE: this has a significant performance impact on native (not so much on web). + builder = builder.feature("unsafe_remove_boundchecks"); + + if target_arch == "wasm32" { + builder = builder.feature("web-compat"); + } + + builder.build(output_dir); +} diff --git a/crates/nexus_mpm3d/Cargo.toml b/crates/nexus_mpm3d/Cargo.toml new file mode 100644 index 00000000..9350c180 --- /dev/null +++ b/crates/nexus_mpm3d/Cargo.toml @@ -0,0 +1,55 @@ +[package] +name = "nexus_mpm3d" +authors = { workspace = true } +description = "Cross-platform 3D GPU-accelerated MPM physics." +repository = { workspace = true } +version = { workspace = true } +edition = { workspace = true } +license = { workspace = true } + +[lib] +name = "nexus_mpm3d" +path = "../../src_mpm/lib.rs" +required-features = ["dim3"] + +[lints] +rust.unexpected_cfgs = { level = "warn", check-cfg = [ + 'cfg(feature, values("dim2"))', +] } + +[features] +default = ["dim3", "f32", "webgpu"] +dim3 = [] +f32 = [] +f64 = [] +webgpu = ["khal/webgpu"] +metal = ["khal/metal", "nexus_rbd3d/metal"] +cpu = ["nexus_mpm_shaders3d/cpu", "nexus_rbd3d/cpu", "vortx/cpu"] +cpu-parallel = ["cpu", "nexus_mpm_shaders3d/cpu-parallel", "nexus_rbd3d/cpu-parallel", "vortx/cpu-parallel"] +cuda = ["khal/cuda", "khal-builder/cuda", "nexus_mpm_shaders3d/cuda"] + +[dependencies] +nexus_mpm_shaders3d = { workspace = true } +nexus_rbd3d = { workspace = true, features = ["dim3"] } +glamx = { workspace = true } +include_dir = { workspace = true } +khal = { workspace = true } +vortx = { workspace = true } +bytemuck = { workspace = true } +web-time = { workspace = true } +static_assertions = { workspace = true } +bvh = { workspace = true } + +# Optional dependencies for rapier/parry interop +rapier3d = { workspace = true, features = ["default"] } +parry3d = { workspace = true, features = ["default"] } + +[dev-dependencies] +futures-test = { workspace = true } +serial_test = { workspace = true } +approx = { workspace = true } +rand = { workspace = true } + +[build-dependencies] +khal-builder = {workspace = true} +nexus_mpm_shaders3d = { workspace = true } diff --git a/crates/nexus_mpm3d/build.rs b/crates/nexus_mpm3d/build.rs new file mode 100644 index 00000000..9197a975 --- /dev/null +++ b/crates/nexus_mpm3d/build.rs @@ -0,0 +1,19 @@ +use khal_builder::KhalBuilder; +use std::path::PathBuf; + +fn main() { + let output_dir = PathBuf::from(std::env::var_os("OUT_DIR").expect("OUT_DIR not set by cargo")) + .join("shaders-spirv"); + let target_arch = std::env::var("CARGO_CFG_TARGET_ARCH").unwrap(); + + let mut builder = KhalBuilder::from_dependency("nexus_mpm_shaders3d", true).feature("dim3"); + + // NOTE: this has a significant performance impact on native (not so much on web). + builder = builder.feature("unsafe_remove_boundchecks"); + + if target_arch == "wasm32" { + builder = builder.feature("web-compat"); + } + + builder.build(output_dir); +} diff --git a/crates/nexus_mpm_shaders2d/Cargo.toml b/crates/nexus_mpm_shaders2d/Cargo.toml new file mode 100644 index 00000000..34a45ded --- /dev/null +++ b/crates/nexus_mpm_shaders2d/Cargo.toml @@ -0,0 +1,49 @@ +[package] +name = "nexus_mpm_shaders2d" +authors = { workspace = true } +description = "2D GPU shaders for nexus physics engine (rust-gpu)." +repository = { workspace = true } +version = { workspace = true } +edition = { workspace = true } +license = { workspace = true } +links = "nexus_mpm_shaders2d" +build = "build.rs" + + +[lib] +path = "../../src_mpm_shaders/lib.rs" + +[lints] +rust.unexpected_cfgs = { level = "warn", check-cfg = [ + 'cfg(feature, values("dim3"))', + 'cfg(target_arch, values("spirv", "nvptx64"))', + 'cfg(target_arch_is_gpu)', +] } + +[features] +default = ["dim2"] +dim2 = [] +unsafe_remove_boundchecks = ["vortx-shaders/unsafe_remove_boundchecks"] +push_constants = [] +# Enables some changes in the shaders for compatibility with web platforms. +web-compat = [] +cpu = [] +cpu-parallel = ["cpu", "vortx-shaders/cpu-parallel"] +cuda = [] + +[dependencies] +nexus_rbd_shaders2d = { workspace = true } +vortx-shaders = { workspace = true } +khal-std = { workspace = true } +unroll = { workspace = true } +crunchy = { workspace = true } +glamx = { workspace = true } +parry2d = { workspace = true, features = ["dim2","f32"] } + +[build-dependencies] +khal-std = { workspace = true } + +# Host-only dependencies (excluded on GPU targets: spirv, nvptx64). +[target.'cfg(not(any(target_arch = "spirv", target_arch = "nvptx64")))'.dependencies] +bytemuck = { workspace = true } +khal = { workspace = true } diff --git a/crates/nexus_mpm_shaders2d/build.rs b/crates/nexus_mpm_shaders2d/build.rs new file mode 100644 index 00000000..d50740ce --- /dev/null +++ b/crates/nexus_mpm_shaders2d/build.rs @@ -0,0 +1,3 @@ +fn main() { + khal_std::setup_shader_crate_build(); +} diff --git a/crates/nexus_mpm_shaders3d/Cargo.toml b/crates/nexus_mpm_shaders3d/Cargo.toml new file mode 100644 index 00000000..f3ef990a --- /dev/null +++ b/crates/nexus_mpm_shaders3d/Cargo.toml @@ -0,0 +1,48 @@ +[package] +name = "nexus_mpm_shaders3d" +authors = { workspace = true } +description = "3D GPU shaders for nexus physics engine (rust-gpu)." +repository = { workspace = true } +version = { workspace = true } +edition = { workspace = true } +license = { workspace = true } +links = "nexus_mpm_shaders3d" +build = "build.rs" + +[lib] +path = "../../src_mpm_shaders/lib.rs" + +[lints] +rust.unexpected_cfgs = { level = "warn", check-cfg = [ + 'cfg(feature, values("dim2"))', + 'cfg(target_arch, values("spirv", "nvptx64"))', + 'cfg(target_arch_is_gpu)', +] } + +[features] +default = ["dim3"] +dim3 = [] +unsafe_remove_boundchecks = ["vortx-shaders/unsafe_remove_boundchecks"] +push_constants = [] +# Enables some changes in the shaders for compatibility with web platforms. +web-compat = [] +cpu = [] +cpu-parallel = ["cpu", "vortx-shaders/cpu-parallel"] +cuda = [] + +[dependencies] +nexus_rbd_shaders3d = { workspace = true } +vortx-shaders = { workspace = true } +khal-std = { workspace = true } +unroll = { workspace = true } +crunchy = { workspace = true } +glamx = { workspace = true } +parry3d = { workspace = true, features = ["dim3","f32"] } + +[build-dependencies] +khal-std = { workspace = true } + +# Host-only dependencies (excluded on GPU targets: spirv, nvptx64). +[target.'cfg(not(any(target_arch = "spirv", target_arch = "nvptx64")))'.dependencies] +bytemuck = { workspace = true } +khal = { workspace = true } \ No newline at end of file diff --git a/crates/nexus_mpm_shaders3d/build.rs b/crates/nexus_mpm_shaders3d/build.rs new file mode 100644 index 00000000..d50740ce --- /dev/null +++ b/crates/nexus_mpm_shaders3d/build.rs @@ -0,0 +1,3 @@ +fn main() { + khal_std::setup_shader_crate_build(); +} diff --git a/src_mpm/grid/grid.rs b/src_mpm/grid/grid.rs new file mode 100644 index 00000000..e5cad982 --- /dev/null +++ b/src_mpm/grid/grid.rs @@ -0,0 +1,547 @@ +//! Grid data structures and GPU kernels for sparse grid management. + +use crate::grid::sort::WgSort; +use crate::mpm_shaders::grid::grid::{ + ActiveBlockHeader, GpuCaptureNumActiveBlocks, GpuInitIndirectWorkgroups, GpuResetHmap, Grid, + GridHashMapEntry, Node, +}; +use crate::solver::{GpuParticles, GpuRigidParticles}; +use khal::backend::{Encoder, GpuBackend, GpuBackendError, GpuEncoder, GpuPass, GpuTimestamps}; +use khal::{BufferUsages, Shader}; +use nexus_rbd::utils::{GpuPrefixSum, PrefixSumWorkspace}; +use vortx::tensor::Tensor; + +/// GPU kernels for grid initialization and management. +/// +/// Handles sparse grid allocation, reset, and indirect dispatch setup. +#[derive(Shader)] +pub struct WgGrid { + reset_hmap: GpuResetHmap, + capture_num_active_blocks: GpuCaptureNumActiveBlocks, + init_indirect_workgroups: GpuInitIndirectWorkgroups, +} + +impl WgGrid { + /// Sorts particles into grid cells and allocates sparse grid blocks. + pub fn launch_sort( + &self, + backend: &GpuBackend, + pass: &mut GpuPass, + particles: &mut GpuParticles, + mut rigid_particles: Option<&mut GpuRigidParticles>, + grid: &mut GpuGrid, + prefix_sum: &mut PrefixSumWorkspace, + sort_module: &WgSort, + prefix_sum_module: &GpuPrefixSum, + ) -> Result<(), GpuBackendError> { + let particles_len = particles.len() as u32; + let hmap_capacity = grid.cpu_meta.hmap_capacity; + + // Retry until we allocated enough room on the sparse grid for all the blocks. + let mut sparse_grid_has_the_correct_size = false; + while !sparse_grid_has_the_correct_size { + // - Reset next grid's hashmap. + // - Reset grid.num_active_blocks to 0. + // - Run touch_particle_blocks on the next grid. + // - Readback num_active_blocks. + // - Update the hashmap & grid buffer sizes if its occupancy is too high. + + // NOTE: num_active_blocks := 0 is set in reset_hmap. + self.reset_hmap + .call(pass, hmap_capacity, &mut grid.meta, &mut grid.hmap_entries)?; + + // Block activation in two passes (cheaper than every particle inserting + // all NUM_ASSOC_BLOCKS of its stencil): + // 1. Each particle activates only its primary (base) block. + // 2. Each base block activates its +1 neighbour blocks (once per block, + // not once per particle). `active_blocks_snapshot` records the base-block + // count so pass 2 doesn't reprocess the neighbours it appends. + sort_module.touch_primary_blocks.call( + pass, + particles_len, + &mut grid.meta, + &mut grid.hmap_entries, + &mut grid.active_blocks, + &particles.positions, + &particles.gpu_len, + )?; + + self.capture_num_active_blocks.call( + pass, + 1u32, + &grid.meta, + &mut grid.active_blocks_snapshot, + )?; + + // Indirect dispatch sized for the base blocks (the neighbour pass runs one + // thread per base block). + self.init_indirect_workgroups.call( + pass, + 1u32, + &grid.meta, + &mut grid.indirect_n_blocks_groups, + &mut grid.indirect_n_g2p_p2g_groups, + )?; + + sort_module.touch_neighbor_blocks.call( + pass, + indirect_dispatch_tensor(&grid.indirect_n_blocks_groups), + &mut grid.meta, + &mut grid.hmap_entries, + &mut grid.active_blocks, + &grid.active_blocks_snapshot, + )?; + + // Ensure blocks exist wherever we have rigid particles that might affect + // other blocks. This is done in two passes: + // 1. Mark all rigid particles that need to ensure its associated block exists + // 2. Touch the blocks with marked rigid particles. + if let Some(rigid_particles) = rigid_particles.as_deref_mut() + && !rigid_particles.is_empty() + { + let rigid_particles_len = rigid_particles.len() as u32; + sort_module.mark_rigid_particles_needing_block.call( + pass, + rigid_particles_len, + &grid.meta, + &grid.hmap_entries, + &rigid_particles.sample_points, + &mut rigid_particles.rigid_particle_needs_block, + )?; + + sort_module.touch_rigid_particle_blocks.call( + pass, + rigid_particles_len, + &mut grid.meta, + &mut grid.hmap_entries, + &mut grid.active_blocks, + &rigid_particles.sample_points, + &rigid_particles.rigid_particle_needs_block, + )?; + } + + // TODO: handle grid buffer resizing + sparse_grid_has_the_correct_size = true; + } + + // - Launch update_block_particle_count + // - Launch copy_particle_len_to_scan_value + // - Launch cumulated sum. + // - Launch copy_scan_values_to_first_particles + // - Launch finalize_particles_sort + // - Launch write_blocks_multiplicity_to_scan_value + // - Launch cumulated sum + + // Prepare workgroups for indirect dispatches based on the number of active blocks. + self.init_indirect_workgroups.call( + pass, + 1u32, + &grid.meta, + &mut grid.indirect_n_blocks_groups, + &mut grid.indirect_n_g2p_p2g_groups, + )?; + + sort_module.update_nbh_block_ids.call( + pass, + indirect_dispatch_tensor(&grid.indirect_n_blocks_groups), + &grid.meta, + &grid.hmap_entries, + &mut grid.active_blocks, + )?; + + sort_module.update_block_particle_count.call( + pass, + particles_len, + &grid.meta, + &grid.hmap_entries, + &particles.positions, + &particles.gpu_len, + &mut grid.active_blocks, + )?; + + sort_module.copy_particles_len_to_scan_value.call( + pass, + indirect_dispatch_tensor(&grid.indirect_n_blocks_groups), + &grid.meta, + &grid.active_blocks, + &mut grid.scan_values, + )?; + prefix_sum_module.launch(backend, pass, prefix_sum, &mut grid.scan_values, 1)?; + + sort_module.copy_scan_values_to_first_particles.call( + pass, + indirect_dispatch_tensor(&grid.indirect_n_blocks_groups), + &grid.meta, + &grid.scan_values, + &mut grid.active_blocks, + )?; + + sort_module.finalize_particles_sort.call( + pass, + particles_len, + &grid.meta, + &grid.hmap_entries, + &particles.positions, + &particles.gpu_len, + &mut grid.active_blocks, + &mut particles.sorted_ids, + )?; + + Ok(()) + } + + /// Test helper: resets the hashmap and activates blocks for `particles` using + /// either the legacy single-pass touch (`two_pass = false`) or the new two-pass + /// touch (`two_pass = true`), leaving `grid.meta.num_active_blocks` readable. + /// + /// Both paths must activate the identical set of blocks; a benchmark compares the + /// resulting `num_active_blocks` to validate the two-pass touch without depending + /// on CPU/GPU rounding agreement. + #[doc(hidden)] + pub fn launch_touch_for_test( + &self, + pass: &mut GpuPass, + particles: &GpuParticles, + grid: &mut GpuGrid, + sort_module: &WgSort, + two_pass: bool, + ) -> Result<(), GpuBackendError> { + let particles_len = particles.len() as u32; + let hmap_capacity = grid.cpu_meta.hmap_capacity; + + self.reset_hmap + .call(pass, hmap_capacity, &mut grid.meta, &mut grid.hmap_entries)?; + + if two_pass { + sort_module.touch_primary_blocks.call( + pass, + particles_len, + &mut grid.meta, + &mut grid.hmap_entries, + &mut grid.active_blocks, + &particles.positions, + &particles.gpu_len, + )?; + self.capture_num_active_blocks.call( + pass, + 1u32, + &grid.meta, + &mut grid.active_blocks_snapshot, + )?; + self.init_indirect_workgroups.call( + pass, + 1u32, + &grid.meta, + &mut grid.indirect_n_blocks_groups, + &mut grid.indirect_n_g2p_p2g_groups, + )?; + sort_module.touch_neighbor_blocks.call( + pass, + indirect_dispatch_tensor(&grid.indirect_n_blocks_groups), + &mut grid.meta, + &mut grid.hmap_entries, + &mut grid.active_blocks, + &grid.active_blocks_snapshot, + )?; + } else { + sort_module.touch_particle_blocks.call( + pass, + particles_len, + &mut grid.meta, + &mut grid.hmap_entries, + &mut grid.active_blocks, + &particles.positions, + &particles.gpu_len, + )?; + } + + Ok(()) + } + + /// Per-kernel-profiled variant of [`launch_sort`](Self::launch_sort) for the + /// regular-particle (CPIC-disabled) path. + /// + /// Runs the same kernel sequence as `launch_sort` with `rigid_particles = None`, + /// but wraps each sub-kernel in its own timestamp scope so a benchmark can see + /// where the sort spends its time. Diagnostics only: the production pipeline + /// uses `launch_sort`. + #[doc(hidden)] + pub fn launch_sort_profiled( + &self, + backend: &GpuBackend, + encoder: &mut GpuEncoder, + timestamps: &mut GpuTimestamps, + particles: &mut GpuParticles, + grid: &mut GpuGrid, + prefix_sum: &mut PrefixSumWorkspace, + sort_module: &WgSort, + prefix_sum_module: &GpuPrefixSum, + ) -> Result<(), GpuBackendError> { + let particles_len = particles.len() as u32; + let hmap_capacity = grid.cpu_meta.hmap_capacity; + + { + let mut pass = encoder.begin_pass("sort:reset_hmap", Some(timestamps)); + self.reset_hmap.call( + &mut pass, + hmap_capacity, + &mut grid.meta, + &mut grid.hmap_entries, + )?; + } + { + let mut pass = encoder.begin_pass("sort:touch_primary_blocks", Some(timestamps)); + sort_module.touch_primary_blocks.call( + &mut pass, + particles_len, + &mut grid.meta, + &mut grid.hmap_entries, + &mut grid.active_blocks, + &particles.positions, + &particles.gpu_len, + )?; + } + { + let mut pass = encoder.begin_pass("sort:capture_num_active_blocks", Some(timestamps)); + self.capture_num_active_blocks.call( + &mut pass, + 1u32, + &grid.meta, + &mut grid.active_blocks_snapshot, + )?; + } + { + let mut pass = encoder.begin_pass("sort:init_indirect_workgroups", Some(timestamps)); + self.init_indirect_workgroups.call( + &mut pass, + 1u32, + &grid.meta, + &mut grid.indirect_n_blocks_groups, + &mut grid.indirect_n_g2p_p2g_groups, + )?; + } + { + let mut pass = encoder.begin_pass("sort:touch_neighbor_blocks", Some(timestamps)); + sort_module.touch_neighbor_blocks.call( + &mut pass, + indirect_dispatch_tensor(&grid.indirect_n_blocks_groups), + &mut grid.meta, + &mut grid.hmap_entries, + &mut grid.active_blocks, + &grid.active_blocks_snapshot, + )?; + } + { + let mut pass = encoder.begin_pass("sort:init_indirect_workgroups2", Some(timestamps)); + self.init_indirect_workgroups.call( + &mut pass, + 1u32, + &grid.meta, + &mut grid.indirect_n_blocks_groups, + &mut grid.indirect_n_g2p_p2g_groups, + )?; + } + { + let mut pass = encoder.begin_pass("sort:update_nbh_block_ids", Some(timestamps)); + sort_module.update_nbh_block_ids.call( + &mut pass, + indirect_dispatch_tensor(&grid.indirect_n_blocks_groups), + &grid.meta, + &grid.hmap_entries, + &mut grid.active_blocks, + )?; + } + { + let mut pass = encoder.begin_pass("sort:update_block_particle_count", Some(timestamps)); + sort_module.update_block_particle_count.call( + &mut pass, + particles_len, + &grid.meta, + &grid.hmap_entries, + &particles.positions, + &particles.gpu_len, + &mut grid.active_blocks, + )?; + } + { + let mut pass = + encoder.begin_pass("sort:copy_particles_len_to_scan_value", Some(timestamps)); + sort_module.copy_particles_len_to_scan_value.call( + &mut pass, + indirect_dispatch_tensor(&grid.indirect_n_blocks_groups), + &grid.meta, + &grid.active_blocks, + &mut grid.scan_values, + )?; + } + { + let mut pass = encoder.begin_pass("sort:prefix_sum", Some(timestamps)); + prefix_sum_module.launch(backend, &mut pass, prefix_sum, &mut grid.scan_values, 1)?; + } + { + let mut pass = + encoder.begin_pass("sort:copy_scan_values_to_first_particles", Some(timestamps)); + sort_module.copy_scan_values_to_first_particles.call( + &mut pass, + indirect_dispatch_tensor(&grid.indirect_n_blocks_groups), + &grid.meta, + &grid.scan_values, + &mut grid.active_blocks, + )?; + } + + { + let mut pass = encoder.begin_pass("sort:finalize_particles_sort", Some(timestamps)); + sort_module.finalize_particles_sort.call( + &mut pass, + particles_len, + &grid.meta, + &grid.hmap_entries, + &particles.positions, + &particles.gpu_len, + &mut grid.active_blocks, + &mut particles.sorted_ids, + )?; + } + + Ok(()) + } +} + +/// Reinterprets a `Tensor` (with 3 elements) as a `Tensor<[u32; 3]>` for indirect dispatch. +/// +/// # Safety +/// The underlying GPU buffer is just raw bytes; `Tensor` with 3 elements has identical +/// memory layout to `Tensor<[u32; 3]>` with 1 element. Both `u32` and `[u32; 3]` are `Pod`. +pub(crate) fn indirect_dispatch_tensor(tensor: &Tensor) -> &Tensor<[u32; 3]> { + unsafe { &*(tensor as *const Tensor as *const Tensor<[u32; 3]>) } +} + +/// GPU-resident sparse grid structure. +/// +/// The MPM grid uses a sparse representation with a hashmap to efficiently +/// store only active blocks (blocks containing particles). This dramatically +/// reduces memory usage for spatially localized simulations. +pub struct GpuGrid { + /// CPU copy of grid metadata for readback. + pub cpu_meta: Grid, + /// GPU buffer containing grid metadata. + pub meta: Tensor, + /// Pong buffer for grid metadata. + pub prev_meta: Tensor, + /// Hash map entries for virtual-to-physical block mapping. + pub hmap_entries: Tensor, + /// Pong buffer for hmap entries. + pub prev_hmap_entries: Tensor, + /// Grid node data (momentum, mass, CDF). + pub nodes: Tensor, + /// Active block headers tracking particle ranges. + pub active_blocks: Tensor, + /// Workspace for prefix sum operations. + pub scan_values: Tensor, + /// Single-element snapshot of `num_active_blocks` taken after the primary-block + /// touch pass, so the neighbour-block touch pass only iterates over base blocks. + pub active_blocks_snapshot: Tensor, + /// Indirect dispatch arguments for block-parallel kernels. + /// + /// Stored as `Tensor` with 3 elements so it can be written by + /// `init_indirect_workgroups` (which operates on `&mut [u32]`). + /// Use [`indirect_n_blocks_dispatch`](Self::indirect_n_blocks_dispatch) + /// to obtain a `DispatchGrid` for indirect dispatch. + pub indirect_n_blocks_groups: Tensor, + /// Indirect dispatch arguments for node-parallel kernels. + /// + /// Same layout as `indirect_n_blocks_groups`. Use + /// [`indirect_n_g2p_p2g_dispatch`](Self::indirect_n_g2p_p2g_dispatch) + /// for indirect dispatch. + pub indirect_n_g2p_p2g_groups: Tensor, + /// Debug buffer for GPU-side diagnostics. + pub debug: Tensor, +} + +impl GpuGrid { + /// Returns indirect dispatch arguments for block-parallel kernels. + /// + /// This reinterprets a `Tensor` (with 3 elements) as a `Tensor<[u32; 3]>` + /// (with 1 element). This is sound because the memory layout is identical and both + /// types are `Pod`. + pub fn indirect_n_blocks_dispatch(&self) -> &Tensor<[u32; 3]> { + indirect_dispatch_tensor(&self.indirect_n_blocks_groups) + } + + /// Returns indirect dispatch arguments for node-parallel (G2P/P2G) kernels. + /// + /// See [`indirect_n_blocks_dispatch`](Self::indirect_n_blocks_dispatch) for safety rationale. + pub fn indirect_n_g2p_p2g_dispatch(&self) -> &Tensor<[u32; 3]> { + indirect_dispatch_tensor(&self.indirect_n_g2p_p2g_groups) + } + + /// Creates a new sparse grid with the specified capacity. + pub fn with_capacity( + backend: &GpuBackend, + capacity: u32, + cell_width: f32, + ) -> Result { + const NODES_PER_BLOCK: u32 = 64; // 8 * 8 in 2D and 4 * 4 * 4 in 3D. + let capacity = capacity.next_power_of_two(); + let cpu_meta = Grid { + num_active_blocks: 0, + cell_width, + hmap_capacity: capacity, + capacity, + }; + let meta = Tensor::scalar( + backend, + cpu_meta, + BufferUsages::UNIFORM | BufferUsages::STORAGE | BufferUsages::COPY_SRC, + )?; + let prev_meta = Tensor::scalar( + backend, + cpu_meta, + BufferUsages::UNIFORM | BufferUsages::STORAGE | BufferUsages::COPY_SRC, + )?; + let default_entry = GridHashMapEntry { + state: 0xFFFFFFFF, + key: Default::default(), + value: Default::default(), + ownership: 0, + padding: [0; _], + }; + let default_entries = vec![default_entry; capacity as usize]; + let prev_hmap_entries = Tensor::vector(backend, &default_entries, BufferUsages::STORAGE)?; + let hmap_entries = Tensor::vector(backend, &default_entries, BufferUsages::STORAGE)?; + let nodes = + Tensor::vector_uninit(backend, capacity * NODES_PER_BLOCK, BufferUsages::STORAGE)?; + let active_blocks = Tensor::vector_uninit(backend, capacity, BufferUsages::STORAGE)?; + let scan_values = Tensor::vector_uninit(backend, capacity, BufferUsages::STORAGE)?; + let active_blocks_snapshot = Tensor::vector(backend, [0u32], BufferUsages::STORAGE)?; + let indirect_n_blocks_groups = + Tensor::vector_uninit(backend, 3, BufferUsages::STORAGE | BufferUsages::INDIRECT)?; + let indirect_n_g2p_p2g_groups = Tensor::vector_uninit( + backend, + 3, + BufferUsages::STORAGE | BufferUsages::INDIRECT | BufferUsages::COPY_SRC, + )?; + let debug = Tensor::vector(backend, [0u32, 0], BufferUsages::STORAGE)?; + + Ok(Self { + cpu_meta, + meta, + prev_meta, + hmap_entries, + prev_hmap_entries, + nodes, + active_blocks, + scan_values, + active_blocks_snapshot, + indirect_n_blocks_groups, + indirect_n_g2p_p2g_groups, + debug, + }) + } + + pub fn swap_buffers(&mut self) { + std::mem::swap(&mut self.meta, &mut self.prev_meta); + std::mem::swap(&mut self.prev_hmap_entries, &mut self.hmap_entries); + } +} diff --git a/src_mpm/grid/mod.rs b/src_mpm/grid/mod.rs new file mode 100644 index 00000000..b9464809 --- /dev/null +++ b/src_mpm/grid/mod.rs @@ -0,0 +1,4 @@ +//! Spatial grid data structures and operations. + +pub mod grid; +pub mod sort; diff --git a/src_mpm/grid/sort.rs b/src_mpm/grid/sort.rs new file mode 100644 index 00000000..13a47317 --- /dev/null +++ b/src_mpm/grid/sort.rs @@ -0,0 +1,103 @@ +//! Particle sorting kernels for spatial acceleration. +//! +//! These kernels handle spatial hashing and sorting to group particles by grid block +//! for efficient neighbor queries during P2G/G2P. + +use crate::grid::grid::{GpuGrid, indirect_dispatch_tensor}; +use crate::mpm_shaders::grid::sort::{ + GpuCopyParticlesLenToScanValue, GpuCopyRigidParticlesLenToScanValue, + GpuCopyScanValuesToFirstParticles, GpuCopyScanValuesToFirstRigidParticles, + GpuFinalizeParticlesSort, GpuFinalizeRigidParticlesSort, GpuMarkRigidParticlesNeedingBlock, + GpuTouchNeighborBlocks, GpuTouchParticleBlocks, GpuTouchPrimaryBlocks, + GpuTouchRigidParticleBlocks, GpuUpdateBlockParticleCount, GpuUpdateBlockRigidParticleCount, + GpuUpdateNbhBlockIds, +}; +use crate::solver::GpuRigidParticles; +use khal::Shader; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use nexus_rbd::utils::{GpuPrefixSum, PrefixSumWorkspace}; + +/// GPU compute kernels for sorting particles into grid cells. +/// +/// Implements spatial hashing and sorting to group particles by grid block +/// for efficient neighbor queries during P2G/G2P. +#[derive(Shader)] +pub struct WgSort { + /// Legacy single-pass block activation, kept for the `launch_touch_for_test` + /// correctness check against the two-pass `touch_primary_blocks`/`touch_neighbor_blocks`. + pub(crate) touch_particle_blocks: GpuTouchParticleBlocks, + pub(crate) touch_primary_blocks: GpuTouchPrimaryBlocks, + pub(crate) touch_neighbor_blocks: GpuTouchNeighborBlocks, + pub(crate) touch_rigid_particle_blocks: GpuTouchRigidParticleBlocks, + pub(crate) mark_rigid_particles_needing_block: GpuMarkRigidParticlesNeedingBlock, + pub(crate) update_block_particle_count: GpuUpdateBlockParticleCount, + pub(crate) update_nbh_block_ids: GpuUpdateNbhBlockIds, + pub(crate) copy_particles_len_to_scan_value: GpuCopyParticlesLenToScanValue, + pub(crate) copy_scan_values_to_first_particles: GpuCopyScanValuesToFirstParticles, + pub(crate) finalize_particles_sort: GpuFinalizeParticlesSort, + pub(crate) update_block_rigid_particle_count: GpuUpdateBlockRigidParticleCount, + pub(crate) copy_rigid_particles_len_to_scan_value: GpuCopyRigidParticlesLenToScanValue, + pub(crate) copy_scan_values_to_first_rigid_particles: GpuCopyScanValuesToFirstRigidParticles, + pub(crate) finalize_rigid_particles_sort: GpuFinalizeRigidParticlesSort, +} + +impl WgSort { + /// Sorts rigid body particles by grid block. + /// + /// Runs the same count / prefix-sum / finalize sequence as the regular particle + /// sort, reusing `grid.scan_values` (which the regular sort no longer needs once + /// `launch_sort` returned). The result feeds the scatter-style P2G-CDF kernel. + pub fn launch_sort_rigid_particles( + &self, + backend: &GpuBackend, + pass: &mut GpuPass, + rigid_particles: &mut GpuRigidParticles, + grid: &mut GpuGrid, + prefix_sum: &mut PrefixSumWorkspace, + prefix_sum_module: &GpuPrefixSum, + ) -> Result<(), GpuBackendError> { + if rigid_particles.is_empty() { + return Ok(()); + } + + let rigid_particles_len = rigid_particles.len() as u32; + + self.update_block_rigid_particle_count.call( + pass, + rigid_particles_len, + &grid.meta, + &grid.hmap_entries, + &rigid_particles.sample_points, + &mut grid.active_blocks, + )?; + + self.copy_rigid_particles_len_to_scan_value.call( + pass, + indirect_dispatch_tensor(&grid.indirect_n_blocks_groups), + &grid.meta, + &grid.active_blocks, + &mut grid.scan_values, + )?; + prefix_sum_module.launch(backend, pass, prefix_sum, &mut grid.scan_values, 1)?; + + self.copy_scan_values_to_first_rigid_particles.call( + pass, + indirect_dispatch_tensor(&grid.indirect_n_blocks_groups), + &grid.meta, + &grid.scan_values, + &mut grid.active_blocks, + )?; + + self.finalize_rigid_particles_sort.call( + pass, + rigid_particles_len, + &grid.meta, + &grid.hmap_entries, + &rigid_particles.sample_points, + &mut grid.active_blocks, + &mut rigid_particles.sorted_ids, + )?; + + Ok(()) + } +} diff --git a/src_mpm/lib.rs b/src_mpm/lib.rs new file mode 100644 index 00000000..7217c61f --- /dev/null +++ b/src_mpm/lib.rs @@ -0,0 +1,34 @@ +//! GPU-accelerated Material Point Method simulation with rigid body coupling. + +#![allow(clippy::too_many_arguments)] +#![allow(clippy::module_inception)] +#![allow(missing_docs)] + +#[cfg(feature = "dim2")] +pub use nexus_mpm_shaders2d as mpm_shaders; +#[cfg(feature = "dim3")] +pub use nexus_mpm_shaders3d as mpm_shaders; + +#[cfg(feature = "dim2")] +pub extern crate nexus_rbd2d as nexus_rbd; +#[cfg(feature = "dim3")] +pub extern crate nexus_rbd3d as nexus_rbd; + +#[cfg(feature = "dim2")] +pub extern crate rapier2d as rapier; +#[cfg(feature = "dim3")] +pub extern crate rapier3d as rapier; + +use khal::re_exports::include_dir::{Dir, include_dir}; + +/// Embedded SPIR-V shader directory. +pub static SPIRV_DIR: Dir<'static> = include_dir!("$OUT_DIR/shaders-spirv"); + +pub mod grid; +pub mod models; +pub mod pipeline; +pub use pipeline::MpmCapacities; +pub(crate) mod sampling; +pub mod solver; +#[cfg(feature = "dim3")] +pub mod trimesh; diff --git a/src_mpm/models/drucker_prager.rs b/src_mpm/models/drucker_prager.rs new file mode 100644 index 00000000..298f8039 --- /dev/null +++ b/src_mpm/models/drucker_prager.rs @@ -0,0 +1,57 @@ +use crate::models::lame_lambda_mu; +use crate::mpm_shaders::models::drucker_prager::DruckerPragerPlasticity; + +/// CPU-side convenience wrapper for constructing `DruckerPragerPlasticity`. +pub struct DruckerPrager; + +impl DruckerPrager { + /// Creates a Drucker-Prager model with default sand parameters. + // Factory constructing the configured plasticity struct; kept as `new` for + // API stability. + #[allow(clippy::new_ret_no_self)] + pub fn new(young_modulus: f32, poisson_ratio: f32) -> DruckerPragerPlasticity { + let (lambda, mu) = if young_modulus > 0.0 { + lame_lambda_mu(young_modulus, poisson_ratio) + } else { + (-1.0, -1.0) + }; + + Self::from_lame(lambda, mu) + } + + /// Friction coefficient `alpha` at the initial hardening state. + /// + /// Mirrors `DruckerPragerPlasticity::alpha` on the GPU, evaluated at the + /// accumulated plastic strain a fresh particle starts with. It is the factor + /// relating a cohesion strain to the shear strength it implies. + pub fn initial_alpha() -> f32 { + let plasticity = Self::from_lame(1.0, 1.0); + let q = 1.0f32; + let angle = + plasticity.ha + (plasticity.hb * q - plasticity.hd) * (-plasticity.hc * q).exp(); + let s = angle.sin(); + (2.0f32 / 3.0).sqrt() * (2.0 * s) / (3.0 - s) + } + + /// Creates a Drucker-Prager model from Lamé parameters with default plasticity settings. + pub fn from_lame(lambda: f32, mu: f32) -> DruckerPragerPlasticity { + Self::from_lame_with_cohesion(lambda, mu, 0.0) + } + + /// Creates a Drucker-Prager model from Lamé parameters and a cohesion. + /// + /// `cohesion` is the volumetric log-strain the material sustains in tension + /// before separating; 0 gives dry sand, and a few 1e-3 already lets a pile + /// stand at a much steeper angle than its friction angle allows. + pub fn from_lame_with_cohesion(lambda: f32, mu: f32, cohesion: f32) -> DruckerPragerPlasticity { + DruckerPragerPlasticity { + ha: 35.0f32.to_radians(), + hb: 9.0f32.to_radians(), + hc: 0.2, + hd: 10.0f32.to_radians(), + lambda, + mu, + cohesion, + } + } +} diff --git a/src_mpm/models/mod.rs b/src_mpm/models/mod.rs new file mode 100644 index 00000000..c7a28080 --- /dev/null +++ b/src_mpm/models/mod.rs @@ -0,0 +1,51 @@ +//! Material constitutive models for MPM particles. +//! +//! This module provides material models that define how particles respond to deformation. +//! The actual model implementations live in the shader crate; this module re-exports +//! them and provides CPU-side convenience constructors. + +pub use crate::mpm_shaders::models::drucker_prager::{ + DruckerPragerPlasticState, DruckerPragerPlasticity, +}; +pub use crate::mpm_shaders::models::fluid::FluidModel; +pub use crate::mpm_shaders::models::linear_elasticity::LinearElasticModel; +pub use crate::mpm_shaders::models::snow::{SnowPlasticState, SnowPlasticity}; + +pub use drucker_prager::DruckerPrager; + +mod drucker_prager; + +/// Computes Lamé parameters (λ, μ) from Young's modulus and Poisson's ratio. +pub(crate) fn lame_lambda_mu(young_modulus: f32, poisson_ratio: f32) -> (f32, f32) { + ( + young_modulus * poisson_ratio / ((1.0 + poisson_ratio) * (1.0 - 2.0 * poisson_ratio)), + shear_modulus(young_modulus, poisson_ratio), + ) +} + +/// Computes shear modulus μ (also called G) from Young's modulus and Poisson's ratio. +fn shear_modulus(young_modulus: f32, poisson_ratio: f32) -> f32 { + young_modulus / (2.0 * (1.0 + poisson_ratio)) +} + +/// Lamé parameters for linear elastic materials. +/// +/// CPU-side convenience type wrapping `LinearElasticModel`. +pub type ElasticCoefficients = LinearElasticModel; + +/// Extension trait for creating `LinearElasticModel` from engineering parameters. +pub trait ElasticCoefficientsExt { + /// Creates elastic coefficients from engineering parameters. + fn from_young_modulus(young_modulus: f32, poisson_ratio: f32) -> Self; +} + +impl ElasticCoefficientsExt for LinearElasticModel { + fn from_young_modulus(young_modulus: f32, poisson_ratio: f32) -> Self { + let (lambda, mu) = lame_lambda_mu(young_modulus, poisson_ratio); + Self { + lambda, + mu, + cfl_coeff: 0.5, + } + } +} diff --git a/src_mpm/pipeline.rs b/src_mpm/pipeline.rs new file mode 100644 index 00000000..ba199b31 --- /dev/null +++ b/src_mpm/pipeline.rs @@ -0,0 +1,558 @@ +//! High-level MPM simulation pipeline orchestration. +//! +//! This module provides the main entry point for running MPM simulations. The pipeline +//! coordinates the execution of all MPM algorithm stages on the GPU. + +use crate::grid::grid::{GpuGrid, WgGrid}; +use crate::grid::sort::WgSort; +use crate::solver::{ + BoundaryCondition, GpuImpulses, GpuMaterials, GpuParticles, GpuRigidParticles, + GpuSimulationParams, GpuTimestepBounds, Particle, SimulationParams, WgG2P, WgG2PCdf, + WgGridUpdate, WgGridUpdateCdf, WgIntegrateBodies, WgP2G, WgP2GCdf, WgParticleUpdate, + WgRigidParticleUpdate, WgTimestepBounds, +}; +use khal::backend::{Backend, Encoder, GpuBackend, GpuBackendError, GpuTimestamps}; +use khal::{BufferUsages, Shader}; +use nexus_rbd::dynamics::GpuBodySet; +use nexus_rbd::math::{Pose, Vector}; +use nexus_rbd::utils::{GpuPrefixSum, PrefixSumWorkspace}; +use vortx::tensor::Tensor; + +use nexus_rbd::dynamics::body::{BodyCoupling, RapierBodyCouplingEntry}; + +/// Initial capacities of an MPM scene's GPU-resident buffers. +/// +/// `NexusState` holds one of these and forwards it on the first particle +/// insertion, when the sub-state is created. +#[derive(Copy, Clone, Debug)] +pub struct MpmCapacities { + /// Number of grid cells reserved for the background grid. + pub grid_size: u32, + /// Number of particles to reserve buffer space for up front, so the initial + /// particle upload (and early emitter growth) doesn't reallocate. + pub particles_capacity: u32, +} + +impl Default for MpmCapacities { + fn default() -> Self { + Self { + grid_size: 32768, + particles_capacity: 65536, + } + } +} + +/// GPU compute pipeline for Material Point Method simulation. +pub struct MpmPipeline { + grid: WgGrid, + prefix_sum: GpuPrefixSum, + sort: WgSort, + p2g: WgP2G, + p2g_cdf: WgP2GCdf, + grid_update_cdf: WgGridUpdateCdf, + grid_update: WgGridUpdate, + particles_update: WgParticleUpdate, + g2p: WgG2P, + g2p_cdf: WgG2PCdf, + rigid_particles_update: WgRigidParticleUpdate, + /// Maximum timestep bound calculation. + pub timestep_bounds: WgTimestepBounds, + /// Rigid body impulse computation kernel (publicly accessible for external use). + pub integrate_bodies: WgIntegrateBodies, +} + +/// GPU-resident simulation state for MPM. +pub struct MpmState { + /// The simulation timestep. + pub base_dt: f32, + pub gravity: Vector, + pub use_cpic: bool, + /// Global simulation parameters (gravity, timestep). + pub sim_params: GpuSimulationParams, + /// Spatial grid for momentum transfer. + pub grid: GpuGrid, + /// MPM particles (positions, velocities, masses, material properties). + pub particles: GpuParticles, + /// Particles sampled from rigid body collider surfaces for two-way coupling. + pub rigid_particles: GpuRigidParticles, + /// Rigid bodies coupled with the MPM simulation. + pub bodies: GpuBodySet, + /// MPM materials associated to each rigid-body. + pub body_materials: GpuMaterials, + /// Accumulated impulses to apply to rigid bodies from MPM interactions. + pub impulses: GpuImpulses, + /// Staging buffer for reading rigid body poses back to CPU. + pub poses_staging: Tensor, + /// For each coupled body, the rigid-body pipeline slot it mirrors. Empty + /// unless the coupling was built through the `NexusState` path, which is + /// the only one that knows the rigid-body slot layout. Consumed by + /// [`MpmPipeline::writeback_body_poses`]. + pub rbd_body_slots: Tensor, + /// The timestep estimate computed from particles and their models. + pub timestep_bounds: Tensor, + /// Staging buffer for reading the timestep bound estimate. + pub timestep_bounds_staging: Tensor, + prefix_sum: PrefixSumWorkspace, + coupling: Vec, +} + +impl MpmState { + /// Creates an empty MPM state with no particles and no coupled bodies. + /// + /// The grid is preallocated to hold `grid_capacity` cells. Physical + /// parameters (`gravity`, `base_dt`, the grid `cell_width`) are left at + /// neutral defaults, so set them before stepping. Particles and coupled + /// bodies are appended lazily through `NexusState`, growing the GPU buffers + /// on demand. + pub fn empty( + backend: &GpuBackend, + capacities: &MpmCapacities, + ) -> Result { + const DEFAULT_CELL_WIDTH: f32 = 1.0; + let grid_capacity = capacities.grid_size; + let params = SimulationParams { + gravity: Vector::ZERO, + #[cfg(feature = "dim2")] + padding: 0.0, + dt: 1.0 / 60.0, + }; + let sim_params = GpuSimulationParams::new(backend, params)?; + let mut particles = GpuParticles::from_particles(backend, &[])?; + // Reserve room up front so the initial particle upload / early emitter + // growth doesn't reallocate the per-particle buffers. + particles.reserve(backend, capacities.particles_capacity as usize)?; + let rigid_particles = GpuRigidParticles::new(backend)?; + let bodies = GpuBodySet::empty(backend); + let body_materials = GpuMaterials::new(backend, &[])?; + let grid = GpuGrid::with_capacity(backend, grid_capacity, DEFAULT_CELL_WIDTH)?; + let prefix_sum = PrefixSumWorkspace::with_capacity(backend, grid_capacity); + let impulses = GpuImpulses::new(backend)?; + let poses_staging = Tensor::vector_uninit( + backend, + bodies.len(), + BufferUsages::COPY_DST | BufferUsages::MAP_READ, + )?; + let bounds = GpuTimestepBounds::default(); + let timestep_bounds = Tensor::scalar( + backend, + bounds, + BufferUsages::STORAGE | BufferUsages::COPY_SRC, + )?; + let timestep_bounds_staging = Tensor::scalar( + backend, + bounds, + BufferUsages::COPY_DST | BufferUsages::MAP_READ, + )?; + + Ok(Self { + base_dt: params.dt, + gravity: params.gravity, + use_cpic: false, + sim_params, + grid, + particles, + rigid_particles, + bodies, + body_materials, + impulses, + poses_staging, + rbd_body_slots: Tensor::vector(backend, [], BufferUsages::STORAGE)?, + timestep_bounds, + timestep_bounds_staging, + prefix_sum, + coupling: Vec::new(), + }) + } + + /// Updates the global simulation parameters (gravity, timestep) and uploads + /// them to the GPU. + pub fn set_simulation_params( + &mut self, + backend: &GpuBackend, + params: SimulationParams, + ) -> Result<(), GpuBackendError> { + self.gravity = params.gravity; + self.base_dt = params.dt; + self.sim_params = GpuSimulationParams::new(backend, params)?; + Ok(()) + } + + /// Uploads the per-substep parameters: the visible timestep `base_dt` divided + /// by `num_substeps`, keeping the current gravity. Cheap (one buffer write), + /// called each frame by the `NexusState` substep loop. + pub fn write_substep_params( + &mut self, + backend: &GpuBackend, + num_substeps: u32, + ) -> Result<(), GpuBackendError> { + let params = SimulationParams { + gravity: self.gravity, + dt: self.base_dt / num_substeps.max(1) as f32, + #[cfg(feature = "dim2")] + padding: 0.0, + }; + backend.write_buffer(self.sim_params.params.buffer_mut(), 0, &[params])?; + Ok(()) + } + + /// Reallocates the background grid with a new cell width (and capacity). Must + /// be called before particles are added, since it discards grid state. + pub fn set_cell_width( + &mut self, + backend: &GpuBackend, + cell_width: f32, + grid_capacity: u32, + ) -> Result<(), GpuBackendError> { + self.grid = GpuGrid::with_capacity(backend, grid_capacity, cell_width)?; + self.prefix_sum = PrefixSumWorkspace::with_capacity(backend, grid_capacity); + Ok(()) + } + + /// (Re)builds the rigid-body coupling: uploads the coupled bodies, samples + /// rigid particles from their collider surfaces, and stores the per-collider + /// boundary materials. + /// + /// `rbd_body_slots[i]` is the rigid-body pipeline slot mirroring coupling + /// entry `i`; it lets [`MpmPipeline::writeback_body_poses`] push the poses + /// MPM integrates back to the buffer rendering reads. + /// + /// Leaves the MPM particles / grid / sim-params untouched. + pub fn set_coupling( + &mut self, + backend: &GpuBackend, + bodies: &rapier::dynamics::RigidBodySet, + colliders: &rapier::geometry::ColliderSet, + coupling: Vec, + materials: &[BoundaryCondition], + rbd_body_slots: &[u32], + cell_width: f32, + ) -> Result<(), GpuBackendError> { + assert_eq!(coupling.len(), materials.len()); + assert_eq!(coupling.len(), rbd_body_slots.len()); + let gpu_bodies = GpuBodySet::from_rapier(backend, bodies, colliders, &coupling); + let rigid_particles = + GpuRigidParticles::from_rapier(backend, colliders, &gpu_bodies, &coupling, cell_width)?; + self.body_materials = GpuMaterials::new(backend, materials)?; + self.poses_staging = Tensor::vector_uninit( + backend, + gpu_bodies.len(), + BufferUsages::COPY_DST | BufferUsages::MAP_READ, + )?; + self.rbd_body_slots = Tensor::vector(backend, rbd_body_slots, BufferUsages::STORAGE)?; + self.use_cpic = !coupling.is_empty(); + self.bodies = gpu_bodies; + self.rigid_particles = rigid_particles; + self.coupling = coupling; + Ok(()) + } +} + +impl MpmState { + /// Creates new MPM simulation data with default two-way coupling for all colliders. + pub fn new( + backend: &GpuBackend, + params: SimulationParams, + particles: &[Particle], + bodies: &rapier::dynamics::RigidBodySet, + colliders: &rapier::geometry::ColliderSet, + materials: &[(rapier::geometry::ColliderHandle, BoundaryCondition)], + cell_width: f32, + grid_capacity: u32, + ) -> Result { + let coupling: Vec<_> = colliders + .iter() + .filter_map(|(co_handle, co)| { + let rb_handle = co.parent()?; + Some(RapierBodyCouplingEntry { + body: rb_handle, + collider: co_handle, + mode: BodyCoupling::OneWay, + }) + }) + .collect(); + let materials: Vec<_> = coupling + .iter() + .map(|c| { + materials + .iter() + .find(|e| e.0 == c.collider) + .map(|e| e.1) + .unwrap_or(BoundaryCondition::separate(1.0)) + }) + .collect(); + Self::with_select_coupling( + backend, + params, + particles, + bodies, + colliders, + coupling, + &materials, + cell_width, + grid_capacity, + ) + } + + /// Creates new MPM simulation data with custom rigid body coupling configuration. + pub fn with_select_coupling( + backend: &GpuBackend, + params: SimulationParams, + particles: &[Particle], + bodies: &rapier::dynamics::RigidBodySet, + colliders: &rapier::geometry::ColliderSet, + coupling: Vec, + materials: &[BoundaryCondition], + cell_width: f32, + grid_capacity: u32, + ) -> Result { + assert_eq!(coupling.len(), materials.len()); + + let sampling_step = cell_width; + let bodies = GpuBodySet::from_rapier(backend, bodies, colliders, &coupling); + let body_materials = GpuMaterials::new(backend, materials)?; + let sim_params = GpuSimulationParams::new(backend, params)?; + let particles = GpuParticles::from_particles(backend, particles)?; + let rigid_particles = + GpuRigidParticles::from_rapier(backend, colliders, &bodies, &coupling, sampling_step)?; + let grid = GpuGrid::with_capacity(backend, grid_capacity, cell_width)?; + let prefix_sum = PrefixSumWorkspace::with_capacity(backend, grid_capacity); + let impulses = GpuImpulses::new(backend)?; + let poses_staging = Tensor::vector_uninit( + backend, + bodies.len(), + BufferUsages::COPY_DST | BufferUsages::MAP_READ, + )?; + let bounds = GpuTimestepBounds::default(); + let timestep_bounds = Tensor::scalar( + backend, + bounds, + BufferUsages::STORAGE | BufferUsages::COPY_SRC, + )?; + let timestep_bounds_staging = Tensor::scalar( + backend, + bounds, + BufferUsages::COPY_DST | BufferUsages::MAP_READ, + )?; + + Ok(Self { + sim_params, + particles, + gravity: params.gravity, + use_cpic: true, + rigid_particles, + bodies, + body_materials, + impulses, + grid, + prefix_sum, + poses_staging, + // Standalone MPM: no rigid-body pipeline to write poses back to. + rbd_body_slots: Tensor::vector(backend, [], BufferUsages::STORAGE)?, + coupling, + timestep_bounds, + timestep_bounds_staging, + base_dt: params.dt, + }) + } + + /// Returns the list of rigid body coupling entries. + pub fn coupling(&self) -> &[RapierBodyCouplingEntry] { + &self.coupling + } +} + +impl MpmPipeline { + /// Creates a new MPM compute pipeline by compiling all necessary shaders. + pub fn new(backend: &GpuBackend) -> Result { + Ok(Self { + grid: WgGrid::from_backend(backend)?, + prefix_sum: GpuPrefixSum::from_backend(backend)?, + sort: WgSort::from_backend(backend)?, + p2g: WgP2G::from_backend(backend)?, + p2g_cdf: WgP2GCdf::from_backend(backend)?, + grid_update: WgGridUpdate::from_backend(backend)?, + grid_update_cdf: WgGridUpdateCdf::from_backend(backend)?, + particles_update: WgParticleUpdate::from_backend(backend)?, + rigid_particles_update: WgRigidParticleUpdate::from_backend(backend)?, + g2p: WgG2P::from_backend(backend)?, + g2p_cdf: WgG2PCdf::from_backend(backend)?, + integrate_bodies: WgIntegrateBodies::from_backend(backend)?, + timestep_bounds: WgTimestepBounds::from_backend(backend)?, + }) + } + + /// Executes one complete MPM simulation timestep. + pub fn step( + &self, + backend: &GpuBackend, + data: &mut MpmState, + mut timestamps: Option<&mut GpuTimestamps>, + ) -> Result<(), GpuBackendError> { + let mut encoder = backend.begin_encoding(); + + { + let mut pass = encoder.begin_pass("[MPM] Rigid update", timestamps.as_deref_mut()); + self.integrate_bodies.launch_update_world_mass_properties( + &mut pass, + &mut data.impulses, + &mut data.bodies, + )?; + self.rigid_particles_update.launch( + &mut pass, + &mut data.bodies, + &mut data.rigid_particles, + )?; + } + + { + let mut pass = encoder.begin_pass("[MPM] Grid sort", timestamps.as_deref_mut()); + data.grid.swap_buffers(); + self.grid.launch_sort( + backend, + &mut pass, + &mut data.particles, + data.use_cpic.then_some(&mut data.rigid_particles), + &mut data.grid, + &mut data.prefix_sum, + &self.sort, + &self.prefix_sum, + )?; + + if data.use_cpic { + self.sort.launch_sort_rigid_particles( + backend, + &mut pass, + &mut data.rigid_particles, + &mut data.grid, + &mut data.prefix_sum, + &self.prefix_sum, + )?; + } + } + + if data.use_cpic { + { + let mut pass = + encoder.begin_pass("[MPM] CDF grid update", timestamps.as_deref_mut()); + self.grid_update_cdf + .launch(&mut pass, &mut data.grid, &data.bodies)?; + } + + { + let mut pass = encoder.begin_pass("[MPM] CDF P2G", timestamps.as_deref_mut()); + self.p2g_cdf.launch( + &mut pass, + &mut data.grid, + &data.rigid_particles, + &data.bodies, + )?; + } + + { + let mut pass = encoder.begin_pass("[MPM] CDF G2P", timestamps.as_deref_mut()); + self.g2p_cdf.launch( + &mut pass, + &data.sim_params, + &data.grid, + &mut data.particles, + )?; + } + } + + { + let mut pass = encoder.begin_pass("[MPM] P2G", timestamps.as_deref_mut()); + self.p2g.launch( + &mut pass, + data.use_cpic, + &mut data.grid, + &data.particles, + &mut data.impulses, + &data.bodies, + &data.body_materials, + )?; + } + + { + let mut pass = encoder.begin_pass("[MPM] Grid update", timestamps.as_deref_mut()); + self.grid_update.launch( + &mut pass, + data.use_cpic, + &data.sim_params, + &mut data.grid, + &data.bodies, + &data.body_materials, + )?; + } + + { + let mut pass = encoder.begin_pass("[MPM] G2P", timestamps.as_deref_mut()); + self.g2p.launch( + &mut pass, + data.use_cpic, + &data.sim_params, + &data.grid, + &mut data.particles, + &data.bodies, + &data.body_materials, + )?; + } + + { + let mut pass = encoder.begin_pass("[MPM] Particle update", timestamps.as_deref_mut()); + self.particles_update.launch( + &mut pass, + &data.sim_params, + &data.grid, + &mut data.particles, + )?; + } + + { + let mut pass = encoder.begin_pass("[MPM] Integrate bodies", timestamps.as_deref_mut()); + self.integrate_bodies.launch( + &mut pass, + &data.grid, + &data.sim_params, + &mut data.impulses, + &mut data.bodies, + )?; + } + + if let Some(timestamps) = timestamps { + timestamps.resolve(&mut encoder); + } + + backend.submit(encoder) + } + + /// Pushes the coupled bodies' MPM-integrated poses into `rbd_poses`, the + /// rigid-body pipeline's body-pose buffer. + /// + /// MPM integrates its own copy of every coupled body (see + /// [`WgIntegrateBodies::launch`]) while the rigid-body pipeline treats them + /// as static, so the two copies diverge as soon as a body moves. Call this + /// once per visible frame, after the substep loop, so rendering and the next + /// step's broad phase see where the body really is. + pub fn writeback_body_poses( + &self, + backend: &GpuBackend, + data: &MpmState, + rbd_poses: &mut Tensor, + ) -> Result<(), GpuBackendError> { + if data.bodies.is_empty() || data.rbd_body_slots.is_empty() { + return Ok(()); + } + + let mut encoder = backend.begin_encoding(); + { + let mut pass = encoder.begin_pass("[MPM] Body pose writeback", None); + self.integrate_bodies.launch_writeback_body_poses( + &mut pass, + &data.bodies, + &data.rbd_body_slots, + rbd_poses, + )?; + } + backend.submit(encoder) + } +} diff --git a/src_mpm/sampling/mod.rs b/src_mpm/sampling/mod.rs new file mode 100644 index 00000000..4032e2e5 --- /dev/null +++ b/src_mpm/sampling/mod.rs @@ -0,0 +1,15 @@ +//! Surface sampling for rigid body coupling. +//! +//! Samples particles on the surfaces of rigid body colliders for two-way +//! MPM-rigid body coupling. In 2D, samples polyline edges; in 3D, samples +//! triangle mesh surfaces. + +#[cfg(feature = "dim2")] +pub use sample_polyline::*; +#[cfg(feature = "dim3")] +pub use sample_trimesh::*; + +#[cfg(feature = "dim2")] +mod sample_polyline; +#[cfg(feature = "dim3")] +mod sample_trimesh; diff --git a/src_mpm/sampling/sample_polyline.rs b/src_mpm/sampling/sample_polyline.rs new file mode 100644 index 00000000..3f8ae97b --- /dev/null +++ b/src_mpm/sampling/sample_polyline.rs @@ -0,0 +1,57 @@ +use crate::mpm_shaders::solver::particle::{Position, RigidParticleIndices}; +use glamx::UVec2; +use rapier::geometry::{Polyline, Segment}; + +/// Type alias for backward compatibility. +pub type GpuSampleIds = RigidParticleIndices; + +#[derive(Copy, Clone, Debug)] +#[repr(C)] +pub struct SamplingParams { + pub base_vid: u32, + pub collider_id: u32, + pub sampling_step: f32, +} + +#[derive(Default, Clone)] +pub struct SamplingBuffers { + pub samples: Vec, + pub samples_ids: Vec, +} + +pub fn sample_polyline( + polyline: &Polyline, + params: &SamplingParams, + buffers: &mut SamplingBuffers, +) { + for seg_idx in polyline.indices() { + let seg = Segment::new( + polyline.vertices()[seg_idx[0] as usize], + polyline.vertices()[seg_idx[1] as usize], + ); + let sample_id = GpuSampleIds { + segment: UVec2::new(params.base_vid + seg_idx[0], params.base_vid + seg_idx[1]), + collider: params.collider_id, + _pad: 0, + }; + buffers.samples.push(Position { pt: seg.a }); + buffers.samples_ids.push(sample_id); + + if let Some(dir) = seg.direction() { + for i in 0.. { + let shift = (i as f32) * params.sampling_step; + if shift > seg.length() { + break; + } + + buffers.samples.push(Position { + pt: seg.a + dir * shift, + }); + buffers.samples_ids.push(sample_id); + } + + buffers.samples.push(Position { pt: seg.b }); + buffers.samples_ids.push(sample_id); + } + } +} diff --git a/src_mpm/sampling/sample_trimesh.rs b/src_mpm/sampling/sample_trimesh.rs new file mode 100644 index 00000000..a3cda82e --- /dev/null +++ b/src_mpm/sampling/sample_trimesh.rs @@ -0,0 +1,231 @@ +use crate::mpm_shaders::solver::particle::{Position, RigidParticleIndices}; +use glamx::UVec3; +use nexus_rbd::math::Vector; +use rapier::geometry::{Segment, TriMesh, Triangle}; +use std::collections::HashSet; + +/// Type alias for backward compatibility. +pub type GpuSampleIds = RigidParticleIndices; + +// Epsilon used as a length threshold in various steps of the sampling. In particular, this avoids +// degenerate geometries from generating invalid samples. +const EPS: f32 = 1.0e-5; + +pub struct TriangleSample { + pub triangle_id: u32, + pub point: Vector, +} + +#[derive(Copy, Clone, Debug)] +pub struct SamplingParams { + pub base_vid: u32, + pub collider_id: u32, + pub sampling_step: f32, +} + +#[derive(Default, Clone)] +pub struct SamplingBuffers { + pub samples: Vec, + pub samples_ids: Vec, +} + +pub fn sample_trimesh(trimesh: &TriMesh, params: &SamplingParams, buffers: &mut SamplingBuffers) { + let samples = sample_mesh(trimesh.vertices(), trimesh.indices(), params.sampling_step); + + for sample in samples { + let tri_idx = trimesh.indices()[sample.triangle_id as usize]; + let sample_id = GpuSampleIds { + triangle: UVec3::new( + params.base_vid + tri_idx[0], + params.base_vid + tri_idx[1], + params.base_vid + tri_idx[2], + ), + collider: params.collider_id, + }; + buffers.samples.push(Position { + pt: sample.point, + padding: 0, + }); + buffers.samples_ids.push(sample_id); + } + + println!( + "Num rigid particles: {}, num triangles: {}", + buffers.samples.len(), + trimesh.indices().len() + ); +} + +/// Samples a triangle mesh with a set of points such that at least one point is generated +/// inside each cell on a grid on the x-y plane with cells sized by `xy_spacing`. +pub fn sample_mesh( + vertices: &[Vector], + indices: &[[u32; 3]], + xy_spacing: f32, +) -> Vec { + let mut samples = vec![]; + // TODO: switch to a matrix of boolean to avoid hashing if + // this proves to be a perf bottleneck. + let mut visited_segs = HashSet::new(); + + let mut seg_needs_sampling = |mut ia: u32, mut ib: u32| { + if ib > ia { + std::mem::swap(&mut ia, &mut ib); + } + + visited_segs.insert([ia, ib]) + }; + + for (tri_id, idx) in indices.iter().enumerate() { + let tri = Triangle::new( + vertices[idx[0] as usize], + vertices[idx[1] as usize], + vertices[idx[2] as usize], + ); + sample_triangle(tri, &mut samples, xy_spacing, tri_id as u32); + + if seg_needs_sampling(idx[0], idx[1]) { + let seg = Segment::new(vertices[idx[0] as usize], vertices[idx[1] as usize]); + sample_edge(seg, &mut samples, xy_spacing, tri_id as u32); + } + + if seg_needs_sampling(idx[1], idx[2]) { + let seg = Segment::new(vertices[idx[1] as usize], vertices[idx[2] as usize]); + sample_edge(seg, &mut samples, xy_spacing, tri_id as u32); + } + + if seg_needs_sampling(idx[2], idx[0]) { + let seg = Segment::new(vertices[idx[2] as usize], vertices[idx[0] as usize]); + sample_edge(seg, &mut samples, xy_spacing, tri_id as u32); + } + } + + samples +} + +/// Samples a triangle edge with a set of points such that at least one point is generated +/// inside each cell on a grid on the x-y plane with cells sized by `xy_spacing`. +/// +/// The returned samples will not contain `edge.a`. It might contain `edge.b` (but it is unlikely) +/// if it aligns exactly with the internal sampling spacing. +pub fn sample_edge( + edge: Segment, + samples: &mut Vec, + xy_spacing: f32, + triangle_id: u32, +) { + let ab = edge.b - edge.a; + let edge_length = ab.length(); + + if edge_length > EPS { + let edge_dir = ab / edge_length; + let spacing = xy_spacing / 2.0f32.sqrt(); + let nsteps = (edge_length / spacing).ceil() as usize; + + // Start at one so we don't push edge.a. + for i in 1..nsteps { + let point = edge.a + edge_dir * (spacing * i as f32); + samples.push(TriangleSample { point, triangle_id }) + } + } +} + +/// Samples a triangle with a set of points such that at least one point is generated +/// inside each cell on a grid on the x-y plane with cells sized by `xy_spacing`. +/// +/// No sample is placed on the base or any of the triangle vertices. Because this +/// does not attempt to sample the edges of the triangles, small or thin triangles +/// might not result in any samples. Edges should be sampled separately with [`sample_edge`]. +pub fn sample_triangle( + triangle: Triangle, + samples: &mut Vec, + xy_spacing: f32, + triangle_id: u32, +) { + // select the longest edge as the base + let distance_ab = triangle.b.distance(triangle.a); + let distance_bc = triangle.c.distance(triangle.b); + let distance_ca = triangle.a.distance(triangle.c); + let max = distance_ab.max(distance_bc).max(distance_ca); + + let triangle = if max == distance_bc { + Triangle { + a: triangle.b, + b: triangle.c, + c: triangle.a, + } + } else if max == distance_ca { + Triangle { + a: triangle.c, + b: triangle.a, + c: triangle.b, + } + } else { + triangle + }; + + let ac = triangle.c - triangle.a; + let base = triangle.b - triangle.a; + let base_length = base.length(); + let base_dir = base / base_length; + + // Adjust the spacing so it matches the required spacing on the x-y plane. + // For simplicity, we just divide by sqrt(2) so that the spacing in any direction is guaranteed + // to be smaller or equal to the inner-circle diameter of any cell from the implicit grid with + // spacing `xy_spacing`. + // We could use a more fine-grained adjustment that depends on the angle between the base-dir + // and the world x-y axes. But this doesn't make a significant difference in point count or + // computation times. However, the sampling looks worse (less uniform in practice). So we stick + // to the simple sqrt(2) approach. + let spacing = xy_spacing / 2.0f32.sqrt(); + + // Calculate the step increment on the base. + let base_step_count = (base_length / spacing).ceil(); + let base_step = base_dir * spacing; + + // Project C on the base AB. + let ac_offset_length = ac.dot(base_dir); + let bc_offset_length = base_length - ac_offset_length; + + if ac_offset_length < EPS || bc_offset_length < EPS || base_length < EPS { + return; + } + + // Compute the triangle's height vector. + let height = ac - base_dir * ac_offset_length; + let height_length = height.length(); + let height_dir = height / height_length; + // Calculate the tangents. + let tan_alpha = height_length / ac_offset_length; + let tan_beta = height_length / bc_offset_length; + + // Start at 1 so we don't sample the perpendicular edge if it's at a right angle + // with `triangle.a`. + for i in 1..base_step_count as u32 { + let base_position = triangle.a + (i as f32) * base_step; + + // Compute the height at the current base_position. The point at the + // end of that height is either in the line (AC) or (BC), whichever is closer. + let height_ac = tan_alpha * triangle.a.distance(base_position); + let height_bc = tan_beta * triangle.b.distance(base_position); + let height_length = height_ac.min(height_bc); + + // Calculate the step increment on the height. + let height_step_count = (height_length / spacing).ceil(); + let height_step = height_dir * spacing; + + // Start at 1 so we don't sample the basis edge. + for j in 1..height_step_count as u32 { + let particle_position = base_position + (j as f32) * height_step; + + if !particle_position.is_finite() { + continue; + } + + samples.push(TriangleSample { + point: particle_position, + triangle_id, + }); + } + } +} diff --git a/src_mpm/solver/boundary_condition.rs b/src_mpm/solver/boundary_condition.rs new file mode 100644 index 00000000..a72c9bb9 --- /dev/null +++ b/src_mpm/solver/boundary_condition.rs @@ -0,0 +1,39 @@ +pub use crate::mpm_shaders::solver::boundary_condition::{ + BodyMaterials, BoundaryCondition, MAX_COLLISION_BODIES, +}; +use khal::BufferUsages; +use khal::backend::{GpuBackend, GpuBackendError}; +use vortx::tensor::Tensor; + +/// GPU buffer storing the per-rigid-body boundary conditions. +/// +/// Held as a single **uniform** `BodyMaterials` (a fixed `MAX_COLLISION_BODIES` +/// array) rather than a storage buffer, so the MPM kernels that read it stay +/// within the 8-storage-buffer WebGPU limit. +pub struct GpuMaterials { + pub materials: Tensor, +} + +impl GpuMaterials { + /// Creates the boundary-condition uniform buffer. + /// + /// Allocates space for up to `MAX_COLLISION_BODIES` bodies (CPIC limitation). + pub fn new( + backend: &GpuBackend, + materials: &[BoundaryCondition], + ) -> Result { + assert!( + materials.len() <= MAX_COLLISION_BODIES, + "CPIC only supports up to {MAX_COLLISION_BODIES} colliders" + ); + let mut mats = [BoundaryCondition::default(); MAX_COLLISION_BODIES]; + mats[..materials.len()].copy_from_slice(materials); + Ok(Self { + materials: Tensor::scalar( + backend, + BodyMaterials { mats }, + BufferUsages::UNIFORM | BufferUsages::COPY_DST, + )?, + }) + } +} diff --git a/src_mpm/solver/g2p.rs b/src_mpm/solver/g2p.rs new file mode 100644 index 00000000..e74d22ba --- /dev/null +++ b/src_mpm/solver/g2p.rs @@ -0,0 +1,70 @@ +//! Grid-to-Particle (G2P) transfer kernel. +//! +//! Interpolates grid velocities back to particles and updates particle velocity +//! gradients. This happens after grid forces have been applied. + +use crate::grid::grid::{GpuGrid, indirect_dispatch_tensor}; +use crate::mpm_shaders::solver::g2p::{GpuG2p, GpuG2pCpic}; +use crate::solver::{GpuMaterials, GpuParticles, GpuSimulationParams}; +use khal::Shader; +use khal::backend::{GpuBackendError, GpuPass}; +use nexus_rbd::dynamics::GpuBodySet; + +/// GPU compute kernel for Grid-to-Particle (G2P) velocity interpolation. +/// +/// Samples grid velocities at particle positions using quadratic B-spline weights +/// and updates particle velocity gradients for deformation tracking (APIC method). +#[derive(Shader)] +pub struct WgG2P { + /// Compiled G2P compute shader. + g2p: GpuG2p, + g2p_cpic: GpuG2pCpic, +} + +impl WgG2P { + /// Launches the G2P kernel to update particle velocities from grid. + pub fn launch( + &self, + pass: &mut GpuPass, + use_cpic: bool, + sim_params: &GpuSimulationParams, + grid: &GpuGrid, + particles: &mut GpuParticles, + bodies: &GpuBodySet, + body_materials: &GpuMaterials, + ) -> Result<(), GpuBackendError> { + if use_cpic { + self.g2p_cpic.call( + pass, + indirect_dispatch_tensor(&grid.indirect_n_g2p_p2g_groups), + &sim_params.params, + &grid.meta, + &grid.hmap_entries, + &grid.active_blocks, + &grid.nodes, + &particles.sorted_ids, + &particles.positions, + &mut particles.kinematics, + &bodies.vels, + &bodies.mprops, + &body_materials.materials, + ) + } else { + self.g2p.call( + pass, + indirect_dispatch_tensor(&grid.indirect_n_g2p_p2g_groups), + &sim_params.params, + &grid.meta, + &grid.hmap_entries, + &grid.active_blocks, + &grid.nodes, + &particles.sorted_ids, + &particles.positions, + &mut particles.kinematics, + &bodies.vels, + &bodies.mprops, + &body_materials.materials, + ) + } + } +} diff --git a/src_mpm/solver/g2p_cdf.rs b/src_mpm/solver/g2p_cdf.rs new file mode 100644 index 00000000..68b63e32 --- /dev/null +++ b/src_mpm/solver/g2p_cdf.rs @@ -0,0 +1,41 @@ +//! Grid-to-Particle transfer with Collision Detection Field updates. + +use crate::grid::grid::{GpuGrid, indirect_dispatch_tensor}; +use crate::mpm_shaders::solver::g2p_cdf::GpuG2pCdf; +use crate::solver::{GpuParticles, GpuSimulationParams}; +use khal::Shader; +use khal::backend::{GpuBackendError, GpuPass}; + +/// GPU kernel for G2P transfer with CDF updates for rigid body coupling. +/// +/// Updates particle CDF (Collision Detection Field) data based on proximity +/// to rigid bodies during the G2P phase. +#[derive(Shader)] +pub struct WgG2PCdf { + /// Compiled G2P-CDF compute shader. + g2p_cdf: GpuG2pCdf, +} + +impl WgG2PCdf { + /// Launches G2P with CDF updates for MPM particles. + pub fn launch( + &self, + pass: &mut GpuPass, + sim_params: &GpuSimulationParams, + grid: &GpuGrid, + particles: &mut GpuParticles, + ) -> Result<(), GpuBackendError> { + self.g2p_cdf.call( + pass, + indirect_dispatch_tensor(&grid.indirect_n_g2p_p2g_groups), + &sim_params.params, + &grid.meta, + &grid.hmap_entries, + &grid.active_blocks, + &grid.nodes, + &particles.sorted_ids, + &particles.positions, + &mut particles.kinematics, + ) + } +} diff --git a/src_mpm/solver/grid_update.rs b/src_mpm/solver/grid_update.rs new file mode 100644 index 00000000..d62c7896 --- /dev/null +++ b/src_mpm/solver/grid_update.rs @@ -0,0 +1,68 @@ +//! Grid node update kernel. +//! +//! Updates grid node velocities by applying forces (gravity, boundary conditions) +//! and solving momentum equations on the grid. + +use crate::grid::grid::{GpuGrid, indirect_dispatch_tensor}; +use crate::mpm_shaders::solver::grid_update::GpuGridUpdate; +use crate::mpm_shaders::solver::grid_update_collide::GpuGridUpdateCollide; +use crate::solver::{GpuMaterials, GpuSimulationParams}; +use khal::Shader; +use khal::backend::{GpuBackendError, GpuPass}; +use nexus_rbd::dynamics::GpuBodySet; + +/// GPU compute kernel for updating grid node velocities. +/// +/// Applies external forces (gravity), boundary conditions (sticky/slip walls), +/// and solves momentum equations on grid nodes. Runs between P2G and G2P stages. +#[derive(Shader)] +pub struct WgGridUpdate { + /// Compiled grid update compute shader. + grid_update: GpuGridUpdate, + grid_update_collide: GpuGridUpdateCollide, +} + +impl WgGridUpdate { + /// Launches the grid update kernel. + pub fn launch( + &self, + pass: &mut GpuPass, + use_cpic: bool, + sim_params: &GpuSimulationParams, + grid: &mut GpuGrid, + bodies: &GpuBodySet, + body_materials: &GpuMaterials, + ) -> Result<(), GpuBackendError> { + self.grid_update.call( + pass, + indirect_dispatch_tensor(&grid.indirect_n_g2p_p2g_groups), + &sim_params.params, + &grid.meta, + &grid.active_blocks, + &mut grid.nodes, + )?; + + if !use_cpic { + self.grid_update_collide.call( + pass, + indirect_dispatch_tensor(&grid.indirect_n_g2p_p2g_groups), + &sim_params.params, + &grid.meta, + &grid.active_blocks, + &bodies.shapes, + &bodies.poses, + // The trimesh path projects a shape-local point, so it needs the + // local-space vertices (with a valid interleaved BVH), not the + // world-space buffer maintained for `p2g_cdf`. + &bodies.shapes_local_vertex_buffers, + &bodies.shapes_index_buffers, + &bodies.vels, + &bodies.mprops, + &body_materials.materials, + &mut grid.nodes, + )?; + } + + Ok(()) + } +} diff --git a/src_mpm/solver/grid_update_cdf.rs b/src_mpm/solver/grid_update_cdf.rs new file mode 100644 index 00000000..34dc8840 --- /dev/null +++ b/src_mpm/solver/grid_update_cdf.rs @@ -0,0 +1,41 @@ +//! Grid CDF (Collision Detection Field) update for rigid body coupling. + +use crate::grid::grid::{GpuGrid, indirect_dispatch_tensor}; +use crate::mpm_shaders::solver::grid_update_cdf::GpuGridUpdateCdf; +use khal::Shader; +use khal::backend::{GpuBackendError, GpuPass}; +use nexus_rbd::dynamics::GpuBodySet; + +/// GPU kernel for updating grid node CDF data from rigid bodies. +/// +/// Computes signed distance fields and closest points on rigid body surfaces +/// for each active grid node. +#[derive(Shader)] +pub struct WgGridUpdateCdf { + /// Compiled grid CDF update shader. + grid_update: GpuGridUpdateCdf, +} + +impl WgGridUpdateCdf { + /// Launches grid CDF update from rigid body geometries. + pub fn launch( + &self, + pass: &mut GpuPass, + grid: &mut GpuGrid, + bodies: &GpuBodySet, + ) -> Result<(), GpuBackendError> { + if bodies.is_empty() { + return Ok(()); + } + + self.grid_update.call( + pass, + indirect_dispatch_tensor(&grid.indirect_n_g2p_p2g_groups), + &grid.meta, + &grid.active_blocks, + &bodies.shapes, + &bodies.poses, + &mut grid.nodes, + ) + } +} diff --git a/src_mpm/solver/mod.rs b/src_mpm/solver/mod.rs new file mode 100644 index 00000000..eaae03a9 --- /dev/null +++ b/src_mpm/solver/mod.rs @@ -0,0 +1,35 @@ +//! Core MPM solver algorithms and GPU kernels. + +pub use boundary_condition::{BoundaryCondition, GpuMaterials}; +pub use g2p::WgG2P; +pub use g2p_cdf::WgG2PCdf; +pub use grid_update::WgGridUpdate; +pub use grid_update_cdf::WgGridUpdateCdf; +pub use p2g::WgP2G; +pub use p2g_cdf::WgP2GCdf; +pub use params::{GpuSimulationParams, SimulationParams}; +pub use particle::*; +pub use particle_model::*; +pub use particle_update::WgParticleUpdate; +pub use rigid_integrate::{GpuImpulses, WgIntegrateBodies}; +pub use rigid_particle_update::WgRigidParticleUpdate; +pub use timestep_bound::WgTimestepBounds; + +pub use crate::mpm_shaders::solver::p2g::IntegerImpulse; +pub use crate::mpm_shaders::solver::timestep_bound::GpuTimestepBounds; + +mod boundary_condition; +mod g2p; +mod g2p_cdf; +mod grid_update; +mod grid_update_cdf; +mod p2g; +mod p2g_cdf; +mod params; +mod particle; +mod particle_model; +mod particle_update; +pub mod prep_readback; +mod rigid_integrate; +mod rigid_particle_update; +mod timestep_bound; diff --git a/src_mpm/solver/p2g.rs b/src_mpm/solver/p2g.rs new file mode 100644 index 00000000..37de003b --- /dev/null +++ b/src_mpm/solver/p2g.rs @@ -0,0 +1,66 @@ +//! Particle-to-Grid (P2G) transfer kernel. +//! +//! Transfers particle mass, momentum, and forces to nearby grid nodes using +//! interpolation weights. This is the first major step of each MPM timestep. + +use crate::grid::grid::{GpuGrid, indirect_dispatch_tensor}; +use crate::mpm_shaders::solver::p2g::{GpuP2g, GpuP2gCpic}; +use crate::solver::{GpuImpulses, GpuMaterials, GpuParticles}; +use khal::Shader; +use khal::backend::{GpuBackendError, GpuPass}; +use nexus_rbd::dynamics::GpuBodySet; + +/// GPU compute kernel for Particle-to-Grid (P2G) momentum transfer. +/// +/// Rasterizes particle mass and momentum onto the background grid using quadratic +/// B-spline interpolation. Also handles impulse accumulation for rigid body coupling. +#[derive(Shader)] +pub struct WgP2G { + /// Compiled P2G compute shader. + p2g: GpuP2g, + /// Compiled P2G compute shader with CPIC enabled. + p2g_cpic: GpuP2gCpic, +} + +impl WgP2G { + /// Launches the P2G kernel to transfer particle data to grid nodes. + pub fn launch( + &self, + pass: &mut GpuPass, + use_cpic: bool, + grid: &mut GpuGrid, + particles: &GpuParticles, + impulses: &mut GpuImpulses, + bodies: &GpuBodySet, + body_materials: &GpuMaterials, + ) -> Result<(), GpuBackendError> { + // Scatter-style P2G: one workgroup per active block, one thread per grid node, + // streaming the block's particles (primaries + extras) from `sorted_particle_ids`. + if use_cpic { + self.p2g_cpic.call( + pass, + indirect_dispatch_tensor(&grid.indirect_n_g2p_p2g_groups), + &grid.meta, + &grid.active_blocks, + &particles.sorted_ids, + particles.positions(), + particles.kinematics(), + &mut grid.nodes, + &bodies.vels, + &body_materials.materials, + &mut impulses.incremental_impulses, + ) + } else { + self.p2g.call( + pass, + indirect_dispatch_tensor(&grid.indirect_n_g2p_p2g_groups), + &grid.meta, + &grid.active_blocks, + &particles.sorted_ids, + particles.positions(), + particles.kinematics(), + &mut grid.nodes, + ) + } + } +} diff --git a/src_mpm/solver/p2g_cdf.rs b/src_mpm/solver/p2g_cdf.rs new file mode 100644 index 00000000..61181a6d --- /dev/null +++ b/src_mpm/solver/p2g_cdf.rs @@ -0,0 +1,45 @@ +//! Particle-to-Grid transfer with Collision Detection Field for rigid bodies. + +use crate::grid::grid::{GpuGrid, indirect_dispatch_tensor}; +use crate::mpm_shaders::solver::p2g_cdf::GpuP2gCdf; +use crate::solver::GpuRigidParticles; +use khal::Shader; +use khal::backend::{GpuBackendError, GpuPass}; +use nexus_rbd::dynamics::GpuBodySet; + +/// GPU kernel for P2G transfer from rigid body particles. +/// +/// Transfers momentum from rigid body surface particles to grid nodes, +/// enabling two-way coupling between MPM and rigid bodies. +#[derive(Shader)] +pub struct WgP2GCdf { + /// Compiled P2G-CDF compute shader. + p2g_cdf: GpuP2gCdf, +} + +impl WgP2GCdf { + /// Launches P2G transfer from rigid body particles to grid. + pub fn launch( + &self, + pass: &mut GpuPass, + grid: &mut GpuGrid, + rigid_particles: &GpuRigidParticles, + bodies: &GpuBodySet, + ) -> Result<(), GpuBackendError> { + if rigid_particles.is_empty() { + return Ok(()); + } + + self.p2g_cdf.call( + pass, + indirect_dispatch_tensor(&grid.indirect_n_g2p_p2g_groups), + &grid.meta, + &grid.active_blocks, + &rigid_particles.sorted_ids, + &rigid_particles.sample_points, + &bodies.shapes_vertex_buffers, + &rigid_particles.sample_ids, + &mut grid.nodes, + ) + } +} diff --git a/src_mpm/solver/params.rs b/src_mpm/solver/params.rs new file mode 100644 index 00000000..622cce0f --- /dev/null +++ b/src_mpm/solver/params.rs @@ -0,0 +1,23 @@ +pub use crate::mpm_shaders::solver::params::SimulationParams; +use khal::BufferUsages; +use khal::backend::{GpuBackend, GpuBackendError}; +use vortx::tensor::Tensor; + +/// GPU-resident simulation parameters. +pub struct GpuSimulationParams { + /// Uniform buffer containing simulation parameters. + pub params: Tensor, +} + +impl GpuSimulationParams { + /// Uploads simulation parameters to GPU memory. + pub fn new(backend: &GpuBackend, params: SimulationParams) -> Result { + Ok(Self { + params: Tensor::scalar( + backend, + params, + BufferUsages::UNIFORM | BufferUsages::COPY_DST, + )?, + }) + } +} diff --git a/src_mpm/solver/particle.rs b/src_mpm/solver/particle.rs new file mode 100644 index 00000000..8c4ff6f4 --- /dev/null +++ b/src_mpm/solver/particle.rs @@ -0,0 +1,648 @@ +use crate::mpm_shaders::solver::particle::{ + Kinematics, ParticleProperties, Position, RigidParticleIndices, +}; +use crate::mpm_shaders::{PaddedMatrix, PaddingExt}; +use khal::BufferUsages; +use khal::backend::{Backend, Encoder, GpuBackend, GpuBackendError}; +use nexus_rbd::dynamics::GpuBodySet; +use nexus_rbd::math::{DIM, Matrix, Vector}; +use std::ops::RangeBounds; +use vortx::tensor::Tensor; + +use crate::solver::{GpuParticleModel, ParticleModel}; +use { + crate::sampling::{self, SamplingBuffers, SamplingParams}, + nexus_rbd::dynamics::body::RapierBodyCouplingEntry, +}; + +/// Particle position type used on the GPU. +/// +/// In 2D: `Position` contains a Vec2. +/// In 3D: `Position` contains a Vec3. +pub type ParticlePosition = Position; + +/// A single MPM particle with position, dynamics, and material model. +#[derive(Copy, Clone, Debug)] +pub struct Particle { + /// Spatial position. + pub position: Vector, + /// Physical state (velocity, deformation, mass, etc.). + pub dynamics: ParticleDynamics, + /// Material model defining constitutive behavior. + pub model: ParticleModel, +} + +impl Particle { + /// Creates a new particle with the given properties. + pub fn new(position: Vector, radius: f32, density: f32, model: ParticleModel) -> Self { + Particle { + position, + dynamics: ParticleDynamics::new(radius, density), + model, + } + } + + /// Creates a new particle belonging to the given render group. + /// + /// See [`ParticleDynamics::group_id`]. + pub fn with_group( + position: Vector, + radius: f32, + density: f32, + model: ParticleModel, + group_id: u32, + ) -> Self { + let mut result = Self::new(position, radius, density, model); + result.dynamics.group_id = group_id; + result + } +} + +/// CPU-side particle dynamics for initialization. +/// +/// Splits into GPU `Kinematics`, `Cdf`, deformation gradient, and `ParticleProperties` buffers on upload. +#[derive(Copy, Clone, Debug, PartialEq)] +pub struct ParticleDynamics { + /// Current velocity. + pub velocity: Vector, + /// Deformation gradient. + pub def_grad: Matrix, + /// APIC affine velocity matrix. + pub affine: Matrix, + /// Additional force * dt. + pub force_dt: Vector, + /// Determinant of velocity gradient. + pub vel_grad_det: f32, + /// Collision detection field. + pub cdf: crate::mpm_shaders::solver::particle::Cdf, + /// Initial particle volume. + pub init_volume: f32, + /// Initial particle radius. + pub init_radius: f32, + /// Particle mass. + pub mass: f32, + /// Damping coefficient. + pub damping: f32, + /// Particle phase. + pub phase: f32, + /// Multiplier on a collider's friction for this particle. See + /// [`Kinematics::boundary_friction`](crate::mpm_shaders::solver::particle::Kinematics::boundary_friction). + pub boundary_friction: f32, + /// Group this particle belongs to. Carries no physics: the viewer looks the + /// group up in its color palette to shade the particle. + pub group_id: u32, + /// Whether this particle is active. + pub enabled: u32, + /// Whether this particle is fixed. + pub fixed: u32, +} + +impl ParticleDynamics { + /// Creates new particle dynamics from radius and material density. + pub fn new(radius: f32, density: f32) -> Self { + let exponent = if cfg!(feature = "dim2") { 2 } else { 3 }; + let init_volume = (radius * 2.0).powi(exponent); + Self { + velocity: Vector::ZERO, + def_grad: Matrix::IDENTITY, + affine: Matrix::ZERO, + force_dt: Vector::ZERO, + vel_grad_det: 0.0, + init_volume, + init_radius: radius, + mass: init_volume * density, + damping: 0.0, + cdf: crate::mpm_shaders::solver::particle::Cdf::zero(), + phase: 1.0, + boundary_friction: 1.0, + group_id: 0, + enabled: 1, + fixed: 0, + } + } + + /// Sets whether this particle is fixed. + pub fn set_fixed(&mut self, fixed: bool) { + self.fixed = fixed as u32; + } + + /// Sets how much of a collider's friction this particle feels. + /// + /// 1 (the default) takes the collider's friction as given; 0 lets the + /// particle slide freely along every surface without passing through it, + /// which is what a fluid sharing a floor with granular material wants. + /// + /// Requires CPIC coupling, which is on by default: see + /// [`Kinematics::boundary_friction`](crate::mpm_shaders::solver::particle::Kinematics::boundary_friction). + pub fn set_boundary_friction(&mut self, factor: f32) { + self.boundary_friction = factor; + } + + /// Sets the group this particle belongs to. + pub fn set_group_id(&mut self, group_id: u32) { + self.group_id = group_id; + } + + /// Sets the damping coefficient. + pub fn set_damping(&mut self, damping: f32) { + self.damping = damping; + } + + /// Updates the particle mass based on a new density. + pub fn set_density(&mut self, density: f32) { + self.mass = self.init_volume * density; + } + + /// Converts to the GPU `Kinematics` struct. + fn to_gpu_kinematics(self) -> Kinematics { + Kinematics { + affine: PaddedMatrix::add_padding(self.affine), + velocity: self.velocity, + vel_grad_det: self.vel_grad_det, + force_dt: self.force_dt, + mass: self.mass, + enabled: self.enabled, + boundary_friction: self.boundary_friction, + #[cfg(feature = "dim3")] + _padding: Default::default(), + cdf: self.cdf, + #[cfg(feature = "dim2")] + _tail_padding: Default::default(), + } + } + + /// Converts the deformation gradient to a GPU `PaddedMatrix`. + fn to_gpu_def_grad(self) -> PaddedMatrix { + PaddedMatrix::add_padding(self.def_grad) + } + + /// Converts to the GPU `ParticleProperties` struct. + fn to_gpu_properties(self) -> ParticleProperties { + ParticleProperties { + init_volume: self.init_volume, + init_radius: self.init_radius, + damping: self.damping, + phase: self.phase, + fixed: self.fixed, + group_id: self.group_id, + padding: Default::default(), + } + } +} + +struct SoAParticles { + positions: Vec, + kinematics: Vec, + def_grad: Vec, + properties: Vec, + models: Vec, +} + +impl SoAParticles { + pub fn new(particles: &[Particle]) -> Self { + let positions: Vec<_> = particles + .iter() + .map(|p| Position::new(p.position)) + .collect(); + let kinematics: Vec<_> = particles + .iter() + .map(|p| p.dynamics.to_gpu_kinematics()) + .collect(); + let def_grad: Vec<_> = particles + .iter() + .map(|p| p.dynamics.to_gpu_def_grad()) + .collect(); + let properties: Vec<_> = particles + .iter() + .map(|p| p.dynamics.to_gpu_properties()) + .collect(); + let models: Vec<_> = particles + .iter() + .map(|p| GpuParticleModel::from(p.model)) + .collect(); + + Self { + positions, + kinematics, + def_grad, + properties, + models, + } + } +} + +/// GPU buffers for particles sampled from rigid body surfaces. +pub struct GpuRigidParticles { + /// Sample points in local (body-relative) coordinates. + pub local_sample_points: Tensor, + /// Sample points transformed to world coordinates. + pub sample_points: Tensor, + /// Bitmask indicating which rigid particles need grid cell blocking. + pub rigid_particle_needs_block: Tensor, + /// Rigid particle indices sorted by grid block (with room for per-block "extras"). + pub sorted_ids: Tensor, + /// Metadata associating each sample with its source collider and body. + pub sample_ids: Tensor, +} + +impl GpuRigidParticles { + /// Creates an empty set of rigid particles. + pub fn new(backend: &GpuBackend) -> Result { + let empty_positions: &[Position] = &[]; + let empty_ids: &[RigidParticleIndices] = &[]; + Ok(Self { + local_sample_points: Tensor::vector(backend, empty_positions, BufferUsages::STORAGE)?, + sample_points: Tensor::vector(backend, empty_positions, BufferUsages::STORAGE)?, + sorted_ids: Tensor::vector_uninit(backend, 0, BufferUsages::STORAGE)?, + sample_ids: Tensor::vector(backend, empty_ids, BufferUsages::STORAGE)?, + rigid_particle_needs_block: Tensor::vector_uninit(backend, 0, BufferUsages::STORAGE)?, + }) + } + + fn from_buffers( + backend: &GpuBackend, + sampling_buffers: &SamplingBuffers, + ) -> Result { + Ok(Self { + local_sample_points: Tensor::vector( + backend, + &sampling_buffers.samples, + BufferUsages::STORAGE, + )?, + sample_points: Tensor::vector( + backend, + &sampling_buffers.samples, + BufferUsages::STORAGE, + )?, + sorted_ids: Tensor::vector_uninit( + backend, + sampling_buffers.samples.len() as u32 * 2_u32.pow(DIM as u32), + BufferUsages::STORAGE, + )?, + sample_ids: Tensor::vector( + backend, + &sampling_buffers.samples_ids, + BufferUsages::STORAGE, + )?, + rigid_particle_needs_block: Tensor::vector_uninit( + backend, + sampling_buffers.samples.len().div_ceil(32) as u32, + BufferUsages::STORAGE, + )?, + }) + } + + /// Samples particles from collider surfaces for MPM coupling. + pub fn from_rapier( + backend: &GpuBackend, + colliders: &rapier::geometry::ColliderSet, + gpu_bodies: &GpuBodySet, + coupling: &[RapierBodyCouplingEntry], + sampling_step: f32, + ) -> Result { + let mut sampling_buffers = SamplingBuffers::default(); + + for (collider_id, (coupling, gpu_data)) in coupling + .iter() + .zip(gpu_bodies.shapes_data().iter()) + .enumerate() + { + let collider = &colliders[coupling.collider]; + + #[cfg(feature = "dim2")] + if let Some(polyline) = collider.shape().as_polyline() { + // Use polyline_vertex_start() to get the correct base index, + // which accounts for BVH AABB data preceding the actual vertices. + let sampling_params = SamplingParams { + collider_id: collider_id as u32, + base_vid: gpu_data.polyline_vertex_start(), + sampling_step, + }; + sampling::sample_polyline(polyline, &sampling_params, &mut sampling_buffers) + } + + #[cfg(feature = "dim3")] + if let Some(trimesh) = collider.shape().as_trimesh() { + // Use trimesh_vertex_start() to get the correct base index, + // which accounts for BVH AABB data preceding the actual vertices. + let sampling_params = SamplingParams { + collider_id: collider_id as u32, + base_vid: gpu_data.trimesh_vertex_start(), + sampling_step, + }; + sampling::sample_trimesh(trimesh, &sampling_params, &mut sampling_buffers) + } else if let Some(heightfield) = collider.shape().as_heightfield() { + let (vtx, idx) = heightfield.to_trimesh(); + let trimesh = rapier::geometry::TriMesh::new(vtx, idx).unwrap(); + // Use trimesh_vertex_start() to get the correct base index, + // which accounts for BVH AABB data preceding the actual vertices. + let sampling_params = SamplingParams { + collider_id: collider_id as u32, + base_vid: gpu_data.trimesh_vertex_start(), + sampling_step, + }; + sampling::sample_trimesh(&trimesh, &sampling_params, &mut sampling_buffers) + } + } + + Self::from_buffers(backend, &sampling_buffers) + } + + /// Returns the number of rigid body particles. + pub fn len(&self) -> u64 { + self.sample_points.len() + } + + /// Returns true if there are no rigid body particles. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +/// GPU buffers storing all MPM particle data in Structure-of-Arrays layout. +pub struct GpuParticles { + len: usize, + pub gpu_len: Tensor, + pub positions: Tensor, + pub kinematics: Tensor, + pub def_grad: Tensor, + pub properties: Tensor, + pub models: Tensor, + pub sorted_ids: Tensor, +} + +impl GpuParticles { + /// Returns true if there are no particles. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Returns the number of particles. + pub fn len(&self) -> usize { + self.len + } + + /// Returns reference to GPU buffer containing particle count. + pub fn gpu_len(&self) -> &Tensor { + &self.gpu_len + } + + /// Uploads CPU-side particles to GPU buffers. + pub fn from_particles( + backend: &GpuBackend, + particles: &[Particle], + ) -> Result { + let data = SoAParticles::new(particles); + let resizeable = BufferUsages::STORAGE | BufferUsages::COPY_SRC | BufferUsages::COPY_DST; + Ok(Self { + len: particles.len(), + gpu_len: Tensor::scalar( + backend, + particles.len() as u32, + BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, + )?, + positions: Tensor::vector(backend, &data.positions, resizeable)?, + kinematics: Tensor::vector(backend, &data.kinematics, resizeable)?, + def_grad: Tensor::vector(backend, &data.def_grad, resizeable)?, + properties: Tensor::vector(backend, &data.properties, resizeable)?, + models: Tensor::vector(backend, &data.models, resizeable)?, + sorted_ids: Tensor::vector_uninit( + backend, + particles.len() as u32 * 2_u32.pow(DIM as u32), + resizeable, + )?, + }) + } + + /// Reserves GPU buffer room for at least `capacity` particles, so a + /// subsequent [`append`](Self::append) of up to `capacity` particles doesn't + /// reallocate. Must be called while empty (right after construction, say): + /// it replaces the per-particle buffers with empty, capacity-sized ones. + pub fn reserve( + &mut self, + backend: &GpuBackend, + capacity: usize, + ) -> Result<(), GpuBackendError> { + assert_eq!( + self.len, 0, + "GpuParticles::reserve must be called while empty" + ); + let cap = capacity as u32; + let resizeable = BufferUsages::STORAGE | BufferUsages::COPY_SRC | BufferUsages::COPY_DST; + self.positions = Tensor::with_capacity(backend, cap, resizeable)?; + self.kinematics = Tensor::with_capacity(backend, cap, resizeable)?; + self.def_grad = Tensor::with_capacity(backend, cap, resizeable)?; + self.properties = Tensor::with_capacity(backend, cap, resizeable)?; + self.models = Tensor::with_capacity(backend, cap, resizeable)?; + // `sorted_ids` holds one entry per particle per touched grid node. + self.sorted_ids = Tensor::with_capacity(backend, cap * 2u32.pow(DIM as u32), resizeable)?; + Ok(()) + } + + /// Returns reference to material model buffer. + pub fn models(&self) -> &Tensor { + &self.models + } + + /// Returns mutable reference to material model buffer. + pub fn models_mut(&mut self) -> &mut Tensor { + &mut self.models + } + + /// Returns reference to position buffer. + pub fn positions(&self) -> &Tensor { + &self.positions + } + + /// Returns mutable reference to position buffer. + pub fn positions_mut(&mut self) -> &mut Tensor { + &mut self.positions + } + + /// Returns reference to kinematics buffer. + pub fn kinematics(&self) -> &Tensor { + &self.kinematics + } + + /// Returns mutable reference to kinematics buffer. + pub fn kinematics_mut(&mut self) -> &mut Tensor { + &mut self.kinematics + } + + /// Returns reference to deformation gradient buffer. + pub fn def_grad(&self) -> &Tensor { + &self.def_grad + } + + /// Returns mutable reference to deformation gradient buffer. + pub fn def_grad_mut(&mut self) -> &mut Tensor { + &mut self.def_grad + } + + /// Returns reference to particle properties buffer (read-only on GPU). + pub fn properties(&self) -> &Tensor { + &self.properties + } + + /// Returns mutable reference to particle properties buffer. + pub fn properties_mut(&mut self) -> &mut Tensor { + &mut self.properties + } + + /// Returns reference to sorted particle ID buffer. + pub fn sorted_ids(&self) -> &Tensor { + &self.sorted_ids + } + + /// Returns mutable reference to sorted particle ID buffer. + pub fn sorted_ids_mut(&mut self) -> &mut Tensor { + &mut self.sorted_ids + } + + /// Removes a range of particles from the GPU buffers, shifting elements to fill the gap. + /// + /// Returns the number of removed particles on success. + pub fn shift_remove( + &mut self, + backend: &GpuBackend, + range: impl RangeBounds + Clone, + ) -> Result { + let Self { + len, + gpu_len, + positions, + kinematics, + def_grad, + properties, + models, + sorted_ids: _, + } = self; + + let removed = positions.shift_remove(backend, range.clone())?; + kinematics.shift_remove(backend, range.clone())?; + def_grad.shift_remove(backend, range.clone())?; + properties.shift_remove(backend, range.clone())?; + models.shift_remove(backend, range)?; + + *len -= removed; + backend.write_buffer(gpu_len.buffer_mut(), 0, &[*len as u32])?; + Ok(removed) + } + + /// Appends particles at the end of the GPU buffers. + pub fn append( + &mut self, + backend: &GpuBackend, + particles: &[Particle], + ) -> Result<(), GpuBackendError> { + let Self { + len, + gpu_len, + positions, + kinematics, + def_grad, + properties, + models, + sorted_ids, + } = self; + + let data = SoAParticles::new(particles); + // `sorted_ids` is the spatial-sort scratch; it must stay sized + // `total_particles * 2^DIM` to match `from_particles` (one entry per + // particle per touched grid node). Undersizing it corrupts the sort. + let zeros = vec![0u32; particles.len() * 2usize.pow(DIM as u32)]; + + positions.append(backend, &data.positions)?; + kinematics.append(backend, &data.kinematics)?; + def_grad.append(backend, &data.def_grad)?; + properties.append(backend, &data.properties)?; + models.append(backend, &data.models)?; + sorted_ids.append(backend, &zeros)?; + + *len += particles.len(); + backend.write_buffer(gpu_len.buffer_mut(), 0, &[*len as u32])?; + Ok(()) + } + + /// Removes the given particle slots by *swap-removing* them: the live tail + /// particles are moved into the freed slots and the buffers are truncated. + /// This is `O(number of removed slots)` with no full-buffer shift, which + /// makes runtime removal cheap. The reordering is invisible to the solver + /// because particles are spatially re-sorted at the start of every step. + /// + /// Returns the relocations performed as `(from, to)` pairs (a tail particle + /// moved from slot `from` down to freed slot `to`) so callers can patch any + /// slot→handle maps. `from` is always a now-truncated tail slot and `to` a + /// freed slot below the new length. + pub fn swap_remove( + &mut self, + backend: &GpuBackend, + slots: &[u32], + ) -> Result, GpuBackendError> { + // Process descending so truncating the tail never disturbs a + // not-yet-processed (lower) slot. + let mut targets: Vec = slots.to_vec(); + targets.sort_unstable_by(|a, b| b.cmp(a)); + targets.dedup(); + + let mut remaps = Vec::new(); + for slot in targets { + let slot = slot as usize; + if slot >= self.len { + continue; + } + let last = self.len - 1; + if slot != last { + // Relocate the live tail particle into the freed slot. A staging + // buffer avoids same-buffer overlapping copies. + macro_rules! relocate { + ($t:expr) => {{ + let mut staging = backend.uninit_buffer( + 1, + BufferUsages::STORAGE | BufferUsages::COPY_SRC | BufferUsages::COPY_DST, + )?; + let mut enc = backend.begin_encoding(); + enc.copy_buffer_to_buffer($t.buffer(), last, &mut staging, 0, 1)?; + enc.copy_buffer_to_buffer(&staging, 0, $t.buffer_mut(), slot, 1)?; + backend.submit(enc)?; + }}; + } + relocate!(self.positions); + relocate!(self.kinematics); + relocate!(self.def_grad); + relocate!(self.properties); + relocate!(self.models); + remaps.push((last as u32, slot as u32)); + } + + // Drop the (now-duplicated) tail element. Removing the last element + // shifts nothing, so this is O(1). `sorted_ids` is solver scratch + // (resized by the spatial sort) and is left untouched. + self.positions.shift_remove(backend, last..)?; + self.kinematics.shift_remove(backend, last..)?; + self.def_grad.shift_remove(backend, last..)?; + self.properties.shift_remove(backend, last..)?; + self.models.shift_remove(backend, last..)?; + self.len -= 1; + } + + backend.write_buffer(self.gpu_len.buffer_mut(), 0, &[self.len as u32])?; + Ok(remaps) + } + + /// Reads the current particle world positions back to the CPU. + pub async fn read_positions( + &self, + backend: &GpuBackend, + ) -> Result, GpuBackendError> { + // `append` grows the positions buffer to a power-of-two capacity and + // `slow_read_buffer` reads all of it, so size the destination to the + // capacity and keep only the `len` live particles. + let mut data = vec![Position::default(); self.positions.capacity() as usize]; + backend + .slow_read_buffer(self.positions.buffer(), &mut data) + .await?; + data.truncate(self.len); + Ok(data.iter().map(|p| p.pt).collect()) + } +} diff --git a/src_mpm/solver/particle_model.rs b/src_mpm/solver/particle_model.rs new file mode 100644 index 00000000..8e0febcf --- /dev/null +++ b/src_mpm/solver/particle_model.rs @@ -0,0 +1,325 @@ +use crate::models::{ + DruckerPrager, DruckerPragerPlasticState, DruckerPragerPlasticity, ElasticCoefficients, + ElasticCoefficientsExt, FluidModel, SnowPlasticState, SnowPlasticity, +}; +pub use crate::mpm_shaders::models::default::{GpuParticleModel, MODEL_DATA_WORDS}; +use nexus_rbd::math::DIM; + +/// Material model for MPM particles. +/// +/// Defines the constitutive behavior (how stress relates to deformation) for particles. +#[derive(Copy, Clone, Debug, PartialEq)] +pub enum ParticleModel { + /// Linear elastic material (St. Venant-Kirchhoff). + ElasticLinear(ElasticCoefficients), + /// Neo-Hookean hyperelastic material (better for large deformations). + ElasticNeoHookean(ElasticCoefficients), + /// Sand/granular material with linear elasticity and Drucker-Prager plasticity. + SandLinear(SandModel), + /// Sand with Neo-Hookean elasticity and Drucker-Prager plasticity. + SandNeoHookean(SandModel), + /// Weakly-compressible Newtonian fluid (Tait equation of state). + Fluid(FluidModel), + /// Snow: elasticity with singular-value clamping and compaction hardening. + Snow(SnowModel), +} + +impl Default for ParticleModel { + fn default() -> Self { + Self::elastic(Self::DEFAULT_YOUNG_MODULUS, Self::DEFAULT_POISSON_RATIO) + } +} + +impl ParticleModel { + /// Default Young's modulus for elastic materials (Pa). + pub const DEFAULT_YOUNG_MODULUS: f32 = 1_000.0; + /// Default Poisson's ratio for elastic materials (dimensionless). + pub const DEFAULT_POISSON_RATIO: f32 = 0.2; + /// Default tensile stiffness of a fluid, as a fraction of its bulk modulus. + /// + /// Enough to stop the volume drifting at a free surface without pulling the + /// fluid into blobs: raising it to 1 barely improves the drift further but + /// visibly clumps a settled column. + pub const DEFAULT_FLUID_TENSILE_STIFFNESS: f32 = 0.25; + /// Default snow compression yield threshold (Stomakhin et al. 2013). + pub const DEFAULT_SNOW_CRITICAL_COMPRESSION: f32 = 2.5e-2; + /// Default snow stretch yield threshold (Stomakhin et al. 2013). + pub const DEFAULT_SNOW_CRITICAL_STRETCH: f32 = 7.5e-3; + /// Default snow hardening coefficient (Stomakhin et al. 2013). + pub const DEFAULT_SNOW_HARDENING: f32 = 10.0; + + /// Creates a linear elastic material model. + pub fn elastic(young_modulus: f32, poisson_ratio: f32) -> Self { + Self::ElasticLinear(ElasticCoefficients::from_young_modulus( + young_modulus, + poisson_ratio, + )) + } + + pub fn elastic_neo_hookean(young_modulus: f32, poisson_ratio: f32) -> Self { + Self::ElasticNeoHookean(ElasticCoefficients::from_young_modulus( + young_modulus, + poisson_ratio, + )) + } + + /// Creates a sand/granular material model with Drucker-Prager plasticity. + pub fn sand(young_modulus: f32, poisson_ratio: f32) -> Self { + ParticleModel::SandLinear(SandModel { + plastic_state: DruckerPragerPlasticState { + plastic_deformation_gradient_det: 1.0, + plastic_hardening: 1.0, + log_vol_gain: 0.0, + }, + plastic: DruckerPrager::new(young_modulus, poisson_ratio), + elastic: ElasticCoefficients::from_young_modulus(young_modulus, poisson_ratio), + }) + } + + /// Cohesion parameter that gives a material the requested cohesive shear + /// strength, in Pascals. + /// + /// [`Self::cohesive_sand`] takes cohesion as a *strain*, so the strength it + /// produces scales with the elastic moduli: the same value on stiffer sand + /// is a far stronger material. This inverts + /// `tau_c = cohesion * (d*lambda + 2*mu) * alpha` so the material can be + /// specified by the shear stress it should sustain at zero confining + /// pressure. Damp sand is a few kPa; approaching the material's own weight + /// stress makes it immovable. + pub fn sand_cohesion_for_strength( + young_modulus: f32, + poisson_ratio: f32, + shear_strength: f32, + ) -> f32 { + let (lambda, mu) = crate::models::lame_lambda_mu(young_modulus, poisson_ratio); + let alpha = DruckerPrager::initial_alpha(); + let scale = (DIM as f32 * lambda + 2.0 * mu) * alpha; + shear_strength / scale.max(1.0e-6) + } + + /// Creates a cohesive granular material from the shear strength it should + /// sustain at zero confining pressure, in Pascals. + /// + /// Prefer this over [`Self::cohesive_sand`] unless you already know what + /// cohesion strain the material needs; see + /// [`Self::sand_cohesion_for_strength`]. + pub fn cohesive_sand_with_strength( + young_modulus: f32, + poisson_ratio: f32, + shear_strength: f32, + ) -> Self { + Self::cohesive_sand( + young_modulus, + poisson_ratio, + Self::sand_cohesion_for_strength(young_modulus, poisson_ratio, shear_strength), + ) + } + + /// Creates a cohesive granular material (wet sand, mud, packed snow). + /// + /// `cohesion` is the volumetric log-strain the material sustains in tension + /// before separating: 0 reproduces [`Self::sand`], while positive values let + /// the material hold a shape with no confining pressure. Being a strain, the + /// strength it implies depends on the elastic moduli; use + /// [`Self::cohesive_sand_with_strength`] to specify a stress instead. + pub fn cohesive_sand(young_modulus: f32, poisson_ratio: f32, cohesion: f32) -> Self { + let (lambda, mu) = if young_modulus > 0.0 { + crate::models::lame_lambda_mu(young_modulus, poisson_ratio) + } else { + (-1.0, -1.0) + }; + ParticleModel::SandLinear(SandModel { + plastic_state: DruckerPragerPlasticState { + plastic_deformation_gradient_det: 1.0, + plastic_hardening: 1.0, + log_vol_gain: 0.0, + }, + plastic: DruckerPrager::from_lame_with_cohesion(lambda, mu, cohesion), + elastic: ElasticCoefficients::from_young_modulus(young_modulus, poisson_ratio), + }) + } + + /// Creates a sand/granular material model with Neo-Hookean elasticity. + pub fn sand_neo_hookean(young_modulus: f32, poisson_ratio: f32) -> Self { + ParticleModel::SandNeoHookean(SandModel { + plastic_state: DruckerPragerPlasticState { + plastic_deformation_gradient_det: 1.0, + plastic_hardening: 1.0, + log_vol_gain: 0.0, + }, + plastic: DruckerPrager::new(young_modulus, poisson_ratio), + elastic: ElasticCoefficients::from_young_modulus(young_modulus, poisson_ratio), + }) + } + + /// Creates a weakly-compressible fluid. + /// + /// `bulk_modulus` sets how strongly compression is resisted (and therefore + /// how small the stable timestep gets), `gamma` how sharply that resistance + /// grows (7 is the usual value for water), and `viscosity` the dynamic + /// viscosity in Pa.s (~0.001 for water, orders of magnitude more for honey). + pub fn fluid(bulk_modulus: f32, gamma: f32, viscosity: f32) -> Self { + ParticleModel::Fluid(FluidModel { + bulk_modulus, + gamma, + viscosity, + cfl_coeff: 0.5, + tensile_stiffness: Self::DEFAULT_FLUID_TENSILE_STIFFNESS, + }) + } + + /// Creates a water-like fluid with the given bulk modulus. + /// + /// Prefer [`Self::water_for_depth`] unless you already know what bulk modulus + /// the scene needs: too low a value is the usual cause of a fluid that + /// visibly loses volume under its own weight. + pub fn water(bulk_modulus: f32) -> Self { + Self::fluid(bulk_modulus, 7.0, 0.001) + } + + /// Creates a water-like fluid stiff enough to stay nearly incompressible + /// under its own weight. + /// + /// A weakly-compressible fluid trades incompressibility for an explicit + /// solve: the deeper the pool, the higher the pressure at the bottom, and the + /// more the equation of state lets it compress. `depth` is the deepest the + /// fluid will get, and `max_compression` the volume loss tolerated there + /// (0.01 is a good default, since 1% is invisible). + /// + /// The cost is the stable timestep, which scales as `1 / sqrt(bulk_modulus)`. + pub fn water_for_depth(density: f32, gravity: f32, depth: f32, max_compression: f32) -> Self { + Self::water(Self::fluid_bulk_modulus( + density, + gravity, + depth, + max_compression, + 7.0, + )) + } + + /// Bulk modulus that limits compression to `max_compression` at the bottom of + /// a column of fluid `depth` deep. + /// + /// Inverts the equation of state at the hydrostatic pressure `rho g h`. + pub fn fluid_bulk_modulus( + density: f32, + gravity: f32, + depth: f32, + max_compression: f32, + gamma: f32, + ) -> f32 { + let pressure = density * gravity * depth.max(0.0); + let j = (1.0 - max_compression.clamp(1.0e-4, 0.5)).max(1.0e-4); + // `p = k (J^-gamma - 1)` solved for `k`. + let response = j.powf(-gamma) - 1.0; + pressure / response.max(1.0e-6) + } + + /// Creates a snow material with the default yield box and hardening from + /// Stomakhin et al. 2013. + pub fn snow(young_modulus: f32, poisson_ratio: f32) -> Self { + Self::snow_with_params( + young_modulus, + poisson_ratio, + Self::DEFAULT_SNOW_CRITICAL_COMPRESSION, + Self::DEFAULT_SNOW_CRITICAL_STRETCH, + Self::DEFAULT_SNOW_HARDENING, + ) + } + + /// Creates a snow material with an explicit yield box and hardening. + /// + /// Widening `critical_stretch` makes the snow hold together under tension + /// (wet, packing snow); shrinking it makes it powdery. `hardening` controls + /// how much stiffer compacted snow becomes than loose snow. + pub fn snow_with_params( + young_modulus: f32, + poisson_ratio: f32, + critical_compression: f32, + critical_stretch: f32, + hardening: f32, + ) -> Self { + ParticleModel::Snow(SnowModel { + plastic_state: SnowPlasticState { plastic_det: 1.0 }, + plastic: SnowPlasticity { + critical_compression, + critical_stretch, + hardening, + }, + elastic: ElasticCoefficients::from_young_modulus(young_modulus, poisson_ratio), + }) + } +} + +/// Combined elastic-plastic model for snow. +#[derive(Copy, Clone, Debug, PartialEq, bytemuck::Pod, bytemuck::Zeroable)] +#[repr(C)] +pub struct SnowModel { + /// Current plastic compaction state. + pub plastic_state: SnowPlasticState, + /// Yield box and hardening parameters. + pub plastic: SnowPlasticity, + /// Elastic coefficients (Lamé parameters) before hardening. + pub elastic: ElasticCoefficients, +} + +/// Combined elastic-plastic model for sand and granular materials. +#[derive(Copy, Clone, Debug, PartialEq, bytemuck::Pod, bytemuck::Zeroable)] +#[repr(C)] +pub struct SandModel { + /// Current plastic deformation state. + pub plastic_state: DruckerPragerPlasticState, + /// Drucker-Prager plasticity model parameters. + pub plastic: DruckerPragerPlasticity, + /// Elastic coefficients (Lamé parameters). + pub elastic: ElasticCoefficients, +} + +// NOTE: keeps `GpuParticleModel` (tag + [u32; MODEL_DATA_WORDS]) in step with +// the GPU-side layout. +static_assertions::assert_eq_size!(GpuParticleModel, [u8; 4 + 4 * MODEL_DATA_WORDS]); + +impl From for GpuParticleModel { + fn from(val: ParticleModel) -> Self { + let mut data = [0u32; MODEL_DATA_WORDS]; + let tag = match val { + ParticleModel::ElasticLinear(elastic) => { + let bytes = bytemuck::bytes_of(&elastic); + bytemuck::cast_slice_mut::(&mut data)[..bytes.len()] + .copy_from_slice(bytes); + 0 + } + ParticleModel::ElasticNeoHookean(elastic) => { + let bytes = bytemuck::bytes_of(&elastic); + bytemuck::cast_slice_mut::(&mut data)[..bytes.len()] + .copy_from_slice(bytes); + 1 + } + ParticleModel::SandLinear(sand) => { + let bytes = bytemuck::bytes_of(&sand); + bytemuck::cast_slice_mut::(&mut data)[..bytes.len()] + .copy_from_slice(bytes); + 2 + } + ParticleModel::SandNeoHookean(sand) => { + let bytes = bytemuck::bytes_of(&sand); + bytemuck::cast_slice_mut::(&mut data)[..bytes.len()] + .copy_from_slice(bytes); + 3 + } + ParticleModel::Fluid(fluid) => { + let bytes = bytemuck::bytes_of(&fluid); + bytemuck::cast_slice_mut::(&mut data)[..bytes.len()] + .copy_from_slice(bytes); + 4 + } + ParticleModel::Snow(snow) => { + let bytes = bytemuck::bytes_of(&snow); + bytemuck::cast_slice_mut::(&mut data)[..bytes.len()] + .copy_from_slice(bytes); + 5 + } + }; + GpuParticleModel { tag, data } + } +} diff --git a/src_mpm/solver/particle_update.rs b/src_mpm/solver/particle_update.rs new file mode 100644 index 00000000..34f0b967 --- /dev/null +++ b/src_mpm/solver/particle_update.rs @@ -0,0 +1,45 @@ +//! Particle state update kernel. +//! +//! Updates particle positions, deformation gradients, and material state after +//! grid velocities have been transferred back to particles. + +use crate::grid::grid::GpuGrid; +use crate::mpm_shaders::solver::particle_update::GpuParticleUpdate; +use crate::solver::{GpuParticles, GpuSimulationParams}; +use khal::Shader; +use khal::backend::{GpuBackendError, GpuPass}; + +/// GPU compute kernel for updating particle state. +/// +/// Integrates particle positions using updated velocities, updates deformation +/// gradients, and applies constitutive models (elasticity, plasticity). +#[derive(Shader)] +pub struct WgParticleUpdate { + /// Compiled particle update compute shader. + particle_update: GpuParticleUpdate, +} + +impl WgParticleUpdate { + /// Launches the particle update kernel. + pub fn launch( + &self, + pass: &mut GpuPass, + sim_params: &GpuSimulationParams, + grid: &GpuGrid, + particles: &mut GpuParticles, + ) -> Result<(), GpuBackendError> { + let len = particles.len() as u32; + self.particle_update.call( + pass, + [len, 1, 1], + &sim_params.params, + &grid.meta, + &mut particles.models, + &mut particles.positions, + &mut particles.kinematics, + &mut particles.def_grad, + &particles.properties, + &particles.gpu_len, + ) + } +} diff --git a/src_mpm/solver/prep_readback.rs b/src_mpm/solver/prep_readback.rs new file mode 100644 index 00000000..e4fdd55d --- /dev/null +++ b/src_mpm/solver/prep_readback.rs @@ -0,0 +1,251 @@ +//! GPU readback preparation kernel and associated data structures. +//! +//! Computes per-particle render data on the GPU, reducing the amount of data +//! transferred back to the CPU compared to reading raw positions and dynamics. + +use crate::grid::grid::GpuGrid; +use crate::mpm_shaders::solver::prep_readback::{ + GpuMpmPrepRender, GpuPrepReadback, GpuPrepReadbackRigid, +}; +pub use crate::mpm_shaders::solver::prep_readback::{ReadbackData, RenderConfig}; +use crate::solver::{GpuParticles, GpuRigidParticles, GpuSimulationParams}; +use glamx::Vec4; +use khal::backend::{ + Encoder, GpuBackend, GpuBackendError, GpuBufferSliceMut, GpuEncoder, GpuTimestamps, +}; +use khal::{BufferUsages, Shader}; +use vortx::tensor::Tensor; + +/// GPU compute kernel for preparing per-particle readback data. +/// +/// Runs a compute shader that transforms particle positions and dynamics +/// into compact `ReadbackData` suitable for rendering, then copies the +/// result to a staging buffer for CPU readback. +#[derive(Shader)] +pub struct WgPrepReadback { + prep_readback: GpuPrepReadback, + prep_readback_rigid: GpuPrepReadbackRigid, + prep_render: GpuMpmPrepRender, +} + +/// GPU-resident buffers for the readback preparation pipeline. +/// +/// Contains the render configuration, group palette, and output buffers +/// for the readback shader. +pub struct GpuReadbackData { + /// Render mode configuration (uniform, written by CPU). + pub mode: Tensor, + /// Color palette indexed by `ParticleProperties::group_id`. + pub group_colors: Tensor, + /// Number of entries in `group_colors`, mirrored into `RenderConfig`. + pub num_groups: u32, + /// Shader output buffer (written by GPU, source for staging copy). + pub instances: Tensor, + /// Staging buffer for CPU readback (MAP_READ). + pub instances_staging: Tensor, + /// Per-rigid-particle base colors. + pub rigid_base_colors: Tensor, + /// Rigid particle shader output buffer. + pub rigid_instances: Tensor, + /// Staging buffer for rigid particle CPU readback (MAP_READ). + pub rigid_instances_staging: Tensor, + /// Rigid particle count uniform for the shader. + pub rigid_len: Tensor, +} + +impl GpuReadbackData { + /// Fallback palette for a scene that supplies none. + pub const DEFAULT_GROUP_COLORS: [Vec4; 6] = [ + Vec4::new(124.0 / 255.0, 144.0 / 255.0, 1.0, 1.0), + Vec4::new(8.0 / 255.0, 144.0 / 255.0, 1.0, 1.0), + Vec4::new(124.0 / 255.0, 7.0 / 255.0, 1.0, 1.0), + Vec4::new(124.0 / 255.0, 144.0 / 255.0, 7.0 / 255.0, 1.0), + Vec4::new(200.0 / 255.0, 37.0 / 255.0, 1.0, 1.0), + Vec4::new(124.0 / 255.0, 230.0 / 255.0, 25.0 / 255.0, 1.0), + ]; + + /// Creates new readback data buffers for the given number of particles. + /// + /// `group_colors` is the palette particle group ids index into; an empty + /// slice falls back to [`Self::DEFAULT_GROUP_COLORS`]. + pub fn new( + backend: &GpuBackend, + num_particles: usize, + num_rigid_particles: usize, + mode: u32, + group_colors: &[Vec4], + ) -> Result { + let group_colors: Vec = if group_colors.is_empty() { + Self::DEFAULT_GROUP_COLORS.to_vec() + } else { + group_colors.to_vec() + }; + let num_groups = group_colors.len() as u32; + let rigid_base_colors: Vec = (0..num_rigid_particles) + .map(|i| group_colors[i % group_colors.len()]) + .collect(); + + // Use at least 1 element for GPU buffers to avoid zero-sized allocations. + let rigid_buf_len = num_rigid_particles.max(1) as u32; + + Ok(Self { + num_groups, + mode: Tensor::scalar( + backend, + RenderConfig { + mode, + num_groups, + ..Default::default() + }, + // STORAGE for the readback kernel, UNIFORM for the render kernel. + BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, + )?, + group_colors: Tensor::vector(backend, group_colors, BufferUsages::STORAGE)?, + instances: Tensor::vector_uninit( + backend, + num_particles as u32, + BufferUsages::STORAGE | BufferUsages::COPY_SRC, + )?, + instances_staging: Tensor::vector_uninit( + backend, + num_particles as u32, + BufferUsages::COPY_DST | BufferUsages::MAP_READ, + )?, + rigid_base_colors: Tensor::vector(backend, rigid_base_colors, BufferUsages::STORAGE)?, + rigid_instances: Tensor::vector_uninit( + backend, + rigid_buf_len, + BufferUsages::STORAGE | BufferUsages::COPY_SRC, + )?, + rigid_instances_staging: Tensor::vector_uninit( + backend, + rigid_buf_len, + BufferUsages::COPY_DST | BufferUsages::MAP_READ, + )?, + rigid_len: Tensor::scalar( + backend, + num_rigid_particles as u32, + BufferUsages::STORAGE | BufferUsages::UNIFORM, + )?, + }) + } + + /// Recreates all buffers for a new particle count. + pub fn resize( + &mut self, + backend: &GpuBackend, + num_particles: usize, + num_rigid_particles: usize, + mode: u32, + group_colors: &[Vec4], + ) -> Result<(), GpuBackendError> { + *self = Self::new( + backend, + num_particles, + num_rigid_particles, + mode, + group_colors, + )?; + Ok(()) + } +} + +impl WgPrepReadback { + /// Launches the readback preparation shader and copies results to staging. + /// + /// This runs a compute pass that writes `ReadbackData` into `instances`, + /// then copies `instances` → `instances_staging` for CPU readback. + /// Also dispatches the rigid particle readback shader if there are rigid particles. + pub fn launch( + &self, + encoder: &mut GpuEncoder, + timestamps: Option<&mut GpuTimestamps>, + readback: &mut GpuReadbackData, + sim_params: &GpuSimulationParams, + grid: &GpuGrid, + particles: &GpuParticles, + rigid_particles: &GpuRigidParticles, + ) -> Result<(), GpuBackendError> { + let len = particles.len() as u32; + let rigid_len = rigid_particles.len() as u32; + { + let mut pass = encoder.begin_pass("prep-readback", timestamps); + self.prep_readback.call( + &mut pass, + [len, 1, 1], + &mut readback.instances, + &particles.positions, + &particles.kinematics, + &particles.def_grad, + &particles.properties, + &grid.meta, + &sim_params.params, + &readback.mode, + &particles.gpu_len, + &readback.group_colors, + )?; + + if rigid_len > 0 { + self.prep_readback_rigid.call( + &mut pass, + [rigid_len, 1, 1], + &mut readback.rigid_instances, + &rigid_particles.sample_points, + &readback.rigid_base_colors, + &grid.meta, + &readback.rigid_len, + )?; + } + } + readback + .instances_staging + .copy_from_view(encoder, &readback.instances)?; + if rigid_len > 0 { + readback + .rigid_instances_staging + .copy_from_view(encoder, &readback.rigid_instances)?; + } + Ok(()) + } + + /// Zero-readback variant of [`Self::launch`]: writes per-particle render data + /// straight into a renderer's SoA instance buffers (`positions`, + /// `deformations`, `colors`), reading particle state directly on the GPU. No + /// staging copy and no CPU readback. `readback` is reused only for its group + /// palette and render-mode uniform. + #[allow(clippy::too_many_arguments)] + pub fn launch_render( + &self, + encoder: &mut GpuEncoder, + positions: &mut GpuBufferSliceMut, + deformations: &mut GpuBufferSliceMut, + colors: &mut GpuBufferSliceMut, + readback: &GpuReadbackData, + sim_params: &GpuSimulationParams, + grid: &GpuGrid, + particles: &GpuParticles, + ) -> Result<(), GpuBackendError> { + let len = particles.len() as u32; + if len == 0 { + return Ok(()); + } + let mut pass = encoder.begin_pass("mpm-prep-render", None); + self.prep_render.call( + &mut pass, + [len, 1, 1], + positions, + deformations, + colors, + &particles.positions, + &particles.kinematics, + &particles.def_grad, + &particles.properties, + &grid.meta, + &sim_params.params, + &readback.mode, + &particles.gpu_len, + &readback.group_colors, + )?; + Ok(()) + } +} diff --git a/src_mpm/solver/rigid_integrate.rs b/src_mpm/solver/rigid_integrate.rs new file mode 100644 index 00000000..5a0e73c4 --- /dev/null +++ b/src_mpm/solver/rigid_integrate.rs @@ -0,0 +1,121 @@ +//! Impulse accumulation and application for MPM-rigid body coupling. + +use crate::grid::grid::GpuGrid; +use crate::mpm_shaders::solver::p2g::IntegerImpulse; +use crate::mpm_shaders::solver::rigid_impulses::{ + GpuRigidImpulsesUpdate, GpuUpdateWorldMassProperties, GpuWritebackBodyPoses, +}; +use crate::solver::GpuSimulationParams; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::{BufferUsages, Shader}; +use nexus_rbd::dynamics::GpuBodySet; +use nexus_rbd::math::Pose; +use vortx::tensor::Tensor; + +/// GPU kernels for computing and applying impulses to rigid bodies from MPM. +/// +/// Accumulates forces from MPM particles and applies them as impulses to +/// coupled rigid bodies for two-way interaction. +#[derive(Shader)] +pub struct WgIntegrateBodies { + /// Kernel for computing and applying impulses. + update: GpuRigidImpulsesUpdate, + /// Kernel for updating world-space mass properties. + update_world_mass_properties: GpuUpdateWorldMassProperties, + /// Kernel copying the coupled bodies' poses back to the rigid-body pipeline. + writeback_body_poses: GpuWritebackBodyPoses, +} + +/// GPU buffers for storing impulses from MPM to rigid bodies. +pub struct GpuImpulses { + /// Per-timestep incremental impulses (uses atomic integer operations). + pub incremental_impulses: Tensor, +} + +impl GpuImpulses { + /// Creates impulse buffers for rigid bodies. + /// + /// Allocates space for up to 16 bodies (CPIC limitation). + pub fn new(backend: &GpuBackend) -> Result { + const MAX_BODY_COUNT: usize = 16; // CPIC doesn't support more. + let impulses = [IntegerImpulse::default(); MAX_BODY_COUNT]; + Ok(Self { + incremental_impulses: Tensor::vector(backend, impulses, BufferUsages::STORAGE)?, + }) + } +} + +impl WgIntegrateBodies { + /// Computes and applies impulses to rigid bodies from MPM grid. + pub fn launch( + &self, + pass: &mut GpuPass, + grid: &GpuGrid, + sim_params: &GpuSimulationParams, + impulses: &mut GpuImpulses, + bodies: &mut GpuBodySet, + ) -> Result<(), GpuBackendError> { + if bodies.is_empty() { + return Ok(()); + } + + self.update.call( + pass, + 1u32, + &sim_params.params, + &grid.meta, + &bodies.local_mprops, + &mut bodies.poses, + &mut bodies.vels, + &mut bodies.mprops, + &mut impulses.incremental_impulses, + ) + } + + /// Updates world-space mass properties for rigid bodies. + /// + /// Transforms local inertia tensors to world coordinates based on current poses. + pub fn launch_update_world_mass_properties( + &self, + pass: &mut GpuPass, + impulses: &mut GpuImpulses, + bodies: &mut GpuBodySet, + ) -> Result<(), GpuBackendError> { + if bodies.is_empty() { + return Ok(()); + } + + let len = bodies.len(); + self.update_world_mass_properties.call( + pass, + [len, 1, 1], + &bodies.poses, + &bodies.local_mprops, + &mut bodies.mprops, + &mut impulses.incremental_impulses, + ) + } + + /// Copies the coupled bodies' MPM-integrated poses into `rbd_poses`, the + /// rigid-body pipeline's body-pose buffer, at the slots given by + /// `rbd_slots` (one per coupled body). + /// + /// See `gpu_writeback_body_poses`: MPM owns the pose of every body it is + /// coupled to, so the rigid-body copy that rendering and the broad phase + /// read has to be refreshed from it once per frame. + pub fn launch_writeback_body_poses( + &self, + pass: &mut GpuPass, + bodies: &GpuBodySet, + rbd_slots: &Tensor, + rbd_poses: &mut Tensor, + ) -> Result<(), GpuBackendError> { + let len = rbd_slots.len() as u32; + if bodies.is_empty() || len == 0 { + return Ok(()); + } + + self.writeback_body_poses + .call(pass, [len, 1, 1], &bodies.poses, rbd_slots, rbd_poses) + } +} diff --git a/src_mpm/solver/rigid_particle_update.rs b/src_mpm/solver/rigid_particle_update.rs new file mode 100644 index 00000000..afd6da2e --- /dev/null +++ b/src_mpm/solver/rigid_particle_update.rs @@ -0,0 +1,60 @@ +//! Rigid body particle transformation kernels. + +use crate::mpm_shaders::solver::rigid_particle_update::{ + GpuTransformSamplePoints, GpuTransformShapePoints, +}; +use crate::solver::GpuRigidParticles; +use khal::Shader; +use khal::backend::{GpuBackendError, GpuPass}; +use nexus_rbd::dynamics::GpuBodySet; +/// GPU kernels for updating rigid body particle positions. +/// +/// Transforms surface-sampled particles from local to world coordinates +/// as rigid bodies move. +#[derive(Shader)] +pub struct WgRigidParticleUpdate { + /// Kernel for transforming sample points. + transform_sample_points: GpuTransformSamplePoints, + /// Kernel for transforming collider mesh vertices. + transform_shape_points: GpuTransformShapePoints, +} + +impl WgRigidParticleUpdate { + /// Transforms rigid body particles from local to world space. + pub fn launch( + &self, + pass: &mut GpuPass, + bodies: &mut GpuBodySet, + rigid_particles: &mut GpuRigidParticles, + ) -> Result<(), GpuBackendError> { + if rigid_particles.is_empty() { + return Ok(()); + } + + let sample_len = rigid_particles.local_sample_points.len() as u32; + self.transform_sample_points.call( + pass, + [sample_len, 1, 1], + &rigid_particles.sample_ids, + &bodies.poses, + &rigid_particles.local_sample_points, + &mut rigid_particles.sample_points, + )?; + + // Keep the world-space vertex buffer in sync with the body poses; `p2g_cdf` + // projects world-space grid nodes on these vertices. The BVH AABB entries + // interleaved in this buffer get transformed as plain points, which is fine: + // consumers needing the BVH read `shapes_local_vertex_buffers` instead. + let vtx_len = bodies.shapes_vertex_buffers.len() as u32; + self.transform_shape_points.call( + pass, + [vtx_len, 1, 1], + &bodies.shapes_vertex_collider_id, + &bodies.poses, + &bodies.shapes_local_vertex_buffers, + &mut bodies.shapes_vertex_buffers, + )?; + + Ok(()) + } +} diff --git a/src_mpm/solver/timestep_bound.rs b/src_mpm/solver/timestep_bound.rs new file mode 100644 index 00000000..db1eddf4 --- /dev/null +++ b/src_mpm/solver/timestep_bound.rs @@ -0,0 +1,70 @@ +//! Timestep bound estimation kernels. +//! +//! Computes a CFL-based upper bound on the simulation timestep to prevent +//! numerical instability. Uses material sound speed and particle velocities +//! to determine the maximum safe timestep. + +use crate::grid::grid::GpuGrid; +use crate::mpm_shaders::solver::timestep_bound::{ + GpuEstimateTimestepBound, GpuResetTimestepBound, GpuTimestepBounds, +}; +use crate::solver::GpuParticles; +use khal::Shader; +use khal::backend::{Backend, Encoder, GpuBackend, GpuBackendError, GpuPass, GpuTimestamps}; +use vortx::tensor::Tensor; + +/// GPU kernel for estimating the maximum stable timestep (best-effort, does not eliminate all divergence risk). +#[derive(Shader)] +pub struct WgTimestepBounds { + reset_timestep_bound: GpuResetTimestepBound, + estimate_timestep_bound: GpuEstimateTimestepBound, +} + +impl WgTimestepBounds { + /// Launches the timestep bounds estimation and returns the estimated maximum timestep length. + pub async fn compute_bounds( + &self, + backend: &GpuBackend, + timestamps: Option<&mut GpuTimestamps>, + grid: &GpuGrid, + particles: &GpuParticles, + bounds: &mut Tensor, + bounds_staging: &mut Tensor, + ) -> Result { + let mut encoder = backend.begin_encoding(); + let mut pass = encoder.begin_pass("timestep-bounds", timestamps); + self.launch(&mut pass, grid, particles, bounds)?; + drop(pass); + bounds_staging.copy_from_view(&mut encoder, &*bounds)?; + backend.submit(encoder)?; + + let mut result = [GpuTimestepBounds::default()]; + backend + .read_buffer(bounds_staging.buffer(), &mut result) + .await?; + Ok(result[0].computed_max_dt_as_uint as f32 / GpuTimestepBounds::FLOAT_TO_INT) + } + + fn launch( + &self, + pass: &mut GpuPass, + grid: &GpuGrid, + particles: &GpuParticles, + bounds: &mut Tensor, + ) -> Result<(), GpuBackendError> { + self.reset_timestep_bound.call(pass, 1u32, bounds)?; + + let len = particles.len() as u32; + self.estimate_timestep_bound.call( + pass, + [len, 1, 1], + &grid.meta, + &particles.models, + &particles.kinematics, + &particles.def_grad, + &particles.properties, + &particles.gpu_len, + bounds, + ) + } +} diff --git a/src_mpm/trimesh.rs b/src_mpm/trimesh.rs new file mode 100644 index 00000000..2b73103a --- /dev/null +++ b/src_mpm/trimesh.rs @@ -0,0 +1,132 @@ +// TODO: move this to nexus? + +use crate::mpm_shaders::PaddedVector; +use nexus_rbd::math::Vector; +use rapier::geometry::TriMesh; + +/// Convert a glamx Vec3 to bvh's Point3 (different glam versions). +fn to_bvh_point(v: Vector) -> bvh::Point3 { + bvh::Point3::new(v.x, v.y, v.z) +} + +/// Convert a bvh Point3 to a glamx Vec3. +fn from_bvh_point(v: bvh::Point3) -> Vector { + Vector::new(v.x, v.y, v.z) +} + +#[derive(Default, Clone, Debug)] +pub struct ShapeBuffers { + /// Vertex buffer for polylines and triangle meshes. + /// + /// Polyline and TriMesh shapes reference ranges within this buffer. + /// The shape stores the start and end indices of its vertices in this buffer. + pub vertices: Vec, + /// Index buffers for polylines, triangle meshes, and convex polyhedrons. + pub indices: Vec, +} + +#[derive(Copy, Clone)] +pub struct GpuTriMesh { + /// Index of the root AABB in the vertex buffer. + pub bvh_vtx_root_id: u32, + /// The root AABB left-child index. + pub bvh_idx_root_id: u32, + // The number of BVH nodes. Triangle indices are stored after the last bvh node. + pub bvh_node_len: u32, + // The total number of triangles in the mesh. + pub num_triangles: u32, + // The total number of vertices in the mesh. + pub num_vertices: u32, +} + +pub fn convert_trimesh_to_gpu(shape: &TriMesh, buffers: &mut ShapeBuffers) -> GpuTriMesh { + let bvh_vtx_root_id = buffers.vertices.len(); + let bvh_idx_root_id = buffers.indices.len(); + // Append the BVH data to the vertex/index buffers. + // TODO: we are constructing a BVH using the `bvh` crate. + // While the TriMesh shape technically already has a BVH, parry's BVH + // doesn't provide explicit access to the BVH topology. So, for now, + // let's just build a new BVH that exposes its internal. + struct BvhObject { + aabb: bvh::aabb::AABB, + node_index: usize, + } + + impl bvh::aabb::Bounded for BvhObject { + fn aabb(&self) -> bvh::aabb::AABB { + self.aabb + } + } + + impl bvh::bounding_hierarchy::BHShape for BvhObject { + fn set_bh_node_index(&mut self, index: usize) { + self.node_index = index; + } + + fn bh_node_index(&self) -> usize { + self.node_index + } + } + + let mut objects: Vec<_> = shape + .triangles() + .map(|tri| { + let aabb = tri.local_aabb(); + BvhObject { + aabb: bvh::aabb::AABB::with_bounds( + to_bvh_point(aabb.mins), + to_bvh_point(aabb.maxs), + ), + node_index: 0, + } + }) + .collect(); + + let bvh = bvh::bvh::BVH::build(&mut objects); + let flat_bvh = bvh.flatten(); + buffers.vertices.extend(flat_bvh.iter().flat_map(|n| { + [ + PaddedVector::new(from_bvh_point(n.aabb.min)), + PaddedVector::new(from_bvh_point(n.aabb.max)), + ] + })); + let bvh_node_len = flat_bvh.len(); + buffers.indices.extend( + flat_bvh + .iter() + .flat_map(|n| [n.entry_index, n.exit_index, n.shape_index]), + ); + + // Append the actual mesh vertex/index buffers. + #[cfg(feature = "dim3")] + { + let pn = shape + .pseudo_normals() + .expect("trimeshes without pseudo-normals are not supported"); + buffers + .vertices + .extend(shape.vertices().iter().map(|v| PaddedVector::new(*v))); + buffers.vertices.extend( + pn.vertices_pseudo_normal + .iter() + .map(|v| PaddedVector::new(*v)), + ); + assert_eq!(shape.vertices().len(), pn.vertices_pseudo_normal.len()); + buffers.vertices.extend( + pn.edges_pseudo_normal + .iter() + .flat_map(|n| *n) + .map(PaddedVector::new), + ); + } + buffers + .indices + .extend(shape.indices().iter().flat_map(|tri| tri.iter().copied())); + GpuTriMesh { + bvh_vtx_root_id: bvh_vtx_root_id as u32, + bvh_idx_root_id: bvh_idx_root_id as u32, + bvh_node_len: bvh_node_len as u32, + num_triangles: shape.indices().len() as u32, + num_vertices: shape.vertices().len() as u32, + } +} diff --git a/src_mpm_shaders/collision/collide.rs b/src_mpm_shaders/collision/collide.rs new file mode 100644 index 00000000..3d400e50 --- /dev/null +++ b/src_mpm_shaders/collision/collide.rs @@ -0,0 +1,60 @@ +//! Collision detection for MPM grid nodes against collider shapes. + +use crate::grid::grid::{AffinityBits, NONE, NodeCdf}; +use crate::nexus_rbd_shaders::shapes::{SHAPE_TYPE_POLYLINE, SHAPE_TYPE_TRIMESH, Shape}; +use crate::{Pose, Vector}; +use khal_std::index::MaybeIndexUnchecked; + +/// Tests a point against all collision shapes and returns the nearest CDF data. +/// +/// For each non-polyline, non-trimesh shape, projects the point onto +/// the shape boundary and tracks distance, affinity bits, and closest collider ID. +pub fn collide( + collision_shapes: &[Shape], + collision_shape_poses: &[Pose], + cell_width: f32, + point: Vector, +) -> NodeCdf { + let mut cdf = NodeCdf::NONE; + + let dist_cap = Vector::splat(cell_width * 1.5); + + // TODO: don't rely on the array length, e.g., if the user wants to + // preallocate the array to add more dynamically. + let num_shapes = collision_shapes.len(); + + for i in 0..num_shapes as u32 { + // FIXME: figure out a way to support more than 16 colliders. + let shape = collision_shapes.read(i as usize); + let shape_pose = collision_shape_poses.read(i as usize); + let shape_type = shape.shape_type(); + + if shape_type != SHAPE_TYPE_POLYLINE && shape_type != SHAPE_TYPE_TRIMESH { + let proj = shape.project_point_on_boundary(shape_pose, point); + let dpt = proj.point - point; + + // Check if the projection is inside or within the distance cap. + // `all(abs(dpt) <= dist_cap)` means every component of abs(dpt) + // is <= the corresponding component of dist_cap. + let abs_dpt = dpt.abs(); + #[cfg(feature = "dim2")] + let within_cap = abs_dpt.x <= dist_cap.x && abs_dpt.y <= dist_cap.y; + #[cfg(feature = "dim3")] + let within_cap = + abs_dpt.x <= dist_cap.x && abs_dpt.y <= dist_cap.y && abs_dpt.z <= dist_cap.z; + + if proj.is_inside || within_cap { + let dist = dpt.length(); + // TODO: take is_inside into account to select the deepest + // penetration as the closest collider? + if dist < cdf.distance { + cdf.closest_id = i; + } + cdf.distance = cdf.distance.min(dist); + cdf.affinities.set_bit(i, proj.is_inside); + } + } + } + + cdf +} diff --git a/src_mpm_shaders/collision/mod.rs b/src_mpm_shaders/collision/mod.rs new file mode 100644 index 00000000..0c2cca38 --- /dev/null +++ b/src_mpm_shaders/collision/mod.rs @@ -0,0 +1 @@ +pub mod collide; diff --git a/src_mpm_shaders/grid/grid.rs b/src_mpm_shaders/grid/grid.rs new file mode 100644 index 00000000..6938eb51 --- /dev/null +++ b/src_mpm_shaders/grid/grid.rs @@ -0,0 +1,744 @@ +//! Sparse grid data structures and hashmap for MPM. +//! +//! The MPM grid is stored as a sparse set of active blocks. Each block contains +//! a fixed number of grid nodes (8x8 in 2D, 4x4x4 in 3D = 64 nodes per block). +//! Active blocks are tracked via a GPU hashmap that maps virtual block +//! coordinates to physical storage indices. + +use crate::nexus_rbd_shaders::utils::udiv_ceil; +use crate::{IVector, Vector}; +use core::ops::BitOrAssign; +use glamx::*; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::{ + macros::{spirv, spirv_bindgen}, + sync::atomic_add_u32, +}; +use nexus_rbd_shaders::MAX_FLT; + +/* + * Constants. + */ + +/// Number of cells (nodes) per block. +/// 8 * 8 = 64 in 2D, 4 * 4 * 4 = 64 in 3D. +const NUM_CELL_PER_BLOCK: u32 = 64; + +/// Workgroup size for grid-level operations. +/// Must match `NUM_CELL_PER_BLOCK` because some kernels (like reset) rely on it. +const GRID_WORKGROUP_SIZE: u32 = NUM_CELL_PER_BLOCK; + +/// Sentinel value indicating "no entry" / "empty slot" in the hashmap and linked lists. +pub const NONE: u32 = 0xFFFFFFFF; + +/// Number of blocks associated with each particle/point. +/// In 2D a particle can straddle 4 blocks (2x2), in 3D it can straddle 8 blocks (2x2x2). +#[cfg(feature = "dim2")] +pub const NUM_ASSOC_BLOCKS: usize = 4; +/// Number of blocks associated with each particle/point. +#[cfg(feature = "dim3")] +pub const NUM_ASSOC_BLOCKS: usize = 8; +pub const NUM_NBH_BLOCKS: usize = NUM_ASSOC_BLOCKS - 1; + +/// Number of "slab" buckets for the within-block counting sort of regular +/// particles (those whose primary block is the block being sorted). +#[cfg(feature = "dim2")] +pub const NUM_PRIMARY_SORT_BUCKETS: usize = 8; +/// Number of "slab" buckets for the within-block counting sort of regular particles. +#[cfg(feature = "dim3")] +pub const NUM_PRIMARY_SORT_BUCKETS: usize = 4; + +/// Number of slab buckets for "extras" (particles spilling in from a neighbour +/// block). +#[cfg(feature = "dim2")] +pub const NUM_EXTRA_SORT_BUCKETS: usize = 10; +/// Number of slab buckets for "extras". +#[cfg(feature = "dim3")] +pub const NUM_EXTRA_SORT_BUCKETS: usize = 6; + +/// Total number of within-block sort buckets. Primary buckets come first so that +/// primaries end up contiguous in `[first_particle, first_particle + num_particles)`, +/// which G2P relies on. +pub const NUM_SORT_BUCKETS: usize = NUM_PRIMARY_SORT_BUCKETS + NUM_EXTRA_SORT_BUCKETS; + +/* + * Index newtypes. + */ + +/// Virtual (logical) block coordinate in the sparse grid. +/// +/// This is an integer vector (IVec2 in 2D, IVec3 in 3D) identifying a block's +/// position in the infinite virtual grid. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct BlockVirtualId { + pub id: IVector, + #[cfg(feature = "dim3")] + pub padding: u32, +} + +impl BlockVirtualId { + pub fn new(id: IVector) -> Self { + Self { + id, + #[cfg(feature = "dim3")] + padding: 0, + } + } + + /// Packs a virtual block ID into a single u32 for use as a hashmap key. + /// + /// In 2D: 16 bits for X, 16 bits for Y. + /// In 3D: 11 bits for X, 10 bits for Y, 11 bits for Z (Y gets fewer bits + /// assuming Y-up, since the vertical extent is typically smaller). + #[cfg(feature = "dim2")] + fn pack(&self) -> u32 { + ((self.id.x + 0x00007FFF) as u32 & 0x0000FFFF) + | (((self.id.y + 0x00007FFF) as u32 & 0x0000FFFF) << 16) + } + + #[cfg(feature = "dim3")] + fn pack(&self) -> u32 { + ((self.id.x + 0x000003FF) as u32 & 0x000007FF) + | (((self.id.y + 0x000001FF) as u32 & 0x000003FF) << 11) + | (((self.id.z + 0x000003FF) as u32 & 0x000007FF) << 21) + } + + /// Returns the primary block associated with a world-space point. + #[cfg(feature = "dim2")] + #[inline] + pub fn block_associated_to_point(cell_width: f32, pt: Vector) -> BlockVirtualId { + let assoc_cell = (pt / cell_width).round() - Vector::ONE; + let assoc_block = (assoc_cell / 8.0).floor(); + BlockVirtualId { + id: IVec2::new(assoc_block.x as i32, assoc_block.y as i32), + } + } + + /// Returns the primary block associated with a world-space point. + #[cfg(feature = "dim3")] + #[inline] + pub fn block_associated_to_point(cell_width: f32, pt: Vector) -> BlockVirtualId { + let assoc_cell = (pt / cell_width).round() - Vector::ONE; + let assoc_block = (assoc_cell / 4.0).floor(); + BlockVirtualId { + id: IVec3::new( + assoc_block.x as i32, + assoc_block.y as i32, + assoc_block.z as i32, + ), + #[cfg(feature = "dim3")] + padding: 0, + } + } + + /// Returns all blocks associated with a world-space point (the ones a + /// particle's quadratic kernel stencil can reach): 4 blocks in 2D, 8 in 3D. + #[inline] + pub fn blocks_associated_to_point( + cell_width: f32, + pt: Vector, + ) -> [BlockVirtualId; NUM_ASSOC_BLOCKS] { + let main_block = Self::block_associated_to_point(cell_width, pt); + Self::blocks_associated_to_block(&main_block) + } + + /// Returns all blocks neighboring a given block (including itself). + /// + /// For a main block at position B, returns all blocks in the 2x2 (2D) or 2x2x2 (3D) + /// neighborhood starting at B. + #[cfg(feature = "dim2")] + #[inline] + pub fn blocks_associated_to_block( + block: &BlockVirtualId, + ) -> [BlockVirtualId; NUM_ASSOC_BLOCKS] { + [ + BlockVirtualId { + id: block.id + IVec2::new(0, 0), + }, + BlockVirtualId { + id: block.id + IVec2::new(0, 1), + }, + BlockVirtualId { + id: block.id + IVec2::new(1, 0), + }, + BlockVirtualId { + id: block.id + IVec2::new(1, 1), + }, + ] + } + + /// Returns all blocks neighboring a given block (including itself). + #[cfg(feature = "dim3")] + #[inline] + pub fn blocks_associated_to_block( + block: &BlockVirtualId, + ) -> [BlockVirtualId; NUM_ASSOC_BLOCKS] { + [ + BlockVirtualId { + id: block.id + IVec3::new(0, 0, 0), + padding: 0, + }, + BlockVirtualId { + id: block.id + IVec3::new(0, 0, 1), + padding: 0, + }, + BlockVirtualId { + id: block.id + IVec3::new(0, 1, 0), + padding: 0, + }, + BlockVirtualId { + id: block.id + IVec3::new(0, 1, 1), + padding: 0, + }, + BlockVirtualId { + id: block.id + IVec3::new(1, 0, 0), + padding: 0, + }, + BlockVirtualId { + id: block.id + IVec3::new(1, 0, 1), + padding: 0, + }, + BlockVirtualId { + id: block.id + IVec3::new(1, 1, 0), + padding: 0, + }, + BlockVirtualId { + id: block.id + IVec3::new(1, 1, 1), + padding: 0, + }, + ] + } +} + +/// Index into the active block headers array. +/// +/// After insertion into the hashmap, each active block is assigned a header ID +/// that serves as its index in the `active_blocks` array. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct BlockHeaderId { + pub id: u32, +} + +impl BlockHeaderId { + /// Converts a block header ID to a physical storage ID. + /// + /// The physical ID is the index of the block's first node in the flat node arrays. + #[inline] + pub fn physical_id(self) -> BlockPhysicalId { + BlockPhysicalId { + id: self.id * NUM_CELL_PER_BLOCK, + } + } +} + +/// Physical (storage) index for a block's first node. +/// +/// Computed as `header_id * NUM_CELL_PER_BLOCK`, so it indexes straight into +/// the flat node arrays. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct BlockPhysicalId { + pub id: u32, +} + +impl BlockPhysicalId { + /// Computes the physical node ID from a block's physical ID and a local offset within the block. + /// + /// In 2D: nodes are laid out in row-major order within 8x8 blocks. + /// In 3D: nodes are laid out in row-major order within 4x4x4 blocks. + #[cfg(feature = "dim2")] + #[inline] + pub fn node_id(self, shift_in_block: UVec2) -> NodePhysicalId { + NodePhysicalId { + id: self.id + shift_in_block.x + shift_in_block.y * 8, + } + } + + /// Computes the physical node ID from a block's physical ID and a local offset within the block. + #[cfg(feature = "dim3")] + #[inline] + pub fn node_id(self, shift_in_block: UVec3) -> NodePhysicalId { + NodePhysicalId { + id: self.id + shift_in_block.x + shift_in_block.y * 4 + shift_in_block.z * 4 * 4, + } + } +} + +/// Physical (storage) index for a single grid node. +/// +/// Computed as `block_physical_id + local_offset_in_block`. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct NodePhysicalId { + pub id: u32, +} + +/* + * Data structures. + */ + +/// An entry in the GPU hashmap that maps block virtual IDs to header IDs. +/// +/// The hashmap uses open addressing with linear probing. The `state` field +/// serves double duty: `NONE` means the slot is empty, otherwise it stores +/// the packed key for comparison during probing. +/// +/// NOTE: changing this struct (including its layout) means changing the +/// host-side struct to match, or the hashmap breaks. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct GridHashMapEntry { + /// The virtual block ID key. + pub key: BlockVirtualId, + /// The associated block header ID value. + pub value: BlockHeaderId, + /// The packed key stored in this slot, or `NONE` if the slot is empty. + pub state: u32, + /// Ownership flag for weak-CAS correctness. + /// Reset to 0 each frame; the first thread to `atomic_exchange` it to 1 + /// becomes the slot's owner and allocates the block header. + pub ownership: u32, + pub padding: [u32; 1], +} + +/// Header for an active block in the sparse grid. +/// +/// Stores the virtual ID (for computing world-space positions) and +/// particle sorting information. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct ActiveBlockHeader { + /// The virtual block coordinate needed to compute world-space node positions. + pub virtual_id: BlockVirtualId, + /// Index of the first particle belonging to this block in the sorted array. + pub first_particle: u32, + /// Number of particles whose primary (base) block is this block. + pub num_particles: u32, + /// Total number of particles contributing to this block, including those whose + /// quadratic stencil only spills in from a neighbouring block ("extras"). + pub num_particles_with_extras: u32, + /// Index of the first rigid particle belonging to this block in the sorted rigid + /// particle array. + pub first_rigid_particle: u32, + /// Total number of rigid particles contributing to this block, extras included. + pub num_rigid_particles_with_extras: u32, + /// Per-slab-bucket cursors for the within-block counting sort of regular + /// particles. + pub sort_bucket_cursors: [u32; NUM_SORT_BUCKETS], + /// Header IDs of adjacent blocks to avoid repeated hmap lookup + /// in particle sorts. + pub nbh_block_ids: [BlockHeaderId; NUM_NBH_BLOCKS], + /// Padding. + pub padding: [u32; 2], +} + +/// Top-level grid metadata. +/// +/// Contains the current number of active blocks and configuration parameters. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct Grid { + /// Current number of active blocks (modified atomically during insertion). + pub num_active_blocks: u32, + /// The uniform cell width (grid spacing). + pub cell_width: f32, + /// Capacity of the hashmap (must be a power of 2). + pub hmap_capacity: u32, + /// Maximum number of blocks that can be stored. + pub capacity: u32, +} + +/// Contact distance field data stored per grid node. +/// +/// Carries what CPIC (Compatible Particle-In-Cell) rigid coupling needs: a +/// signed distance and the affinity bits its compatibility checks compare. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct NodeCdf { + /// Signed distance to the closest collider surface. + pub distance: f32, + /// Affinity bits: lower 16 bits are affinity flags, upper 16 bits are sign flags. + /// Two bits per collider. + pub affinities: AffinityBits, + /// Index of the closest collider, or `NONE` if no collider is nearby. + pub closest_id: u32, +} + +impl NodeCdf { + pub const NONE: NodeCdf = NodeCdf { + distance: MAX_FLT, + affinities: AffinityBits(0), + closest_id: NONE, + }; + + /// Creates a new `NodeCdf` with the given values. + #[inline] + pub fn new(distance: f32, affinities: AffinityBits, closest_id: u32) -> Self { + Self { + distance, + affinities, + closest_id, + } + } +} + +/// A single grid node's state. +/// +/// Stores momentum/velocity packed with mass, plus CDF data for rigid body coupling. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct Node { + /// Contains either momentum or velocity (depending on context). + pub momentum_velocity: Vector, + /// The node’s mass. + #[cfg(feature = "dim3")] // The field ordering is different in 2D and 3D to reduce padding. + pub mass: f32, + /// Momentum/velocity for particles that are incompatible with this node + /// (per CPIC's affinity-based compatibility). This ensures P2G/G2P transfers + /// on incompatible nodes still work properly without losing contributions from + /// other compatible particles. + pub momentum_velocity_incompatible: Vector, + /// The node’s mass. + #[cfg(feature = "dim2")] // The field ordering is different in 2D and 3D to reduce padding. + pub mass: f32, + /// Mass for particles that are incompatible with this node. + pub mass_incompatible: f32, + /// Contact distance field data for rigid body coupling. + pub cdf: NodeCdf, + /// SPIR-V padding. + pub _padding: u32, +} + +/* + * Hashmap functions. + */ + +/// Computes a Murmur3-based hash of a packed key, which the hashmap probes +/// from. +#[inline] +fn hash(packed_key: u32) -> u32 { + let mut key = packed_key; + key = key.wrapping_mul(0xCC9E2D51); + key = key.rotate_left(15); + key = key.wrapping_mul(0x1B873593); + key +} + +// TODO: refactor the hash-map code into something that doesn’t depends on the grid types? +impl Grid { + /// Attempts to insert a block into the hashmap using atomic compare-exchange. + /// + /// Returns the slot index if a new entry was created, or `NONE` if the key + /// already exists or the hashmap is full. Handles weak-CAS semantics (as found + /// on WebGPU/WGSL/Metal targets), so a single unique thread is elected as the + /// inserter even under concurrent same-key insertions. + #[inline] + pub fn insertion_index( + &self, + hmap_entries: &mut [GridHashMapEntry], + key: &BlockVirtualId, + ) -> u32 { + let packed_key = key.pack(); + let mut slot = hash(packed_key) & (self.hmap_capacity - 1); + + // NOTE: if there is no more room in the hashmap to store the data, we just do nothing. + // It is up to the user to detect the high occupancy, resize the hashmap, and re-run + // the failed insertion. + for _ in 0..self.hmap_capacity { + let old_value = khal_std::sync::atomic_compare_exchange_u32( + &mut hmap_entries.at_mut(slot as usize).state, + packed_key, + NONE, + ); + + if old_value == packed_key { + // The entry already exists. + return NONE; + } + + if old_value != NONE { + // Slot occupied by a different key. Probe next. + slot = (slot + 1) & (self.hmap_capacity - 1); + continue; + } + + // CAS returned NONE. Either we wrote successfully, or it was a spurious + // failure (weak CAS on WGSL/Metal). Verify with atomic_load (which is always strong). + let current = + khal_std::sync::atomic_load_u32_shared(&hmap_entries.at(slot as usize).state); + + if current == packed_key { + // Slot contains our key (we wrote it, or a same-key thread did). + // Use atomic_exchange on ownership to determine the unique owner. + // atomic_exchange is always strong (no weak variant in WGSL). + hmap_entries.at_mut(slot as usize).key = *key; + let prev = khal_std::sync::atomic_exchange_u32( + &mut hmap_entries.at_mut(slot as usize).ownership, + 1, + ); + if prev == 0 { + return slot; // We are the owner (new insertion). + } + return NONE; // Another thread owns this slot. + } + + if current != NONE { + // A different key was written between our CAS and load. Probe next. + slot = (slot + 1) & (self.hmap_capacity - 1); + continue; + } + + // current == NONE: spurious CAS failure. Retry the same slot on the + // next iteration (slot is not advanced). This wastes one iteration of + // the capacity-bounded loop but spurious failures are extremely rare. + } + + NONE + } + + /// Looks up a block's header ID in the hashmap. + /// + /// Returns the `BlockHeaderId` for the given virtual block coordinate, + /// or a `BlockHeaderId` with `id == NONE` if the block is not active. + #[inline] + pub fn find_block_header_id( + &self, + hmap_entries: &[GridHashMapEntry], + key: &BlockVirtualId, + ) -> BlockHeaderId { + let packed_key = key.pack(); + let capacity = self.hmap_capacity; + let mut slot = hash(packed_key) & (capacity - 1); + + for _ in 0..capacity { + let state = hmap_entries.at(slot as usize).state; + if state == packed_key { + return hmap_entries.at(slot as usize).value; + } else if state == NONE { + break; + } + + slot = (slot + 1) & (capacity - 1); + } + + BlockHeaderId { id: NONE } + } + + /// Marks a block as active by inserting it into the hashmap and allocating a header. + /// + /// If the block is successfully inserted (i.e., it was not already active), + /// a new `ActiveBlockHeader` entry is created and the hashmap entry is linked + /// to it via an atomically-assigned header ID. + #[inline] + pub fn mark_block_as_active( + &mut self, + hmap_entries: &mut [GridHashMapEntry], + active_blocks: &mut [ActiveBlockHeader], + block: &BlockVirtualId, + ) { + let slot = self.insertion_index(hmap_entries, block); + + if slot != NONE { + let block_header_id = atomic_add_u32(&mut self.num_active_blocks, 1); + active_blocks.at_mut(block_header_id as usize).virtual_id = *block; + active_blocks + .at_mut(block_header_id as usize) + .first_particle = 0; + active_blocks.at_mut(block_header_id as usize).num_particles = 0; + active_blocks + .at_mut(block_header_id as usize) + .num_particles_with_extras = 0; + active_blocks + .at_mut(block_header_id as usize) + .first_rigid_particle = 0; + active_blocks + .at_mut(block_header_id as usize) + .num_rigid_particles_with_extras = 0; + for k in 0..NUM_SORT_BUCKETS { + active_blocks + .at_mut(block_header_id as usize) + .sort_bucket_cursors + .write(k, 0); + } + hmap_entries.at_mut(slot as usize).value = BlockHeaderId { + id: block_header_id, + }; + } + } +} + +/* + * Affinity functions for CPIC. + */ + +/// Affinity bits: lower 16 bits are affinity flags, upper 16 bits are sign flags. +/// Two bits per collider. +#[derive(Clone, Copy, Default)] +#[cfg_attr( + not(target_arch_is_gpu), + derive(Debug, PartialEq, bytemuck::Pod, bytemuck::Zeroable) +)] +#[repr(C)] +pub struct AffinityBits(pub u32); + +impl AffinityBits { + pub const EMPTY: AffinityBits = AffinityBits(0); + /// Mask for the lower 16 affinity bits in the CDF affinity field. + pub const AFFINITY_BITS_MASK: u32 = 0x0000FFFF; + /// Bit shift to access the sign bits in the upper 16 bits of the affinity field. + pub const SIGN_BITS_SHIFT: u32 = 16; + + /// Checks if a specific collider's affinity bit is set. + #[inline] + pub fn bit(self, i_collider: u32) -> bool { + (self.0 & (1 << i_collider)) != 0 + } + + /// Checks if a specific collider's sign bit is set. + #[inline] + pub fn sign_bit(self, i_collider: u32) -> bool { + ((self.0 >> Self::SIGN_BITS_SHIFT) & (1 << i_collider)) != 0 + } + + pub fn set_unsigned_bits(&mut self, other: Self) { + self.0 |= other.0 & Self::AFFINITY_BITS_MASK; + } + + pub fn set_bit(&mut self, i_collider: u32, signed: bool) { + if signed { + self.0 |= 0x00010001u32 << i_collider; + } else { + self.0 |= 0x00000001u32 << i_collider; + } + } + + pub fn set_sign_bit(&mut self, i_collider: u32) { + self.0 |= 0x00010000u32 << i_collider; + } + + pub fn or_sign_bit(&mut self, affinity2: Self, i_collider: u32) { + self.0 |= affinity2.0 & (0x00010000u32 << i_collider); + } + + /// Checks if two affinity fields are compatible (same sign for all shared affinities). + /// + /// Two nodes/particles are compatible if, for every collider they both have affinity to, + /// they agree on the sign (i.e., they are on the same side of the collider surface). + #[inline] + pub fn is_compatible(self, affinity2: Self) -> bool { + let affinities_in_common = self.0 & affinity2.0 & Self::AFFINITY_BITS_MASK; + let signs1 = (self.0 >> Self::SIGN_BITS_SHIFT) & affinities_in_common; + let signs2 = (affinity2.0 >> Self::SIGN_BITS_SHIFT) & affinities_in_common; + signs1 == signs2 + } +} + +impl BitOrAssign for AffinityBits { + fn bitor_assign(&mut self, rhs: Self) { + self.0 |= rhs.0; + } +} + +/* + * Entry points. + */ + +/// Resets all hashmap entries to the empty state and clears the active block count. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_reset_hmap( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] grid_data: &mut Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + hmap_entries: &mut [GridHashMapEntry], +) { + let id = invocation_id.x; + + if id < grid_data.hmap_capacity { + let entry = hmap_entries.at_mut(id as usize); + entry.state = NONE; + // Reset ownership so the next frame's insertions can claim slots. + entry.ownership = 0; + // Resetting the following isn't necessary for correctness, + // but it makes debugging easier. + entry.key = BlockVirtualId { + id: IVector::ZERO, + #[cfg(feature = "dim3")] + padding: 0, + }; + entry.value = BlockHeaderId { id: 0 }; + } + if id == 0 { + grid_data.num_active_blocks = 0; + } +} + +/// Snapshots the current `num_active_blocks` into a single-element buffer. +/// +/// This is the barrier in the two-pass block activation: it freezes the base +/// block count after `gpu_touch_primary_blocks`, so `gpu_touch_neighbor_blocks` +/// iterates over base blocks only and not over the neighbours it appends. +#[spirv_bindgen] +#[spirv(compute(threads(1)))] +pub fn gpu_capture_num_active_blocks( + #[spirv(global_invocation_id)] _invocation_id: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] num_base_blocks: &mut [u32], +) { + num_base_blocks.write(0, grid.num_active_blocks); +} + +/// Computes indirect dispatch sizes based on the number of active blocks. +/// +/// Produces two sets of dispatch arguments: +/// - `n_block_groups`: for per-block dispatches (ceil(num_active_blocks / GRID_WORKGROUP_SIZE)) +/// - `n_g2p_p2g_groups`: for P2G/G2P dispatches (one workgroup per active block) +#[spirv_bindgen] +#[spirv(compute(threads(1)))] +pub fn gpu_init_indirect_workgroups( + #[spirv(global_invocation_id)] _invocation_id: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] n_block_groups: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] n_g2p_p2g_groups: &mut [u32], +) { + let num_active_blocks = grid.num_active_blocks; + n_block_groups.write(0, udiv_ceil(num_active_blocks, GRID_WORKGROUP_SIZE)); + n_block_groups.write(1, 1); + n_block_groups.write(2, 1); + n_g2p_p2g_groups.write(0, num_active_blocks); + n_g2p_p2g_groups.write(1, 1); + n_g2p_p2g_groups.write(2, 1); +} + +/// Resets all grid nodes for the current set of active blocks (momentum, +/// velocity, mass, and CDF data). +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_reset( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] nodes: &mut [Node], +) { + let i = invocation_id.x; + let num_nodes = grid.num_active_blocks * NUM_CELL_PER_BLOCK; + if i < num_nodes { + let idx = i as usize; + let node = nodes.at_mut(idx); + node.momentum_velocity = Vector::ZERO; + node.mass = 0.0; + node.momentum_velocity_incompatible = Vector::ZERO; + node.mass_incompatible = 0.0; + node.cdf = NodeCdf::NONE; + } +} diff --git a/src_mpm_shaders/grid/kernel.rs b/src_mpm_shaders/grid/kernel.rs new file mode 100644 index 00000000..b18fa7e6 --- /dev/null +++ b/src_mpm_shaders/grid/kernel.rs @@ -0,0 +1,180 @@ +//! Quadratic B-spline kernel for MPM grid transfers. +//! +//! Provides the quadratic B-spline weights and the neighbour stencil. The kernel +//! spans 3 nodes per dimension, giving 9 neighbors in 2D and 27 in 3D. + +use crate::{DIM_USIZE, Vector, abs}; +use glamx::*; + +/* + * Neighborhood stencil constants. + * + * These define the iteration order and shared-memory indices for the + * 3x3 (2D) or 3x3x3 (3D) neighborhood around a particle. + */ + +/// Number of neighbor nodes in the kernel stencil. +#[cfg(feature = "dim2")] +pub const NBH_LEN: usize = 9; +/// Number of neighbor nodes in the kernel stencil. +#[cfg(feature = "dim3")] +pub const NBH_LEN: usize = 27; + +/// Returns the stencil offset for neighbor `i` as a UVector. +/// +/// In 2D, these are UVec2 offsets into a 3x3 grid. +/// In 3D, these are UVec3 offsets into a 3x3x3 grid. +#[cfg(feature = "dim2")] +pub const NBH_SHIFTS: [UVec2; 9] = [ + UVec2::new(2, 2), + UVec2::new(2, 0), + UVec2::new(2, 1), + UVec2::new(0, 2), + UVec2::new(0, 0), + UVec2::new(0, 1), + UVec2::new(1, 2), + UVec2::new(1, 0), + UVec2::new(1, 1), +]; + +/// Returns the stencil offset for neighbor `i` as a UVector. +#[cfg(feature = "dim3")] +pub const NBH_SHIFTS: [UVec3; 27] = [ + UVec3::new(2, 2, 2), + UVec3::new(2, 0, 2), + UVec3::new(2, 1, 2), + UVec3::new(0, 2, 2), + UVec3::new(0, 0, 2), + UVec3::new(0, 1, 2), + UVec3::new(1, 2, 2), + UVec3::new(1, 0, 2), + UVec3::new(1, 1, 2), + UVec3::new(2, 2, 0), + UVec3::new(2, 0, 0), + UVec3::new(2, 1, 0), + UVec3::new(0, 2, 0), + UVec3::new(0, 0, 0), + UVec3::new(0, 1, 0), + UVec3::new(1, 2, 0), + UVec3::new(1, 0, 0), + UVec3::new(1, 1, 0), + UVec3::new(2, 2, 1), + UVec3::new(2, 0, 1), + UVec3::new(2, 1, 1), + UVec3::new(0, 2, 1), + UVec3::new(0, 0, 1), + UVec3::new(0, 1, 1), + UVec3::new(1, 2, 1), + UVec3::new(1, 0, 1), + UVec3::new(1, 1, 1), +]; + +/// Flattens the 2D/3D stencil offset of neighbor `i` into a workgroup +/// shared-memory index. +#[cfg(feature = "dim2")] +pub const NBH_SHIFT_SHARED: [u32; 9] = [22, 2, 12, 20, 0, 10, 21, 1, 11]; +#[cfg(feature = "dim3")] +pub const NBH_SHIFT_SHARED: [u32; 27] = [ + 86, 74, 80, 84, 72, 78, 85, 73, 79, 14, 2, 8, 12, 0, 6, 13, 1, 7, 50, 38, 44, 48, 36, 42, 49, + 37, 43, +]; + +/// Extracts a component from a Vec3 by dynamic index. +/// +/// This avoids `Vec3::Index` which generates SPIR-V pointer phi nodes +/// (requiring the VariablePointers capability). Instead, this function +/// produces value phi nodes which are always valid in SPIR-V. +#[inline] +pub fn vec3_extract(v: Vec3, index: u32) -> f32 { + if index == 0 { + v.x + } else if index == 1 { + v.y + } else { + v.z + } +} + +/// Quadratic B-spline kernel. +/// +/// This kernel provides quadratic (degree 2) B-spline basis functions for the +/// MPM particle-grid transfers. Each basis function has support over 3 cells, +/// giving a smooth C1-continuous interpolation. +pub struct QuadraticKernel; + +impl QuadraticKernel { + /// Computes the inverse of the D matrix diagonal for APIC transfers. + /// + /// For the quadratic B-spline, `inv_d = 4 / h^2` where `h` is the cell width. + #[inline] + pub fn inv_d(cell_width: f32) -> f32 { + 4.0 / (cell_width * cell_width) + } + + /// Evaluates all three quadratic B-spline basis functions at position `x`, + /// returning `Vec3(w_left, w_center, w_right)`. `x` is the distance from the + /// associated (leftmost) grid node, in cell units. + #[inline] + pub fn eval_all(x: f32) -> Vec3 { + Vec3::new( + 0.5 * (1.5 - x) * (1.5 - x), + 0.75 - (x - 1.0) * (x - 1.0), + 0.5 * (x - 0.5) * (x - 0.5), + ) + } + + /// Evaluates a single quadratic B-spline basis function at position `x`. + /// + /// Uses absolute value of `x` and selects the appropriate piece of the + /// piecewise-quadratic function based on distance from center. + #[inline] + pub fn eval(x: f32) -> f32 { + // Branchless: rust-gpu lowers `if/else if/else` to real branches, which are + // expensive in the per-particle inner loop of the transfer kernels. Computing both + // pieces and masking compiles to selects/arithmetic instead (matching what the + // slang version produces). + let x_abs = abs(x); + let part1 = 0.75 - x_abs * x_abs; + let part2 = 0.5 * (1.5 - x_abs) * (1.5 - x_abs); + let lt05 = (x_abs < 0.5) as u32 as f32; + let lt15 = (x_abs < 1.5) as u32 as f32; + lt05 * part1 + (1.0 - lt05) * lt15 * part2 + } + + /// Evaluates the derivative of a single quadratic B-spline basis function at position `x`. + #[inline] + pub fn eval_derivative(x: f32) -> f32 { + // Branchless, see `eval`. + let x_abs = abs(x); + let sign = (x >= 0.0) as u32 as f32 * 2.0 - 1.0; + let part1 = -2.0 * sign * x_abs; + let part2 = -sign * (1.5 - x_abs); + let lt05 = (x_abs < 0.5) as u32 as f32; + let lt15 = (x_abs < 1.5) as u32 as f32; + lt05 * part1 + (1.0 - lt05) * lt15 * part2 + } + + /// Precomputes all kernel weights for a particle at position `ref_pos` relative + /// to the associated grid node, with cell width `h`. + /// + /// Returns an array of `DIM_USIZE` Vec3 values, one per spatial dimension. + /// Each Vec3 contains the three basis function weights for that dimension. + #[inline] + pub fn precompute_weights(ref_elt_pos_minus_particle_pos: Vector, h: f32) -> [Vec3; DIM_USIZE] { + #[cfg(feature = "dim2")] + { + [ + Self::eval_all(-ref_elt_pos_minus_particle_pos.x / h), + Self::eval_all(-ref_elt_pos_minus_particle_pos.y / h), + ] + } + #[cfg(feature = "dim3")] + { + [ + Self::eval_all(-ref_elt_pos_minus_particle_pos.x / h), + Self::eval_all(-ref_elt_pos_minus_particle_pos.y / h), + Self::eval_all(-ref_elt_pos_minus_particle_pos.z / h), + ] + } + } +} diff --git a/src_mpm_shaders/grid/mod.rs b/src_mpm_shaders/grid/mod.rs new file mode 100644 index 00000000..47945849 --- /dev/null +++ b/src_mpm_shaders/grid/mod.rs @@ -0,0 +1,3 @@ +pub mod grid; +pub mod kernel; +pub mod sort; diff --git a/src_mpm_shaders/grid/sort.rs b/src_mpm_shaders/grid/sort.rs new file mode 100644 index 00000000..f4d1858f --- /dev/null +++ b/src_mpm_shaders/grid/sort.rs @@ -0,0 +1,667 @@ +//! Particle sorting kernels for the sparse MPM grid. +//! +//! These kernels handle: +//! 1. Marking blocks as active based on particle positions. +//! 2. Counting particles per block. +//! 3. Building sorted particle arrays. +//! +//! The sorting pipeline runs in multiple passes: +//! 1. `touch_particle_blocks` / `touch_rigid_particle_blocks` - mark active blocks +//! 2. `mark_rigid_particles_needing_block` - flag rigid particles near block boundaries +//! 3. `update_block_particle_count` - count particles per active block +//! 4. `copy_particles_len_to_scan_value` - prepare prefix sum input +//! 5. prefix sum (external) - compute exclusive scan of particle counts +//! 6. `copy_scan_values_to_first_particles` - write back sorted offsets +//! 7. `finalize_particles_sort` - place particles in sorted order +//! +//! Rigid particles go through the same count / prefix-sum / finalize sequence +//! (`update_block_rigid_particle_count`, `copy_rigid_particles_len_to_scan_value`, +//! `copy_scan_values_to_first_rigid_particles`, `finalize_rigid_particles_sort`), +//! reusing the scan workspace after the regular particle sort completed. + +// Indexed loops on purpose throughout this module: iterator-based loops over +// storage buffers are fragile under rust-gpu's SPIR-V codegen. +#![allow(clippy::needless_range_loop)] +use crate::grid::grid::*; +use crate::solver::particle::Position; +use crate::{IVector, UVector}; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use khal_std::sync::atomic_add_u32; + +#[cfg(feature = "dim2")] +const EXTRA_PARTICLE_MIN_SHIFT: u32 = 6; +#[cfg(feature = "dim3")] +const EXTRA_PARTICLE_MIN_SHIFT: u32 = 2; + +/// Returns the within-block sort bucket for a particle counted/inserted into its +/// primary block: one bucket per associated-cell slab along the slowest-varying node +/// axis (y in 2D, z in 3D). +#[inline] +fn primary_sort_bucket(assoc: UVector) -> usize { + #[cfg(feature = "dim2")] + { + assoc.y as usize + } + #[cfg(feature = "dim3")] + { + assoc.z as usize + } +} + +/// Returns the within-block sort bucket for a particle counted/inserted as an "extra" +/// into the neighbour block shifted by `bshift` from its primary block. Slabs below +/// -2 are clamped into the -2 bucket. +#[inline] +fn extra_sort_bucket(assoc: UVector, bshift: IVector) -> usize { + #[cfg(feature = "dim2")] + let local = assoc.y as i32 - bshift.y * 8; + #[cfg(feature = "dim3")] + let local = assoc.z as i32 - bshift.z * 4; + NUM_PRIMARY_SORT_BUCKETS + (local.max(-2) + 2) as usize +} + +/// Marks all blocks associated with each particle as active. +/// +/// For each particle, computes the set of blocks whose stencil could overlap +/// the particle, and inserts them into the hashmap. This must be run before +/// any per-block operations. +// TODO HACK: enabling spirv-passthrough for this shader since naga panics +// on the spv backend because of https://github.com/gfx-rs/wgpu/issues/7315 +// (in our case, it’s caused by the lines involving the atomic compare-exchange). +#[spirv_bindgen(spirv_passthrough)] +#[spirv(compute(threads(64)))] +pub fn gpu_touch_particle_blocks( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] grid: &mut Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + hmap_entries: &mut [GridHashMapEntry], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + active_blocks: &mut [ActiveBlockHeader], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] particles_pos: &[Position], + #[spirv(uniform, descriptor_set = 0, binding = 4)] particles_len: &u32, +) { + let id = invocation_id.x; + if id < *particles_len { + let cell_width = grid.cell_width; + let particle = particles_pos.read(id as usize); + let blocks = BlockVirtualId::blocks_associated_to_point(cell_width, particle.pt); + for i in 0..NUM_ASSOC_BLOCKS { + grid.mark_block_as_active(hmap_entries, active_blocks, &blocks[i]); + } + } +} + +/// Marks only each particle's **primary** (base) block as active. +/// +/// First half of the two-pass block activation (with `gpu_touch_neighbor_blocks`) +/// that replaces `gpu_touch_particle_blocks`; the union of activated blocks is +/// identical. +// TODO HACK: spirv_passthrough because naga panics on the atomic compare-exchange +// in `mark_block_as_active` (see `gpu_touch_particle_blocks`). +#[spirv_bindgen(spirv_passthrough)] +#[spirv(compute(threads(64)))] +pub fn gpu_touch_primary_blocks( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] grid: &mut Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + hmap_entries: &mut [GridHashMapEntry], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + active_blocks: &mut [ActiveBlockHeader], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] particles_pos: &[Position], + #[spirv(uniform, descriptor_set = 0, binding = 4)] particles_len: &u32, +) { + let id = invocation_id.x; + if id < *particles_len { + let cell_width = grid.cell_width; + let particle = particles_pos.read(id as usize); + let block = BlockVirtualId::block_associated_to_point(cell_width, particle.pt); + grid.mark_block_as_active(hmap_entries, active_blocks, &block); + } +} + +/// Marks the +1 neighbour blocks of every already-active base block as active. +/// +/// Second half of the two-pass block activation (see `gpu_touch_primary_blocks`). +/// `num_base_blocks` must be a snapshot of `grid.num_active_blocks` taken *before* +/// this pass runs, so that the neighbour blocks appended during the pass are not +/// themselves processed. +// TODO HACK: spirv_passthrough because naga panics on the atomic compare-exchange +// in `mark_block_as_active` (see `gpu_touch_particle_blocks`). +#[spirv_bindgen(spirv_passthrough)] +#[spirv(compute(threads(64)))] +pub fn gpu_touch_neighbor_blocks( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] grid: &mut Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + hmap_entries: &mut [GridHashMapEntry], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + active_blocks: &mut [ActiveBlockHeader], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] num_base_blocks: &[u32], +) { + let id = invocation_id.x; + if id < num_base_blocks.read(0) { + let raw = active_blocks.at(id as usize).virtual_id; + // Rematerialize the block id into a fresh, register-aligned vector. Reading + // `IVec3` straight out of the storage buffer keeps naga's packed-`int3` type, + // and the vector arithmetic in `blocks_associated_to_block` then emits illegal + // packed↔aligned `as_type` casts under the spirv-passthrough path (Metal). + #[cfg(feature = "dim2")] + let vid = BlockVirtualId { + id: IVector::new(raw.id.x, raw.id.y), + }; + #[cfg(feature = "dim3")] + let vid = BlockVirtualId { + id: IVector::new(raw.id.x, raw.id.y, raw.id.z), + padding: 0, + }; + let blocks = BlockVirtualId::blocks_associated_to_block(&vid); + for i in 1..NUM_ASSOC_BLOCKS { + grid.mark_block_as_active(hmap_entries, active_blocks, &blocks[i]); + } + } +} + +/// Marks all blocks associated with each rigid particle as active. +/// +/// Similar to `gpu_touch_particle_blocks`, but operates on rigid body surface +/// particles. Only touches blocks for rigid particles that are flagged as needing +/// a block (via the `rigid_particle_needs_block` bitfield). +// TODO HACK: enabling spirv-passthrough for this shader since naga panics +// on the spv backend because of https://github.com/gfx-rs/wgpu/issues/7315 +// (in our case, it’s caused by the lines involving the atomic compare-exchange). +#[spirv_bindgen(spirv_passthrough)] +#[spirv(compute(threads(64)))] +pub fn gpu_touch_rigid_particle_blocks( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] grid: &mut Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + hmap_entries: &mut [GridHashMapEntry], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + active_blocks: &mut [ActiveBlockHeader], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] rigid_particles_pos: &[Position], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] rigid_particle_needs_block: &[u32], +) { + let id = invocation_id.x; + if id < rigid_particles_pos.len() as u32 { + let cell_width = grid.cell_width; + let entry_id = (id / 32) as usize; + let entry_bit = 1u32 << (id % 32); + let needs_block = (rigid_particle_needs_block.read(entry_id) & entry_bit) != 0; + + if needs_block { + let particle = rigid_particles_pos.read(id as usize); + let block = BlockVirtualId::block_associated_to_point(cell_width, particle.pt); + grid.mark_block_as_active(hmap_entries, active_blocks, &block); + } + } +} + +/// Flags rigid particles that need their own block activated. +/// +/// A rigid particle needs its own block if at least one (but not all) of its +/// associated blocks are already active. This means the particle is near a +/// block boundary and its contributions would be lost without an additional block. +/// +/// The result is stored as a bitfield in `rigid_particle_needs_block`, where +/// each u32 holds flags for 32 rigid particles. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mark_rigid_particles_needing_block( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] hmap_entries: &[GridHashMapEntry], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] rigid_particles_pos: &[Position], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] + rigid_particle_needs_block: &mut [u32], +) { + let id = invocation_id.x; + if id < rigid_particles_pos.len() as u32 { + let cell_width = grid.cell_width; + let particle = rigid_particles_pos.read(id as usize); + let blocks = BlockVirtualId::blocks_associated_to_point(cell_width, particle.pt); + + // Find the first block that already has a header in the hashmap. + let mut i = 0u32; + for _ in 0..NUM_ASSOC_BLOCKS { + if grid + .find_block_header_id(hmap_entries, &blocks[i as usize]) + .id + != NONE + { + break; + } + i += 1; + } + + let entry_id = (id / 32) as usize; + let entry_bit = 1u32 << (id % 32); + + // If some but not all associated blocks are active, the rigid particle + // needs its own block to ensure proper grid transfers. + if i > 0 && i < NUM_ASSOC_BLOCKS as u32 { + // Set the bit atomically. + khal_std::sync::atomic_or_u32(rigid_particle_needs_block.at_mut(entry_id), entry_bit); + } else { + // Clear the bit atomically. + khal_std::sync::atomic_and_u32(rigid_particle_needs_block.at_mut(entry_id), !entry_bit); + } + } +} + +/// Precomputes, for each active block, the header IDs of its +1 neighbour blocks, +/// caching them in `ActiveBlockHeader::nbh_block_ids`. Inactive neighbours are +/// stored as `NONE`. +/// +/// Must run after all blocks have been touched (so `num_active_blocks` is final and every +/// neighbour that exists is in the hashmap) and before the particle/rigid count and +/// finalize passes that read `nbh_block_ids`. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_update_nbh_block_ids( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] hmap_entries: &[GridHashMapEntry], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + active_blocks: &mut [ActiveBlockHeader], +) { + let id = invocation_id.x; + if id < grid.num_active_blocks { + let block0 = active_blocks.at_mut(id as usize); + let vid = &block0.virtual_id; + let assoc = BlockVirtualId::blocks_associated_to_block(vid); + for nbh in 0..NUM_ASSOC_BLOCKS - 1 { + let nbh_vid = assoc[nbh + 1]; + let nbh_hid = grid.find_block_header_id(hmap_entries, &nbh_vid); + block0.nbh_block_ids.write(nbh, nbh_hid); + } + } +} + +/// Counts the number of particles in each active block. +/// +/// Each thread processes one particle, finds its associated block, and +/// atomically increments that block's `num_particles` counter. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_update_block_particle_count( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] hmap_entries: &[GridHashMapEntry], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] particles_pos: &[Position], + #[spirv(uniform, descriptor_set = 0, binding = 3)] particles_len: &u32, + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] + active_blocks: &mut [ActiveBlockHeader], +) { + let id = invocation_id.x; + if id < *particles_len { + let cell_width = grid.cell_width; + let particle = particles_pos.read(id as usize); + let blocks = BlockVirtualId::blocks_associated_to_point(cell_width, particle.pt); + + // The particle's primary (base) block gets it as a regular particle and + // as an "extra". Only the per-slab-bucket counter is incremented here: + // `num_particles` and `num_particles_with_extras` are derived per block in + // the copy passes, avoiding two heavily contended atomics per particle. + let assoc = particle.associated_cell_index_in_block_off_by_one(cell_width); + let block0 = grid.find_block_header_id(hmap_entries, &blocks[0]); + atomic_add_u32( + active_blocks + .at_mut(block0.id as usize) + .sort_bucket_cursors + .at_mut(primary_sort_bucket(assoc)), + 1, + ); + + // Each +1 neighbour block also receives the particle as an "extra" if the + // quadratic stencil actually spills into it, i.e. the local base-cell index is + // >= EXTRA_PARTICLE_MIN_SHIFT along every axis where that block is the +1 neighbour. + let id0 = blocks[0].id; + for i in 1..NUM_ASSOC_BLOCKS { + let bshift = blocks[i].id - id0; + #[cfg(feature = "dim2")] + let spills = (bshift.x == 0 || assoc.x >= EXTRA_PARTICLE_MIN_SHIFT) + && (bshift.y == 0 || assoc.y >= EXTRA_PARTICLE_MIN_SHIFT); + #[cfg(feature = "dim3")] + let spills = (bshift.x == 0 || assoc.x >= EXTRA_PARTICLE_MIN_SHIFT) + && (bshift.y == 0 || assoc.y >= EXTRA_PARTICLE_MIN_SHIFT) + && (bshift.z == 0 || assoc.z >= EXTRA_PARTICLE_MIN_SHIFT); + if spills { + // The header IDs of the +1 neighbour blocks were precomputed by + // `gpu_update_nbh_block_ids`, so we read them from the primary block + // instead of doing a hashmap lookup per particle. Only the extra slab + // bucket is incremented; the neighbour's `num_particles_with_extras` is + // recovered as the sum of its buckets in the copy passes. + let block_i = active_blocks + .at(block0.id as usize) + .nbh_block_ids + .read(i - 1); + atomic_add_u32( + active_blocks + .at_mut(block_i.id as usize) + .sort_bucket_cursors + .at_mut(extra_sort_bucket(assoc, bshift)), + 1, + ); + } + } + } +} + +/// Copies each active block's particle count into the scan_values buffer. +/// +/// This prepares the input for the prefix sum pass. After the prefix sum, +/// `scan_values[i]` will contain the global offset for block `i`'s particles. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_copy_particles_len_to_scan_value( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] active_blocks: &[ActiveBlockHeader], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] scan_values: &mut [u32], +) { + let id = invocation_id.x; + if id < grid.num_active_blocks { + // The sorted array reserves room for every particle a block touches, extras + // included. `num_particles_with_extras` is the sum of all slab buckets (the count + // pass no longer maintains it as a separate atomic). + let mut total = 0u32; + for k in 0..NUM_SORT_BUCKETS { + total += active_blocks.at(id as usize).sort_bucket_cursors.read(k); + } + scan_values.write(id as usize, total); + } +} + +/// Writes the prefix sum results back as `first_particle` offsets and resets particle counts. +/// +/// After the prefix sum, `scan_values[i]` contains the exclusive scan result. +/// This kernel copies it into `active_blocks[i].first_particle`. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_copy_scan_values_to_first_particles( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] scan_values: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + active_blocks: &mut [ActiveBlockHeader], +) { + let id = invocation_id.x; + if id < grid.num_active_blocks { + let idx = id as usize; + let first = scan_values.read(idx); + active_blocks.at_mut(idx).first_particle = first; + // Convert the per-bucket counts accumulated by the count pass into running + // insertion cursors. The cursors are *absolute* offsets into the sorted array + // (i.e. `first_particle` is baked in), so the finalize pass can scatter each + // particle with a single atomic and no per-contribution `first_particle` read. + // Primary buckets come first, so primaries land in + // [first_particle, first_particle + num_particles) as G2P expects, with the + // extras after them; both segments end up ordered by slab key. + // + // The running total advanced past the primary buckets is `num_particles` (primaries + // only land in primary buckets), and the grand total is `num_particles_with_extras`. + // Both fields are derived here rather than maintained as per-particle atomics in + // the count pass. + let mut running = first; + for k in 0..NUM_SORT_BUCKETS { + if k == NUM_PRIMARY_SORT_BUCKETS { + active_blocks.at_mut(idx).num_particles = running - first; + } + let count = active_blocks.at(idx).sort_bucket_cursors.read(k); + active_blocks + .at_mut(idx) + .sort_bucket_cursors + .write(k, running); + running += count; + } + active_blocks.at_mut(idx).num_particles_with_extras = running - first; + } +} + +/// Places particles into their sorted positions. +/// +/// Each thread processes one particle: +/// 1. Finds the particle's active block via the hashmap. +/// 2. Atomically claims a slot in the sorted array (using `scan_values` as a counter). +/// 3. Writes the particle's original index into `sorted_particle_ids`. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_finalize_particles_sort( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] hmap_entries: &[GridHashMapEntry], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] particles_pos: &[Position], + #[spirv(uniform, descriptor_set = 0, binding = 3)] particles_len: &u32, + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] + active_blocks: &mut [ActiveBlockHeader], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] sorted_particle_ids: &mut [u32], +) { + let id = invocation_id.x; + if id < *particles_len { + let cell_width = grid.cell_width; + let particle = particles_pos.read(id as usize); + let blocks = BlockVirtualId::blocks_associated_to_point(cell_width, particle.pt); + + // Place the particle in its primary block's range. The prepare pass turned the + // bucket counts into absolute insertion cursors (first_particle baked in), so the + // atomically-claimed slot is already the final sorted index. + let assoc = particle.associated_cell_index_in_block_off_by_one(cell_width); + let block0 = grid.find_block_header_id(hmap_entries, &blocks[0]); + let slot0 = atomic_add_u32( + active_blocks + .at_mut(block0.id as usize) + .sort_bucket_cursors + .at_mut(primary_sort_bucket(assoc)), + 1, + ); + sorted_particle_ids.write(slot0 as usize, id); + + // Place the particle as an "extra" into each +1 neighbour block whose stencil it + // spills into, using the extra slab bucket cursors (extras land after the + // primaries because their buckets come last). + let id0 = blocks[0].id; + for i in 1..NUM_ASSOC_BLOCKS { + let bshift = blocks[i].id - id0; + #[cfg(feature = "dim2")] + let spills = (bshift.x == 0 || assoc.x >= EXTRA_PARTICLE_MIN_SHIFT) + && (bshift.y == 0 || assoc.y >= EXTRA_PARTICLE_MIN_SHIFT); + #[cfg(feature = "dim3")] + let spills = (bshift.x == 0 || assoc.x >= EXTRA_PARTICLE_MIN_SHIFT) + && (bshift.y == 0 || assoc.y >= EXTRA_PARTICLE_MIN_SHIFT) + && (bshift.z == 0 || assoc.z >= EXTRA_PARTICLE_MIN_SHIFT); + if spills { + // Reuse the neighbour header IDs precomputed by `gpu_update_nbh_block_ids` + // rather than re-querying the hashmap. + let block_i = active_blocks + .at(block0.id as usize) + .nbh_block_ids + .read(i - 1); + let slot_i = atomic_add_u32( + active_blocks + .at_mut(block_i.id as usize) + .sort_bucket_cursors + .at_mut(extra_sort_bucket(assoc, bshift)), + 1, + ); + sorted_particle_ids.write(slot_i as usize, id); + } + } + } +} + +/// Counts the number of rigid particles contributing to each active block. +/// +/// Mirrors `gpu_update_block_particle_count`, with two differences: rigid particles +/// whose primary block is not active are silently skipped (they can't affect the +/// simulation), and neighbour blocks receiving an "extra" may be inactive too (only +/// the primary block's activation is guaranteed by the touch passes). +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_update_block_rigid_particle_count( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] hmap_entries: &[GridHashMapEntry], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] rigid_particles_pos: &[Position], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] + active_blocks: &mut [ActiveBlockHeader], +) { + let id = invocation_id.x; + if id < rigid_particles_pos.len() as u32 { + let cell_width = grid.cell_width; + let particle = rigid_particles_pos.read(id as usize); + let blocks = BlockVirtualId::blocks_associated_to_point(cell_width, particle.pt); + + let block0 = grid.find_block_header_id(hmap_entries, &blocks[0]); + if block0.id != NONE { + atomic_add_u32( + &mut active_blocks + .at_mut(block0.id as usize) + .num_rigid_particles_with_extras, + 1, + ); + + // Each +1 neighbour block also receives the particle as an "extra" if its + // 3-cell influence range actually spills into it. + let assoc = particle.associated_cell_index_in_block_off_by_one(cell_width); + let id0 = blocks[0].id; + for i in 1..NUM_ASSOC_BLOCKS { + let bshift = blocks[i].id - id0; + #[cfg(feature = "dim2")] + let spills = (bshift.x == 0 || assoc.x >= EXTRA_PARTICLE_MIN_SHIFT) + && (bshift.y == 0 || assoc.y >= EXTRA_PARTICLE_MIN_SHIFT); + #[cfg(feature = "dim3")] + let spills = (bshift.x == 0 || assoc.x >= EXTRA_PARTICLE_MIN_SHIFT) + && (bshift.y == 0 || assoc.y >= EXTRA_PARTICLE_MIN_SHIFT) + && (bshift.z == 0 || assoc.z >= EXTRA_PARTICLE_MIN_SHIFT); + if spills { + // The neighbour header IDs (or NONE for inactive neighbours) were + // precomputed by `gpu_update_nbh_block_ids`. + let block_i = active_blocks + .at(block0.id as usize) + .nbh_block_ids + .read(i - 1); + if block_i.id != NONE { + atomic_add_u32( + &mut active_blocks + .at_mut(block_i.id as usize) + .num_rigid_particles_with_extras, + 1, + ); + } + } + } + } + } +} + +/// Copies each active block's rigid particle count into the scan_values buffer. +/// +/// Prepares the prefix sum input for the rigid particle sort. Must run after the +/// regular particle sort no longer needs `scan_values`. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_copy_rigid_particles_len_to_scan_value( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] active_blocks: &[ActiveBlockHeader], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] scan_values: &mut [u32], +) { + let id = invocation_id.x; + if id < grid.num_active_blocks { + scan_values.write( + id as usize, + active_blocks + .at(id as usize) + .num_rigid_particles_with_extras, + ); + } +} + +/// Writes the prefix sum results back as `first_rigid_particle` offsets. +/// +/// Also resets `num_rigid_particles_with_extras` to 0 so the finalize pass can +/// re-purpose it as the running insertion cursor (it ends up back at the count). +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_copy_scan_values_to_first_rigid_particles( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] scan_values: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + active_blocks: &mut [ActiveBlockHeader], +) { + let id = invocation_id.x; + if id < grid.num_active_blocks { + let idx = id as usize; + active_blocks.at_mut(idx).first_rigid_particle = scan_values.read(idx); + active_blocks.at_mut(idx).num_rigid_particles_with_extras = 0; + } +} + +/// Places rigid particles into their sorted positions. +/// +/// Mirrors `gpu_finalize_particles_sort` with the rigid-specific inactive-block +/// skips; each particle must contribute to the exact same set of blocks as in the +/// count pass so every claimed slot stays within its block's range. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_finalize_rigid_particles_sort( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] hmap_entries: &[GridHashMapEntry], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] rigid_particles_pos: &[Position], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] + active_blocks: &mut [ActiveBlockHeader], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] + sorted_rigid_particle_ids: &mut [u32], +) { + let id = invocation_id.x; + if id < rigid_particles_pos.len() as u32 { + let cell_width = grid.cell_width; + let particle = rigid_particles_pos.read(id as usize); + let blocks = BlockVirtualId::blocks_associated_to_point(cell_width, particle.pt); + + let block0 = grid.find_block_header_id(hmap_entries, &blocks[0]); + if block0.id != NONE { + let first0 = active_blocks.at(block0.id as usize).first_rigid_particle; + let slot0 = atomic_add_u32( + &mut active_blocks + .at_mut(block0.id as usize) + .num_rigid_particles_with_extras, + 1, + ); + sorted_rigid_particle_ids.write((first0 + slot0) as usize, id); + + let assoc = particle.associated_cell_index_in_block_off_by_one(cell_width); + let id0 = blocks[0].id; + for i in 1..NUM_ASSOC_BLOCKS { + let bshift = blocks[i].id - id0; + #[cfg(feature = "dim2")] + let spills = (bshift.x == 0 || assoc.x >= EXTRA_PARTICLE_MIN_SHIFT) + && (bshift.y == 0 || assoc.y >= EXTRA_PARTICLE_MIN_SHIFT); + #[cfg(feature = "dim3")] + let spills = (bshift.x == 0 || assoc.x >= EXTRA_PARTICLE_MIN_SHIFT) + && (bshift.y == 0 || assoc.y >= EXTRA_PARTICLE_MIN_SHIFT) + && (bshift.z == 0 || assoc.z >= EXTRA_PARTICLE_MIN_SHIFT); + if spills { + // Reuse the neighbour header IDs precomputed by `gpu_update_nbh_block_ids`. + let block_i = active_blocks + .at(block0.id as usize) + .nbh_block_ids + .read(i - 1); + if block_i.id != NONE { + let first_i = active_blocks.at(block_i.id as usize).first_rigid_particle; + let slot_i = atomic_add_u32( + &mut active_blocks + .at_mut(block_i.id as usize) + .num_rigid_particles_with_extras, + 1, + ); + sorted_rigid_particle_ids.write((first_i + slot_i) as usize, id); + } + } + } + } + } +} diff --git a/src_mpm_shaders/lib.rs b/src_mpm_shaders/lib.rs new file mode 100644 index 00000000..ba09d838 --- /dev/null +++ b/src_mpm_shaders/lib.rs @@ -0,0 +1,160 @@ +//! Nexus MPM (Material Point Method) GPU shaders. +//! +//! This crate contains Rust GPU shaders for the nexus_mpm solver, +//! providing GPU-accelerated MPM simulation. + +#![cfg_attr(target_arch_is_gpu, no_std)] +#![cfg_attr(target_arch = "spirv", feature(asm_experimental_arch))] +#![allow(clippy::too_many_arguments)] +#![allow(clippy::manual_range_contains)] +#![allow(clippy::module_inception)] +#![allow(unused_imports)] + +// Re-export the nexus_rbd_shaders crate for shared types. +#[cfg(feature = "dim2")] +pub extern crate nexus_rbd_shaders2d as nexus_rbd_shaders; +#[cfg(feature = "dim3")] +pub extern crate nexus_rbd_shaders3d as nexus_rbd_shaders; + +// Re-export parry for collision shapes. +#[cfg(feature = "dim2")] +pub extern crate parry2d as parry; +#[cfg(feature = "dim3")] +pub extern crate parry3d as parry; + +// Re-export glamx for convenience. +pub use glamx; +use glamx::*; + +// Re-export key types and utilities from nexus_rbd_shaders. +pub use nexus_rbd_shaders::{AngVector, DIM, Pad, PaddedVector, Pose, RotMatrix, Rotation, Vector}; +pub use nexus_rbd_shaders::{abs, acos, asin, atan2, cos, safe_div, sin, sqrt}; +pub use nexus_rbd_shaders::{gcross, gcross_av, gdot, maybe_inv, rotation_to_matrix}; +pub use nexus_rbd_shaders::{udiv, umod}; + +// NOTE: we disable two-ways coupling for now since it’s quite expensive and not very +// stable anyway. We’d need some deeper integration with Nexus’ solver for it +// to look really good. +pub const TWO_WAYS_COUPLING_ENABLED: bool = false; + +// +// MPM-specific type aliases +// + +/// Signed integer vector type (IVec2 in 2D, IVec3 in 3D). +#[cfg(feature = "dim2")] +pub type IVector = IVec2; +/// Signed integer vector type (IVec2 in 2D, IVec3 in 3D). +#[cfg(feature = "dim3")] +pub type IVector = IVec3; + +/// Unsigned integer vector type (UVec2 in 2D, UVec3 in 3D). +#[cfg(feature = "dim2")] +pub type UVector = UVec2; +/// Unsigned integer vector type (UVec2 in 2D, UVec3 in 3D). +#[cfg(feature = "dim3")] +pub type UVector = UVec3; + +/// The square matrix type for the current dimension (Mat2 in 2D, Mat3 in 3D). +#[cfg(feature = "dim2")] +pub type Matrix = Mat2; +/// The square matrix type for the current dimension (Mat2 in 2D, Mat3 in 3D). +#[cfg(feature = "dim3")] +pub type Matrix = Mat3; + +/// The square matrix type for the current dimension (Mat2 in 2D, Mat3 in 3D). +#[cfg(feature = "dim2")] +pub type PaddedMatrix = Mat2; +/// The square matrix type for the current dimension (Mat2 in 2D, Mat3 in 3D), with explicit padding. +#[cfg(feature = "dim3")] +pub type PaddedMatrix = Mat4; + +/// The dimension constant as usize (for array indexing). +#[cfg(feature = "dim2")] +pub const DIM_USIZE: usize = 2; +/// The dimension constant as usize (for array indexing). +#[cfg(feature = "dim3")] +pub const DIM_USIZE: usize = 3; + +// +// Helper function: construct a diagonal matrix. +// +#[cfg(feature = "dim2")] +#[inline] +pub fn diag(v: Vector) -> Matrix { + Mat2::from_cols(Vec2::new(v.x, 0.0), Vec2::new(0.0, v.y)) +} + +#[cfg(feature = "dim3")] +#[inline] +pub fn diag(v: Vector) -> Matrix { + Mat3::from_cols( + Vec3::new(v.x, 0.0, 0.0), + Vec3::new(0.0, v.y, 0.0), + Vec3::new(0.0, 0.0, v.z), + ) +} + +/// Helper to compute the trace of a matrix. +#[cfg(feature = "dim2")] +#[inline] +pub fn trace(m: Matrix) -> f32 { + m.x_axis.x + m.y_axis.y +} + +#[cfg(feature = "dim3")] +#[inline] +pub fn trace(m: Matrix) -> f32 { + m.x_axis.x + m.y_axis.y + m.z_axis.z +} + +/// The length of a vector as AngVector. In 2D this is abs, in 3D it's the standard length. +#[cfg(feature = "dim2")] +#[inline] +pub fn ang_length(v: AngVector) -> f32 { + abs(v) +} + +#[cfg(feature = "dim3")] +#[inline] +pub fn ang_length(v: AngVector) -> f32 { + v.length() +} + +pub trait PaddingExt { + type WithoutPadding; + fn remove_padding(self) -> Self::WithoutPadding; + fn add_padding(without_padding: Self::WithoutPadding) -> Self; +} + +impl PaddingExt for Mat2 { + type WithoutPadding = Mat2; + #[inline] + fn remove_padding(self) -> Mat2 { + self + } + #[inline] + fn add_padding(without_padding: Mat2) -> Mat2 { + without_padding + } +} + +impl PaddingExt for Mat4 { + type WithoutPadding = Mat3; + #[inline] + fn remove_padding(self) -> Mat3 { + Mat3::from_cols(self.x_axis.xyz(), self.y_axis.xyz(), self.z_axis.xyz()) + } + #[inline] + fn add_padding(without_padding: Mat3) -> Mat4 { + Mat4::from_mat3(without_padding) + } +} + +// +// Modules +// +pub mod collision; +pub mod grid; +pub mod models; +pub mod solver; diff --git a/src_mpm_shaders/models/default.rs b/src_mpm_shaders/models/default.rs new file mode 100644 index 00000000..520eda31 --- /dev/null +++ b/src_mpm_shaders/models/default.rs @@ -0,0 +1,343 @@ +//! Default particle model: a tagged union dispatching to different constitutive models. + +use super::drucker_prager::*; +use super::fluid::FluidModel; +use super::interfaces::*; +use super::linear_elasticity::LinearElasticModel; +use super::neo_hookean_elasticity::NeoHookeanModel; +use super::snow::{SnowPlasticState, SnowPlasticity}; +use crate::{Matrix, PaddedMatrix, PaddingExt, Vector}; + +/// Number of `u32` words of per-particle storage available to a constitutive +/// model. Holds both the model parameters and any state the model mutates +/// (plastic state, damage, viscous strain, temperature, ...). +pub const MODEL_DATA_WORDS: usize = 24; + +/// The byte stride of `GpuParticleModel` in GPU buffers. Has to match the size +/// of `GpuParticleModel` on the host side. +pub const DEFAULT_MODEL_BYTES_STRIDE: usize = 100; + +/// Model type tag: corotated linear elasticity. +pub const MODEL_ELASTIC_LINEAR: u32 = 0; +/// Model type tag: neo-Hookean elasticity. +pub const MODEL_ELASTIC_NEO_HOOKEAN: u32 = 1; +/// Model type tag: Drucker-Prager sand with linear elastic backbone. +pub const MODEL_SAND_LINEAR: u32 = 2; +/// Model type tag: Drucker-Prager sand with neo-Hookean elastic backbone. +pub const MODEL_SAND_NEO_HOOKEAN: u32 = 3; +/// Model type tag: weakly-compressible fluid. +pub const MODEL_FLUID: u32 = 4; +/// Model type tag: snow (singular-value clamping with compaction hardening). +pub const MODEL_SNOW: u32 = 5; + +/// GPU particle model stored as a tagged union in a fixed-size buffer. +/// +/// The `tag` field selects the constitutive model variant, and `data` holds the +/// model parameters in a flat `[u32; 24]` (96 bytes) that is reinterpreted +/// as the appropriate model struct. +#[derive(Clone, Copy)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct GpuParticleModel { + pub tag: u32, + pub data: [u32; MODEL_DATA_WORDS], +} + +/// Sand model combining Drucker-Prager plasticity with an elastic backbone. +/// +/// The `plastic_state` field must be first because the state offset +/// assumes offset 0 within `SandModel`. +#[derive(Clone, Copy)] +#[repr(C)] +struct SandModel { + plastic_state: DruckerPragerPlasticState, + plastic: DruckerPragerPlasticity, + elastic: E, +} + +/// Loads a `LinearElasticModel` from the raw data array (3 words at offset 0). +#[inline] +fn load_elastic(data: &[u32; MODEL_DATA_WORDS], offset: usize) -> LinearElasticModel { + LinearElasticModel { + lambda: f32::from_bits(data[offset]), + mu: f32::from_bits(data[offset + 1]), + cfl_coeff: f32::from_bits(data[offset + 2]), + } +} + +/// Loads a `NeoHookeanModel` from the raw data array (3 words at offset 0). +#[inline] +fn load_neo_hookean(data: &[u32; MODEL_DATA_WORDS], offset: usize) -> NeoHookeanModel { + NeoHookeanModel { + lambda: f32::from_bits(data[offset]), + mu: f32::from_bits(data[offset + 1]), + cfl_coeff: f32::from_bits(data[offset + 2]), + } +} + +/// Snow model: plastic state, plasticity parameters, and elastic backbone. +/// +/// As for `SandModel`, `plastic_state` must come first: the store helper assumes +/// it sits at offset 0. +#[derive(Clone, Copy)] +#[repr(C)] +struct SnowModel { + plastic_state: SnowPlasticState, + plastic: SnowPlasticity, + elastic: LinearElasticModel, +} + +/// Loads a `SnowModel` from the raw data array (7 words). +#[inline] +fn load_snow(data: &[u32; MODEL_DATA_WORDS]) -> SnowModel { + SnowModel { + plastic_state: SnowPlasticState { + plastic_det: f32::from_bits(data[0]), + }, + plastic: SnowPlasticity { + critical_compression: f32::from_bits(data[1]), + critical_stretch: f32::from_bits(data[2]), + hardening: f32::from_bits(data[3]), + }, + elastic: load_elastic(data, 4), + } +} + +/// Writes a `SnowPlasticState` back into the raw data array at offset 0. +#[inline] +fn store_snow_state(data: &mut [u32; MODEL_DATA_WORDS], state: SnowPlasticState) { + data[0] = state.plastic_det.to_bits(); +} + +/// Scales the Lamé parameters of an elastic model by a hardening factor. +#[inline] +fn harden(elastic: LinearElasticModel, hardening: f32) -> LinearElasticModel { + LinearElasticModel { + lambda: elastic.lambda * hardening, + mu: elastic.mu * hardening, + cfl_coeff: elastic.cfl_coeff, + } +} + +/// Loads a `FluidModel` from the raw data array (5 words). +#[inline] +fn load_fluid(data: &[u32; MODEL_DATA_WORDS], offset: usize) -> FluidModel { + FluidModel { + bulk_modulus: f32::from_bits(data[offset]), + gamma: f32::from_bits(data[offset + 1]), + viscosity: f32::from_bits(data[offset + 2]), + cfl_coeff: f32::from_bits(data[offset + 3]), + tensile_stiffness: f32::from_bits(data[offset + 4]), + } +} + +/// Loads a `DruckerPragerPlasticState` from the raw data array (3 words). +#[inline] +fn load_plastic_state(data: &[u32; MODEL_DATA_WORDS], offset: usize) -> DruckerPragerPlasticState { + DruckerPragerPlasticState { + plastic_deformation_gradient_det: f32::from_bits(data[offset]), + plastic_hardening: f32::from_bits(data[offset + 1]), + log_vol_gain: f32::from_bits(data[offset + 2]), + } +} + +/// Loads a `DruckerPragerPlasticity` from the raw data array (7 words). +#[inline] +fn load_plasticity(data: &[u32; MODEL_DATA_WORDS], offset: usize) -> DruckerPragerPlasticity { + DruckerPragerPlasticity { + ha: f32::from_bits(data[offset]), + hb: f32::from_bits(data[offset + 1]), + hc: f32::from_bits(data[offset + 2]), + hd: f32::from_bits(data[offset + 3]), + lambda: f32::from_bits(data[offset + 4]), + mu: f32::from_bits(data[offset + 5]), + cohesion: f32::from_bits(data[offset + 6]), + } +} + +/// Loads a `SandModel` from the raw data array (13 words). +#[inline] +fn load_sand_linear(data: &[u32; MODEL_DATA_WORDS]) -> SandModel { + SandModel { + plastic_state: load_plastic_state(data, 0), + plastic: load_plasticity(data, 3), + elastic: load_elastic(data, 10), + } +} + +/// Loads a `SandModel` from the raw data array (13 words). +#[inline] +fn load_sand_neo_hookean(data: &[u32; MODEL_DATA_WORDS]) -> SandModel { + SandModel { + plastic_state: load_plastic_state(data, 0), + plastic: load_plasticity(data, 3), + elastic: load_neo_hookean(data, 10), + } +} + +/// Writes a `DruckerPragerPlasticState` back into the raw data array at offset 0. +#[inline] +fn store_plastic_state(data: &mut [u32; MODEL_DATA_WORDS], state: DruckerPragerPlasticState) { + data[0] = state.plastic_deformation_gradient_det.to_bits(); + data[1] = state.plastic_hardening.to_bits(); + data[2] = state.log_vol_gain.to_bits(); +} + +/// Default particle model dispatcher. +/// +/// Reads the model tag from the `GpuParticleModel` array and dispatches +/// to the appropriate constitutive model implementation. +pub struct DefaultParticleModel; + +impl DefaultParticleModel { + /// Returns the model flags for a given particle. + #[inline] + pub fn model_flags(models: &[GpuParticleModel], particle_id: u32) -> u32 { + if models[particle_id as usize].tag == MODEL_FLUID { + MODEL_FLAGS_FLUID + } else { + MODEL_FLAGS_NONE + } + } + + /// Runs the constitutive model update for a particle. + /// + /// Reads the model data, computes the Kirchoff stress, and for plastic models + /// also updates the plastic state and deformation gradient in place. + #[inline] + pub fn update( + models: &mut [GpuParticleModel], + data: &ParticleUpdateData, + def_grad_padded: &mut PaddedMatrix, + ) -> ModelUpdateResult { + let model = &mut models[data.particle_id as usize]; + let tag = model.tag; + let def_grad = def_grad_padded.remove_padding(); + + match tag { + MODEL_ELASTIC_LINEAR => { + let elastic = load_elastic(&model.data, 0); + let stress = elastic.kirchoff_stress(def_grad); + ModelUpdateResult::new(stress) + } + MODEL_ELASTIC_NEO_HOOKEAN => { + let elastic = load_neo_hookean(&model.data, 0); + let stress = elastic.kirchoff_stress(def_grad); + ModelUpdateResult::new(stress) + } + MODEL_SAND_LINEAR => { + let sand = load_sand_linear(&model.data); + let projection = sand.plastic.project(sand.plastic_state, def_grad); + store_plastic_state(&mut model.data, projection.state); + *def_grad_padded = PaddedMatrix::add_padding(projection.deformation_gradient); + let stress = sand + .elastic + .kirchoff_stress(projection.deformation_gradient); + ModelUpdateResult::new(stress) + } + MODEL_SAND_NEO_HOOKEAN => { + let sand = load_sand_neo_hookean(&model.data); + let projection = sand.plastic.project(sand.plastic_state, def_grad); + store_plastic_state(&mut model.data, projection.state); + *def_grad_padded = PaddedMatrix::add_padding(projection.deformation_gradient); + let stress = sand + .elastic + .kirchoff_stress(projection.deformation_gradient); + ModelUpdateResult::new(stress) + } + MODEL_FLUID => { + let fluid = load_fluid(&model.data, 0); + let stress = fluid.kirchoff_stress(def_grad, data.strain_rate()); + ModelUpdateResult::new(stress) + } + MODEL_SNOW => { + let snow = load_snow(&model.data); + let projection = snow.plastic.project(snow.plastic_state, def_grad); + store_snow_state(&mut model.data, projection.state); + *def_grad_padded = projection.deformation_gradient; + let stress = harden(snow.elastic, projection.hardening) + .kirchoff_stress(projection.deformation_gradient.remove_padding()); + ModelUpdateResult::new(stress) + } + _ => ModelUpdateResult::new(Matrix::ZERO), + } + } + + /// Computes the CFL-based timestep bound for a given particle's model. + #[inline] + pub fn timestep_bound( + models: &[GpuParticleModel], + particle_id: u32, + particle_density0: f32, + def_grad: Matrix, + particle_velocity: Vector, + cell_width: f32, + ) -> f32 { + let model = &models[particle_id as usize]; + let tag = model.tag; + let def_grad_det = def_grad.determinant(); + + match tag { + MODEL_ELASTIC_LINEAR => { + let elastic = load_elastic(&model.data, 0); + elastic.timestep_bound( + particle_density0, + particle_velocity, + def_grad_det, + 1.0, + cell_width, + ) + } + MODEL_ELASTIC_NEO_HOOKEAN => { + let elastic = load_neo_hookean(&model.data, 0); + elastic.timestep_bound( + particle_density0, + particle_velocity, + def_grad_det, + 1.0, + cell_width, + ) + } + MODEL_SAND_LINEAR => { + let sand = load_sand_linear(&model.data); + sand.elastic.timestep_bound( + particle_density0, + particle_velocity, + def_grad_det, + 1.0, + cell_width, + ) + } + MODEL_SAND_NEO_HOOKEAN => { + let sand = load_sand_neo_hookean(&model.data); + sand.elastic.timestep_bound( + particle_density0, + particle_velocity, + def_grad_det, + 1.0, + cell_width, + ) + } + MODEL_FLUID => { + let fluid = load_fluid(&model.data, 0); + fluid.timestep_bound( + particle_density0, + particle_velocity, + def_grad_det, + cell_width, + ) + } + MODEL_SNOW => { + let snow = load_snow(&model.data); + snow.elastic.timestep_bound( + particle_density0, + particle_velocity, + def_grad_det, + snow.plastic.hardening_factor(snow.plastic_state), + cell_width, + ) + } + _ => 0.0, + } + } +} diff --git a/src_mpm_shaders/models/drucker_prager.rs b/src_mpm_shaders/models/drucker_prager.rs new file mode 100644 index 00000000..101bfab1 --- /dev/null +++ b/src_mpm_shaders/models/drucker_prager.rs @@ -0,0 +1,268 @@ +//! Drucker-Prager plasticity model. + +use crate::glamx::MatExt; +use crate::{Matrix, Vector, diag, sin, sqrt}; +use khal_std::num_traits::Float; + +/// Persistent plastic state for a Drucker-Prager particle. +#[derive(Clone, Copy)] +#[cfg_attr( + not(target_arch_is_gpu), + derive(Debug, PartialEq, bytemuck::Pod, bytemuck::Zeroable) +)] +#[repr(C)] +pub struct DruckerPragerPlasticState { + pub plastic_deformation_gradient_det: f32, + pub plastic_hardening: f32, + pub log_vol_gain: f32, +} + +/// Result of a Drucker-Prager plasticity projection. +/// +/// Contains the updated plastic state and the projected deformation gradient. +#[derive(Clone, Copy)] +pub struct DruckerPragerResult { + pub state: DruckerPragerPlasticState, + pub deformation_gradient: Matrix, +} + +/// Intermediate result of the return mapping on singular values. +#[derive(Clone, Copy)] +pub struct DruckerPragerProjectionResult { + pub singular_values: Vector, + pub plastic_hardening: f32, + pub valid: bool, +} + +/// Drucker-Prager plasticity model with hardening. +/// +/// The hardening law is parameterized by (ha, hb, hc, hd) which control +/// how the friction angle evolves with accumulated plastic strain. +#[derive(Clone, Copy)] +#[cfg_attr( + not(target_arch_is_gpu), + derive(Debug, PartialEq, bytemuck::Pod, bytemuck::Zeroable) +)] +#[repr(C)] +pub struct DruckerPragerPlasticity { + pub ha: f32, + pub hb: f32, + pub hc: f32, + pub hd: f32, + pub lambda: f32, + pub mu: f32, + /// Cohesion, expressed as the volumetric log-strain the material can sustain + /// in tension before it separates. Zero recovers cohesionless sand; positive + /// values let the material hold a shape at zero confining pressure (wet sand, + /// mud, snow that packs). + pub cohesion: f32, +} + +impl DruckerPragerPlasticity { + /// Computes the friction-angle-based alpha parameter given accumulated plastic strain `q`. + #[inline] + pub fn alpha(&self, q: f32) -> f32 { + let angle = self.ha + (self.hb * q - self.hd) * (-self.hc * q).exp(); + let s_angle = sin(angle); + sqrt(2.0 / 3.0) * (2.0 * s_angle) / (3.0 - s_angle) + } + + /// Projects the singular values of the deformation gradient onto the yield surface (2D). + #[cfg(feature = "dim2")] + #[inline] + fn project_deformation_gradient( + &self, + singular_values: Vector, + log_vol_gain: f32, + alpha: f32, + ) -> DruckerPragerProjectionResult { + let d = 2.0; + let strain = glamx::Vec2::new(singular_values.x.ln(), singular_values.y.ln()) + + Vector::splat(log_vol_gain / d); + let strain_trace = strain.x + strain.y; + let deviatoric_strain = strain - Vector::splat(strain_trace / d); + // Cohesion shifts the yield surface along the pressure axis: the material + // behaves as if it were under `cohesion` more compression than it is. + let shifted_trace = strain_trace - self.cohesion; + + if shifted_trace > 0.0 || deviatoric_strain == Vector::ZERO { + // Past the tensile capacity: project onto the tensile limit rather + // than all the way back to the undeformed state. + return DruckerPragerProjectionResult { + singular_values: Vector::splat((self.cohesion / d).exp()), + plastic_hardening: strain.length(), + valid: true, + }; + } + + let deviatoric_strain_norm = deviatoric_strain.length(); + let gamma = deviatoric_strain_norm + + (d * self.lambda + 2.0 * self.mu) / (2.0 * self.mu) * shifted_trace * alpha; + + if gamma <= 0.0 { + return DruckerPragerProjectionResult { + singular_values: Vector::ZERO, + plastic_hardening: 0.0, + valid: false, + }; + } + + let h = strain - deviatoric_strain * (gamma / deviatoric_strain_norm); + DruckerPragerProjectionResult { + singular_values: glamx::Vec2::new(h.x.exp(), h.y.exp()), + plastic_hardening: gamma, + valid: true, + } + } + + /// Projects the singular values of the deformation gradient onto the yield surface (3D). + #[cfg(feature = "dim3")] + #[inline] + fn project_deformation_gradient( + &self, + singular_values: Vector, + log_vol_gain: f32, + alpha: f32, + ) -> DruckerPragerProjectionResult { + let d = 3.0; + let strain = glamx::Vec3::new( + singular_values.x.ln(), + singular_values.y.ln(), + singular_values.z.ln(), + ) + Vector::splat(log_vol_gain / d); + let strain_trace = strain.x + strain.y + strain.z; + let deviatoric_strain = strain - Vector::splat(strain_trace / d); + // Cohesion shifts the yield surface along the pressure axis: the material + // behaves as if it were under `cohesion` more compression than it is. + let shifted_trace = strain_trace - self.cohesion; + + if shifted_trace > 0.0 || deviatoric_strain == Vector::ZERO { + // Past the tensile capacity: project onto the tensile limit rather + // than all the way back to the undeformed state. + return DruckerPragerProjectionResult { + singular_values: Vector::splat((self.cohesion / d).exp()), + plastic_hardening: strain.length(), + valid: true, + }; + } + + let deviatoric_strain_norm = deviatoric_strain.length(); + let gamma = deviatoric_strain_norm + + (d * self.lambda + 2.0 * self.mu) / (2.0 * self.mu) * shifted_trace * alpha; + + if gamma <= 0.0 { + return DruckerPragerProjectionResult { + singular_values: Vector::ZERO, + plastic_hardening: 0.0, + valid: false, + }; + } + + let h = strain - deviatoric_strain * (gamma / deviatoric_strain_norm); + DruckerPragerProjectionResult { + singular_values: glamx::Vec3::new(h.x.exp(), h.y.exp(), h.z.exp()), + plastic_hardening: gamma, + valid: true, + } + } + + /// Projects the deformation gradient through the Drucker-Prager yield surface (2D). + /// + /// If plasticity is disabled (lambda == 0), returns the input unchanged. + /// Otherwise, performs SVD, projects the singular values, and recomposes. + #[cfg(feature = "dim2")] + #[inline] + pub fn project( + &self, + state: DruckerPragerPlasticState, + deformation_gradient: Matrix, + ) -> DruckerPragerResult { + if self.lambda == 0.0 { + // Plasticity is disabled on this particle. + return DruckerPragerResult { + state, + deformation_gradient, + }; + } + + let svd = deformation_gradient.svd(); + let alpha = self.alpha(state.plastic_hardening); + let projection = self.project_deformation_gradient(svd.s, state.log_vol_gain, alpha); + + if projection.valid { + let prev_det = svd.s.x * svd.s.y; + let new_det = projection.singular_values.x * projection.singular_values.y; + + let new_plastic_deformation_gradient_det = + state.plastic_deformation_gradient_det * prev_det / new_det; + let new_log_vol_gain = state.log_vol_gain + prev_det.ln() - new_det.ln(); + let new_plastic_hardening = state.plastic_hardening + projection.plastic_hardening; + let new_deformation_gradient = svd.u * diag(projection.singular_values) * svd.vt; + + DruckerPragerResult { + state: DruckerPragerPlasticState { + plastic_deformation_gradient_det: new_plastic_deformation_gradient_det, + plastic_hardening: new_plastic_hardening, + log_vol_gain: new_log_vol_gain, + }, + deformation_gradient: new_deformation_gradient, + } + } else { + DruckerPragerResult { + state, + deformation_gradient, + } + } + } + + /// Projects the deformation gradient through the Drucker-Prager yield surface (3D). + /// + /// If plasticity is disabled (lambda == 0), returns the input unchanged. + /// Otherwise, performs SVD, projects the singular values, and recomposes. + #[cfg(feature = "dim3")] + #[inline] + pub fn project( + &self, + state: DruckerPragerPlasticState, + deformation_gradient: Matrix, + ) -> DruckerPragerResult { + if self.lambda == 0.0 { + // Plasticity is disabled on this particle. + return DruckerPragerResult { + state, + deformation_gradient, + }; + } + + let svd = deformation_gradient.svd(); + let alpha = self.alpha(state.plastic_hardening); + let projection = self.project_deformation_gradient(svd.s, state.log_vol_gain, alpha); + + if projection.valid { + let prev_det = svd.s.x * svd.s.y * svd.s.z; + let new_det = projection.singular_values.x + * projection.singular_values.y + * projection.singular_values.z; + + let new_plastic_deformation_gradient_det = + state.plastic_deformation_gradient_det * prev_det / new_det; + let new_log_vol_gain = state.log_vol_gain + prev_det.ln() - new_det.ln(); + let new_plastic_hardening = state.plastic_hardening + projection.plastic_hardening; + let new_deformation_gradient = svd.u * diag(projection.singular_values) * svd.vt; + + DruckerPragerResult { + state: DruckerPragerPlasticState { + plastic_deformation_gradient_det: new_plastic_deformation_gradient_det, + plastic_hardening: new_plastic_hardening, + log_vol_gain: new_log_vol_gain, + }, + deformation_gradient: new_deformation_gradient, + } + } else { + DruckerPragerResult { + state, + deformation_gradient, + } + } + } +} diff --git a/src_mpm_shaders/models/fluid.rs b/src_mpm_shaders/models/fluid.rs new file mode 100644 index 00000000..b02344dd --- /dev/null +++ b/src_mpm_shaders/models/fluid.rs @@ -0,0 +1,107 @@ +//! Weakly-compressible fluid model: Tait equation of state plus viscosity. + +use super::utils::{ElasticitySoundSpeedTimestepBound, deviatoric_part}; +use crate::glamx::MatExt; +use crate::{Matrix, Vector}; +use khal_std::num_traits::Float; + +/// Weakly-compressible Newtonian fluid. +/// +/// The pressure comes from the Tait equation of state, so the fluid resists +/// compression stiffly without requiring a pressure solve; the deviatoric part +/// of the stress is a plain Newtonian viscous term. +/// +/// Particles using this model only track the volumetric part of the deformation +/// gradient (see [`MODEL_FLAGS_FLUID`](super::interfaces::MODEL_FLAGS_FLUID)), +/// which avoids the drift a full tensor would accumulate under the large shear +/// a fluid undergoes. +#[derive(Clone, Copy)] +#[cfg_attr( + not(target_arch_is_gpu), + derive(Debug, PartialEq, bytemuck::Pod, bytemuck::Zeroable) +)] +#[repr(C)] +pub struct FluidModel { + /// Bulk modulus `k` of the equation of state (Pa). Higher values make the + /// fluid less compressible but shorten the stable timestep. + pub bulk_modulus: f32, + /// Stiffness exponent `gamma` of the equation of state (7 for water). + pub gamma: f32, + /// Dynamic viscosity (Pa.s). + pub viscosity: f32, + /// CFL coefficient scaling the stable timestep. + pub cfl_coeff: f32, + /// Stiffness of the tensile branch, as a fraction of `bulk_modulus`. + /// + /// A free surface biases the divergence the grid reports (the stencil + /// reaches into empty cells), so an explicitly integrated volume ratio + /// ratchets upward every step and the fluid slowly inflates. A soft + /// pull-back at `J > 1` holds the volume without making the surface clump + /// the way a full tensile equation of state would. + pub tensile_stiffness: f32, +} + +impl FluidModel { + /// Pressure as a function of the volume ratio `J`. + /// + /// Compression follows the Tait equation of state, + /// `k ((rho / rho0)^gamma - 1)` with `rho / rho0 = 1 / J`, which stiffens + /// sharply and keeps the fluid nearly incompressible. Expansion uses a much + /// softer linear branch instead: the Tait curve is far too strong in tension + /// and would pull the fluid into blobs, but some restoring force is still + /// needed (see [`Self::tensile_stiffness`]). + #[inline] + pub fn pressure(&self, j: f32) -> f32 { + let j = f32::max(j, 1.0e-6); + if j <= 1.0 { + // `exp(-gamma * ln(j))` rather than `powf`, matching the + // transcendentals the other models already rely on. + let ratio = (-self.gamma * j.ln()).exp(); + self.bulk_modulus * (ratio - 1.0) + } else { + -self.bulk_modulus * self.tensile_stiffness * (j - 1.0) + } + } + + /// Computes the Kirchoff stress `J * (-p I + 2 mu dev(strain_rate))`. + #[inline] + pub fn kirchoff_stress(&self, deformation_gradient: Matrix, strain_rate: Matrix) -> Matrix { + let j = f32::max(deformation_gradient.determinant(), 1.0e-6); + // Only the deviatoric part of the strain rate contributes: the + // volumetric response is entirely governed by the equation of state. + let mut stress = deviatoric_part(strain_rate) * (2.0 * self.viscosity * j); + let diag_val = -self.pressure(j) * j; + + stress.x_axis.x += diag_val; + stress.y_axis.y += diag_val; + #[cfg(feature = "dim3")] + { + stress.z_axis.z += diag_val; + } + + stress + } + + /// Computes the CFL-based timestep bound from the speed of sound of the + /// equation of state, `sqrt(gamma * k / rho)`. + #[inline] + pub fn timestep_bound( + &self, + particle_density0: f32, + particle_velocity: Vector, + particle_def_grad_det: f32, + cell_width: f32, + ) -> f32 { + let bound = ElasticitySoundSpeedTimestepBound::new( + self.cfl_coeff, + self.bulk_modulus * self.gamma, + 0.0, + ); + bound.timestep_bound( + particle_density0, + particle_def_grad_det, + particle_velocity, + cell_width, + ) + } +} diff --git a/src_mpm_shaders/models/interfaces.rs b/src_mpm_shaders/models/interfaces.rs new file mode 100644 index 00000000..ce97fa32 --- /dev/null +++ b/src_mpm_shaders/models/interfaces.rs @@ -0,0 +1,48 @@ +use crate::Matrix; + +/// Result of a constitutive model update, containing the Kirchoff stress tensor. +#[derive(Clone, Copy, Default)] +pub struct ModelUpdateResult { + pub kirchoff_stress: Matrix, +} + +impl ModelUpdateResult { + #[inline] + pub fn new(kirchoff_stress: Matrix) -> Self { + Self { kirchoff_stress } + } +} + +/// Data passed to the particle model update function. +#[derive(Clone, Copy)] +pub struct ParticleUpdateData { + pub dt: f32, + pub cell_width: f32, + pub particle_id: u32, + /// Velocity gradient at the particle, as gathered by G2P. Rate-dependent + /// models (viscosity, viscoplasticity, viscoelasticity) need it; purely + /// deformation-driven models can ignore it. + pub velocity_gradient: Matrix, +} + +impl ParticleUpdateData { + #[inline] + pub fn new(dt: f32, cell_width: f32, particle_id: u32, velocity_gradient: Matrix) -> Self { + Self { + dt, + cell_width, + particle_id, + velocity_gradient, + } + } + + /// Symmetric part of the velocity gradient (the strain rate). + #[inline] + pub fn strain_rate(&self) -> Matrix { + (self.velocity_gradient + self.velocity_gradient.transpose()) * 0.5 + } +} + +/// Model behavior flags (bitflags stored as u32). +pub const MODEL_FLAGS_NONE: u32 = 0; +pub const MODEL_FLAGS_FLUID: u32 = 1; diff --git a/src_mpm_shaders/models/linear_elasticity.rs b/src_mpm_shaders/models/linear_elasticity.rs new file mode 100644 index 00000000..4832652f --- /dev/null +++ b/src_mpm_shaders/models/linear_elasticity.rs @@ -0,0 +1,99 @@ +//! Linear (corotated) elasticity model. + +use super::utils::{ + ElasticitySoundSpeedTimestepBound, bulk_modulus_from_lame, shear_modulus_from_lame, +}; +use crate::glamx::MatExt; +use crate::{Matrix, Vector, diag}; + +/// Corotated linear elastic constitutive model. +/// +/// Uses SVD-based corotated formulation for computing the Kirchoff stress. +#[derive(Clone, Copy)] +#[cfg_attr( + not(target_arch_is_gpu), + derive(Debug, PartialEq, bytemuck::Pod, bytemuck::Zeroable) +)] +#[repr(C)] +pub struct LinearElasticModel { + pub lambda: f32, + pub mu: f32, + pub cfl_coeff: f32, +} + +impl LinearElasticModel { + /// Computes the Kirchoff stress tensor using corotated linear elasticity. + /// + /// Performs SVD on the deformation gradient, computes the strain from the + /// singular values, and assembles the stress tensor. + #[inline] + pub fn kirchoff_stress(&self, deformation_gradient: Matrix) -> Matrix { + let svd = deformation_gradient.svd(); + + #[cfg(feature = "dim2")] + let j = svd.s.x * svd.s.y; + #[cfg(feature = "dim3")] + let j = svd.s.x * svd.s.y * svd.s.z; + + let modified_s = svd.s - Vector::ONE; + + // Recompose with modified singular values: U * diag(modified_s) * Vt + let recomposed = svd.u * diag(modified_s) * svd.vt; + + let diag_val = self.lambda * (j - 1.0) * j; + let mut result = recomposed * deformation_gradient.transpose() * (2.0 * self.mu); + + result.x_axis.x += diag_val; + result.y_axis.y += diag_val; + #[cfg(feature = "dim3")] + { + result.z_axis.z += diag_val; + } + + result + } + + /// Computes the CFL-based timestep bound for this elastic model. + #[inline] + pub fn timestep_bound( + &self, + particle_density0: f32, + particle_velocity: Vector, + particle_def_grad_det: f32, + elastic_hardening: f32, + cell_width: f32, + ) -> f32 { + let bulk_modulus = bulk_modulus_from_lame(self.lambda, self.mu); + let shear_modulus = shear_modulus_from_lame(self.lambda, self.mu); + + let bound = ElasticitySoundSpeedTimestepBound::new( + self.cfl_coeff, + bulk_modulus * elastic_hardening, + shear_modulus * elastic_hardening, + ); + bound.timestep_bound( + particle_density0, + particle_def_grad_det, + particle_velocity, + cell_width, + ) + } + + /// Computes the positive part of the elastic energy density, keeping only + /// the tensile contributions. + #[inline] + pub fn pos_energy_density(&self, def_grad: Matrix, elastic_hardening: f32) -> f32 { + let j = def_grad.determinant(); + let svd = def_grad.svd(); + + let sig = (svd.s - Vector::ONE).max(Vector::ZERO); + let pos_dev_part = self.mu * elastic_hardening * sig.dot(sig); + let spherical_part = self.lambda * elastic_hardening * 0.5 * (j - 1.0) * (j - 1.0); + + if j < 1.0 { + pos_dev_part + } else { + pos_dev_part + spherical_part + } + } +} diff --git a/src_mpm_shaders/models/mod.rs b/src_mpm_shaders/models/mod.rs new file mode 100644 index 00000000..68fb0a60 --- /dev/null +++ b/src_mpm_shaders/models/mod.rs @@ -0,0 +1,9 @@ +pub mod default; +pub mod drucker_prager; +pub mod fluid; +pub mod interfaces; +pub mod linear_elasticity; +pub mod neo_hookean_elasticity; +pub mod snow; +pub mod specializations; +pub mod utils; diff --git a/src_mpm_shaders/models/neo_hookean_elasticity.rs b/src_mpm_shaders/models/neo_hookean_elasticity.rs new file mode 100644 index 00000000..64b2d96b --- /dev/null +++ b/src_mpm_shaders/models/neo_hookean_elasticity.rs @@ -0,0 +1,84 @@ +//! Neo-Hookean elasticity model. + +use super::utils::{ + ElasticitySoundSpeedTimestepBound, bulk_modulus_from_lame, shear_modulus_from_lame, +}; +use crate::glamx::MatExt; +use crate::{Matrix, Vector}; +use khal_std::num_traits::Float; + +/// Neo-Hookean hyperelastic constitutive model. +/// +/// Computes stress based on the deformation gradient using the +/// neo-Hookean strain energy density function. +#[derive(Clone, Copy)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct NeoHookeanModel { + pub lambda: f32, + pub mu: f32, + pub cfl_coeff: f32, +} + +impl NeoHookeanModel { + /// Computes the Kirchoff stress tensor for the neo-Hookean model. + #[inline] + pub fn kirchoff_stress(&self, deformation_gradient: Matrix) -> Matrix { + let j = f32::max(deformation_gradient.determinant(), 1.0e-10); + let diag_val = self.lambda * j.ln() - self.mu; + let mut stress = deformation_gradient * deformation_gradient.transpose() * self.mu; + + stress.x_axis.x += diag_val; + stress.y_axis.y += diag_val; + #[cfg(feature = "dim3")] + { + stress.z_axis.z += diag_val; + } + + stress + } + + /// Computes the CFL-based timestep bound for this elastic model. + #[inline] + pub fn timestep_bound( + &self, + particle_density0: f32, + particle_velocity: Vector, + particle_def_grad_det: f32, + elastic_hardening: f32, + cell_width: f32, + ) -> f32 { + let bulk_modulus = bulk_modulus_from_lame(self.lambda, self.mu); + let shear_modulus = shear_modulus_from_lame(self.lambda, self.mu); + + let bound = ElasticitySoundSpeedTimestepBound::new( + self.cfl_coeff, + bulk_modulus * elastic_hardening, + shear_modulus * elastic_hardening, + ); + bound.timestep_bound( + particle_density0, + particle_def_grad_det, + particle_velocity, + cell_width, + ) + } + + /// Computes the positive part of the elastic energy density, keeping only + /// the tensile contributions. + #[inline] + pub fn pos_energy_density(&self, def_grad: Matrix, elastic_hardening: f32) -> f32 { + let j = def_grad.determinant(); + let svd = def_grad.svd(); + + let sig = (svd.s - Vector::ONE).max(Vector::ZERO); + let pos_dev_part = self.mu * elastic_hardening * sig.dot(sig); + let spherical_part = self.lambda * elastic_hardening * 0.5 * (j - 1.0) * (j - 1.0); + + if j < 1.0 { + pos_dev_part + } else { + pos_dev_part + spherical_part + } + } +} diff --git a/src_mpm_shaders/models/snow.rs b/src_mpm_shaders/models/snow.rs new file mode 100644 index 00000000..be013567 --- /dev/null +++ b/src_mpm_shaders/models/snow.rs @@ -0,0 +1,114 @@ +//! Snow elastoplasticity (Stomakhin et al. 2013). + +use crate::glamx::MatExt; +use crate::{Matrix, PaddedMatrix, PaddingExt, Vector, diag}; +use khal_std::num_traits::Float; + +/// Persistent plastic state for a snow particle. +#[derive(Clone, Copy)] +#[cfg_attr( + not(target_arch_is_gpu), + derive(Debug, PartialEq, bytemuck::Pod, bytemuck::Zeroable) +)] +#[repr(C)] +pub struct SnowPlasticState { + /// Determinant of the plastic part of the deformation gradient. Below 1 the + /// particle has been compacted, which hardens it. + pub plastic_det: f32, +} + +/// Snow plasticity: singular-value clamping with compaction hardening. +/// +/// The elastic deformation is confined to a box in singular-value space and +/// whatever is clamped away becomes plastic. Compaction is tracked separately, +/// so a packed snowball becomes stiffer than the loose snow it was made from. +#[derive(Clone, Copy)] +#[cfg_attr( + not(target_arch_is_gpu), + derive(Debug, PartialEq, bytemuck::Pod, bytemuck::Zeroable) +)] +#[repr(C)] +pub struct SnowPlasticity { + /// Compression the elastic response sustains before yielding (`theta_c`). + pub critical_compression: f32, + /// Stretch the elastic response sustains before yielding (`theta_s`). + /// Small values make the snow break apart readily under tension. + pub critical_stretch: f32, + /// Hardening coefficient (`xi`): how sharply the moduli grow with compaction. + pub hardening: f32, +} + +/// Result of a snow return mapping. +/// +/// The deformation gradient is a [`PaddedMatrix`]: a bare `Matrix` member is laid +/// out with 36-byte offsets in 3D while SPIR-V pads each column of a 3x3 to 16 +/// bytes, so any member after it overlaps and the module fails validation. The +/// `#[repr(C)]` keeps Rust from reordering it back into that position. +#[derive(Clone, Copy)] +#[repr(C)] +pub struct SnowResult { + pub state: SnowPlasticState, + pub deformation_gradient: PaddedMatrix, + /// Multiplier to apply to the elastic moduli for this particle. + pub hardening: f32, +} + +/// Product of the components of a vector, i.e. the determinant of `diag(v)`. +#[inline] +fn component_product(v: Vector) -> f32 { + #[cfg(feature = "dim2")] + { + v.x * v.y + } + #[cfg(feature = "dim3")] + { + v.x * v.y * v.z + } +} + +impl SnowPlasticity { + /// Lowest and highest plastic compaction tracked. Left unbounded, a particle + /// squeezed in a corner can harden without limit and blow up the timestep. + const MIN_PLASTIC_DET: f32 = 0.1; + const MAX_PLASTIC_DET: f32 = 4.0; + + /// Multiplier applied to the Lamé parameters for the given compaction. + #[inline] + pub fn hardening_factor(&self, state: SnowPlasticState) -> f32 { + (self.hardening * (1.0 - state.plastic_det)).exp() + } + + /// Clamps the elastic singular values into the yield box, moving the excess + /// into the plastic determinant. + #[inline] + pub fn project(&self, state: SnowPlasticState, deformation_gradient: Matrix) -> SnowResult { + if self.critical_compression <= 0.0 && self.critical_stretch <= 0.0 { + // Plasticity disabled: purely elastic snow. + return SnowResult { + state, + deformation_gradient: PaddedMatrix::add_padding(deformation_gradient), + hardening: self.hardening_factor(state), + }; + } + + let svd = deformation_gradient.svd(); + let lo = 1.0 - self.critical_compression; + let hi = 1.0 + self.critical_stretch; + let clamped = svd.s.clamp(Vector::splat(lo), Vector::splat(hi)); + + let prev_det = component_product(svd.s); + let new_det = component_product(clamped); + + // Everything the clamp removed is absorbed into the plastic part, so the + // total deformation is preserved. + let plastic_det = (state.plastic_det * prev_det / new_det) + .clamp(Self::MIN_PLASTIC_DET, Self::MAX_PLASTIC_DET); + let state = SnowPlasticState { plastic_det }; + + SnowResult { + state, + deformation_gradient: PaddedMatrix::add_padding(svd.u * diag(clamped) * svd.vt), + hardening: self.hardening_factor(state), + } + } +} diff --git a/src_mpm_shaders/models/specializations.rs b/src_mpm_shaders/models/specializations.rs new file mode 100644 index 00000000..2527949c --- /dev/null +++ b/src_mpm_shaders/models/specializations.rs @@ -0,0 +1 @@ +pub use super::default::DefaultParticleModel as ParticleModel; diff --git a/src_mpm_shaders/models/utils.rs b/src_mpm_shaders/models/utils.rs new file mode 100644 index 00000000..d8055d70 --- /dev/null +++ b/src_mpm_shaders/models/utils.rs @@ -0,0 +1,176 @@ +use crate::trace; +use crate::{DIM, DIM_USIZE, Matrix, Vector, sqrt}; + +/// Computes the Lame parameters (lambda, mu) from the Young modulus and Poisson ratio. +/// Returns (lambda, mu). +#[inline] +pub fn lame_lambda_mu(young_modulus: f32, poisson_ratio: f32) -> (f32, f32) { + let lambda = + young_modulus * poisson_ratio / ((1.0 + poisson_ratio) * (1.0 - 2.0 * poisson_ratio)); + let mu = hook_to_shear_modulus(young_modulus, poisson_ratio); + (lambda, mu) +} + +/// Computes the shear modulus from the Young modulus and Poisson ratio. +#[inline] +pub fn hook_to_shear_modulus(young_modulus: f32, poisson_ratio: f32) -> f32 { + young_modulus / (2.0 * (1.0 + poisson_ratio)) +} + +/// Computes the bulk modulus from the Young modulus and Poisson ratio. +#[inline] +pub fn hook_to_bulk_modulus(young_modulus: f32, poisson_ratio: f32) -> f32 { + young_modulus / (3.0 * (1.0 - 2.0 * poisson_ratio)) +} + +/// Returns the shear modulus from the Lame parameters. +#[inline] +pub fn shear_modulus_from_lame(_lambda: f32, mu: f32) -> f32 { + mu +} + +/// Computes the bulk modulus from the Lame parameters. +#[inline] +pub fn bulk_modulus_from_lame(lambda: f32, mu: f32) -> f32 { + lambda + 2.0 * mu / 3.0 +} + +/// Solves the quadratic equation `a*x^2 + b*x + c = 0`. +/// Returns the two roots as (x1, x2). +#[inline] +pub fn solve_quadratic(a: f32, b: f32, c: f32) -> (f32, f32) { + let discr_sqr = sqrt(b * b - 4.0 * a * c); + ((-b + discr_sqr) / (2.0 * a), (-b - discr_sqr) / (2.0 * a)) +} + +/// Computes the spin tensor (antisymmetric part) of a velocity gradient. +#[inline] +pub fn spin_tensor(velocity_gradient: Matrix) -> Matrix { + (velocity_gradient - velocity_gradient.transpose()) * 0.5 +} + +/// Computes the strain rate (symmetric part) of a velocity gradient. +#[inline] +pub fn strain_rate(velocity_gradient: Matrix) -> Matrix { + (velocity_gradient + velocity_gradient.transpose()) * 0.5 +} + +/// Computes the deviatoric part of a tensor. +#[inline] +pub fn deviatoric_part(tensor: Matrix) -> Matrix { + DecomposedTensor::new(tensor).deviatoric_part +} + +/// Computes the spherical part (mean diagonal value) of a tensor. +#[inline] +pub fn spherical_part(tensor: Matrix) -> f32 { + trace(tensor) / DIM as f32 +} + +/// A tensor decomposed into its deviatoric and spherical parts. +#[derive(Clone, Copy)] +pub struct DecomposedTensor { + pub deviatoric_part: Matrix, + pub spherical_part: f32, +} + +impl DecomposedTensor { + /// Decomposes a tensor into its deviatoric and spherical parts. + #[inline] + pub fn new(tensor: Matrix) -> Self { + let spherical_part = trace(tensor) / DIM as f32; + let mut deviatoric_part = tensor; + + #[cfg(feature = "dim2")] + { + deviatoric_part.x_axis.x -= spherical_part; + deviatoric_part.y_axis.y -= spherical_part; + } + #[cfg(feature = "dim3")] + { + deviatoric_part.x_axis.x -= spherical_part; + deviatoric_part.y_axis.y -= spherical_part; + deviatoric_part.z_axis.z -= spherical_part; + } + + Self { + deviatoric_part, + spherical_part, + } + } + + /// Recomposes the tensor from its deviatoric and spherical parts. + #[inline] + pub fn recompose(&self) -> Matrix { + let mut result = self.deviatoric_part; + + #[cfg(feature = "dim2")] + { + result.x_axis.x += self.spherical_part; + result.y_axis.y += self.spherical_part; + } + #[cfg(feature = "dim3")] + { + result.x_axis.x += self.spherical_part; + result.y_axis.y += self.spherical_part; + result.z_axis.z += self.spherical_part; + } + + result + } +} + +/// CFL-based timestep bound using the speed of sound in an elastic material. +#[derive(Clone, Copy)] +pub struct ElasticitySoundSpeedTimestepBound { + pub alpha: f32, + pub bulk_modulus: f32, + pub shear_modulus: f32, +} + +impl ElasticitySoundSpeedTimestepBound { + /// Creates a new timestep bound from the CFL coefficient and Lame-derived moduli. + #[inline] + pub fn new(alpha: f32, bulk_modulus: f32, shear_modulus: f32) -> Self { + Self { + alpha, + bulk_modulus, + shear_modulus, + } + } + + /// Creates a new timestep bound from the CFL coefficient, Young modulus, and Poisson ratio. + #[inline] + pub fn from_elasticity(alpha: f32, young_modulus: f32, poisson_ratio: f32) -> Self { + Self { + alpha, + bulk_modulus: hook_to_bulk_modulus(young_modulus, poisson_ratio), + shear_modulus: hook_to_shear_modulus(young_modulus, poisson_ratio), + } + } + + /// Computes the CFL-based timestep bound. + /// + /// Uses the speed of sound for pressure waves and the physical velocity + /// to determine the maximum stable timestep. + #[inline] + pub fn timestep_bound( + &self, + density0: f32, + def_grad_det: f32, + velocity: Vector, + cell_width: f32, + ) -> f32 { + // Avoid division by zero. + let curr_density = density0 / f32::max(def_grad_det, 1.0e-6); + + // Speed of sound of pressure waves. + let sound_speed = sqrt((self.bulk_modulus + self.shear_modulus * 4.0 / 3.0) / curr_density); + + // Take the max between sound speed and physical velocity. + // Note that we don't calculate the speed of sound for the shear waves since that's + // always slower than for the pressure wave. + let max_speed = f32::max(velocity.length(), sound_speed); + self.alpha * cell_width / max_speed + } +} diff --git a/src_mpm_shaders/solver/boundary_condition.rs b/src_mpm_shaders/solver/boundary_condition.rs new file mode 100644 index 00000000..c183fc98 --- /dev/null +++ b/src_mpm_shaders/solver/boundary_condition.rs @@ -0,0 +1,114 @@ +use crate::Vector; + +/// Boundary condition type constants. +/// +/// Represented as `u32` for GPU compatibility instead of a Rust enum. +pub const BOUNDARY_CONDITION_STICK: u32 = 0; +pub const BOUNDARY_CONDITION_SLIP: u32 = 1; +pub const BOUNDARY_CONDITION_SEPARATE: u32 = 2; +pub const BOUNDARY_CONDITION_NON_REFLECTING: u32 = 3; + +/// A boundary condition applied to grid nodes at domain boundaries or collider surfaces. +/// +/// The `ty` field should be one of the `BOUNDARY_CONDITION_*` constants. +#[derive(Clone, Copy, Default, PartialEq, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct BoundaryCondition { + /// The type of boundary condition (see `BOUNDARY_CONDITION_*` constants). + pub ty: u32, + /// Friction coefficient. Only meaningful when `ty` is `Separate`. + pub friction: f32, + /// Pads the struct to 16 bytes so it can be an element of the `BodyMaterials` + /// uniform array (std140 requires a 16-byte array stride). Kept on the host + /// (`bytemuck`) type too so both layouts stay identical. + pub _pad0: u32, + pub _pad1: u32, +} + +impl BoundaryCondition { + pub const fn stick() -> BoundaryCondition { + BoundaryCondition::new(0, 0.0) + } + + pub const fn slip() -> BoundaryCondition { + BoundaryCondition::new(1, 0.0) + } + + pub const fn separate(friction: f32) -> BoundaryCondition { + BoundaryCondition::new(2, friction) + } +} + +/// Maximum number of collision bodies coupled to the MPM domain (CPIC limit). +pub const MAX_COLLISION_BODIES: usize = 16; + +/// Per-body boundary conditions, passed as a **uniform** (read-only, ≤16 bodies) +/// so the MPM kernels consuming it stay within the 8-storage-buffer WebGPU +/// limit. Indexed by body id. +#[derive(Clone, Copy)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct BodyMaterials { + pub mats: [BoundaryCondition; MAX_COLLISION_BODIES], +} + +impl BodyMaterials { + /// All-zero materials, usable as a `const` placeholder for kernels that + /// don't use rigid-body coupling (e.g. the non-CPIC P2G path). + pub const EMPTY: BodyMaterials = BodyMaterials { + mats: [BoundaryCondition::stick(); MAX_COLLISION_BODIES], + }; +} + +impl BoundaryCondition { + /// Creates a new boundary condition. + pub const fn new(ty: u32, friction: f32) -> Self { + Self { + ty, + friction, + _pad0: 0, + _pad1: 0, + } + } + + /// Projects a velocity according to this boundary condition. + /// `n` is the boundary normal (pointing inward). + pub fn project_velocity(&self, vel: Vector, n: Vector) -> Vector { + if self.ty == BOUNDARY_CONDITION_STICK { + return Vector::ZERO; + } + + if self.ty == BOUNDARY_CONDITION_SLIP { + let normal_vel = vel.dot(n); + let tangent_vel = vel - n * normal_vel; + return tangent_vel; + } + + if self.ty == BOUNDARY_CONDITION_SEPARATE { + let normal_vel = vel.dot(n); + + if normal_vel < 0.0 { + let tangent_vel = vel - n * normal_vel; + let tangent_vel_len = tangent_vel.length(); + let tangent_vel_dir = if tangent_vel_len > 1.0e-8 { + tangent_vel / tangent_vel_len + } else { + Vector::ZERO + }; + let projected_len = tangent_vel_len + self.friction * normal_vel; + let projected_len = if projected_len > 0.0 { + projected_len + } else { + 0.0 + }; + return tangent_vel_dir * projected_len; + } else { + return vel; + } + } + + // BOUNDARY_CONDITION_NON_REFLECTING or unknown: pass through. + vel + } +} diff --git a/src_mpm_shaders/solver/g2p.rs b/src_mpm_shaders/solver/g2p.rs new file mode 100644 index 00000000..ab927de5 --- /dev/null +++ b/src_mpm_shaders/solver/g2p.rs @@ -0,0 +1,474 @@ +//! Grid-to-Particle (G2P) transfer kernel. +//! +//! This kernel transfers grid node velocities back to particles using APIC +//! (Affine Particle-In-Cell) interpolation. It handles CPIC compatibility +//! checks, computes velocity gradients for the affine matrix, and accumulates +//! rigid body velocities for particles near colliders. + +use crate::PaddingExt; +use crate::grid::grid::*; +use crate::grid::kernel::*; +use crate::nexus_rbd_shaders::dynamics::{ + Velocity as BodyVelocity, WorldMassProperties as BodyMassProperties, +}; +use crate::solver::boundary_condition::{BodyMaterials, BoundaryCondition}; +use crate::solver::params::SimulationParams; +use crate::solver::particle::{Kinematics, Position}; +use crate::{Matrix, PaddedMatrix, Vector}; +use glamx::*; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::{ + macros::{spirv, spirv_bindgen}, + sync::workgroup_memory_barrier_with_group_sync, +}; +use unroll::unroll_for_loops; +/* + * Constants. + */ + +#[cfg(feature = "dim2")] +const NUM_SHARED_CELLS: usize = 10 * 10; +#[cfg(feature = "dim3")] +const NUM_SHARED_CELLS: usize = 6 * 6 * 6; + +const WORKGROUP_SIZE: u32 = 64; + +/* + * Global -> shared memory transfer. + */ + +#[inline] +#[unroll_for_loops] +fn global_shared_memory_transfers( + grid: &Grid, + hmap_entries: &[GridHashMapEntry], + nodes: &[Node], + tid: khal_std::glamx::UVec3, + active_block_vid: BlockVirtualId, + shared_nodes_vel: &mut [Vector; NUM_SHARED_CELLS], + shared_nodes_vel_incompatible: &mut [Vector; NUM_SHARED_CELLS], + shared_nodes_cdf: &mut [NodeCdf; NUM_SHARED_CELLS], +) { + let base_block_pos_int = active_block_vid.id; + + #[cfg(feature = "dim2")] + { + for i_loop in 0..2 { + for j_loop in 0..2 { + if !((i_loop == 1 && tid.x > 1) || (j_loop == 1 && tid.y > 1)) { + let octant = UVec2::new(i_loop as u32, j_loop as u32); + let octant_hid = grid.find_block_header_id( + hmap_entries, + &BlockVirtualId { + id: base_block_pos_int + IVec2::new(octant.x as i32, octant.y as i32), + }, + ); + let shared_index = octant * 8 + UVec2::new(tid.x, tid.y); + let flat_shared_index = + flatten_shared_index(shared_index.x, shared_index.y) as usize; + + if octant_hid.id != NONE { + let global_chunk_id = octant_hid.physical_id(); + let tid_xy = UVec2::new(tid.x, tid.y); + let global_node_id = global_chunk_id.node_id(tid_xy); + let node = nodes.read(global_node_id.id as usize); + shared_nodes_vel.write(flat_shared_index, node.momentum_velocity); + + if USE_CPIC { + shared_nodes_vel_incompatible + .write(flat_shared_index, node.momentum_velocity_incompatible); + shared_nodes_cdf.write(flat_shared_index, node.cdf); + } + } else { + shared_nodes_vel.write(flat_shared_index, Vector::ZERO); + + if USE_CPIC { + shared_nodes_vel_incompatible.write(flat_shared_index, Vector::ZERO); + shared_nodes_cdf.write(flat_shared_index, NodeCdf::NONE); + } + } + } + } + } + } + + #[cfg(feature = "dim3")] + { + for i_loop in 0..2 { + for j_loop in 0..2 { + for k_loop in 0..2 { + if !((i_loop == 1 && tid.x > 1) + || (j_loop == 1 && tid.y > 1) + || (k_loop == 1 && tid.z > 1)) + { + let octant = UVec3::new(i_loop as u32, j_loop as u32, k_loop as u32); + let octant_hid = grid.find_block_header_id( + hmap_entries, + &BlockVirtualId::new( + base_block_pos_int + + IVec3::new(octant.x as i32, octant.y as i32, octant.z as i32), + ), + ); + let tid_xyz = UVec3::new(tid.x, tid.y, tid.z); + let shared_index = octant * 4 + tid_xyz; + let flat_shared_index = + flatten_shared_index(shared_index.x, shared_index.y, shared_index.z) + as usize; + + if octant_hid.id != NONE { + let global_chunk_id = octant_hid.physical_id(); + let global_node_id = global_chunk_id.node_id(tid_xyz); + let node = nodes.read(global_node_id.id as usize); + shared_nodes_vel.write(flat_shared_index, node.momentum_velocity); + if USE_CPIC { + shared_nodes_vel_incompatible + .write(flat_shared_index, node.momentum_velocity_incompatible); + shared_nodes_cdf.write(flat_shared_index, node.cdf); + } + } else { + shared_nodes_vel.write(flat_shared_index, Vector::ZERO); + if USE_CPIC { + shared_nodes_vel_incompatible + .write(flat_shared_index, Vector::ZERO); + shared_nodes_cdf.write(flat_shared_index, NodeCdf::NONE); + } + } + } + } + } + } + } +} + +/* + * Per-particle G2P interpolation. + */ + +#[inline] +#[allow(clippy::too_many_arguments)] +#[unroll_for_loops] +fn particle_g2p( + body_vels: &[BodyVelocity], + body_mprops: &[BodyMassProperties], + body_materials: &BodyMaterials, + particles_pos: &[Position], + particles_kin: &mut [Kinematics], + particle_id: u32, + cell_width: f32, + _dt: f32, + shared_nodes_vel: &[Vector; NUM_SHARED_CELLS], + shared_nodes_vel_incompatible: &[Vector; NUM_SHARED_CELLS], + shared_nodes_cdf: &[NodeCdf; NUM_SHARED_CELLS], +) { + let mut rigid_vel = Vector::ZERO; + let mut velocity = Vector::ZERO; + let mut velocity_gradient = Matrix::ZERO; + let mut vel_grad_det = 0.0f32; + + // G2P + if particles_kin.at(particle_id as usize).enabled != 0 { + let particle_pos = particles_pos.read(particle_id as usize); + let particle_cdf = particles_kin.at(particle_id as usize).cdf; + let boundary_friction = particles_kin.at(particle_id as usize).boundary_friction; + + let inv_d = QuadraticKernel::inv_d(cell_width); + let ref_elt_pos_minus_particle_pos = particle_pos.dir_to_associated_grid_node(cell_width); + let w = QuadraticKernel::precompute_weights(ref_elt_pos_minus_particle_pos, cell_width); + + let assoc_cell_index_in_block = + particle_pos.associated_cell_index_in_block_off_by_one(cell_width); + + #[cfg(feature = "dim2")] + let packed_cell_index_in_block = + flatten_shared_index(assoc_cell_index_in_block.x, assoc_cell_index_in_block.y); + #[cfg(feature = "dim3")] + let packed_cell_index_in_block = flatten_shared_index( + assoc_cell_index_in_block.x, + assoc_cell_index_in_block.y, + assoc_cell_index_in_block.z, + ); + + for i in 0..27 { + // For loop unrolling, use the fixed bound (the maximum one between 2D and 3D). + if i < NBH_LEN { + let shift = NBH_SHIFTS.read(i); + let packed_shift = NBH_SHIFT_SHARED.read(i); + let shared_id = (packed_cell_index_in_block + packed_shift) as usize; + let mut cell_vel = shared_nodes_vel.read(shared_id); + + #[cfg(feature = "dim2")] + let dpt = ref_elt_pos_minus_particle_pos + + Vec2::new(shift.x as f32, shift.y as f32) * cell_width; + #[cfg(feature = "dim3")] + let dpt = ref_elt_pos_minus_particle_pos + + Vec3::new(shift.x as f32, shift.y as f32, shift.z as f32) * cell_width; + + if USE_CPIC { + let cell_cdf = shared_nodes_cdf.read(shared_id); + let is_compatible = particle_cdf.affinity.is_compatible(cell_cdf.affinities); + + if !is_compatible { + cell_vel = shared_nodes_vel_incompatible.read(shared_id); + + if cell_cdf.closest_id != NONE { + let body_vel = body_vels.read(cell_cdf.closest_id as usize); + let body_com = body_mprops.at(cell_cdf.closest_id as usize).com; + let body_material = body_materials.mats[cell_cdf.closest_id as usize]; + // Scale the collider's friction by this particle's own + // factor, so one surface can grip sand and let water + // slide. Applied here because CPIC resolves the + // boundary per particle, not per grid node. + let material = BoundaryCondition::new( + body_material.ty, + body_material.friction * boundary_friction, + ); + let cell_center = dpt + particle_pos.pt; + let body_pt_vel = body_vel.velocity_at_point(body_com, cell_center); + + cell_vel = body_pt_vel + + material + .project_velocity(cell_vel - body_pt_vel, particle_cdf.normal); + } + } + } + + #[cfg(feature = "dim2")] + let weight = vec3_extract(w[0], shift.x) * vec3_extract(w[1], shift.y); + #[cfg(feature = "dim3")] + let weight = vec3_extract(w[0], shift.x) + * vec3_extract(w[1], shift.y) + * vec3_extract(w[2], shift.z); + + velocity += cell_vel * weight; + velocity_gradient += outer_product(cell_vel, dpt) * (weight * inv_d); + vel_grad_det += weight * inv_d * cell_vel.dot(dpt); + } + } + + if USE_CPIC { + // Accumulate rigid body velocities for all affinity-linked colliders. + for i_collider in 0..16 { + if particle_cdf.affinity.bit(i_collider as u32) { + let body_vel = body_vels.read(i_collider); + let body_com = body_mprops.at(i_collider).com; + rigid_vel += body_vel.velocity_at_point(body_com, particle_pos.pt); + } + } + } + } + + if USE_CPIC { + particles_kin.at_mut(particle_id as usize).cdf.rigid_vel = rigid_vel; + } + + // Set the particle velocity, and store the velocity gradient into the affine matrix. + // The rest will be dealt with in the particle update kernel(s). + particles_kin.at_mut(particle_id as usize).affine = + PaddedMatrix::add_padding(velocity_gradient); + particles_kin.at_mut(particle_id as usize).vel_grad_det = vel_grad_det; + particles_kin.at_mut(particle_id as usize).velocity = velocity; +} + +/* + * GPU entry points. + */ + +/// GPU kernel: G2P transfer (2D). +/// +/// Transfers grid node velocities back to particles using APIC interpolation. +/// Dispatched with one workgroup per active block. +#[unroll_for_loops] +pub fn gpu_g2p_generic( + block_id: khal_std::glamx::UVec3, + tid: khal_std::glamx::UVec3, + tid_flat: u32, + params: &SimulationParams, + grid: &Grid, + hmap_entries: &[GridHashMapEntry], + active_blocks: &[ActiveBlockHeader], + nodes: &[Node], + sorted_particle_ids: &[u32], + particles_pos: &[Position], + particles_kin: &mut [Kinematics], + body_vels: &[BodyVelocity], + body_mprops: &[BodyMassProperties], + body_materials: &BodyMaterials, + shared_nodes_vel: &mut [Vector; NUM_SHARED_CELLS], + shared_nodes_vel_incompatible: &mut [Vector; NUM_SHARED_CELLS], + shared_nodes_cdf: &mut [NodeCdf; NUM_SHARED_CELLS], +) { + let bid = block_id.x; + // Force copy of the virtual ID (naga bug workaround). + let vid_ = active_blocks.at(bid as usize).virtual_id.id; + let vid = BlockVirtualId::new(vid_); + + // Block -> shared memory transfer. + global_shared_memory_transfers::( + grid, + hmap_entries, + nodes, + tid, + vid, + shared_nodes_vel, + shared_nodes_vel_incompatible, + shared_nodes_cdf, + ); + + // Sync after shared memory initialization. + workgroup_memory_barrier_with_group_sync(); + + // Particle update. Runs g2p on shared memory only. + let first_particle = active_blocks.at(bid as usize).first_particle; + let max_particle_id = first_particle + active_blocks.at(bid as usize).num_particles; + + let num_block_particles = max_particle_id - first_particle; + let max_iters = num_block_particles.div_ceil(WORKGROUP_SIZE); + let mut sorted_particle_id = first_particle + tid_flat; + for _ in 0..max_iters { + if sorted_particle_id >= max_particle_id { + break; + } + let particle_id = sorted_particle_ids.read(sorted_particle_id as usize); + particle_g2p::( + body_vels, + body_mprops, + body_materials, + particles_pos, + particles_kin, + particle_id, + grid.cell_width, + params.dt, + shared_nodes_vel, + shared_nodes_vel_incompatible, + shared_nodes_cdf, + ); + sorted_particle_id += WORKGROUP_SIZE; + } +} + +/* + * Shared memory flatten helpers for G2P. + * Note: different from P2G -- no shift subtraction since the truncated blocks + * are in the higher-index quadrants. + */ + +#[cfg(feature = "dim2")] +#[inline] +fn flatten_shared_index(x: u32, y: u32) -> u32 { + x + y * 10 +} + +#[cfg(feature = "dim3")] +#[inline] +fn flatten_shared_index(x: u32, y: u32, z: u32) -> u32 { + x + y * 6 + z * 6 * 6 +} + +/* + * Outer product helper. + */ + +#[cfg(feature = "dim2")] +#[inline] +fn outer_product(a: Vec2, b: Vec2) -> Mat2 { + Mat2::from_cols(a * b.x, a * b.y) +} + +#[cfg(feature = "dim3")] +#[inline] +fn outer_product(a: Vec3, b: Vec3) -> Mat3 { + Mat3::from_cols(a * b.x, a * b.y, a * b.z) +} + +/* + * Specialized entry points. + */ +#[spirv_bindgen] +#[cfg_attr(feature = "dim2", spirv(compute(threads(8, 8))))] +#[cfg_attr(feature = "dim3", spirv(compute(threads(4, 4, 4))))] +#[unroll_for_loops] +pub fn gpu_g2p( + #[spirv(workgroup_id)] block_id: khal_std::glamx::UVec3, + #[spirv(local_invocation_id)] tid: khal_std::glamx::UVec3, + #[spirv(local_invocation_index)] tid_flat: u32, + #[spirv(uniform, descriptor_set = 0, binding = 0)] params: &SimulationParams, + #[spirv(uniform, descriptor_set = 0, binding = 1)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] hmap_entries: &[GridHashMapEntry], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] active_blocks: &[ActiveBlockHeader], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] nodes: &[Node], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] sorted_particle_ids: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] particles_pos: &[Position], + #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] particles_kin: &mut [Kinematics], + #[spirv(storage_buffer, descriptor_set = 0, binding = 8)] body_vels: &[BodyVelocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 9)] body_mprops: &[BodyMassProperties], + #[spirv(uniform, descriptor_set = 0, binding = 10)] body_materials: &BodyMaterials, + // Shared memory. + #[spirv(workgroup)] shared_nodes_vel: &mut [Vector; NUM_SHARED_CELLS], + #[spirv(workgroup)] shared_nodes_vel_incompatible: &mut [Vector; NUM_SHARED_CELLS], + #[spirv(workgroup)] shared_nodes_cdf: &mut [NodeCdf; NUM_SHARED_CELLS], +) { + gpu_g2p_generic::( + block_id, + tid, + tid_flat, + params, + grid, + hmap_entries, + active_blocks, + nodes, + sorted_particle_ids, + particles_pos, + particles_kin, + body_vels, + body_mprops, + body_materials, + shared_nodes_vel, + shared_nodes_vel_incompatible, + shared_nodes_cdf, + ) +} + +#[spirv_bindgen] +#[cfg_attr(feature = "dim2", spirv(compute(threads(8, 8))))] +#[cfg_attr(feature = "dim3", spirv(compute(threads(4, 4, 4))))] +#[unroll_for_loops] +pub fn gpu_g2p_cpic( + #[spirv(workgroup_id)] block_id: khal_std::glamx::UVec3, + #[spirv(local_invocation_id)] tid: khal_std::glamx::UVec3, + #[spirv(local_invocation_index)] tid_flat: u32, + #[spirv(uniform, descriptor_set = 0, binding = 0)] params: &SimulationParams, + #[spirv(uniform, descriptor_set = 0, binding = 1)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] hmap_entries: &[GridHashMapEntry], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] active_blocks: &[ActiveBlockHeader], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] nodes: &[Node], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] sorted_particle_ids: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] particles_pos: &[Position], + #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] particles_kin: &mut [Kinematics], + #[spirv(storage_buffer, descriptor_set = 0, binding = 8)] body_vels: &[BodyVelocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 9)] body_mprops: &[BodyMassProperties], + #[spirv(uniform, descriptor_set = 0, binding = 10)] body_materials: &BodyMaterials, + // Shared memory. + #[spirv(workgroup)] shared_nodes_vel: &mut [Vector; NUM_SHARED_CELLS], + #[spirv(workgroup)] shared_nodes_vel_incompatible: &mut [Vector; NUM_SHARED_CELLS], + #[spirv(workgroup)] shared_nodes_cdf: &mut [NodeCdf; NUM_SHARED_CELLS], +) { + gpu_g2p_generic::( + block_id, + tid, + tid_flat, + params, + grid, + hmap_entries, + active_blocks, + nodes, + sorted_particle_ids, + particles_pos, + particles_kin, + body_vels, + body_mprops, + body_materials, + shared_nodes_vel, + shared_nodes_vel_incompatible, + shared_nodes_cdf, + ) +} diff --git a/src_mpm_shaders/solver/g2p_cdf.rs b/src_mpm_shaders/solver/g2p_cdf.rs new file mode 100644 index 00000000..abb2c79a --- /dev/null +++ b/src_mpm_shaders/solver/g2p_cdf.rs @@ -0,0 +1,407 @@ +//! Grid-to-Particle CDF (Contact Distance Field) transfer kernel. +//! +//! Transfers grid-level contact distance field data back to particles, computing +//! per-particle signed distances, contact normals, and the affinity/sign bits used +//! by CPIC in subsequent timesteps. + +use crate::grid::grid::*; +use crate::grid::kernel::*; +use crate::solver::params::SimulationParams; +use crate::solver::particle::{Cdf, Kinematics, Position}; +use crate::{Vector, abs}; +use crunchy::unroll; +use glamx::*; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::{ + macros::{spirv, spirv_bindgen}, + sync::workgroup_memory_barrier_with_group_sync, +}; +use unroll::unroll_for_loops; +/* + * Constants. + */ + +#[cfg(feature = "dim2")] +const NUM_SHARED_CELLS: usize = 10 * 10; +#[cfg(feature = "dim3")] +const NUM_SHARED_CELLS: usize = 6 * 6 * 6; + +const WORKGROUP_SIZE: u32 = 64; + +/* + * Shared memory flatten helpers for G2P (no subtraction, same as g2p.rs). + */ + +#[cfg(feature = "dim2")] +#[inline] +fn flatten_shared_index(x: u32, y: u32) -> u32 { + x + y * 10 +} + +#[cfg(feature = "dim3")] +#[inline] +fn flatten_shared_index(x: u32, y: u32, z: u32) -> u32 { + x + y * 6 + z * 6 * 6 +} + +/* + * Outer product helpers for the MLS reconstruction. + * In 2D: Vec3 x Vec3 -> Mat3 (homogeneous coordinates). + * In 3D: Vec4 x Vec4 -> Mat4 (homogeneous coordinates). + */ + +#[cfg(feature = "dim2")] +#[inline] +fn outer_product_3(a: Vec3, b: Vec3) -> Mat3 { + Mat3::from_cols(a * b.x, a * b.y, a * b.z) +} + +#[cfg(feature = "dim3")] +#[inline] +fn outer_product_4(a: Vec4, b: Vec4) -> Mat4 { + Mat4::from_cols(a * b.x, a * b.y, a * b.z, a * b.w) +} + +/* + * Helper: whether a shape has a solid interior (for sign bit computation). + */ + +#[inline] +fn shape_has_solid_interior(_i_collider: u32) -> bool { + // TODO: needs to be false for unoriented trimeshes and polylines, + // true for geometric primitives. + false +} + +/* + * Global -> shared memory transfer. + */ + +#[inline] +#[unroll_for_loops] +fn global_shared_memory_transfers( + grid: &Grid, + hmap_entries: &[GridHashMapEntry], + nodes: &[Node], + tid: khal_std::glamx::UVec3, + active_block_vid: BlockVirtualId, + shared_nodes: &mut [NodeCdf; NUM_SHARED_CELLS], +) { + let base_block_pos_int = active_block_vid.id; + + #[cfg(feature = "dim2")] + { + for i_loop in 0..2 { + for j_loop in 0..2 { + if !((i_loop == 1 && tid.x > 1) || (j_loop == 1 && tid.y > 1)) { + let octant = UVec2::new(i_loop as u32, j_loop as u32); + let octant_hid = grid.find_block_header_id( + hmap_entries, + &BlockVirtualId { + id: base_block_pos_int + IVec2::new(octant.x as i32, octant.y as i32), + }, + ); + let shared_index = octant * 8 + UVec2::new(tid.x, tid.y); + let flat_id = flatten_shared_index(shared_index.x, shared_index.y) as usize; + + if octant_hid.id != NONE { + let global_chunk_id = octant_hid.physical_id(); + let tid_xy = UVec2::new(tid.x, tid.y); + let global_node_id = global_chunk_id.node_id(tid_xy); + shared_nodes.write(flat_id, nodes.at(global_node_id.id as usize).cdf); + } else { + shared_nodes.write(flat_id, NodeCdf::NONE); + } + } + } + } + } + + #[cfg(feature = "dim3")] + { + for i_loop in 0..2 { + for j_loop in 0..2 { + for k_loop in 0..2 { + if !((i_loop == 1 && tid.x > 1) + || (j_loop == 1 && tid.y > 1) + || (k_loop == 1 && tid.z > 1)) + { + let octant = UVec3::new(i_loop as u32, j_loop as u32, k_loop as u32); + let octant_hid = grid.find_block_header_id( + hmap_entries, + &BlockVirtualId::new( + base_block_pos_int + + IVec3::new(octant.x as i32, octant.y as i32, octant.z as i32), + ), + ); + let tid_xyz = UVec3::new(tid.x, tid.y, tid.z); + let shared_index = octant * 4 + tid_xyz; + let shared_node_id = + flatten_shared_index(shared_index.x, shared_index.y, shared_index.z) + as usize; + + if octant_hid.id != NONE { + let global_chunk_id = octant_hid.physical_id(); + let global_node_id = global_chunk_id.node_id(tid_xyz); + shared_nodes + .write(shared_node_id, nodes.at(global_node_id.id as usize).cdf); + } else { + shared_nodes.write(shared_node_id, NodeCdf::NONE); + } + } + } + } + } + } +} + +/* + * Per-particle G2P CDF interpolation. + */ + +#[inline] +#[unroll_for_loops] +fn particle_g2p( + particles_pos: &[Position], + particles_kin: &mut [Kinematics], + particle_id: u32, + cell_width: f32, + _dt: f32, + shared_nodes: &[NodeCdf; NUM_SHARED_CELLS], +) { + let mut particle_affinity = AffinityBits::EMPTY; + let mut affinity_signs = [0.0f32; 16]; + + let prev_affinity = particles_kin.at(particle_id as usize).cdf.affinity; + let particle_pos = particles_pos.read(particle_id as usize); + let ref_elt_pos_minus_particle_pos = particle_pos.dir_to_associated_grid_node(cell_width); + let w = QuadraticKernel::precompute_weights(ref_elt_pos_minus_particle_pos, cell_width); + + let assoc_cell_index_in_block = + particle_pos.associated_cell_index_in_block_off_by_one(cell_width); + + #[cfg(feature = "dim2")] + let packed_cell_index_in_block = + flatten_shared_index(assoc_cell_index_in_block.x, assoc_cell_index_in_block.y); + #[cfg(feature = "dim3")] + let packed_cell_index_in_block = flatten_shared_index( + assoc_cell_index_in_block.x, + assoc_cell_index_in_block.y, + assoc_cell_index_in_block.z, + ); + + // Pass 1: Determine sign bits (Eqn. 21) and combine affinity masks. + for i in 0..27 { + // For loop unrolling, use the fixed bound (the maximum one between 2D and 3D). + if i < NBH_LEN { + let shift = NBH_SHIFTS.read(i); + let packed_shift = NBH_SHIFT_SHARED.read(i); + let cell_data = shared_nodes[(packed_cell_index_in_block + packed_shift) as usize]; + particle_affinity.set_unsigned_bits(cell_data.affinities); + + #[cfg(feature = "dim2")] + let weight = vec3_extract(w[0], shift.x) * vec3_extract(w[1], shift.y); + #[cfg(feature = "dim3")] + let weight = vec3_extract(w[0], shift.x) + * vec3_extract(w[1], shift.y) + * vec3_extract(w[2], shift.z); + + // Unrolled inner loop over 16 colliders. + // NOTE: `unroll_for_loops` doesn’t see through the closure so we use crunchy::unroll instead. + unroll! { + for i_collider in 0..16 { + let compatible = if cell_data.affinities.bit(i_collider as u32) { + 1.0f32 + } else { + 0.0f32 + }; + let sign = if cell_data.affinities.sign_bit(i_collider as u32) + && !shape_has_solid_interior(i_collider as u32) + { + -1.0f32 + } else { + 1.0f32 + }; + affinity_signs[i_collider] += compatible * weight * sign * cell_data.distance; + } + } + } + } + + // Convert the affinity signs to bits. + for i_collider in 0..16 { + if !prev_affinity.bit(i_collider as u32) { + // Only set the sign bit for affinities that didn't exist before. + if affinity_signs[i_collider] < 0.0 { + particle_affinity.set_sign_bit(i_collider as u32); + } + } else { + particle_affinity.or_sign_bit(prev_affinity, i_collider as u32); + } + } + + // Pass 2: MLS reconstruction of the contact distance/normal (Eq. 4). + #[cfg(feature = "dim2")] + let mut qtq = Mat3::ZERO; + #[cfg(feature = "dim2")] + let mut qtu = Vec3::ZERO; + #[cfg(feature = "dim3")] + let mut qtq = Mat4::ZERO; + #[cfg(feature = "dim3")] + let mut qtu = Vec4::ZERO; + + for i in 0..27 { + // For loop unrolling, use the fixed bound (the maximum one between 2D and 3D). + if i < NBH_LEN { + let shift = NBH_SHIFTS.read(i); + let packed_shift = NBH_SHIFT_SHARED.read(i); + let cell_data = shared_nodes[(packed_cell_index_in_block + packed_shift) as usize]; + + #[cfg(feature = "dim2")] + let dpt = ref_elt_pos_minus_particle_pos + + Vec2::new(shift.x as f32, shift.y as f32) * cell_width; + #[cfg(feature = "dim2")] + let weight = vec3_extract(w[0], shift.x) * vec3_extract(w[1], shift.y); + #[cfg(feature = "dim3")] + let dpt = ref_elt_pos_minus_particle_pos + + Vec3::new(shift.x as f32, shift.y as f32, shift.z as f32) * cell_width; + #[cfg(feature = "dim3")] + let weight = vec3_extract(w[0], shift.x) + * vec3_extract(w[1], shift.y) + * vec3_extract(w[2], shift.z); + + let combined_affinity = + cell_data.affinities.0 & particle_affinity.0 & AffinityBits::AFFINITY_BITS_MASK; + let sign_differences = ((cell_data.affinities.0 >> AffinityBits::SIGN_BITS_SHIFT) + ^ (particle_affinity.0 >> AffinityBits::SIGN_BITS_SHIFT)) + & combined_affinity; + + #[cfg(feature = "dim2")] + let p = Vec3::new(dpt.x, dpt.y, 1.0); + #[cfg(feature = "dim3")] + let p = Vec4::new(dpt.x, dpt.y, dpt.z, 1.0); + + if combined_affinity != 0 { + if sign_differences == 0 { + // All signs match: positive distance. + #[cfg(feature = "dim2")] + { + qtq += outer_product_3(p, p) * weight; + qtu += p * (weight * cell_data.distance); + } + #[cfg(feature = "dim3")] + { + qtq += outer_product_4(p, p) * weight; + qtu += p * (weight * cell_data.distance); + } + } else { + // Sign difference: negative distance. + #[cfg(feature = "dim2")] + { + qtq += outer_product_3(p, p) * weight; + qtu += p * (weight * -cell_data.distance); + } + #[cfg(feature = "dim3")] + { + qtq += outer_product_4(p, p) * weight; + qtu += p * (weight * -cell_data.distance); + } + } + } + } + } + + if qtq.determinant() > 1.0e-8 { + #[cfg(feature = "dim2")] + { + let result = qtq.inverse() * qtu; + let len = Vec2::new(result.x, result.y).length(); + let normal = if len > 1.0e-6 { + Vec2::new(result.x, result.y) / len + } else { + Vec2::ZERO + }; + particles_kin.at_mut(particle_id as usize).cdf = Cdf::new( + normal, + Vec2::ZERO, // PERF: init the rigid-velocities here instead of in g2p? + result.z, + particle_affinity, + ); + } + #[cfg(feature = "dim3")] + { + let result = qtq.inverse() * qtu; + let normal_vec = Vec3::new(result.x, result.y, result.z); + let len = normal_vec.length(); + let normal = if len > 1.0e-6 { + normal_vec / len + } else { + Vec3::ZERO + }; + particles_kin.at_mut(particle_id as usize).cdf = + Cdf::new(normal, Vec3::ZERO, result.w, particle_affinity); + } + } else { + // TODO: store the affinity in this case too? + particles_kin.at_mut(particle_id as usize).cdf = Cdf::zero(); + } +} + +/* + * GPU entry points. + */ + +/// GPU kernel: G2P CDF transfer (2D). +/// +/// Transfers grid CDF data back to particles using MLS reconstruction. +#[spirv_bindgen] +#[cfg_attr(feature = "dim2", spirv(compute(threads(8, 8))))] +#[cfg_attr(feature = "dim3", spirv(compute(threads(4, 4, 4))))] +pub fn gpu_g2p_cdf( + #[spirv(workgroup_id)] block_id: khal_std::glamx::UVec3, + #[spirv(local_invocation_id)] tid: khal_std::glamx::UVec3, + #[spirv(local_invocation_index)] tid_flat: u32, + #[spirv(uniform, descriptor_set = 0, binding = 0)] params: &SimulationParams, + #[spirv(uniform, descriptor_set = 0, binding = 1)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] hmap_entries: &[GridHashMapEntry], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] active_blocks: &[ActiveBlockHeader], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] nodes: &[Node], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] sorted_particle_ids: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] particles_pos: &[Position], + #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] particles_kin: &mut [Kinematics], + // Shared memory. + #[spirv(workgroup)] shared_nodes: &mut [NodeCdf; NUM_SHARED_CELLS], +) { + let bid = block_id.x; + let vid_ = active_blocks.at(bid as usize).virtual_id.id; + let vid = BlockVirtualId::new(vid_); + + // Block -> shared memory transfer. + global_shared_memory_transfers(grid, hmap_entries, nodes, tid, vid, shared_nodes); + + // Sync after shared memory initialization. + workgroup_memory_barrier_with_group_sync(); + + // Particle update. Runs g2p on shared memory only. + let first_particle = active_blocks.at(bid as usize).first_particle; + let max_particle_id = first_particle + active_blocks.at(bid as usize).num_particles; + + let num_block_particles = max_particle_id - first_particle; + let max_iters = num_block_particles.div_ceil(WORKGROUP_SIZE); + let mut sorted_particle_id = first_particle + tid_flat; + for _ in 0..max_iters { + if sorted_particle_id >= max_particle_id { + break; + } + let particle_id = sorted_particle_ids.read(sorted_particle_id as usize); + particle_g2p( + particles_pos, + particles_kin, + particle_id, + grid.cell_width, + params.dt, + shared_nodes, + ); + sorted_particle_id += WORKGROUP_SIZE; + } +} diff --git a/src_mpm_shaders/solver/grid_update.rs b/src_mpm_shaders/solver/grid_update.rs new file mode 100644 index 00000000..ccd59a17 --- /dev/null +++ b/src_mpm_shaders/solver/grid_update.rs @@ -0,0 +1,88 @@ +//! Grid update kernel: converts grid momentum to velocity and applies gravity. +//! +//! After P2G transfers momentum onto the grid, this kernel converts momentum to +//! velocity (dividing by mass), applies gravity, and clamps velocities so no node +//! moves more than one cell width per timestep. + +use crate::Vector; +use crate::grid::grid::*; +use crate::solver::params::SimulationParams; +use glamx::*; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; + +/// GPU kernel: grid update. +/// +/// Converts grid momentum to velocity, applies gravity, and clamps velocities. +/// Dispatched with one workgroup per active block, one thread per node. +#[spirv_bindgen] +#[cfg_attr(feature = "dim2", spirv(compute(threads(8, 8))))] +#[cfg_attr(feature = "dim3", spirv(compute(threads(4, 4, 4))))] +pub fn gpu_grid_update( + #[spirv(workgroup_id)] block_id: khal_std::glamx::UVec3, + #[spirv(local_invocation_id)] tid: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] sim_params: &SimulationParams, + #[spirv(uniform, descriptor_set = 0, binding = 1)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] active_blocks: &[ActiveBlockHeader], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] nodes: &mut [Node], +) { + let bid = block_id.x; + let vid = active_blocks.at(bid as usize).virtual_id; + let cell_width = grid.cell_width; + + let global_chunk_id = BlockHeaderId { id: bid }.physical_id(); + + #[cfg(feature = "dim2")] + let global_node_id = global_chunk_id.node_id(UVec2::new(tid.x, tid.y)); + #[cfg(feature = "dim3")] + let global_node_id = global_chunk_id.node_id(UVec3::new(tid.x, tid.y, tid.z)); + + #[cfg(feature = "dim2")] + let cell_pos = Vec2::new( + (vid.id.x * 8 + tid.x as i32) as f32, + (vid.id.y * 8 + tid.y as i32) as f32, + ) * cell_width; + #[cfg(feature = "dim3")] + let cell_pos = Vec3::new( + (vid.id.x * 4 + tid.x as i32) as f32, + (vid.id.y * 4 + tid.y as i32) as f32, + (vid.id.z * 4 + tid.z as i32) as f32, + ) * cell_width; + + let global_id = global_node_id.id as usize; + let momentum = nodes.at(global_id).momentum_velocity; + let mass = nodes.at(global_id).mass; + let momentum_incompatible = nodes.at(global_id).momentum_velocity_incompatible; + let mass_incompatible = nodes.at(global_id).mass_incompatible; + nodes.at_mut(global_id).momentum_velocity = + update_single_cell(sim_params, cell_width, cell_pos, momentum, mass); + nodes.at_mut(global_id).momentum_velocity_incompatible = update_single_cell( + sim_params, + cell_width, + cell_pos, + momentum_incompatible, + mass_incompatible, + ); +} + +/// Updates a single cell's momentum to velocity. +/// +/// Converts momentum to velocity by dividing by mass, adds gravity, +/// and clamps velocity to at most one cell width per timestep. +#[inline] +fn update_single_cell( + sim_params: &SimulationParams, + cell_width: f32, + _cell_pos: Vector, + momentum: Vector, + mass: f32, +) -> Vector { + let inv_mass = if mass > 0.0 { 1.0 / mass } else { 0.0 }; + let mut velocity = (momentum + sim_params.gravity * (mass * sim_params.dt)) * inv_mass; + + // Clamp the velocity so it doesn't exceed 1 grid cell in one step. + let vel_limit = Vector::splat(cell_width / sim_params.dt); + velocity = velocity.clamp(-vel_limit, vel_limit); + + velocity +} diff --git a/src_mpm_shaders/solver/grid_update_cdf.rs b/src_mpm_shaders/solver/grid_update_cdf.rs new file mode 100644 index 00000000..ca6fd774 --- /dev/null +++ b/src_mpm_shaders/solver/grid_update_cdf.rs @@ -0,0 +1,132 @@ +//! Grid update CDF kernel: runs collision detection on each grid node. +//! +//! This kernel detects collisions between the grid nodes and the collision shapes, +//! storing the resulting contact distance field (CDF) data in each node's `cdf` field. + +use crate::grid::grid::*; +use crate::nexus_rbd_shaders::dynamics::Velocity as BodyVelocity; +use crate::nexus_rbd_shaders::shapes::Shape; +use crate::{Pose, Vector}; +use glamx::*; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; + +/// Performs collision detection for a single grid node against all collision shapes. +/// +/// Returns a `NodeCdf` with the closest collider distance, affinity bits, and collider ID. +#[inline] +fn collide( + collision_shapes: &[Shape], + collision_shape_poses: &[Pose], + cell_width: f32, + point: Vector, +) -> NodeCdf { + let mut cdf = NodeCdf::NONE; + + let dist_cap = Vector::splat(cell_width * 1.5); + + // NOTE: we iterate using a fixed upper bound to avoid dynamic buffer length queries + // that may not be available in all SPIR-V environments. The caller must ensure + // the shapes buffer length matches the actual number of shapes. + for i in 0..collision_shapes.len() { + let shape = collision_shapes.read(i); + let shape_pose = collision_shape_poses.read(i); + let shape_type = shape.shape_type(); + + use crate::nexus_rbd_shaders::shapes::{SHAPE_TYPE_POLYLINE, SHAPE_TYPE_TRIMESH}; + if shape_type != SHAPE_TYPE_POLYLINE && shape_type != SHAPE_TYPE_TRIMESH { + let proj = shape.project_point_on_boundary(shape_pose, point); + let dpt = proj.point - point; + + let abs_dpt = dpt.abs(); + #[cfg(feature = "dim2")] + let within_cap = abs_dpt.x <= dist_cap.x && abs_dpt.y <= dist_cap.y; + #[cfg(feature = "dim3")] + let within_cap = + abs_dpt.x <= dist_cap.x && abs_dpt.y <= dist_cap.y && abs_dpt.z <= dist_cap.z; + + if proj.is_inside || within_cap { + let dist = dpt.length(); + if dist < cdf.distance { + cdf.closest_id = i as u32; + cdf.distance = dist; + } + cdf.affinities.set_bit(i as u32, proj.is_inside); + } + } + } + + cdf +} + +/// GPU kernel: grid update CDF. +/// +/// For each active grid node, runs collision detection against all collision shapes +/// and writes the resulting `NodeCdf` (distance, affinity bits, closest collider ID). +/// +/// Dispatched with one workgroup per active block, one thread per node in the block. +#[cfg(feature = "dim2")] +#[spirv_bindgen] +#[spirv(compute(threads(8, 8)))] +pub fn gpu_grid_update_cdf( + #[spirv(workgroup_id)] block_id: khal_std::glamx::UVec3, + #[spirv(local_invocation_id)] tid: khal_std::glamx::UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] active_blocks: &[ActiveBlockHeader], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] collision_shapes: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] collision_shape_poses: &[Pose], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] nodes: &mut [Node], +) { + let bid = block_id.x; + let vid = active_blocks.at(bid as usize).virtual_id; + + let global_chunk_id = BlockHeaderId { id: bid }.physical_id(); + let tid_xy = UVec2::new(tid.x, tid.y); + let global_node_id = global_chunk_id.node_id(tid_xy); + let cell_pos = Vec2::new( + (vid.id.x * 8 + tid.x as i32) as f32, + (vid.id.y * 8 + tid.y as i32) as f32, + ) * grid.cell_width; + + let global_id = global_node_id.id; + nodes.at_mut(global_id as usize).cdf = collide( + collision_shapes, + collision_shape_poses, + grid.cell_width, + cell_pos, + ); +} + +/// GPU kernel: grid update CDF (3D version). +#[cfg(feature = "dim3")] +#[spirv_bindgen] +#[spirv(compute(threads(4, 4, 4)))] +pub fn gpu_grid_update_cdf( + #[spirv(workgroup_id)] block_id: khal_std::glamx::UVec3, + #[spirv(local_invocation_id)] tid: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] active_blocks: &[ActiveBlockHeader], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] collision_shapes: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] collision_shape_poses: &[Pose], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] nodes: &mut [Node], +) { + let bid = block_id.x; + let vid = active_blocks.at(bid as usize).virtual_id; + + let global_chunk_id = BlockHeaderId { id: bid }.physical_id(); + let tid_xyz = UVec3::new(tid.x, tid.y, tid.z); + let global_node_id = global_chunk_id.node_id(tid_xyz); + let cell_pos = Vec3::new( + (vid.id.x * 4 + tid.x as i32) as f32, + (vid.id.y * 4 + tid.y as i32) as f32, + (vid.id.z * 4 + tid.z as i32) as f32, + ) * grid.cell_width; + + let global_id = global_node_id.id; + nodes.at_mut(global_id as usize).cdf = collide( + collision_shapes, + collision_shape_poses, + grid.cell_width, + cell_pos, + ); +} diff --git a/src_mpm_shaders/solver/grid_update_collide.rs b/src_mpm_shaders/solver/grid_update_collide.rs new file mode 100644 index 00000000..be03b97d --- /dev/null +++ b/src_mpm_shaders/solver/grid_update_collide.rs @@ -0,0 +1,183 @@ +//! Grid update CDF kernel: runs collision detection on each grid node. +//! +//! This kernel detects collisions between the grid nodes and the collision shapes, +//! storing the resulting contact distance field (CDF) data in each node's `cdf` field. + +use crate::grid::grid::*; +use crate::nexus_rbd_shaders::dynamics::{ + Velocity as BodyVelocity, WorldMassProperties as BodyMassProperties, +}; +use crate::nexus_rbd_shaders::shapes::Shape; +use crate::solver::boundary_condition::{BodyMaterials, BoundaryCondition}; +use crate::solver::params::SimulationParams; +use crate::{Pose, Vector}; +use glamx::*; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use nexus_rbd_shaders::{MAX_FLT, PaddedVector}; + +struct Collision { + normal: Vector, + distance: f32, + closest_id: usize, +} + +/// Performs collision detection for a single grid node against all collision shapes. +/// +/// Returns a `NodeCdf` with the closest collider distance, affinity bits, and collider ID. +#[inline] +#[cfg_attr(not(feature = "dim3"), allow(unused_variables))] +fn collide( + collision_shapes: &[Shape], + collision_shape_poses: &[Pose], + collision_shape_indices: &[u32], + collision_shape_vertices: &[PaddedVector], + cell_width: f32, + point: Vector, +) -> Collision { + let dist_cap = Vector::splat(cell_width * 1.5); + let mut collision = Collision { + normal: Vector::ZERO, + distance: MAX_FLT, + closest_id: 0, + }; + + // NOTE: we iterate using a fixed upper bound to avoid dynamic buffer length queries + // that may not be available in all SPIR-V environments. The caller must ensure + // the shapes buffer length matches the actual number of shapes. + for i in 0..collision_shapes.len() { + let shape = collision_shapes.read(i); + let shape_pose = collision_shape_poses.read(i); + let shape_type = shape.shape_type(); + + use crate::nexus_rbd_shaders::shapes::{SHAPE_TYPE_POLYLINE, SHAPE_TYPE_TRIMESH}; + + #[cfg(feature = "dim3")] + let (proj, valid) = if shape_type == SHAPE_TYPE_TRIMESH { + let mesh = shape.to_trimesh(); + let local_pt = shape_pose.inverse() * point; + let (mut proj, valid) = mesh.project_local_point( + collision_shape_indices, + collision_shape_vertices, + local_pt, + dist_cap.x, + ); + // Transform the projected point back to world space. + proj.point = shape_pose * proj.point; + (proj, valid) + } else { + (shape.project_point_on_boundary(shape_pose, point), true) + }; + + #[cfg(feature = "dim2")] + let (proj, valid) = (shape.project_point_on_boundary(shape_pose, point), true); + + if valid { + let dpt = proj.point - point; + let abs_dpt = dpt.abs(); + #[cfg(feature = "dim2")] + let within_cap = abs_dpt.x <= dist_cap.x && abs_dpt.y <= dist_cap.y; + #[cfg(feature = "dim3")] + let within_cap = + abs_dpt.x <= dist_cap.x && abs_dpt.y <= dist_cap.y && abs_dpt.z <= dist_cap.z; + + if proj.is_inside || within_cap { + let sign = if proj.is_inside { -1.0 } else { 1.0 }; + let distance = dpt.length(); + let normal = dpt / (distance * -sign); + let signed_dist = sign * distance; + if signed_dist < collision.distance { + collision.distance = signed_dist; + collision.normal = normal; + collision.closest_id = i; + } + } + } + } + + collision +} + +// TODO(PERF): merge with gpu_grid_update. +#[spirv_bindgen] +#[cfg_attr(feature = "dim2", spirv(compute(threads(8, 8))))] +#[cfg_attr(feature = "dim3", spirv(compute(threads(4, 4, 4))))] +pub fn gpu_grid_update_collide( + #[spirv(workgroup_id)] block_id: khal_std::glamx::UVec3, + #[spirv(local_invocation_id)] tid: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] params: &SimulationParams, + #[spirv(uniform, descriptor_set = 0, binding = 1)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] active_blocks: &[ActiveBlockHeader], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] collision_shapes: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] collision_shape_poses: &[Pose], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] + collision_shape_vertices: &[PaddedVector], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] collision_shape_indices: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] body_vels: &[BodyVelocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 8)] body_mprops: &[BodyMassProperties], + #[spirv(uniform, descriptor_set = 0, binding = 9)] body_materials: &BodyMaterials, + #[spirv(storage_buffer, descriptor_set = 0, binding = 10)] nodes: &mut [Node], +) { + let dt = params.dt; + let bid = block_id.x; + let vid = active_blocks.at(bid as usize).virtual_id; + + let global_chunk_id = BlockHeaderId { id: bid }.physical_id(); + + let global_node_id; + let cell_pt; + + #[cfg(feature = "dim2")] + { + let tid_xy = UVec2::new(tid.x, tid.y); + global_node_id = global_chunk_id.node_id(tid_xy); + cell_pt = Vec2::new( + (vid.id.x * 8 + tid.x as i32) as f32, + (vid.id.y * 8 + tid.y as i32) as f32, + ) * grid.cell_width; + } + + #[cfg(feature = "dim3")] + { + let tid_xyz = UVec3::new(tid.x, tid.y, tid.z); + global_node_id = global_chunk_id.node_id(tid_xyz); + cell_pt = Vec3::new( + (vid.id.x * 4 + tid.x as i32) as f32, + (vid.id.y * 4 + tid.y as i32) as f32, + (vid.id.z * 4 + tid.z as i32) as f32, + ) * grid.cell_width; + } + + let global_id = global_node_id.id; + let cell_width = grid.cell_width; + let collision = collide( + collision_shapes, + collision_shape_poses, + collision_shape_indices, + collision_shape_vertices, + cell_width, + cell_pt, + ); + + if collision.distance != MAX_FLT { + // Found a collision, apply the boundary condition. + let body_vel = body_vels.at(collision.closest_id); + let body_com = body_mprops.at(collision.closest_id).com; + let body_vel_at_grid_pos = body_vel.velocity_at_point(body_com, cell_pt); + let node_vel = nodes.at(global_id as usize).momentum_velocity; + let body_material = body_materials.mats[collision.closest_id]; + let delta_vel = node_vel - body_vel_at_grid_pos; + let normal_vel = delta_vel.dot(collision.normal); + let margin = cell_width; + + if collision.distance <= margin { + let corrected_vel = + body_vel_at_grid_pos + body_material.project_velocity(delta_vel, collision.normal); + + nodes.at_mut(global_id as usize).momentum_velocity = corrected_vel; + } else if -normal_vel * dt > collision.distance - margin { + let excess_vel = (normal_vel + (collision.distance - margin) / dt) * collision.normal; + nodes.at_mut(global_id as usize).momentum_velocity -= excess_vel; + } + } +} diff --git a/src_mpm_shaders/solver/mod.rs b/src_mpm_shaders/solver/mod.rs new file mode 100644 index 00000000..68f6a62b --- /dev/null +++ b/src_mpm_shaders/solver/mod.rs @@ -0,0 +1,15 @@ +pub mod boundary_condition; +pub mod g2p; +pub mod g2p_cdf; +pub mod grid_update; +pub mod grid_update_cdf; +pub mod grid_update_collide; +pub mod p2g; +pub mod p2g_cdf; +pub mod params; +pub mod particle; +pub mod particle_update; +pub mod prep_readback; +pub mod rigid_impulses; +pub mod rigid_particle_update; +pub mod timestep_bound; diff --git a/src_mpm_shaders/solver/p2g.rs b/src_mpm_shaders/solver/p2g.rs new file mode 100644 index 00000000..4b362c35 --- /dev/null +++ b/src_mpm_shaders/solver/p2g.rs @@ -0,0 +1,449 @@ +//! Particle-to-Grid (P2G) transfer kernel (scatter style). +//! +//! The core MPM kernel that transfers particle data (momentum, mass, affine matrix) +//! onto the grid nodes. Dispatched with one workgroup per active block. +//! +//! The CPIC (Compatible Particle-In-Cell) variant also handles affinity checks: particles +//! incompatible with a node (different side of a collider) contribute to the node's +//! `incompatible` momentum field instead, and impulses are accumulated for the rigid body +//! coupling. + +use crate::grid::grid::*; +use crate::grid::kernel::QuadraticKernel; +use crate::nexus_rbd_shaders::dynamics::Velocity as BodyVelocity; +use crate::solver::boundary_condition::{BodyMaterials, BoundaryCondition}; +use crate::solver::particle::{Kinematics, Position}; +use crate::{AngVector, Matrix, PaddingExt, TWO_WAYS_COUPLING_ENABLED, Vector}; +use glamx::*; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use khal_std::sync::{atomic_add_i32, workgroup_memory_barrier_with_group_sync}; + +/// Workgroup size: one thread per grid node of a block (8*8 in 2D, 4*4*4 in 3D). +const WORKGROUP_SIZE: usize = 64; + +/// Integer impulse atomic struct for accumulating impulses across threads. +/// +/// Uses integer atomics to avoid floating-point atomic limitations on GPU. +/// The COM (center of mass) is stored alongside to reduce binding count. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct IntegerImpulse { + pub com: Vector, + pub linear_x: i32, + pub linear_y: i32, + #[cfg(feature = "dim3")] + pub linear_z: i32, + #[cfg(feature = "dim3")] + pub _padding_a: i32, + #[cfg(feature = "dim2")] + pub angular: i32, + #[cfg(feature = "dim2")] + pub _padding: i32, + #[cfg(feature = "dim3")] + pub angular_x: i32, + #[cfg(feature = "dim3")] + pub angular_y: i32, + #[cfg(feature = "dim3")] + pub angular_z: i32, + #[cfg(feature = "dim3")] + pub _padding_b: [i32; 2], +} + +const FLOAT_TO_INT_FACTOR: f32 = 1e5; + +/// Converts a float to an integer for atomic accumulation. +#[inline] +fn flt2int(flt: f32) -> i32 { + (flt * FLOAT_TO_INT_FACTOR) as i32 +} + +/// Generic scatter-style P2G shared by the plain and CPIC entry points. +#[allow(clippy::too_many_arguments)] +pub fn gpu_p2g_generic( + block_id: khal_std::glamx::UVec3, + tid: khal_std::glamx::UVec3, + tid_flat: u32, + grid: &Grid, + active_blocks: &[ActiveBlockHeader], + sorted_particle_ids: &[u32], + particles_pos: &[Position], + particles_kin: &[Kinematics], + nodes: &mut [Node], + body_vels: &[BodyVelocity], + body_impulses: &mut [IntegerImpulse], + body_materials: &BodyMaterials, + // Shared memory: one chunk of particles, loaded cooperatively (one per thread). + shared_pos: &mut [Position; WORKGROUP_SIZE], + shared_vel_mass: &mut [(Vector, f32); WORKGROUP_SIZE], + shared_affine: &mut [Matrix; WORKGROUP_SIZE], + // NOTE: these are only read/written under CPIC, but rust-gpu can't coerce a workgroup + // `&mut [T; N]` to `&mut [T]`, so both entry points pass fixed-size arrays. + shared_affinities: &mut [AffinityBits; WORKGROUP_SIZE], + shared_normals: &mut [Vector; WORKGROUP_SIZE], + // Per-particle slab key: associated-cell coordinate along the slowest node + // axis, relative to the block. The per-chunk culling bounds derive from it. + shared_zkey: &mut [i32; WORKGROUP_SIZE], +) { + let bid = block_id.x; + let cell_width = grid.cell_width; + let inv_cell_width = 1.0 / cell_width; + + // Force copy of the virtual ID (naga bug workaround, as in the original kernel). + let vid = active_blocks.at(bid as usize).virtual_id.id; + + // This thread owns one grid node of the block. + #[cfg(feature = "dim2")] + let (local_cell, cell_pos) = { + let lc = UVec2::new(tid.x, tid.y); + let c = vid * 8 + IVec2::new(tid.x as i32, tid.y as i32); + (lc, Vec2::new(c.x as f32, c.y as f32) * cell_width) + }; + #[cfg(feature = "dim3")] + let (local_cell, cell_pos) = { + let lc = UVec3::new(tid.x, tid.y, tid.z); + let c = vid * 4 + IVec3::new(tid.x as i32, tid.y as i32, tid.z as i32); + ( + lc, + Vec3::new(c.x as f32, c.y as f32, c.z as f32) * cell_width, + ) + }; + + let gid = BlockHeaderId { id: bid } + .physical_id() + .node_id(local_cell) + .id as usize; + + // Per-node CDF data (computed by an earlier pass), needed only for CPIC. + let node_affinity = nodes.at(gid).cdf.affinities; + let collider_id = if USE_CPIC { + nodes.at(gid).cdf.closest_id + } else { + NONE + }; + + let mut acc_mv = Vector::ZERO; + let mut acc_mass = 0.0f32; + let mut acc_mv_incompatible = Vector::ZERO; + let mut acc_mass_incompatible = 0.0f32; + let mut impulse = Vector::ZERO; + #[cfg(feature = "dim2")] + let mut ang_impulse: AngVector = 0.0; + #[cfg(feature = "dim3")] + let mut ang_impulse: AngVector = Vec3::ZERO; + + let first = active_blocks.at(bid as usize).first_particle; + let num = active_blocks.at(bid as usize).num_particles_with_extras; + let last = first + num; + // End of the primaries segment: the sorted slab keys are ascending within the + // primaries and within the extras, so chunk bounds need this boundary. + let primaries_end = first + active_blocks.at(bid as usize).num_particles; + + // This thread's node slab along the sort axis (the slowest-varying node axis). + #[cfg(feature = "dim2")] + let node_slab = tid.y as i32; + #[cfg(feature = "dim3")] + let node_slab = tid.z as i32; + + // Number of workgroup-sized chunks. Capped on the web (bounded loop with a per-chunk + // guard) so the workgroup barriers stay in uniform control flow; off the web, the + // exact count is used and the guard is always true. + // We set it to 128, which would be exceeded in quite degenerate situations (for example + // if we end up with more than 28 particles per cell in the entire block and its neighborhood. + // The typical particle count per cell is 8). We could make the limit bigger, but starting + // at 256 we’ve seen it result in a measurable negative performance impact. + #[cfg(feature = "web-compat")] + let num_chunks = 128u32; + #[cfg(not(feature = "web-compat"))] + let num_chunks = num.div_ceil(WORKGROUP_SIZE as u32); + + for chunk in 0..num_chunks { + let chunk_base = first + chunk * WORKGROUP_SIZE as u32; + let active_chunk = chunk_base < last; + + // Wait for the previous chunk's readers before overwriting shared memory. + workgroup_memory_barrier_with_group_sync(); + + if active_chunk { + let load_idx = chunk_base + tid_flat; + let slot = tid_flat as usize; + if load_idx < last { + let pid = sorted_particle_ids.read(load_idx as usize); + let pos = particles_pos.read(pid as usize); + + // Slab key along the sort axis, relative to the block. Must match + // the sort's bucket key (same associated-cell rounding, same clamp + // at -2) so the shared keys stay ascending within each segment. + let assoc_cell = (pos.pt / cell_width).round() - Vector::ONE; + #[cfg(feature = "dim2")] + let zkey = (assoc_cell.y as i32 - vid.y * 8).max(-2); + #[cfg(feature = "dim3")] + let zkey = (assoc_cell.z as i32 - vid.z * 4).max(-2); + shared_zkey.write(slot, zkey); + + let pkin = particles_kin.at(pid as usize); + if pkin.enabled != 0 { + // The first component holds the raw velocity when CPIC is on (the + // impulse computation needs it) or the precomputed momentum + // (velocity * mass) otherwise, so the inner loop never recomputes it. + let vel_or_momentum = if USE_CPIC { + pkin.velocity + } else { + pkin.velocity * pkin.mass + }; + shared_pos.write(slot, pos); + shared_vel_mass.write(slot, (vel_or_momentum, pkin.mass)); + shared_affine.write(slot, pkin.affine.remove_padding()); + if USE_CPIC { + shared_affinities.write(slot, pkin.cdf.affinity); + shared_normals.write(slot, pkin.cdf.normal); + } + } else { + // Disabled particle: contribute nothing (mass = 0). + shared_pos.at_mut(slot).pt = Vector::ZERO; + shared_vel_mass.write(slot, (Vector::ZERO, 0.0)); + shared_affine.write(slot, Matrix::ZERO); + if USE_CPIC { + shared_affinities.write(slot, AffinityBits::EMPTY); + shared_normals.write(slot, Vector::ZERO); + } + } + } + } + + workgroup_memory_barrier_with_group_sync(); + + if active_chunk { + // `chunk_len` is uniform across the workgroup. + let chunk_len = (last - chunk_base).min(WORKGROUP_SIZE as u32); + + // Per-chunk slab bounds, exact thanks to the within-block sort: the keys + // are ascending within the primaries and within the extras, so the range + // is given by the chunk's first/last key, plus the two values around the + // primaries/extras boundary if the chunk straddles it. + let mut zmin = shared_zkey.read(0); + let mut zmax = shared_zkey.read((chunk_len - 1) as usize); + if primaries_end > chunk_base && primaries_end < chunk_base + chunk_len { + let b = (primaries_end - chunk_base) as usize; + zmin = zmin.min(shared_zkey.read(b)); + zmax = zmax.max(shared_zkey.read(b - 1)); + } + + // A particle with slab key `a` only influences nodes in slabs [a, a + 2]: + // skip the whole chunk if this thread's node slab is outside the chunk's + // dilated slab range. When that holds for every thread of a warp (e.g. a + // chunk of extras below the block vs. the upper-half warp), the warp skips + // the chunk entirely. + let in_range = node_slab >= zmin && node_slab <= zmax + 2; + let culled_len = if in_range { chunk_len } else { 0 }; + for p in 0..culled_len { + let p = p as usize; + let pos = shared_pos.read(p); + // `vel_or_momentum` is the precomputed momentum (non-CPIC) or the raw + // velocity (CPIC); see the chunk load above. + let (vel_or_momentum, mass) = shared_vel_mass.read(p); + let dpt = cell_pos - pos.pt; + + #[cfg(feature = "dim2")] + let weight = QuadraticKernel::eval(dpt.x * inv_cell_width) + * QuadraticKernel::eval(dpt.y * inv_cell_width); + #[cfg(feature = "dim3")] + let weight = QuadraticKernel::eval(dpt.x * inv_cell_width) + * QuadraticKernel::eval(dpt.y * inv_cell_width) + * QuadraticKernel::eval(dpt.z * inv_cell_width); + + // The quadratic kernel is exactly zero outside the 3-node support, the + // common case for the dense node x particle cross product. + if weight != 0.0 { + let affine = shared_affine.at(p); + let momentum = if USE_CPIC { + vel_or_momentum * mass + } else { + vel_or_momentum + }; + let vel_contribution = (affine * dpt + momentum) * weight; + let mass_contribution = mass * weight; + + if USE_CPIC { + let particle_affinity = shared_affinities.read(p); + if !particle_affinity.is_compatible(node_affinity) { + if TWO_WAYS_COUPLING_ENABLED && collider_id != NONE { + let particle_normal = shared_normals.read(p); + let body_vel = body_vels.read(collider_id as usize); + let body_com = body_impulses.at(collider_id as usize).com; + let body_material = body_materials.mats[collider_id as usize]; + let cell_center = cell_pos; + let body_pt_vel = body_vel.velocity_at_point(body_com, cell_center); + let particle_ghost_vel = body_pt_vel + + body_material.project_velocity( + vel_or_momentum - body_pt_vel, + particle_normal, + ); + let delta_impulse = + (vel_or_momentum - particle_ghost_vel) * (weight * mass); + let lever_arm = body_com - cell_center; + + #[cfg(feature = "dim2")] + { + ang_impulse += + delta_impulse.dot(Vec2::new(lever_arm.y, -lever_arm.x)); + } + #[cfg(feature = "dim3")] + { + ang_impulse += delta_impulse.cross(lever_arm); + } + + impulse += delta_impulse; + } + + acc_mv_incompatible += vel_contribution; + acc_mass_incompatible += mass_contribution; + } else { + acc_mv += vel_contribution; + acc_mass += mass_contribution; + } + } else { + acc_mv += vel_contribution; + acc_mass += mass_contribution; + } + } + } + } + } + + // Write the node state to global memory (one write per node, no atomics). + nodes.at_mut(gid).momentum_velocity = acc_mv; + nodes.at_mut(gid).mass = acc_mass; + nodes.at_mut(gid).momentum_velocity_incompatible = acc_mv_incompatible; + nodes.at_mut(gid).mass_incompatible = acc_mass_incompatible; + + if USE_CPIC { + // Apply the accumulated impulse to the closest body using integer atomics. + if TWO_WAYS_COUPLING_ENABLED && collider_id != NONE { + let ci = collider_id as usize; + #[cfg(feature = "dim2")] + { + atomic_add_i32(&mut body_impulses.at_mut(ci).linear_x, flt2int(impulse.x)); + atomic_add_i32(&mut body_impulses.at_mut(ci).linear_y, flt2int(impulse.y)); + atomic_add_i32(&mut body_impulses.at_mut(ci).angular, flt2int(ang_impulse)); + } + #[cfg(feature = "dim3")] + { + atomic_add_i32(&mut body_impulses.at_mut(ci).linear_x, flt2int(impulse.x)); + atomic_add_i32(&mut body_impulses.at_mut(ci).linear_y, flt2int(impulse.y)); + atomic_add_i32(&mut body_impulses.at_mut(ci).linear_z, flt2int(impulse.z)); + atomic_add_i32( + &mut body_impulses.at_mut(ci).angular_x, + flt2int(ang_impulse.x), + ); + atomic_add_i32( + &mut body_impulses.at_mut(ci).angular_y, + flt2int(ang_impulse.y), + ); + atomic_add_i32( + &mut body_impulses.at_mut(ci).angular_z, + flt2int(ang_impulse.z), + ); + } + } + } +} + +/* + * GPU entry points. + */ + +/// GPU kernel: scatter-style P2G transfer (no CPIC). +/// +/// Dispatched with one workgroup per active block. +#[spirv_bindgen] +#[cfg_attr(feature = "dim2", spirv(compute(threads(8, 8))))] +#[cfg_attr(feature = "dim3", spirv(compute(threads(4, 4, 4))))] +pub fn gpu_p2g( + #[spirv(workgroup_id)] block_id: khal_std::glamx::UVec3, + #[spirv(local_invocation_id)] tid: khal_std::glamx::UVec3, + #[spirv(local_invocation_index)] tid_flat: u32, + #[spirv(uniform, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] active_blocks: &[ActiveBlockHeader], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] sorted_particle_ids: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] particles_pos: &[Position], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] particles_kin: &[Kinematics], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] nodes: &mut [Node], + #[spirv(workgroup)] shared_pos: &mut [Position; WORKGROUP_SIZE], + #[spirv(workgroup)] shared_vel_mass: &mut [(Vector, f32); WORKGROUP_SIZE], + #[spirv(workgroup)] shared_affine: &mut [Matrix; WORKGROUP_SIZE], + #[spirv(workgroup)] shared_affinities: &mut [AffinityBits; WORKGROUP_SIZE], + #[spirv(workgroup)] shared_normals: &mut [Vector; WORKGROUP_SIZE], + #[spirv(workgroup)] shared_zkey: &mut [i32; WORKGROUP_SIZE], +) { + gpu_p2g_generic::( + block_id, + tid, + tid_flat, + grid, + active_blocks, + sorted_particle_ids, + particles_pos, + particles_kin, + nodes, + &[], + &mut [], + &BodyMaterials::EMPTY, + shared_pos, + shared_vel_mass, + shared_affine, + shared_affinities, + shared_normals, + shared_zkey, + ); +} + +/// GPU kernel: scatter-style P2G transfer with CPIC rigid-body coupling. +/// +/// Dispatched with one workgroup per active block. +#[spirv_bindgen] +#[cfg_attr(feature = "dim2", spirv(compute(threads(8, 8))))] +#[cfg_attr(feature = "dim3", spirv(compute(threads(4, 4, 4))))] +pub fn gpu_p2g_cpic( + #[spirv(workgroup_id)] block_id: khal_std::glamx::UVec3, + #[spirv(local_invocation_id)] tid: khal_std::glamx::UVec3, + #[spirv(local_invocation_index)] tid_flat: u32, + #[spirv(uniform, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] active_blocks: &[ActiveBlockHeader], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] sorted_particle_ids: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] particles_pos: &[Position], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] particles_kin: &[Kinematics], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] nodes: &mut [Node], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] body_vels: &[BodyVelocity], + #[spirv(uniform, descriptor_set = 0, binding = 7)] body_materials: &BodyMaterials, + #[spirv(storage_buffer, descriptor_set = 0, binding = 8)] + body_impulses: &mut [IntegerImpulse], + #[spirv(workgroup)] shared_pos: &mut [Position; WORKGROUP_SIZE], + #[spirv(workgroup)] shared_vel_mass: &mut [(Vector, f32); WORKGROUP_SIZE], + #[spirv(workgroup)] shared_affine: &mut [Matrix; WORKGROUP_SIZE], + #[spirv(workgroup)] shared_affinities: &mut [AffinityBits; WORKGROUP_SIZE], + #[spirv(workgroup)] shared_normals: &mut [Vector; WORKGROUP_SIZE], + #[spirv(workgroup)] shared_zkey: &mut [i32; WORKGROUP_SIZE], +) { + gpu_p2g_generic::( + block_id, + tid, + tid_flat, + grid, + active_blocks, + sorted_particle_ids, + particles_pos, + particles_kin, + nodes, + body_vels, + body_impulses, + body_materials, + shared_pos, + shared_vel_mass, + shared_affine, + shared_affinities, + shared_normals, + shared_zkey, + ); +} diff --git a/src_mpm_shaders/solver/p2g_cdf.rs b/src_mpm_shaders/solver/p2g_cdf.rs new file mode 100644 index 00000000..f93f454f --- /dev/null +++ b/src_mpm_shaders/solver/p2g_cdf.rs @@ -0,0 +1,259 @@ +//! Particle-to-Grid CDF (Contact Distance Field) transfer kernel (scatter style). +//! +//! Transfers collision primitives (segments in 2D, triangles in 3D) from rigid body +//! surface particles onto nearby grid nodes, computing the signed distance and the +//! CPIC affinity bits. Dispatched with one workgroup per active block. + +use crate::grid::grid::*; +use crate::solver::particle::{Position, RigidParticleIndices}; +use crate::{IVector, Vector, abs}; +use glamx::*; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use khal_std::sync::workgroup_memory_barrier_with_group_sync; +use nexus_rbd_shaders::PaddedVector; + +/// Workgroup size: one thread per grid node of a block (8*8 in 2D, 4*4*4 in 3D). +const WORKGROUP_SIZE: usize = 64; + +/// A collision primitive stored in shared memory. +/// In 2D: segment (two endpoints). In 3D: triangle (three vertices). +#[derive(Clone, Copy, Default)] +#[repr(C)] +pub struct SharedPrimitive { + a: Vector, + b: Vector, + #[cfg(feature = "dim3")] + c: Vector, +} + +/* + * Segment projection helper (2D). + */ + +#[cfg(feature = "dim2")] +#[inline] +fn project_local_point_on_segment(a: Vec2, b: Vec2, point: Vec2) -> Vec2 { + let ab = b - a; + let ap = point - a; + let ab_sqnorm = ab.dot(ab); + + if ab_sqnorm < 1.0e-10 { + return a; + } + + let t = ap.dot(ab) / ab_sqnorm; + let t = t.clamp(0.0, 1.0); + a + ab * t +} + +/* + * GPU entry points. + */ + +/// GPU kernel: P2G CDF transfer. +/// +/// Dispatched with one workgroup per active block. +#[spirv_bindgen] +#[cfg_attr(feature = "dim2", spirv(compute(threads(8, 8))))] +#[cfg_attr(feature = "dim3", spirv(compute(threads(4, 4, 4))))] +pub fn gpu_p2g_cdf( + #[spirv(workgroup_id)] block_id: khal_std::glamx::UVec3, + #[spirv(local_invocation_id)] tid: khal_std::glamx::UVec3, + #[spirv(local_invocation_index)] tid_flat: u32, + #[spirv(uniform, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] active_blocks: &[ActiveBlockHeader], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] sorted_rigid_particle_ids: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] rigid_particles_pos: &[Position], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] collider_vertices: &[PaddedVector], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] + rigid_particle_indices: &[RigidParticleIndices], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] nodes: &mut [Node], + // Shared memory: one chunk of rigid particles, loaded cooperatively (one per thread). + #[spirv(workgroup)] shared_primitives: &mut [SharedPrimitive; WORKGROUP_SIZE], + #[spirv(workgroup)] shared_collider_ids: &mut [u32; WORKGROUP_SIZE], + #[spirv(workgroup)] shared_assoc_cells: &mut [IVector; WORKGROUP_SIZE], +) { + let bid = block_id.x; + let cell_width = grid.cell_width; + + // Force copy of the virtual ID (naga bug workaround, as in the original kernel). + let vid = active_blocks.at(bid as usize).virtual_id.id; + + // This thread owns one grid node of the block. + #[cfg(feature = "dim2")] + let (local_cell, cell_int) = { + let lc = UVec2::new(tid.x, tid.y); + (lc, vid * 8 + IVec2::new(tid.x as i32, tid.y as i32)) + }; + #[cfg(feature = "dim3")] + let (local_cell, cell_int) = { + let lc = UVec3::new(tid.x, tid.y, tid.z); + ( + lc, + vid * 4 + IVec3::new(tid.x as i32, tid.y as i32, tid.z as i32), + ) + }; + #[cfg(feature = "dim2")] + let cell_pos = Vec2::new(cell_int.x as f32, cell_int.y as f32) * cell_width; + #[cfg(feature = "dim3")] + let cell_pos = Vec3::new(cell_int.x as f32, cell_int.y as f32, cell_int.z as f32) * cell_width; + + let gid = BlockHeaderId { id: bid } + .physical_id() + .node_id(local_cell) + .id as usize; + + // Merge into the CDF computed by the analytical-shapes pass (`grid_update_cdf`). + let mut node_cdf = nodes.at(gid).cdf; + + let first = active_blocks.at(bid as usize).first_rigid_particle; + let num = active_blocks + .at(bid as usize) + .num_rigid_particles_with_extras; + let last = first + num; + + // Number of workgroup-sized chunks. Capped on the web (bounded loop with a per-chunk + // guard) so the workgroup barriers stay in uniform control flow; off the web, the + // exact count is used and the guard is always true. Rigid particles are surface + // samples spaced roughly one cell apart, so 128 chunks (8192 particles per block) + // is far beyond anything reachable in practice. + #[cfg(feature = "web-compat")] + let num_chunks = 16u32; + #[cfg(not(feature = "web-compat"))] + let num_chunks = num.div_ceil(WORKGROUP_SIZE as u32); + + for chunk in 0..num_chunks { + let chunk_base = first + chunk * WORKGROUP_SIZE as u32; + let active_chunk = chunk_base < last; + + // Wait for the previous chunk's readers before overwriting shared memory. + workgroup_memory_barrier_with_group_sync(); + + if active_chunk { + let load_idx = chunk_base + tid_flat; + let slot = tid_flat as usize; + if load_idx < last { + let pid = sorted_rigid_particle_ids.read(load_idx as usize); + let rigid_idx = rigid_particle_indices.read(pid as usize); + shared_collider_ids.write(slot, rigid_idx.collider); + + #[cfg(feature = "dim2")] + shared_primitives.write( + slot, + SharedPrimitive { + a: collider_vertices.read(rigid_idx.segment.x as usize).0, + b: collider_vertices.read(rigid_idx.segment.y as usize).0, + }, + ); + #[cfg(feature = "dim3")] + shared_primitives.write( + slot, + SharedPrimitive { + a: collider_vertices.read(rigid_idx.triangle.x as usize).0, + b: collider_vertices.read(rigid_idx.triangle.y as usize).0, + c: collider_vertices.read(rigid_idx.triangle.z as usize).0, + }, + ); + + // The cell the particle is associated with (off-by-one convention): the + // primitive only influences nodes in the 3-cell range starting there. + // NOTE: must divide (not multiply by the inverse) to round exactly like + // the sort kernels' block association. + let assoc = + (rigid_particles_pos.read(pid as usize).pt / cell_width).round() - Vector::ONE; + #[cfg(feature = "dim2")] + shared_assoc_cells.write(slot, IVec2::new(assoc.x as i32, assoc.y as i32)); + #[cfg(feature = "dim3")] + shared_assoc_cells.write( + slot, + IVec3::new(assoc.x as i32, assoc.y as i32, assoc.z as i32), + ); + } + } + + workgroup_memory_barrier_with_group_sync(); + + if active_chunk { + // `chunk_len` is uniform across the workgroup. + let chunk_len = (last - chunk_base).min(WORKGROUP_SIZE as u32); + for p in 0..chunk_len { + let p = p as usize; + + // Restrict each primitive's influence to the quadratic-stencil-shaped + // 3-cell neighbourhood of its associated cell, matching the original + // gather implementation. + let shift = cell_int - shared_assoc_cells.read(p); + #[cfg(feature = "dim2")] + let in_range = shift.x >= 0 && shift.x <= 2 && shift.y >= 0 && shift.y <= 2; + #[cfg(feature = "dim3")] + let in_range = shift.x >= 0 + && shift.x <= 2 + && shift.y >= 0 + && shift.y <= 2 + && shift.z >= 0 + && shift.z <= 2; + + if in_range { + let collider_id = shared_collider_ids.read(p); + let primitive = shared_primitives.read(p); + + #[cfg(feature = "dim2")] + { + // Project on Segment. + let proj = + project_local_point_on_segment(primitive.a, primitive.b, cell_pos); + // Check if this is a valid projection (not clamped to an endpoint). + let not_at_a = proj.x != primitive.a.x || proj.y != primitive.a.y; + let not_at_b = proj.x != primitive.b.x || proj.y != primitive.b.y; + if not_at_a && not_at_b { + let dpt = cell_pos - proj; + let distance = dpt.length(); + let ab = primitive.b - primitive.a; + let sign = dpt.dot(Vec2::new(-ab.y, ab.x)) < 0.0; + node_cdf.affinities.set_bit(collider_id, sign); + + if distance < node_cdf.distance { + node_cdf.distance = distance; + node_cdf.closest_id = collider_id; + } + } + } + + #[cfg(feature = "dim3")] + { + // Project on Triangle. + let ap = cell_pos - primitive.a; + let bp = cell_pos - primitive.b; + let cp = cell_pos - primitive.c; + let ab = primitive.b - primitive.a; + let ac = primitive.c - primitive.a; + let bc = primitive.c - primitive.b; + let n = ab.cross(ac); + let n_length = n.length(); + + if n_length != 0.0 + && ab.cross(n).dot(ap) <= 0.0 + && bc.cross(n).dot(bp) <= 0.0 + && ac.cross(n).dot(cp) >= 0.0 + // Positive sign due to `ac` instead of `ca`. + { + // Valid projection on the face interior. + let signed_dist = n.dot(ap) / n_length; + let distance = abs(signed_dist); + node_cdf.affinities.set_bit(collider_id, signed_dist < 0.0); + + if distance < node_cdf.distance { + node_cdf.distance = distance; + node_cdf.closest_id = collider_id; + } + } + } + } + } + } + } + + // Write the node cdf to global memory. + nodes.at_mut(gid).cdf = node_cdf; +} diff --git a/src_mpm_shaders/solver/params.rs b/src_mpm_shaders/solver/params.rs new file mode 100644 index 00000000..892cb55e --- /dev/null +++ b/src_mpm_shaders/solver/params.rs @@ -0,0 +1,17 @@ +use crate::Vector; + +/// Parameters for the MPM simulation. +/// +/// In 2D, a padding field is added after gravity to satisfy uniform size/alignment requirements. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct SimulationParams { + /// Gravity vector (Vec2 in 2D, Vec3 in 3D). + pub gravity: Vector, + /// Padding required in 2D due to uniform size limits. + #[cfg(feature = "dim2")] + pub padding: f32, + /// The simulation timestep. + pub dt: f32, +} diff --git a/src_mpm_shaders/solver/particle.rs b/src_mpm_shaders/solver/particle.rs new file mode 100644 index 00000000..fadb2e22 --- /dev/null +++ b/src_mpm_shaders/solver/particle.rs @@ -0,0 +1,238 @@ +use crate::grid::grid::AffinityBits; +use crate::{Matrix, PaddedMatrix, UVector, Vector}; + +/// A particle position in the MPM grid. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct Position { + /// The particle's world-space position. + pub pt: Vector, + #[cfg(feature = "dim3")] + pub padding: u32, +} + +impl Position { + pub fn new(pt: Vector) -> Self { + Self { + pt, + #[cfg(feature = "dim3")] + padding: 0, + } + } +} + +/// Contact distance field data for a particle. +/// +/// Stores the result of the collision detection between a particle and the +/// nearest rigid collider surface. +#[derive(Clone, Copy, Default)] +#[cfg_attr( + not(target_arch_is_gpu), + derive(Debug, PartialEq, bytemuck::Pod, bytemuck::Zeroable) +)] +#[repr(C)] +pub struct Cdf { + /// The contact normal direction. + pub normal: Vector, + // NOTE: to avoid padding, the location of this field in the struct depends on whether + // we are in 2D or 3D. + /// The signed distance from the particle to the closest collider surface. + #[cfg(feature = "dim3")] + pub signed_distance: f32, + /// The velocity of the rigid body at the closest surface point. + pub rigid_vel: Vector, + /// The signed distance from the particle to the closest collider surface. + #[cfg(feature = "dim2")] + pub signed_distance: f32, + /// Affinity bits for CPIC compatibility checks. + pub affinity: AffinityBits, +} + +impl Cdf { + /// Creates a new zeroed Cdf. + pub fn zero() -> Self { + Self { + normal: Vector::ZERO, + rigid_vel: Vector::ZERO, + signed_distance: 0.0, + affinity: AffinityBits::default(), + } + } + + /// Creates a new Cdf with the given values. + pub fn new( + normal: Vector, + rigid_vel: Vector, + signed_distance: f32, + affinity: AffinityBits, + ) -> Self { + Self { + normal, + rigid_vel, + signed_distance, + affinity, + } + } +} + +/// Indices referencing the rigid body element closest to a particle. +/// +/// In 2D, this references a segment (edge) by its two vertex indices. +/// In 3D, this references a triangle by its three vertex indices. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct RigidParticleIndices { + /// The vertex indices of the closest segment (2D) or triangle (3D). + #[cfg(feature = "dim2")] + pub segment: UVector, + /// The vertex indices of the closest segment (2D) or triangle (3D). + #[cfg(feature = "dim3")] + pub triangle: UVector, + /// The collider index this element belongs to. + pub collider: u32, + /// SPIR-V padding: UVec2 has align(8) in SPIR-V, so stride must be a multiple of 8. + #[cfg(feature = "dim2")] + pub _pad: u32, +} + +/// Core kinematic state for APIC particle-grid transfers. +/// +/// Contains the fields needed by P2G and G2P kernels: velocity, mass, affine matrix, +/// external forces, and particle status. Separated from deformation and material +/// properties. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct Kinematics { + /// During `particle_update`, this contains the velocity gradient. + /// After `particle_update`, this contains the affine matrix for APIC transfer. + pub affine: PaddedMatrix, + /// The particle's velocity. + pub velocity: Vector, + /// Determinant of the velocity gradient. Only the fluid models track it. + #[cfg(feature = "dim3")] + pub vel_grad_det: f32, + /// Additional user-defined force applied to the particle, multiplied by dt. + /// Reset at each `particle_update` invocation. + /// Stored as force * dt so that dt is not needed during p2g. + pub force_dt: Vector, + /// Determinant of the velocity gradient. Only the fluid models track it. + #[cfg(feature = "dim2")] + pub vel_grad_det: f32, + /// The particle's mass. + pub mass: f32, + /// Whether this particle is enabled (non-zero = enabled). + pub enabled: u32, + /// Multiplier applied to a collider's friction for this particle. + /// + /// Boundary friction is a property of the *pair*, not of the collider alone: + /// water running over the same floor that sand piles up on should barely feel + /// it. 1 uses the collider's friction as given, 0 makes the particle slide + /// freely along the surface while still being stopped from passing through it. + /// + /// Only has an effect under CPIC (the default), which resolves the boundary + /// per particle. Without it the boundary condition is applied to grid nodes, + /// which hold a blend of every material touching them, and a single node + /// cannot be frictional for the sand and slippery for the water at once. + /// + /// Lives here rather than in `ParticleProperties` because both transfer + /// kernels already bind the kinematics buffer, and the one that needs it is + /// at its storage-buffer limit. + pub boundary_friction: f32, + /// Alignment padding before the CDF field. + #[cfg(feature = "dim3")] + pub _padding: [u32; 2], + /// Contact distance field data for CPIC rigid body coupling. + pub cdf: Cdf, + /// Tail padding so the struct size is a multiple of its alignment. + #[cfg(feature = "dim2")] + pub _tail_padding: [u32; 2], +} + +/// Static per-particle properties that are read-only on the GPU. +/// +/// These fields are set once during particle creation and never modified by any +/// GPU shader. +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct ParticleProperties { + /// The particle's initial volume (reference configuration). + pub init_volume: f32, + /// The particle's initial radius. + pub init_radius: f32, + /// Rayleigh mass-proportional damping coefficient (1/s). + pub damping: f32, + /// Phase value for multi-material mixing. + pub phase: f32, + /// Whether this particle is fixed in place (non-zero = fixed). + pub fixed: u32, + /// Index of the group this particle belongs to. Carries no physics; the + /// render kernels use it to look up a color in the viewer's group palette. + pub group_id: u32, + /// Pad to 32 bytes for GPU cache line alignment. + pub padding: [u32; 2], +} + +/* + * + * Grid-related position helper functions. + * + */ + +impl Position { + /// Returns the position of the grid node closest to the particle. + /// + /// This rounds the particle position to the nearest cell center. + #[inline] + pub fn closest_grid_pos(&self, cell_width: f32) -> Vector { + (self.pt / cell_width).round() * cell_width + } + + /// Returns the position of the "associated" grid node for the particle. + /// + /// The associated node is one cell before the closest node in each dimension, + /// which is the base node for the 3-node (quadratic) B-spline stencil. + #[inline] + pub fn associated_grid_pos(&self, cell_width: f32) -> Vector { + ((self.pt / cell_width).round() - Vector::ONE) * cell_width + } + + /// Returns the index of the associated cell within its block, offset by one. + /// + /// This is what maps a particle to its block in the sparse grid. Blocks are + /// 8x8 in 2D and 4x4x4 in 3D. + #[inline] + pub fn associated_cell_index_in_block_off_by_one(&self, cell_width: f32) -> UVector { + let assoc_cell = (self.pt / cell_width).round() - Vector::ONE; + #[cfg(feature = "dim2")] + let assoc_block = (assoc_cell / 8.0).floor() * 8.0; + #[cfg(feature = "dim3")] + let assoc_block = (assoc_cell / 4.0).floor() * 4.0; + // The result is always non-negative, so the cast to unsigned is safe. + #[cfg(feature = "dim2")] + { + let diff = assoc_cell - assoc_block; + UVector::new(diff.x as u32, diff.y as u32) + } + #[cfg(feature = "dim3")] + { + let diff = assoc_cell - assoc_block; + UVector::new(diff.x as u32, diff.y as u32, diff.z as u32) + } + } + + /// Returns the direction vector from the particle to the closest grid node. + #[inline] + pub fn dir_to_closest_grid_node(&self, cell_width: f32) -> Vector { + self.closest_grid_pos(cell_width) - self.pt + } + + /// Returns the direction vector from the particle to the associated grid node. + #[inline] + pub fn dir_to_associated_grid_node(&self, cell_width: f32) -> Vector { + self.associated_grid_pos(cell_width) - self.pt + } +} diff --git a/src_mpm_shaders/solver/particle_update.rs b/src_mpm_shaders/solver/particle_update.rs new file mode 100644 index 00000000..32ce42db --- /dev/null +++ b/src_mpm_shaders/solver/particle_update.rs @@ -0,0 +1,201 @@ +//! Particle update kernel: advection, deformation gradient update, constitutive model, +//! and APIC affine matrix computation. +//! +//! The main per-particle update kernel, run after G2P has transferred grid velocities +//! back to particles. Diverged particles (NaN position or non-positive deformation +//! gradient determinant) are automatically disabled. + +use crate::PaddingExt; +use crate::grid::grid::Grid; +use crate::grid::kernel::QuadraticKernel; +use crate::models::default::{DefaultParticleModel, GpuParticleModel}; +use crate::models::interfaces::{MODEL_FLAGS_FLUID, ParticleUpdateData}; +use crate::solver::boundary_condition::{BOUNDARY_CONDITION_SLIP, BoundaryCondition}; +use crate::solver::params::SimulationParams; +use crate::solver::particle::{Kinematics, ParticleProperties, Position}; +use crate::{DIM, Matrix, PaddedMatrix, Vector, diag}; + +/// Largest relative volume change a fluid particle may undergo in one substep. +const MAX_VOLUME_RATE: f32 = 0.1; +/// Bounds on the per-axis stretch of a fluid particle, i.e. on `J^(1/DIM)`. +/// A particle that leaves this range has diverged, not deformed. +const MIN_FLUID_STRETCH: f32 = 0.5; +const MAX_FLUID_STRETCH: f32 = 2.0; +use glamx::*; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; + +/// Phase data for multi-material mixing (currently unused placeholder). +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct Phase { + pub phase: f32, + pub max_stretch: f32, +} + +/// Penalty coefficient for collision response. +const PENALTY_COEFF: f32 = 1.0e3; + +/// Checks if a Vector contains any NaN components. +#[inline] +fn vector_has_nan(v: Vector) -> bool { + #[cfg(feature = "dim2")] + { + v.x.is_nan() || v.y.is_nan() + } + #[cfg(feature = "dim3")] + { + v.x.is_nan() || v.y.is_nan() || v.z.is_nan() + } +} + +/// Main particle update kernel. +/// +/// Each thread processes one particle, performing the full update cycle: +/// advection, deformation gradient update, constitutive model evaluation, +/// and APIC affine matrix computation. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_particle_update( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] params: &SimulationParams, + #[spirv(uniform, descriptor_set = 0, binding = 1)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + particles_model: &mut [GpuParticleModel], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] particles_pos: &mut [Position], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] particles_kin: &mut [Kinematics], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] + particles_def_grad: &mut [PaddedMatrix], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] + particles_props: &[ParticleProperties], + #[spirv(uniform, descriptor_set = 0, binding = 7)] particles_len: &u32, +) { + let particle_id = invocation_id.x; + + if particle_id >= *particles_len { + return; + } + + let flags = DefaultParticleModel::model_flags(particles_model, particle_id); + let dt = params.dt; + let cell_width = grid.cell_width; + let mut kin = particles_kin.read(particle_id as usize); + let cdf = kin.cdf; + let mut def_grad = particles_def_grad.read(particle_id as usize); + let props = particles_props.read(particle_id as usize); + let particle_pos = particles_pos.at(particle_id as usize).pt; + + /* + * Update velocity. + */ + // Reproject velocity if the particle is penetrating a rigid collider. + // TODO: double check that we never need the reprojection below. + // This reprojection isn't part of the original MPM-MLS/CPIC paper but we added + // it at some point as it appeared that we'd still get some penetrating particles. + // However, that might have been caused by other bugs so it is unsure if we need to + // keep it now. + if cdf.signed_distance < -0.05 * cell_width { + let slip = BoundaryCondition::new(BOUNDARY_CONDITION_SLIP, 0.0); + kin.velocity = + cdf.rigid_vel + slip.project_velocity(kin.velocity - cdf.rigid_vel, cdf.normal); + } + + // Clamp the max velocity a particle can get. + // TODO: clamp the grid velocities instead? + let vel_len = kin.velocity.length(); + if vel_len > cell_width / dt { + kin.velocity = kin.velocity / vel_len * cell_width / dt; + } + + // Apply Rayleigh mass-proportional damping (implicit integration for stability). + // v_new = v / (1 + damping * dt) + kin.velocity /= 1.0 + props.damping * dt; + + // If the particle is fixed, clear its velocity. + // This isn't ideal (this should typically be handled on the grid) but we sometimes + // need sub-grid-sized fixed particles. + if props.fixed != 0 { + kin.velocity = Vector::ZERO; + } + + /* + * Update position. + */ + let new_particle_pos = particle_pos + kin.velocity * dt; + + /* + * Penalty impulse. + */ + // TODO: apply the penalty impulse as an extra force on the grid instead of + // changing the particle velocity directly? + if cdf.signed_distance < -0.05 * cell_width { + let corrected_dist = cdf.signed_distance.max(-0.3 * cell_width); + let impulse = (dt * -corrected_dist * PENALTY_COEFF) * cdf.normal; + kin.velocity += impulse; + } + + /* + * Deformation gradient update. + */ + if (flags & MODEL_FLAGS_FLUID) == 0 { + // Solid path: F_new = F + (vel_grad * dt) * F + // NOTE: the velocity gradient was stored in the affine buffer. + def_grad = def_grad + (kin.affine * dt) * def_grad; + } else { + // Fluid path: only the volumetric part is tracked, as `F = d * I`. + // + // `vel_grad_det` is the divergence, so the volume ratio obeys + // `J' = J tr(grad v)`. Since `J = d^DIM`, the diagonal entry grows at + // `1 / DIM` of that rate; dividing here is what keeps + // `def_grad.determinant()` equal to `J`, which the constitutive model, + // the CFL bound and the rendering all assume. + let def_grad0 = def_grad.x_axis.x; + let rate = kin.vel_grad_det * dt / DIM as f32; + // Bound the per-substep change so one bad divergence estimate at a free + // surface cannot send a particle's volume to zero or to infinity. + let rate = rate.clamp(-MAX_VOLUME_RATE, MAX_VOLUME_RATE); + let new_def_grad_diag_elt = + (def_grad0 + rate * def_grad0).clamp(MIN_FLUID_STRETCH, MAX_FLUID_STRETCH); + def_grad = PaddedMatrix::add_padding(diag(Vector::splat(new_def_grad_diag_elt))); + } + + /* + * Constitutive model. + */ + // `kin.affine` still holds the velocity gradient at this point; it is + // overwritten with the APIC affine matrix further down. + let velocity_gradient = kin.affine.remove_padding(); + let update_data = ParticleUpdateData::new(dt, cell_width, particle_id, velocity_gradient); + let update_result = DefaultParticleModel::update(particles_model, &update_data, &mut def_grad); + + /* + * Affine matrix for APIC transfer. + */ + let inv_d = QuadraticKernel::inv_d(cell_width); + // NOTE: the velocity gradient was stored in the affine buffer. + let affine = kin.affine * kin.mass + - PaddedMatrix::add_padding( + update_result.kirchoff_stress * (props.init_volume * inv_d * dt), + ); + + /* + * Write back the new particle properties. + */ + // Check for NaN and invalid deformation gradients. + if !vector_has_nan(new_particle_pos) && def_grad.determinant() > 0.0 { + particles_pos.at_mut(particle_id as usize).pt = new_particle_pos; + kin.affine = affine; + } else { + // This particle diverged, disable it. + kin.enabled = 0; + kin.velocity = Vector::ZERO; + def_grad = PaddedMatrix::IDENTITY; + kin.affine = PaddedMatrix::ZERO; + kin.mass = 0.0; + } + kin.force_dt = Vector::ZERO; + + particles_kin.write(particle_id as usize, kin); + particles_def_grad.write(particle_id as usize, def_grad); +} diff --git a/src_mpm_shaders/solver/prep_readback.rs b/src_mpm_shaders/solver/prep_readback.rs new file mode 100644 index 00000000..2c3d5eeb --- /dev/null +++ b/src_mpm_shaders/solver/prep_readback.rs @@ -0,0 +1,469 @@ +//! Readback preparation shader: computes per-particle render data on the GPU. +//! +//! This shader transforms raw particle positions and dynamics into `ReadbackData` +//! suitable for CPU rendering, avoiding the need to transfer full `Position` and +//! `Dynamics` buffers back to the CPU. + +use crate::glamx::MatExt; +use crate::grid::grid::Grid; +use crate::solver::params::SimulationParams; +use crate::solver::particle::{Cdf, Kinematics, ParticleProperties, Position}; +use crate::{Matrix, PaddedMatrix, PaddingExt, Vector, abs, acos, cos, diag, sqrt}; +use glamx::*; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; + +const RENDER_MODE_VOLUME: u32 = 1; +const RENDER_MODE_VELOCITY: u32 = 2; +const RENDER_MODE_PHASE: u32 = 3; +const RENDER_MODE_CDF_NORMALS: u32 = 4; +const RENDER_MODE_CDF_DISTANCES: u32 = 5; +const RENDER_MODE_CDF_SIGNS: u32 = 6; + +/// Looks up a particle's base color in the group palette. +/// +/// Group ids wrap around the palette, so an id past its end still resolves. +#[inline] +fn group_color(group_colors: &[Vec4], num_groups: u32, group_id: u32) -> Vec4 { + let len = if num_groups == 0 { 1 } else { num_groups }; + *group_colors.at((group_id % len) as usize) +} + +/// Render configuration for the readback shader. +/// +/// Padded to 16 bytes so it satisfies the uniform-buffer minimum binding size +/// (the render kernel binds it as a uniform to stay within Firefox's 8 storage +/// buffers/stage limit; the readback kernel still binds it as a storage array). +#[derive(Clone, Copy, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct RenderConfig { + pub mode: u32, + /// Number of entries in the group-color palette. Group ids are taken modulo + /// this, so any id resolves to a color. + pub num_groups: u32, + // Scalar padding (not an array) so the 16-byte struct satisfies std140 + // uniform layout: array members would require 16-byte alignment. + pub _pad0: u32, + pub _pad1: u32, +} + +/// Per-particle data prepared on the GPU for CPU-side rendering. +/// +/// This struct is written by the GPU readback shader and read back to the CPU. +/// Uses explicit padding to ensure layout matches between host (with SIMD alignment +/// for Vec4) and GPU (scalar block layout). +#[cfg(feature = "dim2")] +#[derive(Clone, Copy, Default)] +#[cfg_attr( + not(target_arch_is_gpu), + derive(Debug, bytemuck::Pod, bytemuck::Zeroable) +)] +#[repr(C)] +pub struct ReadbackData { + pub color: Vec4, + pub deformation: Mat2, + pub position: Vec2, + // Explicit padding: Vec4 has 16-byte alignment on host (SIMD), so the struct + // size must be a multiple of 16. Without padding: 16+16+8 = 40, rounded to 48. + // We add 8 bytes of explicit padding to satisfy bytemuck::Pod (no implicit padding). + pub _pad: [f32; 2], +} + +/// Per-particle data prepared on the GPU for CPU-side rendering. +/// +/// Uses `PaddedMatrix` (Mat4 in 3D) instead of Mat3 to avoid SPIR-V storage buffer +/// alignment issues (Vec3 straddles 16-byte boundaries with std430 layout). +/// On the host side, use `deformation.remove_padding()` to get the Mat3. +#[cfg(feature = "dim3")] +#[derive(Clone, Copy, Default)] +#[cfg_attr( + not(target_arch_is_gpu), + derive(Debug, bytemuck::Pod, bytemuck::Zeroable) +)] +#[repr(C)] +pub struct ReadbackData { + pub color: Vec4, + pub deformation: PaddedMatrix, + pub position: Vec3, + // Vec4(16) + Mat4(64) + Vec3(12) = 92 bytes. Struct align is 16 (Vec4/Mat4). + // Padded to 96 (next multiple of 16). + pub _pad: f32, +} + +/// Compute the clamped and scaled deformation matrix for rendering. +#[cfg(feature = "dim2")] +#[inline] +fn compute_deformation(def_grad: PaddedMatrix, init_radius: f32) -> Mat2 { + let init_def = diag(Vector::splat(init_radius * 2.0)); + let clamped = Mat2::from_cols( + def_grad.x_axis.clamp(Vec2::splat(-4.0), Vec2::splat(4.0)), + def_grad.y_axis.clamp(Vec2::splat(-4.0), Vec2::splat(4.0)), + ); + init_def * clamped +} + +/// Compute the clamped and scaled deformation matrix for rendering. +/// Returns `PaddedMatrix` (Mat4) to avoid SPIR-V alignment issues in storage buffers. +#[cfg(feature = "dim3")] +#[inline] +fn compute_deformation(def_grad: PaddedMatrix, init_radius: f32) -> PaddedMatrix { + let init_def = diag(Vector::splat(init_radius * 2.0)); + let def3 = def_grad.remove_padding(); + let clamped = Mat3::from_cols( + def3.x_axis.clamp(Vec3::splat(-4.0), Vec3::splat(4.0)), + def3.y_axis.clamp(Vec3::splat(-4.0), Vec3::splat(4.0)), + def3.z_axis.clamp(Vec3::splat(-4.0), Vec3::splat(4.0)), + ); + PaddedMatrix::add_padding(init_def * clamped) +} + +/// Compute the color for a particle based on the render mode. +#[cfg(feature = "dim2")] +#[inline] +fn compute_color( + kin: &Kinematics, + cdf: &Cdf, + def_grad: &PaddedMatrix, + props: &ParticleProperties, + base_color: Vec4, + mode: u32, + cell_width: f32, + dt: f32, +) -> Vec4 { + if mode == RENDER_MODE_VELOCITY { + let vel = kin.velocity; + let c = Vec2::new(abs(vel.x), abs(vel.y)) * dt * 100.0 + Vec2::splat(0.2); + Vec4::new(c.x, c.y, base_color.z, base_color.w) + } else if mode == RENDER_MODE_VOLUME { + let sv = def_grad.svd().s; + let c = (Vec2::ONE - sv) / 0.005 + Vec2::splat(0.2); + Vec4::new(c.x, c.y, base_color.z, base_color.w) + } else if mode == RENDER_MODE_PHASE { + let phase = props.phase; + Vec4::new(0.0, 0.4 * phase, 0.4 * (1.0 - phase), base_color.w) + } else if mode == RENDER_MODE_CDF_NORMALS { + let normal = cdf.normal; + if normal == Vec2::ZERO { + Vec4::new(0.0, 0.0, 0.0, base_color.w) + } else { + let n = (normal + Vec2::ONE) * 0.5; + Vec4::new(n.x, n.y, 0.0, base_color.w) + } + } else if mode == RENDER_MODE_CDF_DISTANCES { + let d = cdf.signed_distance / (cell_width * 1.5); + if d > 0.0 { + Vec4::new(0.0, abs(d), 0.0, base_color.w) + } else { + Vec4::new(abs(d), 0.0, 0.0, base_color.w) + } + } else if mode == RENDER_MODE_CDF_SIGNS { + let d = cdf.affinity; + let a = (d.0 >> 16) & (d.0 & 0x0000ffff); + if d.0 == 0 { + Vec4::new(0.0, 0.0, 0.0, base_color.w) + } else if a == 0 { + Vec4::new(0.0, 1.0, 0.0, base_color.w) + } else { + Vec4::new(1.0, 0.0, 0.0, base_color.w) + } + } else { + // Default mode. + base_color + } +} + +/// Compute the color for a particle based on the render mode. +#[cfg(feature = "dim3")] +#[inline] +fn compute_color( + kin: &Kinematics, + cdf: &Cdf, + def_grad: &PaddedMatrix, + props: &ParticleProperties, + base_color: Vec4, + mode: u32, + cell_width: f32, + dt: f32, +) -> Vec4 { + let failed = kin.enabled == 0; + + let color = if mode == RENDER_MODE_VELOCITY { + let vel = kin.velocity; + let c = Vec3::new(abs(vel.x), abs(vel.y), abs(vel.z)) * dt * 100.0 + Vec3::splat(0.2); + Vec4::new(c.x, c.y, c.z, base_color.w) + } else if mode == RENDER_MODE_VOLUME { + let sv = def_grad.remove_padding().svd().s; + let c = (Vec3::ONE - sv) / 0.005 + Vec3::splat(0.2); + Vec4::new(c.x, c.y, c.z, base_color.w) + } else if mode == RENDER_MODE_PHASE { + let phase = props.phase; + Vec4::new(0.0, 0.4 * phase, 0.4 * (1.0 - phase), base_color.w) + } else if mode == RENDER_MODE_CDF_NORMALS { + let normal = cdf.normal; + if normal == Vec3::ZERO { + Vec4::new(0.0, 0.0, 0.0, base_color.w) + } else { + let n = (normal + Vec3::ONE) * 0.5; + Vec4::new(n.x, n.y, n.z, base_color.w) + } + } else if mode == RENDER_MODE_CDF_DISTANCES { + let d = cdf.signed_distance / (cell_width * 1.5); + if d > 0.0 { + Vec4::new(0.0, abs(d), 0.0, base_color.w) + } else { + Vec4::new(abs(d), 0.0, 0.0, base_color.w) + } + } else if mode == RENDER_MODE_CDF_SIGNS { + let d = cdf.affinity; + let a = (d.0 >> 16) & (d.0 & 0x0000ffff); + if d.0 == 0 { + Vec4::new(0.0, 0.0, 0.0, base_color.w) + } else if a == 0 { + Vec4::new(0.0, 1.0, 0.0, base_color.w) + } else { + Vec4::new(1.0, 0.0, 0.0, base_color.w) + } + } else { + // Default mode. + base_color + }; + + // Mark disabled (failed) particles red. + if failed { + Vec4::new(1.0, 0.0, 0.0, 1.0) + } else { + color + } +} + +/// GPU kernel: prepare per-particle readback data for rendering. +/// +/// Reads particle positions and dynamics, computes render color and scaled +/// deformation matrix, and writes the result to the `instances` buffer. +/// Dispatched with one thread per particle. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_prep_readback( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] instances: &mut [ReadbackData], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] particles_pos: &[Position], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] particles_kin: &[Kinematics], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] particles_def_grad: &[PaddedMatrix], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] + particles_props: &[ParticleProperties], + #[spirv(uniform, descriptor_set = 0, binding = 5)] grid: &Grid, + #[spirv(uniform, descriptor_set = 0, binding = 6)] params: &SimulationParams, + #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] config: &[RenderConfig], + #[spirv(uniform, descriptor_set = 0, binding = 8)] particles_len: &u32, + #[spirv(storage_buffer, descriptor_set = 0, binding = 9)] group_colors: &[Vec4], +) { + let particle_id = invocation_id.x; + + if particle_id >= *particles_len { + return; + } + + let pid = particle_id as usize; + let kin = particles_kin.at(pid); + let cdf = &kin.cdf; + let def_grad = particles_def_grad.at(pid); + let props = particles_props.at(pid); + let pos = particles_pos.at(pid); + let cfg = config.at(0); + let base_color = group_color(group_colors, cfg.num_groups, props.group_id); + let cell_width = grid.cell_width; + let mode = cfg.mode; + let dt = params.dt; + + let deformation = compute_deformation(*def_grad, props.init_radius); + let color = compute_color(kin, cdf, def_grad, props, base_color, mode, cell_width, dt); + + #[cfg(feature = "dim2")] + { + *instances.at_mut(pid) = ReadbackData { + color, + deformation, + position: pos.pt, + _pad: [0.0; 2], + }; + } + #[cfg(feature = "dim3")] + { + *instances.at_mut(pid) = ReadbackData { + color, + deformation, + position: pos.pt, + _pad: 0.0, + }; + } +} + +/// GPU kernel: write per-particle render data straight into a renderer's +/// instance buffers. +#[cfg(feature = "dim3")] +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mpm_prep_render( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] positions: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] deformations: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] colors: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] particles_pos: &[Position], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] particles_kin: &[Kinematics], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] particles_def_grad: &[PaddedMatrix], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] + particles_props: &[ParticleProperties], + #[spirv(uniform, descriptor_set = 0, binding = 7)] grid: &Grid, + #[spirv(uniform, descriptor_set = 0, binding = 8)] params: &SimulationParams, + #[spirv(uniform, descriptor_set = 0, binding = 9)] config: &RenderConfig, + #[spirv(uniform, descriptor_set = 0, binding = 10)] particles_len: &u32, + #[spirv(storage_buffer, descriptor_set = 0, binding = 11)] group_colors: &[Vec4], +) { + let particle_id = invocation_id.x; + if particle_id >= *particles_len { + return; + } + + let pid = particle_id as usize; + let kin = particles_kin.at(pid); + let cdf = &kin.cdf; + let def_grad = particles_def_grad.at(pid); + let props = particles_props.at(pid); + let pos = particles_pos.at(pid); + let base_color = group_color(group_colors, config.num_groups, props.group_id); + let cell_width = grid.cell_width; + let mode = config.mode; + let dt = params.dt; + + let deformation = compute_deformation(*def_grad, props.init_radius).remove_padding(); + let color = compute_color(kin, cdf, def_grad, props, base_color, mode, cell_width, dt); + + let pb = pid * 3; + *positions.at_mut(pb) = pos.pt.x; + *positions.at_mut(pb + 1) = pos.pt.y; + *positions.at_mut(pb + 2) = pos.pt.z; + + let db = pid * 9; + *deformations.at_mut(db) = deformation.x_axis.x; + *deformations.at_mut(db + 1) = deformation.x_axis.y; + *deformations.at_mut(db + 2) = deformation.x_axis.z; + *deformations.at_mut(db + 3) = deformation.y_axis.x; + *deformations.at_mut(db + 4) = deformation.y_axis.y; + *deformations.at_mut(db + 5) = deformation.y_axis.z; + *deformations.at_mut(db + 6) = deformation.z_axis.x; + *deformations.at_mut(db + 7) = deformation.z_axis.y; + *deformations.at_mut(db + 8) = deformation.z_axis.z; + + let cb = pid * 4; + *colors.at_mut(cb) = color.x; + *colors.at_mut(cb + 1) = color.y; + *colors.at_mut(cb + 2) = color.z; + *colors.at_mut(cb + 3) = color.w; +} + +/// GPU kernel: write per-particle render data straight into a renderer's +/// instance buffers (2D zero-readback path). +#[cfg(feature = "dim2")] +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mpm_prep_render( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] positions: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] deformations: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] colors: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] particles_pos: &[Position], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] particles_kin: &[Kinematics], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] particles_def_grad: &[PaddedMatrix], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] + particles_props: &[ParticleProperties], + #[spirv(uniform, descriptor_set = 0, binding = 7)] grid: &Grid, + #[spirv(uniform, descriptor_set = 0, binding = 8)] params: &SimulationParams, + #[spirv(uniform, descriptor_set = 0, binding = 9)] config: &RenderConfig, + #[spirv(uniform, descriptor_set = 0, binding = 10)] particles_len: &u32, + #[spirv(storage_buffer, descriptor_set = 0, binding = 11)] group_colors: &[Vec4], +) { + let particle_id = invocation_id.x; + if particle_id >= *particles_len { + return; + } + + let pid = particle_id as usize; + let kin = particles_kin.at(pid); + let cdf = &kin.cdf; + let def_grad = particles_def_grad.at(pid); + let props = particles_props.at(pid); + let pos = particles_pos.at(pid); + let base_color = group_color(group_colors, config.num_groups, props.group_id); + let cell_width = grid.cell_width; + let mode = config.mode; + let dt = params.dt; + + let deformation = compute_deformation(*def_grad, props.init_radius); + let color = compute_color(kin, cdf, def_grad, props, base_color, mode, cell_width, dt); + + let pb = pid * 2; + *positions.at_mut(pb) = pos.pt.x; + *positions.at_mut(pb + 1) = pos.pt.y; + + let db = pid * 4; + *deformations.at_mut(db) = deformation.x_axis.x; + *deformations.at_mut(db + 1) = deformation.x_axis.y; + *deformations.at_mut(db + 2) = deformation.y_axis.x; + *deformations.at_mut(db + 3) = deformation.y_axis.y; + + let cb = pid * 4; + *colors.at_mut(cb) = color.x; + *colors.at_mut(cb + 1) = color.y; + *colors.at_mut(cb + 2) = color.z; + *colors.at_mut(cb + 3) = color.w; +} + +/// GPU kernel: prepare per-rigid-particle readback data for rendering. +/// +/// Reads rigid particle world positions and writes `ReadbackData` with a fixed +/// deformation scale (no deformation gradient) and base color from a palette. +/// Dispatched with one thread per rigid particle. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_prep_readback_rigid( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] instances: &mut [ReadbackData], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] particles_pos: &[Position], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] base_colors: &[Vec4], + #[spirv(uniform, descriptor_set = 0, binding = 3)] grid: &Grid, + #[spirv(uniform, descriptor_set = 0, binding = 4)] particles_len: &u32, +) { + let particle_id = invocation_id.x; + + if particle_id >= *particles_len { + return; + } + + let pid = particle_id as usize; + let pos = particles_pos.at(pid); + let base_color = *base_colors.at(pid); + let cell_width = grid.cell_width; + let scale = cell_width * 0.4; + + #[cfg(feature = "dim2")] + { + let deformation = diag(Vector::splat(scale)); + *instances.at_mut(pid) = ReadbackData { + color: base_color, + deformation, + position: pos.pt, + _pad: [0.0; 2], + }; + } + #[cfg(feature = "dim3")] + { + let deformation = PaddedMatrix::add_padding(diag(Vector::splat(scale))); + *instances.at_mut(pid) = ReadbackData { + color: base_color, + deformation, + position: pos.pt, + _pad: 0.0, + }; + } +} diff --git a/src_mpm_shaders/solver/rigid_impulses.rs b/src_mpm_shaders/solver/rigid_impulses.rs new file mode 100644 index 00000000..50535acd --- /dev/null +++ b/src_mpm_shaders/solver/rigid_impulses.rs @@ -0,0 +1,185 @@ +//! Rigid body impulse accumulation and integration kernels for coupling MPM particles +//! with rigid bodies. + +use crate::grid::grid::Grid; +use crate::nexus_rbd_shaders::dynamics::{ + Impulse, LocalMassProperties, Velocity, WorldMassProperties, +}; +use crate::solver::p2g::IntegerImpulse; +use crate::solver::params::SimulationParams; +use crate::{AngVector, IVector, Pose, Vector, ang_length}; +use glamx::*; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; + +/// Scaling factor for float-to-integer impulse conversion. +pub const FLOAT_TO_INT_FACTOR: f32 = 1e5; + +/// Converts a float value to its integer-quantized representation. +#[inline] +pub fn flt2int(flt: f32) -> i32 { + (flt * FLOAT_TO_INT_FACTOR) as i32 +} + +/// Converts an integer-quantized value back to float. +#[inline] +pub fn int2flt(i: i32) -> f32 { + i as f32 / FLOAT_TO_INT_FACTOR +} + +impl IntegerImpulse { + /// Converts this integer-quantized impulse to a floating-point [`Impulse`]. + #[inline] + pub fn to_float(&self) -> Impulse { + #[cfg(feature = "dim2")] + { + Impulse::new( + Vec2::new(int2flt(self.linear_x), int2flt(self.linear_y)), + int2flt(self.angular), + ) + } + #[cfg(feature = "dim3")] + { + Impulse::new( + Vec3::new( + int2flt(self.linear_x), + int2flt(self.linear_y), + int2flt(self.linear_z), + ), + Vec3::new( + int2flt(self.angular_x), + int2flt(self.angular_y), + int2flt(self.angular_z), + ), + ) + } + } +} + +/// Updates rigid body velocities and poses by applying accumulated impulses, then resets +/// the impulse accumulator for the next substep. +/// +/// NOTE: numthreads(16) because we are currently limited to 16 bodies +/// due to the CPIC affinity bitmask size. +#[spirv_bindgen] +#[spirv(compute(threads(16)))] +pub fn gpu_rigid_impulses_update( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] sim_params: &SimulationParams, + #[spirv(uniform, descriptor_set = 0, binding = 1)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + local_mprops: &[LocalMassProperties], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] poses: &mut [Pose], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] vels: &mut [Velocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] mprops: &mut [WorldMassProperties], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] + incremental_impulses: &mut [IntegerImpulse], +) { + let id = invocation_id.x; + + if id < vels.len() as u32 { + let idx = id as usize; + let inc_impulse = incremental_impulses.at(idx).to_float(); + + // Reset the incremental impulse to zero for the next substep. + *incremental_impulses.at_mut(idx) = IntegerImpulse::default(); + + // Apply impulse and integrate. + let current_vel = vels.read(idx); + let current_mprops = mprops.read(idx); + let mut new_vel = current_vel.apply_impulse(¤t_mprops, &inc_impulse); + + // Cap the velocities to not move more than a fraction of a cell-width in a given substep. + let linvel_norm = new_vel.linear.length(); + let angvel_norm = ang_length(new_vel.angular); + let lin_limit = 0.1 * grid.cell_width / sim_params.dt; + let ang_limit = 1.0; // TODO: what's a good angular limit? + + let impulse_linear_len = inc_impulse.linear.length(); + let impulse_angular_len = ang_length(inc_impulse.angular); + + if impulse_linear_len != 0.0 || impulse_angular_len != 0.0 { + if linvel_norm > lin_limit { + new_vel.linear *= lin_limit / linvel_norm; + } + if angvel_norm > ang_limit { + new_vel.angular *= ang_limit / angvel_norm; + } + } + + let current_pose = poses.read(idx); + let local_mp = local_mprops.read(idx); + let new_pose = new_vel.integrate(¤t_pose, local_mp.com, sim_params.dt); + + // Apply gravity. + // Construct a mask: 1.0 where inv_mass != 0.0, 0.0 otherwise. + #[cfg(feature = "dim2")] + let mass_mask = Vec2::new( + (current_mprops.inv_mass.x != 0.0) as u32 as f32, + (current_mprops.inv_mass.y != 0.0) as u32 as f32, + ); + #[cfg(feature = "dim3")] + let mass_mask = Vec3::new( + (current_mprops.inv_mass.x != 0.0) as u32 as f32, + (current_mprops.inv_mass.y != 0.0) as u32 as f32, + (current_mprops.inv_mass.z != 0.0) as u32 as f32, + ); + new_vel.linear += sim_params.gravity * mass_mask * sim_params.dt; + + vels.write(idx, new_vel); + poses.write(idx, new_pose); + } +} + +/// Copies the MPM-integrated poses of the coupled bodies into the rigid-body +/// pipeline's body-pose buffer. +/// +/// MPM keeps its own copy of every coupled body and integrates it each substep +/// in [`gpu_rigid_impulses_update`]; that copy is the one the sand collides +/// with. The rigid-body pipeline treats those bodies as static (zero inverse +/// mass), so without this writeback the copy that rendering and the broad phase +/// read stays frozen at the insertion pose. +/// +/// `rbd_slots[i]` is the rigid-body slot mirroring coupled body `i`. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_writeback_body_poses( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] poses: &[Pose], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] rbd_slots: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] rbd_poses: &mut [Pose], +) { + let id = invocation_id.x; + + if id < rbd_slots.len() as u32 { + let idx = id as usize; + let slot = rbd_slots.read(idx) as usize; + rbd_poses.write(slot, poses.read(idx)); + } +} + +/// Updates world-space mass properties from local properties and current poses. +/// +/// Also writes the updated center of mass into the incremental impulse buffer +/// so that P2G can access it without an extra binding. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_update_world_mass_properties( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] poses: &[Pose], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + local_mprops: &[LocalMassProperties], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] mprops: &mut [WorldMassProperties], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] + incremental_impulses: &mut [IntegerImpulse], +) { + let id = invocation_id.x; + + if id < mprops.len() as u32 { + let idx = id as usize; + let local_mp = local_mprops.read(idx); + let new_mprops = local_mp.to_world(poses.at(idx)); + incremental_impulses.at_mut(idx).com = new_mprops.com; + mprops.write(idx, new_mprops); + } +} diff --git a/src_mpm_shaders/solver/rigid_particle_update.rs b/src_mpm_shaders/solver/rigid_particle_update.rs new file mode 100644 index 00000000..34ad4fe5 --- /dev/null +++ b/src_mpm_shaders/solver/rigid_particle_update.rs @@ -0,0 +1,54 @@ +//! Rigid particle update kernels: transforms sample/shape points from local to world space. + +use crate::solver::particle::{Position, RigidParticleIndices}; +use crate::{Pose, Vector}; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use nexus_rbd_shaders::PaddedVector; + +/// Transforms rigid body sample points from local space to world space. +/// +/// Each thread transforms one sample point by applying the pose of the collider +/// that owns the corresponding rigid particle. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_transform_sample_points( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + rigid_particle_indices: &[RigidParticleIndices], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] poses: &[Pose], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] local_pts: &[Position], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] world_pts: &mut [Position], +) { + let id = invocation_id.x; + + if id < local_pts.len() as u32 { + let collider_id = rigid_particle_indices.read(id as usize).collider; + let pose = poses.read(collider_id as usize); + let local_pt = local_pts.read(id as usize); + world_pts.write(id as usize, Position::new(pose * local_pt.pt)); + } +} + +/// Transforms rigid body shape (collider mesh) vertices from local space to world space. +/// +/// Each thread transforms one vertex by applying the pose of the collider +/// identified by the vertex-to-collider mapping. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_transform_shape_points( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] vertex_collider_ids: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] poses: &[Pose], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] local_pts: &[PaddedVector], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] world_pts: &mut [PaddedVector], +) { + let id = invocation_id.x; + + if id < local_pts.len() as u32 { + let collider_id = vertex_collider_ids.read(id as usize); + let pose = poses.read(collider_id as usize); + let local_pt = local_pts.read(id as usize); + world_pts.write(id as usize, (pose * local_pt.0).into()); + } +} diff --git a/src_mpm_shaders/solver/timestep_bound.rs b/src_mpm_shaders/solver/timestep_bound.rs new file mode 100644 index 00000000..7ebbfc97 --- /dev/null +++ b/src_mpm_shaders/solver/timestep_bound.rs @@ -0,0 +1,123 @@ +//! Timestep bound estimation kernels. +//! +//! Computes a CFL-based timestep bound across all particles. Each thread computes +//! a per-particle bound and atomically reduces it to find the global minimum. + +use crate::PaddingExt; +use crate::grid::grid::Grid; +use crate::models::default::{DefaultParticleModel, GpuParticleModel}; +use crate::solver::particle::{Kinematics, ParticleProperties}; +use crate::{DIM, Matrix, PaddedMatrix, sqrt}; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use khal_std::sync::atomic_min_u32; + +/// GPU-side timestep bound result. +/// +/// Uses an atomic unsigned integer to store the minimum timestep across all particles. +/// The float timestep is converted to an integer via a fixed-point scaling factor +/// so that atomic min operations can be used. +#[derive(Clone, Copy, Default, Debug)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct GpuTimestepBounds { + pub computed_max_dt_as_uint: u32, +} + +impl GpuTimestepBounds { + /// Conversion factor from seconds to integer representation. + pub const FLOAT_TO_INT: f32 = 1.0e12; + + /// Converts a timestep in seconds to its integer representation. + /// + /// Since `secs` is always positive, truncation via `as u32` is equivalent to floor. + #[inline] + pub fn secs_to_int(secs: f32) -> u32 { + (secs * Self::FLOAT_TO_INT) as u32 + } +} + +/// Resets the timestep bound to the maximum possible value. +#[spirv_bindgen] +#[spirv(compute(threads(1)))] +pub fn gpu_reset_timestep_bound( + #[spirv(global_invocation_id)] _invocation_id: khal_std::glamx::UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] result: &mut [GpuTimestepBounds], +) { + result.at_mut(0).computed_max_dt_as_uint = 0xFFFFFFFF; +} + +/// Estimates the CFL-based timestep bound across all particles. +/// +/// Each thread computes a per-particle timestep bound based on: +/// 1. Material model sound speed (model-specific). +/// 2. Particle velocity and APIC affine matrix contribution. +/// +/// The minimum across all particles is stored atomically in `result`. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_estimate_timestep_bound( + #[spirv(global_invocation_id)] invocation_id: khal_std::glamx::UVec3, + #[spirv(uniform, descriptor_set = 0, binding = 0)] grid: &Grid, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + particles_model: &[GpuParticleModel], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] particles_kin: &[Kinematics], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] particles_def_grad: &[PaddedMatrix], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] + particles_props: &[ParticleProperties], + #[spirv(uniform, descriptor_set = 0, binding = 5)] particles_len: &u32, + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] result: &mut [GpuTimestepBounds], +) { + let particle_id = invocation_id.x; + + if particle_id >= *particles_len { + return; + } + + let kin = particles_kin.read(particle_id as usize); + + if kin.enabled == 0 { + return; + } + + let def_grad = particles_def_grad.read(particle_id as usize); + let props = particles_props.read(particle_id as usize); + let cell_width = grid.cell_width; + + // Model-specific restrictions (usually based on sound speed, section 4.1). + let density0 = kin.mass / props.init_volume; + let def_grad = def_grad.remove_padding(); + let velocity = kin.velocity; + let affine = kin.affine.remove_padding(); + let mass = kin.mass; + + let mut dt = DefaultParticleModel::timestep_bound( + particles_model, + particle_id, + density0, + def_grad, + velocity, + cell_width, + ); + + // Velocity-based restrictions (section 4.2). + let norm_affine_squared = frobenius_norm_squared(affine); + + let d = (cell_width * cell_width) / 4.0; + let norm_b = d * sqrt(norm_affine_squared) / mass; + let apic_v = norm_b * 6.0 * sqrt(DIM as f32) / cell_width; + let v = velocity.length() + apic_v; + dt = dt.min(cell_width / v); + + let candidate = GpuTimestepBounds::secs_to_int(dt); + atomic_min_u32(&mut result.at_mut(0).computed_max_dt_as_uint, candidate); +} + +/// Computes the squared Frobenius norm of a matrix (sum of squares of all elements). +#[inline] +fn frobenius_norm_squared(m: Matrix) -> f32 { + #[cfg(feature = "dim2")] + return m.x_axis.length_squared() + m.y_axis.length_squared(); + #[cfg(feature = "dim3")] + return m.x_axis.length_squared() + m.y_axis.length_squared() + m.z_axis.length_squared(); +} From bc2100898e27abf56a0a9dc6e6d228a980300a51 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 14 Aug 2026 23:47:14 +0200 Subject: [PATCH 2/7] feat: couple MPM with rigid bodies through NexusState --- crates/nexus2d/Cargo.toml | 12 +- crates/nexus3d/Cargo.toml | 12 +- src/lib.rs | 5 + src/pipeline.rs | 30 ++ src/state.rs | 497 +++++++++++++++++++++++++++++----- src_rbd/dynamics/body.rs | 420 ++++++++++++++++++++++++++++ src_rbd/dynamics/mod.rs | 2 + src_rbd/pipeline/rbd_state.rs | 9 + src_rbd/utils/prefix_sum.rs | 5 +- 9 files changed, 905 insertions(+), 87 deletions(-) create mode 100644 src_rbd/dynamics/body.rs diff --git a/crates/nexus2d/Cargo.toml b/crates/nexus2d/Cargo.toml index 199cfbaa..f8b2c299 100644 --- a/crates/nexus2d/Cargo.toml +++ b/crates/nexus2d/Cargo.toml @@ -23,17 +23,19 @@ default = ["dim2", "f32", "webgpu"] dim2 = [] f32 = [] f64 = [] -webgpu = ["nexus_rbd2d?/webgpu"] -metal = ["nexus_rbd2d?/metal"] -cpu = ["nexus_rbd2d?/cpu"] -cpu-parallel = ["cpu", "nexus_rbd2d?/cpu-parallel"] -cuda = ["nexus_rbd2d?/cuda"] +webgpu = ["nexus_rbd2d?/webgpu", "nexus_mpm2d?/webgpu"] +metal = ["nexus_rbd2d?/metal", "nexus_mpm2d?/metal"] +cpu = ["nexus_rbd2d?/cpu", "nexus_mpm2d?/cpu"] +cpu-parallel = ["cpu", "nexus_rbd2d?/cpu-parallel", "nexus_mpm2d?/cpu-parallel"] +cuda = ["nexus_rbd2d?/cuda", "nexus_mpm2d?/cuda"] rbd = ["dep:nexus_rbd2d"] +mpm = ["dep:nexus_mpm2d"] [dependencies] khal = { workspace = true } nexus_rbd2d = { workspace = true, optional = true, features = ["default"] } +nexus_mpm2d = { workspace = true, optional = true } bitflags = { workspace = true } web-time = { workspace = true } diff --git a/crates/nexus3d/Cargo.toml b/crates/nexus3d/Cargo.toml index 6262afb6..ef496c2b 100644 --- a/crates/nexus3d/Cargo.toml +++ b/crates/nexus3d/Cargo.toml @@ -23,17 +23,19 @@ default = ["dim3", "f32", "webgpu"] dim3 = [] f32 = [] f64 = [] -webgpu = ["nexus_rbd3d?/webgpu"] -metal = ["nexus_rbd3d?/metal"] -cpu = ["nexus_rbd3d?/cpu"] -cpu-parallel = ["cpu", "nexus_rbd3d?/cpu-parallel"] -cuda = ["nexus_rbd3d?/cuda"] +webgpu = ["nexus_rbd3d?/webgpu", "nexus_mpm3d?/webgpu"] +metal = ["nexus_rbd3d?/metal", "nexus_mpm3d?/metal"] +cpu = ["nexus_rbd3d?/cpu", "nexus_mpm3d?/cpu"] +cpu-parallel = ["cpu", "nexus_rbd3d?/cpu-parallel", "nexus_mpm3d?/cpu-parallel"] +cuda = ["nexus_rbd3d?/cuda", "nexus_mpm3d?/cuda"] rbd = ["dep:nexus_rbd3d"] +mpm = ["dep:nexus_mpm3d"] [dependencies] bitflags = { workspace = true } web-time = { workspace = true } khal = { workspace = true } nexus_rbd3d = { workspace = true, optional = true, features = ["default"] } +nexus_mpm3d = { workspace = true, optional = true } diff --git a/src/lib.rs b/src/lib.rs index ea5f35cb..164fa05a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -11,6 +11,11 @@ pub use nexus_rbd2d as rbd; #[cfg(all(feature = "dim3", feature = "rbd"))] pub use nexus_rbd3d as rbd; +#[cfg(all(feature = "dim2", feature = "mpm"))] +pub use nexus_mpm2d as mpm; +#[cfg(all(feature = "dim3", feature = "mpm"))] +pub use nexus_mpm3d as mpm; + #[cfg(feature = "rbd")] pub use rbd::{parry, rapier}; diff --git a/src/pipeline.rs b/src/pipeline.rs index 8a9e42ca..a330e721 100644 --- a/src/pipeline.rs +++ b/src/pipeline.rs @@ -1,3 +1,4 @@ +use crate::mpm::pipeline::MpmPipeline; use crate::rbd::pipeline::RbdPipeline; use crate::state::NexusState; use khal::backend::{GpuBackend, GpuBackendError, GpuTimestamps}; @@ -7,12 +8,14 @@ bitflags::bitflags! { #[derive(Copy, Clone, PartialEq, Eq, Debug)] pub struct NexusPipelineMask: u8 { const RBD = 1 << 0; + const MPM = 1 << 1; } } #[derive(Default)] pub struct NexusPipeline { pub rbd_pipeline: Option, + pub mpm_pipeline: Option, } impl NexusPipeline { @@ -24,6 +27,9 @@ impl NexusPipeline { if pipelines.contains(NexusPipelineMask::RBD) && self.rbd_pipeline.is_none() { self.rbd_pipeline = Some(RbdPipeline::new(backend)?); } + if pipelines.contains(NexusPipelineMask::MPM) && self.mpm_pipeline.is_none() { + self.mpm_pipeline = Some(MpmPipeline::new(backend)?); + } Ok(()) } @@ -66,6 +72,30 @@ impl NexusPipeline { pipeline.auto_resize_buffers(backend, rbd)?; } + // MPM pipeline + if let Some(mpm) = state.mpm.as_mut() { + self.preload_pipelines(backend, NexusPipelineMask::MPM)?; + let pipeline = self.mpm_pipeline.as_mut().unwrap_or_else(|| unreachable!()); + + // MPM needs many small substeps per visible frame for stability. + // Upload the per-substep dt once, then run the substep loop. + let substeps = state.mpm_substeps.max(1); + let _ = mpm.write_substep_params(backend, substeps); + for _ in 0..substeps { + let _ = pipeline.step(backend, mpm, timestamps.as_deref_mut()); + } + } + + // MPM owns the pose of every body it is coupled to: it integrates its + // own copy each substep while the rigid-body pipeline treats those + // bodies as static. Push that copy back so rendering and the next + // step's broad phase see a boundary that actually moved. + // FIXME: the RBD pipeline should remain in charge of moving the bodies. + if let (Some(rbd), Some(mpm)) = (state.rbd.as_mut(), state.mpm.as_ref()) { + let pipeline = self.mpm_pipeline.as_ref().unwrap_or_else(|| unreachable!()); + pipeline.writeback_body_poses(backend, mpm, rbd.body_poses_mut())?; + } + state.run_stats.encoding_time = t0.elapsed(); // If we recorded this frame, kick off the non-blocking readback of the diff --git a/src/state.rs b/src/state.rs index 3e24bcc3..76ab81db 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,9 +1,14 @@ -use crate::rapier::data::{Coarena, Index}; +use crate::mpm::pipeline::{MpmCapacities, MpmState}; +use crate::mpm::solver::{BoundaryCondition, Particle, SimulationParams}; +use crate::rapier::data::{Arena, Coarena, Index}; use crate::rapier::prelude::{ Collider, ColliderHandle, GenericJoint, ImpulseJointHandle, MultibodyJointHandle, PhysicsWorld, RigidBody, RigidBodyHandle, }; -use crate::rbd::dynamics::RbdSimParams; +use crate::rbd::dynamics::{ + RbdSimParams, + body::{BodyCoupling, RapierBodyCouplingEntry}, +}; use crate::rbd::pipeline::{RbdCapacities, RbdResizePolicy, RbdState, RunStats}; use khal::backend::{GpuBackend, GpuBackendError}; @@ -11,15 +16,29 @@ use khal::backend::{GpuBackend, GpuBackendError}; #[derive(Copy, Clone, PartialEq, Eq, Debug, Hash)] pub struct NexusRbdHandle(Index); -/// Initial capacities used when allocating the GPU-resident physics states. +/// Handle referencing a *chunk* of MPM particles managed by a [`NexusState`]. /// -/// Groups the per-subsystem capacities, each owned by its own crate -/// ([`RbdCapacities`] in nexus-rbd) and forwarded to that subsystem when its -/// scene is first created. +/// Particles are addressed by chunk rather than individually (a per-particle +/// handle map would be prohibitive at MPM scale). A chunk is mutable: particles +/// can be appended to it ([`NexusState::extend_chunk`]) or removed from it +/// ([`NexusState::remove_particles_from_chunk`] / [`NexusState::remove_chunk`]). +#[derive(Copy, Clone, PartialEq, Eq, Debug, Hash)] +pub struct NexusParticleChunk(Index); + +#[derive(Copy, Clone, PartialEq, Debug)] +pub enum RbdCoupling { + None, + MpmOneWay(BoundaryCondition), + MpmTwoWay(BoundaryCondition), +} + +/// Initial capacities used when allocating the GPU-resident physics states. #[derive(Copy, Clone, Debug, Default)] pub struct NexusCapacities { - /// Rigid-body subsystem capacities. + /// Rigid-body solver capacities. pub rbd: RbdCapacities, + /// MPM solver capacities. + pub mpm: MpmCapacities, } impl NexusCapacities { @@ -38,20 +57,34 @@ impl NexusCapacities { self } + pub fn mpm_grid_size(mut self, num_chunks: u32) -> Self { + self.mpm.grid_size = num_chunks; + self + } + pub fn rbd_resize_policy(mut self, resize_policy: RbdResizePolicy) -> Self { self.rbd.collisions_resize_policy = resize_policy; self } + + pub fn mpm_particles(mut self, capacity: u32) -> Self { + self.mpm.particles_capacity = capacity; + self + } } #[derive(Copy, Clone, Debug)] pub struct GpuRigidBodyRef { + pub coupling: RbdCoupling, pub gpu_id: u32, } impl Default for GpuRigidBodyRef { fn default() -> Self { - Self { gpu_id: u32::MAX } + Self { + coupling: RbdCoupling::None, + gpu_id: u32::MAX, + } } } @@ -67,17 +100,21 @@ pub struct NexusCounts { pub multibody_dofs: usize, pub collision_pairs: usize, pub collision_pairs_capacity: usize, + pub particles: usize, } -/// High-level, GPU-resident state of a physics simulation. +/// High-level, GPU-resident state of a multiphysics simulation. /// -/// The rigid-body sub-state is lazily allocated the first time content is -/// added. The `rbd2gpu` maps translate the stable public handles into the -/// (unstable) GPU buffer slots, which shift around as bodies are inserted and -/// removed. +/// Each sub-state (`rbd`/`mpm`) is lazily allocated the first time content +/// of the corresponding kind is added. The `*2gpu` maps translate the stable +/// public handles into the (unstable) GPU buffer slots, which shift around as +/// bodies/particles are inserted and removed. pub struct NexusState { /// Rigid-body sub-state, allocated on the first [`Self::add_rigid_bodies`]. pub rbd: Option, + /// MPM sub-state, allocated on the first [`Self::add_particles`] (or the + /// first coupled rigid-body insertion). + pub mpm: Option, pub run_stats: RunStats, @@ -85,6 +122,27 @@ pub struct NexusState { /// (batch). pub rbd2gpu: Vec>, + /// Live particle count per MPM chunk (the arena key is the public + /// [`NexusParticleChunk`] handle). + mpm_chunks: Arena, + /// Owning chunk for each GPU particle slot, kept in sync under the + /// swap-removal performed by [`Self::remove_chunk`] / + /// [`Self::remove_particles_from_chunk`]. + slot2chunk: Vec, + /// MPM simulation params / grid cell width requested before the MPM + /// sub-state is lazily created. + mpm_params: Option, + mpm_cell_width: f32, + /// Number of MPM substeps run per [`NexusPipeline::simulate`](crate::pipeline::NexusPipeline::simulate) call. + pub mpm_substeps: u32, + /// Desired CPIC rigid-coupling flag, kept here so it survives until the MPM + /// sub-state is lazily created (and is what [`Self::mpm_use_cpic`] reports + /// meanwhile). + mpm_use_cpic: bool, + /// Set when particles or MPM-coupled bodies change; consumed by + /// [`Self::finalize`] to rebuild the MPM↔rapier coupling. + mpm_dirty: bool, + // Initial capacities used to allocate the states lazily. capacities: NexusCapacities, @@ -121,6 +179,7 @@ impl NexusState { pub fn new(capacities: NexusCapacities) -> Self { Self { rbd: None, + mpm: None, run_stats: RunStats::default(), rbd_envs: vec![PhysicsWorld::default()], rbd_sim_params: vec![RbdSimParams::tgs_soft()], @@ -128,17 +187,105 @@ impl NexusState { rbd_steps_per_frame: 1, rbd_reserve_per_env: 0, rbd2gpu: vec![Coarena::new()], + mpm_chunks: Arena::new(), + slot2chunk: Vec::new(), + mpm_params: None, + mpm_cell_width: 1.0, + mpm_substeps: 20, + mpm_use_cpic: true, + mpm_dirty: false, capacities, } } /// Reserves additional capacity in the handle maps to avoid reallocations - /// when a known number of bodies is about to be inserted. + /// when a known number of bodies/particles is about to be inserted. pub fn reserve(&mut self, additional: NexusCapacities) { for env in &mut self.rbd2gpu { env.reserve(additional.rbd.body_capacity as usize); } - // TODO: resize the GPU buffers too. + // TODO: reserve the MPM handle maps and resize the GPU buffers too. + } + + /// Sets the MPM simulation parameters (gravity, timestep) and grid cell + /// width. Call before the first [`Self::add_particles`]; the values are + /// applied when the MPM sub-state is created. If MPM already exists they are + /// applied immediately (the grid is reset, so prefer calling this first). + pub fn set_mpm_params( + &mut self, + backend: &GpuBackend, + params: SimulationParams, + cell_width: f32, + ) -> Result<(), GpuBackendError> { + self.mpm_params = Some(params); + self.mpm_cell_width = cell_width; + if let Some(mpm) = self.mpm.as_mut() { + mpm.set_cell_width(backend, cell_width, self.capacities.mpm.grid_size)?; + mpm.set_simulation_params(backend, params)?; + } + Ok(()) + } + + /// Sets the number of MPM substeps run per [`NexusPipeline::simulate`](crate::pipeline::NexusPipeline::simulate) call (default + /// 20). More substeps → smaller timestep → more stable but slower. + pub fn set_mpm_substeps(&mut self, substeps: u32) { + self.mpm_substeps = substeps.max(1); + } + + /// Number of MPM substeps run per [`NexusPipeline::simulate`](crate::pipeline::NexusPipeline::simulate) call. + pub fn mpm_substeps(&self) -> u32 { + self.mpm_substeps + } + + /// Enables/disables CPIC (compatible particle-in-cell) rigid coupling. The + /// preference is stored so it survives until MPM is lazily allocated. Not + /// overwritten by [`Self::finalize`] unless the coupling set changes. + pub fn set_mpm_use_cpic(&mut self, enabled: bool) { + self.mpm_use_cpic = enabled; + if let Some(mpm) = self.mpm.as_mut() { + mpm.use_cpic = enabled; + } + } + + /// Whether CPIC rigid coupling is enabled. Falls back to the stored + /// preference before MPM is lazily allocated. + pub fn mpm_use_cpic(&self) -> bool { + self.mpm + .as_ref() + .map(|m| m.use_cpic) + .unwrap_or(self.mpm_use_cpic) + } + + /// Whether this state uses the MPM solver. True once MPM has been configured + /// via [`Self::set_mpm_params`], even before the sub-state is lazily + /// allocated on the first [`Self::add_particles`], so a particle emitter + /// that starts empty still reports its MPM usage. + pub fn has_mpm(&self) -> bool { + self.mpm.is_some() || self.mpm_params.is_some() + } + + /// Sets the MPM gravity vector. Applied on the next [`NexusPipeline::simulate`](crate::pipeline::NexusPipeline::simulate) (the + /// per-substep params are re-uploaded each frame), so this is cheap. + pub fn set_mpm_gravity(&mut self, gravity: crate::rbd::math::Vector) { + // Keep the stored params authoritative so the gravity survives until MPM + // is lazily allocated (and is what `mpm_gravity` reports meanwhile). + if let Some(params) = self.mpm_params.as_mut() { + params.gravity = gravity; + } + if let Some(mpm) = self.mpm.as_mut() { + mpm.gravity = gravity; + } + } + + /// Current MPM gravity vector. Falls back to the gravity configured via + /// [`Self::set_mpm_params`] before MPM is lazily allocated, and only to zero + /// if no params were ever set. + pub fn mpm_gravity(&self) -> crate::rbd::math::Vector { + self.mpm + .as_ref() + .map(|m| m.gravity) + .or_else(|| self.mpm_params.map(|p| p.gravity)) + .unwrap_or(crate::rbd::math::Vector::ZERO) } /// Sets the rigid-body gravity vector, e.g. `[0.0, 0.0, -9.81]` for a Z-up @@ -166,9 +313,9 @@ impl NexusState { self.rbd_steps_per_frame } - /// Current entity counts (rigid bodies, colliders, joints, multibody DOFs) - /// for display in the UI. Rigid-body counts are summed across all - /// environments. + /// Current entity counts (rigid bodies, colliders, joints, multibody DOFs, + /// particles) for display in the UI. Rigid-body + /// counts are summed across all environments. pub fn counts(&self) -> NexusCounts { let mut c = NexusCounts { num_environments: self.rbd_envs.len(), @@ -187,9 +334,29 @@ impl NexusState { c.collision_pairs = rbd.collision_pairs_len() as usize; c.collision_pairs_capacity = rbd.collision_pairs_capacity() as usize; } + if let Some(mpm) = self.mpm.as_ref() { + c.particles = mpm.particles.len(); + } c } + /// Returns a mutable reference to the MPM sub-state, allocating an empty one + /// (sized from the stored capacities, configured from [`Self::set_mpm_params`]) + /// if it doesn’t exist yet. + fn mpm_or_insert(&mut self, backend: &GpuBackend) -> Result<&mut MpmState, GpuBackendError> { + if self.mpm.is_none() { + let grid_capacity = self.capacities.mpm.grid_size; + let mut mpm = MpmState::empty(backend, &self.capacities.mpm)?; + mpm.set_cell_width(backend, self.mpm_cell_width, grid_capacity)?; + if let Some(params) = self.mpm_params { + mpm.set_simulation_params(backend, params)?; + } + mpm.use_cpic = self.mpm_use_cpic; + self.mpm = Some(mpm); + } + Ok(self.mpm.as_mut().unwrap()) + } + /// Adds a new (empty) simulation environment (batch) and returns its index. /// /// Environment 0 always exists; batched demos call this once per extra @@ -243,13 +410,18 @@ impl NexusState { /// state. Use this to run rapier-side helpers whose output you then push /// through a runtime setter — e.g. driving an MJCF actuator model and /// forwarding the resulting motors with - /// `GpuMultibodySet::set_motors` (3D only, hence no intra-doc link here). + /// `GpuMultibodySet::set_motors`. pub fn rbd_world_mut_untracked(&mut self, env: usize) -> &mut PhysicsWorld { &mut self.rbd_envs[env] } - pub fn insert_rigid_body(&mut self, body: RigidBody, collider: Collider) -> RigidBodyHandle { - self.insert_rigid_body_in(0, body, collider) + pub fn insert_rigid_body( + &mut self, + body: RigidBody, + collider: Collider, + coupling: RbdCoupling, + ) -> RigidBodyHandle { + self.insert_rigid_body_in(0, body, collider, coupling) } /// Inserts a body + collider into environment `env`. @@ -258,10 +430,22 @@ impl NexusState { env: usize, body: RigidBody, collider: Collider, + coupling: RbdCoupling, ) -> RigidBodyHandle { let (handle, _) = self.rbd_envs[env].insert(body, collider); - self.rbd2gpu[env].insert(handle.0, GpuRigidBodyRef { gpu_id: u32::MAX }); + self.rbd2gpu[env].insert( + handle.0, + GpuRigidBodyRef { + coupling, + gpu_id: u32::MAX, + }, + ); self.rbd_dirty = true; + // MPM-coupled boundary colliders live only in environment 0 and feed the + // MPM coupling rebuild in `finalize`. + if env == 0 && coupling != RbdCoupling::None { + self.mpm_dirty = true; + } handle } @@ -289,8 +473,9 @@ impl NexusState { backend: &GpuBackend, body: RigidBody, collider: Collider, + coupling: RbdCoupling, ) -> Result { - let handles = self.add_rigid_bodies(backend, [(body, collider)])?; + let handles = self.add_rigid_bodies(backend, [(body, collider, coupling)])?; Ok(handles[0]) } @@ -307,15 +492,17 @@ impl NexusState { pub fn add_rigid_bodies( &mut self, backend: &GpuBackend, - bodies: impl IntoIterator, + bodies: impl IntoIterator, ) -> Result, GpuBackendError> { // Keep copies for the GPU append before the rapier world consumes them. let mut gpu_pairs: Vec<(RigidBody, Collider)> = Vec::new(); let mut handles: Vec = Vec::new(); - for (body, collider) in bodies { + let mut couplings: Vec = Vec::new(); + for (body, collider, coupling) in bodies { gpu_pairs.push((body.clone(), collider.clone())); let (handle, _) = self.rbd_envs[0].insert(body, collider); handles.push(handle); + couplings.push(coupling); } if handles.is_empty() { return Ok(handles); @@ -328,10 +515,11 @@ impl NexusState { { let range = rbd.append_bodies(backend, &gpu_pairs)?; // Single environment: the per-batch local slot is the gpu_id. - for (i, &handle) in handles.iter().enumerate() { + for (i, (&handle, &coupling)) in handles.iter().zip(&couplings).enumerate() { self.rbd2gpu[0].insert( handle.0, GpuRigidBodyRef { + coupling, gpu_id: range.start + i as u32, }, ); @@ -344,26 +532,46 @@ impl NexusState { if !appended { // No GPU state yet, or not enough room for the whole batch: fall back // to a full rebuild on the next `finalize`. - for &handle in handles.iter() { - self.rbd2gpu[0].insert(handle.0, GpuRigidBodyRef { gpu_id: u32::MAX }); + for (&handle, &coupling) in handles.iter().zip(&couplings) { + self.rbd2gpu[0].insert( + handle.0, + GpuRigidBodyRef { + coupling, + gpu_id: u32::MAX, + }, + ); } self.rbd_dirty = true; } + if couplings.iter().any(|c| *c != RbdCoupling::None) { + self.mpm_dirty = true; + } Ok(handles) } /// Inserts a rigid-body without any attached collider (e.g. a joint anchor). - pub fn insert_body(&mut self, body: RigidBody) -> RigidBodyHandle { - self.insert_body_in(0, body) + pub fn insert_body(&mut self, body: RigidBody, coupling: RbdCoupling) -> RigidBodyHandle { + self.insert_body_in(0, body, coupling) } // TODO: remove this. Inserting a collider should insert into all envs. // (though we should also have a variant that allows specifying different // shapes per env). /// Inserts a collider-less rigid-body into environment `env`. - pub fn insert_body_in(&mut self, env: usize, body: RigidBody) -> RigidBodyHandle { + pub fn insert_body_in( + &mut self, + env: usize, + body: RigidBody, + coupling: RbdCoupling, + ) -> RigidBodyHandle { let handle = self.rbd_envs[env].insert_body(body); - self.rbd2gpu[env].insert(handle.0, GpuRigidBodyRef { gpu_id: u32::MAX }); + self.rbd2gpu[env].insert( + handle.0, + GpuRigidBodyRef { + coupling, + gpu_id: u32::MAX, + }, + ); self.rbd_dirty = true; handle } @@ -453,7 +661,120 @@ impl NexusState { Ok(()) } + /// Appends a new chunk of MPM particles (`O(added)`) and returns its handle. + pub fn add_particles( + &mut self, + backend: &GpuBackend, + particles: Vec, + ) -> Result { + let n = particles.len(); + let chunk = self.mpm_chunks.insert(n); + { + let mpm = self.mpm_or_insert(backend)?; + mpm.particles.append(backend, &particles)?; + } + self.slot2chunk.extend(std::iter::repeat_n(chunk, n)); + self.mpm_dirty = true; + Ok(NexusParticleChunk(chunk)) + } + + /// Appends more particles to an existing chunk (`O(added)`). + pub fn extend_chunk( + &mut self, + backend: &GpuBackend, + chunk: NexusParticleChunk, + particles: Vec, + ) -> Result<(), GpuBackendError> { + let n = particles.len(); + { + let mpm = self.mpm_or_insert(backend)?; + mpm.particles.append(backend, &particles)?; + } + self.slot2chunk.extend(std::iter::repeat_n(chunk.0, n)); + if let Some(c) = self.mpm_chunks.get_mut(chunk.0) { + *c += n; + } + self.mpm_dirty = true; + Ok(()) + } + + /// MPM background-grid cell width. + pub fn mpm_cell_width(&self) -> f32 { + self.mpm_cell_width + } + + /// Removes every particle of a chunk (`O(removed)`) and drops the handle. + pub fn remove_chunk( + &mut self, + backend: &GpuBackend, + chunk: NexusParticleChunk, + ) -> Result<(), GpuBackendError> { + let slots: Vec = self + .slot2chunk + .iter() + .enumerate() + .filter(|(_, c)| **c == chunk.0) + .map(|(i, _)| i as u32) + .collect(); + self.swap_remove_particle_slots(backend, &slots)?; + self.mpm_chunks.remove(chunk.0); + Ok(()) + } + + /// Removes up to `count` particles from a chunk (`O(removed)`), returning the + /// number actually removed. The chunk itself is kept (even if emptied). + pub fn remove_particles_from_chunk( + &mut self, + backend: &GpuBackend, + chunk: NexusParticleChunk, + count: usize, + ) -> Result { + let mut slots: Vec = self + .slot2chunk + .iter() + .enumerate() + .filter(|(_, c)| **c == chunk.0) + .map(|(i, _)| i as u32) + .collect(); + // Remove the highest GPU slots first, which keeps the swap-removal cheap. + slots.sort_unstable_by(|a, b| b.cmp(a)); + slots.truncate(count); + let removed = slots.len(); + self.swap_remove_particle_slots(backend, &slots)?; + if let Some(c) = self.mpm_chunks.get_mut(chunk.0) { + *c = c.saturating_sub(removed); + } + Ok(removed) + } + + /// Swap-removes the given GPU particle slots and patches `slot2chunk` to + /// follow the relocations the GPU performed. + fn swap_remove_particle_slots( + &mut self, + backend: &GpuBackend, + slots: &[u32], + ) -> Result<(), GpuBackendError> { + if slots.is_empty() { + return Ok(()); + } + let remaps = { + let Some(mpm) = self.mpm.as_mut() else { + return Ok(()); + }; + mpm.particles.swap_remove(backend, slots)? + }; + // Each `(from, to)`: the tail particle at `from` was moved down to the + // freed slot `to`, so its chunk ownership moves with it. + for (from, to) in remaps { + self.slot2chunk[to as usize] = self.slot2chunk[from as usize]; + } + let new_len = self.mpm.as_ref().unwrap().particles.len(); + self.slot2chunk.truncate(new_len); + Ok(()) + } + pub async fn finalize(&mut self, backend: &GpuBackend) -> Result<(), GpuBackendError> { + let rbd_was_dirty = self.rbd_dirty; if self.rbd_dirty { // Finalize each body's mass properties so additional (``) // mass combined with its colliders is reflected in `local_mprops`. @@ -523,14 +844,12 @@ impl NexusState { RbdState::from_rapier(backend, &environments, self.capacities.rbd) }; - // Rebuild the per-environment handle → GPU-slot maps. A handle's - // `gpu_id` is its BODY slot (which indexes the body-keyed buffers - // such as `body_poses`), NOT a collider slot — a body may own - // several colliders. Body slots are assigned in the same order - // `from_rapier` uses: the first time each parent body is seen while - // iterating colliders (a parentless collider consumes a synthetic - // body slot, matching `from_rapier`). Bodies are laid out env-major - // with stride `num_colliders_per_batch`. + // Rebuild the per-environment handle to GPU-slot maps. A handle's + // `gpu_id` is its *body* slot, not a collider slot, since a body may + // own several colliders. Body slots are assigned in the order + // `from_rapier` uses (the first time each parent body is seen while + // iterating colliders) and are laid out env-major with stride + // `num_colliders_per_batch`. let stride = rbd_state.num_colliders_per_batch(); for (env_idx, world) in self.rbd_envs.iter().enumerate() { let mut body_slot: std::collections::HashMap<_, u32> = @@ -552,9 +871,14 @@ impl NexusState { next_slot += 1; s }); + let coupling = self.rbd2gpu[env_idx] + .get(body_handle.0) + .map(|r| r.coupling) + .unwrap_or(RbdCoupling::None); self.rbd2gpu[env_idx].insert( body_handle.0, GpuRigidBodyRef { + coupling, gpu_id: env_idx as u32 * stride + slot, }, ); @@ -573,9 +897,14 @@ impl NexusState { let slot = next_slot; next_slot += 1; body_slot.insert(body_handle, slot); + let coupling = self.rbd2gpu[env_idx] + .get(body_handle.0) + .map(|r| r.coupling) + .unwrap_or(RbdCoupling::None); self.rbd2gpu[env_idx].insert( body_handle.0, GpuRigidBodyRef { + coupling, gpu_id: env_idx as u32 * stride + slot, }, ); @@ -585,41 +914,59 @@ impl NexusState { self.rbd = Some(rbd_state); self.rbd_dirty = false; } + + // MPM/rapier coupling. Boundary colliders are inserted into environment 0 + // as rigid bodies tagged `RbdCoupling::Mpm*`; rebuild the coupling + // (sampled rigid particles, uploaded body set) whenever those bodies or + // the particle set changed. + if (rbd_was_dirty || self.mpm_dirty) && self.mpm.is_some() { + let world = &self.rbd_envs[0]; + let mut coupling = Vec::new(); + let mut materials = Vec::new(); + // Rigid-body slot mirroring each coupling entry, so the MPM-owned + // poses can be written back to the buffer rendering reads. + let mut rbd_body_slots = Vec::new(); + for (collider_handle, collider) in world.colliders.iter() { + let Some(body_handle) = collider.parent() else { + continue; + }; + let Some(gpu_ref) = self.rbd2gpu[0].get(body_handle.0) else { + continue; + }; + + let (boundary_condition, mode) = match gpu_ref.coupling { + RbdCoupling::None => continue, + RbdCoupling::MpmOneWay(boundary_condition) => { + (boundary_condition, BodyCoupling::OneWay) + } + RbdCoupling::MpmTwoWay(boundary_condition) => { + (boundary_condition, BodyCoupling::TwoWays) + } + }; + + coupling.push(RapierBodyCouplingEntry { + body: body_handle, + collider: collider_handle, + mode, + }); + materials.push(boundary_condition); + rbd_body_slots.push(gpu_ref.gpu_id); + } + if !coupling.is_empty() { + let cell_width = self.mpm_cell_width; + let mpm = self.mpm.as_mut().unwrap_or_else(|| unreachable!()); + mpm.set_coupling( + backend, + &world.bodies, + &world.colliders, + coupling, + &materials, + &rbd_body_slots, + cell_width, + )?; + } + self.mpm_dirty = false; + } Ok(()) } - - // /// Removes the given rigid-bodies from the simulation. - // pub fn remove_rigid_bodies( - // &mut self, - // backend: &GpuBackend, - // bodies: &[NexusRbdHandle], - // ) -> Result<(), GpuBackendError> { - // // 1. Resolve the rbd GPU slots of the bodies to remove. - // let gpu_ids: Vec = bodies - // .iter() - // .filter_map(|h| self.rbd2gpu.get(h.0).copied()) - // .collect(); - // - // // 2. Swap-remove them from the rbd GPU buffers. `remove_bodies` returns - // // the slot relocations it performed (`(from, to)`) so we can patch the - // // handle map. - // let remaps = match self.rbd.as_mut() { - // Some(rbd) => rbd.remove_bodies(backend, &gpu_ids)?, - // None => Vec::new(), - // }; - // - // // 3. Drop the removed handles, then patch the relocated slots. - // for h in bodies { - // self.rbd2gpu.remove(h.0); - // } - // for (from, to) in remaps { - // for (_, slot) in self.rbd2gpu.iter_mut() { - // if *slot == from { - // *slot = to; - // } - // } - // } - // - // Ok(()) - // } } diff --git a/src_rbd/dynamics/body.rs b/src_rbd/dynamics/body.rs new file mode 100644 index 00000000..e08239e9 --- /dev/null +++ b/src_rbd/dynamics/body.rs @@ -0,0 +1,420 @@ +//! Rigid-body definitions, mass properties, velocities, and GPU storage. +//! +//! This module provides the core data structures for representing rigid bodies on the GPU, +//! including their poses, velocities, forces, and mass properties. It also provides +//! [`GpuBodySet`] for managing collections of rigid bodies in GPU memory. + +use crate::math::{Pose, Vector}; +use crate::shapes::ShapeBuffers; + +use crate::shaders::dynamics::{LocalMassProperties, Velocity, WorldMassProperties}; +use crate::shaders::shapes::Shape; +use khal::BufferUsages; +use khal::backend::{GpuBackend, GpuBackendError}; +use vortx::tensor::Tensor; + +use crate::shaders::PaddedVector; +/// Re-export types from the shader crate for convenience. +pub use crate::shaders::dynamics::{ + Force, Impulse, LocalMassProperties as GpuLocalMassProperties, Velocity as GpuVelocity, + WorldMassProperties as GpuWorldMassProperties, +}; +use { + crate::rapier::dynamics::{RigidBodyHandle, RigidBodySet}, + crate::rapier::geometry::{ColliderHandle, ColliderSet}, + crate::rapier::prelude::MassProperties, + crate::shapes::shape_from_parry, +}; + +/// A set of rigid-bodies stored on the gpu. +pub struct GpuBodySet { + len: u32, + shapes_data: Vec, + /// World-space mass properties for each body. + pub mprops: Tensor, + /// Local-space mass properties for each body. + pub local_mprops: Tensor, + /// Velocities (linear + angular) for each body. + pub vels: Tensor, + /// Poses (position + orientation) for each body. + pub poses: Tensor, + /// Shape descriptors for each collider. + pub shapes: Tensor, + /// Vertex positions in local space for each collider shape. + pub shapes_local_vertex_buffers: Tensor, + /// Vertex positions in world space for each collider shape. + pub shapes_vertex_buffers: Tensor, + /// Triangle index buffers for each collider shape. + pub shapes_index_buffers: Tensor, + /// Collider ID for each vertex. + pub shapes_vertex_collider_id: Tensor, +} + +#[derive(Copy, Clone)] +/// Helper struct for defining a rigid-body to be added to a [`GpuBodySet`]. +pub struct BodyDesc { + /// The rigid-body's mass-properties in local-space. + pub local_mprops: LocalMassProperties, + /// The rigid-body's mass-properties in world-space. + pub mprops: WorldMassProperties, + /// The rigid-body's linear and angular velocities. + pub vel: Velocity, + /// The rigid-body's world-space pose. + pub pose: Pose, + /// The rigid-body's shape. + pub shape: Shape, +} + +impl Default for BodyDesc { + fn default() -> Self { + Self { + local_mprops: Default::default(), + mprops: Default::default(), + vel: Default::default(), + pose: Default::default(), + shape: Shape::cuboid(Vector::splat(0.5)), + } + } +} + +/// Coupling mode between a GPU body and the physics simulation. +/// +/// This controls whether a body is affected by physics forces or acts as a kinematic body. +#[derive(Copy, Clone, Debug, PartialEq, Eq, Default)] +pub enum BodyCoupling { + /// One-way coupling: the body affects other bodies but is not affected by them. + /// + /// This is useful for kinematic bodies that move independently of physics forces. + OneWay, + /// Two-way coupling: the body both affects and is affected by other bodies. + /// + /// This is the standard mode for dynamic rigid bodies. + #[default] + TwoWays, +} + +/// Associates a body/collider pair with a coupling mode, one entry per body a +/// [`GpuBodySet`] couples. +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +pub struct BodyCouplingEntry { + /// The rigid body index. + pub body: usize, + /// The collider index. + pub collider: usize, + /// The coupling mode for this body. + pub mode: BodyCoupling, +} + +/// Associates a Rapier body/collider pair with a coupling mode. The rapier-typed +/// counterpart of [`BodyCouplingEntry`]. +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +pub struct RapierBodyCouplingEntry { + /// The Rapier rigid body handle. + pub body: RigidBodyHandle, + /// The Rapier collider handle. + pub collider: ColliderHandle, + /// The coupling mode for this body. + pub mode: BodyCoupling, +} + +impl GpuBodySet { + /// Returns `true` if this set contains no rigid bodies. + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// Returns the number of rigid bodies in this set. + pub fn len(&self) -> u32 { + self.len + } + + /// Creates a new GPU body set from Rapier rigid bodies and colliders. + pub fn from_rapier( + backend: &GpuBackend, + bodies: &RigidBodySet, + colliders: &ColliderSet, + coupling: &[RapierBodyCouplingEntry], + ) -> Self { + let mut shape_buffers = ShapeBuffers::default(); + let mut gpu_bodies = vec![]; + let mut pt_collider_ids = vec![]; + + for (co_id, coupling) in coupling.iter().enumerate() { + let co = &colliders[coupling.collider]; + let rb = &bodies[coupling.body]; + + let prev_len = shape_buffers.vertices.len(); + let shape = + shape_from_parry(co.shape(), &mut shape_buffers).expect("Unsupported shape type"); + for _ in prev_len..shape_buffers.vertices.len() { + pt_collider_ids.push(co_id as u32); + } + + let zero_mprops = MassProperties::default(); + let two_ways_coupling = rb.is_dynamic() && coupling.mode == BodyCoupling::TwoWays; + let desc = BodyDesc { + vel: Velocity::new( + rb.linvel(), + #[cfg(feature = "dim2")] + rb.angvel(), + #[cfg(feature = "dim3")] + rb.angvel(), + ), + pose: *rb.position(), + shape, + local_mprops: if two_ways_coupling { + convert_local_mprops(&rb.mass_properties().local_mprops) + } else { + convert_local_mprops(&zero_mprops) + }, + mprops: Default::default(), + }; + gpu_bodies.push(desc); + } + + Self::new(backend, &gpu_bodies, &pt_collider_ids, &shape_buffers) + } + + /// Create a set of `bodies` on the gpu. + pub fn new( + backend: &GpuBackend, + bodies: &[BodyDesc], + pt_collider_ids: &[u32], + shape_buffers: &ShapeBuffers, + ) -> Self { + #[allow(clippy::type_complexity)] + let (local_mprops, (mprops, (vels, (poses, shapes_data)))): ( + Vec<_>, + (Vec<_>, (Vec<_>, (Vec<_>, Vec<_>))), + ) = bodies + .iter() + .copied() + .map(|b| (b.local_mprops, (b.mprops, (b.vel, (b.pose, b.shape))))) + .collect(); + + // Avoid empty buffer bindings. + let vertex_buffer = if !shape_buffers.vertices.is_empty() { + &shape_buffers.vertices[..] + } else { + &[PaddedVector::default()] + }; + let index_buffer = if !shape_buffers.indices.is_empty() { + &shape_buffers.indices[..] + } else { + &[0, 0, 0] + }; + + // All per-body buffers carry COPY_SRC | COPY_DST so the incremental + // `append` / `shift_remove` paths can write and relocate slots in place. + let resizeable = BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC; + Self { + len: bodies.len() as u32, + mprops: Tensor::vector(backend, &mprops, resizeable).unwrap(), + local_mprops: Tensor::vector(backend, &local_mprops, resizeable).unwrap(), + vels: Tensor::vector(backend, &vels, resizeable).unwrap(), + poses: Tensor::vector(backend, &poses, resizeable).unwrap(), + shapes: Tensor::vector(backend, &shapes_data, resizeable).unwrap(), + shapes_local_vertex_buffers: Tensor::vector(backend, vertex_buffer, resizeable) + .unwrap(), + shapes_vertex_buffers: Tensor::vector(backend, vertex_buffer, resizeable).unwrap(), + shapes_index_buffers: Tensor::vector(backend, index_buffer, resizeable).unwrap(), + shapes_vertex_collider_id: Tensor::vector(backend, pt_collider_ids, resizeable) + .unwrap(), + shapes_data, + } + } + + /// Creates an empty body set. + pub fn empty(backend: &GpuBackend) -> Self { + Self::new(backend, &[], &[], &ShapeBuffers::default()) + } + + /// Removes a range of body slots from this set, shifting later bodies down + /// to fill the gap. Returns the number of removed bodies. + /// + /// NOTE: this updates the per-body buffers only. Shapes referencing the + /// shared vertex/index buffers (trimesh, heightfield, polyline) are *not* + /// compacted, so the orphaned vertices remain allocated. Primitive + /// (vertex-less) colliders are fully handled. + pub fn shift_remove( + &mut self, + backend: &GpuBackend, + range: impl std::ops::RangeBounds + Clone, + ) -> Result { + let removed = self.poses.shift_remove(backend, range.clone())?; + self.vels.shift_remove(backend, range.clone())?; + self.mprops.shift_remove(backend, range.clone())?; + self.local_mprops.shift_remove(backend, range.clone())?; + self.shapes.shift_remove(backend, range.clone())?; + + // Mirror the CPU-side shape cache. + let start = match range.start_bound() { + std::ops::Bound::Included(i) => *i, + std::ops::Bound::Excluded(i) => *i + 1, + std::ops::Bound::Unbounded => 0, + }; + self.shapes_data.drain(start..start + removed); + self.len -= removed as u32; + Ok(removed) + } + + /// GPU storage buffer containing the poses of every rigid-body. + pub fn poses(&self) -> &Tensor { + &self.poses + } + + /// GPU storage buffer containing the velocities of every rigid-body. + pub fn vels(&self) -> &Tensor { + &self.vels + } + + /// GPU storage buffer containing the world-space mass-properties of every rigid-body. + pub fn mprops(&self) -> &Tensor { + &self.mprops + } + + /// GPU storage buffer containing the local-space mass-properties of every rigid-body. + pub fn local_mprops(&self) -> &Tensor { + &self.local_mprops + } + + /// GPU storage buffer containing the shape of every rigid-body. + pub fn shapes(&self) -> &Tensor { + &self.shapes + } + + /// Mutable reference to the GPU storage buffer containing the poses of every rigid-body. + pub fn poses_mut(&mut self) -> &mut Tensor { + &mut self.poses + } + + /// Mutable reference to the GPU storage buffer containing the velocities of every rigid-body. + pub fn vels_mut(&mut self) -> &mut Tensor { + &mut self.vels + } + + /// Mutable reference to the GPU storage buffer containing the world-space mass-properties of every rigid-body. + pub fn mprops_mut(&mut self) -> &mut Tensor { + &mut self.mprops + } + + /// Returns the GPU buffer containing shape vertices in world-space. + /// + /// This buffer is updated each frame as bodies move. + pub fn shapes_vertex_buffers(&self) -> &Tensor { + &self.shapes_vertex_buffers + } + + /// Mutable reference to the GPU buffer containing shape vertices in world-space. + pub fn shapes_vertex_buffers_mut(&mut self) -> &mut Tensor { + &mut self.shapes_vertex_buffers + } + + /// Returns the GPU buffer containing shape vertices in local-space. + /// + /// These are the original vertex positions before transformation. + pub fn shapes_local_vertex_buffers(&self) -> &Tensor { + &self.shapes_local_vertex_buffers + } + + /// Returns the GPU buffer mapping each vertex to its collider ID. + pub fn shapes_vertex_collider_id(&self) -> &Tensor { + &self.shapes_vertex_collider_id + } + + /// Returns a CPU-side slice of the shape data. + /// + /// Useful for accessing shape information without GPU readback. + pub fn shapes_data(&self) -> &[Shape] { + &self.shapes_data + } +} + +impl GpuBodySet { + /// Appends rigid-bodies (converted from rapier) to this set and returns the + /// indices of the newly inserted bodies. + /// + /// Only primitive (vertex-less) colliders are supported by this incremental + /// path; mesh-based colliders (trimesh / heightfield / polyline) would + /// require growing the shared vertex/index buffers and offsetting the shape + /// references, which isn't handled yet. + pub fn append_rapier( + &mut self, + backend: &GpuBackend, + bodies: &[( + crate::rapier::dynamics::RigidBody, + crate::rapier::geometry::Collider, + BodyCoupling, + )], + ) -> Result, GpuBackendError> { + let start = self.len; + let mut poses = Vec::with_capacity(bodies.len()); + let mut vels = Vec::with_capacity(bodies.len()); + let mut mprops = Vec::with_capacity(bodies.len()); + let mut local_mprops = Vec::with_capacity(bodies.len()); + let mut shapes = Vec::with_capacity(bodies.len()); + + for (rb, co, coupling) in bodies { + let mut shape_buffers = ShapeBuffers::default(); + let shape = + shape_from_parry(co.shape(), &mut shape_buffers).expect("Unsupported shape type"); + assert!( + shape_buffers.vertices.is_empty(), + "GpuBodySet::append_rapier currently supports primitive (vertex-less) colliders only." + ); + + let two_ways_coupling = rb.is_dynamic() && *coupling == BodyCoupling::TwoWays; + let local = if two_ways_coupling { + convert_local_mprops(&rb.mass_properties().local_mprops) + } else { + convert_local_mprops(&MassProperties::default()) + }; + + poses.push(*rb.position()); + vels.push(Velocity::new( + rb.linvel(), + #[cfg(feature = "dim2")] + rb.angvel(), + #[cfg(feature = "dim3")] + rb.angvel(), + )); + mprops.push(WorldMassProperties::default()); + local_mprops.push(local); + shapes.push(shape); + } + + self.poses.append(backend, &poses)?; + self.vels.append(backend, &vels)?; + self.mprops.append(backend, &mprops)?; + self.local_mprops.append(backend, &local_mprops)?; + self.shapes.append(backend, &shapes)?; + self.shapes_data.extend_from_slice(&shapes); + self.len += bodies.len() as u32; + + Ok((start..self.len).collect()) + } +} + +fn convert_local_mprops(mprops: &MassProperties) -> LocalMassProperties { + #[cfg(feature = "dim2")] + { + LocalMassProperties { + inv_mass: glamx::Vec2::splat(mprops.inv_mass), + com: mprops.local_com, + padding2: 0, + inv_inertia: mprops.inv_principal_inertia, + } + } + #[cfg(feature = "dim3")] + { + LocalMassProperties { + inertia_ref_frame: mprops.principal_inertia_local_frame, + inv_principal_inertia: mprops.inv_principal_inertia, + padding0: 0, + inv_mass: glamx::Vec3::splat(mprops.inv_mass), + padding1: 0, + com: mprops.local_com, + padding2: 0, + } + } +} diff --git a/src_rbd/dynamics/mod.rs b/src_rbd/dynamics/mod.rs index e67b91e1..c20e1f41 100644 --- a/src_rbd/dynamics/mod.rs +++ b/src_rbd/dynamics/mod.rs @@ -1,6 +1,7 @@ //! Rigid-body dynamics: forces, velocities, constraints, and solvers. pub use crate::shaders::dynamics::RbdSimParams; +pub use body::{BodyCoupling, BodyCouplingEntry, BodyDesc, GpuBodySet}; pub use coloring::{ColorBucketsArgs, ColoringArgs, GpuColoring}; pub use joint::{GpuImpulseJointSet, GpuJointSolver, JointSolverArgs, convert_joint_motor}; pub use mprops_update::{GpuMpropsUpdate, GpuSyncColliderPosesShader}; @@ -10,6 +11,7 @@ pub use prep_render::{RbdInstanceDesc, WgRbdPrepRender}; pub use solver::{GpuSolver, SolverArgs}; pub use warmstart::{GpuWarmstart, WarmstartArgs}; +pub mod body; mod coloring; mod joint; mod mprops_update; diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 2016e005..ac691fac 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -322,6 +322,15 @@ impl RbdState { &self.body_poses } + /// Mutable access to the per-body world-origin poses, for bodies this + /// pipeline doesn't own and that another solver integrates itself. + /// + /// Only valid between steps: the solver works on `solver_body_poses` and + /// overwrites this buffer wholesale in its finalize pass. + pub fn body_poses_mut(&mut self) -> &mut Tensor { + &mut self.body_poses + } + /// Live collision-pair count (batch 0) most recently harvested by the /// non-blocking readback in [`RbdPipeline::auto_resize_buffers`](crate::pipeline::RbdPipeline::auto_resize_buffers). Lags the GPU by a /// frame or two; `0` until the first readback completes. diff --git a/src_rbd/utils/prefix_sum.rs b/src_rbd/utils/prefix_sum.rs index ab2727ab..a89afcf4 100644 --- a/src_rbd/utils/prefix_sum.rs +++ b/src_rbd/utils/prefix_sum.rs @@ -118,8 +118,9 @@ impl GpuPrefixSum { let batch_data = &mut slice[start..start + batch_stride]; // Inclusive prefix sum. Uses wrapping_add to match GPU u32 - // semantics: callers may leave tail entries uninitialized, and the - // corresponding summed results are never read back downstream. + // semantics: callers may leave tail entries uninitialized, + // and the corresponding summed results are never read back + // downstream. for i in 0..batch_data.len() - 1 { batch_data[i + 1] = batch_data[i + 1].wrapping_add(batch_data[i]); } From 668170dc5995ca40b83d7452a087e4e19b0e98f0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 14 Aug 2026 23:47:14 +0200 Subject: [PATCH 3/7] feat: render MPM particles in the viewer --- crates/nexus_viewer2d/Cargo.toml | 2 +- crates/nexus_viewer3d/Cargo.toml | 2 +- src_viewer/graphics.rs | 9 +- src_viewer/lib.rs | 1 + src_viewer/rbd/backend/cpu.rs | 137 ------------- src_viewer/rbd/backend/gpu.rs | 196 ------------------- src_viewer/rbd/backend/mod.rs | 102 ---------- src_viewer/rbd/mod.rs | 292 ---------------------------- src_viewer/ui.rs | 178 +++++------------ src_viewer/viewer.rs | 320 ++++++++++++++++++++++++++++++- 10 files changed, 372 insertions(+), 867 deletions(-) delete mode 100644 src_viewer/rbd/backend/cpu.rs delete mode 100644 src_viewer/rbd/backend/gpu.rs delete mode 100644 src_viewer/rbd/backend/mod.rs delete mode 100644 src_viewer/rbd/mod.rs diff --git a/crates/nexus_viewer2d/Cargo.toml b/crates/nexus_viewer2d/Cargo.toml index d24c3634..29b41a86 100644 --- a/crates/nexus_viewer2d/Cargo.toml +++ b/crates/nexus_viewer2d/Cargo.toml @@ -29,7 +29,7 @@ cpu-parallel = ["cpu", "nexus2d/cpu-parallel", "vortx/cpu-parallel"] cuda = ["nexus2d/cuda"] [dependencies] -nexus2d = { workspace = true, features = ["rbd"]} +nexus2d = { workspace = true, features = ["rbd", "mpm"]} glamx = { workspace = true } khal = { workspace = true } vortx = { workspace = true } diff --git a/crates/nexus_viewer3d/Cargo.toml b/crates/nexus_viewer3d/Cargo.toml index 0a4b9a78..b14824c0 100644 --- a/crates/nexus_viewer3d/Cargo.toml +++ b/crates/nexus_viewer3d/Cargo.toml @@ -29,7 +29,7 @@ cpu-parallel = ["cpu", "nexus3d/cpu-parallel", "vortx/cpu-parallel"] cuda = ["nexus3d/cuda"] [dependencies] -nexus3d = { workspace = true, features = ["rbd"]} +nexus3d = { workspace = true, features = ["rbd", "mpm"]} glamx = { workspace = true } khal = { workspace = true } vortx = { workspace = true } diff --git a/src_viewer/graphics.rs b/src_viewer/graphics.rs index 29d0246e..4c4d8be4 100644 --- a/src_viewer/graphics.rs +++ b/src_viewer/graphics.rs @@ -157,10 +157,16 @@ impl InstancedNode { /// opaque and translucent instances of the same shape must live in separate /// nodes because the transparent/opaque phase split is per-node, not /// per-instance. + /// + /// Backface culling is off on every node: instance transforms may mirror a + /// shape (a negative scale flips the winding), demos routinely look at thin + /// boundary colliders and open geometry from inside, and the translucent + /// nodes need their back faces to read as solid volumes. #[cfg(feature = "dim2")] - fn new(node: SceneNode2d, _transparent: bool) -> Self { + fn new(mut node: SceneNode2d, _transparent: bool) -> Self { // 2D rendering blends per-fragment regardless of node, so the split only // keeps the keying symmetric with 3D — no extra styling needed here. + node.enable_backface_culling_recursive(false); Self { node, entries: vec![], @@ -177,6 +183,7 @@ impl InstancedNode { if transparent { node.set_color(Color::new(1.0, 1.0, 1.0, TRANSPARENT_NODE_ALPHA)); } + node.enable_backface_culling_recursive(false); Self { node, entries: vec![], diff --git a/src_viewer/lib.rs b/src_viewer/lib.rs index b7abb1d3..b608d285 100644 --- a/src_viewer/lib.rs +++ b/src_viewer/lib.rs @@ -39,6 +39,7 @@ pub struct UiSections { #[derive(Copy, Clone, PartialEq, Eq, Debug)] pub enum DemoKind { Rbd, + Mpm, } /// A loop transition requested from the UI: stop entirely, or switch to another diff --git a/src_viewer/rbd/backend/cpu.rs b/src_viewer/rbd/backend/cpu.rs deleted file mode 100644 index ecc20734..00000000 --- a/src_viewer/rbd/backend/cpu.rs +++ /dev/null @@ -1,137 +0,0 @@ -use super::SimulationBackend; -use crate::rbd::SimulationState; -use khal::backend::GpuBackend as KhalGpuBackend; -use nexus::rbd::math::Pose; -use nexus::rbd::pipeline::RunStats; -use rapier::dynamics::{CCDSolver, IntegrationParameters, IslandManager}; -use rapier::geometry::{BroadPhaseBvh, ColliderSet, NarrowPhase}; -use rapier::prelude::{ - ImpulseJointSet, JointAxis, MultibodyJointSet, PhysicsPipeline, RigidBodySet, -}; - -/// CPU-based physics backend using rapier -pub struct CpuBackend { - pipeline: PhysicsPipeline, - integration_parameters: IntegrationParameters, - islands: IslandManager, - broad_phase: BroadPhaseBvh, - narrow_phase: NarrowPhase, - bodies: RigidBodySet, - colliders: ColliderSet, - impulse_joints: ImpulseJointSet, - multibody_joints: MultibodyJointSet, - ccd_solver: CCDSolver, - poses_cache: Vec, -} - -impl CpuBackend { - /// See [`super::PhysicsBackend::set_multibody_motor_velocity`]. - pub fn set_multibody_motor_velocity( - &mut self, - _batch: u32, - link_id: u32, - axis: JointAxis, - target_vel: f32, - ) { - // Walk every multibody and apply on the link whose body's local id - // matches `link_id`. For a one-collider-per-body world the local id is - // also the body id. - let mut handle = None; - for (h, _) in self.bodies.iter() { - if h.into_raw_parts().0 as u32 == link_id { - handle = Some(h); - break; - } - } - let Some(handle) = handle else { return }; - if let Some(link) = self.multibody_joints.rigid_body_link(handle) { - let multibody_handle = link.multibody; - let link_id_in_mb = link.id; - if let Some(mb) = self.multibody_joints.get_multibody_mut(multibody_handle) { - if let Some(link) = mb.link_mut(link_id_in_mb) { - link.joint.data.set_motor_velocity(axis, target_vel, 1.0); - } - } - } - } - - pub fn new(phys: &SimulationState) -> Self { - let env = &phys.environments[0]; - let mut poses_cache = Vec::new(); - let mut shapes_cache = Vec::new(); - - // Build initial poses and shapes from the first environment. - for (_, co) in env.colliders.iter() { - poses_cache.push(*co.position()); - shapes_cache.push(co.shared_shape().clone()); - } - - let mut params = IntegrationParameters::default(); - params.dt = env.sim_params.dt; - Self { - pipeline: PhysicsPipeline::new(), - integration_parameters: params, - islands: IslandManager::new(), - broad_phase: BroadPhaseBvh::new(), - narrow_phase: NarrowPhase::new(), - bodies: env.bodies.clone(), - colliders: env.colliders.clone(), - impulse_joints: env.impulse_joints.clone(), - multibody_joints: env.multibody_joints.clone(), - ccd_solver: CCDSolver::new(), - poses_cache, - } - } -} - -impl SimulationBackend for CpuBackend { - fn poses(&self) -> &[Pose] { - &self.poses_cache - } - fn num_bodies(&self) -> usize { - self.poses().len() - } - fn num_joints(&self) -> usize { - self.impulse_joints.len() - } - - fn num_batches(&self) -> usize { - 1 - } - - async fn step(&mut self, _gpu: Option<&KhalGpuBackend>) -> RunStats { - let t0 = web_time::Instant::now(); - - #[cfg(feature = "dim2")] - let gravity = glamx::Vec2::Y * -9.81; - #[cfg(feature = "dim3")] - let gravity = glamx::Vec3::Y * -9.81; - - self.pipeline.step( - gravity, - &self.integration_parameters, - &mut self.islands, - &mut self.broad_phase, - &mut self.narrow_phase, - &mut self.bodies, - &mut self.colliders, - &mut self.impulse_joints, - &mut self.multibody_joints, - &mut self.ccd_solver, - &(), - &(), - ); - let total_sim_time = t0.elapsed(); - - // Update poses cache - self.poses_cache.clear(); - for (_, co) in self.colliders.iter() { - self.poses_cache.push(*co.position()); - } - - RunStats { - total_simulation_time_with_readback: total_sim_time, - ..Default::default() - } - } -} diff --git a/src_viewer/rbd/backend/gpu.rs b/src_viewer/rbd/backend/gpu.rs deleted file mode 100644 index c41b47f6..00000000 --- a/src_viewer/rbd/backend/gpu.rs +++ /dev/null @@ -1,196 +0,0 @@ -use super::SimulationBackend; -use crate::rbd::SimulationState; -use khal::backend::{Backend, GpuBackend as KhalGpuBackend, GpuTimestamps}; -use nexus::rbd::math::Pose; -use nexus::rbd::pipeline::{RbdPipeline, RbdState, RunStats}; -use rapier::prelude::JointAxis; - -/// GPU-based physics backend using nexus -pub struct GpuBackend { - gpu: KhalGpuBackend, - pipeline: RbdPipeline, - state: RbdState, - poses_cache: Vec, - timestamps: GpuTimestamps, -} - -impl GpuBackend { - /// Reads poses from GPU buffer, handling dimension-specific conversion. - async fn read_poses( - gpu: &KhalGpuBackend, - state: &RbdState, - ) -> Result, String> { - gpu.slow_read_vec(state.poses().buffer()) - .await - .map_err(|e| format!("Failed to read poses: {:?}", e)) - } - - /// Reads poses into an existing buffer, handling dimension-specific conversion. - async fn read_poses_into( - gpu: &KhalGpuBackend, - state: &RbdState, - poses_cache: &mut Vec, - ) { - poses_cache.resize(state.poses().len() as usize, Pose::default()); - let _ = gpu - .slow_read_buffer(state.poses().buffer(), poses_cache) - .await; - } - - /// Attempts to create a new GPU backend, returning an error if initialization fails. - /// - /// This method can fail if: - /// - Shader compilation fails - /// - GPU device doesn't support required features - /// - Memory allocation fails - pub async fn try_new(gpu: &KhalGpuBackend, phys: &SimulationState) -> Result { - let pipeline = RbdPipeline::from_backend(gpu); - let envs: Vec<_> = phys - .environments - .iter() - .map(|e| { - ( - &e.bodies, - &e.colliders, - &e.impulse_joints, - &e.multibody_joints, - &e.sim_params, - ) - }) - .collect(); - let state = RbdState::from_rapier(gpu, &envs); - let poses_cache = Self::read_poses(gpu, &state).await?; - let timestamps = GpuTimestamps::new(gpu, 2048); - - Ok(Self { - gpu: gpu.clone(), - pipeline, - state, - poses_cache, - timestamps, - }) - } - - /// Creates a new GPU backend with a pre-compiled pipeline. - /// - /// This is faster than [`try_new`](Self::try_new) when switching demos because - /// it reuses the existing pipeline instead of recompiling shaders. - pub async fn with_pipeline( - gpu: &KhalGpuBackend, - pipeline: RbdPipeline, - phys: &SimulationState, - ) -> Self { - let envs: Vec<_> = phys - .environments - .iter() - .map(|e| { - ( - &e.bodies, - &e.colliders, - &e.impulse_joints, - &e.multibody_joints, - &e.sim_params, - ) - }) - .collect(); - let state = RbdState::from_rapier(gpu, &envs); - let poses_cache = Self::read_poses(gpu, &state).await.unwrap_or_default(); - let timestamps = GpuTimestamps::new(gpu, 2048); - - Self { - gpu: gpu.clone(), - pipeline, - state, - poses_cache, - timestamps, - } - } - - /// Extracts the pipeline from this backend, consuming it. - /// - /// Useful for reusing the pipeline when switching demos. - pub fn into_pipeline(self) -> RbdPipeline { - self.pipeline - } - - /// Creates a new GPU backend, panicking if initialization fails. - /// - /// Use [`try_new`](Self::try_new) for error handling. - pub async fn new(gpu: &KhalGpuBackend, phys: &SimulationState) -> Self { - Self::try_new(gpu, phys).await.unwrap() - } - - /// See [`super::PhysicsBackend::set_multibody_motor_velocity`]. - pub fn set_multibody_motor_velocity( - &mut self, - batch: u32, - link_id: u32, - axis: JointAxis, - target_vel: f32, - ) { - #[cfg(feature = "dim3")] - let _ = self - .state - .multibodies_mut() - .set_motor_velocity(&self.gpu, batch, link_id, axis, target_vel); - #[cfg(feature = "dim2")] - { - let _ = (batch, link_id, axis, target_vel); - todo!() - } - } -} - -impl SimulationBackend for GpuBackend { - fn poses(&self) -> &[Pose] { - &self.poses_cache - } - fn num_bodies(&self) -> usize { - self.state.num_colliders_per_batch() as usize - } - fn num_joints(&self) -> usize { - self.state.joints().len() - } - fn num_batches(&self) -> usize { - self.state.num_batches() as usize - } - - async fn step(&mut self, _gpu: Option<&KhalGpuBackend>) -> RunStats { - let gpu = &self.gpu; - - self.timestamps.reset(); - - let t0 = web_time::Instant::now(); - let mut run_stats = self - .pipeline - .step(gpu, &mut self.state, Some(&mut self.timestamps)) - .unwrap(); - - // Read back poses (synchronizes with the GPU when using WebGPU backend). - gpu.synchronize().unwrap(); - run_stats.total_simulation_time_without_readback = t0.elapsed(); - - self.pipeline - .auto_resize_buffers(gpu, &mut self.state) - .await; - Self::read_poses_into(gpu, &self.state, &mut self.poses_cache).await; - - // Read timestamp results. - if let Ok(results) = self.timestamps.read(gpu).await { - let mut aggregated: Vec<(String, f64)> = Vec::new(); - for r in &results { - if let Some(existing) = aggregated.iter_mut().find(|(label, _)| label == &r.label) { - existing.1 += r.duration_ms; - } else { - aggregated.push((r.label.clone(), r.duration_ms)); - } - } - run_stats.gpu_total_time = aggregated.iter().map(|e| e.1).sum(); - run_stats.gpu_pass_times = aggregated; - } - - run_stats.total_simulation_time_with_readback = t0.elapsed(); - - run_stats - } -} diff --git a/src_viewer/rbd/backend/mod.rs b/src_viewer/rbd/backend/mod.rs deleted file mode 100644 index 32f0a4fd..00000000 --- a/src_viewer/rbd/backend/mod.rs +++ /dev/null @@ -1,102 +0,0 @@ -mod cpu; -mod gpu; - -pub use cpu::CpuBackend; -pub use gpu::GpuBackend; - -use khal::backend::GpuBackend as KhalGpuBackend; -use nexus::rbd::math::Pose; -use nexus::rbd::pipeline::RunStats; -use rapier::prelude::JointAxis; - -#[derive(Copy, Clone, Debug, PartialEq, Eq)] -pub enum BackendType { - /// CPU physics using rapier. - Rapier, - /// GPU-accelerated physics using nexus + WebGPU. - Gpu, - /// CPU physics using nexus (same pipeline as GPU, executed on CPU). - Cpu, - /// GPU-accelerated physics using nexus + CUDA. - #[cfg(feature = "cuda")] - Cuda, - /// GPU-accelerated physics using nexus + native Metal (macOS only). - #[cfg(feature = "metal")] - Metal, -} - -/// Trait for physics simulation backends (CPU or GPU) -pub trait SimulationBackend { - /// Get the current poses for rendering - fn poses(&self) -> &[Pose]; - fn num_bodies(&self) -> usize; - fn num_joints(&self) -> usize; - fn num_batches(&self) -> usize; - - /// Step the simulation - #[allow(async_fn_in_trait)] - async fn step(&mut self, gpu: Option<&KhalGpuBackend>) -> RunStats; -} - -#[allow(clippy::large_enum_variant)] -pub enum PhysicsBackend { - Cpu(CpuBackend), - Gpu(GpuBackend), -} - -impl PhysicsBackend { - pub async fn step(&mut self, gpu: Option<&KhalGpuBackend>) -> RunStats { - match self { - PhysicsBackend::Cpu(backend) => backend.step(gpu).await, - PhysicsBackend::Gpu(backend) => backend.step(gpu).await, - } - } - - pub fn poses(&self) -> &[Pose] { - match self { - PhysicsBackend::Cpu(backend) => backend.poses(), - PhysicsBackend::Gpu(backend) => backend.poses(), - } - } - - pub fn num_bodies(&self) -> usize { - match self { - PhysicsBackend::Cpu(backend) => backend.num_bodies(), - PhysicsBackend::Gpu(backend) => backend.num_bodies(), - } - } - - pub fn num_joints(&self) -> usize { - match self { - PhysicsBackend::Cpu(backend) => backend.num_joints(), - PhysicsBackend::Gpu(backend) => backend.num_joints(), - } - } - - pub fn num_batches(&self) -> usize { - match self { - PhysicsBackend::Cpu(backend) => backend.num_batches(), - PhysicsBackend::Gpu(backend) => backend.num_batches(), - } - } - - /// Set a multibody joint motor's target velocity for the link `link_id` - /// in batch `batch`, on joint axis `axis` (0..=2 = linear, 3..=5 = angular). - /// The motor on that axis is enabled automatically. - pub fn set_multibody_motor_velocity( - &mut self, - batch: u32, - link_id: u32, - axis: JointAxis, - target_vel: f32, - ) { - match self { - PhysicsBackend::Cpu(backend) => { - backend.set_multibody_motor_velocity(batch, link_id, axis, target_vel) - } - PhysicsBackend::Gpu(backend) => { - backend.set_multibody_motor_velocity(batch, link_id, axis, target_vel) - } - } - } -} diff --git a/src_viewer/rbd/mod.rs b/src_viewer/rbd/mod.rs deleted file mode 100644 index ff34cb4d..00000000 --- a/src_viewer/rbd/mod.rs +++ /dev/null @@ -1,292 +0,0 @@ -pub mod backend; -pub mod graphics; - -pub use backend::{BackendType, CpuBackend, GpuBackend, PhysicsBackend}; -pub use graphics::{RenderContext}; - -use crate::RunState; -use khal::backend::GpuBackend as KhalGpuBackend; -use nexus::rbd::dynamics::RbdSimParams; -use nexus::rbd::math::Pose; -use nexus::rbd::pipeline::{RbdPipeline, RunStats}; -use rapier::geometry::{ColliderHandle, ColliderSet, SharedShape}; -use rapier::prelude::{ImpulseJointSet, MultibodyJointSet, RigidBodySet}; -use std::collections::HashMap; -use nexus::state::NexusState; - -/// Custom visual shape that overrides a collider's default rendering. The shape is -/// drawn at the collider's world pose composed with [`Self::local_pose`] (handy when -/// the collider's shape is a proxy — e.g. an OBB — and the visual mesh sits in a -/// different local frame). -#[derive(Clone)] -pub struct VisualShape { - pub shape: SharedShape, - pub local_pose: Pose, -} - -impl VisualShape { - pub fn new(shape: SharedShape) -> Self { - Self { - shape, - local_pose: Pose::IDENTITY, - } - } - - pub fn with_local_pose(shape: SharedShape, local_pose: Pose) -> Self { - Self { shape, local_pose } - } -} - -pub struct BatchEnvironment { - pub bodies: RigidBodySet, - pub colliders: ColliderSet, - pub impulse_joints: ImpulseJointSet, - pub multibody_joints: MultibodyJointSet, - pub sim_params: RbdSimParams, - /// Optional per-collider visual override. When a collider handle is present in - /// this map its [`VisualShape`] is rendered instead of the collider's own shape. - pub visuals: HashMap, -} - -pub struct SimulationState { - pub environments: Vec, - /// Number of physics steps run per render frame (default: `1`). - pub num_steps_per_frame: u32, -} - -impl SimulationState { - pub fn from_environments(environments: Vec) -> Self { - Self { - environments, - num_steps_per_frame: 1, - } - } - - pub fn single( - bodies: RigidBodySet, - colliders: ColliderSet, - impulse_joints: ImpulseJointSet, - ) -> Self { - Self::single_with_multibody(bodies, colliders, impulse_joints, MultibodyJointSet::new()) - } - - pub fn single_with_multibody( - bodies: RigidBodySet, - colliders: ColliderSet, - impulse_joints: ImpulseJointSet, - multibody_joints: MultibodyJointSet, - ) -> Self { - Self::single_with_multibody_and_visuals( - bodies, - colliders, - impulse_joints, - multibody_joints, - HashMap::new(), - ) - } - - pub fn single_with_multibody_and_visuals( - bodies: RigidBodySet, - colliders: ColliderSet, - impulse_joints: ImpulseJointSet, - multibody_joints: MultibodyJointSet, - visuals: HashMap, - ) -> Self { - Self::from_environments(vec![BatchEnvironment { - bodies, - colliders, - impulse_joints, - multibody_joints, - sim_params: RbdSimParams::default(), - visuals, - }]) - } - - /// Sets `sim_params.dt` on every environment to `dt`. - pub fn with_dt(mut self, dt: f32) -> Self { - for env in &mut self.environments { - env.sim_params.dt = dt; - } - self - } - - /// Sets the number of physics steps run between two renders. - pub fn with_num_steps_per_frame(mut self, num_steps_per_frame: u32) -> Self { - self.num_steps_per_frame = num_steps_per_frame; - self - } -} - -pub struct PhysicsContext { - pub backend: PhysicsBackend, -} - -impl PhysicsContext { - pub fn new(backend: PhysicsBackend) -> Self { - Self { backend } - } -} - -/// A rigid-body scene: GPU/CPU physics state plus its rendering instances. -/// -/// Built via [`crate::NexusViewer::set_rbd`]. The example owns this and drives the -/// loop with [`RbdScene::simulate`]. -pub struct RbdScene { - pub physics: PhysicsContext, - pub render_ctx: RenderContext, - /// Total elapsed simulated time, accumulated across non-paused steps. - pub sim_time: f64, - /// Per-step timestep length (copied from the scene's `sim_params.dt`). - pub dt: f32, - /// Number of physics steps run per render frame. - pub num_steps_per_frame: u32, - /// Backend the scene was created with, used to decide whether the GPU - /// pipeline can be cached on teardown. - pub(crate) created_backend: BackendType, -} - -impl RbdScene { - /// Mutable access to the physics backend (e.g. to drive joint motors). - pub fn backend_mut(&mut self) -> &mut PhysicsBackend { - &mut self.physics.backend - } - - /// Runs a single physics step. Self-contained (uses the backend's own GPU - /// device); does not render. This is the headless/Python entry point. - pub async fn step(&mut self) -> RunStats { - self.physics.backend.step(None).await - } - - /// Pushes the latest poses into the kiss3d render instances. - pub fn sync_graphics(&mut self, state: &NexusState) { - self.render_ctx.update_instances(state, &self.physics.backend); - } - - /// Advances the simulation for one render frame (honoring pause/step) and - /// syncs graphics. Call this inside the example's loop body. - pub async fn simulate(&mut self, viewer: &mut crate::NexusViewer) { - todo!() - // if viewer.ui.run_state != RunState::Paused { - // for _ in 0..self.num_steps_per_frame { - // viewer.ui.run_stats = self.step().await; - // self.sim_time += self.dt as f64; - // } - // } - // self.sync_graphics(); - // if viewer.ui.run_state == RunState::Step { - // viewer.ui.run_state = RunState::Paused; - // } - } - - /// Detaches the render nodes and, when the backend is unchanged, caches the - /// compiled GPU pipeline in the viewer for reuse by the next RBD scene. - pub fn detach(self, viewer: &mut crate::NexusViewer) { - let RbdScene { - mut render_ctx, - physics, - created_backend, - .. - } = self; - render_ctx.clear(); - if created_backend == viewer.ui.backend_type { - if let PhysicsBackend::Gpu(gpu_backend) = physics.backend { - viewer.cache_pipeline(gpu_backend.into_pipeline()); - } - } - } -} - -pub async fn setup_physics( - gpu: Option<&KhalGpuBackend>, - phys: &SimulationState, - backend_type: BackendType, - gpu_error: &mut Option, - cached_pipeline: &mut Option, -) -> PhysicsContext { - let backend = match backend_type { - BackendType::Gpu => { - let gpu = gpu.unwrap(); - - if let Some(pipeline) = cached_pipeline.take() { - let gpu_backend = GpuBackend::with_pipeline(gpu, pipeline, phys).await; - PhysicsBackend::Gpu(gpu_backend) - } else { - match GpuBackend::try_new(gpu, phys).await { - Ok(gpu_backend) => PhysicsBackend::Gpu(gpu_backend), - Err(e) => { - *gpu_error = Some(format!( - "GPU backend initialization failed: {}. Using CPU backend.", - e - )); - PhysicsBackend::Cpu(CpuBackend::new(phys)) - } - } - } - } - BackendType::Cpu => { - #[cfg(feature = "cpu")] - { - let cpu_backend = KhalGpuBackend::Cpu; - match GpuBackend::try_new(&cpu_backend, phys).await { - Ok(gpu_backend) => PhysicsBackend::Gpu(gpu_backend), - Err(e) => { - *gpu_error = Some(format!( - "Nexus CPU backend initialization failed: {}. Using rapier CPU backend.", - e - )); - PhysicsBackend::Cpu(CpuBackend::new(phys)) - } - } - } - #[cfg(not(feature = "cpu"))] - { - *gpu_error = - Some("CPU backend not available (compiled without 'cpu' feature).".to_string()); - PhysicsBackend::Cpu(CpuBackend::new(phys)) - } - } - #[cfg(feature = "cuda")] - BackendType::Cuda => { - let gpu = gpu.expect("Cuda device initialization failed"); - - if let Some(pipeline) = cached_pipeline.take() { - let gpu_backend = GpuBackend::with_pipeline(gpu, pipeline, phys).await; - PhysicsBackend::Gpu(gpu_backend) - } else { - match GpuBackend::try_new(gpu, phys).await { - Ok(gpu_backend) => PhysicsBackend::Gpu(gpu_backend), - Err(e) => { - *gpu_error = Some(format!( - "CUDA backend initialization failed: {}. Using CPU backend.", - e - )); - PhysicsBackend::Cpu(CpuBackend::new(phys)) - } - } - } - } - #[cfg(feature = "metal")] - BackendType::Metal => { - let gpu = gpu.expect("Metal device initialization failed"); - - if let Some(pipeline) = cached_pipeline.take() { - let gpu_backend = GpuBackend::with_pipeline(gpu, pipeline, phys).await; - PhysicsBackend::Gpu(gpu_backend) - } else { - match GpuBackend::try_new(gpu, phys).await { - Ok(gpu_backend) => PhysicsBackend::Gpu(gpu_backend), - Err(e) => { - *gpu_error = Some(format!( - "Metal backend initialization failed: {}. Using CPU backend.", - e - )); - PhysicsBackend::Cpu(CpuBackend::new(phys)) - } - } - } - } - BackendType::Rapier => PhysicsBackend::Cpu(CpuBackend::new(phys)), - }; - - PhysicsContext::new(backend) -} diff --git a/src_viewer/ui.rs b/src_viewer/ui.rs index fb706abf..3f8886d1 100644 --- a/src_viewer/ui.rs +++ b/src_viewer/ui.rs @@ -1,13 +1,12 @@ -use std::time::Duration; -// use crate::rbd::BackendType; -use crate::viewer::UiState; +use crate::viewer::{MpmRenderMode, UiState}; use crate::{DemoKind, RunState, Transition}; use kiss3d::egui; use nexus::rbd::pipeline::RunStats; use nexus::state::NexusCounts; +use std::time::Duration; use crate::backend::BackendType; -use egui::{Button, CollapsingHeader, Color32, CornerRadius, RichText, Stroke}; +use egui::{Button, CollapsingHeader, Color32, ComboBox, CornerRadius, RichText, Stroke}; /// Sets up a custom warm theme that complements the app's off-white background. pub fn setup_custom_theme(ctx: &egui::Context) { @@ -194,8 +193,8 @@ pub fn main_panel(ctx: &egui::Context, state: &mut UiState, gpu_available: bool) /// [`crate::NexusViewer::sync`]). Only the groups relevant to the current scene /// are shown. fn simulation_settings(ui: &mut egui::Ui, state: &mut UiState) { - let has_rbd = state.has_rbd; - if !has_rbd { + let (has_mpm, has_rbd) = (state.has_mpm, state.has_rbd); + if !(has_mpm || has_rbd) { return; } @@ -203,8 +202,40 @@ fn simulation_settings(ui: &mut egui::Ui, state: &mut UiState) { ui.add_space(2.0); let s = &mut state.sim_settings; - ui.label("Rigid bodies"); - ui.add(egui::Slider::new(&mut s.rbd_steps_per_frame, 1..=20).text("steps / frame")); + if has_mpm { + ui.label("MPM"); + ui.add(egui::Slider::new(&mut s.mpm_substeps, 1..=200).text("substeps")); + ui.checkbox(&mut s.mpm_use_cpic, "Use CPIC") + .on_hover_text("Compatible particle-in-cell coupling with rigid colliders"); + gravity_drag(ui, "gravity", &mut s.mpm_gravity); + // View-only coloring mode (not part of `sim_settings`). + ComboBox::from_label("coloring") + .selected_text(state.mpm_render_mode.text()) + .show_ui(ui, |ui| { + for mode in MpmRenderMode::ALL { + ui.selectable_value(&mut state.mpm_render_mode, *mode, mode.text()); + } + }); + } + + if has_rbd { + if has_mpm { + ui.add_space(4.0); + } + ui.label("Rigid bodies"); + ui.add(egui::Slider::new(&mut s.rbd_steps_per_frame, 1..=20).text("steps / frame")); + } +} + +/// A labelled per-component drag editor for a gravity vector (2D or 3D). +fn gravity_drag(ui: &mut egui::Ui, label: &str, g: &mut nexus::rbd::math::Vector) { + ui.horizontal(|ui| { + ui.label(label); + ui.add(egui::DragValue::new(&mut g.x).speed(0.1).prefix("x ")); + ui.add(egui::DragValue::new(&mut g.y).speed(0.1).prefix("y ")); + #[cfg(feature = "dim3")] + ui.add(egui::DragValue::new(&mut g.z).speed(0.1).prefix("z ")); + }); } fn performance_ui( @@ -236,6 +267,10 @@ fn performance_ui( row("Multibodies:", counts.multibodies); row("Multibody DOFs:", counts.multibody_dofs); } + + if counts.particles > 0 { + row("Particles:", counts.particles); + } }); ui.add_space(8.0); @@ -284,13 +319,13 @@ fn performance_ui( } impl UiState { - /// Order in which demos appear in the picker: grouped by kind, preserving - /// each group's listing order. Prev/Next walks this sequence so it matches - /// the visible list rather than the raw (lexicographically-sorted) `demos` - /// index order. + /// Order in which demos appear in the picker: grouped by kind (Rbd, Mpm), + /// preserving each group's listing order. Prev/Next walks this sequence so it + /// matches the visible list rather than the raw (lexicographically-sorted) + /// `demos` index order. fn demo_display_order(&self) -> Vec { let mut order = Vec::with_capacity(self.demos.len()); - for kind in [DemoKind::Rbd] { + for kind in [DemoKind::Rbd, DemoKind::Mpm] { for (i, (_, k)) in self.demos.iter().enumerate() { if *k == kind { order.push(i); @@ -341,6 +376,7 @@ fn examples_section(ui: &mut egui::Ui, state: &mut UiState) { ui.add_space(4.0); demo_group(ui, state, DemoKind::Rbd, "Rigid Bodies"); + demo_group(ui, state, DemoKind::Mpm, "MPM"); } fn demo_group(ui: &mut egui::Ui, state: &mut UiState, kind: DemoKind, label: &str) { @@ -380,7 +416,8 @@ fn demo_group(ui: &mut egui::Ui, state: &mut UiState, kind: DemoKind, label: &st }); } -/// Unified backend selector. +/// Unified backend selector. The "CPU (rapier)" option is only shown for RBD +/// scenes; for MPM scenes a current Rapier selection is shown as CPU (nexus). fn backend_selector(ui: &mut egui::Ui, state: &mut UiState, gpu_available: bool) { ui.label(RichText::new("Physics Backend").strong()); ui.add_space(2.0); @@ -432,116 +469,3 @@ fn backend_selector(ui: &mut egui::Ui, state: &mut UiState, gpu_available: bool) state.transition = Some(Transition::Switch); } } - -// // =========================================================================== -// // Per-scene UI, implemented through the `Scene` trait. -// // =========================================================================== -// -// impl Scene for crate::rbd::RbdScene { -// fn is_rbd(&self) -> bool { -// true -// } -// -// fn performance_ui(&mut self, ui: &mut egui::Ui, run_stats: &RunStats, backend_type: BackendType) { -// let physics = &self.physics; -// -// // Scene info. -// ui.label(RichText::new("Scene").strong()); -// ui.add_space(2.0); -// -// egui::Grid::new("rbd_scene_grid") -// .num_columns(2) -// .spacing([20.0, 2.0]) -// .show(ui, |ui| { -// ui.label("Bodies:"); -// ui.label(format!("{}", physics.backend.num_bodies())); -// ui.end_row(); -// -// ui.label("Joints:"); -// ui.label(format!("{}", physics.backend.num_joints())); -// ui.end_row(); -// -// ui.label("Batches:"); -// ui.label(format!("{}", physics.backend.num_batches())); -// ui.end_row(); -// }); -// -// ui.add_space(8.0); -// ui.separator(); -// ui.add_space(4.0); -// -// // Timing. -// let total_ms_with_readback = run_stats.total_simulation_time_with_readback_ms(); -// let total_ms_without_readback = run_stats.total_simulation_time_without_readback_ms(); -// let total_readback_time = total_ms_with_readback - total_ms_without_readback; -// let fps = if total_ms_with_readback > 0.0 { -// (1000.0f32 / total_ms_with_readback).round() -// } else { -// 0.0 -// }; -// -// ui.label( -// RichText::new(format!( -// "Total: {:.2}ms (+ readback: {:.2}ms) - {:.0} FPS", -// total_ms_without_readback, total_readback_time, fps -// )) -// .strong(), -// ); -// ui.add_space(4.0); -// -// if !matches!(backend_type, BackendType::Rapier) { -// CollapsingHeader::new("Simulation details") -// .id_salt("rbd_sim_details") -// .default_open(false) -// .show(ui, |ui| { -// ui.label(format!("Colors: {}", run_stats.num_colors)); -// ui.label(format!( -// "Coloring: {:.2}ms", -// run_stats.coloring_time.as_secs_f32() * 1000.0 -// )); -// ui.label(format!( -// "Coloring iterations: {} x 10", -// run_stats.coloring_iterations -// )); -// ui.label(format!( -// "Start to pairs count: {:.2}ms", -// run_stats.start_to_pairs_count_time.as_secs_f32() * 1000.0 -// )); -// ui.label(format!( -// "Coloring fallback: {:.2}ms", -// run_stats.coloring_fallback_time.as_secs_f32() * 1000.0 -// )); -// }); -// -// if !run_stats.gpu_pass_times.is_empty() { -// CollapsingHeader::new(format!("GPU passes: {:.2}ms", run_stats.gpu_total_time)) -// .id_salt("rbd_gpu_passes") -// .default_open(false) -// .show(ui, |ui| { -// egui::Grid::new("rbd_timestamp_grid") -// .num_columns(2) -// .spacing([20.0, 2.0]) -// .show(ui, |ui| { -// for (label, ms) in &run_stats.gpu_pass_times { -// ui.label(format!("{}:", label)); -// ui.label(format!("{:.2}ms", ms)); -// ui.end_row(); -// } -// }); -// }); -// } -// -// // Slow performance warning. -// if run_stats.total_simulation_time_with_readback.as_secs_f32() > 0.1 { -// ui.add_space(4.0); -// ui.colored_label( -// Color32::from_rgb(180, 120, 60), -// #[cfg(not(target_arch = "wasm32"))] -// "Running slow? If you have both an integrated and discrete GPU, ensure the discrete GPU is in use.", -// #[cfg(target_arch = "wasm32")] -// "Running slow? If you have both an integrated and discrete GPU, ensure your browser runs exclusively on the discrete GPU.", -// ); -// } -// } -// } -// } diff --git a/src_viewer/viewer.rs b/src_viewer/viewer.rs index 770ea0c0..c07b8d54 100644 --- a/src_viewer/viewer.rs +++ b/src_viewer/viewer.rs @@ -17,7 +17,8 @@ use glamx::Vec4; use khal::Shader; use khal::backend::{ - Backend, GpuBackend as KhalGpuBackend, GpuBackendError, GpuTimestamps, WebGpu, + Backend, GpuBackend as KhalGpuBackend, GpuBackendError, GpuBufferSliceMut, GpuTimestamps, + WebGpu, }; use khal::re_exports::wgpu::{Features, Limits}; use std::time::Duration; @@ -25,25 +26,75 @@ use std::time::Duration; use kiss3d::prelude::Color; #[cfg(feature = "dim3")] use kiss3d::renderer::RayTracer; +#[cfg(feature = "dim2")] +use kiss3d::scene::InstanceData2d; +#[cfg(feature = "dim3")] +use kiss3d::scene::InstanceData3d; use kiss3d::scene::{SceneNode2d, SceneNode3d}; use kiss3d::window::{NumSamples, Window}; +/// Viewer-owned scene node type for the active dimension. +#[cfg(feature = "dim2")] +type SceneNodeX = SceneNode2d; +#[cfg(feature = "dim3")] +type SceneNodeX = SceneNode3d; + #[cfg(feature = "dim3")] use kiss3d::camera::{FixedView2d, OrbitCamera3d}; #[cfg(feature = "dim2")] use kiss3d::camera::{FixedView3d, PanZoomCamera2d}; +use nexus::mpm::solver::prep_readback::{ + GpuReadbackData, ReadbackData, RenderConfig, WgPrepReadback, +}; use nexus::rbd::dynamics::WgRbdPrepRender; -use nexus::rbd::math::Pose; +use nexus::rbd::math::{Pose, Vector}; use nexus::rbd::pipeline::RunStats; use nexus::state::{NexusCounts, NexusState}; use rapier::prelude::{RigidBodyHandle, SharedShape}; -// use crate::rbd::{ -// BackendType, RbdScene, RenderContext, SimulationState, setup_physics, -// }; + use crate::backend::BackendType; use crate::graphics::RenderContext; use crate::{DemoKind, RunState, Transition, UiSections}; +/// Per-particle coloring mode for MPM rendering, written into the +/// `WgPrepReadback` render config. Mirrors the `mode` values understood by the +/// `gpu_prep_readback` shader. +#[derive(Copy, Clone, PartialEq, Eq, Debug, Default)] +pub enum MpmRenderMode { + #[default] + Default = 0, + Volume = 1, + Velocity = 2, + Phase = 3, + CdfNormals = 4, + CdfDistances = 5, + CdfSigns = 6, +} + +impl MpmRenderMode { + pub const ALL: &'static [MpmRenderMode] = &[ + Self::Default, + Self::Volume, + Self::Velocity, + Self::Phase, + Self::CdfNormals, + Self::CdfDistances, + Self::CdfSigns, + ]; + + pub fn text(self) -> &'static str { + match self { + Self::Default => "default", + Self::Volume => "volume", + Self::Velocity => "velocity", + Self::Phase => "phase", + Self::CdfNormals => "cdf (normals)", + Self::CdfDistances => "cdf (distances)", + Self::CdfSigns => "cdf (signs)", + } + } +} + /// UI / runtime state that is independent from the GPU/window resources. Kept in /// its own struct so [`NexusViewer::render_frame`] can split-borrow it from `window`. pub struct UiState { @@ -67,9 +118,13 @@ pub struct UiState { /// otherwise — so restarts and backend switches keep the user's edits. pub(crate) settings_demo: Option, /// Which sub-systems the current scene contains (drives which settings show). + pub(crate) has_mpm: bool, pub(crate) has_rbd: bool, /// Current scene entity counts, refreshed every `sync` for the UI. pub(crate) counts: NexusCounts, + /// Per-particle coloring mode for MPM rendering (view-only; not a sim + /// setting). Drives the `WgPrepReadback` render config in `sync`. + pub mpm_render_mode: MpmRenderMode, } /// Editable simulation settings exposed in the viewer UI. The viewer pulls @@ -77,6 +132,12 @@ pub struct UiState { /// frame (see [`NexusViewer::sync`]). #[derive(Clone)] pub struct SimSettings { + /// MPM substeps per rendered frame. + pub mpm_substeps: u32, + /// MPM CPIC rigid-body coupling toggle. + pub mpm_use_cpic: bool, + /// Gravity applied to the MPM particles. + pub mpm_gravity: Vector, /// Rigid-body solver steps advanced per rendered frame. pub rbd_steps_per_frame: u32, } @@ -84,6 +145,9 @@ pub struct SimSettings { impl Default for SimSettings { fn default() -> Self { Self { + mpm_substeps: 20, + mpm_use_cpic: true, + mpm_gravity: Vector::ZERO, rbd_steps_per_frame: 1, } } @@ -135,10 +199,25 @@ pub struct NexusViewer { /// instance buffers on the shared-device (WebGPU) path. Compiled lazily on /// the first direct sync; reused across demos. rbd_prep_render: Option, - /// Backend the cached render-prep shaders/buffers (`rbd_prep_render`, …) - /// were built for. When the active backend changes, those resources belong - /// to a different device and must be dropped and rebuilt — otherwise using - /// them crashes. + /// Viewer-owned point-cloud node for MPM particles (lazily created in `sync`). + mpm_node: Option, + /// GPU kernel that turns raw MPM particle state into per-particle render data + /// (deformed position + mode-dependent color). Compiled lazily on the first + /// MPM frame; reused across demos. + mpm_readback: Option, + /// Output/staging buffers for [`Self::mpm_readback`], sized to the current + /// particle count. Recreated when the particle count changes (emitters). + mpm_readback_data: Option, + /// `(num_particles, num_rigid_particles)` the `mpm_readback_data` buffers + /// are sized for. A mismatch means they need reallocating. + mpm_readback_counts: (usize, usize), + /// Palette indexed by each particle's group id, set by the scene through + /// [`Self::set_particle_group_colors`]. Empty means the built-in palette. + mpm_group_colors: Vec, + /// Backend the cached render-prep shaders/buffers (`rbd_prep_render`, + /// `mpm_readback*`, …) were built for. When the active + /// backend changes, those resources belong to a different device and must be + /// dropped and rebuilt, otherwise using them crashes. render_resources_backend: Option, /// Last GPU pass timings harvested from the non-blocking timestamp readback. /// Re-applied to `run_stats` every frame so the profiler UI keeps showing the @@ -243,6 +322,11 @@ impl NexusViewer { }, nexus_render: RenderContext::new(), rbd_prep_render: None, + mpm_node: None, + mpm_readback: None, + mpm_readback_data: None, + mpm_readback_counts: (0, 0), + mpm_group_colors: Vec::new(), render_resources_backend: None, last_gpu_pass_times: Vec::new(), last_gpu_total_time_ms: 0.0, @@ -265,8 +349,10 @@ impl NexusViewer { transition: None, sim_settings: SimSettings::default(), settings_demo: None, + has_mpm: false, has_rbd: false, counts: NexusCounts::default(), + mpm_render_mode: MpmRenderMode::default(), }, }; @@ -635,6 +721,87 @@ impl NexusViewer { } } + if let Some(mpm) = state.mpm.as_ref() { + let num_particles = mpm.particles.len(); + if num_particles > 0 { + let backend = self.backend().clone(); + let num_rigid = mpm.rigid_particles.len() as usize; + let mode = self.ui.mpm_render_mode as u32; + + // Lazily compile the readback kernel (reused across demos). + if self.mpm_readback.is_none() { + self.mpm_readback = WgPrepReadback::from_backend(&backend).ok(); + } + // (Re)allocate the readback buffers when the particle count + // changes (emitters growing the particle set, for instance). + if self.mpm_readback_data.is_none() + || self.mpm_readback_counts != (num_particles, num_rigid) + { + self.mpm_readback_data = GpuReadbackData::new( + &backend, + num_particles, + num_rigid, + mode, + &self.mpm_group_colors, + ) + .ok(); + self.mpm_readback_counts = (num_particles, num_rigid); + } + + let instances = if let (Some(shader), Some(readback)) = + (self.mpm_readback.as_ref(), self.mpm_readback_data.as_mut()) + { + // Push the current coloring mode (cheap; applies UI switches). + let _ = backend.write_buffer( + readback.mode.buffer_mut(), + 0, + &[RenderConfig { + mode, + num_groups: readback.num_groups, + ..Default::default() + }], + ); + let mut enc = backend.begin_encoding(); + let launched = shader + .launch( + &mut enc, + None, + readback, + &mpm.sim_params, + &mpm.grid, + &mpm.particles, + &mpm.rigid_particles, + ) + .is_ok(); + if launched && backend.submit(enc).is_ok() { + let _ = backend.synchronize(); + let mut v = vec![ReadbackData::default(); num_particles]; + if backend + .read_buffer(readback.instances_staging.buffer(), v.as_mut_slice()) + .await + .is_ok() + { + Some(v) + } else { + None + } + } else { + None + } + } else { + None + }; + + if let Some(instances) = instances { + if self.mpm_node.is_none() { + self.mpm_node = Some(self.new_point_node()); + } + let data = build_mpm_instances(&instances); + self.mpm_node.as_mut().unwrap().set_instances(&data); + } + } + } + Ok(()) } async fn sync_without_readback( @@ -670,11 +837,80 @@ impl NexusViewer { } } + if let Some(mpm) = state.mpm.as_ref() { + let num_particles = mpm.particles.len(); + if num_particles > 0 { + let backend = self.backend().clone(); + let num_rigid = mpm.rigid_particles.len() as usize; + let mode = self.ui.mpm_render_mode as u32; + + // Lazily compile the readback kernel (reused across demos). + if self.mpm_readback.is_none() { + self.mpm_readback = WgPrepReadback::from_backend(&backend).ok(); + } + // (Re)allocate the readback buffers when the particle count + // changes (emitters growing the particle set, for instance). + if self.mpm_readback_data.is_none() + || self.mpm_readback_counts != (num_particles, num_rigid) + { + self.mpm_readback_data = GpuReadbackData::new( + &backend, + num_particles, + num_rigid, + mode, + &self.mpm_group_colors, + ) + .ok(); + self.mpm_readback_counts = (num_particles, num_rigid); + } + + // Zero-readback: a compute kernel writes per-particle render + // data straight into the point node's GPU instance buffers. + if self.mpm_node.is_none() { + self.mpm_node = Some(self.new_point_node()); + } + let bufs = self + .mpm_node + .as_mut() + .unwrap() + .instance_compute_buffers(num_particles); + if let (Some(shader), Some(readback)) = + (self.mpm_readback.as_ref(), self.mpm_readback_data.as_mut()) + { + let _ = backend.write_buffer( + readback.mode.buffer_mut(), + 0, + &[RenderConfig { + mode, + num_groups: readback.num_groups, + ..Default::default() + }], + ); + let mut positions = GpuBufferSliceMut::::from_wgpu(&bufs.positions); + let mut deformations = GpuBufferSliceMut::::from_wgpu(&bufs.deformations); + let mut colors = GpuBufferSliceMut::::from_wgpu(&bufs.colors); + let mut enc = backend.begin_encoding(); + let _ = shader.launch_render( + &mut enc, + &mut positions, + &mut deformations, + &mut colors, + readback, + &mpm.sim_params, + &mpm.grid, + &mpm.particles, + ); + let _ = backend.submit(enc); + } + } + } + Ok(()) } /// Reads the latest state from a [`NexusState`] back from the GPU and pushes - /// it into the viewer-owned render instances (rigid-body collider poses). + /// it into the viewer-owned render instances: rigid-body collider poses and + /// the MPM particle point cloud. pub async fn sync( &mut self, state: &mut NexusState, @@ -701,10 +937,17 @@ impl NexusViewer { // current settings and push them back into the freshly-built scene. if self.ui.settings_demo != Some(self.ui.selected_demo) { self.ui.has_rbd = state.rbd.is_some(); + self.ui.has_mpm = state.has_mpm(); + self.ui.sim_settings.mpm_substeps = state.mpm_substeps(); + self.ui.sim_settings.mpm_use_cpic = state.mpm_use_cpic(); + self.ui.sim_settings.mpm_gravity = state.mpm_gravity(); self.ui.sim_settings.rbd_steps_per_frame = state.rbd_steps_per_frame(); self.ui.settings_demo = Some(self.ui.selected_demo); } else { let s = self.ui.sim_settings.clone(); + state.set_mpm_substeps(s.mpm_substeps); + state.set_mpm_use_cpic(s.mpm_use_cpic); + state.set_mpm_gravity(s.mpm_gravity); state.set_rbd_steps_per_frame(s.rbd_steps_per_frame); } @@ -731,6 +974,17 @@ impl NexusViewer { Ok(()) } + /// Creates a unit point-cloud base node (a cube in 3D, a rectangle in 2D) + /// that subsequent per-particle/per-vertex instances are drawn from. + fn new_point_node(&mut self) -> SceneNodeX { + #[cfg(feature = "dim2")] + let mut node = self.scene2d.add_rectangle(1.0, 1.0); + #[cfg(feature = "dim3")] + let mut node = self.scene3d.add_cube(1.0, 1.0, 1.0); + node.enable_backface_culling_recursive(false); + node + } + /// The backend currently selected in the UI. pub fn backend_type(&self) -> BackendType { self.ui.backend_type @@ -769,6 +1023,21 @@ impl NexusViewer { /// those resources belong to the previous backend's device. fn invalidate_render_resources(&mut self) { self.rbd_prep_render = None; + self.mpm_readback = None; + self.mpm_readback_data = None; + self.mpm_readback_counts = (0, 0); + } + + /// Sets the color palette that MPM particle group ids index into. + /// + /// A particle's group is set with [`Particle::with_group`](nexus::mpm::solver::Particle::with_group); + /// ids past the end of the palette wrap around. Passing an empty slice + /// restores the built-in palette. + pub fn set_particle_group_colors(&mut self, colors: &[Vec4]) { + self.mpm_group_colors = colors.to_vec(); + // Force a reallocation so the new palette reaches the GPU. + self.mpm_readback_data = None; + self.mpm_readback_counts = (0, 0); } /// Tears down the viewer-owned `NexusState` render nodes. A no-op for legacy @@ -777,6 +1046,7 @@ impl NexusViewer { self.scene3d = SceneNode3d::empty(); self.scene2d = SceneNode2d::empty(); self.nexus_render.clear(); + self.mpm_node = None; } /// Whether the simulation should advance this frame, honoring the @@ -958,3 +1228,33 @@ impl NexusViewer { self.window.draw_ui(ui_fn); } } + +/// Builds MPM particle render instances from `WgPrepReadback` output: each +/// particle's position, mode-dependent color, and deformation transform (so +/// stretched/sheared particles render deformed rather than as fixed dots). +#[cfg(feature = "dim3")] +fn build_mpm_instances(instances: &[ReadbackData]) -> Vec { + use nexus::mpm::mpm_shaders::PaddingExt; + instances + .iter() + .map(|d| InstanceData3d { + position: d.position, + color: Color::new(d.color.x, d.color.y, d.color.z, d.color.w), + deformation: d.deformation.remove_padding(), + ..Default::default() + }) + .collect() +} + +#[cfg(feature = "dim2")] +fn build_mpm_instances(instances: &[ReadbackData]) -> Vec { + instances + .iter() + .map(|d| InstanceData2d { + position: d.position, + color: [d.color.x, d.color.y, d.color.z, d.color.w], + deformation: d.deformation, + ..Default::default() + }) + .collect() +} From ab92393abd096ecf773e6293abdd991876743ee9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 14 Aug 2026 23:47:14 +0200 Subject: [PATCH 4/7] chore: prefix the demo file names with their subsystem --- crates/examples2d/{balls2.rs => rbd_balls2.rs} | 11 ++++++----- crates/examples2d/{boxes2.rs => rbd_boxes2.rs} | 11 ++++++----- .../{boxes_and_balls2.rs => rbd_boxes_and_balls2.rs} | 11 ++++++----- crates/examples2d/{compound2.rs => rbd_compound2.rs} | 11 ++++++----- .../examples2d/{dynamic_rbd2.rs => rbd_dynamic2.rs} | 11 ++++++----- .../examples2d/{joint_ball2.rs => rbd_joint_ball2.rs} | 5 +++-- .../{joint_fixed2.rs => rbd_joint_fixed2.rs} | 5 +++-- .../{joint_prismatic2.rs => rbd_joint_prismatic2.rs} | 7 ++++--- crates/examples2d/{polyline2.rs => rbd_polyline2.rs} | 7 ++++--- .../examples2d/{primitives2.rs => rbd_primitives2.rs} | 11 ++++++----- crates/examples2d/{pyramid2.rs => rbd_pyramid2.rs} | 7 ++++--- crates/examples3d/{balls3.rs => rbd_balls3.rs} | 7 ++++--- crates/examples3d/{boxes3.rs => rbd_boxes3.rs} | 7 ++++--- .../{boxes_and_balls3.rs => rbd_boxes_and_balls3.rs} | 7 ++++--- crates/examples3d/{compound3.rs => rbd_compound3.rs} | 11 ++++++----- .../examples3d/{dynamic_rbd3.rs => rbd_dynamic3.rs} | 11 ++++++----- .../examples3d/{joint_ball3.rs => rbd_joint_ball3.rs} | 7 ++++--- .../{joint_fixed3.rs => rbd_joint_fixed3.rs} | 6 ++++-- .../{joint_prismatic3.rs => rbd_joint_prismatic3.rs} | 7 ++++--- .../{joint_revolute3.rs => rbd_joint_revolute3.rs} | 7 ++++--- ...evolute_batch3.rs => rbd_joint_revolute_batch3.rs} | 8 +++++--- crates/examples3d/{joints3.rs => rbd_joints3.rs} | 4 ++-- crates/examples3d/{keva3.rs => rbd_keva3.rs} | 4 ++-- .../{many_pyramids3.rs => rbd_many_pyramids3.rs} | 4 ++-- ...pyramids_batch3.rs => rbd_many_pyramids_batch3.rs} | 7 ++++--- ...{mujoco_menagerie3.rs => rbd_mujoco_menagerie3.rs} | 10 +++++----- ...tibody_pendulum3.rs => rbd_multibody_pendulum3.rs} | 11 ++++++----- .../examples3d/{primitives3.rs => rbd_primitives3.rs} | 7 ++++--- crates/examples3d/{pyramid3.rs => rbd_pyramid3.rs} | 4 ++-- crates/examples3d/{trimesh3.rs => rbd_trimesh3.rs} | 7 ++++--- crates/examples3d/{urdf3.rs => rbd_urdf3.rs} | 7 ++++--- 31 files changed, 134 insertions(+), 106 deletions(-) rename crates/examples2d/{balls2.rs => rbd_balls2.rs} (88%) rename crates/examples2d/{boxes2.rs => rbd_boxes2.rs} (88%) rename crates/examples2d/{boxes_and_balls2.rs => rbd_boxes_and_balls2.rs} (88%) rename crates/examples2d/{compound2.rs => rbd_compound2.rs} (91%) rename crates/examples2d/{dynamic_rbd2.rs => rbd_dynamic2.rs} (91%) rename crates/examples2d/{joint_ball2.rs => rbd_joint_ball2.rs} (95%) rename crates/examples2d/{joint_fixed2.rs => rbd_joint_fixed2.rs} (95%) rename crates/examples2d/{joint_prismatic2.rs => rbd_joint_prismatic2.rs} (94%) rename crates/examples2d/{polyline2.rs => rbd_polyline2.rs} (93%) rename crates/examples2d/{primitives2.rs => rbd_primitives2.rs} (91%) rename crates/examples2d/{pyramid2.rs => rbd_pyramid2.rs} (91%) rename crates/examples3d/{balls3.rs => rbd_balls3.rs} (94%) rename crates/examples3d/{boxes3.rs => rbd_boxes3.rs} (95%) rename crates/examples3d/{boxes_and_balls3.rs => rbd_boxes_and_balls3.rs} (95%) rename crates/examples3d/{compound3.rs => rbd_compound3.rs} (91%) rename crates/examples3d/{dynamic_rbd3.rs => rbd_dynamic3.rs} (93%) rename crates/examples3d/{joint_ball3.rs => rbd_joint_ball3.rs} (97%) rename crates/examples3d/{joint_fixed3.rs => rbd_joint_fixed3.rs} (93%) rename crates/examples3d/{joint_prismatic3.rs => rbd_joint_prismatic3.rs} (94%) rename crates/examples3d/{joint_revolute3.rs => rbd_joint_revolute3.rs} (95%) rename crates/examples3d/{joint_revolute_batch3.rs => rbd_joint_revolute_batch3.rs} (95%) rename crates/examples3d/{joints3.rs => rbd_joints3.rs} (99%) rename crates/examples3d/{keva3.rs => rbd_keva3.rs} (98%) rename crates/examples3d/{many_pyramids3.rs => rbd_many_pyramids3.rs} (97%) rename crates/examples3d/{many_pyramids_batch3.rs => rbd_many_pyramids_batch3.rs} (96%) rename crates/examples3d/{mujoco_menagerie3.rs => rbd_mujoco_menagerie3.rs} (98%) rename crates/examples3d/{multibody_pendulum3.rs => rbd_multibody_pendulum3.rs} (92%) rename crates/examples3d/{primitives3.rs => rbd_primitives3.rs} (96%) rename crates/examples3d/{pyramid3.rs => rbd_pyramid3.rs} (96%) rename crates/examples3d/{trimesh3.rs => rbd_trimesh3.rs} (96%) rename crates/examples3d/{urdf3.rs => rbd_urdf3.rs} (96%) diff --git a/crates/examples2d/balls2.rs b/crates/examples2d/rbd_balls2.rs similarity index 88% rename from crates/examples2d/balls2.rs rename to crates/examples2d/rbd_balls2.rs index 8977e52b..84694b56 100644 --- a/crates/examples2d/balls2.rs +++ b/crates/examples2d/rbd_balls2.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer2d::NexusViewer; -use nexus2d::prelude::{NexusCapacities, NexusPipeline, NexusState}; +use nexus2d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; use rapier2d::prelude::*; pub async fn run( @@ -9,6 +9,7 @@ pub async fn run( ) -> anyhow::Result { let capacities = NexusCapacities::default().rbd_collisions(250_000); let mut state = NexusState::new(capacities); + let no_coupling = RbdCoupling::None; /* * Ground @@ -18,7 +19,7 @@ pub async fn run( let body = RigidBodyBuilder::fixed().build(); let collider = ColliderBuilder::cuboid(ground_size, 1.5).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); let body = RigidBodyBuilder::fixed() @@ -27,7 +28,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::cuboid(ground_size * 1.2, 1.5).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); let body = RigidBodyBuilder::fixed() @@ -36,7 +37,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::cuboid(ground_size * 1.2, 1.5).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); /* @@ -60,7 +61,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::ball(rad).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); } } diff --git a/crates/examples2d/boxes2.rs b/crates/examples2d/rbd_boxes2.rs similarity index 88% rename from crates/examples2d/boxes2.rs rename to crates/examples2d/rbd_boxes2.rs index 7d13db90..31944557 100644 --- a/crates/examples2d/boxes2.rs +++ b/crates/examples2d/rbd_boxes2.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer2d::NexusViewer; -use nexus2d::prelude::{NexusCapacities, NexusPipeline, NexusState}; +use nexus2d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; use rapier2d::prelude::*; pub async fn run( @@ -9,6 +9,7 @@ pub async fn run( ) -> anyhow::Result { let capacities = NexusCapacities::default().rbd_collisions(250_000); let mut state = NexusState::new(capacities); + let no_coupling = RbdCoupling::None; /* * Ground @@ -18,7 +19,7 @@ pub async fn run( let body = RigidBodyBuilder::fixed().build(); let collider = ColliderBuilder::cuboid(ground_size, 1.5).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); let body = RigidBodyBuilder::fixed() @@ -27,7 +28,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::cuboid(ground_size * 1.2, 1.5).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); let body = RigidBodyBuilder::fixed() @@ -36,7 +37,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::cuboid(ground_size * 1.2, 1.5).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); /* @@ -59,7 +60,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::cuboid(rad, rad).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); } } diff --git a/crates/examples2d/boxes_and_balls2.rs b/crates/examples2d/rbd_boxes_and_balls2.rs similarity index 88% rename from crates/examples2d/boxes_and_balls2.rs rename to crates/examples2d/rbd_boxes_and_balls2.rs index 340c0097..880d7117 100644 --- a/crates/examples2d/boxes_and_balls2.rs +++ b/crates/examples2d/rbd_boxes_and_balls2.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer2d::NexusViewer; -use nexus2d::prelude::{NexusCapacities, NexusPipeline, NexusState}; +use nexus2d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; use rapier2d::prelude::*; pub async fn run( @@ -9,6 +9,7 @@ pub async fn run( ) -> anyhow::Result { let capacities = NexusCapacities::default().rbd_collisions(250_000); let mut state = NexusState::new(capacities); + let no_coupling = RbdCoupling::None; /* * Ground @@ -18,7 +19,7 @@ pub async fn run( let body = RigidBodyBuilder::fixed().build(); let collider = ColliderBuilder::cuboid(ground_size, 1.5).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); let body = RigidBodyBuilder::fixed() @@ -27,7 +28,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::cuboid(ground_size * 1.2, 1.5).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); let body = RigidBodyBuilder::fixed() @@ -36,7 +37,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::cuboid(ground_size * 1.2, 1.5).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); /* @@ -64,7 +65,7 @@ pub async fn run( } .build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); } } diff --git a/crates/examples2d/compound2.rs b/crates/examples2d/rbd_compound2.rs similarity index 91% rename from crates/examples2d/compound2.rs rename to crates/examples2d/rbd_compound2.rs index 34a79027..ab21d5ee 100644 --- a/crates/examples2d/compound2.rs +++ b/crates/examples2d/rbd_compound2.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer2d::NexusViewer; -use nexus2d::prelude::{NexusCapacities, NexusPipeline, NexusState}; +use nexus2d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; use rapier2d::prelude::*; /// 2D analogue of rapier's `compound3` demo. @@ -10,6 +10,7 @@ pub async fn run( ) -> anyhow::Result { let capacities = NexusCapacities::default().rbd_collisions(40_000); let mut state = NexusState::new(capacities); + let no_coupling = RbdCoupling::None; /* * Ground @@ -19,7 +20,7 @@ pub async fn run( let body = RigidBodyBuilder::fixed().build(); let collider = ColliderBuilder::cuboid(ground_size, 1.5).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); let body = RigidBodyBuilder::fixed() @@ -28,7 +29,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::cuboid(ground_size * 2.1, 1.5).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); let body = RigidBodyBuilder::fixed() @@ -37,7 +38,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::cuboid(ground_size * 2.1, 1.5).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); /* @@ -79,7 +80,7 @@ pub async fn run( let base = ColliderBuilder::cuboid(base_he.x, base_he.y) .translation(base_offset) .build(); - let handle = state.insert_rigid_body(body, base); + let handle = state.insert_rigid_body(body, base, no_coupling); for (offset, he) in &parts[1..] { let collider = ColliderBuilder::cuboid(he.x, he.y) .translation(*offset) diff --git a/crates/examples2d/dynamic_rbd2.rs b/crates/examples2d/rbd_dynamic2.rs similarity index 91% rename from crates/examples2d/dynamic_rbd2.rs rename to crates/examples2d/rbd_dynamic2.rs index 25459952..af00d80b 100644 --- a/crates/examples2d/dynamic_rbd2.rs +++ b/crates/examples2d/rbd_dynamic2.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer2d::NexusViewer; -use nexus2d::prelude::{NexusPipeline, NexusState}; +use nexus2d::prelude::{NexusPipeline, NexusState, RbdCoupling}; use rapier2d::prelude::*; /// Demonstrates adding rigid-bodies to a live scene WITHOUT rebuilding the whole @@ -25,6 +25,7 @@ pub async fn run( const ROW: usize = 50; let mut state = NexusState::default(); + let no_coupling = RbdCoupling::None; /* * A boxed ground: a floor plus two side walls to keep the pile contained. @@ -42,7 +43,7 @@ pub async fn run( .translation(pos) .build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::from_translation(pos)); } @@ -58,8 +59,8 @@ pub async fn run( while viewer.render_frame().await { if viewer.simulating() { - // Drop in a whole batch of bodies periodically — all inserted with a - // single batched append, without rebuilding the scene. + // Drop in a whole batch of bodies periodically, all inserted with a + // single batched append and without rebuilding the scene. if frame.is_multiple_of(SPAWN_PERIOD) && added < MAX_BODIES { let n = BODIES_PER_SPAWN.min(MAX_BODIES - added); let mut batch = Vec::with_capacity(n); @@ -82,7 +83,7 @@ pub async fn run( ColliderBuilder::cuboid(0.4, 0.4).build() }; shapes.push(collider.shared_shape().clone()); - batch.push((body, collider)); + batch.push((body, collider, no_coupling)); } let handles = state.add_rigid_bodies(viewer.backend(), batch)?; diff --git a/crates/examples2d/joint_ball2.rs b/crates/examples2d/rbd_joint_ball2.rs similarity index 95% rename from crates/examples2d/joint_ball2.rs rename to crates/examples2d/rbd_joint_ball2.rs index 2096164e..3e47d368 100644 --- a/crates/examples2d/joint_ball2.rs +++ b/crates/examples2d/rbd_joint_ball2.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer2d::NexusViewer; -use nexus2d::prelude::{NexusPipeline, NexusState}; +use nexus2d::prelude::{NexusPipeline, NexusState, RbdCoupling}; use rapier2d::prelude::*; pub async fn run( @@ -8,6 +8,7 @@ pub async fn run( pipeline: &mut NexusPipeline, ) -> anyhow::Result { let mut state = NexusState::default(); + let no_coupling = RbdCoupling::None; /* * Create the balls @@ -35,7 +36,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::ball(rad).build(); let shape = collider.shared_shape().clone(); - let child_handle = state.insert_rigid_body(body, collider); + let child_handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(child_handle, &shape, Pose::IDENTITY); // Vertical joint. diff --git a/crates/examples2d/joint_fixed2.rs b/crates/examples2d/rbd_joint_fixed2.rs similarity index 95% rename from crates/examples2d/joint_fixed2.rs rename to crates/examples2d/rbd_joint_fixed2.rs index 0ee384db..8b881ffb 100644 --- a/crates/examples2d/joint_fixed2.rs +++ b/crates/examples2d/rbd_joint_fixed2.rs @@ -1,7 +1,7 @@ use glamx::Pose2; use khal::backend::GpuTimestamps; use nexus_viewer2d::NexusViewer; -use nexus2d::prelude::{NexusPipeline, NexusState}; +use nexus2d::prelude::{NexusPipeline, NexusState, RbdCoupling}; use rapier2d::prelude::*; pub async fn run( @@ -9,6 +9,7 @@ pub async fn run( pipeline: &mut NexusPipeline, ) -> anyhow::Result { let mut state = NexusState::default(); + let no_coupling = RbdCoupling::None; /* * Create the balls @@ -41,7 +42,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::ball(rad).build(); let shape = collider.shared_shape().clone(); - let child_handle = state.insert_rigid_body(body, collider); + let child_handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(child_handle, &shape, Pose::IDENTITY); // Vertical joint. diff --git a/crates/examples2d/joint_prismatic2.rs b/crates/examples2d/rbd_joint_prismatic2.rs similarity index 94% rename from crates/examples2d/joint_prismatic2.rs rename to crates/examples2d/rbd_joint_prismatic2.rs index a1d2cdd6..bab05dfe 100644 --- a/crates/examples2d/joint_prismatic2.rs +++ b/crates/examples2d/rbd_joint_prismatic2.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer2d::NexusViewer; -use nexus2d::prelude::{NexusPipeline, NexusState}; +use nexus2d::prelude::{NexusPipeline, NexusState, RbdCoupling}; use rapier2d::prelude::*; pub async fn run( @@ -8,6 +8,7 @@ pub async fn run( pipeline: &mut NexusPipeline, ) -> anyhow::Result { let mut state = NexusState::default(); + let no_coupling = RbdCoupling::None; /* * Create the balls @@ -27,7 +28,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::cuboid(rad, rad).build(); let shape = collider.shared_shape().clone(); - let mut curr_parent = state.insert_rigid_body(body, collider); + let mut curr_parent = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(curr_parent, &shape, Pose::IDENTITY); for i in 0..num { @@ -38,7 +39,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::cuboid(rad, rad).density(density).build(); let shape = collider.shared_shape().clone(); - let curr_child = state.insert_rigid_body(body, collider); + let curr_child = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(curr_child, &shape, Pose::IDENTITY); let axis = if i % 2 == 0 { diff --git a/crates/examples2d/polyline2.rs b/crates/examples2d/rbd_polyline2.rs similarity index 93% rename from crates/examples2d/polyline2.rs rename to crates/examples2d/rbd_polyline2.rs index a89f536b..f5543e7c 100644 --- a/crates/examples2d/polyline2.rs +++ b/crates/examples2d/rbd_polyline2.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer2d::NexusViewer; -use nexus2d::prelude::{NexusPipeline, NexusState}; +use nexus2d::prelude::{NexusPipeline, NexusState, RbdCoupling}; use rapier2d::prelude::*; pub async fn run( @@ -8,6 +8,7 @@ pub async fn run( pipeline: &mut NexusPipeline, ) -> anyhow::Result { let mut state = NexusState::default(); + let no_coupling = RbdCoupling::None; /* * Ground @@ -28,7 +29,7 @@ pub async fn run( let body = RigidBodyBuilder::fixed().build(); let collider = ColliderBuilder::polyline(points, None).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); // Create 5 predefined convex polygon shapes (so we can render @@ -85,7 +86,7 @@ pub async fn run( } .build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); } } diff --git a/crates/examples2d/primitives2.rs b/crates/examples2d/rbd_primitives2.rs similarity index 91% rename from crates/examples2d/primitives2.rs rename to crates/examples2d/rbd_primitives2.rs index fb1c00ba..680c0f11 100644 --- a/crates/examples2d/primitives2.rs +++ b/crates/examples2d/rbd_primitives2.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer2d::NexusViewer; -use nexus2d::prelude::{NexusCapacities, NexusPipeline, NexusState}; +use nexus2d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; use rapier2d::prelude::*; pub async fn run( @@ -9,6 +9,7 @@ pub async fn run( ) -> anyhow::Result { let capacities = NexusCapacities::default().rbd_collisions(250_000); let mut state = NexusState::new(capacities); + let no_coupling = RbdCoupling::None; /* * Ground @@ -18,7 +19,7 @@ pub async fn run( let body = RigidBodyBuilder::fixed().build(); let collider = ColliderBuilder::cuboid(ground_size, 1.5).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); let body = RigidBodyBuilder::fixed() @@ -27,7 +28,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::cuboid(ground_size * 1.2, 1.5).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); let body = RigidBodyBuilder::fixed() @@ -36,7 +37,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::cuboid(ground_size * 1.2, 1.5).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); /* @@ -93,7 +94,7 @@ pub async fn run( } .build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); } } diff --git a/crates/examples2d/pyramid2.rs b/crates/examples2d/rbd_pyramid2.rs similarity index 91% rename from crates/examples2d/pyramid2.rs rename to crates/examples2d/rbd_pyramid2.rs index bab7d155..12ede5ff 100644 --- a/crates/examples2d/pyramid2.rs +++ b/crates/examples2d/rbd_pyramid2.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer2d::NexusViewer; -use nexus2d::prelude::{NexusPipeline, NexusState}; +use nexus2d::prelude::{NexusPipeline, NexusState, RbdCoupling}; use rapier2d::prelude::*; pub async fn run( @@ -8,6 +8,7 @@ pub async fn run( pipeline: &mut NexusPipeline, ) -> anyhow::Result { let mut state = NexusState::default(); + let no_coupling = RbdCoupling::None; /* * Ground @@ -18,7 +19,7 @@ pub async fn run( let body = RigidBodyBuilder::fixed().build(); let collider = ColliderBuilder::cuboid(ground_size, ground_thickness).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); /* @@ -47,7 +48,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::cuboid(rad, rad).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); } } diff --git a/crates/examples3d/balls3.rs b/crates/examples3d/rbd_balls3.rs similarity index 94% rename from crates/examples3d/balls3.rs rename to crates/examples3d/rbd_balls3.rs index 94a0598b..266157e6 100644 --- a/crates/examples3d/balls3.rs +++ b/crates/examples3d/rbd_balls3.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; use rapier3d::prelude::*; pub async fn run( @@ -12,6 +12,7 @@ pub async fn run( let capacities = NexusCapacities::default().rbd_collisions(300_000); let mut state = NexusState::new(capacities); + let no_coupling = RbdCoupling::None; /* * Floor made of large cuboids. @@ -40,7 +41,7 @@ pub async fn run( .translation(wall_pos) .build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape_with_color( handle, &shape, @@ -65,7 +66,7 @@ pub async fn run( let body = RigidBodyBuilder::dynamic().translation(pos).build(); let collider = ColliderBuilder::ball(0.5).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider.clone()); + let handle = state.insert_rigid_body(body, collider.clone(), no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); } } diff --git a/crates/examples3d/boxes3.rs b/crates/examples3d/rbd_boxes3.rs similarity index 95% rename from crates/examples3d/boxes3.rs rename to crates/examples3d/rbd_boxes3.rs index 3646ebfa..75b3c6c6 100644 --- a/crates/examples3d/boxes3.rs +++ b/crates/examples3d/rbd_boxes3.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; use rapier3d::prelude::*; pub async fn run( @@ -12,6 +12,7 @@ pub async fn run( let capacities = NexusCapacities::default().rbd_collisions(500_000); let mut state = NexusState::new(capacities); + let no_coupling = RbdCoupling::None; /* * Falling dynamic objects. @@ -28,7 +29,7 @@ pub async fn run( let body = RigidBodyBuilder::dynamic().translation(pos).build(); let collider = ColliderBuilder::cuboid(0.5, 0.5, 0.5).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); } } @@ -61,7 +62,7 @@ pub async fn run( .translation(wall_pos) .build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape_with_color( handle, &shape, diff --git a/crates/examples3d/boxes_and_balls3.rs b/crates/examples3d/rbd_boxes_and_balls3.rs similarity index 95% rename from crates/examples3d/boxes_and_balls3.rs rename to crates/examples3d/rbd_boxes_and_balls3.rs index 36cd93da..242906be 100644 --- a/crates/examples3d/boxes_and_balls3.rs +++ b/crates/examples3d/rbd_boxes_and_balls3.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; use rapier3d::prelude::*; pub async fn run( @@ -12,6 +12,7 @@ pub async fn run( let capacities = NexusCapacities::default().rbd_collisions(400_000); let mut state = NexusState::new(capacities); + let no_coupling = RbdCoupling::None; /* * Falling dynamic objects. @@ -33,7 +34,7 @@ pub async fn run( } .build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); } } @@ -66,7 +67,7 @@ pub async fn run( .translation(wall_pos) .build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape_with_color( handle, &shape, diff --git a/crates/examples3d/compound3.rs b/crates/examples3d/rbd_compound3.rs similarity index 91% rename from crates/examples3d/compound3.rs rename to crates/examples3d/rbd_compound3.rs index 448ae2c2..d8486b19 100644 --- a/crates/examples3d/compound3.rs +++ b/crates/examples3d/rbd_compound3.rs @@ -1,17 +1,18 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; use rapier3d::prelude::*; /// Port of rapier's `compound3` demo, but every "U"-shaped body is assembled -/// from THREE separate colliders attached to one rigid body — exercising -/// multiple-colliders-per-body support — instead of a single compound collider. +/// from three separate colliders attached to one rigid body instead of a single +/// compound collider, which exercises multiple-colliders-per-body support. pub async fn run( viewer: &mut NexusViewer, pipeline: &mut NexusPipeline, ) -> anyhow::Result { let capacities = NexusCapacities::default().rbd_collisions(280_000); let mut state = NexusState::new(capacities); + let no_coupling = RbdCoupling::None; /* * Floor made of large cuboids. @@ -40,7 +41,7 @@ pub async fn run( .translation(wall_pos) .build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape_with_color( handle, &shape, @@ -86,7 +87,7 @@ pub async fn run( .translation(Vec3::new(x, y, z)) .build(); - let handle = state.insert_body_in(0, body); + let handle = state.insert_body_in(0, body, no_coupling); for (offset, he) in parts { let collider = ColliderBuilder::cuboid(he.x, he.y, he.z) .translation(offset) diff --git a/crates/examples3d/dynamic_rbd3.rs b/crates/examples3d/rbd_dynamic3.rs similarity index 93% rename from crates/examples3d/dynamic_rbd3.rs rename to crates/examples3d/rbd_dynamic3.rs index 6f7edea3..7f6ea698 100644 --- a/crates/examples3d/dynamic_rbd3.rs +++ b/crates/examples3d/rbd_dynamic3.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; use rapier3d::prelude::*; /// Demonstrates adding rigid-bodies to a live scene WITHOUT rebuilding the whole @@ -26,6 +26,7 @@ pub async fn run( let capacities = NexusCapacities::default().rbd_collisions(310_000); let mut state = NexusState::new(capacities); + let no_coupling = RbdCoupling::None; /* * A boxed ground: a floor plus four low walls to keep the pile contained. @@ -61,7 +62,7 @@ pub async fn run( .translation(pos) .build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape_with_color(handle, &shape, Pose::from_translation(pos), walls_color); } @@ -80,8 +81,8 @@ pub async fn run( while viewer.render_frame().await { if viewer.simulating() { - // Drop in a whole batch of bodies periodically — all inserted with a - // single batched append, without rebuilding the scene. + // Drop in a whole batch of bodies periodically, all inserted with a + // single batched append and without rebuilding the scene. if frame.is_multiple_of(SPAWN_PERIOD) && added < MAX_BODIES { let n = BODIES_PER_SPAWN.min(MAX_BODIES - added); let mut batch = Vec::with_capacity(n); @@ -105,7 +106,7 @@ pub async fn run( ColliderBuilder::cuboid(0.4, 0.4, 0.4).build() }; shapes.push(collider.shared_shape().clone()); - batch.push((body, collider)); + batch.push((body, collider, no_coupling)); } let handles = state.add_rigid_bodies(viewer.backend(), batch)?; diff --git a/crates/examples3d/joint_ball3.rs b/crates/examples3d/rbd_joint_ball3.rs similarity index 97% rename from crates/examples3d/joint_ball3.rs rename to crates/examples3d/rbd_joint_ball3.rs index 67ccebb6..ee088dcd 100644 --- a/crates/examples3d/joint_ball3.rs +++ b/crates/examples3d/rbd_joint_ball3.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; use rapier3d::prelude::*; pub async fn run( @@ -12,6 +12,7 @@ pub async fn run( */ let capacities = NexusCapacities::default().rbd_collisions(350_000); let mut state = NexusState::new(capacities); + let no_coupling = RbdCoupling::None; let rad = 0.4; let ni = 200; @@ -44,7 +45,7 @@ pub async fn run( ColliderBuilder::ball(rad).density(10.0).build() }; let shape = collider.shared_shape().clone(); - let child_handle = state.insert_rigid_body(rigid_body, collider); + let child_handle = state.insert_rigid_body(rigid_body, collider, no_coupling); viewer.insert_shape(child_handle, &shape, Pose::IDENTITY); // Vertical joint. @@ -84,7 +85,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::cuboid(rad, rad, rad).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); } } diff --git a/crates/examples3d/joint_fixed3.rs b/crates/examples3d/rbd_joint_fixed3.rs similarity index 93% rename from crates/examples3d/joint_fixed3.rs rename to crates/examples3d/rbd_joint_fixed3.rs index fe3a2500..167ae73c 100644 --- a/crates/examples3d/joint_fixed3.rs +++ b/crates/examples3d/rbd_joint_fixed3.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; use rapier3d::prelude::*; pub async fn run( @@ -11,6 +11,7 @@ pub async fn run( * World */ let mut state = NexusState::default(); + let no_coupling = RbdCoupling::None; let rad = 0.4; let num = 10; @@ -47,7 +48,8 @@ pub async fn run( .build(); let collider = ColliderBuilder::ball(rad).build(); let shape = collider.shared_shape().clone(); - let child_handle = state.insert_rigid_body(rigid_body, collider); + let child_handle = + state.insert_rigid_body(rigid_body, collider, no_coupling); viewer.insert_shape(child_handle, &shape, Pose::IDENTITY); // Vertical joint. diff --git a/crates/examples3d/joint_prismatic3.rs b/crates/examples3d/rbd_joint_prismatic3.rs similarity index 94% rename from crates/examples3d/joint_prismatic3.rs rename to crates/examples3d/rbd_joint_prismatic3.rs index f4e95aff..db5717d8 100644 --- a/crates/examples3d/joint_prismatic3.rs +++ b/crates/examples3d/rbd_joint_prismatic3.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; use rapier3d::prelude::*; pub async fn run( @@ -11,6 +11,7 @@ pub async fn run( * World */ let mut state = NexusState::default(); + let no_coupling = RbdCoupling::None; let rad = 0.4; let num = 10; @@ -30,7 +31,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::cuboid(rad, rad, rad).build(); let shape = collider.shared_shape().clone(); - let mut curr_parent = state.insert_rigid_body(ground, collider); + let mut curr_parent = state.insert_rigid_body(ground, collider, no_coupling); viewer.insert_shape(curr_parent, &shape, Pose::IDENTITY); for i in 0..num { @@ -43,7 +44,7 @@ pub async fn run( .density(density) .build(); let shape = collider.shared_shape().clone(); - let curr_child = state.insert_rigid_body(rigid_body, collider); + let curr_child = state.insert_rigid_body(rigid_body, collider, no_coupling); viewer.insert_shape(curr_child, &shape, Pose::IDENTITY); let axis = if i % 2 == 0 { diff --git a/crates/examples3d/joint_revolute3.rs b/crates/examples3d/rbd_joint_revolute3.rs similarity index 95% rename from crates/examples3d/joint_revolute3.rs rename to crates/examples3d/rbd_joint_revolute3.rs index 26e03bfa..56a3038d 100644 --- a/crates/examples3d/joint_revolute3.rs +++ b/crates/examples3d/rbd_joint_revolute3.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; use rapier3d::prelude::*; pub async fn run( @@ -11,6 +11,7 @@ pub async fn run( * World */ let mut state = NexusState::default(); + let no_coupling = RbdCoupling::None; let rad = 0.4; let num = 10; @@ -31,7 +32,7 @@ pub async fn run( .build(); let collider = ColliderBuilder::cuboid(rad, rad, rad).build(); let shape = collider.shared_shape().clone(); - let mut curr_parent = state.insert_rigid_body(ground, collider); + let mut curr_parent = state.insert_rigid_body(ground, collider, no_coupling); viewer.insert_shape(curr_parent, &shape, Pose::IDENTITY); for i in 0..num { @@ -52,7 +53,7 @@ pub async fn run( .density(density) .build(); let shape = collider.shared_shape().clone(); - handles[k] = state.insert_rigid_body(rigid_body, collider); + handles[k] = state.insert_rigid_body(rigid_body, collider, no_coupling); viewer.insert_shape(handles[k], &shape, Pose::IDENTITY); } diff --git a/crates/examples3d/joint_revolute_batch3.rs b/crates/examples3d/rbd_joint_revolute_batch3.rs similarity index 95% rename from crates/examples3d/joint_revolute_batch3.rs rename to crates/examples3d/rbd_joint_revolute_batch3.rs index 1f211f16..cf91ab91 100644 --- a/crates/examples3d/joint_revolute_batch3.rs +++ b/crates/examples3d/rbd_joint_revolute_batch3.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; use rapier3d::prelude::*; pub async fn run( @@ -9,6 +9,7 @@ pub async fn run( ) -> anyhow::Result { let capacities = NexusCapacities::default().rbd_collisions(32); let mut state = NexusState::new(capacities); + let no_coupling = RbdCoupling::None; let rad = 0.4; let num = 10; @@ -39,7 +40,8 @@ pub async fn run( .build(); let ground_collider = ColliderBuilder::cuboid(rad, rad, rad).build(); let ground_shape = ground_collider.shared_shape().clone(); - let mut curr_parent = state.insert_rigid_body_in(env, ground, ground_collider); + let mut curr_parent = + state.insert_rigid_body_in(env, ground, ground_collider, no_coupling); viewer.insert_shape_in( env as u32, curr_parent, @@ -66,7 +68,7 @@ pub async fn run( .density(density) .build(); let shape = collider.shared_shape().clone(); - handles[k] = state.insert_rigid_body_in(env, body, collider); + handles[k] = state.insert_rigid_body_in(env, body, collider, no_coupling); viewer.insert_shape_in( env as u32, handles[k], diff --git a/crates/examples3d/joints3.rs b/crates/examples3d/rbd_joints3.rs similarity index 99% rename from crates/examples3d/joints3.rs rename to crates/examples3d/rbd_joints3.rs index 478f377b..303aa48a 100644 --- a/crates/examples3d/joints3.rs +++ b/crates/examples3d/rbd_joints3.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; use rapier3d::prelude::*; /// Inserts a body + collider into the state and registers its render shape. @@ -11,7 +11,7 @@ fn add_body( collider: Collider, ) -> RigidBodyHandle { let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, RbdCoupling::None); viewer.insert_shape(handle, &shape, Pose::IDENTITY); handle } diff --git a/crates/examples3d/keva3.rs b/crates/examples3d/rbd_keva3.rs similarity index 98% rename from crates/examples3d/keva3.rs rename to crates/examples3d/rbd_keva3.rs index 90e6b885..0487379a 100644 --- a/crates/examples3d/keva3.rs +++ b/crates/examples3d/rbd_keva3.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; use rapier3d::prelude::*; /// Inserts a body + collider into the state and registers its render shape. @@ -11,7 +11,7 @@ fn add_body( collider: Collider, ) -> RigidBodyHandle { let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, RbdCoupling::None); viewer.insert_shape(handle, &shape, Pose::IDENTITY); handle } diff --git a/crates/examples3d/many_pyramids3.rs b/crates/examples3d/rbd_many_pyramids3.rs similarity index 97% rename from crates/examples3d/many_pyramids3.rs rename to crates/examples3d/rbd_many_pyramids3.rs index bb709775..a259c420 100644 --- a/crates/examples3d/many_pyramids3.rs +++ b/crates/examples3d/rbd_many_pyramids3.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; use rapier3d::prelude::*; /// Inserts a body + collider into the state and registers its render shape. @@ -11,7 +11,7 @@ fn add_body( collider: Collider, ) -> RigidBodyHandle { let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, RbdCoupling::None); viewer.insert_shape(handle, &shape, Pose::IDENTITY); handle } diff --git a/crates/examples3d/many_pyramids_batch3.rs b/crates/examples3d/rbd_many_pyramids_batch3.rs similarity index 96% rename from crates/examples3d/many_pyramids_batch3.rs rename to crates/examples3d/rbd_many_pyramids_batch3.rs index 578a8d7b..5d484d73 100644 --- a/crates/examples3d/many_pyramids_batch3.rs +++ b/crates/examples3d/rbd_many_pyramids_batch3.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; use rapier3d::prelude::*; fn create_pyramid( @@ -25,7 +25,7 @@ fn create_pyramid( .build(); let collider = ColliderBuilder::cuboid(rad, rad, rad).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body_in(env, body, collider); + let handle = state.insert_rigid_body_in(env, body, collider, RbdCoupling::None); viewer.insert_shape_in(env as u32, handle, &shape, Pose::IDENTITY, None); } } @@ -37,6 +37,7 @@ pub async fn run( ) -> anyhow::Result { let capacities = NexusCapacities::default().rbd_collisions(11_000); let mut state = NexusState::new(capacities); + let no_coupling = RbdCoupling::None; let pyramid_count = 40; for pyramid_index in 0..pyramid_count { @@ -66,7 +67,7 @@ pub async fn run( ) .build(); let shape = collider.shared_shape().clone(); - let ground_handle = state.insert_rigid_body_in(env, body, collider); + let ground_handle = state.insert_rigid_body_in(env, body, collider, no_coupling); viewer.insert_shape_in(env as u32, ground_handle, &shape, Pose::IDENTITY, None); /* diff --git a/crates/examples3d/mujoco_menagerie3.rs b/crates/examples3d/rbd_mujoco_menagerie3.rs similarity index 98% rename from crates/examples3d/mujoco_menagerie3.rs rename to crates/examples3d/rbd_mujoco_menagerie3.rs index 9c27130a..d2159b1f 100644 --- a/crates/examples3d/mujoco_menagerie3.rs +++ b/crates/examples3d/rbd_mujoco_menagerie3.rs @@ -1,7 +1,7 @@ use khal::backend::GpuTimestamps; use kiss3d::egui; use nexus_viewer3d::{NexusViewer, RenderMaterial}; -use nexus3d::prelude::{NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; use nexus3d::rbd::dynamics::convert_joint_motor; use nexus3d::rbd::shaders::dynamics::JointMotor; use rapier3d::prelude::*; @@ -31,11 +31,11 @@ const MAX_MB_DOFS: u32 = 64; /// Cheap pre-flight check: build the model's multibodies *without reading any /// mesh files* (collider/visual shape creation disabled) and return the largest /// per-multibody DoF count. `None` if the model fails to parse (those are kept -/// in the list — `load_scene` reports the error gracefully instead of crashing). +/// in the list: `load_scene` reports the error gracefully instead of crashing). /// Used to drop models the GPU solver can't handle from the picker. fn scene_max_dofs(scene: &Path) -> Option { // Same structural options as the real load (so the DoF count matches), but - // with every collider/visual shape skipped — only bodies and joints, which + // with every collider/visual shape skipped: only bodies and joints, which // determine the multibody DoFs, are needed here. let options = MjcfLoaderOptions { create_colliders_from_collision_shapes: false, @@ -456,7 +456,7 @@ async fn load_scene( let body = RigidBodyBuilder::fixed().translation(center).build(); let collider = ColliderBuilder::cuboid(he.x, he.y, he.z).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, RbdCoupling::None); viewer.insert_shape(handle, &shape, Pose::IDENTITY); } @@ -587,7 +587,7 @@ async fn select_scene( && settings.use_multibody { return Ok(Err(format!( - "{} needs {dofs} DoFs (max {MAX_MB_DOFS}) — not supported by the GPU solver.", + "{} needs {dofs} DoFs (max {MAX_MB_DOFS}), not supported by the GPU solver.", scene_label(scene) ))); } diff --git a/crates/examples3d/multibody_pendulum3.rs b/crates/examples3d/rbd_multibody_pendulum3.rs similarity index 92% rename from crates/examples3d/multibody_pendulum3.rs rename to crates/examples3d/rbd_multibody_pendulum3.rs index 993b4ab6..6e6de320 100644 --- a/crates/examples3d/multibody_pendulum3.rs +++ b/crates/examples3d/rbd_multibody_pendulum3.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; use rapier3d::prelude::*; pub async fn run( @@ -8,6 +8,7 @@ pub async fn run( pipeline: &mut NexusPipeline, ) -> anyhow::Result { let mut state = NexusState::default(); + let no_coupling = RbdCoupling::None; /* * The ground @@ -32,8 +33,8 @@ pub async fn run( * - Under gravity alone, the chain should swing in the YZ plane. * * The GPU pipeline picks up the multibody set from `SimulationState::environments` - * and runs `GpuMultibodySolver::step` each frame — no contacts or constraints - * with multibodies are involved. + * and runs `GpuMultibodySolver::step` each frame. No contacts or + * constraints with multibodies are involved. */ let rad = 0.4; let link_len = 2.0; @@ -43,7 +44,7 @@ pub async fn run( let root_body = RigidBodyBuilder::fixed().build(); let root_collider = ColliderBuilder::cuboid(rad, rad, rad).build(); let root_shape = root_collider.shared_shape().clone(); - let mut parent_handle = state.insert_rigid_body(root_body, root_collider); + let mut parent_handle = state.insert_rigid_body(root_body, root_collider, no_coupling); viewer.insert_shape(parent_handle, &root_shape, Pose::IDENTITY); for i in 0..num_links { @@ -56,7 +57,7 @@ pub async fn run( .collision_groups(InteractionGroups::none()) .build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(rigid_body, collider); + let handle = state.insert_rigid_body(rigid_body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); // Revolute joint about X: anchor on parent is at its bottom diff --git a/crates/examples3d/primitives3.rs b/crates/examples3d/rbd_primitives3.rs similarity index 96% rename from crates/examples3d/primitives3.rs rename to crates/examples3d/rbd_primitives3.rs index b0c6c765..14279361 100644 --- a/crates/examples3d/primitives3.rs +++ b/crates/examples3d/rbd_primitives3.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; use rapier3d::prelude::*; pub async fn run( @@ -12,6 +12,7 @@ pub async fn run( let capacities = NexusCapacities::default().rbd_collisions(150_000); let mut state = NexusState::new(capacities); + let no_coupling = RbdCoupling::None; /* * Falling dynamic objects. @@ -72,7 +73,7 @@ pub async fn run( let body = RigidBodyBuilder::dynamic().translation(pos).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); } } @@ -105,7 +106,7 @@ pub async fn run( .translation(wall_pos) .build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape_with_color( handle, &shape, diff --git a/crates/examples3d/pyramid3.rs b/crates/examples3d/rbd_pyramid3.rs similarity index 96% rename from crates/examples3d/pyramid3.rs rename to crates/examples3d/rbd_pyramid3.rs index 47c0edf5..e794f0f6 100644 --- a/crates/examples3d/pyramid3.rs +++ b/crates/examples3d/rbd_pyramid3.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; use rapier3d::prelude::*; /// Inserts a body + collider into the state and registers its render shape. @@ -11,7 +11,7 @@ fn add_body( collider: Collider, ) -> RigidBodyHandle { let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, RbdCoupling::None); viewer.insert_shape(handle, &shape, Pose::IDENTITY); handle } diff --git a/crates/examples3d/trimesh3.rs b/crates/examples3d/rbd_trimesh3.rs similarity index 96% rename from crates/examples3d/trimesh3.rs rename to crates/examples3d/rbd_trimesh3.rs index 8f3dfdee..baf5c143 100644 --- a/crates/examples3d/trimesh3.rs +++ b/crates/examples3d/rbd_trimesh3.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusCapacities, NexusPipeline, NexusState, RbdCoupling}; use rapier3d::parry::utils::Array2; use rapier3d::prelude::*; @@ -13,6 +13,7 @@ pub async fn run( let capacities = NexusCapacities::default().rbd_collisions(150_000); let mut state = NexusState::new(capacities); + let no_coupling = RbdCoupling::None; /* * Falling dynamic objects. @@ -74,7 +75,7 @@ pub async fn run( let body = RigidBodyBuilder::dynamic().translation(pos).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); } } @@ -114,7 +115,7 @@ pub async fn run( .unwrap() .build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, no_coupling); viewer.insert_shape(handle, &shape, Pose::IDENTITY); let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); diff --git a/crates/examples3d/urdf3.rs b/crates/examples3d/rbd_urdf3.rs similarity index 96% rename from crates/examples3d/urdf3.rs rename to crates/examples3d/rbd_urdf3.rs index de2c7007..b878f35a 100644 --- a/crates/examples3d/urdf3.rs +++ b/crates/examples3d/rbd_urdf3.rs @@ -1,6 +1,6 @@ use khal::backend::GpuTimestamps; use nexus_viewer3d::NexusViewer; -use nexus3d::prelude::{NexusPipeline, NexusState}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; use rapier3d::prelude::*; use rapier3d_urdf::{UrdfLoaderOptions, UrdfMultibodyOptions, UrdfRobot}; use std::path::PathBuf; @@ -23,6 +23,7 @@ pub async fn run( use rand::RngExt; let mut state = NexusState::default(); + let _ = RbdCoupling::None; /* * Robot loaded from URDF. @@ -60,8 +61,8 @@ pub async fn run( Ok((mut robot, _)) => { // Switch every joint's `AngX` motor to acceleration-based mode so the // per-frame motor target velocity feels right regardless of link mass. - // Initial target velocity is 0 — the loop below re-randomizes it every - // 5 simulated seconds. + // Initial target velocity is 0; the loop below re-randomizes it + // every 5 simulated seconds. for urdf_joint in &mut robot.joints { urdf_joint .joint From 10888debc6d74e141161a44d3c8c04ddd84a92c3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 14 Aug 2026 23:47:14 +0200 Subject: [PATCH 5/7] feat: add the MPM demos --- crates/examples2d/Cargo.toml | 3 +- crates/examples2d/all_examples2.rs | 65 +++++--- crates/examples2d/mpm_centilever_beam2.rs | 70 +++++++++ crates/examples2d/mpm_cohesion_sweep2.rs | 99 ++++++++++++ crates/examples2d/mpm_dam_break2.rs | 138 +++++++++++++++++ crates/examples2d/mpm_elastic_cut2.rs | 116 ++++++++++++++ crates/examples2d/mpm_elasticity2.rs | 96 ++++++++++++ crates/examples2d/mpm_emitter2.rs | 160 +++++++++++++++++++ crates/examples2d/mpm_hourglass2.rs | 171 ++++++++++++++++++++ crates/examples2d/mpm_sand2.rs | 161 +++++++++++++++++++ crates/examples2d/mpm_snowball2.rs | 158 +++++++++++++++++++ crates/examples3d/Cargo.toml | 2 +- crates/examples3d/all_examples3.rs | 101 +++++++----- crates/examples3d/mpm_centilever_beam3.rs | 76 +++++++++ crates/examples3d/mpm_dam_break3.rs | 155 ++++++++++++++++++ crates/examples3d/mpm_elastic_cut3.rs | 89 +++++++++++ crates/examples3d/mpm_emitter3.rs | 134 ++++++++++++++++ crates/examples3d/mpm_erosion3.rs | 181 ++++++++++++++++++++++ crates/examples3d/mpm_heightfield3.rs | 76 +++++++++ crates/examples3d/mpm_jelly_drop3.rs | 120 ++++++++++++++ crates/examples3d/mpm_sand3.rs | 101 ++++++++++++ crates/examples3d/mpm_snow3.rs | 155 ++++++++++++++++++ 22 files changed, 2363 insertions(+), 64 deletions(-) create mode 100644 crates/examples2d/mpm_centilever_beam2.rs create mode 100644 crates/examples2d/mpm_cohesion_sweep2.rs create mode 100644 crates/examples2d/mpm_dam_break2.rs create mode 100644 crates/examples2d/mpm_elastic_cut2.rs create mode 100644 crates/examples2d/mpm_elasticity2.rs create mode 100644 crates/examples2d/mpm_emitter2.rs create mode 100644 crates/examples2d/mpm_hourglass2.rs create mode 100644 crates/examples2d/mpm_sand2.rs create mode 100644 crates/examples2d/mpm_snowball2.rs create mode 100644 crates/examples3d/mpm_centilever_beam3.rs create mode 100644 crates/examples3d/mpm_dam_break3.rs create mode 100644 crates/examples3d/mpm_elastic_cut3.rs create mode 100644 crates/examples3d/mpm_emitter3.rs create mode 100644 crates/examples3d/mpm_erosion3.rs create mode 100644 crates/examples3d/mpm_heightfield3.rs create mode 100644 crates/examples3d/mpm_jelly_drop3.rs create mode 100644 crates/examples3d/mpm_sand3.rs create mode 100644 crates/examples3d/mpm_snow3.rs diff --git a/crates/examples2d/Cargo.toml b/crates/examples2d/Cargo.toml index 188d5582..717537ee 100644 --- a/crates/examples2d/Cargo.toml +++ b/crates/examples2d/Cargo.toml @@ -14,7 +14,7 @@ metal = ["nexus_viewer2d/metal"] [dependencies] glamx = { workspace = true } -nexus2d = { workspace = true, features = [ "rbd" ] } +nexus2d = { workspace = true, features = [ "rbd", "mpm" ] } nexus_viewer2d = { workspace = true } nexus_rbd2d = { workspace = true, features = ["default"] } rapier2d = { workspace = true, features = ["default"] } @@ -24,6 +24,7 @@ kiss3d = { workspace = true } oorandom = { workspace = true } rand = "0.10" anyhow = { workspace = true } +pollster = { workspace = true } [target.'cfg(target_arch = "wasm32")'.dependencies] getrandom = { workspace = true } diff --git a/crates/examples2d/all_examples2.rs b/crates/examples2d/all_examples2.rs index 66114c28..6014f1e0 100644 --- a/crates/examples2d/all_examples2.rs +++ b/crates/examples2d/all_examples2.rs @@ -2,17 +2,28 @@ use inflector::Inflector; use nexus_viewer2d::{BackendType, DemoKind, NexusViewer}; use nexus2d::prelude::{NexusPipeline, NexusPipelineMask}; -mod balls2; -mod boxes2; -mod boxes_and_balls2; -mod compound2; -mod dynamic_rbd2; -mod joint_ball2; -mod joint_fixed2; -mod joint_prismatic2; -mod polyline2; -mod primitives2; -mod pyramid2; +mod rbd_balls2; +mod rbd_boxes2; +mod rbd_boxes_and_balls2; +mod rbd_compound2; +mod rbd_dynamic2; +mod rbd_joint_ball2; +mod rbd_joint_fixed2; +mod rbd_joint_prismatic2; +mod rbd_polyline2; +mod rbd_primitives2; +mod rbd_pyramid2; + +// MPM examples. +mod mpm_centilever_beam2; +mod mpm_cohesion_sweep2; +mod mpm_dam_break2; +mod mpm_elastic_cut2; +mod mpm_elasticity2; +mod mpm_emitter2; +mod mpm_hourglass2; +mod mpm_sand2; +mod mpm_snowball2; /// Declares the demo registry: a `(name, kind)` list for the picker UI and a /// name -> `run()` dispatcher. Keeping both in one macro keeps them in sync. @@ -42,17 +53,27 @@ macro_rules! demos { } demos! { - "Balls" => Rbd : balls2, - "Boxes" => Rbd : boxes2, - "Boxes & balls" => Rbd : boxes_and_balls2, - "Compound" => Rbd : compound2, - "Dynamic insertion" => Rbd : dynamic_rbd2, - "Pyramid" => Rbd : pyramid2, - "Primitives" => Rbd : primitives2, - "Polyline" => Rbd : polyline2, - "Joints (spherical)" => Rbd : joint_ball2, - "Joints (prismatic)" => Rbd : joint_prismatic2, - "Joints (fixed)" => Rbd : joint_fixed2, + "Balls" => Rbd : rbd_balls2, + "Boxes" => Rbd : rbd_boxes2, + "Boxes & balls" => Rbd : rbd_boxes_and_balls2, + "Compound" => Rbd : rbd_compound2, + "Dynamic insertion" => Rbd : rbd_dynamic2, + "Pyramid" => Rbd : rbd_pyramid2, + "Primitives" => Rbd : rbd_primitives2, + "Polyline" => Rbd : rbd_polyline2, + "Joints (spherical)" => Rbd : rbd_joint_ball2, + "Joints (prismatic)" => Rbd : rbd_joint_prismatic2, + "Joints (fixed)" => Rbd : rbd_joint_fixed2, + // MPM demos. + "Cantilever beam" => Mpm : mpm_centilever_beam2, + "Sand" => Mpm : mpm_sand2, + "Sand emitter" => Mpm : mpm_emitter2, + "Elasticity" => Mpm : mpm_elasticity2, + "Elastic cut" => Mpm : mpm_elastic_cut2, + "Dam break" => Mpm : mpm_dam_break2, + "Cohesion sweep" => Mpm : mpm_cohesion_sweep2, + "Snowballs" => Mpm : mpm_snowball2, + "Hourglass" => Mpm : mpm_hourglass2, } struct CliOptions { diff --git a/crates/examples2d/mpm_centilever_beam2.rs b/crates/examples2d/mpm_centilever_beam2.rs new file mode 100644 index 00000000..22503ace --- /dev/null +++ b/crates/examples2d/mpm_centilever_beam2.rs @@ -0,0 +1,70 @@ +use khal::backend::GpuTimestamps; +use nexus_viewer2d::NexusViewer; +use nexus2d::mpm::solver::{BoundaryCondition, Particle, ParticleModel, SimulationParams}; +use nexus2d::prelude::{NexusPipeline, NexusState, RbdCoupling}; + +use rapier2d::prelude::{ColliderBuilder, Pose, RigidBodyBuilder}; + +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> anyhow::Result { + let mut state = NexusState::default(); + + let width = 10.0; + let height = 2.0; + let fixed_part = 1.0; + let cell_width = 0.2; + let particle_per_cell_dim = 2; + let young_modulus = 1.0e8; + let poisson_ratio = 0.3; + + let diameter = cell_width / particle_per_cell_dim as f32; + let ni = ((width + fixed_part) / diameter).ceil() as usize; + let nj = (height / diameter).ceil() as usize; + + let mut particles = vec![]; + for i in 0..ni { + for j in 0..nj { + let position = glamx::vec2(i as f32, j as f32) * diameter; + let density = 1000.0; + let radius = diameter / 2.0; + let model = ParticleModel::elastic_neo_hookean(young_modulus, poisson_ratio); + particles.push(Particle::new(position, radius, density, model)); + } + } + + let params = SimulationParams { + gravity: glamx::vec2(0.0, -9.81), + padding: 0.0, + dt: 1.0 / 60.0, + }; + state.set_mpm_params(viewer.backend(), params, cell_width)?; + state.set_mpm_substeps(150); + state.add_particles(viewer.backend(), particles)?; + + // Fixed anchor the beam is cantilevered from (boundary coupled to the MPM + // continuum). + let body = RigidBodyBuilder::fixed() + .translation(glamx::vec2(0.0, height / 2.0)) + .build(); + let collider = ColliderBuilder::cuboid(fixed_part, height).build(); + let shape = collider.shared_shape().clone(); + let stick = BoundaryCondition::stick(); + let handle = state.insert_rigid_body(body, collider, RbdCoupling::MpmOneWay(stick)); + viewer.insert_shape(handle, &shape, Pose::IDENTITY); + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + state.finalize(viewer.backend()).await?; + + while viewer.render_frame().await { + if viewer.simulating() { + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} diff --git a/crates/examples2d/mpm_cohesion_sweep2.rs b/crates/examples2d/mpm_cohesion_sweep2.rs new file mode 100644 index 00000000..9c935fa9 --- /dev/null +++ b/crates/examples2d/mpm_cohesion_sweep2.rs @@ -0,0 +1,99 @@ +//! Five identical granular columns released at once, differing only in cohesion. + +use khal::backend::GpuTimestamps; +use nexus_viewer2d::NexusViewer; +use nexus2d::mpm::solver::{BoundaryCondition, Particle, ParticleModel, SimulationParams}; +use nexus2d::prelude::{NexusPipeline, NexusState, RbdCoupling}; + +use glamx::{Vec4, vec2}; +use rapier2d::prelude::{ColliderBuilder, Pose, RigidBodyBuilder}; + +const DENSITY: f32 = 1600.0; +const YOUNG_MODULUS: f32 = 1.0e7; +const POISSON_RATIO: f32 = 0.2; + +/// Cohesion of each column, left to right. +const COHESIONS: [f32; 5] = [0.0, 0.001, 0.003, 0.01, 0.03]; + +/// Half-width of each column. +const COLUMN_HALF_WIDTH: f32 = 3.0; +/// Height of each column. +const COLUMN_HEIGHT: f32 = 20.0; +/// Distance between column centers. +const COLUMN_PITCH: f32 = 22.0; + +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> anyhow::Result { + let mut state = NexusState::default(); + + let cell_width = 0.2; + let radius = cell_width / 4.0; + let spacing = radius * 2.0; + + let mut particles = vec![]; + let nx = (COLUMN_HALF_WIDTH * 2.0 / spacing) as i32; + let ny = (COLUMN_HEIGHT / spacing) as i32; + let count = COHESIONS.len(); + // Dry sand is pale, the most cohesive column is dark. + let shades: Vec<_> = (0..count) + .map(|c| { + let t = c as f32 / (count - 1) as f32; + Vec4::new(0.92 - 0.5 * t, 0.80 - 0.48 * t, 0.62 - 0.42 * t, 1.0) + }) + .collect(); + viewer.set_particle_group_colors(&shades); + for (c, cohesion) in COHESIONS.iter().enumerate() { + let center_x = (c as f32 - (count - 1) as f32 / 2.0) * COLUMN_PITCH; + let model = ParticleModel::cohesive_sand(YOUNG_MODULUS, POISSON_RATIO, *cohesion); + + for i in 0..nx { + for j in 0..ny { + let position = vec2( + center_x - COLUMN_HALF_WIDTH + (i as f32 + 0.5) * spacing, + 0.2 + (j as f32 + 0.5) * spacing, + ); + particles.push(Particle::with_group( + position, radius, DENSITY, model, c as u32, + )); + } + } + } + + let params = SimulationParams { + gravity: vec2(0.0, -9.81), + padding: 0.0, + dt: 1.0 / 60.0, + }; + state.set_mpm_params(viewer.backend(), params, cell_width)?; + state.set_mpm_substeps(15); + state.add_particles(viewer.backend(), particles)?; + + /* + * Setup the floor. + */ + let half_span = COLUMN_PITCH * count as f32 / 2.0 + 10.0; + let collider = ColliderBuilder::cuboid(half_span, 1.0).build(); + let shape = collider.shared_shape().clone(); + let body = RigidBodyBuilder::fixed() + .translation(vec2(0.0, -1.0)) + .build(); + let coupling = RbdCoupling::MpmOneWay(BoundaryCondition::separate(1.0)); + let handle = state.insert_rigid_body(body, collider, coupling); + viewer.insert_shape(handle, &shape, Pose::IDENTITY); + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + state.finalize(viewer.backend()).await?; + + while viewer.render_frame().await { + if viewer.simulating() { + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} diff --git a/crates/examples2d/mpm_dam_break2.rs b/crates/examples2d/mpm_dam_break2.rs new file mode 100644 index 00000000..74019188 --- /dev/null +++ b/crates/examples2d/mpm_dam_break2.rs @@ -0,0 +1,138 @@ +//! The classic dam break: a column of water collapses and runs along a tank. + +use khal::backend::GpuTimestamps; +use nexus_viewer2d::NexusViewer; +use nexus2d::mpm::solver::{BoundaryCondition, Particle, ParticleModel, SimulationParams}; +use nexus2d::prelude::{NexusPipeline, NexusState, RbdCoupling}; + +use glamx::{Vec4, vec2}; +use rapier2d::prelude::{Collider, ColliderBuilder, Pose, RigidBody, RigidBodyBuilder}; + +const DENSITY: f32 = 1000.0; +/// Deepest the water gets. The bulk modulus follows from it. +const MAX_DEPTH: f32 = 22.0; +/// Volume loss tolerated at that depth. Water's real bulk modulus (~2.2 GPa) +/// would force an impractically small timestep; a weakly compressible scheme +/// instead picks the softest fluid whose compression is still invisible. +const MAX_COMPRESSION: f32 = 0.01; + +const TANK_HALF_WIDTH: f32 = 30.0; +const TANK_HEIGHT: f32 = 26.0; + +fn insert_boundary( + state: &mut NexusState, + viewer: &mut NexusViewer, + body: RigidBody, + collider: Collider, +) { + let shape = collider.shared_shape().clone(); + let coupling = RbdCoupling::MpmOneWay(BoundaryCondition::separate(0.0)); + let handle = state.insert_rigid_body(body, collider, coupling); + viewer.insert_shape(handle, &shape, Pose::IDENTITY); +} + +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> anyhow::Result { + let mut state = NexusState::default(); + + let cell_width = 0.2; + let radius = cell_width / 4.0; + let spacing = radius * 2.0; + let model = ParticleModel::water_for_depth(DENSITY, 9.81, MAX_DEPTH, MAX_COMPRESSION); + + /* + * The water column, held against the left wall. + */ + let mut particles = vec![]; + let column_width = 32.0; + let column_height = 25.0; + let nx = (column_width / spacing) as i32; + let ny = (column_height / spacing) as i32; + let shades: Vec<_> = (0..ny) + .map(|j| { + let t = j as f32 / ny as f32; + Vec4::new(0.10 + 0.35 * t, 0.45 + 0.35 * t, 0.85 + 0.15 * t, 1.0) + }) + .collect(); + viewer.set_particle_group_colors(&shades); + for i in 0..nx { + for j in 0..ny { + let position = vec2( + -TANK_HALF_WIDTH + 0.5 + (i as f32 + 0.5) * spacing, + 0.5 + (j as f32 + 0.5) * spacing, + ); + // One group per row: shading by initial height makes the + // overturning of the surge front visible once the column collapses. + particles.push(Particle::with_group( + position, radius, DENSITY, model, j as u32, + )); + } + } + + let params = SimulationParams { + gravity: vec2(0.0, -9.81), + padding: 0.0, + dt: 1.0 / 60.0, + }; + state.set_mpm_params(viewer.backend(), params, cell_width)?; + state.set_mpm_substeps(25); + state.add_particles(viewer.backend(), particles)?; + + /* + * Tank. + */ + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(vec2(0.0, -1.0)) + .build(), + ColliderBuilder::cuboid(TANK_HALF_WIDTH + 1.0, 1.0).build(), + ); + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(vec2(0.0, TANK_HEIGHT + 1.0)) + .build(), + ColliderBuilder::cuboid(TANK_HALF_WIDTH + 1.0, 1.0).build(), + ); + for side in [-1.0f32, 1.0] { + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(vec2(side * (TANK_HALF_WIDTH + 1.0), TANK_HEIGHT / 2.0)) + .build(), + ColliderBuilder::cuboid(1.0, TANK_HEIGHT / 2.0 + 1.0).build(), + ); + } + + /* + * An obstacle in the path of the surge, which the front breaks over. + */ + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(vec2(8.0, 2.0)) + .build(), + ColliderBuilder::cuboid(1.5, 2.0).build(), + ); + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + state.finalize(viewer.backend()).await?; + + while viewer.render_frame().await { + if viewer.simulating() { + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} diff --git a/crates/examples2d/mpm_elastic_cut2.rs b/crates/examples2d/mpm_elastic_cut2.rs new file mode 100644 index 00000000..e9e3530b --- /dev/null +++ b/crates/examples2d/mpm_elastic_cut2.rs @@ -0,0 +1,116 @@ +use khal::backend::GpuTimestamps; +use nexus_viewer2d::NexusViewer; +use nexus2d::mpm::solver::{BoundaryCondition, Particle, ParticleModel, SimulationParams}; +use nexus2d::prelude::{NexusPipeline, NexusState, RbdCoupling}; + +use glamx::Vec2; +use rapier2d::prelude::{Collider, ColliderBuilder, Pose, RigidBody, RigidBodyBuilder}; + +/// Inserts a boundary collider coupled (one-way) to the MPM particles and +/// registers it for rendering. +fn insert_boundary( + state: &mut NexusState, + viewer: &mut NexusViewer, + body: RigidBody, + collider: Collider, +) { + let shape = collider.shared_shape().clone(); + let friction = BoundaryCondition::separate(1.0); + let handle = state.insert_rigid_body(body, collider, RbdCoupling::MpmOneWay(friction)); + viewer.insert_shape(handle, &shape, Pose::IDENTITY); +} + +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> anyhow::Result { + let mut state = NexusState::default(); + + let offset_y = 46.0; + // let cell_width = 0.1; + let cell_width = 0.2; + + let mut particles = vec![]; + for i in 0..700 { + for j in 0..700 { + let position = + glamx::vec2(i as f32 + 0.5, j as f32 + 0.5) * cell_width / 2.0 + Vec2::Y * offset_y; + + let density = 1000.0; + let radius = cell_width / 4.0; + let model = ParticleModel::elastic(5.0e6, 0.2); + particles.push(Particle::new(position, radius, density, model)); + } + } + + let params = SimulationParams { + gravity: glamx::vec2(0.0, -9.81), + padding: 0.0, + dt: 1.0 / 60.0, + }; + state.set_mpm_params(viewer.backend(), params, cell_width)?; + state.set_mpm_substeps(15); + state.add_particles(viewer.backend(), particles)?; + + // const ANGVEL: f32 = 1.0; // 2.0; + + /* + * Static platforms. + */ + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(glamx::vec2(35.0, 20.0)) + .build(), + ColliderBuilder::cuboid(70.0, 1.0).build(), + ); + + let mut polyline = vec![]; + let subdivs = 100; + let length = 84.0; + let start = glamx::vec2(35.0, 70.0) - glamx::vec2(length / 2.0, 0.0); + + for i in 0..=subdivs { + let step = length / (subdivs as f32); + let dx = i as f32 * step; + polyline.push(start + glamx::vec2(dx, dx.sin())) + } + + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed().build(), + ColliderBuilder::polyline(polyline, None).build(), + ); + + for k in 0..6 { + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed().build(), + ColliderBuilder::polyline( + vec![ + glamx::vec2(0.0 + k as f32 * 15.0, 20.0), + glamx::vec2(-10.0 + k as f32 * 15.0, 45.0), + ], + None, + ) + .build(), + ); + } + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + state.finalize(viewer.backend()).await?; + + while viewer.render_frame().await { + if viewer.simulating() { + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} diff --git a/crates/examples2d/mpm_elasticity2.rs b/crates/examples2d/mpm_elasticity2.rs new file mode 100644 index 00000000..640fd358 --- /dev/null +++ b/crates/examples2d/mpm_elasticity2.rs @@ -0,0 +1,96 @@ +use khal::backend::GpuTimestamps; +use nexus_viewer2d::NexusViewer; +use nexus2d::mpm::solver::{BoundaryCondition, Particle, ParticleModel, SimulationParams}; +use nexus2d::prelude::{NexusPipeline, NexusState, RbdCoupling}; + +use glamx::Vec2; +use rapier2d::prelude::{Collider, ColliderBuilder, Pose, RigidBody, RigidBodyBuilder}; + +/// Inserts a boundary collider coupled (one-way) to the MPM particles and +/// registers it for rendering. +fn insert_boundary( + state: &mut NexusState, + viewer: &mut NexusViewer, + body: RigidBody, + collider: Collider, +) { + let shape = collider.shared_shape().clone(); + let separate = BoundaryCondition::separate(1.0); + let handle = state.insert_rigid_body(body, collider, RbdCoupling::MpmOneWay(separate)); + viewer.insert_shape(handle, &shape, Pose::IDENTITY); +} + +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> anyhow::Result { + let mut state = NexusState::default(); + + let offset_y = 10.0; + // let cell_width = 0.1; + let cell_width = 0.2; + + let mut particles = vec![]; + for i in 0..700 { + for j in 0..700 { + let position = glamx::vec2(i as f32 + 0.5 + (i / 50) as f32 * 2.0, j as f32 + 0.5) + * cell_width + / 2.0 + + Vec2::Y * offset_y; + let density = 1000.0; + let radius = cell_width / 4.0; + let model = ParticleModel::elastic(5.0e6, 0.2); + particles.push(Particle::new(position, radius, density, model)); + } + } + + let params = SimulationParams { + gravity: glamx::vec2(0.0, -9.81) * 2.0, + padding: 0.0, + dt: 1.0 / 60.0, + }; + state.set_mpm_params(viewer.backend(), params, cell_width)?; + state.set_mpm_substeps(15); + state.add_particles(viewer.backend(), particles)?; + + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(glamx::vec2(0.0, -1.0)) + .build(), + ColliderBuilder::cuboid(1000.0, 1.0).build(), + ); + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(glamx::vec2(-20.0, 0.0)) + .rotation(0.5) + .build(), + ColliderBuilder::cuboid(1.0, 60.0).build(), + ); + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(glamx::vec2(90.0, 0.0)) + .rotation(-0.5) + .build(), + ColliderBuilder::cuboid(1.0, 60.0).build(), + ); + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + state.finalize(viewer.backend()).await?; + + while viewer.render_frame().await { + if viewer.simulating() { + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} diff --git a/crates/examples2d/mpm_emitter2.rs b/crates/examples2d/mpm_emitter2.rs new file mode 100644 index 00000000..275c5758 --- /dev/null +++ b/crates/examples2d/mpm_emitter2.rs @@ -0,0 +1,160 @@ +use khal::backend::GpuTimestamps; +use nexus_viewer2d::NexusViewer; +use nexus2d::mpm::solver::{BoundaryCondition, Particle, ParticleModel, SimulationParams}; +use nexus2d::prelude::{NexusParticleChunk, NexusPipeline, NexusState, RbdCoupling}; + +use glamx::{Vec4, vec2}; +use rapier2d::prelude::{Collider, ColliderBuilder, Pose, RigidBody, RigidBodyBuilder}; + +use std::collections::VecDeque; + +const DENSITY: f32 = 1000.0; +const YOUNG_MODULUS: f32 = 1.0e7; +const POISSON_RATIO: f32 = 0.2; + +/// Emit a fresh patch of particles every `EMIT_EVERY` substeps. +const EMIT_EVERY: u64 = 10; +/// Edge length (in particles) of the emitted block. +const EMIT_BLOCK: i32 = 120; +const EMIT_BLOCK_Y: i32 = 20; +/// Cap on the live particle count. Once reached, every emit removes as many of +/// the oldest particles as it adds, so the total stays at this budget. +const MAX_PARTICLES: usize = 250_000; + +/// Inserts a boundary collider coupled (one-way) to the MPM particles and +/// registers it for rendering. +fn insert_boundary( + state: &mut NexusState, + viewer: &mut NexusViewer, + body: RigidBody, + collider: Collider, +) { + let shape = collider.shared_shape().clone(); + let friction = BoundaryCondition::separate(1.0); + let handle = state.insert_rigid_body(body, collider, RbdCoupling::MpmOneWay(friction)); + viewer.insert_shape(handle, &shape, Pose::IDENTITY); +} + +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> anyhow::Result { + let mut state = NexusState::default(); + + let cell_width = 0.2; + let dt = 1.0 / 60.0; + + let params = SimulationParams { + gravity: vec2(0.0, -9.81), + padding: 0.0, + dt, + }; + state.set_mpm_params(viewer.backend(), params, cell_width)?; + state.set_mpm_substeps(10); + // No particles up-front: they are emitted dynamically in the loop below. + + /* + * Boundary colliders: a floor, two side walls, and a pair of angled ramps + * in the middle for the sand stream to cascade over. + */ + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(vec2(40.0, -1.0)) + .build(), + ColliderBuilder::cuboid(45.0, 1.0).build(), + ); + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(vec2(-4.0, 30.0)) + .build(), + ColliderBuilder::cuboid(1.0, 32.0).build(), + ); + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(vec2(84.0, 30.0)) + .build(), + ColliderBuilder::cuboid(1.0, 32.0).build(), + ); + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(vec2(25.0, 20.0)) + .rotation(-0.5) + .build(), + ColliderBuilder::cuboid(12.0, 0.8).build(), + ); + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(vec2(58.0, 12.0)) + .rotation(0.5) + .build(), + ColliderBuilder::cuboid(12.0, 0.8).build(), + ); + viewer.set_particle_group_colors(&[Vec4::new(0.95, 0.75, 0.35, 1.0)]); + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + state.finalize(viewer.backend()).await?; + + /* + * Dynamic emitter: a small block of sand spawned at a point that sweeps + * left and right across the top, pouring a moving curtain of sand. + */ + let radius = cell_width / 4.0; + let spacing = radius * 2.0; + let model = ParticleModel::sand(YOUNG_MODULUS, POISSON_RATIO); + let emit_height = 55.0; + let sweep_center = 40.0; + let sweep_amplitude = 30.0; + let sweep_speed = 0.9; // rad/s + + // Oldest-first queue of live chunks paired with their particle counts, plus + // the running total `MAX_PARTICLES` is enforced against. + let mut chunks: VecDeque<(NexusParticleChunk, usize)> = VecDeque::new(); + let mut total_particles: usize = 0; + let mut t: f32 = 0.0; + let mut step: u64 = 0; + + while viewer.render_frame().await { + if viewer.simulating() { + if step.is_multiple_of(EMIT_EVERY) && total_particles < MAX_PARTICLES { + let phase = t * sweep_speed; + let center = vec2(sweep_center + sweep_amplitude * phase.sin(), emit_height); + let velocity = vec2(0.0, -12.0); + + let mut particles = Vec::with_capacity((EMIT_BLOCK * EMIT_BLOCK_Y) as usize); + for i in 0..EMIT_BLOCK { + for j in 0..EMIT_BLOCK_Y { + let offset = + vec2((i - EMIT_BLOCK / 2) as f32, (j - EMIT_BLOCK_Y / 2) as f32) + * spacing; + let mut particle = Particle::new(center + offset, radius, DENSITY, model); + particle.dynamics.velocity = velocity; + particles.push(particle); + } + } + let n = particles.len(); + let chunk = state.add_particles(viewer.backend(), particles)?; + chunks.push_back((chunk, n)); + total_particles += n; + } + + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + t += dt; + step += 1; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} diff --git a/crates/examples2d/mpm_hourglass2.rs b/crates/examples2d/mpm_hourglass2.rs new file mode 100644 index 00000000..16d88f4a --- /dev/null +++ b/crates/examples2d/mpm_hourglass2.rs @@ -0,0 +1,171 @@ +//! Sand draining through the neck of an hourglass. + +use khal::backend::GpuTimestamps; +use nexus_viewer2d::NexusViewer; +use nexus2d::mpm::solver::{BoundaryCondition, Particle, ParticleModel, SimulationParams}; +use nexus2d::prelude::{NexusPipeline, NexusState, RbdCoupling}; + +use glamx::{Vec4, vec2}; +use rapier2d::prelude::{Collider, ColliderBuilder, Pose, RigidBodyBuilder}; + +const DENSITY: f32 = 1500.0; +const YOUNG_MODULUS: f32 = 1.0e7; +const POISSON_RATIO: f32 = 0.2; + +/// Half-width of the neck the sand drains through. +const NECK_HALF_WIDTH: f32 = 0.9; +/// Height at which the funnel walls meet the vertical chamber walls. +const FUNNEL_TOP: f32 = 12.0; +/// Half-width of both chambers. +const CHAMBER_HALF_WIDTH: f32 = 12.0; + +fn insert_boundary( + state: &mut NexusState, + viewer: &mut NexusViewer, + collider: Collider, + center: glamx::Vec2, + angle: f32, +) { + let shape = collider.shared_shape().clone(); + let body = RigidBodyBuilder::fixed() + .translation(center) + .rotation(angle) + .build(); + let coupling = RbdCoupling::MpmOneWay(BoundaryCondition::separate(0.6)); + let handle = state.insert_rigid_body(body, collider, coupling); + viewer.insert_shape_with_color( + handle, + &shape, + Pose::IDENTITY, + Vec4::new(0.4, 0.6, 0.8, 0.8), + ); +} + +/// Inserts a bar spanning `start` -> `end` with the given half-thickness. +fn insert_bar( + state: &mut NexusState, + viewer: &mut NexusViewer, + start: glamx::Vec2, + end: glamx::Vec2, + half_thickness: f32, +) { + let delta = end - start; + let half_len = delta.length() / 2.0; + let angle = delta.y.atan2(delta.x); + let center = (start + end) / 2.0; + insert_boundary( + state, + viewer, + ColliderBuilder::cuboid(half_len, half_thickness).build(), + center, + angle, + ); +} + +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> anyhow::Result { + let mut state = NexusState::default(); + + let cell_width = 0.2; + let radius = cell_width / 4.0; + let spacing = radius * 2.0; + + /* + * Sand column filling the upper chamber. + */ + let mut particles = vec![]; + let fill_min = vec2(-CHAMBER_HALF_WIDTH + 1.0, FUNNEL_TOP + 1.0); + let fill_max = vec2(CHAMBER_HALF_WIDTH - 1.0, FUNNEL_TOP + 23.0); + let nx = ((fill_max.x - fill_min.x) / spacing) as i32; + let ny = ((fill_max.y - fill_min.y) / spacing) as i32; + let model = ParticleModel::sand(YOUNG_MODULUS, POISSON_RATIO); + // One group per row, shaded by height, so the draining order stays + // readable: the top of the pile ends up on top of the heap below. + let shades: Vec<_> = (0..ny) + .map(|j| { + let t = j as f32 / ny as f32; + Vec4::new(0.95, 0.75 - 0.35 * t, 0.35 - 0.25 * t, 1.0) + }) + .collect(); + viewer.set_particle_group_colors(&shades); + for i in 0..nx { + for j in 0..ny { + let position = fill_min + vec2(i as f32 + 0.5, j as f32 + 0.5) * spacing; + particles.push(Particle::with_group( + position, radius, DENSITY, model, j as u32, + )); + } + } + + let params = SimulationParams { + gravity: vec2(0.0, -9.81), + padding: 0.0, + dt: 1.0 / 60.0, + }; + state.set_mpm_params(viewer.backend(), params, cell_width)?; + state.set_mpm_substeps(15); + state.add_particles(viewer.backend(), particles)?; + + /* + * Hourglass walls. + */ + let thickness = 0.4; + // Upper funnel. + insert_bar( + &mut state, + viewer, + vec2(NECK_HALF_WIDTH, 0.0), + vec2(CHAMBER_HALF_WIDTH, FUNNEL_TOP), + thickness, + ); + insert_bar( + &mut state, + viewer, + vec2(-NECK_HALF_WIDTH, 0.0), + vec2(-CHAMBER_HALF_WIDTH, FUNNEL_TOP), + thickness, + ); + // Upper chamber sides. + for side in [-1.0f32, 1.0] { + insert_bar( + &mut state, + viewer, + vec2(side * CHAMBER_HALF_WIDTH, FUNNEL_TOP), + vec2(side * CHAMBER_HALF_WIDTH, FUNNEL_TOP + 26.0), + thickness, + ); + } + // Lower chamber sides and floor. + for side in [-1.0f32, 1.0] { + insert_bar( + &mut state, + viewer, + vec2(side * CHAMBER_HALF_WIDTH, -18.0), + vec2(side * CHAMBER_HALF_WIDTH, 0.0), + thickness, + ); + } + insert_bar( + &mut state, + viewer, + vec2(-CHAMBER_HALF_WIDTH, -18.0), + vec2(CHAMBER_HALF_WIDTH, -18.0), + thickness, + ); + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + state.finalize(viewer.backend()).await?; + + while viewer.render_frame().await { + if viewer.simulating() { + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} diff --git a/crates/examples2d/mpm_sand2.rs b/crates/examples2d/mpm_sand2.rs new file mode 100644 index 00000000..2b994ec6 --- /dev/null +++ b/crates/examples2d/mpm_sand2.rs @@ -0,0 +1,161 @@ +use khal::backend::GpuTimestamps; +use nexus_viewer2d::NexusViewer; +use nexus2d::mpm::solver::{BoundaryCondition, Particle, ParticleModel, SimulationParams}; +use nexus2d::prelude::{NexusPipeline, NexusState, RbdCoupling}; + +use glamx::{Vec2, Vec4}; +use rapier2d::prelude::{Collider, ColliderBuilder, Pose, RigidBody, RigidBodyBuilder}; + +/// Inserts a boundary collider coupled (one-way) to the MPM particles and +/// registers it for rendering. +fn insert_boundary( + state: &mut NexusState, + viewer: &mut NexusViewer, + body: RigidBody, + collider: Collider, +) { + let shape = collider.shared_shape().clone(); + let friction = BoundaryCondition::separate(1.0); + let handle = state.insert_rigid_body(body, collider, RbdCoupling::MpmOneWay(friction)); + viewer.insert_shape(handle, &shape, Pose::IDENTITY); +} + +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> anyhow::Result { + let mut state = NexusState::default(); + + let offset_y = 46.0; + // let cell_width = 0.1; + let cell_width = 0.2; + let ny = 700; + + let shades: Vec<_> = (0..ny) + .map(|j| { + let t = j as f32 / ny as f32; + Vec4::new(0.95, 0.75 - 0.35 * t, 0.35 - 0.25 * t, 1.0) + }) + .rev() + .collect(); + viewer.set_particle_group_colors(&shades); + + let mut particles = vec![]; + for i in 0..700 { + for j in 0..ny { + let position = + glamx::vec2(i as f32 + 0.5, j as f32 + 0.5) * cell_width / 2.0 + Vec2::Y * offset_y; + let density = 1000.0; + let radius = cell_width / 4.0; + let young_modulus = 1.0e7; + let poisson_ratio = 0.2; + let model = ParticleModel::sand(young_modulus, poisson_ratio); + + particles.push(Particle::with_group(position, radius, density, model, j)); + } + } + + let params = SimulationParams { + gravity: glamx::vec2(0.0, -9.81), + padding: 0.0, + dt: 1.0 / 60.0, + }; + state.set_mpm_params(viewer.backend(), params, cell_width)?; + state.set_mpm_substeps(10); + state.add_particles(viewer.backend(), particles)?; + + const ANGVEL: f32 = 2.0; + + /* + * Static platforms. + */ + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(glamx::vec2(35.0, -1.0)) + .build(), + ColliderBuilder::cuboid(42.0, 1.0).build(), + ); + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(glamx::vec2(-25.0, 45.0)) + .rotation(0.5) + .build(), + ColliderBuilder::cuboid(1.0, 52.0).build(), + ); + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(glamx::vec2(95.0, 45.0)) + .rotation(-0.5) + .build(), + ColliderBuilder::cuboid(1.0, 52.0).build(), + ); + + /* + * Rotating platforms. + */ + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::kinematic_velocity_based() + .translation(glamx::vec2(5.0, 35.0)) + .angvel(ANGVEL) + .build(), + ColliderBuilder::cuboid(1.0, 10.0).build(), + ); + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::kinematic_velocity_based() + .translation(glamx::vec2(35.0, 35.0)) + .angvel(-ANGVEL) + .build(), + ColliderBuilder::cuboid(10.0, 1.0).build(), + ); + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::kinematic_velocity_based() + .translation(glamx::vec2(65.0, 35.0)) + .angvel(ANGVEL) + .build(), + ColliderBuilder::cuboid(1.0, 10.0).build(), + ); + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::kinematic_velocity_based() + .translation(glamx::vec2(20.0, 20.0)) + .angvel(-ANGVEL) + .build(), + ColliderBuilder::ball(5.0).build(), + ); + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::kinematic_velocity_based() + .translation(glamx::vec2(50.0, 20.0)) + .angvel(-ANGVEL) + .build(), + ColliderBuilder::capsule_y(5.0, 3.0).build(), + ); + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + state.finalize(viewer.backend()).await?; + + while viewer.render_frame().await { + if viewer.simulating() { + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} diff --git a/crates/examples2d/mpm_snowball2.rs b/crates/examples2d/mpm_snowball2.rs new file mode 100644 index 00000000..4623f83c --- /dev/null +++ b/crates/examples2d/mpm_snowball2.rs @@ -0,0 +1,158 @@ +//! Snowballs thrown into each other and into a snow bank. + +use khal::backend::GpuTimestamps; +use nexus_viewer2d::NexusViewer; +use nexus2d::mpm::solver::{BoundaryCondition, Particle, ParticleModel, SimulationParams}; +use nexus2d::prelude::{NexusPipeline, NexusState, RbdCoupling}; + +use glamx::{Vec2, Vec4, vec2}; +use rapier2d::prelude::{ColliderBuilder, Pose, RigidBodyBuilder}; + +const DENSITY: f32 = 400.0; +const YOUNG_MODULUS: f32 = 1.4e6; +const POISSON_RATIO: f32 = 0.2; + +const BANK: u32 = 0; +const BLUE_BALL: u32 = 1; +const RED_BALL: u32 = 2; +const GREEN_BALL: u32 = 3; + +/// Adds a disc of snow particles centered on `center`, moving at `velocity`. +#[allow(clippy::too_many_arguments)] +fn add_snowball( + particles: &mut Vec, + center: Vec2, + velocity: Vec2, + group_id: u32, + particle_radius: f32, + snowball_radius: f32, + spacing: f32, + model: ParticleModel, +) { + let n = (snowball_radius / spacing).ceil() as i32; + for i in -n..=n { + for j in -n..=n { + let offset = vec2(i as f32, j as f32) * spacing; + if offset.length() > snowball_radius { + continue; + } + let mut particle = + Particle::with_group(center + offset, particle_radius, DENSITY, model, group_id); + particle.dynamics.velocity = velocity; + particles.push(particle); + } + } +} + +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> anyhow::Result { + let mut state = NexusState::default(); + + let cell_width = 0.2; + let radius = cell_width / 4.0; + let spacing = radius * 2.0; + let model = ParticleModel::snow(YOUNG_MODULUS, POISSON_RATIO); + + // Indexed by BANK / BLUE_BALL / RED_BALL / GREEN_BALL. + viewer.set_particle_group_colors(&[ + Vec4::new(0.70, 0.73, 0.78, 1.0), + Vec4::new(0.35, 0.55, 0.85, 1.0), + Vec4::new(0.85, 0.45, 0.35, 1.0), + Vec4::new(0.45, 0.75, 0.45, 1.0), + ]); + + let mut particles = vec![]; + + /* + * A loose snow bank for the balls to land in. + */ + let bank_half_width = 24.0; + let bank_height = 20.0; + let nx = (bank_half_width * 2.0 / spacing) as i32; + let ny = (bank_height / spacing) as i32; + for i in 0..nx { + for j in 0..ny { + let position = vec2( + -bank_half_width + (i as f32 + 0.5) * spacing, + 0.2 + (j as f32 + 0.5) * spacing, + ); + particles.push(Particle::with_group(position, radius, DENSITY, model, BANK)); + } + } + + /* + * Two balls thrown at each other above the bank, and one dropped straight + * down onto the collision point. + */ + add_snowball( + &mut particles, + vec2(-16.0, 30.0), + vec2(18.0, 0.0), + BLUE_BALL, + radius, + 3.0, + spacing, + model, + ); + add_snowball( + &mut particles, + vec2(16.0, 30.0), + vec2(-18.0, 0.0), + RED_BALL, + radius, + 6.0, + spacing, + model, + ); + add_snowball( + &mut particles, + vec2(0.0, 45.0), + vec2(0.0, -12.0), + GREEN_BALL, + radius, + 9.0, + spacing, + model, + ); + + let params = SimulationParams { + gravity: vec2(0.0, -9.81), + padding: 0.0, + dt: 1.0 / 60.0, + }; + state.set_mpm_params(viewer.backend(), params, cell_width)?; + state.set_mpm_substeps(25); + state.add_particles(viewer.backend(), particles)?; + + /* + * Ground and side walls. + */ + let coupling = RbdCoupling::MpmOneWay(BoundaryCondition::separate(1.0)); + for (pos, half_extents) in [ + (vec2(0.0, -1.0), vec2(bank_half_width + 2.0, 1.0)), + (vec2(-bank_half_width - 1.0, 11.0), vec2(1.0, 13.0)), + (vec2(bank_half_width + 1.0, 11.0), vec2(1.0, 13.0)), + ] { + let collider = ColliderBuilder::cuboid(half_extents.x, half_extents.y).build(); + let shape = collider.shared_shape().clone(); + let body = RigidBodyBuilder::fixed().translation(pos).build(); + let handle = state.insert_rigid_body(body, collider, coupling); + viewer.insert_shape(handle, &shape, Pose::IDENTITY); + } + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + state.finalize(viewer.backend()).await?; + + while viewer.render_frame().await { + if viewer.simulating() { + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} diff --git a/crates/examples3d/Cargo.toml b/crates/examples3d/Cargo.toml index 10221f4e..cb2e05ca 100644 --- a/crates/examples3d/Cargo.toml +++ b/crates/examples3d/Cargo.toml @@ -14,7 +14,7 @@ metal = ["nexus_viewer3d/metal"] [dependencies] glamx = { workspace = true } -nexus3d = { workspace = true, features = [ "rbd" ] } +nexus3d = { workspace = true, features = [ "rbd", "mpm" ] } nexus_viewer3d = { workspace = true } nexus_rbd3d = { workspace = true, features = ["default"] } # TODO: remove this rapier3d = { workspace = true, features = ["default"] } diff --git a/crates/examples3d/all_examples3.rs b/crates/examples3d/all_examples3.rs index ba06d419..b90582e6 100644 --- a/crates/examples3d/all_examples3.rs +++ b/crates/examples3d/all_examples3.rs @@ -2,30 +2,41 @@ use inflector::Inflector; use nexus_viewer3d::{BackendType, DemoKind, NexusViewer}; use nexus3d::prelude::{NexusPipeline, NexusPipelineMask}; -mod balls3; -mod boxes3; -mod boxes_and_balls3; -mod compound3; -mod dynamic_rbd3; -mod joint_ball3; -mod joint_fixed3; -mod joint_prismatic3; -mod joint_revolute3; -mod joint_revolute_batch3; -mod joints3; -mod keva3; -mod many_pyramids3; -mod many_pyramids_batch3; +mod rbd_balls3; +mod rbd_boxes3; +mod rbd_boxes_and_balls3; +mod rbd_compound3; +mod rbd_dynamic3; +mod rbd_joint_ball3; +mod rbd_joint_fixed3; +mod rbd_joint_prismatic3; +mod rbd_joint_revolute3; +mod rbd_joint_revolute_batch3; +mod rbd_joints3; +mod rbd_keva3; +mod rbd_many_pyramids3; +mod rbd_many_pyramids_batch3; // The robot loaders read URDF/MJCF assets from the filesystem, so they are // native-only. #[cfg(not(target_arch = "wasm32"))] -mod mujoco_menagerie3; -mod multibody_pendulum3; -mod primitives3; -mod pyramid3; -mod trimesh3; +mod rbd_mujoco_menagerie3; +mod rbd_multibody_pendulum3; +mod rbd_primitives3; +mod rbd_pyramid3; +mod rbd_trimesh3; #[cfg(not(target_arch = "wasm32"))] -mod urdf3; +mod rbd_urdf3; + +// MPM examples. +mod mpm_centilever_beam3; +// mod mpm_dam_break3; +mod mpm_elastic_cut3; +mod mpm_emitter3; +mod mpm_erosion3; +mod mpm_heightfield3; +mod mpm_jelly_drop3; +mod mpm_sand3; +mod mpm_snow3; /// Declares the demo registry: a `(name, kind)` list for the picker UI and a /// name -> `run()` dispatcher. Keeping both in one macro keeps them in sync. @@ -65,28 +76,38 @@ macro_rules! demos { } demos! { - "Balls" => Rbd : balls3, - "Boxes" => Rbd : boxes3, - "Boxes & balls" => Rbd : boxes_and_balls3, - "Compound" => Rbd : compound3, - "Dynamic insertion" => Rbd : dynamic_rbd3, - "Primitives" => Rbd : primitives3, - "Pyramid" => Rbd : pyramid3, - "Many pyramids" => Rbd : many_pyramids3, - "Many pyramids (batched)" => Rbd : many_pyramids_batch3, - "Keva tower" => Rbd : keva3, - "Joints (multibody)" => Rbd : joints3, - "Joints (Spherical)" => Rbd : joint_ball3, - "Joints (Fixed)" => Rbd : joint_fixed3, - "Joints (Prismatic)" => Rbd : joint_prismatic3, - "Joints (Revolute)" => Rbd : joint_revolute3, - "Joints (Revolute - Batched)" => Rbd : joint_revolute_batch3, - "Multibody (Pendulum)" => Rbd : multibody_pendulum3, - "Trimesh" => Rbd : trimesh3, + "Balls" => Rbd : rbd_balls3, + "Boxes" => Rbd : rbd_boxes3, + "Boxes & balls" => Rbd : rbd_boxes_and_balls3, + "Compound" => Rbd : rbd_compound3, + "Dynamic insertion" => Rbd : rbd_dynamic3, + "Primitives" => Rbd : rbd_primitives3, + "Pyramid" => Rbd : rbd_pyramid3, + "Many pyramids" => Rbd : rbd_many_pyramids3, + "Many pyramids (batched)" => Rbd : rbd_many_pyramids_batch3, + "Keva tower" => Rbd : rbd_keva3, + "Joints (multibody)" => Rbd : rbd_joints3, + "Joints (Spherical)" => Rbd : rbd_joint_ball3, + "Joints (Fixed)" => Rbd : rbd_joint_fixed3, + "Joints (Prismatic)" => Rbd : rbd_joint_prismatic3, + "Joints (Revolute)" => Rbd : rbd_joint_revolute3, + "Joints (Revolute - Batched)" => Rbd : rbd_joint_revolute_batch3, + "Multibody (Pendulum)" => Rbd : rbd_multibody_pendulum3, + "Trimesh" => Rbd : rbd_trimesh3, #[cfg(not(target_arch = "wasm32"))] - "URDF (multibody)" => Rbd : urdf3, + "URDF (multibody)" => Rbd : rbd_urdf3, #[cfg(not(target_arch = "wasm32"))] - "MuJoCo Menagerie" => Rbd : mujoco_menagerie3, + "MuJoCo Menagerie" => Rbd : rbd_mujoco_menagerie3, + // MPM demos. + "Cantilever beam" => Mpm : mpm_centilever_beam3, + "Sand" => Mpm : mpm_sand3, + "Sand emitter" => Mpm : mpm_emitter3, + "Heightfield" => Mpm : mpm_heightfield3, + "Elastic cut" => Mpm : mpm_elastic_cut3, + // "Dam break" => Mpm : mpm_dam_break3, + "Erosion" => Mpm : mpm_erosion3, + "Snow" => Mpm : mpm_snow3, + "Jelly drop" => Mpm : mpm_jelly_drop3, } struct CliOptions { diff --git a/crates/examples3d/mpm_centilever_beam3.rs b/crates/examples3d/mpm_centilever_beam3.rs new file mode 100644 index 00000000..e447113c --- /dev/null +++ b/crates/examples3d/mpm_centilever_beam3.rs @@ -0,0 +1,76 @@ +use khal::backend::GpuTimestamps; +use nexus_viewer3d::NexusViewer; +use nexus3d::mpm::solver::{BoundaryCondition, Particle, ParticleModel, SimulationParams}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; + +use glamx::vec3; +use rapier3d::prelude::{ColliderBuilder, Pose, RigidBodyBuilder}; + +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> anyhow::Result { + let mut state = NexusState::default(); + let coupling = RbdCoupling::MpmOneWay(BoundaryCondition::stick()); + + let width = 10.0; + let height = 2.0; + let fixed_part = 1.0; + let cell_width = 0.2; + let particle_per_cell_dim = 2; + let young_modulus = 1.0e7; + let poisson_ratio = 0.3; + + let diameter = cell_width / particle_per_cell_dim as f32; + let ni = ((width + fixed_part) / diameter).ceil() as usize; + let njk = (height / diameter).ceil() as usize; + + let mut particles = vec![]; + for i in 0..ni { + for j in 0..njk { + for k in 0..njk { + let position = vec3(i as f32, j as f32, k as f32) * diameter; + let density = 1000.0; + let radius = diameter / 2.0; + let model = ParticleModel::elastic_neo_hookean(young_modulus, poisson_ratio); + let mut particle = Particle::new(position, radius, density, model); + particle.dynamics.set_damping(2.0); + particles.push(particle); + } + } + } + + let params = SimulationParams { + gravity: vec3(0.0, -9.81, 0.0), + dt: 1.0 / 60.0, + }; + state.set_mpm_params(viewer.backend(), params, cell_width)?; + state.set_mpm_substeps(20); + state.add_particles(viewer.backend(), particles)?; + + // Fixed block that clamps one end of the beam. + let body = RigidBodyBuilder::fixed() + .translation(vec3(0.0, height / 2.0, height / 2.0)) + .build(); + let collider = ColliderBuilder::cuboid(fixed_part, height, height).build(); + let shape = collider.shared_shape().clone(); + let handle = state.insert_rigid_body(body, collider, coupling); + viewer.insert_shape(handle, &shape, Pose::IDENTITY); + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + viewer + .scene3d_mut() + .add_directional_light(glamx::Vec3::new(1.0, -2.0, 3.0)); + state.finalize(viewer.backend()).await?; + + while viewer.render_frame().await { + if viewer.simulating() { + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} diff --git a/crates/examples3d/mpm_dam_break3.rs b/crates/examples3d/mpm_dam_break3.rs new file mode 100644 index 00000000..bc7e7b56 --- /dev/null +++ b/crates/examples3d/mpm_dam_break3.rs @@ -0,0 +1,155 @@ +//! A three-dimensional dam break running past a square column. +//! +//! The obstacle is offset from the tank centerline, so the surge wraps around it +//! asymmetrically and the two arms collide again downstream. + +use khal::backend::GpuTimestamps; +use nexus_viewer3d::NexusViewer; +use nexus3d::mpm::solver::{BoundaryCondition, Particle, ParticleModel, SimulationParams}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; + +use glamx::{Vec4, vec3}; +use rapier3d::prelude::{Collider, ColliderBuilder, Pose, RigidBody, RigidBodyBuilder}; + +const DENSITY: f32 = 1000.0; +/// Deepest the water gets. The bulk modulus follows from it. +const MAX_DEPTH: f32 = 16.0; +/// See the note in the 2D dam break: the softest fluid whose compression at that +/// depth is still invisible. +const MAX_COMPRESSION: f32 = 0.01; + +const TANK_HALF_X: f32 = 30.0; +const TANK_HALF_Z: f32 = 14.0; +const TANK_HEIGHT: f32 = 28.0; + +fn insert_boundary( + state: &mut NexusState, + viewer: &mut NexusViewer, + body: RigidBody, + collider: Collider, +) { + let shape = collider.shared_shape().clone(); + let coupling = RbdCoupling::MpmOneWay(BoundaryCondition::separate(0.0)); + let handle = state.insert_rigid_body(body, collider, coupling); + viewer.insert_shape(handle, &shape, Pose::IDENTITY); +} + +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> anyhow::Result { + let mut state = NexusState::default(); + + let cell_width = 0.6; + let radius = cell_width / 4.0; + let spacing = radius * 2.0; + let model = ParticleModel::water_for_depth(DENSITY, 9.81, MAX_DEPTH, MAX_COMPRESSION); + + /* + * The water column, held against the -X wall. + */ + let mut particles = vec![]; + let nx = (14.0 / spacing) as i32; + let ny = (16.0 / spacing) as i32; + let nz = (TANK_HALF_Z * 2.0 / spacing) as i32; + // One group per row, shaded by initial height, so the overturning of the + // surge front stays visible. + let shades: Vec<_> = (0..ny) + .map(|j| { + let t = j as f32 / ny as f32; + Vec4::new(0.10 + 0.35 * t, 0.45 + 0.35 * t, 0.85 + 0.15 * t, 1.0) + }) + .collect(); + viewer.set_particle_group_colors(&shades); + for i in 0..nx { + for j in 0..ny { + for k in 0..nz { + let position = vec3( + -TANK_HALF_X + 0.5 + (i as f32 + 0.5) * spacing, + 0.5 + (j as f32 + 0.5) * spacing, + -TANK_HALF_Z + 0.5 + (k as f32 + 0.5) * spacing, + ); + particles.push(Particle::with_group( + position, radius, DENSITY, model, j as u32, + )); + } + } + } + + let params = SimulationParams { + gravity: vec3(0.0, -9.81, 0.0), + dt: 1.0 / 60.0, + }; + state.set_mpm_params(viewer.backend(), params, cell_width)?; + state.set_mpm_substeps(25); + state.add_particles(viewer.backend(), particles)?; + + /* + * Tank. + */ + let walls_color = Vec4::new(0.6, 0.8, 1.0, 0.04); + let walls = [ + ( + vec3(0.0, -1.0, 0.0), + vec3(TANK_HALF_X + 3.0, 1.0, TANK_HALF_Z + 3.0), + ), + ( + vec3(0.0, TANK_HEIGHT + 1.0, 0.0), + vec3(TANK_HALF_X + 3.0, 1.0, TANK_HALF_Z + 3.0), + ), + ( + vec3(-TANK_HALF_X - 1.0, TANK_HEIGHT / 2.0, 0.0), + vec3(1.0, TANK_HEIGHT / 2.0 + 2.0, TANK_HALF_Z + 3.0), + ), + ( + vec3(TANK_HALF_X + 1.0, TANK_HEIGHT / 2.0, 0.0), + vec3(1.0, TANK_HEIGHT / 2.0 + 2.0, TANK_HALF_Z + 3.0), + ), + ( + vec3(0.0, TANK_HEIGHT / 2.0, -TANK_HALF_Z - 1.0), + vec3(TANK_HALF_X + 3.0, TANK_HEIGHT / 2.0 + 2.0, 1.0), + ), + ( + vec3(0.0, TANK_HEIGHT / 2.0, TANK_HALF_Z + 1.0), + vec3(TANK_HALF_X + 3.0, TANK_HEIGHT / 2.0 + 3.0, 1.0), + ), + ]; + for (pos, half_extents) in walls { + let collider = + ColliderBuilder::cuboid(half_extents.x, half_extents.y, half_extents.z).build(); + let shape = collider.shared_shape().clone(); + let body = RigidBodyBuilder::fixed().translation(pos).build(); + let coupling = RbdCoupling::MpmOneWay(BoundaryCondition::separate(0.0)); + let handle = state.insert_rigid_body(body, collider, coupling); + viewer.insert_shape_with_color(handle, &shape, Pose::IDENTITY, walls_color); + } + + /* + * Column for the surge to wrap around. + */ + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(vec3(2.0, 5.0, 0.0)) + .build(), + ColliderBuilder::cuboid(2.0, 5.0, 6.0).build(), + ); + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + viewer + .scene3d_mut() + .add_directional_light(glamx::Vec3::new(1.0, -2.0, 3.0)); + state.finalize(viewer.backend()).await?; + + while viewer.render_frame().await { + if viewer.simulating() { + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} diff --git a/crates/examples3d/mpm_elastic_cut3.rs b/crates/examples3d/mpm_elastic_cut3.rs new file mode 100644 index 00000000..68c11f6d --- /dev/null +++ b/crates/examples3d/mpm_elastic_cut3.rs @@ -0,0 +1,89 @@ +use khal::backend::GpuTimestamps; +use nexus_viewer3d::NexusViewer; +use nexus3d::mpm::solver::{BoundaryCondition, Particle, ParticleModel, SimulationParams}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; + +use glamx::{Pose3, vec3}; +use rapier3d::parry::utils::Array2; +use rapier3d::prelude::{ColliderBuilder, HeightField, Pose, RigidBodyBuilder, TriMeshFlags}; + +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> anyhow::Result { + let mut state = NexusState::default(); + let coupling = RbdCoupling::MpmOneWay(BoundaryCondition::separate(1.0)); + + let nxz = 50; + let cell_width = 1.0; + + let mut particles = vec![]; + for i in 0..nxz { + for j in 0..30 { + for k in 0..nxz { + let position = vec3( + i as f32 + 0.5 - nxz as f32 / 2.0, + j as f32 + 0.5 + 60.0, + k as f32 + 0.5 - nxz as f32 / 2.0, + ) * cell_width + / 2.0; + let density = 2700.0; + let radius = cell_width / 4.0; + let model = ParticleModel::elastic(1.0e7, 0.2); + particles.push(Particle::new(position, radius, density, model)); + } + } + } + + let params = SimulationParams { + gravity: vec3(0.0, -9.81, 0.0) * 4.0, + dt: 1.0 / 60.0, + }; + state.set_mpm_params(viewer.backend(), params, cell_width)?; + state.set_mpm_substeps(20); + state.add_particles(viewer.backend(), particles)?; + + // Floor + let body = RigidBodyBuilder::fixed() + .translation(vec3(0.0, -4.0, 0.0)) + .build(); + let collider = ColliderBuilder::cuboid(100.0, 1.0, 100.0).build(); + let shape = collider.shared_shape().clone(); + let handle = state.insert_rigid_body(body, collider, coupling); + viewer.insert_shape(handle, &shape, Pose::IDENTITY); + + // Cutting planes (3 heightfield trimeshes) + for k in 0..3 { + let heights = Array2::zeros(10, 10); + let heightfield = HeightField::new(heights, vec3(35.0, 1.0, 10.0)); + let (mut vtx, idx) = heightfield.to_trimesh(); + vtx.iter_mut().for_each(|pt| { + *pt = + Pose3::rotation(vec3(1.3, 0.0, 0.0)) * *pt + vec3(0.0, 10.0, k as f32 * 10.0 - 10.0) + }); + let body = RigidBodyBuilder::fixed().build(); + let collider = ColliderBuilder::trimesh_with_flags(vtx, idx, TriMeshFlags::ORIENTED) + .unwrap() + .build(); + let shape = collider.shared_shape().clone(); + let handle = state.insert_rigid_body(body, collider, coupling); + viewer.insert_shape(handle, &shape, Pose::IDENTITY); + } + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + viewer + .scene3d_mut() + .add_directional_light(glamx::Vec3::new(1.0, -2.0, 3.0)); + state.finalize(viewer.backend()).await?; + + while viewer.render_frame().await { + if viewer.simulating() { + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} diff --git a/crates/examples3d/mpm_emitter3.rs b/crates/examples3d/mpm_emitter3.rs new file mode 100644 index 00000000..6c164559 --- /dev/null +++ b/crates/examples3d/mpm_emitter3.rs @@ -0,0 +1,134 @@ +use khal::backend::GpuTimestamps; +use nexus_viewer3d::NexusViewer; +use nexus3d::mpm::solver::{BoundaryCondition, Particle, ParticleModel, SimulationParams}; +use nexus3d::prelude::{NexusParticleChunk, NexusPipeline, NexusState, RbdCoupling}; + +use glamx::{Vec4, vec3}; +use rapier3d::prelude::{ColliderBuilder, Pose, RigidBodyBuilder}; + +use std::collections::VecDeque; + +const DENSITY: f32 = 2700.0; +const YOUNG_MODULUS: f32 = 2.0e8; +const POISSON_RATIO: f32 = 0.2; + +/// Emit a fresh patch of particles every `EMIT_EVERY` substeps. +const EMIT_EVERY: u64 = 10; +/// Edge length (in particles) of the emitted cube. +const EMIT_BLOCK: i32 = 30; +const EMIT_BLOCK_Y: i32 = 2; +/// Cap on the live particle count. Once reached, every emit removes as many of +/// the oldest particles as it adds, so the total stays at this budget (the +/// settled sand erodes behind the orbiting emitter like a comet tail). +const MAX_PARTICLES: usize = 250_000; + +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> anyhow::Result { + let mut state = NexusState::default(); + // MPM boundary colliders are inserted as rigid bodies coupled (one-way) to + // the continuum: they push the particles but aren't pushed back. + let coupling = RbdCoupling::MpmOneWay(BoundaryCondition::separate(1.0)); + + let cell_width = 1.0; + let dt = 1.0 / 60.0; + + let params = SimulationParams { + gravity: vec3(0.0, -9.81, 0.0), + dt, + }; + state.set_mpm_params(viewer.backend(), params, cell_width)?; + state.set_mpm_substeps(10); + // No particles up-front: they are emitted dynamically in the loop below. + + /* + * Boundary colliders: a floor and four walls forming an open box that + * catches the falling sand. + */ + let thickness = 0.5; + let walls_color = Vec4::new(0.6, 0.8, 1.0, 0.3); + let walls = [ + (vec3(0.0, -thickness, 0.0), vec3(30.0, thickness, 30.0)), + (vec3(0.0, 10.0, -30.0), vec3(30.0, 10.0, thickness)), + (vec3(0.0, 10.0, 30.0), vec3(30.0, 10.0, thickness)), + (vec3(-30.0, 10.0, 0.0), vec3(thickness, 10.0, 30.0)), + (vec3(30.0, 10.0, 0.0), vec3(thickness, 10.0, 30.0)), + ]; + for (pos, half_extents) in walls { + let body = RigidBodyBuilder::fixed().translation(pos).build(); + let collider = + ColliderBuilder::cuboid(half_extents.x, half_extents.y, half_extents.z).build(); + let shape = collider.shared_shape().clone(); + let handle = state.insert_rigid_body(body, collider, coupling); + viewer.insert_shape_with_color(handle, &shape, Pose::IDENTITY, walls_color); + } + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + viewer + .scene3d_mut() + .add_directional_light(glamx::Vec3::new(1.0, -2.0, 3.0)); + state.finalize(viewer.backend()).await?; + + /* + * Dynamic emitter: a small cube of sand spawned at a point that orbits the + * center, so the stream paints a moving ring of sand into the box. + */ + let radius = cell_width / 4.0; + let spacing = radius * 2.0; + let model = ParticleModel::sand(YOUNG_MODULUS, POISSON_RATIO); + let emit_height = 40.0; + let orbit_radius = 10.0; + let angular_speed = 1.5; // rad/s + + // Oldest-first queue of live chunks paired with their particle counts, plus + // the running total `MAX_PARTICLES` is enforced against. + let mut chunks: VecDeque<(NexusParticleChunk, usize)> = VecDeque::new(); + let mut total_particles: usize = 0; + let mut t: f32 = 0.0; + let mut step: u64 = 0; + + while viewer.render_frame().await { + if viewer.simulating() { + if step.is_multiple_of(EMIT_EVERY) && total_particles < MAX_PARTICLES { + let angle = t * angular_speed; + let center = vec3( + orbit_radius * angle.cos(), + emit_height, + orbit_radius * angle.sin(), + ); + + let mut particles = + Vec::with_capacity((EMIT_BLOCK as usize).pow(2) * EMIT_BLOCK_Y as usize); + for i in 0..EMIT_BLOCK { + for j in 0..EMIT_BLOCK_Y { + for k in 0..EMIT_BLOCK { + let offset = vec3( + (i - EMIT_BLOCK / 2) as f32, + (j - EMIT_BLOCK_Y / 2) as f32, + (k - EMIT_BLOCK / 2) as f32, + ) * spacing; + let mut particle = + Particle::new(center + offset, radius, DENSITY, model); + particle.dynamics.velocity = vec3(0.0, -8.0, 0.0); + particles.push(particle); + } + } + } + let n = particles.len(); + let chunk = state.add_particles(viewer.backend(), particles)?; + chunks.push_back((chunk, n)); + total_particles += n; + } + + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + t += dt; + step += 1; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} diff --git a/crates/examples3d/mpm_erosion3.rs b/crates/examples3d/mpm_erosion3.rs new file mode 100644 index 00000000..8ec16cff --- /dev/null +++ b/crates/examples3d/mpm_erosion3.rs @@ -0,0 +1,181 @@ +//! Water poured onto a cohesive sand mound, washing it away. + +use khal::backend::GpuTimestamps; +use nexus_viewer3d::NexusViewer; +use nexus3d::mpm::solver::{BoundaryCondition, Particle, ParticleModel, SimulationParams}; +use nexus3d::prelude::{NexusParticleChunk, NexusPipeline, NexusState, RbdCoupling}; + +use glamx::{Vec4, vec3}; +use rapier3d::prelude::{ColliderBuilder, Pose, RigidBodyBuilder}; + +use std::collections::VecDeque; + +const SAND_DENSITY: f32 = 1800.0; +const SAND_YOUNG_MODULUS: f32 = 1.0e7; +const SAND_POISSON_RATIO: f32 = 0.2; +/// Cohesive shear strength of the sand (Pa). +const SAND_SHEAR_STRENGTH: f32 = 3.0e3; + +/// How much of the basin's friction the water feels. +const WATER_BOUNDARY_FRICTION: f32 = 0.05; + +const WATER_DENSITY: f32 = 1000.0; +/// Depth the pooled water reaches in the basin. The bulk modulus follows from it. +const WATER_DEPTH: f32 = 8.0; + +/// Emit a fresh slug of water every `EMIT_EVERY` steps. +const EMIT_EVERY: u64 = 3; +/// Cap on the live water particle count. +const MAX_WATER_PARTICLES: usize = 200_000; + +/// Radius of the sand mound at its base. +const MOUND_RADIUS: f32 = 14.0; +/// Height of the sand mound. +const MOUND_HEIGHT: f32 = 12.0; + +/// The sand bands take `SAND_LIGHT` and `SAND_LIGHT + 1`. +const SAND_LIGHT: u32 = 0; +const WATER: u32 = 2; + +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> anyhow::Result { + let mut state = NexusState::default(); + let coupling = RbdCoupling::MpmOneWay(BoundaryCondition::separate(0.8)); + + let cell_width = 0.6; + let radius = cell_width / 4.0; + let spacing = radius * 2.0; + let dt = 1.0 / 60.0; + + /* + * The sand mound: a cone, sampled by rejection. + */ + let sand_model = ParticleModel::cohesive_sand_with_strength( + SAND_YOUNG_MODULUS, + SAND_POISSON_RATIO, + SAND_SHEAR_STRENGTH, + ); + viewer.set_particle_group_colors(&[ + Vec4::new(0.85, 0.70, 0.45, 1.0), + Vec4::new(0.72, 0.56, 0.34, 1.0), + Vec4::new(0.25, 0.55, 0.95, 1.0), + ]); + + let mut particles = vec![]; + let n = (MOUND_RADIUS / spacing).ceil() as i32; + let ny = (MOUND_HEIGHT / spacing).ceil() as i32; + for i in -n..=n { + for j in 0..ny { + for k in -n..=n { + let y = (j as f32 + 0.5) * spacing; + // Cone: the allowed radius shrinks linearly with height. + let max_radius = MOUND_RADIUS * (1.0 - y / MOUND_HEIGHT); + let x = i as f32 * spacing; + let z = k as f32 * spacing; + if x * x + z * z > max_radius * max_radius { + continue; + } + // Horizontal bands, so the erosion depth stays readable. + let band = ((j / 6) % 2) as u32; + particles.push(Particle::with_group( + vec3(x, y, z), + radius, + SAND_DENSITY, + sand_model, + SAND_LIGHT + band, + )); + } + } + } + + let params = SimulationParams { + gravity: vec3(0.0, -9.81, 0.0), + dt, + }; + state.set_mpm_params(viewer.backend(), params, cell_width)?; + state.set_mpm_substeps(20); + state.add_particles(viewer.backend(), particles)?; + + /* + * Basin. + */ + let walls_color = Vec4::new(0.6, 0.8, 1.0, 0.3); + let basin = 20.0; + let walls = [ + (vec3(0.0, -1.0, 0.0), vec3(basin, 1.0, basin)), + (vec3(0.0, 6.0, -basin), vec3(basin, 7.0, 1.0)), + (vec3(0.0, 6.0, basin), vec3(basin, 7.0, 1.0)), + (vec3(-basin, 6.0, 0.0), vec3(1.0, 7.0, basin)), + (vec3(basin, 6.0, 0.0), vec3(1.0, 7.0, basin)), + ]; + for (pos, half_extents) in walls { + let collider = + ColliderBuilder::cuboid(half_extents.x, half_extents.y, half_extents.z).build(); + let shape = collider.shared_shape().clone(); + let body = RigidBodyBuilder::fixed().translation(pos).build(); + let handle = state.insert_rigid_body(body, collider, coupling); + viewer.insert_shape_with_color(handle, &shape, Pose::IDENTITY, walls_color); + } + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + viewer + .scene3d_mut() + .add_directional_light(glamx::Vec3::new(1.0, -2.0, 3.0)); + state.finalize(viewer.backend()).await?; + + /* + * The jet: a small slab of water spawned above the mound, aimed at the flank + * so the crater migrates instead of drilling straight down. + */ + let water_model = ParticleModel::water_for_depth(WATER_DENSITY, 9.81, WATER_DEPTH, 0.01); + let jet_extent = 5; + let mut chunks: VecDeque<(NexusParticleChunk, usize)> = VecDeque::new(); + let mut live_water = 0usize; + let mut t = 0.0f32; + let mut step = 0u64; + + while viewer.render_frame().await { + if viewer.simulating() { + if step.is_multiple_of(EMIT_EVERY) && live_water < MAX_WATER_PARTICLES { + // Sweep the jet slowly back and forth across the mound. + let center = vec3((t * 1.35).sin() * 8.0, 26.0, 4.0); + let mut water = Vec::new(); + for i in -jet_extent..=jet_extent { + for j in 0..2 { + for k in -jet_extent..=jet_extent { + let offset = vec3(i as f32, j as f32, k as f32) * spacing; + let mut particle = Particle::with_group( + center + offset, + radius, + WATER_DENSITY, + water_model, + WATER, + ); + particle.dynamics.velocity = vec3(0.0, -14.0, 0.0); + particle + .dynamics + .set_boundary_friction(WATER_BOUNDARY_FRICTION); + water.push(particle); + } + } + } + + let n = water.len(); + let chunk = state.add_particles(viewer.backend(), water)?; + chunks.push_back((chunk, n)); + live_water += n; + } + + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + t += dt; + step += 1; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} diff --git a/crates/examples3d/mpm_heightfield3.rs b/crates/examples3d/mpm_heightfield3.rs new file mode 100644 index 00000000..bf8398d3 --- /dev/null +++ b/crates/examples3d/mpm_heightfield3.rs @@ -0,0 +1,76 @@ +use khal::backend::GpuTimestamps; +use nexus_viewer3d::NexusViewer; +use nexus3d::mpm::solver::{BoundaryCondition, Particle, ParticleModel, SimulationParams}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; + +use glamx::vec3; +use rapier3d::parry::utils::Array2; +use rapier3d::prelude::{ColliderBuilder, HeightField, Pose, RigidBodyBuilder, TriMeshFlags}; + +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> anyhow::Result { + let mut state = NexusState::default(); + let coupling = RbdCoupling::MpmOneWay(BoundaryCondition::separate(1.0)); + + let nxz = 45; + let cell_width = 1.0; + + let mut particles = vec![]; + for i in 0..nxz { + for j in 0..100 { + for k in 0..nxz { + let position = vec3( + i as f32 + 0.5 - nxz as f32 / 2.0, + j as f32 + 0.5 + 14.0, + k as f32 + 0.5 - nxz as f32 / 2.0, + ) * cell_width + / 2.0; + let density = 2700.0; + let radius = cell_width / 4.0; + let model = ParticleModel::sand(2.0e9, 0.2); + particles.push(Particle::new(position, radius, density, model)); + } + } + } + + let params = SimulationParams { + gravity: vec3(0.0, -9.81, 0.0), + dt: 1.0 / 60.0, + }; + state.set_mpm_params(viewer.backend(), params, cell_width)?; + state.set_mpm_substeps(20); + state.add_particles(viewer.backend(), particles)?; + + // Sinusoidal heightfield terrain (rendered as the converted trimesh). + let heights = Array2::from_fn(200, 200, |i, j| { + (i as f32 / 10.0).sin() * (j as f32 / 10.0).cos() + }); + let heightfield = HeightField::new(heights, vec3(100.0, 5.0, 100.0)); + let (vtx, idx) = heightfield.to_trimesh(); + let body = RigidBodyBuilder::fixed().build(); + let collider = ColliderBuilder::trimesh_with_flags(vtx, idx, TriMeshFlags::ORIENTED) + .unwrap() + .build(); + let shape = collider.shared_shape().clone(); + let handle = state.insert_rigid_body(body, collider, coupling); + viewer.insert_shape(handle, &shape, Pose::IDENTITY); + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + viewer + .scene3d_mut() + .add_directional_light(glamx::Vec3::new(1.0, -2.0, 3.0)); + state.finalize(viewer.backend()).await?; + + while viewer.render_frame().await { + if viewer.simulating() { + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} diff --git a/crates/examples3d/mpm_jelly_drop3.rs b/crates/examples3d/mpm_jelly_drop3.rs new file mode 100644 index 00000000..42337593 --- /dev/null +++ b/crates/examples3d/mpm_jelly_drop3.rs @@ -0,0 +1,120 @@ +//! A row of neo-Hookean blobs of increasing stiffness dropped onto a floor. + +use khal::backend::GpuTimestamps; +use nexus_viewer3d::NexusViewer; +use nexus3d::mpm::solver::{BoundaryCondition, Particle, ParticleModel, SimulationParams}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; + +use glamx::{Vec4, vec3}; +use rapier3d::prelude::{Collider, ColliderBuilder, Pose, RigidBody, RigidBodyBuilder}; + +const DENSITY: f32 = 1000.0; +const POISSON_RATIO: f32 = 0.3; + +/// Young modulus of each blob, left to right. +const YOUNG_MODULI: [f32; 5] = [2.0e5, 5.0e5, 1.5e6, 5.0e6, 2.0e7]; +/// Radius of each blob. +const BLOB_RADIUS: f32 = 3.0; +/// Horizontal spacing between blob centers. +const BLOB_SPACING: f32 = 8.0; + +fn insert_boundary( + state: &mut NexusState, + viewer: &mut NexusViewer, + body: RigidBody, + collider: Collider, +) { + let shape = collider.shared_shape().clone(); + let coupling = RbdCoupling::MpmOneWay(BoundaryCondition::separate(0.5)); + let handle = state.insert_rigid_body(body, collider, coupling); + viewer.insert_shape(handle, &shape, Pose::IDENTITY); +} + +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> anyhow::Result { + let mut state = NexusState::default(); + + let cell_width = 0.5; + let radius = cell_width / 4.0; + let spacing = radius * 2.0; + + /* + * One blob per stiffness, sampled as a ball of particles. + */ + let mut particles = vec![]; + let n = (BLOB_RADIUS / spacing).ceil() as i32; + let count = YOUNG_MODULI.len(); + // One group per blob: the softest is warm, the stiffest is cold. + let shades: Vec<_> = (0..count) + .map(|b| { + let t = b as f32 / (count - 1) as f32; + Vec4::new(0.95 - 0.7 * t, 0.45, 0.25 + 0.65 * t, 1.0) + }) + .collect(); + viewer.set_particle_group_colors(&shades); + for (b, young_modulus) in YOUNG_MODULI.iter().enumerate() { + let center = vec3( + (b as f32 - (count - 1) as f32 / 2.0) * BLOB_SPACING, + 10.0, + 0.0, + ); + let model = ParticleModel::elastic_neo_hookean(*young_modulus, POISSON_RATIO); + + for i in -n..=n { + for j in -n..=n { + for k in -n..=n { + let offset = vec3(i as f32, j as f32, k as f32) * spacing; + if offset.length() > BLOB_RADIUS { + continue; + } + particles.push(Particle::with_group( + center + offset, + radius, + DENSITY, + model, + b as u32, + )); + } + } + } + } + + let params = SimulationParams { + gravity: vec3(0.0, -9.81, 0.0), + dt: 1.0 / 60.0, + }; + state.set_mpm_params(viewer.backend(), params, cell_width)?; + state.set_mpm_substeps(20); + state.add_particles(viewer.backend(), particles)?; + + /* + * Floor. + */ + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(vec3(0.0, -1.0, 0.0)) + .build(), + ColliderBuilder::cuboid(40.0, 1.0, 20.0).build(), + ); + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + viewer + .scene3d_mut() + .add_directional_light(glamx::Vec3::new(1.0, -2.0, 3.0)); + state.finalize(viewer.backend()).await?; + + while viewer.render_frame().await { + if viewer.simulating() { + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} diff --git a/crates/examples3d/mpm_sand3.rs b/crates/examples3d/mpm_sand3.rs new file mode 100644 index 00000000..587449c2 --- /dev/null +++ b/crates/examples3d/mpm_sand3.rs @@ -0,0 +1,101 @@ +use khal::backend::GpuTimestamps; +use nexus_viewer3d::NexusViewer; +use nexus3d::mpm::solver::{BoundaryCondition, Particle, ParticleModel, SimulationParams}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; + +use glamx::{Vec4, vec3}; +use rapier3d::prelude::{ColliderBuilder, Pose, RigidBodyBuilder}; + +const DENSITY: f32 = 2700.0; +const YOUNG_MODULUS: f32 = 2.0e9; +const POISSON_RATIO: f32 = 0.2; + +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> anyhow::Result { + let mut state = NexusState::default(); + // MPM boundary colliders are inserted as rigid bodies coupled to the + // continuum; they push the particles but aren't pushed back. + let coupling = RbdCoupling::MpmOneWay(BoundaryCondition::separate(1.0)); + + let nxz = 45; + let cell_width = 1.0; + + /* + * Sand particles. + */ + let mut particles = vec![]; + for i in 0..nxz { + for j in 0..100 { + for k in 0..nxz { + let position = vec3( + i as f32 + 0.5 - nxz as f32 / 2.0, + j as f32 + 0.5 + 10.0, + k as f32 + 0.5 - nxz as f32 / 2.0, + ) * cell_width + / 2.0; + let radius = cell_width / 4.0; + let model = ParticleModel::sand(YOUNG_MODULUS, POISSON_RATIO); + particles.push(Particle::new(position, radius, DENSITY, model)); + } + } + } + + let params = SimulationParams { + gravity: vec3(0.0, -9.81, 0.0), + dt: 1.0 / 60.0, + }; + state.set_mpm_params(viewer.backend(), params, cell_width)?; + state.set_mpm_substeps(20); + state.add_particles(viewer.backend(), particles)?; + + /* + * Boundary colliders (floor, walls, rotating blade). + */ + let thickness = 0.5; + let walls_color = Vec4::new(0.6, 0.8, 1.0, 0.3); + let walls = [ + (vec3(0.0, -4.0, 0.0), vec3(100.0, 4.0, 100.0)), + (vec3(0.0, 5.0, -35.0), vec3(35.0, 5.0, thickness)), + (vec3(0.0, 5.0, 35.0), vec3(35.0, 5.0, thickness)), + (vec3(-35.0, 5.0, 0.0), vec3(thickness, 5.0, 35.0)), + (vec3(35.0, 5.0, 0.0), vec3(thickness, 5.0, 35.0)), + ]; + for (pos, half_extents) in walls { + let body = RigidBodyBuilder::fixed().translation(pos).build(); + let collider = + ColliderBuilder::cuboid(half_extents.x, half_extents.y, half_extents.z).build(); + let shape = collider.shared_shape().clone(); + let handle = state.insert_rigid_body(body, collider, coupling); + viewer.insert_shape_with_color(handle, &shape, Pose::IDENTITY, walls_color); + } + + // Rotating blade (kinematic). + let body = RigidBodyBuilder::kinematic_velocity_based() + .translation(vec3(0.0, 2.0, 0.0)) + .rotation(vec3(0.0, 0.0, -0.5)) + .angvel(vec3(0.0, -1.0, 0.0)) + .build(); + let collider = ColliderBuilder::cuboid(thickness, 2.0, 30.0).build(); + let shape = collider.shared_shape().clone(); + let handle = state.insert_rigid_body(body, collider, coupling); + viewer.insert_shape(handle, &shape, Pose::IDENTITY); + + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + viewer + .scene3d_mut() + .add_directional_light(glamx::Vec3::new(1.0, -2.0, 3.0)); + state.finalize(viewer.backend()).await?; + + while viewer.render_frame().await { + if viewer.simulating() { + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} diff --git a/crates/examples3d/mpm_snow3.rs b/crates/examples3d/mpm_snow3.rs new file mode 100644 index 00000000..7067f89d --- /dev/null +++ b/crates/examples3d/mpm_snow3.rs @@ -0,0 +1,155 @@ +//! Snowballs fired into a snow wall, and a plough sweeping the debris. + +use khal::backend::GpuTimestamps; +use nexus_viewer3d::NexusViewer; +use nexus3d::mpm::solver::{BoundaryCondition, Particle, ParticleModel, SimulationParams}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; + +use glamx::{Vec3, Vec4, vec3}; +use rapier3d::prelude::{Collider, ColliderBuilder, Pose, RigidBody, RigidBodyBuilder}; + +const DENSITY: f32 = 400.0; +const YOUNG_MODULUS: f32 = 1.4e5; +const POISSON_RATIO: f32 = 0.2; + +const BALL_RADIUS: f32 = 3.0; + +const WALL: u32 = 0; +/// The three projectiles take `FIRST_BALL`, `FIRST_BALL + 1` and `FIRST_BALL + 2`. +const FIRST_BALL: u32 = 1; + +fn insert_boundary( + state: &mut NexusState, + viewer: &mut NexusViewer, + body: RigidBody, + collider: Collider, +) { + let shape = collider.shared_shape().clone(); + let coupling = RbdCoupling::MpmOneWay(BoundaryCondition::separate(0.7)); + let handle = state.insert_rigid_body(body, collider, coupling); + viewer.insert_shape(handle, &shape, Pose::IDENTITY); +} + +/// Adds a ball of snow particles centered on `center`, moving at `velocity`. +fn add_snowball( + particles: &mut Vec, + center: Vec3, + velocity: Vec3, + group_id: u32, + radius: f32, + spacing: f32, + model: ParticleModel, +) { + let n = (BALL_RADIUS / spacing).ceil() as i32; + for i in -n..=n { + for j in -n..=n { + for k in -n..=n { + let offset = vec3(i as f32, j as f32, k as f32) * spacing; + if offset.length() > BALL_RADIUS { + continue; + } + let mut particle = + Particle::with_group(center + offset, radius, DENSITY, model, group_id); + particle.dynamics.velocity = velocity; + particles.push(particle); + } + } + } +} + +pub async fn run( + viewer: &mut NexusViewer, + pipeline: &mut NexusPipeline, +) -> anyhow::Result { + let mut state = NexusState::default(); + + let cell_width = 0.5; + let radius = cell_width / 4.0; + let spacing = radius * 2.0; + let model = ParticleModel::snow(YOUNG_MODULUS, POISSON_RATIO); + + let mut group_colors = vec![Vec4::new(0.90, 0.93, 0.98, 1.0)]; + group_colors.extend((0..3).map(|n| { + let t = n as f32 / 2.0; + Vec4::new(0.30 + 0.55 * t, 0.55, 0.85 - 0.45 * t, 1.0) + })); + viewer.set_particle_group_colors(&group_colors); + + let mut particles = vec![]; + + /* + * A standing snow wall. + */ + let wall_half_width = 14.0; + let wall_height = 16.0; + let wall_half_depth = 3.0; + let nx = (wall_half_width * 2.0 / spacing) as i32; + let ny = (wall_height / spacing) as i32; + let nz = (wall_half_depth * 2.0 / spacing) as i32; + for i in 0..nx { + for j in 0..ny { + for k in 0..nz { + let position = vec3( + -wall_half_width + (i as f32 + 0.5) * spacing, + 0.2 + (j as f32 + 0.5) * spacing, + -wall_half_depth + (k as f32 + 0.5) * spacing, + ); + particles.push(Particle::with_group(position, radius, DENSITY, model, WALL)); + } + } + } + + /* + * Three projectiles aimed at different heights of the wall. + */ + for (n, (height, speed)) in [(4.0f32, 26.0f32), (10.0, 30.0), (14.0, 22.0)] + .into_iter() + .enumerate() + { + add_snowball( + &mut particles, + vec3((n as f32 - 1.0) * 7.0, height, 26.0), + vec3(0.0, 0.0, -speed), + FIRST_BALL + n as u32, + radius, + spacing, + model, + ); + } + + let params = SimulationParams { + gravity: vec3(0.0, -9.81, 0.0), + dt: 1.0 / 60.0, + }; + state.set_mpm_params(viewer.backend(), params, cell_width)?; + state.set_mpm_substeps(20); + state.add_particles(viewer.backend(), particles)?; + + /* + * Ground. + */ + insert_boundary( + &mut state, + viewer, + RigidBodyBuilder::fixed() + .translation(vec3(0.0, -1.0, 0.0)) + .build(), + ColliderBuilder::cuboid(40.0, 1.0, 40.0).build(), + ); + let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); + viewer + .scene3d_mut() + .add_directional_light(glamx::Vec3::new(1.0, -2.0, 3.0)); + state.finalize(viewer.backend()).await?; + + while viewer.render_frame().await { + if viewer.simulating() { + pipeline + .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) + .await?; + } + viewer.sync(&mut state, Some(&mut timestamps)).await?; + } + + Ok(state) +} From ebab1719d53da71e463c50f3a74249a12e653d48 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 14 Aug 2026 23:47:14 +0200 Subject: [PATCH 6/7] feat: expose MPM through the Python bindings --- crates/nexus_python3d/Cargo.toml | 2 +- crates/nexus_python3d/examples/balls3.py | 6 +- crates/nexus_python3d/examples/boxes3.py | 6 +- .../examples/boxes_and_balls3.py | 6 +- .../examples/centilever_beam3.py | 93 +++++++++++ crates/nexus_python3d/examples/compound3.py | 6 +- .../nexus_python3d/examples/dynamic_rbd3.py | 6 +- .../nexus_python3d/examples/elastic_cut3.py | 131 ++++++++++++++++ .../nexus_python3d/examples/heightfield3.py | 107 +++++++++++++ crates/nexus_python3d/examples/joint_ball3.py | 6 +- .../nexus_python3d/examples/joint_fixed3.py | 4 +- .../examples/joint_prismatic3.py | 6 +- .../examples/joint_revolute3.py | 6 +- .../examples/joint_revolute_batch3.py | 6 +- crates/nexus_python3d/examples/joints3.py | 3 +- crates/nexus_python3d/examples/keva3.py | 3 +- .../nexus_python3d/examples/many_pyramids3.py | 3 +- .../examples/many_pyramids_batch3.py | 6 +- .../nexus_python3d/examples/mpm_emitter3.py | 132 ++++++++++++++++ .../examples/multibody_pendulum3.py | 6 +- crates/nexus_python3d/examples/primitives3.py | 6 +- crates/nexus_python3d/examples/pyramid3.py | 3 +- crates/nexus_python3d/examples/sand3.py | 113 ++++++++++++++ crates/nexus_python3d/examples/trimesh3.py | 6 +- crates/nexus_python3d/run_examples.py | 6 + crates/nexus_python3d/src/lib.rs | 9 ++ crates/nexus_python3d/src/loaders.rs | 3 +- crates/nexus_python3d/src/mpm.rs | 137 ++++++++++++++++ crates/nexus_python3d/src/nexus.rs | 147 ++++++++++++++++-- 29 files changed, 927 insertions(+), 47 deletions(-) create mode 100644 crates/nexus_python3d/examples/centilever_beam3.py create mode 100644 crates/nexus_python3d/examples/elastic_cut3.py create mode 100644 crates/nexus_python3d/examples/heightfield3.py create mode 100644 crates/nexus_python3d/examples/mpm_emitter3.py create mode 100644 crates/nexus_python3d/examples/sand3.py create mode 100644 crates/nexus_python3d/src/mpm.rs diff --git a/crates/nexus_python3d/Cargo.toml b/crates/nexus_python3d/Cargo.toml index 2f7782ec..1114d9e3 100644 --- a/crates/nexus_python3d/Cargo.toml +++ b/crates/nexus_python3d/Cargo.toml @@ -37,7 +37,7 @@ extension-module = ["pyo3/extension-module"] [dependencies] pyo3 = { version = "0.29", features = ["abi3-py39"] } numpy = "0.29" -nexus3d = { workspace = true, features = ["rbd"] } +nexus3d = { workspace = true, features = ["rbd", "mpm"] } nexus_viewer3d = { workspace = true } rapier3d = { workspace = true, features = ["default"] } rapier3d-urdf = { workspace = true, features = ["stl", "wavefront"] } diff --git a/crates/nexus_python3d/examples/balls3.py b/crates/nexus_python3d/examples/balls3.py index c72c8d13..77900880 100644 --- a/crates/nexus_python3d/examples/balls3.py +++ b/crates/nexus_python3d/examples/balls3.py @@ -10,6 +10,7 @@ NexusViewer, NexusPipeline, NexusState, + RbdCoupling, RigidBodyBuilder, ColliderBuilder, GpuTimestamps, @@ -24,6 +25,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: state = NexusState() + no_coupling = RbdCoupling.NONE # Floor made of large cuboids. thick = NXZ * 1.3 @@ -44,7 +46,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: .build() ) shape = collider.shared_shape() - handle = state.insert_rigid_body(body, collider) + handle = state.insert_rigid_body(body, collider, no_coupling) viewer.insert_shape_with_color( handle, shape, Pose.from_translation(wall_pos), walls_color ) @@ -62,7 +64,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: body = RigidBodyBuilder.dynamic().translation(pos).build() collider = ColliderBuilder.ball(0.5).build() shape = collider.shared_shape() - handle = state.insert_rigid_body(body, collider) + handle = state.insert_rigid_body(body, collider, no_coupling) viewer.insert_shape(handle, shape, Pose.IDENTITY) # Optional: render even before starting the simulation. diff --git a/crates/nexus_python3d/examples/boxes3.py b/crates/nexus_python3d/examples/boxes3.py index 0d0ce716..df271a72 100644 --- a/crates/nexus_python3d/examples/boxes3.py +++ b/crates/nexus_python3d/examples/boxes3.py @@ -10,6 +10,7 @@ NexusViewer, NexusPipeline, NexusState, + RbdCoupling, RigidBodyBuilder, ColliderBuilder, GpuTimestamps, @@ -24,6 +25,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: state = NexusState() + no_coupling = RbdCoupling.NONE # Falling dynamic objects. for j in range(NY): @@ -38,7 +40,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: body = RigidBodyBuilder.dynamic().translation(pos).build() collider = ColliderBuilder.cuboid(0.5, 0.5, 0.5).build() shape = collider.shared_shape() - handle = state.insert_rigid_body(body, collider) + handle = state.insert_rigid_body(body, collider, no_coupling) viewer.insert_shape(handle, shape, Pose.IDENTITY) # Floor made of large cuboids. @@ -60,7 +62,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: .build() ) shape = collider.shared_shape() - handle = state.insert_rigid_body(body, collider) + handle = state.insert_rigid_body(body, collider, no_coupling) viewer.insert_shape_with_color( handle, shape, Pose.from_translation(wall_pos), walls_color ) diff --git a/crates/nexus_python3d/examples/boxes_and_balls3.py b/crates/nexus_python3d/examples/boxes_and_balls3.py index 7068ae6f..db6e7399 100644 --- a/crates/nexus_python3d/examples/boxes_and_balls3.py +++ b/crates/nexus_python3d/examples/boxes_and_balls3.py @@ -10,6 +10,7 @@ NexusViewer, NexusPipeline, NexusState, + RbdCoupling, RigidBodyBuilder, ColliderBuilder, GpuTimestamps, @@ -24,6 +25,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: state = NexusState() + no_coupling = RbdCoupling.NONE # Falling dynamic objects. for j in range(NY): @@ -41,7 +43,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: else: collider = ColliderBuilder.ball(0.5).build() shape = collider.shared_shape() - handle = state.insert_rigid_body(body, collider) + handle = state.insert_rigid_body(body, collider, no_coupling) viewer.insert_shape(handle, shape, Pose.IDENTITY) # Floor made of large cuboids. @@ -63,7 +65,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: .build() ) shape = collider.shared_shape() - handle = state.insert_rigid_body(body, collider) + handle = state.insert_rigid_body(body, collider, no_coupling) viewer.insert_shape_with_color( handle, shape, Pose.from_translation(wall_pos), walls_color ) diff --git a/crates/nexus_python3d/examples/centilever_beam3.py b/crates/nexus_python3d/examples/centilever_beam3.py new file mode 100644 index 00000000..95eb0391 --- /dev/null +++ b/crates/nexus_python3d/examples/centilever_beam3.py @@ -0,0 +1,93 @@ +"""Python port of `crates/examples3d/centilever_beam3.rs`. + +A Neo-Hookean elastic MPM beam clamped at one end by a fixed cuboid, sagging +under gravity (an MPM example despite the name). +""" + +import math + +from nexus3d import ( + NexusViewer, + NexusPipeline, + NexusState, + RbdCoupling, + BoundaryCondition, + RigidBodyBuilder, + ColliderBuilder, + GpuTimestamps, + SimulationParams, + ParticleModel, + Particle, + Vec3, + Pose, + vec3, +) + + +def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: + state = NexusState() + coupling = RbdCoupling.mpm_one_way(BoundaryCondition.stick()) + + width = 10.0 + height = 2.0 + fixed_part = 1.0 + cell_width = 0.2 + particle_per_cell_dim = 2 + young_modulus = 1.0e7 + poisson_ratio = 0.3 + + diameter = cell_width / particle_per_cell_dim + ni = int(math.ceil((width + fixed_part) / diameter)) + njk = int(math.ceil(height / diameter)) + + particles = [] + for i in range(ni): + for j in range(njk): + for k in range(njk): + position = vec3(float(i), float(j), float(k)) * diameter + density = 1000.0 + radius = diameter / 2.0 + model = ParticleModel.elastic_neo_hookean(young_modulus, poisson_ratio) + particle = Particle(position, radius, density, model) + particle.set_damping(2.0) + particles.append(particle) + + params = SimulationParams(vec3(0.0, -9.81, 0.0), 1.0 / 60.0) + state.set_mpm_params(viewer, params, cell_width) + state.set_mpm_substeps(20) + state.add_particles(viewer, particles) + + # Fixed block that clamps one end of the beam. + body = RigidBodyBuilder.fixed().translation( + vec3(0.0, height / 2.0, height / 2.0) + ).build() + collider = ColliderBuilder.cuboid(fixed_part, height, height).build() + shape = collider.shared_shape() + handle = state.insert_rigid_body(body, collider, coupling) + viewer.insert_shape(handle, shape, Pose.IDENTITY) + + timestamps = GpuTimestamps(viewer, 2048) + viewer.add_directional_light(Vec3(1.0, -2.0, 3.0)) + state.finalize(viewer) + + while viewer.render_frame(): + if viewer.simulating(): + pipeline.simulate(viewer, state, timestamps) + viewer.sync(state, timestamps) + + return state + + +def main() -> None: + viewer = NexusViewer() + viewer.init_backend() + pipeline = NexusPipeline() + pipeline.preload_pipelines(viewer) + run(viewer, pipeline) + + +if __name__ == "__main__": + main() + import os + + os._exit(0) diff --git a/crates/nexus_python3d/examples/compound3.py b/crates/nexus_python3d/examples/compound3.py index 6209e3fe..949ab05a 100644 --- a/crates/nexus_python3d/examples/compound3.py +++ b/crates/nexus_python3d/examples/compound3.py @@ -13,6 +13,7 @@ NexusViewer, NexusPipeline, NexusState, + RbdCoupling, RigidBodyBuilder, ColliderBuilder, GpuTimestamps, @@ -24,6 +25,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: state = NexusState() + no_coupling = RbdCoupling.NONE # Floor made of large cuboids. thick = 50.0 @@ -44,7 +46,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: .build() ) shape = collider.shared_shape() - handle = state.insert_rigid_body(body, collider) + handle = state.insert_rigid_body(body, collider, no_coupling) viewer.insert_shape_with_color( handle, shape, Pose.from_translation(wall_pos), walls_color ) @@ -75,7 +77,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: body = RigidBodyBuilder.dynamic().translation(Vec3(x, y, z)).build() - handle = state.insert_body_in(0, body) + handle = state.insert_body_in(0, body, no_coupling) for offset, he in parts: collider = ( ColliderBuilder.cuboid(he.x, he.y, he.z) diff --git a/crates/nexus_python3d/examples/dynamic_rbd3.py b/crates/nexus_python3d/examples/dynamic_rbd3.py index cbd8ebd0..87519066 100644 --- a/crates/nexus_python3d/examples/dynamic_rbd3.py +++ b/crates/nexus_python3d/examples/dynamic_rbd3.py @@ -11,6 +11,7 @@ NexusViewer, NexusPipeline, NexusState, + RbdCoupling, RigidBodyBuilder, ColliderBuilder, GpuTimestamps, @@ -34,6 +35,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: state = NexusState() + no_coupling = RbdCoupling.NONE # A boxed ground: a floor plus four low walls to keep the pile contained. floor_half = 50.0 @@ -52,7 +54,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: ColliderBuilder.cuboid(half.x, half.y, half.z).translation(pos).build() ) shape = collider.shared_shape() - handle = state.insert_rigid_body(body, collider) + handle = state.insert_rigid_body(body, collider, no_coupling) viewer.insert_shape_with_color( handle, shape, Pose.from_translation(pos), walls_color ) @@ -97,7 +99,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: colliders.append(collider) shapes.append(collider.shared_shape()) - handles = state.add_rigid_bodies(viewer, bodies, colliders) + handles = state.add_rigid_bodies(viewer, bodies, colliders, no_coupling) for handle, shape in zip(handles, shapes): viewer.insert_shape(handle, shape, Pose.IDENTITY) added += n diff --git a/crates/nexus_python3d/examples/elastic_cut3.py b/crates/nexus_python3d/examples/elastic_cut3.py new file mode 100644 index 00000000..546889ee --- /dev/null +++ b/crates/nexus_python3d/examples/elastic_cut3.py @@ -0,0 +1,131 @@ +"""Python port of `crates/examples3d/elastic_cut3.rs`. + +An elastic MPM block falls onto a floor through three tilted cutting planes +(each a flat heightfield trimesh rotated about X and offset). +""" + +import math + +from nexus3d import ( + NexusViewer, + NexusPipeline, + NexusState, + RbdCoupling, + BoundaryCondition, + RigidBodyBuilder, + ColliderBuilder, + GpuTimestamps, + SimulationParams, + ParticleModel, + Particle, + Vec3, + Pose, + vec3, +) + + +def heightfield_trimesh(nrows, ncols, height_fn, scale): + """Grid trimesh centered at origin. scale=(sx,sy,sz); y = height_fn(i,j)*sy.""" + sx, sy, sz = scale + verts = [] + for i in range(nrows): + for j in range(ncols): + x = (i / (nrows - 1) - 0.5) * sx + z = (j / (ncols - 1) - 0.5) * sz + verts.append([x, height_fn(i, j) * sy, z]) + idx = [] + for i in range(nrows - 1): + for j in range(ncols - 1): + a = i * ncols + j + idx.append([a, a + 1, a + ncols]) + idx.append([a + 1, a + ncols + 1, a + ncols]) + return verts, idx + + +def rotate_x(p, angle): + """Rotate point [x, y, z] about the X axis by `angle` radians.""" + c = math.cos(angle) + s = math.sin(angle) + x, y, z = p + return [x, y * c - z * s, y * s + z * c] + + +def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: + state = NexusState() + coupling = RbdCoupling.mpm_one_way(BoundaryCondition.separate(1.0)) + + nxz = 50 + cell_width = 1.0 + + particles = [] + for i in range(nxz): + for j in range(30): + for k in range(nxz): + position = vec3( + i + 0.5 - nxz / 2.0, + j + 0.5 + 60.0, + k + 0.5 - nxz / 2.0, + ) * (cell_width / 2.0) + density = 2700.0 + radius = cell_width / 4.0 + model = ParticleModel.elastic(1.0e7, 0.2) + particles.append(Particle(position, radius, density, model)) + + params = SimulationParams(vec3(0.0, -9.81, 0.0) * 4.0, 1.0 / 60.0) + state.set_mpm_params(viewer, params, cell_width) + state.set_mpm_substeps(20) + state.add_particles(viewer, particles) + + # Floor + body = RigidBodyBuilder.fixed().translation(vec3(0.0, -4.0, 0.0)).build() + collider = ColliderBuilder.cuboid(100.0, 1.0, 100.0).build() + shape = collider.shared_shape() + handle = state.insert_rigid_body(body, collider, coupling) + viewer.insert_shape(handle, shape, Pose.IDENTITY) + + # Cutting planes (3 heightfield trimeshes), each tilted about X by 1.3 rad + # and offset. The rotation + translation are baked into the vertices here. + for k in range(3): + vtx, idx = heightfield_trimesh( + 10, 10, lambda i, j: 0.0, (35.0, 1.0, 10.0) + ) + offset = [0.0, 10.0, k * 10.0 - 10.0] + vtx = [ + [ + r[0] + offset[0], + r[1] + offset[1], + r[2] + offset[2], + ] + for r in (rotate_x(pt, 1.3) for pt in vtx) + ] + body = RigidBodyBuilder.fixed().build() + collider = ColliderBuilder.trimesh(vtx, idx).build() + shape = collider.shared_shape() + handle = state.insert_rigid_body(body, collider, coupling) + viewer.insert_shape(handle, shape, Pose.IDENTITY) + + timestamps = GpuTimestamps(viewer, 2048) + viewer.add_directional_light(Vec3(1.0, -2.0, 3.0)) + state.finalize(viewer) + + while viewer.render_frame(): + if viewer.simulating(): + pipeline.simulate(viewer, state, timestamps) + viewer.sync(state, timestamps) + + return state + + +def main() -> None: + viewer = NexusViewer() + viewer.init_backend() + pipeline = NexusPipeline() + pipeline.preload_pipelines(viewer) + run(viewer, pipeline) + + +if __name__ == "__main__": + main() + import os + + os._exit(0) diff --git a/crates/nexus_python3d/examples/heightfield3.py b/crates/nexus_python3d/examples/heightfield3.py new file mode 100644 index 00000000..638e97d7 --- /dev/null +++ b/crates/nexus_python3d/examples/heightfield3.py @@ -0,0 +1,107 @@ +"""Python port of `crates/examples3d/heightfield3.rs`. + +MPM sand poured onto a sinusoidal heightfield terrain (rendered as a trimesh). +""" + +import math + +from nexus3d import ( + NexusViewer, + NexusPipeline, + NexusState, + RbdCoupling, + BoundaryCondition, + RigidBodyBuilder, + ColliderBuilder, + GpuTimestamps, + SimulationParams, + ParticleModel, + Particle, + Vec3, + Pose, + vec3, +) + + +def heightfield_trimesh(nrows, ncols, height_fn, scale): + """Grid trimesh centered at origin. scale=(sx,sy,sz); y = height_fn(i,j)*sy.""" + sx, sy, sz = scale + verts = [] + for i in range(nrows): + for j in range(ncols): + x = (i / (nrows - 1) - 0.5) * sx + z = (j / (ncols - 1) - 0.5) * sz + verts.append([x, height_fn(i, j) * sy, z]) + idx = [] + for i in range(nrows - 1): + for j in range(ncols - 1): + a = i * ncols + j + idx.append([a, a + 1, a + ncols]) + idx.append([a + 1, a + ncols + 1, a + ncols]) + return verts, idx + + +def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: + state = NexusState() + coupling = RbdCoupling.mpm_one_way(BoundaryCondition.separate(1.0)) + + nxz = 45 + cell_width = 1.0 + + particles = [] + for i in range(nxz): + for j in range(100): + for k in range(nxz): + position = vec3( + i + 0.5 - nxz / 2.0, + j + 0.5 + 14.0, + k + 0.5 - nxz / 2.0, + ) * (cell_width / 2.0) + density = 2700.0 + radius = cell_width / 4.0 + model = ParticleModel.sand(2.0e9, 0.2) + particles.append(Particle(position, radius, density, model)) + + params = SimulationParams(vec3(0.0, -9.81, 0.0), 1.0 / 60.0) + state.set_mpm_params(viewer, params, cell_width) + state.set_mpm_substeps(20) + state.add_particles(viewer, particles) + + # Sinusoidal heightfield terrain (rendered as the converted trimesh). + vtx, idx = heightfield_trimesh( + 200, + 200, + lambda i, j: math.sin(i / 10.0) * math.cos(j / 10.0), + (100.0, 5.0, 100.0), + ) + body = RigidBodyBuilder.fixed().build() + collider = ColliderBuilder.trimesh(vtx, idx).build() + shape = collider.shared_shape() + handle = state.insert_rigid_body(body, collider, coupling) + viewer.insert_shape(handle, shape, Pose.IDENTITY) + + timestamps = GpuTimestamps(viewer, 2048) + viewer.add_directional_light(Vec3(1.0, -2.0, 3.0)) + state.finalize(viewer) + + while viewer.render_frame(): + if viewer.simulating(): + pipeline.simulate(viewer, state, timestamps) + viewer.sync(state, timestamps) + + return state + + +def main() -> None: + viewer = NexusViewer() + viewer.init_backend() + pipeline = NexusPipeline() + pipeline.preload_pipelines(viewer) + run(viewer, pipeline) + + +if __name__ == "__main__": + main() + import os + + os._exit(0) diff --git a/crates/nexus_python3d/examples/joint_ball3.py b/crates/nexus_python3d/examples/joint_ball3.py index 13c0cf59..17bc22ba 100644 --- a/crates/nexus_python3d/examples/joint_ball3.py +++ b/crates/nexus_python3d/examples/joint_ball3.py @@ -8,6 +8,7 @@ NexusViewer, NexusPipeline, NexusState, + RbdCoupling, RigidBodyBuilder, ColliderBuilder, SphericalJointBuilder, @@ -19,6 +20,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: state = NexusState() + no_coupling = RbdCoupling.NONE rad = 0.4 ni = 200 @@ -54,7 +56,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: collider = ColliderBuilder.ball(rad).density(10.0).build() shape = collider.shared_shape() - child_handle = state.insert_rigid_body(rigid_body, collider) + child_handle = state.insert_rigid_body(rigid_body, collider, no_coupling) viewer.insert_shape(child_handle, shape, Pose.IDENTITY) # Vertical joint. @@ -98,7 +100,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: ) collider = ColliderBuilder.cuboid(rad, rad, rad).build() shape = collider.shared_shape() - handle = state.insert_rigid_body(body, collider) + handle = state.insert_rigid_body(body, collider, no_coupling) viewer.insert_shape(handle, shape, Pose.IDENTITY) timestamps = GpuTimestamps(viewer, 2048) diff --git a/crates/nexus_python3d/examples/joint_fixed3.py b/crates/nexus_python3d/examples/joint_fixed3.py index 308d7dfe..d375168d 100644 --- a/crates/nexus_python3d/examples/joint_fixed3.py +++ b/crates/nexus_python3d/examples/joint_fixed3.py @@ -7,6 +7,7 @@ NexusViewer, NexusPipeline, NexusState, + RbdCoupling, RigidBodyBuilder, ColliderBuilder, FixedJointBuilder, @@ -18,6 +19,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: state = NexusState() + no_coupling = RbdCoupling.NONE rad = 0.4 num = 10 @@ -61,7 +63,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: collider = ColliderBuilder.ball(rad).build() shape = collider.shared_shape() child_handle = state.insert_rigid_body( - rigid_body, collider + rigid_body, collider, no_coupling ) viewer.insert_shape(child_handle, shape, Pose.IDENTITY) diff --git a/crates/nexus_python3d/examples/joint_prismatic3.py b/crates/nexus_python3d/examples/joint_prismatic3.py index 8cec12bd..8a047e53 100644 --- a/crates/nexus_python3d/examples/joint_prismatic3.py +++ b/crates/nexus_python3d/examples/joint_prismatic3.py @@ -9,6 +9,7 @@ NexusViewer, NexusPipeline, NexusState, + RbdCoupling, RigidBodyBuilder, ColliderBuilder, PrismaticJointBuilder, @@ -20,6 +21,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: state = NexusState() + no_coupling = RbdCoupling.NONE rad = 0.4 num = 10 @@ -37,7 +39,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: ground = RigidBodyBuilder.fixed().translation(Vec3(x, y, z)).build() collider = ColliderBuilder.cuboid(rad, rad, rad).build() shape = collider.shared_shape() - curr_parent = state.insert_rigid_body(ground, collider) + curr_parent = state.insert_rigid_body(ground, collider, no_coupling) viewer.insert_shape(curr_parent, shape, Pose.IDENTITY) for i in range(num): @@ -53,7 +55,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: ) shape = collider.shared_shape() curr_child = state.insert_rigid_body( - rigid_body, collider + rigid_body, collider, no_coupling ) viewer.insert_shape(curr_child, shape, Pose.IDENTITY) diff --git a/crates/nexus_python3d/examples/joint_revolute3.py b/crates/nexus_python3d/examples/joint_revolute3.py index ba624be4..646ba5c7 100644 --- a/crates/nexus_python3d/examples/joint_revolute3.py +++ b/crates/nexus_python3d/examples/joint_revolute3.py @@ -7,6 +7,7 @@ NexusViewer, NexusPipeline, NexusState, + RbdCoupling, RigidBodyBuilder, ColliderBuilder, RevoluteJointBuilder, @@ -18,6 +19,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: state = NexusState() + no_coupling = RbdCoupling.NONE rad = 0.4 num = 10 @@ -36,7 +38,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: ground = RigidBodyBuilder.fixed().translation(Vec3(x, y, z0)).build() collider = ColliderBuilder.cuboid(rad, rad, rad).build() shape = collider.shared_shape() - curr_parent = state.insert_rigid_body(ground, collider) + curr_parent = state.insert_rigid_body(ground, collider, no_coupling) viewer.insert_shape(curr_parent, shape, Pose.IDENTITY) for i in range(num): @@ -61,7 +63,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: ) shape = collider.shared_shape() handles[m] = state.insert_rigid_body( - rigid_body, collider + rigid_body, collider, no_coupling ) viewer.insert_shape(handles[m], shape, Pose.IDENTITY) diff --git a/crates/nexus_python3d/examples/joint_revolute_batch3.py b/crates/nexus_python3d/examples/joint_revolute_batch3.py index 6314f982..c8be5726 100644 --- a/crates/nexus_python3d/examples/joint_revolute_batch3.py +++ b/crates/nexus_python3d/examples/joint_revolute_batch3.py @@ -7,6 +7,7 @@ NexusViewer, NexusPipeline, NexusState, + RbdCoupling, RigidBodyBuilder, ColliderBuilder, RevoluteJointBuilder, @@ -18,6 +19,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: state = NexusState() + no_coupling = RbdCoupling.NONE rad = 0.4 num = 10 @@ -42,7 +44,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: ground = RigidBodyBuilder.fixed().translation(Vec3(x, y, z0)).build() ground_collider = ColliderBuilder.cuboid(rad, rad, rad).build() ground_shape = ground_collider.shared_shape() - curr_parent = state.insert_rigid_body_in(env, ground, ground_collider) + curr_parent = state.insert_rigid_body_in(env, ground, ground_collider, no_coupling) viewer.insert_shape_in(env, curr_parent, ground_shape, Pose.IDENTITY) for i in range(num): @@ -59,7 +61,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: body = RigidBodyBuilder.dynamic().pose(positions[m]).build() collider = ColliderBuilder.cuboid(rad, rad, rad).density(1.0).build() shape = collider.shared_shape() - handles[m] = state.insert_rigid_body_in(env, body, collider) + handles[m] = state.insert_rigid_body_in(env, body, collider, no_coupling) viewer.insert_shape_in(env, handles[m], shape, Pose.IDENTITY) ax = Vec3.X diff --git a/crates/nexus_python3d/examples/joints3.py b/crates/nexus_python3d/examples/joints3.py index 5c533dda..1e93f02d 100644 --- a/crates/nexus_python3d/examples/joints3.py +++ b/crates/nexus_python3d/examples/joints3.py @@ -15,6 +15,7 @@ NexusViewer, NexusPipeline, NexusState, + RbdCoupling, RigidBodyBuilder, ColliderBuilder, FixedJointBuilder, @@ -44,7 +45,7 @@ def normalized(v: Vec3) -> Vec3: def add_body(state, viewer, body, collider): """Inserts a body + collider and registers its render shape.""" shape = collider.shared_shape() - handle = state.insert_rigid_body(body, collider) + handle = state.insert_rigid_body(body, collider, RbdCoupling.NONE) viewer.insert_shape(handle, shape, Pose.IDENTITY) return handle diff --git a/crates/nexus_python3d/examples/keva3.py b/crates/nexus_python3d/examples/keva3.py index a22053ce..16ffa57f 100644 --- a/crates/nexus_python3d/examples/keva3.py +++ b/crates/nexus_python3d/examples/keva3.py @@ -10,6 +10,7 @@ NexusViewer, NexusPipeline, NexusState, + RbdCoupling, RigidBodyBuilder, ColliderBuilder, GpuTimestamps, @@ -21,7 +22,7 @@ def add_body(state, viewer, body, collider): """Inserts a body + collider into the state and registers its render shape.""" shape = collider.shared_shape() - handle = state.insert_rigid_body(body, collider) + handle = state.insert_rigid_body(body, collider, RbdCoupling.NONE) viewer.insert_shape(handle, shape, Pose.IDENTITY) return handle diff --git a/crates/nexus_python3d/examples/many_pyramids3.py b/crates/nexus_python3d/examples/many_pyramids3.py index 0f24fcaf..5609b5a7 100644 --- a/crates/nexus_python3d/examples/many_pyramids3.py +++ b/crates/nexus_python3d/examples/many_pyramids3.py @@ -10,6 +10,7 @@ NexusViewer, NexusPipeline, NexusState, + RbdCoupling, RigidBodyBuilder, ColliderBuilder, GpuTimestamps, @@ -21,7 +22,7 @@ def add_body(state, viewer, body, collider): """Inserts a body + collider into the state and registers its render shape.""" shape = collider.shared_shape() - handle = state.insert_rigid_body(body, collider) + handle = state.insert_rigid_body(body, collider, RbdCoupling.NONE) viewer.insert_shape(handle, shape, Pose.IDENTITY) return handle diff --git a/crates/nexus_python3d/examples/many_pyramids_batch3.py b/crates/nexus_python3d/examples/many_pyramids_batch3.py index 152e44ff..975ddc07 100644 --- a/crates/nexus_python3d/examples/many_pyramids_batch3.py +++ b/crates/nexus_python3d/examples/many_pyramids_batch3.py @@ -7,6 +7,7 @@ NexusViewer, NexusPipeline, NexusState, + RbdCoupling, RigidBodyBuilder, ColliderBuilder, GpuTimestamps, @@ -26,12 +27,13 @@ def create_pyramid(state, viewer, env, offset, stack_height, rad): body = RigidBodyBuilder.dynamic().translation(Vec3(x, y, 0.0) + offset).build() collider = ColliderBuilder.cuboid(rad, rad, rad).build() shape = collider.shared_shape() - handle = state.insert_rigid_body_in(env, body, collider) + handle = state.insert_rigid_body_in(env, body, collider, RbdCoupling.NONE) viewer.insert_shape_in(env, handle, shape, Pose.IDENTITY) def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: state = NexusState() + no_coupling = RbdCoupling.NONE spacing = 4.0 for pyramid_index in range(PYRAMID_COUNT): @@ -48,7 +50,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: ground_size, ground_height, PYRAMID_COUNT * spacing / 2.0 + ground_size ).build() shape = collider.shared_shape() - ground_handle = state.insert_rigid_body_in(env, body, collider) + ground_handle = state.insert_rigid_body_in(env, body, collider, no_coupling) viewer.insert_shape_in(env, ground_handle, shape, Pose.IDENTITY) # Cubes. diff --git a/crates/nexus_python3d/examples/mpm_emitter3.py b/crates/nexus_python3d/examples/mpm_emitter3.py new file mode 100644 index 00000000..212a60df --- /dev/null +++ b/crates/nexus_python3d/examples/mpm_emitter3.py @@ -0,0 +1,132 @@ +"""Python port of `crates/examples3d/mpm_emitter3.rs`. + +A dynamic emitter spawns a stream of sand that orbits the center of a walled box. +""" + +import math +from collections import deque + +from nexus3d import ( + NexusViewer, + NexusPipeline, + NexusState, + RbdCoupling, + BoundaryCondition, + RigidBodyBuilder, + ColliderBuilder, + GpuTimestamps, + SimulationParams, + ParticleModel, + Particle, + Vec3, + Vec4, + Pose, + vec3, +) + +DENSITY = 2700.0 +YOUNG_MODULUS = 2.0e8 +POISSON_RATIO = 0.2 + +EMIT_EVERY = 10 +EMIT_BLOCK = 30 +EMIT_BLOCK_Y = 2 +MAX_PARTICLES = 250_000 + + +def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: + state = NexusState() + coupling = RbdCoupling.mpm_one_way(BoundaryCondition.separate(1.0)) + + cell_width = 1.0 + dt = 1.0 / 60.0 + + params = SimulationParams(vec3(0.0, -9.81, 0.0), dt) + state.set_mpm_params(viewer, params, cell_width) + state.set_mpm_substeps(10) + + # Boundary colliders: a floor and four walls forming an open box. + thickness = 0.5 + walls_color = Vec4(0.6, 0.8, 1.0, 0.3) + walls = [ + (vec3(0.0, -thickness, 0.0), vec3(30.0, thickness, 30.0)), + (vec3(0.0, 10.0, -30.0), vec3(30.0, 10.0, thickness)), + (vec3(0.0, 10.0, 30.0), vec3(30.0, 10.0, thickness)), + (vec3(-30.0, 10.0, 0.0), vec3(thickness, 10.0, 30.0)), + (vec3(30.0, 10.0, 0.0), vec3(thickness, 10.0, 30.0)), + ] + for pos, half_extents in walls: + body = RigidBodyBuilder.fixed().translation(pos).build() + collider = ColliderBuilder.cuboid( + half_extents.x, half_extents.y, half_extents.z + ).build() + shape = collider.shared_shape() + handle = state.insert_rigid_body(body, collider, coupling) + viewer.insert_shape_with_color(handle, shape, Pose.IDENTITY, walls_color) + + timestamps = GpuTimestamps(viewer, 2048) + viewer.add_directional_light(Vec3(1.0, -2.0, 3.0)) + state.finalize(viewer) + + # Dynamic emitter: a small cube of sand spawned at an orbiting point. + radius = cell_width / 4.0 + spacing = radius * 2.0 + model = ParticleModel.sand(YOUNG_MODULUS, POISSON_RATIO) + emit_height = 40.0 + orbit_radius = 10.0 + angular_speed = 1.5 # rad/s + + chunks: deque = deque() + total_particles = 0 + t = 0.0 + step = 0 + + while viewer.render_frame(): + if viewer.simulating(): + if step % EMIT_EVERY == 0 and total_particles < MAX_PARTICLES: + angle = t * angular_speed + center = vec3( + orbit_radius * math.cos(angle), + emit_height, + orbit_radius * math.sin(angle), + ) + + particles = [] + for i in range(EMIT_BLOCK): + for j in range(EMIT_BLOCK_Y): + for k in range(EMIT_BLOCK): + offset = vec3( + (i - EMIT_BLOCK // 2) * spacing, + (j - EMIT_BLOCK_Y // 2) * spacing, + (k - EMIT_BLOCK // 2) * spacing, + ) + particle = Particle(center + offset, radius, DENSITY, model) + particle.velocity = vec3(0.0, -8.0, 0.0) + particles.append(particle) + + n = len(particles) + chunk = state.add_particles(viewer, particles) + chunks.append((chunk, n)) + total_particles += n + + pipeline.simulate(viewer, state, timestamps) + t += dt + step += 1 + viewer.sync(state, timestamps) + + return state + + +def main() -> None: + viewer = NexusViewer() + viewer.init_backend() + pipeline = NexusPipeline() + pipeline.preload_pipelines(viewer) + run(viewer, pipeline) + + +if __name__ == "__main__": + main() + import os + + os._exit(0) diff --git a/crates/nexus_python3d/examples/multibody_pendulum3.py b/crates/nexus_python3d/examples/multibody_pendulum3.py index 42dbe528..1ccaf2f3 100644 --- a/crates/nexus_python3d/examples/multibody_pendulum3.py +++ b/crates/nexus_python3d/examples/multibody_pendulum3.py @@ -9,6 +9,7 @@ NexusViewer, NexusPipeline, NexusState, + RbdCoupling, RigidBodyBuilder, ColliderBuilder, InteractionGroups, @@ -21,6 +22,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: state = NexusState() + no_coupling = RbdCoupling.NONE rad = 0.4 link_len = 2.0 @@ -30,7 +32,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: root_body = RigidBodyBuilder.fixed().build() root_collider = ColliderBuilder.cuboid(rad, rad, rad).build() root_shape = root_collider.shared_shape() - parent_handle = state.insert_rigid_body(root_body, root_collider) + parent_handle = state.insert_rigid_body(root_body, root_collider, no_coupling) viewer.insert_shape(parent_handle, root_shape, Pose.IDENTITY) for i in range(num_links): @@ -46,7 +48,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: .build() ) shape = collider.shared_shape() - handle = state.insert_rigid_body(rigid_body, collider) + handle = state.insert_rigid_body(rigid_body, collider, no_coupling) viewer.insert_shape(handle, shape, Pose.IDENTITY) # Revolute joint about Z: anchor on parent is at its bottom diff --git a/crates/nexus_python3d/examples/primitives3.py b/crates/nexus_python3d/examples/primitives3.py index 27a9a9e7..f015011a 100644 --- a/crates/nexus_python3d/examples/primitives3.py +++ b/crates/nexus_python3d/examples/primitives3.py @@ -14,6 +14,7 @@ NexusViewer, NexusPipeline, NexusState, + RbdCoupling, RigidBodyBuilder, ColliderBuilder, GpuTimestamps, @@ -52,6 +53,7 @@ def make_polyhedron_points(): def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: state = NexusState() + no_coupling = RbdCoupling.NONE # Create 5 predefined convex polyhedron point sets (so we can render them # efficiently with instancing). @@ -87,7 +89,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: collider = cb.build() body = RigidBodyBuilder.dynamic().translation(pos).build() shape = collider.shared_shape() - handle = state.insert_rigid_body(body, collider) + handle = state.insert_rigid_body(body, collider, no_coupling) viewer.insert_shape(handle, shape, Pose.IDENTITY) # Floor made of large cuboids. @@ -109,7 +111,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: .build() ) shape = collider.shared_shape() - handle = state.insert_rigid_body(body, collider) + handle = state.insert_rigid_body(body, collider, no_coupling) viewer.insert_shape_with_color( handle, shape, Pose.from_translation(wall_pos), walls_color ) diff --git a/crates/nexus_python3d/examples/pyramid3.py b/crates/nexus_python3d/examples/pyramid3.py index 48e60ad7..58d575b4 100644 --- a/crates/nexus_python3d/examples/pyramid3.py +++ b/crates/nexus_python3d/examples/pyramid3.py @@ -10,6 +10,7 @@ NexusViewer, NexusPipeline, NexusState, + RbdCoupling, RigidBodyBuilder, ColliderBuilder, GpuTimestamps, @@ -21,7 +22,7 @@ def add_body(state, viewer, body, collider): """Inserts a body + collider into the state and registers its render shape.""" shape = collider.shared_shape() - handle = state.insert_rigid_body(body, collider) + handle = state.insert_rigid_body(body, collider, RbdCoupling.NONE) viewer.insert_shape(handle, shape, Pose.IDENTITY) return handle diff --git a/crates/nexus_python3d/examples/sand3.py b/crates/nexus_python3d/examples/sand3.py new file mode 100644 index 00000000..2969f661 --- /dev/null +++ b/crates/nexus_python3d/examples/sand3.py @@ -0,0 +1,113 @@ +"""Python port of `crates/examples3d/sand3.rs`. + +MPM sand poured into a walled box, stirred by a rotating kinematic blade. +""" + +from nexus3d import ( + NexusViewer, + NexusPipeline, + NexusState, + RbdCoupling, + BoundaryCondition, + RigidBodyBuilder, + ColliderBuilder, + GpuTimestamps, + SimulationParams, + ParticleModel, + Particle, + Vec3, + Vec4, + Pose, + vec3, +) + +DENSITY = 2700.0 +YOUNG_MODULUS = 2.0e9 +POISSON_RATIO = 0.2 + + +def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: + state = NexusState() + # MPM boundary colliders are inserted as rigid bodies coupled to the + # continuum; they push the particles but aren't pushed back. + coupling = RbdCoupling.mpm_one_way(BoundaryCondition.separate(1.0)) + + nxz = 45 + cell_width = 1.0 + + # Sand particles. + particles = [] + for i in range(nxz): + for j in range(100): + for k in range(nxz): + position = vec3( + i + 0.5 - nxz / 2.0, + j + 0.5 + 10.0, + k + 0.5 - nxz / 2.0, + ) * (cell_width / 2.0) + radius = cell_width / 4.0 + model = ParticleModel.sand(YOUNG_MODULUS, POISSON_RATIO) + particles.append(Particle(position, radius, DENSITY, model)) + + params = SimulationParams(vec3(0.0, -9.81, 0.0), 1.0 / 60.0) + state.set_mpm_params(viewer, params, cell_width) + state.set_mpm_substeps(20) + state.add_particles(viewer, particles) + + # Boundary colliders (floor, walls). + thickness = 0.5 + walls_color = Vec4(0.6, 0.8, 1.0, 0.3) + walls = [ + (vec3(0.0, -4.0, 0.0), vec3(100.0, 4.0, 100.0)), + (vec3(0.0, 5.0, -35.0), vec3(35.0, 5.0, thickness)), + (vec3(0.0, 5.0, 35.0), vec3(35.0, 5.0, thickness)), + (vec3(-35.0, 5.0, 0.0), vec3(thickness, 5.0, 35.0)), + (vec3(35.0, 5.0, 0.0), vec3(thickness, 5.0, 35.0)), + ] + for pos, half_extents in walls: + body = RigidBodyBuilder.fixed().translation(pos).build() + collider = ColliderBuilder.cuboid( + half_extents.x, half_extents.y, half_extents.z + ).build() + shape = collider.shared_shape() + handle = state.insert_rigid_body(body, collider, coupling) + viewer.insert_shape_with_color(handle, shape, Pose.IDENTITY, walls_color) + + # Rotating blade (kinematic). + body = ( + RigidBodyBuilder.kinematic_velocity_based() + .translation(vec3(0.0, 2.0, 0.0)) + .rotation(vec3(0.0, 0.0, -0.5)) + .angvel(vec3(0.0, -1.0, 0.0)) + .build() + ) + collider = ColliderBuilder.cuboid(thickness, 2.0, 30.0).build() + shape = collider.shared_shape() + handle = state.insert_rigid_body(body, collider, coupling) + viewer.insert_shape(handle, shape, Pose.IDENTITY) + + timestamps = GpuTimestamps(viewer, 2048) + viewer.add_directional_light(Vec3(1.0, -2.0, 3.0)) + state.finalize(viewer) + + while viewer.render_frame(): + if viewer.simulating(): + pipeline.simulate(viewer, state, timestamps) + viewer.sync(state, timestamps) + + return state + + +def main() -> None: + viewer = NexusViewer() + viewer.init_backend() + pipeline = NexusPipeline() + pipeline.preload_pipelines(viewer) + run(viewer, pipeline) + + +if __name__ == "__main__": + main() + import os + + os._exit(0) diff --git a/crates/nexus_python3d/examples/trimesh3.py b/crates/nexus_python3d/examples/trimesh3.py index 307c14f5..4cfd6e2b 100644 --- a/crates/nexus_python3d/examples/trimesh3.py +++ b/crates/nexus_python3d/examples/trimesh3.py @@ -15,6 +15,7 @@ NexusViewer, NexusPipeline, NexusState, + RbdCoupling, RigidBodyBuilder, ColliderBuilder, GpuTimestamps, @@ -70,6 +71,7 @@ def make_polyhedron_points(): def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: state = NexusState() + no_coupling = RbdCoupling.NONE # Create 5 predefined convex polyhedron point sets. polyhedron_points = make_polyhedron_points() @@ -103,7 +105,7 @@ def run(viewer: NexusViewer, pipeline: NexusPipeline) -> NexusState: collider = cb.build() body = RigidBodyBuilder.dynamic().translation(pos).build() shape = collider.shared_shape() - handle = state.insert_rigid_body(body, collider) + handle = state.insert_rigid_body(body, collider, no_coupling) viewer.insert_shape(handle, shape, Pose.IDENTITY) # A trimesh floor built from the mesh representation of a heightfield. @@ -124,7 +126,7 @@ def height_fn(i, j): body = RigidBodyBuilder.fixed().build() collider = ColliderBuilder.trimesh(vertices, indices).build() shape = collider.shared_shape() - handle = state.insert_rigid_body(body, collider) + handle = state.insert_rigid_body(body, collider, no_coupling) viewer.insert_shape(handle, shape, Pose.IDENTITY) timestamps = GpuTimestamps(viewer, 2048) diff --git a/crates/nexus_python3d/run_examples.py b/crates/nexus_python3d/run_examples.py index 7454f868..cf451366 100644 --- a/crates/nexus_python3d/run_examples.py +++ b/crates/nexus_python3d/run_examples.py @@ -54,6 +54,12 @@ "multibody_pendulum3", "joint_revolute_batch3", "many_pyramids_batch3", + # MPM + "mpm_emitter3", + "sand3", + "heightfield3", + "elastic_cut3", + "centilever_beam3", # Robots (need external assets via env vars; see README) "urdf3", "mujoco_menagerie3", diff --git a/crates/nexus_python3d/src/lib.rs b/crates/nexus_python3d/src/lib.rs index e34f1278..e4fe4623 100644 --- a/crates/nexus_python3d/src/lib.rs +++ b/crates/nexus_python3d/src/lib.rs @@ -10,6 +10,7 @@ use pyo3::prelude::*; pub mod loaders; pub mod math; +pub mod mpm; pub mod nexus; pub mod rbd; pub mod viewer; @@ -46,7 +47,9 @@ fn nexus3d(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; // Core simulation + m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; @@ -59,5 +62,11 @@ fn nexus3d(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; + // MPM + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + Ok(()) } diff --git a/crates/nexus_python3d/src/loaders.rs b/crates/nexus_python3d/src/loaders.rs index 56223670..58447299 100644 --- a/crates/nexus_python3d/src/loaders.rs +++ b/crates/nexus_python3d/src/loaders.rs @@ -121,6 +121,7 @@ pub fn insert_mjcf( scene_path: &std::path::Path, render_colliders: bool, ) -> PyResult { + use nexus3d::prelude::RbdCoupling; use pyo3::exceptions::PyRuntimeError; use rapier3d::parry::bounding_volume::BoundingVolume; // for `Aabb::merge` use rapier3d_mjcf::{MjcfLoaderOptions, MjcfMultibodyOptions, MjcfRobot}; @@ -241,7 +242,7 @@ pub fn insert_mjcf( let body = rp::RigidBodyBuilder::fixed().translation(center).build(); let collider = rp::ColliderBuilder::cuboid(he.x, he.y, he.z).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider); + let handle = state.insert_rigid_body(body, collider, RbdCoupling::None); v.insert_shape(handle, &shape, rp::Pose::IDENTITY); } diff --git a/crates/nexus_python3d/src/mpm.rs b/crates/nexus_python3d/src/mpm.rs new file mode 100644 index 00000000..728f8b2b --- /dev/null +++ b/crates/nexus_python3d/src/mpm.rs @@ -0,0 +1,137 @@ +//! Material-Point-Method types (`nexus3d::mpm::solver`). + +use crate::math::Vec3; +use nexus3d::mpm::solver::{ + BoundaryCondition as RBoundaryCondition, Particle as RParticle, + ParticleModel as RParticleModel, SimulationParams as RSimulationParams, +}; +use pyo3::prelude::*; + +/// MPM global simulation parameters. +#[pyclass(name = "SimulationParams", from_py_object)] +#[derive(Clone, Copy)] +pub struct SimulationParams(pub RSimulationParams); + +#[pymethods] +impl SimulationParams { + #[new] + fn new(gravity: Vec3, dt: f32) -> Self { + SimulationParams(RSimulationParams { + gravity: gravity.0, + dt, + }) + } + + #[getter] + fn dt(&self) -> f32 { + self.0.dt + } + #[getter] + fn gravity(&self) -> Vec3 { + Vec3(self.0.gravity) + } +} + +/// A boundary condition applied to grid nodes at a coupled collider surface. +#[pyclass(name = "BoundaryCondition", from_py_object)] +#[derive(Clone, Copy)] +pub struct BoundaryCondition(pub RBoundaryCondition); + +#[pymethods] +impl BoundaryCondition { + /// No-slip: the material sticks to the surface (zero relative velocity). + #[staticmethod] + fn stick() -> Self { + BoundaryCondition(RBoundaryCondition::stick()) + } + /// Free-slip: the normal velocity component is removed, tangential kept. + #[staticmethod] + fn slip() -> Self { + BoundaryCondition(RBoundaryCondition::slip()) + } + /// Separating contact with Coulomb `friction` (`>= 0`): the material can + /// pull away from the surface but is resisted tangentially. + #[staticmethod] + fn separate(friction: f32) -> Self { + BoundaryCondition(RBoundaryCondition::separate(friction)) + } + + /// The Coulomb friction coefficient (only meaningful for `separate`). + #[getter] + fn friction(&self) -> f32 { + self.0.friction + } +} + +/// A particle constitutive model. +#[pyclass(name = "ParticleModel", from_py_object)] +#[derive(Clone, Copy)] +pub struct ParticleModel(pub RParticleModel); + +#[pymethods] +impl ParticleModel { + #[staticmethod] + fn elastic(young_modulus: f32, poisson_ratio: f32) -> Self { + ParticleModel(RParticleModel::elastic(young_modulus, poisson_ratio)) + } + #[staticmethod] + fn elastic_neo_hookean(young_modulus: f32, poisson_ratio: f32) -> Self { + ParticleModel(RParticleModel::elastic_neo_hookean( + young_modulus, + poisson_ratio, + )) + } + #[staticmethod] + fn sand(young_modulus: f32, poisson_ratio: f32) -> Self { + ParticleModel(RParticleModel::sand(young_modulus, poisson_ratio)) + } + #[staticmethod] + fn sand_neo_hookean(young_modulus: f32, poisson_ratio: f32) -> Self { + ParticleModel(RParticleModel::sand_neo_hookean( + young_modulus, + poisson_ratio, + )) + } +} + +/// A single MPM particle. +#[pyclass(name = "Particle", from_py_object)] +#[derive(Clone, Copy)] +pub struct Particle(pub RParticle); + +#[pymethods] +impl Particle { + #[new] + fn new(position: Vec3, radius: f32, density: f32, model: ParticleModel) -> Self { + Particle(RParticle::new(position.0, radius, density, model.0)) + } + + /// The particle velocity (`particle.dynamics.velocity`). + #[getter] + fn velocity(&self) -> Vec3 { + Vec3(self.0.dynamics.velocity) + } + #[setter] + fn set_velocity(&mut self, v: Vec3) { + self.0.dynamics.velocity = v.0; + } + + #[getter] + fn position(&self) -> Vec3 { + Vec3(self.0.position) + } + #[setter] + fn set_position(&mut self, p: Vec3) { + self.0.position = p.0; + } + + fn set_fixed(&mut self, fixed: bool) { + self.0.dynamics.set_fixed(fixed); + } + fn set_damping(&mut self, damping: f32) { + self.0.dynamics.set_damping(damping); + } + fn set_density(&mut self, density: f32) { + self.0.dynamics.set_density(density); + } +} diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index 0cee52f4..ef9648d3 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -1,16 +1,19 @@ -//! Core simulation objects: `NexusState`, `NexusPipeline`, `GpuTimestamps`, -//! and the various entity handles. +//! Core simulation objects: `NexusState`, `NexusPipeline`, `RbdCoupling`, +//! `GpuTimestamps`, and the various entity handles. use crate::loaders::{MjcfSceneInfo, UrdfLoaderOptions, UrdfRobotHandles}; use crate::math::{Pose, Vec3}; +use crate::mpm::{BoundaryCondition, Particle, SimulationParams}; use crate::rbd::{ Collider, ImpulseJointHandle, JointArg, JointAxis, MultibodyJointHandle, RigidBody, RigidBodyHandle, SharedShape, }; use crate::viewer::NexusViewer; use khal::backend::GpuTimestamps as RGpuTimestamps; +use nexus3d::mpm::solver::BoundaryCondition as RBoundaryCondition; use nexus3d::prelude::{ NexusPipeline as RNexusPipeline, NexusPipelineMask, NexusState as RNexusState, + RbdCoupling as RRbdCoupling, }; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; @@ -21,6 +24,39 @@ fn gpu_err(e: E) -> PyErr { PyRuntimeError::new_err(format!("{e:?}")) } +/// Coupling mode between a rigid body and the MPM simulation. +#[pyclass(name = "RbdCoupling", from_py_object)] +#[derive(Clone, Copy)] +pub struct RbdCoupling(pub RRbdCoupling); + +#[pymethods] +impl RbdCoupling { + #[classattr] + const NONE: RbdCoupling = RbdCoupling(RRbdCoupling::None); + // Convenience constants defaulting to a `stick()` boundary; use + // `mpm_one_way` / `mpm_two_way` to pick a specific boundary condition. + #[classattr] + const MPM_ONE_WAY_COUPLING: RbdCoupling = + RbdCoupling(RRbdCoupling::MpmOneWay(RBoundaryCondition::stick())); + #[classattr] + const MPM_TWO_WAY_COUPLING: RbdCoupling = + RbdCoupling(RRbdCoupling::MpmTwoWay(RBoundaryCondition::stick())); + + /// One-way coupling (MPM pushes the rigid body, not vice-versa) using the + /// given boundary condition at the collider surface. + #[staticmethod] + fn mpm_one_way(boundary: BoundaryCondition) -> RbdCoupling { + RbdCoupling(RRbdCoupling::MpmOneWay(boundary.0)) + } + + /// Two-way coupling (MPM and the rigid body affect each other) using the + /// given boundary condition at the collider surface. + #[staticmethod] + fn mpm_two_way(boundary: BoundaryCondition) -> RbdCoupling { + RbdCoupling(RRbdCoupling::MpmTwoWay(boundary.0)) + } +} + /// Entity counts for a `NexusState` (mirrors `NexusCounts`). #[pyclass(name = "NexusCounts", from_py_object)] #[derive(Clone, Copy)] @@ -37,8 +73,15 @@ pub struct NexusCounts { pub multibodies: usize, #[pyo3(get)] pub multibody_dofs: usize, + #[pyo3(get)] + pub particles: usize, } +/// Handle to a chunk of MPM particles added via `NexusState.add_particles`. +#[pyclass(name = "NexusParticleChunk", from_py_object)] +#[derive(Clone, Copy)] +pub struct NexusParticleChunk(pub nexus3d::prelude::NexusParticleChunk); + /// Optional GPU timing-query buffer (`khal::backend::GpuTimestamps`). #[pyclass(name = "GpuTimestamps", unsendable)] pub struct GpuTimestamps(pub RGpuTimestamps); @@ -69,8 +112,12 @@ impl NexusState { &mut self, body: PyRef, collider: PyRef, + coupling: RbdCoupling, ) -> RigidBodyHandle { - RigidBodyHandle(self.0.insert_rigid_body(body.0.clone(), collider.0.clone())) + RigidBodyHandle( + self.0 + .insert_rigid_body(body.0.clone(), collider.0.clone(), coupling.0), + ) } fn insert_rigid_body_in( @@ -78,21 +125,29 @@ impl NexusState { env: usize, body: PyRef, collider: PyRef, + coupling: RbdCoupling, ) -> RigidBodyHandle { - RigidBodyHandle( - self.0 - .insert_rigid_body_in(env, body.0.clone(), collider.0.clone()), - ) + RigidBodyHandle(self.0.insert_rigid_body_in( + env, + body.0.clone(), + collider.0.clone(), + coupling.0, + )) } - fn insert_body(&mut self, body: PyRef) -> RigidBodyHandle { - RigidBodyHandle(self.0.insert_body(body.0.clone())) + fn insert_body(&mut self, body: PyRef, coupling: RbdCoupling) -> RigidBodyHandle { + RigidBodyHandle(self.0.insert_body(body.0.clone(), coupling.0)) } /// Inserts a collider-less body into environment `env`; attach colliders to /// it afterwards with `insert_collider_in` (multiple colliders per body). - fn insert_body_in(&mut self, env: usize, body: PyRef) -> RigidBodyHandle { - RigidBodyHandle(self.0.insert_body_in(env, body.0.clone())) + fn insert_body_in( + &mut self, + env: usize, + body: PyRef, + coupling: RbdCoupling, + ) -> RigidBodyHandle { + RigidBodyHandle(self.0.insert_body_in(env, body.0.clone(), coupling.0)) } /// Attaches a collider to an existing body (`parent`), or inserts a @@ -127,15 +182,19 @@ impl NexusState { viewer: PyRef, bodies: Vec, colliders: Vec, + coupling: RbdCoupling, ) -> PyResult> { if bodies.len() != colliders.len() { return Err(PyRuntimeError::new_err( "bodies and colliders must have the same length", )); } - let pairs = bodies.into_iter().zip(colliders).map(|(b, c)| (b.0, c.0)); + let triples = bodies + .into_iter() + .zip(colliders) + .map(|(b, c)| (b.0, c.0, coupling.0)); self.0 - .add_rigid_bodies(viewer.backend(), pairs) + .add_rigid_bodies(viewer.backend(), triples) .map(|hs| hs.into_iter().map(RigidBodyHandle).collect()) .map_err(gpu_err) } @@ -306,6 +365,65 @@ impl NexusState { .map_err(gpu_err) } + // --- mpm -------------------------------------------------------------- + + fn set_mpm_params( + &mut self, + viewer: PyRef, + params: PyRef, + cell_width: f32, + ) -> PyResult<()> { + self.0 + .set_mpm_params(viewer.backend(), params.0, cell_width) + .map_err(gpu_err) + } + + fn set_mpm_substeps(&mut self, substeps: u32) { + self.0.set_mpm_substeps(substeps); + } + + fn set_mpm_use_cpic(&mut self, enabled: bool) { + self.0.set_mpm_use_cpic(enabled); + } + + fn set_mpm_gravity(&mut self, gravity: Vec3) { + self.0.set_mpm_gravity(gravity.0); + } + + fn add_particles( + &mut self, + viewer: PyRef, + particles: Vec, + ) -> PyResult { + let particles: Vec<_> = particles.into_iter().map(|p| p.0).collect(); + self.0 + .add_particles(viewer.backend(), particles) + .map(NexusParticleChunk) + .map_err(gpu_err) + } + + fn extend_chunk( + &mut self, + viewer: PyRef, + chunk: NexusParticleChunk, + particles: Vec, + ) -> PyResult<()> { + let particles: Vec<_> = particles.into_iter().map(|p| p.0).collect(); + self.0 + .extend_chunk(viewer.backend(), chunk.0, particles) + .map_err(gpu_err) + } + + fn remove_chunk( + &mut self, + viewer: PyRef, + chunk: NexusParticleChunk, + ) -> PyResult<()> { + self.0 + .remove_chunk(viewer.backend(), chunk.0) + .map_err(gpu_err) + } + // --- lifecycle -------------------------------------------------------- fn counts(&self) -> NexusCounts { @@ -317,6 +435,7 @@ impl NexusState { impulse_joints: c.impulse_joints, multibodies: c.multibodies, multibody_dofs: c.multibody_dofs, + particles: c.particles, } } @@ -338,7 +457,7 @@ impl NexusPipeline { NexusPipeline(RNexusPipeline::default()) } - /// Compiles all GPU pipelines up-front. + /// Compiles all GPU pipelines up-front (RBD + MPM). fn preload_pipelines(&mut self, viewer: PyRef) -> PyResult<()> { self.0 .preload_pipelines(viewer.backend(), NexusPipelineMask::all()) From 92c4b3b7590997a042b0a8f4394906f2fdcc4c16 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 14 Aug 2026 23:55:17 +0200 Subject: [PATCH 7/7] chore: remove dead code --- crates/examples3d/bench_joints3.rs | 641 ------------------ .../examples3d/bench_multibody_pendulum3.rs | 338 --------- crates/examples3d/bench_urdf3.rs | 279 -------- 3 files changed, 1258 deletions(-) delete mode 100644 crates/examples3d/bench_joints3.rs delete mode 100644 crates/examples3d/bench_multibody_pendulum3.rs delete mode 100644 crates/examples3d/bench_urdf3.rs diff --git a/crates/examples3d/bench_joints3.rs b/crates/examples3d/bench_joints3.rs deleted file mode 100644 index aec5fc7e..00000000 --- a/crates/examples3d/bench_joints3.rs +++ /dev/null @@ -1,641 +0,0 @@ -//! Headless benchmark mirroring the `joints3` demo (multibody variant). -//! -//! Builds the same scene as `crates/examples3d/joints3.rs` with -//! `use_articulations = true`: prismatic / revolute / fixed / spherical joint -//! chains, including motors and limits, all wired through `MultibodyJointSet`. -//! No rendering, just timed stepping. -//! -//! Usage: -//! -//! ```text -//! cargo run -p nexus_examples_3d --release \ -//! --bin bench_joints3 -- [num_batches] [num_warmup] [num_iters] [num_substeps] -//! ``` -//! -//! Defaults: 1 batch, 10 warmup steps, 200 timed steps, 4 substeps. - -use std::collections::HashMap; -use std::time::Duration; - -use khal::backend::GpuBackend as KhalGpuBackend; -use khal::backend::WebGpu; -use khal::re_exports::wgpu; -use nexus_viewer3d::SimulationState; -use nexus_viewer3d::nexus::rbd::dynamics::RbdSimParams; -use nexus_viewer3d::rbd::BatchEnvironment; -use nexus_viewer3d::rbd::GpuBackend; -use nexus_viewer3d::rbd::backend::SimulationBackend; -use rapier3d::prelude::*; - -fn create_prismatic_joints( - bodies: &mut RigidBodySet, - colliders: &mut ColliderSet, - multibody_joints: &mut MultibodyJointSet, - origin: Vector, - num: usize, -) { - let rad = 0.4; - let shift = 2.0; - - let ground = RigidBodyBuilder::fixed().translation(origin); - let mut curr_parent = bodies.insert(ground); - let collider = ColliderBuilder::cuboid(rad, rad, rad); - colliders.insert_with_parent(collider, curr_parent, bodies); - - for i in 0..num { - let z = origin.z + (i + 1) as f32 * shift; - let rb = RigidBodyBuilder::dynamic().translation(Vector::new(origin.x, origin.y, z)); - let curr_child = bodies.insert(rb); - let collider = ColliderBuilder::cuboid(rad, rad, rad); - colliders.insert_with_parent(collider, curr_child, bodies); - - let axis = if i % 2 == 0 { - Vector::new(1.0f32, 1.0, 0.0).normalize() - } else { - Vector::new(-1.0f32, 1.0, 0.0).normalize() - }; - - let prism = PrismaticJointBuilder::new(axis) - .local_anchor1(Vector::new(0.0, 0.0, 0.0)) - .local_anchor2(Vector::new(0.0, 0.0, -shift)) - .limits([-2.0, 2.0]); - - multibody_joints.insert(curr_parent, curr_child, prism, true); - curr_parent = curr_child; - } -} - -fn create_actuated_prismatic_joints( - bodies: &mut RigidBodySet, - colliders: &mut ColliderSet, - multibody_joints: &mut MultibodyJointSet, - origin: Vector, - num: usize, -) { - let rad = 0.4; - let shift = 2.0; - - let ground = RigidBodyBuilder::fixed().translation(origin); - let mut curr_parent = bodies.insert(ground); - let collider = ColliderBuilder::cuboid(rad, rad, rad); - colliders.insert_with_parent(collider, curr_parent, bodies); - - for i in 0..num { - let z = origin.z + (i + 1) as f32 * shift; - let rb = RigidBodyBuilder::dynamic().translation(Vector::new(origin.x, origin.y, z)); - let curr_child = bodies.insert(rb); - let collider = ColliderBuilder::cuboid(rad, rad, rad); - colliders.insert_with_parent(collider, curr_child, bodies); - - let axis = if i % 2 == 0 { - Vector::new(1.0, 1.0, 0.0).normalize() - } else { - Vector::new(-1.0, 1.0, 0.0).normalize() - }; - - let mut prism = PrismaticJointBuilder::new(axis) - .local_anchor1(Vector::new(0.0, 0.0, shift)) - .local_anchor2(Vector::new(0.0, 0.0, 0.0)) - .build(); - - if i == 0 { - prism - .set_motor_velocity(2.0, 1.0e5) - .set_limits([-2.0, 5.0]) - .set_motor_max_force(100.0); - } else if i == 1 { - prism - .set_limits([-Real::MAX, 5.0]) - .set_motor_velocity(6.0, 1.0e3) - .set_motor_max_force(100.0); - } else if i > 1 { - prism - .set_motor_position(2.0, 1.0e3, 1.0e2) - .set_motor_max_force(60.0); - } - - multibody_joints.insert(curr_parent, curr_child, prism, true); - curr_parent = curr_child; - } -} - -fn create_revolute_joints( - bodies: &mut RigidBodySet, - colliders: &mut ColliderSet, - multibody_joints: &mut MultibodyJointSet, - origin: Vector, - num: usize, -) { - let rad = 0.4; - let shift = 2.0; - - let ground = RigidBodyBuilder::fixed().translation(Vector::new(origin.x, origin.y, 0.0)); - let mut curr_parent = bodies.insert(ground); - let collider = ColliderBuilder::cuboid(rad, rad, rad); - colliders.insert_with_parent(collider, curr_parent, bodies); - - for i in 0..num { - let z = origin.z + i as f32 * shift * 2.0 + shift; - let positions = [ - Pose::from_translation(Vector::new(origin.x, origin.y, z)), - Pose::from_translation(Vector::new(origin.x + shift, origin.y, z)), - Pose::from_translation(Vector::new(origin.x + shift, origin.y, z + shift)), - Pose::from_translation(Vector::new(origin.x, origin.y, z + shift)), - ]; - let mut handles = [curr_parent; 4]; - for k in 0..4 { - let rb = RigidBodyBuilder::dynamic().pose(positions[k]); - handles[k] = bodies.insert(rb); - let collider = ColliderBuilder::cuboid(rad, rad, rad); - colliders.insert_with_parent(collider, handles[k], bodies); - } - let x = Vector::X; - let z = Vector::Z; - let revs = [ - RevoluteJointBuilder::new(z).local_anchor2(Vector::new(0.0, 0.0, -shift)), - RevoluteJointBuilder::new(x).local_anchor2(Vector::new(-shift, 0.0, 0.0)), - RevoluteJointBuilder::new(z).local_anchor2(Vector::new(0.0, 0.0, -shift)), - RevoluteJointBuilder::new(x).local_anchor2(Vector::new(shift, 0.0, 0.0)), - ]; - multibody_joints.insert(curr_parent, handles[0], revs[0], true); - multibody_joints.insert(handles[0], handles[1], revs[1], true); - multibody_joints.insert(handles[1], handles[2], revs[2], true); - multibody_joints.insert(handles[2], handles[3], revs[3], true); - curr_parent = handles[3]; - } -} - -fn create_revolute_joints_with_limits( - bodies: &mut RigidBodySet, - colliders: &mut ColliderSet, - multibody_joints: &mut MultibodyJointSet, - origin: Vector, -) { - let origin_v = origin; - let ground = bodies.insert(RigidBodyBuilder::fixed().translation(origin_v)); - colliders.insert_with_parent(ColliderBuilder::cuboid(0.1, 0.1, 0.1), ground, bodies); - - let shift = Vector::new(0.0, 0.0, 6.0); - let platform1 = bodies.insert(RigidBodyBuilder::dynamic().translation(origin_v + shift)); - colliders.insert_with_parent(ColliderBuilder::cuboid(4.0, 0.2, 2.0), platform1, bodies); - - let platform2 = bodies.insert(RigidBodyBuilder::dynamic().translation(origin_v + shift * 2.0)); - colliders.insert_with_parent(ColliderBuilder::cuboid(4.0, 0.2, 2.0), platform2, bodies); - - let z = Vector::Z; - let joint1 = RevoluteJointBuilder::new(z) - .local_anchor1(shift) - .limits([-0.2, 0.2]); - multibody_joints.insert(ground, platform1, joint1, true); - - let joint2 = RevoluteJointBuilder::new(z) - .local_anchor2(-shift) - .limits([-0.2, 0.2]); - multibody_joints.insert(platform1, platform2, joint2, true); - - let cuboid_body1 = bodies.insert( - RigidBodyBuilder::dynamic().translation(origin_v + shift + Vector::new(-2.0, 4.0, 0.0)), - ); - colliders.insert_with_parent( - ColliderBuilder::cuboid(0.6, 0.6, 0.6).friction(1.0), - cuboid_body1, - bodies, - ); - let cuboid_body2 = bodies.insert( - RigidBodyBuilder::dynamic() - .translation(origin_v + shift * 2.0 + Vector::new(2.0, 16.0, 0.0)), - ); - colliders.insert_with_parent( - ColliderBuilder::cuboid(0.6, 0.6, 0.6).friction(1.0), - cuboid_body2, - bodies, - ); -} - -fn create_fixed_joints( - bodies: &mut RigidBodySet, - colliders: &mut ColliderSet, - impulse_joints: &mut ImpulseJointSet, - multibody_joints: &mut MultibodyJointSet, - origin: Vector, - num: usize, -) { - let rad = 0.4; - let shift = 1.0; - let mut body_handles = Vec::new(); - - for i in 0..num { - for k in 0..num { - let fk = k as f32; - let fi = i as f32; - let status = if i == 0 && (k % 4 == 0 && k != num - 2 || k == num - 1) { - RigidBodyType::Fixed - } else { - RigidBodyType::Dynamic - }; - let rb = RigidBodyBuilder::new(status).translation(Vector::new( - origin.x + fk * shift, - origin.y, - origin.z + fi * shift, - )); - let child = bodies.insert(rb); - let collider = ColliderBuilder::ball(rad); - colliders.insert_with_parent(collider, child, bodies); - - if i > 0 { - let parent_index = body_handles.len() - num; - let parent_handle = body_handles[parent_index]; - let joint = FixedJointBuilder::new().local_anchor2(Vector::new(0.0, 0.0, -shift)); - multibody_joints.insert(parent_handle, child, joint, true); - } - - if k > 0 { - let parent_index = body_handles.len() - 1; - let parent_handle = body_handles[parent_index]; - let joint = FixedJointBuilder::new().local_anchor2(Vector::new(-shift, 0.0, 0.0)); - impulse_joints.insert(parent_handle, child, joint, true); - } - - body_handles.push(child); - } - } -} - -fn create_spherical_joints( - bodies: &mut RigidBodySet, - colliders: &mut ColliderSet, - impulse_joints: &mut ImpulseJointSet, - multibody_joints: &mut MultibodyJointSet, - num: usize, -) { - let rad = 0.4; - let shift = 1.0; - let mut body_handles = Vec::new(); - - for k in 0..num { - for i in 0..num { - let fk = k as f32; - let fi = i as f32; - let status = if i == 0 && (k % 4 == 0 || k == num - 1) { - RigidBodyType::Fixed - } else { - RigidBodyType::Dynamic - }; - let rb = RigidBodyBuilder::new(status).translation(Vector::new( - fk * shift, - 0.0, - fi * shift * 2.0, - )); - let child = bodies.insert(rb); - let collider = ColliderBuilder::capsule_z(rad * 1.25, rad); - colliders.insert_with_parent(collider, child, bodies); - - if i > 0 { - let parent = *body_handles.last().unwrap(); - let joint = - SphericalJointBuilder::new().local_anchor2(Vector::new(0.0, 0.0, -shift * 2.0)); - multibody_joints.insert(parent, child, joint, true); - } - if k > 0 { - let parent = body_handles[body_handles.len() - num]; - let joint = - SphericalJointBuilder::new().local_anchor2(Vector::new(-shift, 0.0, 0.0)); - impulse_joints.insert(parent, child, joint, true); - } - body_handles.push(child); - } - } -} - -fn create_spherical_joints_with_limits( - bodies: &mut RigidBodySet, - colliders: &mut ColliderSet, - multibody_joints: &mut MultibodyJointSet, - origin: Vector, -) { - let shift = Vector::new(0.0, 0.0, 3.0); - let origin_v = origin; - let ground = bodies.insert(RigidBodyBuilder::fixed().translation(origin_v)); - colliders.insert_with_parent(ColliderBuilder::cuboid(0.1, 0.1, 0.1), ground, bodies); - - let ball1 = bodies.insert( - RigidBodyBuilder::dynamic() - .translation(origin_v + shift) - .linvel(Vector::new(20.0, 20.0, 0.0)), - ); - colliders.insert_with_parent(ColliderBuilder::cuboid(1.0, 1.0, 1.0), ball1, bodies); - - let ball2 = bodies.insert(RigidBodyBuilder::dynamic().translation(origin_v + shift * 2.0)); - colliders.insert_with_parent(ColliderBuilder::cuboid(1.0, 1.0, 1.0), ball2, bodies); - - let joint1 = SphericalJointBuilder::new() - .local_anchor2(-shift) - .limits(JointAxis::LinX, [-0.2, 0.2]) - .limits(JointAxis::LinY, [-0.2, 0.2]); - let joint2 = SphericalJointBuilder::new() - .local_anchor2(-shift) - .limits(JointAxis::LinX, [-0.3, 0.3]) - .limits(JointAxis::LinY, [-0.3, 0.3]); - - multibody_joints.insert(ground, ball1, joint1, true); - multibody_joints.insert(ball1, ball2, joint2, true); -} - -fn create_actuated_revolute_joints( - bodies: &mut RigidBodySet, - colliders: &mut ColliderSet, - multibody_joints: &mut MultibodyJointSet, - origin: Vector, - num: usize, -) { - let rad = 0.4; - let shift = 2.0; - let z = Vector::Z; - let joint_template = RevoluteJointBuilder::new(z).local_anchor2(Vector::new(0.0, 0.0, -shift)); - let mut parent_handle = RigidBodyHandle::invalid(); - - for i in 0..num { - let fi = i as f32; - let status = if i == 0 { - RigidBodyType::Fixed - } else { - RigidBodyType::Dynamic - }; - let shifty = (i >= 1) as u32 as f32 * -2.0; - let rb = RigidBodyBuilder::new(status).translation(Vector::new( - origin.x, - origin.y + shifty, - origin.z + fi * shift, - )); - let child = bodies.insert(rb); - let collider = ColliderBuilder::cuboid(rad * 2.0, rad * 6.0 / (fi + 1.0), rad); - colliders.insert_with_parent(collider, child, bodies); - - if i > 0 { - let mut joint = joint_template.motor_model(MotorModel::AccelerationBased); - if i % 3 == 1 { - joint = joint.motor_velocity(-20.0, 100.0); - } else if i == num - 1 { - joint = joint.motor_position(std::f32::consts::FRAC_PI_2, 200.0, 100.0); - } - if i == 1 { - joint = joint - .local_anchor2(Vector::new(0.0, 2.0, -shift)) - .motor_velocity(-2.0, 1000.0); - } - multibody_joints.insert(parent_handle, child, joint, true); - } - parent_handle = child; - } -} - -fn create_actuated_spherical_joints( - bodies: &mut RigidBodySet, - colliders: &mut ColliderSet, - multibody_joints: &mut MultibodyJointSet, - origin: Vector, - num: usize, -) { - let rad = 0.4; - let shift = 2.0; - let joint_template = SphericalJointBuilder::new().local_anchor1(Vector::new(0.0, 0.0, shift)); - let mut parent_handle = RigidBodyHandle::invalid(); - - for i in 0..num { - let fi = i as f32; - let status = if i == 0 { - RigidBodyType::Fixed - } else { - RigidBodyType::Dynamic - }; - let rb = RigidBodyBuilder::new(status).translation(Vector::new( - origin.x, - origin.y, - origin.z + fi * shift, - )); - let child = bodies.insert(rb); - let collider = ColliderBuilder::capsule_y(rad * 2.0 / (fi + 1.0), rad); - colliders.insert_with_parent(collider, child, bodies); - - if i > 0 { - let mut joint = joint_template; - if i == 1 { - joint = joint - .motor_velocity(JointAxis::AngX, 0.0, 0.1) - .motor_velocity(JointAxis::AngY, 0.5, 0.1) - .motor_velocity(JointAxis::AngZ, -2.0, 0.1); - } else if i == num - 1 { - joint = joint - .motor_position(JointAxis::AngX, 0.0, 0.2, 1.0) - .motor_position(JointAxis::AngY, 1.0, 0.2, 1.0) - .motor_position(JointAxis::AngZ, std::f32::consts::FRAC_PI_2, 0.2, 1.0); - } - multibody_joints.insert(parent_handle, child, joint, true); - } - parent_handle = child; - } -} - -fn build_one_batch(num_substeps: u32) -> BatchEnvironment { - let mut bodies = RigidBodySet::new(); - let mut colliders = ColliderSet::new(); - let mut impulse_joints = ImpulseJointSet::new(); - let mut multibody_joints = MultibodyJointSet::new(); - - create_prismatic_joints( - &mut bodies, - &mut colliders, - &mut multibody_joints, - Vector::new(20.0, 5.0, 0.0), - 4, - ); - create_actuated_prismatic_joints( - &mut bodies, - &mut colliders, - &mut multibody_joints, - Vector::new(25.0, 5.0, 0.0), - 4, - ); - create_revolute_joints( - &mut bodies, - &mut colliders, - &mut multibody_joints, - Vector::new(20.0, 0.0, 0.0), - 3, - ); - create_revolute_joints_with_limits( - &mut bodies, - &mut colliders, - &mut multibody_joints, - Vector::new(34.0, 0.0, 0.0), - ); - create_fixed_joints( - &mut bodies, - &mut colliders, - &mut impulse_joints, - &mut multibody_joints, - Vector::new(0.0, 10.0, 0.0), - 10, - ); - create_actuated_revolute_joints( - &mut bodies, - &mut colliders, - &mut multibody_joints, - Vector::new(20.0, 10.0, 0.0), - 6, - ); - create_actuated_spherical_joints( - &mut bodies, - &mut colliders, - &mut multibody_joints, - Vector::new(13.0, 10.0, 0.0), - 3, - ); - create_spherical_joints( - &mut bodies, - &mut colliders, - &mut impulse_joints, - &mut multibody_joints, - 9, - ); - create_spherical_joints_with_limits( - &mut bodies, - &mut colliders, - &mut multibody_joints, - Vector::new(-5.0, 0.0, 0.0), - ); - - let mut sim_params = RbdSimParams::default(); - sim_params.num_solver_iterations = num_substeps; - BatchEnvironment { - bodies, - colliders, - impulse_joints, - multibody_joints, - sim_params, - visuals: HashMap::new(), - } -} - -fn build_scene(num_batches: usize, num_substeps: u32) -> SimulationState { - let envs = (0..num_batches.max(1)) - .map(|_| build_one_batch(num_substeps)) - .collect(); - SimulationState::from_environments(envs) -} - -struct Sample { - label: &'static str, - per_step_avg: Duration, - per_step_p50: Duration, - per_step_min: Duration, - per_step_max: Duration, -} - -impl Sample { - fn fmt_us(d: Duration) -> String { - format!("{:>10.2} µs", d.as_secs_f64() * 1.0e6) - } - - fn print(&self) { - println!( - " {:<10} avg {} p50 {} min {} max {}", - self.label, - Self::fmt_us(self.per_step_avg), - Self::fmt_us(self.per_step_p50), - Self::fmt_us(self.per_step_min), - Self::fmt_us(self.per_step_max), - ); - } -} - -async fn bench_backend( - label: &'static str, - backend: &KhalGpuBackend, - state: &SimulationState, - n_warmup: usize, - n_iters: usize, -) -> Sample { - let mut phys = GpuBackend::try_new(backend, state) - .await - .unwrap_or_else(|e| panic!("{label} backend init failed: {e}")); - - for _ in 0..n_warmup { - let _ = phys.step(None).await; - } - - let mut samples = Vec::with_capacity(n_iters); - let mut last_stats = None; - for i in 0..n_iters { - let stats = phys.step(None).await; - samples.push(stats.total_simulation_time_without_readback); - if i == n_iters - 1 { - last_stats = Some(stats); - } - } - samples.sort(); - - if let Some(stats) = last_stats { - if !stats.gpu_pass_times.is_empty() { - let mut passes = stats.gpu_pass_times.clone(); - passes.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); - println!( - " top passes ({} total, {:.3} ms):", - passes.len(), - stats.gpu_total_time - ); - for (l, ms) in passes.iter().take(15) { - println!(" {:>9.3} ms {}", ms, l); - } - } - } - - let total: Duration = samples.iter().sum(); - let per_step_avg = total / (n_iters as u32); - let per_step_p50 = samples[n_iters / 2]; - let per_step_min = *samples.first().unwrap(); - let per_step_max = *samples.last().unwrap(); - - Sample { - label, - per_step_avg, - per_step_p50, - per_step_min, - per_step_max, - } -} - -async fn webgpu_backend() -> KhalGpuBackend { - let limits = wgpu::Limits { - max_buffer_size: 1_000_000_000, - max_storage_buffer_binding_size: 1_000_000_000, - max_storage_buffers_per_shader_stage: 14, - max_compute_workgroup_storage_size: 19_904, - ..Default::default() - }; - let mut webgpu = WebGpu::new(wgpu::Features::default(), limits) - .await - .expect("Failed to initialize WebGPU backend"); - webgpu.force_buffer_copy_src = true; - KhalGpuBackend::WebGpu(webgpu) -} - -async fn run(num_batches: usize, n_warmup: usize, n_iters: usize, num_substeps: u32) { - println!( - "Joints3 multibody benchmark — {num_batches} batches, num_substeps={num_substeps}, \ - {n_warmup} warmup + {n_iters} timed steps" - ); - let state = build_scene(num_batches, num_substeps); - let webgpu = webgpu_backend().await; - let s = bench_backend("WebGPU", &webgpu, &state, n_warmup, n_iters).await; - s.print(); -} - -fn main() { - let args: Vec = std::env::args().collect(); - let num_batches = args.get(1).and_then(|s| s.parse().ok()).unwrap_or(1); - let n_warmup = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(10); - let n_iters = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(200); - let num_substeps = args.get(4).and_then(|s| s.parse().ok()).unwrap_or(4u32); - pollster::block_on(run(num_batches, n_warmup, n_iters, num_substeps)); -} diff --git a/crates/examples3d/bench_multibody_pendulum3.rs b/crates/examples3d/bench_multibody_pendulum3.rs deleted file mode 100644 index 6bdc1b5a..00000000 --- a/crates/examples3d/bench_multibody_pendulum3.rs +++ /dev/null @@ -1,338 +0,0 @@ -//! Headless benchmark for the 3D multibody pendulum scene. -//! -//! Runs the same scene as `multibody_pendulum3` (a 20-link revolute-joint -//! pendulum) on both the WebGPU and the nexus-CPU backends, with no graphics, -//! and reports per-step wall-clock times so we can track the perf gap and -//! verify that shader optimisations are landing. -//! -//! Usage (the `cpu` feature wires the nexus-CPU backend through nexus_viewer3d): -//! -//! ```text -//! cargo run -p nexus_examples_3d --release \ -//! --features cpu \ -//! --bin bench_multibody_pendulum3 -- [num_links] [num_warmup] [num_iters] -//! ``` -//! -//! Defaults: 20 links, 10 warmup steps, 200 timed steps. - -use std::time::Duration; - -use khal::backend::GpuBackend as KhalGpuBackend; -use khal::backend::WebGpu; -use khal::re_exports::wgpu; -use nexus_viewer3d::SimulationState; -use nexus_viewer3d::nexus::rbd::dynamics::RbdSimParams; -use nexus_viewer3d::rbd::BatchEnvironment; -use nexus_viewer3d::rbd::GpuBackend; -use nexus_viewer3d::rbd::backend::SimulationBackend; -use rapier3d::prelude::*; -use std::collections::HashMap; - -fn build_one_batch(num_links: usize) -> BatchEnvironment { - build_one_batch_with_substeps(num_links, 4) -} - -fn build_one_batch_with_substeps(num_links: usize, num_substeps: u32) -> BatchEnvironment { - let mut bodies = RigidBodySet::new(); - let mut colliders = ColliderSet::new(); - let impulse_joints = ImpulseJointSet::new(); - let mut multibody_joints = MultibodyJointSet::new(); - - let rad = 0.4; - let link_len = 2.0; - - let root_body = RigidBodyBuilder::fixed(); - let mut parent_handle = bodies.insert(root_body); - let root_collider = ColliderBuilder::cuboid(rad, rad, rad); - colliders.insert_with_parent(root_collider, parent_handle, &mut bodies); - - for i in 0..num_links { - let x = (i as f32 + 1.0) * link_len; - let rigid_body = RigidBodyBuilder::dynamic().translation(Vec3::new(x, 0.0, 0.0)); - let handle = bodies.insert(rigid_body); - let collider = ColliderBuilder::cuboid(link_len * 0.5, rad, rad); - colliders.insert_with_parent(collider, handle, &mut bodies); - - let parent_anchor = if i == 0 { - Vec3::ZERO - } else { - Vec3::new(link_len * 0.8, 0.0, 0.0) - }; - let joint = RevoluteJointBuilder::new(Vec3::Z) - .local_anchor1(parent_anchor) - .local_anchor2(Vec3::new(-link_len * 0.8, 0.0, 0.0)) - .build(); - multibody_joints.insert(parent_handle, handle, joint, true); - - parent_handle = handle; - } - - let mut sim_params = RbdSimParams::default(); - sim_params.num_solver_iterations = num_substeps; - BatchEnvironment { - bodies, - colliders, - impulse_joints, - multibody_joints, - sim_params, - visuals: HashMap::new(), - } -} - -fn build_scene(num_links: usize, num_batches: usize) -> SimulationState { - build_scene_with_substeps(num_links, num_batches, 4) -} - -fn build_scene_with_substeps( - num_links: usize, - num_batches: usize, - num_substeps: u32, -) -> SimulationState { - let environments = (0..num_batches.max(1)) - .map(|_| build_one_batch_with_substeps(num_links, num_substeps)) - .collect(); - SimulationState::from_environments(environments) -} - -struct Sample { - label: &'static str, - per_step_avg: Duration, - per_step_p50: Duration, - per_step_min: Duration, - per_step_max: Duration, -} - -impl Sample { - fn fmt_us(d: Duration) -> String { - format!("{:>10.2} µs", d.as_secs_f64() * 1.0e6) - } - - fn print(&self) { - println!( - " {:<10} avg {} p50 {} min {} max {}", - self.label, - Self::fmt_us(self.per_step_avg), - Self::fmt_us(self.per_step_p50), - Self::fmt_us(self.per_step_min), - Self::fmt_us(self.per_step_max), - ); - } -} - -async fn bench_backend( - label: &'static str, - backend: &KhalGpuBackend, - state: &SimulationState, - n_warmup: usize, - n_iters: usize, -) -> Sample { - bench_backend_inner(label, backend, state, n_warmup, n_iters, true).await -} - -async fn bench_backend_inner( - label: &'static str, - backend: &KhalGpuBackend, - state: &SimulationState, - n_warmup: usize, - n_iters: usize, - print_passes: bool, -) -> Sample { - let mut phys = GpuBackend::try_new(backend, state) - .await - .unwrap_or_else(|e| panic!("{label} backend init failed: {e}")); - - // Warmup — first steps include shader compilation and pipeline cache fill, - // we don't want those skewing the measurement. - for _ in 0..n_warmup { - let _ = phys.step(None).await; - } - - // Measure `total_simulation_time_without_readback` instead of wall-clocking - // the whole `phys.step()` — the post-step `auto_resize_buffers` + pose - // readback + timestamp readback each add a CPU-GPU round trip whose latency - // is dominated by driver/OS scheduler jitter, not by the pipeline itself. - let mut samples = Vec::with_capacity(n_iters); - let mut last_stats = None; - for i in 0..n_iters { - let stats = phys.step(None).await; - samples.push(stats.total_simulation_time_without_readback); - if i == n_iters - 1 { - last_stats = Some(stats); - } - } - samples.sort(); - - // Print the top per-pass GPU timings from the last iteration (so the - // user can see which kernel still dominates after warmup). The CPU - // backend reports an empty list — guard with `is_empty()`. - if print_passes { - if let Some(stats) = last_stats { - if !stats.gpu_pass_times.is_empty() { - let mut passes = stats.gpu_pass_times.clone(); - passes.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); - println!( - " top passes ({} total, {:.3} ms):", - passes.len(), - stats.gpu_total_time - ); - for (l, ms) in passes.iter().take(8) { - println!(" {:>9.3} ms {}", ms, l); - } - } - } - } - - let total: Duration = samples.iter().sum(); - let per_step_avg = total / (n_iters as u32); - let per_step_p50 = samples[n_iters / 2]; - let per_step_min = *samples.first().unwrap(); - let per_step_max = *samples.last().unwrap(); - - Sample { - label, - per_step_avg, - per_step_p50, - per_step_min, - per_step_max, - } -} - -async fn webgpu_backend() -> KhalGpuBackend { - // The pendulum scene's narrow-phase shader needs more storage buffers - // and a larger workgroup-storage budget than wgpu's defaults — mirror - // the limits the viewer requests. - let limits = wgpu::Limits { - max_buffer_size: 1_000_000_000, - max_storage_buffer_binding_size: 1_000_000_000, - max_storage_buffers_per_shader_stage: 14, - max_compute_workgroup_storage_size: 19_904, - ..Default::default() - }; - let mut webgpu = WebGpu::new(wgpu::Features::default(), limits) - .await - .expect("Failed to initialize WebGPU backend"); - webgpu.force_buffer_copy_src = true; - KhalGpuBackend::WebGpu(webgpu) -} - -async fn run( - num_links: usize, - num_batches: usize, - n_warmup: usize, - n_iters: usize, - num_substeps: u32, -) { - println!( - "Multibody pendulum benchmark — {num_links} links × {num_batches} batches, \ - num_substeps={num_substeps}, {n_warmup} warmup + {n_iters} timed steps" - ); - - let state = build_scene_with_substeps(num_links, num_batches, num_substeps); - let webgpu = webgpu_backend().await; - - let webgpu_sample = { - let s = bench_backend("WebGPU", &webgpu, &state, n_warmup, n_iters).await; - s.print(); - s - }; - - // Nexus-CPU backend (same pipeline, executed on CPU). Only available when - // built with `--features cpu`. - #[cfg(feature = "cpu")] - let cpu_sample = { - let backend = KhalGpuBackend::Cpu; - let s = bench_backend("Nexus-CPU", &backend, &state, n_warmup, n_iters).await; - s.print(); - s - }; - - #[cfg(feature = "cpu")] - { - let ratio = - webgpu_sample.per_step_avg.as_secs_f64() / cpu_sample.per_step_avg.as_secs_f64(); - println!( - " → WebGPU/CPU ratio (avg): {:.2}× {}", - ratio, - if ratio > 1.0 { - "(GPU slower)" - } else { - "(GPU faster)" - } - ); - } - #[cfg(not(feature = "cpu"))] - { - let _ = webgpu_sample; - println!(" (rebuild with --features cpu to also benchmark the nexus CPU backend)"); - } -} - -/// Sweep mode: hold `num_links` fixed and vary `num_batches` over a power-of-two -/// range, reporting the crossover where GPU catches up to CPU. -#[cfg(feature = "cpu")] -async fn sweep(num_links: usize, max_batches: usize, n_warmup: usize, n_iters: usize) { - println!( - "Multibody pendulum sweep — {num_links} links, batches 1..={max_batches} (×2 each step), \ - {n_warmup} warmup + {n_iters} timed steps each" - ); - println!( - "{:>7} {:>14} {:>14} {:>10} {}", - "batches", "WebGPU avg", "CPU avg", "ratio", "verdict" - ); - - let webgpu = webgpu_backend().await; - let cpu = KhalGpuBackend::Cpu; - - let mut bs = 1; - while bs <= max_batches { - let state = build_scene(num_links, bs); - let g = bench_backend_inner("WebGPU", &webgpu, &state, n_warmup, n_iters, false).await; - let c = bench_backend_inner("Nexus-CPU", &cpu, &state, n_warmup, n_iters, false).await; - let ratio = g.per_step_avg.as_secs_f64() / c.per_step_avg.as_secs_f64(); - let verdict = if ratio < 1.0 { - "GPU faster" - } else if ratio < 4.0 { - "GPU within 4×" - } else { - "GPU slower" - }; - println!( - "{:>7} {:>14} {:>14} {:>9.2}× {}", - bs, - Sample::fmt_us(g.per_step_avg), - Sample::fmt_us(c.per_step_avg), - ratio, - verdict, - ); - bs *= 2; - } -} - -fn main() { - // Args: - // bench_multibody_pendulum3 [num_links] [num_batches] [num_warmup] [num_iters] - // bench_multibody_pendulum3 sweep [num_links] [max_batches] [num_warmup] [num_iters] - let args: Vec = std::env::args().collect(); - - if args.get(1).map(|s| s.as_str()) == Some("sweep") { - #[cfg(feature = "cpu")] - { - let num_links = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(20); - let max_batches = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(256); - let n_warmup = args.get(4).and_then(|s| s.parse().ok()).unwrap_or(5); - let n_iters = args.get(5).and_then(|s| s.parse().ok()).unwrap_or(30); - pollster::block_on(sweep(num_links, max_batches, n_warmup, n_iters)); - } - #[cfg(not(feature = "cpu"))] - eprintln!("sweep mode requires the `cpu` feature"); - return; - } - - let num_links = args.get(1).and_then(|s| s.parse().ok()).unwrap_or(20); - let num_batches = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(1); - let n_warmup = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(10); - let n_iters = args.get(4).and_then(|s| s.parse().ok()).unwrap_or(200); - let num_substeps = args.get(5).and_then(|s| s.parse().ok()).unwrap_or(4u32); - - pollster::block_on(run(num_links, num_batches, n_warmup, n_iters, num_substeps)); -} diff --git a/crates/examples3d/bench_urdf3.rs b/crates/examples3d/bench_urdf3.rs deleted file mode 100644 index 8c9a28e5..00000000 --- a/crates/examples3d/bench_urdf3.rs +++ /dev/null @@ -1,279 +0,0 @@ -//! Headless benchmark mirroring the `urdf3` demo. -//! -//! Same setup as `crates/examples3d/urdf3.rs`: load a URDF (defaulting to the -//! openarm_v10 robot), insert it as a multibody with `make_roots_fixed = true` -//! and self-contacts disabled, switch every joint to acceleration-based motors, -//! and step the simulation. The `apply_random_ang_motors` tick is replicated so -//! every motor's target velocity is re-randomised every 5 simulated seconds. -//! -//! Usage: -//! -//! ```text -//! cargo run -p nexus_examples_3d --release \ -//! --bin bench_urdf3 -- [num_batches] [num_warmup] [num_iters] [num_substeps] [path] -//! ``` -//! -//! Defaults: 1 batch, 10 warmup steps, 200 timed steps, 4 substeps, openarm path. - -use std::collections::HashMap; -use std::path::PathBuf; -use std::time::Duration; - -use khal::backend::GpuBackend as KhalGpuBackend; -use khal::backend::WebGpu; -use khal::re_exports::wgpu; -use nexus_viewer3d::SimulationState; -use nexus_viewer3d::nexus::rbd::dynamics::RbdSimParams; -use nexus_viewer3d::rbd::BatchEnvironment; -use nexus_viewer3d::rbd::GpuBackend; -use nexus_viewer3d::rbd::backend::SimulationBackend; -use rapier3d::prelude::*; -use rapier3d_urdf::{UrdfLoaderOptions, UrdfMultibodyOptions, UrdfRobot}; - -fn default_urdf_path() -> PathBuf { - // Mirrors `urdf3.rs`. The benchmark exits cleanly with an error if missing, - // so users can pass their own URDF as the last CLI argument. - PathBuf::from("/Users/sebcrozet/work/nexus-demos/XoQ/js/examples/assets/openarm_v10.urdf") -} - -fn build_one_batch(path: &PathBuf, num_substeps: u32) -> Option { - let mut bodies = RigidBodySet::new(); - let mut colliders = ColliderSet::new(); - let impulse_joints = ImpulseJointSet::new(); - let mut multibody_joints = MultibodyJointSet::new(); - - let scale = 40.0; - let options = UrdfLoaderOptions { - create_colliders_from_collision_shapes: true, - create_colliders_from_visual_shapes: true, - apply_imported_mass_props: true, - make_roots_fixed: true, - scale, - mesh_converter: None, - shift: Pose::from_parts( - Vec3::new(0.0, scale, 0.0), - Rotation::from_rotation_x(-std::f32::consts::FRAC_PI_2), - ), - collider_blueprint: ColliderBuilder::ball(0.5).collision_groups(InteractionGroups::none()), - ..UrdfLoaderOptions::default() - }; - - let (mut robot, _) = match UrdfRobot::from_file(path, options, None) { - Ok(r) => r, - Err(e) => { - eprintln!("Failed to load URDF file at {}: {e}", path.display()); - return None; - } - }; - - for urdf_joint in &mut robot.joints { - urdf_joint - .joint - .set_motor_model(JointAxis::AngX, MotorModel::AccelerationBased); - urdf_joint - .joint - .set_motor_velocity(JointAxis::AngX, 0.0, 1.0); - } - - let _ = robot.insert_using_multibody_joints( - &mut bodies, - &mut colliders, - &mut multibody_joints, - UrdfMultibodyOptions::DISABLE_SELF_CONTACTS, - ); - - let mut sim_params = RbdSimParams::default(); - sim_params.num_solver_iterations = num_substeps; - Some(BatchEnvironment { - bodies, - colliders, - impulse_joints, - multibody_joints, - sim_params, - visuals: HashMap::new(), - }) -} - -fn build_scene(path: &PathBuf, num_batches: usize, num_substeps: u32) -> Option { - let env = build_one_batch(path, num_substeps)?; - let mut envs = Vec::with_capacity(num_batches.max(1)); - envs.push(env); - for _ in 1..num_batches.max(1) { - envs.push(build_one_batch(path, num_substeps)?); - } - Some(SimulationState::from_environments(envs)) -} - -struct Sample { - label: &'static str, - per_step_avg: Duration, - per_step_p50: Duration, - per_step_min: Duration, - per_step_max: Duration, -} - -impl Sample { - fn fmt_us(d: Duration) -> String { - format!("{:>10.2} µs", d.as_secs_f64() * 1.0e6) - } - - fn print(&self) { - println!( - " {:<10} avg {} p50 {} min {} max {}", - self.label, - Self::fmt_us(self.per_step_avg), - Self::fmt_us(self.per_step_p50), - Self::fmt_us(self.per_step_min), - Self::fmt_us(self.per_step_max), - ); - } -} - -async fn bench_backend( - label: &'static str, - backend: &KhalGpuBackend, - state: &SimulationState, - n_warmup: usize, - n_iters: usize, - num_links: u32, -) -> Sample { - let mut phys = GpuBackend::try_new(backend, state) - .await - .unwrap_or_else(|e| panic!("{label} backend init failed: {e}")); - - // Mirrors the urdf3 demo's `apply_random_ang_motors` tick: every 5 simulated - // seconds, push a fresh random AngX motor target velocity to every link. - use rand::RngExt; - let mut rng = rand::rng(); - let dt = 1.0 / 60.0_f64; - let mut next_change_at = 0.0_f64; - let interval = 5.0; - let n_batches = phys.num_batches() as u32; - - let mut do_tick = |phys: &mut GpuBackend, sim_time: f64| { - if sim_time < next_change_at { - return; - } - next_change_at = sim_time + interval; - for batch in 0..n_batches { - for link_id in 0..num_links { - let target_vel: f32 = rng.random_range(-0.6f32..=0.6); - phys.set_multibody_motor_velocity(batch, link_id, JointAxis::AngX, target_vel); - } - } - }; - - // Warmup. - let mut sim_time = 0.0_f64; - for _ in 0..n_warmup { - do_tick(&mut phys, sim_time); - let _ = phys.step(None).await; - sim_time += dt; - } - - let mut samples = Vec::with_capacity(n_iters); - let mut last_stats = None; - for i in 0..n_iters { - do_tick(&mut phys, sim_time); - let stats = phys.step(None).await; - samples.push(stats.total_simulation_time_without_readback); - if i == n_iters - 1 { - last_stats = Some(stats); - } - sim_time += dt; - } - samples.sort(); - - if let Some(stats) = last_stats { - if !stats.gpu_pass_times.is_empty() { - let mut passes = stats.gpu_pass_times.clone(); - passes.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); - println!( - " top passes ({} total, {:.3} ms):", - passes.len(), - stats.gpu_total_time - ); - for (l, ms) in passes.iter().take(15) { - println!(" {:>9.3} ms {}", ms, l); - } - } - } - - let total: Duration = samples.iter().sum(); - let per_step_avg = total / (n_iters as u32); - let per_step_p50 = samples[n_iters / 2]; - let per_step_min = *samples.first().unwrap(); - let per_step_max = *samples.last().unwrap(); - - Sample { - label, - per_step_avg, - per_step_p50, - per_step_min, - per_step_max, - } -} - -async fn webgpu_backend() -> KhalGpuBackend { - let limits = wgpu::Limits { - max_buffer_size: 1_000_000_000, - max_storage_buffer_binding_size: 1_000_000_000, - max_storage_buffers_per_shader_stage: 14, - max_compute_workgroup_storage_size: 19_904, - ..Default::default() - }; - let mut webgpu = WebGpu::new(wgpu::Features::default(), limits) - .await - .expect("Failed to initialize WebGPU backend"); - webgpu.force_buffer_copy_src = true; - KhalGpuBackend::WebGpu(webgpu) -} - -async fn run( - path: PathBuf, - num_batches: usize, - n_warmup: usize, - n_iters: usize, - num_substeps: u32, -) { - println!( - "URDF3 multibody benchmark — {num_batches} batches, num_substeps={num_substeps}, \ - {n_warmup} warmup + {n_iters} timed steps (URDF: {})", - path.display() - ); - - let state = match build_scene(&path, num_batches, num_substeps) { - Some(s) => s, - None => { - eprintln!("Aborting benchmark — URDF could not be loaded."); - return; - } - }; - let webgpu = webgpu_backend().await; - let num_links = state.environments[0] - .multibody_joints - .multibodies() - .next() - .map(|mb| mb.num_links()) - .unwrap_or(0) as u32; - - let s = bench_backend("WebGPU", &webgpu, &state, n_warmup, n_iters, num_links).await; - s.print(); -} - -fn main() { - // Args: - // bench_urdf3 [num_batches] [num_warmup] [num_iters] [num_substeps] [path] - let args: Vec = std::env::args().collect(); - - let num_batches = args.get(1).and_then(|s| s.parse().ok()).unwrap_or(1); - let n_warmup = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(10); - let n_iters = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(200); - let num_substeps = args.get(4).and_then(|s| s.parse().ok()).unwrap_or(4u32); - let path = args - .get(5) - .map(PathBuf::from) - .unwrap_or_else(default_urdf_path); - - pollster::block_on(run(path, num_batches, n_warmup, n_iters, num_substeps)); -}