1#[derive(Debug, Clone, Copy)]
17pub struct MacCell {
18 pub pressure: f64,
20 pub u_x: f64,
22 pub u_y: f64,
24 pub is_fluid: bool,
26}
27
28impl Default for MacCell {
29 fn default() -> Self {
30 Self {
31 pressure: 0.0,
32 u_x: 0.0,
33 u_y: 0.0,
34 is_fluid: true,
35 }
36 }
37}
38
39#[derive(Debug, Clone)]
44pub struct MacGrid {
45 pub nx: usize,
47 pub ny: usize,
49 pub dx: f64,
51 pub cells: Vec<MacCell>,
53}
54
55impl MacGrid {
56 pub fn new(nx: usize, ny: usize, dx: f64) -> Self {
58 Self {
59 nx,
60 ny,
61 dx,
62 cells: vec![MacCell::default(); nx * ny],
63 }
64 }
65
66 #[inline]
68 pub fn cell(&self, i: usize, j: usize) -> &MacCell {
69 &self.cells[j * self.nx + i]
70 }
71
72 #[inline]
74 pub fn cell_mut(&mut self, i: usize, j: usize) -> &mut MacCell {
75 &mut self.cells[j * self.nx + i]
76 }
77
78 pub fn divergence(&self, i: usize, j: usize) -> f64 {
84 let ux_e = self.cells[j * self.nx + i].u_x;
85 let ux_w = if i > 0 {
86 self.cells[j * self.nx + (i - 1)].u_x
87 } else {
88 0.0
89 };
90 let uy_n = self.cells[j * self.nx + i].u_y;
91 let uy_s = if j > 0 {
92 self.cells[(j - 1) * self.nx + i].u_y
93 } else {
94 0.0
95 };
96 (ux_e - ux_w + uy_n - uy_s) / self.dx
97 }
98
99 pub fn pressure_solve_jacobi(&mut self, n_iter: usize) {
103 let nx = self.nx;
104 let ny = self.ny;
105 let dx2 = self.dx * self.dx;
106 for _ in 0..n_iter {
107 let old = self.cells.clone();
108 for j in 1..ny.saturating_sub(1) {
109 for i in 1..nx.saturating_sub(1) {
110 if !old[j * nx + i].is_fluid {
111 continue;
112 }
113 let rhs = self.divergence(i, j);
114 let p_e = old[j * nx + i + 1].pressure;
115 let p_w = old[j * nx + i - 1].pressure;
116 let p_n = old[(j + 1) * nx + i].pressure;
117 let p_s = old[(j - 1) * nx + i].pressure;
118 self.cells[j * nx + i].pressure = (p_e + p_w + p_n + p_s - dx2 * rhs) / 4.0;
119 }
120 }
121 }
122 }
123
124 pub fn advect_velocity(&mut self, dt: f64) {
129 let nx = self.nx;
130 let ny = self.ny;
131 let dx = self.dx;
132 let old = self.cells.clone();
133
134 for j in 0..ny {
135 for i in 0..nx {
136 if !old[j * nx + i].is_fluid {
137 continue;
138 }
139 let x = (i as f64 + 1.0) * dx;
141 let y = (j as f64 + 0.5) * dx;
142 let [vx, vy] = interpolate_velocity_cells(&old, nx, ny, dx, x, y);
143 let xp = x - dt * vx;
144 let yp = y - dt * vy;
145 let [new_ux, _] = interpolate_velocity_cells(&old, nx, ny, dx, xp, yp);
146
147 let x2 = (i as f64 + 0.5) * dx;
149 let y2 = (j as f64 + 1.0) * dx;
150 let [vx2, vy2] = interpolate_velocity_cells(&old, nx, ny, dx, x2, y2);
151 let xp2 = x2 - dt * vx2;
152 let yp2 = y2 - dt * vy2;
153 let [_, new_uy] = interpolate_velocity_cells(&old, nx, ny, dx, xp2, yp2);
154
155 let idx = j * nx + i;
156 self.cells[idx].u_x = new_ux;
157 self.cells[idx].u_y = new_uy;
158 }
159 }
160 }
161
162 pub fn project(&mut self, dt: f64) {
166 let nx = self.nx;
167 let ny = self.ny;
168 let dx = self.dx;
169 let p: Vec<f64> = self.cells.iter().map(|c| c.pressure).collect();
170 for j in 0..ny {
171 for i in 0..nx {
172 if !self.cells[j * nx + i].is_fluid {
173 continue;
174 }
175 let p_e = if i + 1 < nx { p[j * nx + i + 1] } else { 0.0 };
177 let p_c = p[j * nx + i];
178 self.cells[j * nx + i].u_x -= dt / dx * (p_e - p_c);
179 let p_n = if j + 1 < ny { p[(j + 1) * nx + i] } else { 0.0 };
181 self.cells[j * nx + i].u_y -= dt / dx * (p_n - p_c);
182 }
183 }
184 }
185}
186
187fn interpolate_velocity_cells(
191 cells: &[MacCell],
192 nx: usize,
193 ny: usize,
194 dx: f64,
195 x: f64,
196 y: f64,
197) -> [f64; 2] {
198 let cx = (x / dx - 0.5).max(0.0).min((nx - 1) as f64 - 1e-9);
199 let cy = (y / dx - 0.5).max(0.0).min((ny - 1) as f64 - 1e-9);
200 let i = cx as usize;
201 let j = cy as usize;
202 let fi = cx - i as f64;
203 let fj = cy - j as f64;
204
205 let safe = |ii: usize, jj: usize| -> [f64; 2] {
206 if ii < nx && jj < ny {
207 let c = &cells[jj * nx + ii];
208 [c.u_x, c.u_y]
209 } else {
210 [0.0, 0.0]
211 }
212 };
213
214 let c00 = safe(i, j);
215 let c10 = safe(i + 1, j);
216 let c01 = safe(i, j + 1);
217 let c11 = safe(i + 1, j + 1);
218
219 let ux = (1.0 - fi) * (1.0 - fj) * c00[0]
220 + fi * (1.0 - fj) * c10[0]
221 + (1.0 - fi) * fj * c01[0]
222 + fi * fj * c11[0];
223 let uy = (1.0 - fi) * (1.0 - fj) * c00[1]
224 + fi * (1.0 - fj) * c10[1]
225 + (1.0 - fi) * fj * c01[1]
226 + fi * fj * c11[1];
227 [ux, uy]
228}
229
230pub struct FluidSimGpu {
237 pub grid: MacGrid,
239 pub viscosity: f64,
241 pub density: f64,
243}
244
245impl FluidSimGpu {
246 pub fn new(nx: usize, ny: usize, dx: f64, viscosity: f64, density: f64) -> Self {
248 Self {
249 grid: MacGrid::new(nx, ny, dx),
250 viscosity,
251 density,
252 }
253 }
254
255 pub fn step(&mut self, dt: f64) {
259 self.grid.advect_velocity(dt);
260 self.grid.pressure_solve_jacobi(20);
261 self.grid.project(dt);
262 }
263}
264
265pub struct VorticityConfinement {
272 pub strength: f64,
274}
275
276impl VorticityConfinement {
277 pub fn new(strength: f64) -> Self {
279 Self { strength }
280 }
281
282 pub fn compute_vorticity(&self, grid: &MacGrid) -> Vec<f64> {
286 let nx = grid.nx;
287 let ny = grid.ny;
288 let dx = grid.dx;
289 let mut omega = vec![0.0f64; nx * ny];
290 for j in 1..ny.saturating_sub(1) {
291 for i in 1..nx.saturating_sub(1) {
292 let duy_dx = (grid.cells[j * nx + i].u_y - grid.cells[j * nx + (i - 1)].u_y) / dx;
293 let dux_dy = (grid.cells[j * nx + i].u_x - grid.cells[(j - 1) * nx + i].u_x) / dx;
294 omega[j * nx + i] = duy_dx - dux_dy;
295 }
296 }
297 omega
298 }
299
300 pub fn apply_confinement_force(&self, grid: &mut MacGrid, dt: f64) {
305 let omega = self.compute_vorticity(grid);
306 let nx = grid.nx;
307 let ny = grid.ny;
308 let dx = grid.dx;
309 let omega_abs: Vec<f64> = omega.iter().map(|v| v.abs()).collect();
310 for j in 1..ny.saturating_sub(1) {
311 for i in 1..nx.saturating_sub(1) {
312 let grad_x = (omega_abs[j * nx + i + 1] - omega_abs[j * nx + i - 1]) / (2.0 * dx);
313 let grad_y =
314 (omega_abs[(j + 1) * nx + i] - omega_abs[(j - 1) * nx + i]) / (2.0 * dx);
315 let len = (grad_x * grad_x + grad_y * grad_y).sqrt();
316 if len < 1e-12 {
317 continue;
318 }
319 let nx_hat = grad_x / len;
320 let ny_hat = grad_y / len;
321 let w = omega[j * nx + i];
322 let fx = self.strength * dx * w * ny_hat;
324 let fy = -self.strength * dx * w * nx_hat;
325 grid.cells[j * nx + i].u_x += dt * fx;
326 grid.cells[j * nx + i].u_y += dt * fy;
327 }
328 }
329 }
330}
331
332pub struct FluidObstacle {
339 pub cells: Vec<(usize, usize)>,
341}
342
343impl FluidObstacle {
344 pub fn new(cells: Vec<(usize, usize)>) -> Self {
346 Self { cells }
347 }
348
349 pub fn set_solid(&self, grid: &mut MacGrid) {
351 for &(i, j) in &self.cells {
352 if i < grid.nx && j < grid.ny {
353 grid.cell_mut(i, j).is_fluid = false;
354 }
355 }
356 }
357
358 pub fn apply_no_slip(&self, grid: &mut MacGrid) {
360 for &(i, j) in &self.cells {
361 if i < grid.nx && j < grid.ny {
362 let c = grid.cell_mut(i, j);
363 c.u_x = 0.0;
364 c.u_y = 0.0;
365 c.pressure = 0.0;
366 }
367 }
368 }
369}
370
371pub fn cfl_timestep(grid: &MacGrid, safety: f64) -> f64 {
378 let mut u_max: f64 = 0.0;
379 for c in &grid.cells {
380 u_max = u_max.max(c.u_x.abs()).max(c.u_y.abs());
381 }
382 if u_max < 1e-15 {
383 return safety * grid.dx;
384 }
385 safety * grid.dx / u_max
386}
387
388pub fn reynolds_number_grid(u_max: f64, l: f64, nu: f64) -> f64 {
394 if nu.abs() < 1e-30 {
395 return f64::INFINITY;
396 }
397 u_max * l / nu
398}
399
400pub fn interpolate_velocity(grid: &MacGrid, x: f64, y: f64) -> [f64; 2] {
404 interpolate_velocity_cells(&grid.cells, grid.nx, grid.ny, grid.dx, x, y)
405}
406
407#[cfg(test)]
410mod tests {
411 use super::*;
412
413 fn uniform_grid(nx: usize, ny: usize, ux: f64, uy: f64) -> MacGrid {
414 let mut g = MacGrid::new(nx, ny, 1.0);
415 for c in g.cells.iter_mut() {
416 c.u_x = ux;
417 c.u_y = uy;
418 }
419 g
420 }
421
422 #[test]
425 fn mac_cell_default_is_fluid() {
426 let c = MacCell::default();
427 assert!(c.is_fluid);
428 assert_eq!(c.pressure, 0.0);
429 assert_eq!(c.u_x, 0.0);
430 assert_eq!(c.u_y, 0.0);
431 }
432
433 #[test]
434 fn mac_cell_copy() {
435 let c = MacCell {
436 pressure: 1.0,
437 u_x: 2.0,
438 u_y: 3.0,
439 is_fluid: false,
440 };
441 let d = c;
442 assert_eq!(d.pressure, 1.0);
443 assert!(!d.is_fluid);
444 }
445
446 #[test]
449 fn mac_grid_new_dimensions() {
450 let g = MacGrid::new(4, 5, 0.1);
451 assert_eq!(g.nx, 4);
452 assert_eq!(g.ny, 5);
453 assert_eq!(g.cells.len(), 20);
454 }
455
456 #[test]
457 fn mac_grid_cell_accessor() {
458 let mut g = MacGrid::new(3, 3, 1.0);
459 g.cell_mut(1, 2).u_x = 42.0;
460 assert_eq!(g.cell(1, 2).u_x, 42.0);
461 }
462
463 #[test]
466 fn divergence_uniform_flow_is_zero() {
467 let g = uniform_grid(5, 5, 1.0, 0.0);
469 let div = g.divergence(2, 2);
471 assert!(div.abs() < 1e-12, "div = {div}");
472 }
473
474 #[test]
475 fn divergence_source_cell() {
476 let mut g = MacGrid::new(4, 4, 1.0);
478 g.cell_mut(1, 1).u_x = 2.0; g.cell_mut(0, 1).u_x = 1.0; let div = g.divergence(1, 1);
481 assert!(div > 0.0, "expected positive div, got {div}");
482 }
483
484 #[test]
485 fn divergence_boundary_left_edge() {
486 let mut g = MacGrid::new(4, 4, 1.0);
487 g.cell_mut(0, 0).u_x = 1.0;
488 let div = g.divergence(0, 0);
490 assert!((div - 1.0).abs() < 1e-12, "div = {div}");
491 }
492
493 #[test]
494 fn divergence_all_zeros() {
495 let g = MacGrid::new(4, 4, 1.0);
496 for j in 0..4 {
497 for i in 0..4 {
498 assert_eq!(g.divergence(i, j), 0.0);
499 }
500 }
501 }
502
503 #[test]
506 fn pressure_jacobi_reduces_magnitude_with_source() {
507 let mut g = MacGrid::new(6, 6, 1.0);
508 g.cell_mut(3, 3).u_x = 1.0;
510 g.pressure_solve_jacobi(50);
511 let max_p: f64 = g
513 .cells
514 .iter()
515 .map(|c| c.pressure.abs())
516 .fold(0.0_f64, f64::max);
517 assert!(max_p > 1e-6, "pressure unchanged after Jacobi");
518 }
519
520 #[test]
521 fn pressure_jacobi_zero_field_stays_zero() {
522 let mut g = MacGrid::new(5, 5, 1.0);
523 g.pressure_solve_jacobi(10);
524 for c in &g.cells {
525 assert_eq!(c.pressure, 0.0);
526 }
527 }
528
529 #[test]
530 fn pressure_jacobi_convergence_monotone() {
531 let mut g1 = MacGrid::new(6, 6, 1.0);
533 let mut g2 = MacGrid::new(6, 6, 1.0);
534 g1.cell_mut(3, 3).u_x = 2.0;
535 g2.cell_mut(3, 3).u_x = 2.0;
536 g1.pressure_solve_jacobi(10);
537 g2.pressure_solve_jacobi(100);
538 let rms = |grid: &MacGrid| {
539 let s: f64 = grid.cells.iter().map(|c| c.pressure * c.pressure).sum();
540 (s / grid.cells.len() as f64).sqrt()
541 };
542 let _ = rms(&g1);
544 let _ = rms(&g2);
545 }
547
548 #[test]
551 fn project_corrects_velocity() {
552 let mut g = MacGrid::new(4, 4, 1.0);
553 for j in 0..4 {
555 for i in 0..4 {
556 g.cell_mut(i, j).pressure = i as f64;
557 }
558 }
559 let ux_before = g.cell(1, 1).u_x;
560 g.project(0.1);
561 let ux_after = g.cell(1, 1).u_x;
562 assert!((ux_after - ux_before).abs() > 1e-10);
564 }
565
566 #[test]
567 fn project_solid_cell_unchanged() {
568 let mut g = MacGrid::new(4, 4, 1.0);
569 g.cell_mut(2, 2).is_fluid = false;
570 g.cell_mut(2, 2).u_x = 5.0;
571 for j in 0..4 {
572 for i in 0..4 {
573 g.cell_mut(i, j).pressure = 1.0;
574 }
575 }
576 g.project(0.1);
577 assert_eq!(g.cell(2, 2).u_x, 5.0);
579 }
580
581 #[test]
584 fn cfl_timestep_uniform_flow() {
585 let g = uniform_grid(4, 4, 2.0, 0.0);
586 let dt = cfl_timestep(&g, 0.5);
587 assert!((dt - 0.25).abs() < 1e-12, "dt = {dt}");
589 }
590
591 #[test]
592 fn cfl_timestep_quiescent() {
593 let g = MacGrid::new(4, 4, 0.1);
594 let dt = cfl_timestep(&g, 0.5);
595 assert!((dt - 0.05).abs() < 1e-12, "dt = {dt}");
597 }
598
599 #[test]
600 fn cfl_timestep_uses_max_component() {
601 let mut g = MacGrid::new(4, 4, 1.0);
602 g.cell_mut(1, 1).u_x = 1.0;
603 g.cell_mut(2, 2).u_y = 5.0;
604 let dt = cfl_timestep(&g, 1.0);
605 assert!((dt - 1.0 / 5.0).abs() < 1e-12, "dt = {dt}");
606 }
607
608 #[test]
609 fn cfl_timestep_safety_factor() {
610 let g = uniform_grid(4, 4, 4.0, 0.0);
611 let dt_full = cfl_timestep(&g, 1.0);
612 let dt_half = cfl_timestep(&g, 0.5);
613 assert!((dt_full - 2.0 * dt_half).abs() < 1e-12);
614 }
615
616 #[test]
619 fn reynolds_number_basic() {
620 let re = reynolds_number_grid(1.0, 1.0, 1e-6);
621 assert!((re - 1.0e6).abs() < 1.0, "re = {re}");
622 }
623
624 #[test]
625 fn reynolds_number_zero_viscosity() {
626 let re = reynolds_number_grid(1.0, 1.0, 0.0);
627 assert!(re.is_infinite());
628 }
629
630 #[test]
631 fn reynolds_number_proportional_to_velocity() {
632 let re1 = reynolds_number_grid(1.0, 1.0, 1e-3);
633 let re2 = reynolds_number_grid(2.0, 1.0, 1e-3);
634 assert!((re2 - 2.0 * re1).abs() < 1e-6);
635 }
636
637 #[test]
640 fn interpolate_velocity_at_cell_centre() {
641 let mut g = MacGrid::new(4, 4, 1.0);
642 for c in g.cells.iter_mut() {
643 c.u_x = 3.0;
644 c.u_y = 1.5;
645 }
646 let v = interpolate_velocity(&g, 1.5, 1.5);
647 assert!((v[0] - 3.0).abs() < 0.1, "v[0] = {}", v[0]);
648 }
649
650 #[test]
651 fn interpolate_velocity_clamps_outside() {
652 let g = uniform_grid(4, 4, 1.0, 2.0);
653 let v = interpolate_velocity(&g, -10.0, -10.0);
655 let _ = v;
656 }
657
658 #[test]
659 fn interpolate_velocity_zero_field() {
660 let g = MacGrid::new(4, 4, 1.0);
661 let v = interpolate_velocity(&g, 1.0, 1.0);
662 assert_eq!(v[0], 0.0);
663 assert_eq!(v[1], 0.0);
664 }
665
666 #[test]
667 fn interpolate_velocity_corner() {
668 let g = uniform_grid(4, 4, 1.0, 2.0);
669 let v = interpolate_velocity(&g, 0.0, 0.0);
670 let _ = v; }
672
673 #[test]
676 fn vorticity_zero_flow() {
677 let g = MacGrid::new(5, 5, 1.0);
678 let vc = VorticityConfinement::new(0.1);
679 let omega = vc.compute_vorticity(&g);
680 assert!(omega.iter().all(|&v| v == 0.0));
681 }
682
683 #[test]
684 fn vorticity_shear_flow() {
685 let mut g = MacGrid::new(5, 5, 1.0);
686 for j in 0..5 {
688 for i in 0..5 {
689 g.cell_mut(i, j).u_x = j as f64;
690 }
691 }
692 let vc = VorticityConfinement::new(0.1);
693 let omega = vc.compute_vorticity(&g);
694 let has_nonzero = omega.iter().any(|&v| v.abs() > 1e-12);
696 assert!(has_nonzero);
697 }
698
699 #[test]
700 fn vorticity_confinement_apply_no_panic() {
701 let mut g = MacGrid::new(5, 5, 1.0);
702 for c in g.cells.iter_mut() {
703 c.u_x = 0.1;
704 c.u_y = -0.1;
705 }
706 let vc = VorticityConfinement::new(0.2);
707 vc.apply_confinement_force(&mut g, 0.01);
708 }
710
711 #[test]
712 fn vorticity_len_equals_grid_size() {
713 let g = MacGrid::new(6, 7, 0.5);
714 let vc = VorticityConfinement::new(0.1);
715 let omega = vc.compute_vorticity(&g);
716 assert_eq!(omega.len(), 6 * 7);
717 }
718
719 #[test]
722 fn obstacle_set_solid() {
723 let mut g = MacGrid::new(5, 5, 1.0);
724 let obs = FluidObstacle::new(vec![(2, 2), (3, 2)]);
725 obs.set_solid(&mut g);
726 assert!(!g.cell(2, 2).is_fluid);
727 assert!(!g.cell(3, 2).is_fluid);
728 assert!(g.cell(1, 1).is_fluid);
729 }
730
731 #[test]
732 fn obstacle_apply_no_slip() {
733 let mut g = MacGrid::new(5, 5, 1.0);
734 g.cell_mut(2, 2).u_x = 10.0;
735 g.cell_mut(2, 2).u_y = 5.0;
736 let obs = FluidObstacle::new(vec![(2, 2)]);
737 obs.apply_no_slip(&mut g);
738 assert_eq!(g.cell(2, 2).u_x, 0.0);
739 assert_eq!(g.cell(2, 2).u_y, 0.0);
740 }
741
742 #[test]
743 fn obstacle_out_of_bounds_no_panic() {
744 let mut g = MacGrid::new(4, 4, 1.0);
745 let obs = FluidObstacle::new(vec![(100, 100)]);
746 obs.set_solid(&mut g); obs.apply_no_slip(&mut g);
748 }
749
750 #[test]
751 fn obstacle_multiple_cells() {
752 let mut g = MacGrid::new(10, 10, 1.0);
753 let cells: Vec<_> = (2..5).flat_map(|i| (2..5).map(move |j| (i, j))).collect();
754 let obs = FluidObstacle::new(cells);
755 obs.set_solid(&mut g);
756 assert!(!g.cell(3, 3).is_fluid);
757 assert!(g.cell(0, 0).is_fluid);
758 }
759
760 #[test]
763 fn fluid_sim_gpu_step_no_panic() {
764 let mut sim = FluidSimGpu::new(6, 6, 0.1, 1e-3, 1000.0);
765 sim.grid.cell_mut(3, 3).u_x = 0.5;
766 sim.step(0.01);
767 }
769
770 #[test]
771 fn fluid_sim_gpu_quiescent_stable() {
772 let mut sim = FluidSimGpu::new(8, 8, 0.1, 1e-3, 1000.0);
773 for _ in 0..5 {
774 sim.step(0.001);
775 }
776 for c in &sim.grid.cells {
778 assert!(c.u_x.abs() < 1e-10, "u_x = {}", c.u_x);
779 assert!(c.u_y.abs() < 1e-10, "u_y = {}", c.u_y);
780 }
781 }
782
783 #[test]
784 fn fluid_sim_gpu_fields_accessible() {
785 let sim = FluidSimGpu::new(4, 4, 0.05, 0.001, 998.0);
786 assert_eq!(sim.grid.nx, 4);
787 assert!((sim.viscosity - 0.001).abs() < 1e-12);
788 assert!((sim.density - 998.0).abs() < 1e-12);
789 }
790
791 #[test]
794 fn advect_zero_dt_no_change() {
795 let mut g = uniform_grid(5, 5, 1.0, 0.5);
796 let before: Vec<_> = g.cells.iter().map(|c| [c.u_x, c.u_y]).collect();
797 g.advect_velocity(0.0);
798 let after: Vec<_> = g.cells.iter().map(|c| [c.u_x, c.u_y]).collect();
799 for (b, a) in before.iter().zip(after.iter()) {
800 assert!((b[0] - a[0]).abs() < 1e-10);
801 assert!((b[1] - a[1]).abs() < 1e-10);
802 }
803 }
804
805 #[test]
806 fn advect_does_not_explode() {
807 let mut g = uniform_grid(6, 6, 0.3, 0.2);
808 for _ in 0..10 {
809 g.advect_velocity(0.01);
810 }
811 for c in &g.cells {
812 assert!(c.u_x.is_finite());
813 assert!(c.u_y.is_finite());
814 }
815 }
816}