1use crate::math::{Fix128, Vec3Fix};
30#[cfg(feature = "std")]
31use crate::sdf_collider::SdfCollider;
32use crate::spatial::SpatialGrid;
33
34#[cfg(not(feature = "std"))]
35use alloc::vec;
36#[cfg(not(feature = "std"))]
37use alloc::vec::Vec;
38
39#[derive(Clone, Copy, Debug, PartialEq, Eq)]
45pub struct FluidConfig {
46 pub rest_density: Fix128,
48 pub kernel_radius: Fix128,
50 pub iterations: usize,
52 pub substeps: usize,
54 pub gravity: Vec3Fix,
56 pub damping: Fix128,
59 pub viscosity: Fix128,
61 pub vorticity_strength: Fix128,
63 pub surface_tension: Fix128,
66 pub particle_mass: Fix128,
68 pub relaxation: Fix128,
70}
71
72impl Default for FluidConfig {
73 fn default() -> Self {
74 Self {
75 rest_density: Fix128::from_int(1000),
76 kernel_radius: Fix128::from_ratio(1, 5), iterations: 4,
78 substeps: 2,
79 gravity: Vec3Fix::new(Fix128::ZERO, Fix128::from_int(-10), Fix128::ZERO),
80 damping: Fix128::from_ratio(99, 100),
81 viscosity: Fix128::from_ratio(1, 100),
82 vorticity_strength: Fix128::from_ratio(1, 1000),
83 surface_tension: Fix128::from_ratio(1, 100),
84 particle_mass: Fix128::from_ratio(1, 100),
85 relaxation: Fix128::ONE,
86 }
87 }
88}
89
90#[inline(always)]
96fn poly6(r_sq: Fix128, h: Fix128) -> Fix128 {
97 let h_sq = h * h;
98 if r_sq >= h_sq {
99 return Fix128::ZERO;
100 }
101 let diff = h_sq - r_sq;
102 let h9 = h_sq * h_sq * h_sq * h_sq * h;
106 let coeff = Fix128::from_ratio(315, 64) / (h9 * Fix128::PI);
107 coeff * diff * diff * diff
108}
109
110#[inline(always)]
112fn spiky_grad(r: Fix128, h: Fix128) -> Fix128 {
113 if r >= h || r.is_zero() {
114 return Fix128::ZERO;
115 }
116 let diff = h - r;
117 let h6 = h * h * h * h * h * h;
119 let coeff = Fix128::from_ratio(45, 1) / (h6 * Fix128::PI);
120 -coeff * diff * diff
121}
122
123#[repr(C, align(64))]
129pub struct Fluid {
130 pub positions: Vec<Vec3Fix>,
132 predicted: Vec<Vec3Fix>,
134 pub velocities: Vec<Vec3Fix>,
136 pub densities: Vec<Fix128>,
138 lambdas: Vec<Fix128>,
140 grid: SpatialGrid,
142 pub config: FluidConfig,
144 inv_rest_density: Fix128,
146}
147
148impl Fluid {
149 #[must_use]
151 pub fn new(positions: Vec<Vec3Fix>, config: FluidConfig) -> Self {
152 let n = positions.len();
153 let grid_dim = 32;
154 let inv_rest_density = if config.rest_density.is_zero() {
155 Fix128::ONE
156 } else {
157 Fix128::ONE / config.rest_density
158 };
159
160 Self {
161 predicted: positions.clone(),
162 velocities: vec![Vec3Fix::ZERO; n],
163 densities: vec![Fix128::ZERO; n],
164 lambdas: vec![Fix128::ZERO; n],
165 grid: SpatialGrid::new(config.kernel_radius, grid_dim),
166 positions,
167 config,
168 inv_rest_density,
169 }
170 }
171
172 #[must_use]
183 pub fn new_block(min: Vec3Fix, max: Vec3Fix, spacing: Fix128, config: FluidConfig) -> Self {
184 assert!(
185 spacing > Fix128::ZERO,
186 "Fluid::new_block: spacing must be strictly positive"
187 );
188 let mut positions = Vec::new();
189
190 let mut x = min.x;
191 while x <= max.x {
192 let mut y = min.y;
193 while y <= max.y {
194 let mut z = min.z;
195 while z <= max.z {
196 positions.push(Vec3Fix::new(x, y, z));
197 z = z + spacing;
198 }
199 y = y + spacing;
200 }
201 x = x + spacing;
202 }
203
204 Self::new(positions, config)
205 }
206
207 #[inline(always)]
209 #[must_use]
210 pub fn particle_count(&self) -> usize {
211 self.positions.len()
212 }
213
214 pub fn step(&mut self, dt: Fix128) {
216 let substep_dt = dt / Fix128::from_int(self.config.substeps as i64);
217 for _ in 0..self.config.substeps {
218 self.substep(substep_dt);
219 }
220 self.apply_frame_damping();
221 }
222
223 #[cfg(feature = "std")]
225 pub fn step_with_sdf(&mut self, dt: Fix128, sdf_colliders: &[SdfCollider]) {
226 let substep_dt = dt / Fix128::from_int(self.config.substeps as i64);
227 for _ in 0..self.config.substeps {
228 self.substep(substep_dt);
229 self.resolve_sdf_boundary(sdf_colliders);
230 }
231 self.apply_frame_damping();
232 }
233
234 fn substep(&mut self, dt: Fix128) {
236 let n = self.particle_count();
237 let h = self.config.kernel_radius;
238 let h_sq = h * h;
239
240 for i in 0..n {
244 self.velocities[i] = self.velocities[i] + self.config.gravity * dt;
245 self.predicted[i] = self.positions[i] + self.velocities[i] * dt;
246 }
247
248 self.grid.clear();
250 for i in 0..n {
251 self.grid.insert(i, self.predicted[i]);
252 }
253 self.grid.build();
254
255 let mut neighbors_buf = Vec::new();
257 for _ in 0..self.config.iterations {
258 for i in 0..n {
260 self.grid
261 .query_neighbors_into(self.predicted[i], h_sq, &mut neighbors_buf);
262 let mut density = Fix128::ZERO;
263 let mut sum_grad_sq = Fix128::ZERO;
264 let mut grad_i = Vec3Fix::ZERO;
265
266 for &j in &neighbors_buf {
267 let delta = self.predicted[i] - self.predicted[j];
268 let r_sq = delta.length_squared();
269
270 density = density + self.config.particle_mass * poly6(r_sq, h);
271
272 if i != j && r_sq < h_sq {
273 let r = r_sq.sqrt();
274 let grad_mag = spiky_grad(r, h);
275 if !r.is_zero() {
276 let grad = delta / r * grad_mag;
277 grad_i = grad_i + grad;
278 sum_grad_sq = sum_grad_sq + grad.length_squared();
279 }
280 }
281 }
282
283 self.densities[i] = density;
284
285 let constraint = density * self.inv_rest_density - Fix128::ONE;
288 sum_grad_sq = sum_grad_sq + grad_i.length_squared();
289
290 let epsilon = Fix128::from_ratio(1, 10000);
291 self.lambdas[i] = if sum_grad_sq.is_zero() {
292 Fix128::ZERO
293 } else {
294 -constraint / (sum_grad_sq + epsilon)
295 };
296 }
297
298 for i in 0..n {
300 self.grid
301 .query_neighbors_into(self.predicted[i], h_sq, &mut neighbors_buf);
302 let mut correction = Vec3Fix::ZERO;
303
304 for &j in &neighbors_buf {
305 if i == j {
306 continue;
307 }
308 let delta = self.predicted[i] - self.predicted[j];
309 let r_sq = delta.length_squared();
310 if r_sq >= h_sq || r_sq.is_zero() {
311 continue;
312 }
313
314 let r = r_sq.sqrt();
315 let grad_mag = spiky_grad(r, h);
316 let grad = delta / r * grad_mag;
317
318 let lambda_sum = self.lambdas[i] + self.lambdas[j];
319 correction = correction + grad * lambda_sum;
320 }
321
322 self.predicted[i] =
324 self.predicted[i] + correction * self.inv_rest_density * self.config.relaxation;
325 }
326 }
327
328 let inv_dt = Fix128::ONE / dt;
330 for i in 0..n {
331 self.velocities[i] = (self.predicted[i] - self.positions[i]) * inv_dt;
332 self.positions[i] = self.predicted[i];
333 }
334
335 self.apply_viscosity();
337
338 if !self.config.surface_tension.is_zero() {
340 self.apply_surface_tension(dt);
341 }
342
343 if !self.config.vorticity_strength.is_zero() {
345 self.apply_vorticity_confinement(dt);
346 }
347 }
348
349 fn apply_frame_damping(&mut self) {
351 let d = self.config.damping;
352 for v in &mut self.velocities {
353 *v = *v * d;
354 }
355 }
356
357 #[inline]
362 fn particle_volume(&self) -> Fix128 {
363 self.config.particle_mass * self.inv_rest_density
364 }
365
366 fn apply_viscosity(&mut self) {
368 let n = self.particle_count();
369 let h = self.config.kernel_radius;
370 let h_sq = h * h;
371 let c = self.config.viscosity;
372
373 let velocities_copy: Vec<Vec3Fix> = self.velocities.clone();
374 let mut neighbors_buf = Vec::new();
375
376 for i in 0..n {
377 self.grid
378 .query_neighbors_into(self.positions[i], h_sq, &mut neighbors_buf);
379 let mut avg_vel = Vec3Fix::ZERO;
380 let mut weight_sum = Fix128::ZERO;
381
382 for &j in &neighbors_buf {
383 if i == j {
384 continue;
385 }
386 let delta = self.positions[i] - self.positions[j];
387 let r_sq = delta.length_squared();
388 let w = poly6(r_sq, h);
389 avg_vel = avg_vel + (velocities_copy[j] - velocities_copy[i]) * w;
390 weight_sum = weight_sum + w;
391 }
392
393 if !weight_sum.is_zero() {
394 self.velocities[i] = self.velocities[i] + avg_vel * (c / weight_sum);
395 }
396 }
397 }
398
399 fn apply_surface_tension(&mut self, dt: Fix128) {
409 let n = self.particle_count();
410 let h = self.config.kernel_radius;
411 let h_sq = h * h;
412 let coeff = self.config.surface_tension * dt * self.particle_volume();
413 let mut neighbors_buf = Vec::new();
414
415 for i in 0..n {
416 self.grid
417 .query_neighbors_into(self.positions[i], h_sq, &mut neighbors_buf);
418 let mut force = Vec3Fix::ZERO;
419
420 for &j in &neighbors_buf {
421 if i == j {
422 continue;
423 }
424 let delta = self.positions[j] - self.positions[i];
425 let r_sq = delta.length_squared();
426 if r_sq.is_zero() || r_sq >= h_sq {
427 continue;
428 }
429 let r = r_sq.sqrt();
430 let w = poly6(r_sq, h);
431 force = force + delta / r * w;
432 }
433
434 self.velocities[i] = self.velocities[i] + force * coeff;
435 }
436 }
437
438 fn apply_vorticity_confinement(&mut self, dt: Fix128) {
447 let n = self.particle_count();
448 let h = self.config.kernel_radius;
449 let h_sq = h * h;
450 let vol = self.particle_volume();
451 let epsilon = self.config.vorticity_strength * dt;
452 let mut neighbors_buf = Vec::new();
453
454 let mut curls: Vec<Vec3Fix> = vec![Vec3Fix::ZERO; n];
456 for (i, curl_out) in curls.iter_mut().enumerate() {
457 self.grid
458 .query_neighbors_into(self.positions[i], h_sq, &mut neighbors_buf);
459 let mut curl = Vec3Fix::ZERO;
460 for &j in &neighbors_buf {
461 if i == j {
462 continue;
463 }
464 let delta = self.positions[j] - self.positions[i];
465 let r_sq = delta.length_squared();
466 if r_sq.is_zero() || r_sq >= h_sq {
467 continue;
468 }
469 let r = r_sq.sqrt();
470 let grad_mag = spiky_grad(r, h) * vol;
471 let grad = delta / r * grad_mag;
472 let vel_diff = self.velocities[j] - self.velocities[i];
473 curl = curl + vel_diff.cross(grad);
474 }
475 *curl_out = curl;
476 }
477
478 for i in 0..n {
481 self.grid
482 .query_neighbors_into(self.positions[i], h_sq, &mut neighbors_buf);
483 let mut grad_mag_curl = Vec3Fix::ZERO;
484 for &j in &neighbors_buf {
485 if i == j {
486 continue;
487 }
488 let delta = self.positions[j] - self.positions[i];
489 let r_sq = delta.length_squared();
490 if r_sq.is_zero() || r_sq >= h_sq {
491 continue;
492 }
493 let r = r_sq.sqrt();
494 let grad_w = spiky_grad(r, h) * vol;
495 grad_mag_curl = grad_mag_curl + delta / r * (grad_w * curls[j].length());
497 }
498 if grad_mag_curl.length_squared().is_zero() {
499 continue;
500 }
501 let n_vec = grad_mag_curl.normalize();
502 let force = n_vec.cross(curls[i]) * epsilon;
503 self.velocities[i] = self.velocities[i] + force;
504 }
505 }
506
507 #[cfg(feature = "std")]
509 fn resolve_sdf_boundary(&mut self, sdf_colliders: &[SdfCollider]) {
510 for i in 0..self.particle_count() {
511 for sdf in sdf_colliders {
512 let (lx, ly, lz) = sdf.world_to_local(self.positions[i]);
513 let dist = sdf.field.distance(lx, ly, lz) * sdf.scale_f32;
514
515 if dist > 0.0 {
517 let (nx, ny, nz) = sdf.field.normal(lx, ly, lz);
518 let normal = sdf.local_normal_to_world(nx, ny, nz);
519 let push = Fix128::from_f32(-dist);
520 self.positions[i] = self.positions[i] + normal * push;
521 self.predicted[i] = self.positions[i];
522
523 let vn = normal * self.velocities[i].dot(normal);
525 self.velocities[i] = self.velocities[i] - vn * Fix128::from_ratio(15, 10);
526 }
527 }
528 }
529 }
530}
531
532impl core::fmt::Debug for Fluid {
533 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
534 f.debug_struct("Fluid")
535 .field(
536 "positions",
537 &format_args!("[{} items]", self.positions.len()),
538 )
539 .field(
540 "predicted",
541 &format_args!("[{} items]", self.predicted.len()),
542 )
543 .field(
544 "velocities",
545 &format_args!("[{} items]", self.velocities.len()),
546 )
547 .field(
548 "densities",
549 &format_args!("[{} items]", self.densities.len()),
550 )
551 .field("lambdas", &format_args!("[{} items]", self.lambdas.len()))
552 .field("grid", &"<SpatialGrid>")
553 .field("config", &self.config)
554 .field("inv_rest_density", &self.inv_rest_density)
555 .finish()
556 }
557}
558
559#[cfg(test)]
564mod tests {
565 use super::*;
566
567 #[test]
568 fn test_fluid_creation() {
569 let config = FluidConfig::default();
570 let fluid = Fluid::new_block(
571 Vec3Fix::from_f32(-0.5, 0.0, -0.5),
572 Vec3Fix::from_f32(0.5, 1.0, 0.5),
573 Fix128::from_ratio(1, 5),
574 config,
575 );
576
577 assert!(fluid.particle_count() > 0, "Should create particles");
578 }
579
580 #[test]
581 fn test_fluid_gravity() {
582 let config = FluidConfig {
583 iterations: 2,
584 substeps: 1,
585 ..Default::default()
586 };
587 let mut fluid = Fluid::new(vec![Vec3Fix::from_f32(0.0, 5.0, 0.0)], config);
588
589 let dt = Fix128::from_ratio(1, 60);
590 for _ in 0..60 {
591 fluid.step(dt);
592 }
593
594 let y = fluid.positions[0].y.to_f32();
595 assert!(y < 5.0, "Particle should fall under gravity");
596 }
597
598 fn ball_integral(h: Fix128, f: impl Fn(Fix128) -> Fix128) -> f64 {
600 let n = 4000i64;
601 let dr = h / Fix128::from_int(n);
602 let mut sum = 0.0;
603 for i in 0..n {
604 let r = dr * (Fix128::from_int(i) + Fix128::from_ratio(1, 2));
605 sum += (f(r) * Fix128::from_int(4) * Fix128::PI * r * r * dr).to_f64();
606 }
607 sum
608 }
609
610 #[test]
611 fn poly6_integrates_to_one_over_its_support() {
612 for h in [
613 Fix128::from_ratio(1, 5),
614 Fix128::from_ratio(3, 2),
615 Fix128::ONE,
616 ] {
617 let v = ball_integral(h, |r| poly6(r * r, h));
618 assert!((v - 1.0).abs() < 1e-4, "h = {}: {v}", h.to_f64());
619 }
620 }
621
622 #[test]
623 fn poly6_has_compact_support() {
624 let h = Fix128::from_ratio(1, 5);
625 for k in [10, 11, 15, 30] {
626 let r = h * Fix128::from_ratio(k, 10);
627 assert_eq!(poly6(r * r, h), Fix128::ZERO, "r = {k}/10 h");
628 }
629 }
630
631 #[test]
632 fn spiky_gradient_matches_its_closed_form() {
633 let h = Fix128::from_ratio(1, 5);
634 let hf = 0.2f64;
635 for k in 1..10 {
636 let r = h * Fix128::from_ratio(k, 10);
637 let d = hf - r.to_f64();
638 let h6 = hf * hf * hf * hf * hf * hf;
639 let want = -45.0 / (core::f64::consts::PI * h6) * d * d;
640 let got = spiky_grad(r, h).to_f64();
641 assert!(got < 0.0, "gradient points inward");
642 assert!(
643 ((got - want) / want).abs() < 1e-9,
644 "k = {k}: {got} vs {want}"
645 );
646 }
647 assert_eq!(spiky_grad(h, h), Fix128::ZERO);
648 assert_eq!(spiky_grad(h * Fix128::from_ratio(3, 2), h), Fix128::ZERO);
649 assert_eq!(spiky_grad(Fix128::ZERO, h), Fix128::ZERO);
650 let v = ball_integral(h, |r| -spiky_grad(r, h));
652 assert!((v - 6.0 / hf).abs() / (6.0 / hf) < 1e-3, "{v}");
653 }
654
655 #[test]
656 fn test_poly6_kernel() {
657 let h = Fix128::from_ratio(1, 5);
658 let zero = poly6(Fix128::ZERO, h);
659 assert!(zero > Fix128::ZERO, "Poly6 at r=0 should be positive");
660
661 let far = poly6(h * h, h);
662 assert!(
663 far.is_zero() || far <= Fix128::ZERO,
664 "Poly6 at r=h should be zero"
665 );
666 }
667
668 #[test]
669 fn test_spatial_grid() {
670 let mut grid = SpatialGrid::new(Fix128::from_ratio(1, 5), 32);
671 grid.insert(0, Vec3Fix::ZERO);
672 grid.insert(1, Vec3Fix::from_f32(0.1, 0.0, 0.0));
673 grid.insert(2, Vec3Fix::from_f32(10.0, 0.0, 0.0));
674 grid.build();
675
676 let h_sq = Fix128::from_ratio(1, 5) * Fix128::from_ratio(1, 5);
677 let mut neighbors = Vec::new();
678 grid.query_neighbors_into(Vec3Fix::ZERO, h_sq, &mut neighbors);
679 assert!(neighbors.contains(&0), "Should find self");
680 assert!(neighbors.contains(&1), "Should find nearby particle");
681 }
682
683 #[cfg(feature = "std")]
685 fn unit_sphere_collider() -> crate::sdf_collider::SdfCollider {
686 use crate::sdf_collider::{ClosureSdf, SdfCollider};
687 let field = ClosureSdf::new(
688 |x, y, z| (x * x + y * y + z * z).sqrt() - 1.0,
689 |x, y, z| {
690 let len = (x * x + y * y + z * z).sqrt();
691 if len > 1e-6 {
692 (x / len, y / len, z / len)
693 } else {
694 (0.0, 1.0, 0.0)
695 }
696 },
697 );
698 SdfCollider::new_static(
699 Box::new(field),
700 Vec3Fix::ZERO,
701 crate::math::QuatFix::IDENTITY,
702 )
703 }
704
705 #[cfg(feature = "std")]
707 fn sphere_dist(p: Vec3Fix) -> f32 {
708 let (x, y, z) = p.to_f32();
709 (x * x + y * y + z * z).sqrt() - 1.0
710 }
711
712 #[cfg(feature = "std")]
713 #[test]
714 fn step_with_sdf_contains_fluid_particles_inside_unit_sphere() {
715 let make = || {
718 Fluid::new(
719 vec![
720 Vec3Fix::from_f32(0.0, 0.5, 0.0),
721 Vec3Fix::from_f32(0.3, 0.2, 0.0),
722 Vec3Fix::from_f32(-0.3, 0.1, 0.3),
723 ],
724 FluidConfig::default(),
725 )
726 };
727 let sphere = [unit_sphere_collider()];
728 let mut fluid = make();
729 let dt = Fix128::from_ratio(1, 60);
730 let mut max_dist = f32::MIN;
731 for frame in 0..120 {
732 fluid.step_with_sdf(dt, &sphere);
733 for (i, p) in fluid.positions.iter().enumerate() {
734 let d = sphere_dist(*p);
735 max_dist = max_dist.max(d);
736 assert!(
737 d <= 1e-3,
738 "frame {frame} particle {i} escaped sphere: dist {d}"
739 );
740 }
741 }
742 assert!(
744 max_dist > -0.05,
745 "never reached the wall: max dist {max_dist}"
746 );
747 for (i, p) in fluid.positions.iter().enumerate() {
748 let (_, y, _) = p.to_f32();
749 assert!(y >= -1.0 - 1e-3, "particle {i} below sphere bottom: y {y}");
750 }
751
752 let mut free = make();
754 for _ in 0..120 {
755 free.step(dt);
756 }
757 for (i, p) in free.positions.iter().enumerate() {
758 let d = sphere_dist(*p);
759 assert!(d > 1.0, "particle {i} should have fallen out: dist {d}");
760 }
761
762 let mut a = make();
764 let mut b = make();
765 for _ in 0..10 {
766 a.step_with_sdf(dt, &[]);
767 b.step(dt);
768 }
769 assert_eq!(a.positions, b.positions);
770 assert_eq!(a.velocities, b.velocities);
771 }
772}