Skip to main content

oxiphysics_gpu/fluid_sim_gpu/
functions.rs

1//! Auto-generated module
2//!
3//! 🤖 Generated with [SplitRS](https://github.com/cool-japan/splitrs)
4
5use super::types::{
6    FlipParticle, GpuBoundaryBox, LbmCellType, LbmD2Q9, MacGrid, SphConfig, SphKernels, SphParticle,
7};
8
9pub(super) fn dot3(a: [f64; 3], b: [f64; 3]) -> f64 {
10    a[0] * b[0] + a[1] * b[1] + a[2] * b[2]
11}
12pub(super) fn sub3(a: [f64; 3], b: [f64; 3]) -> [f64; 3] {
13    [a[0] - b[0], a[1] - b[1], a[2] - b[2]]
14}
15pub(super) fn add3(a: [f64; 3], b: [f64; 3]) -> [f64; 3] {
16    [a[0] + b[0], a[1] + b[1], a[2] + b[2]]
17}
18pub(super) fn scale3(v: [f64; 3], s: f64) -> [f64; 3] {
19    [v[0] * s, v[1] * s, v[2] * s]
20}
21pub(super) fn length3(v: [f64; 3]) -> f64 {
22    dot3(v, v).sqrt()
23}
24pub(super) fn cross3(a: [f64; 3], b: [f64; 3]) -> [f64; 3] {
25    [
26        a[1] * b[2] - a[2] * b[1],
27        a[2] * b[0] - a[0] * b[2],
28        a[0] * b[1] - a[1] * b[0],
29    ]
30}
31/// Compute SPH density for all particles (mock GPU kernel dispatch).
32///
33/// For each particle i, density_i = Σ_j m_j * W_poly6(|r_i - r_j|, h)
34pub fn sph_compute_density(particles: &mut [SphParticle], config: &SphConfig) {
35    let n = particles.len();
36    let mut densities = vec![0.0f64; n];
37    for i in 0..n {
38        let mut rho = 0.0;
39        for j in 0..n {
40            let r_vec = sub3(particles[i].position, particles[j].position);
41            let r = length3(r_vec);
42            rho += particles[j].mass * SphKernels::poly6(r, config.h);
43        }
44        densities[i] = rho.max(1e-6);
45    }
46    for (i, p) in particles.iter_mut().enumerate() {
47        p.density = densities[i];
48    }
49}
50/// Compute SPH pressure from density (Tait equation).
51pub fn sph_compute_pressure(particles: &mut [SphParticle], config: &SphConfig) {
52    for p in particles.iter_mut() {
53        let ratio = p.density / config.rest_density;
54        p.pressure = config.pressure_k * (ratio.powi(7) - 1.0);
55    }
56}
57/// Compute SPH forces: pressure gradient + viscosity + gravity + surface tension.
58pub fn sph_compute_forces(particles: &mut [SphParticle], config: &SphConfig) {
59    let n = particles.len();
60    let mut forces = vec![[0.0f64; 3]; n];
61    for i in 0..n {
62        let mut f_pressure = [0.0f64; 3];
63        let mut f_viscosity = [0.0f64; 3];
64        for j in 0..n {
65            if i == j {
66                continue;
67            }
68            let r_vec = sub3(particles[i].position, particles[j].position);
69            let r = length3(r_vec);
70            if r > config.h || r < 1e-15 {
71                continue;
72            }
73            let pressure_factor = particles[i].pressure
74                / (particles[i].density * particles[i].density)
75                + particles[j].pressure / (particles[j].density * particles[j].density);
76            let grad = SphKernels::spiky_grad(r_vec, r, config.h);
77            f_pressure = add3(
78                f_pressure,
79                scale3(grad, -particles[j].mass * pressure_factor),
80            );
81            let v_diff = sub3(particles[j].velocity, particles[i].velocity);
82            let lap = SphKernels::viscosity_laplacian(r, config.h);
83            let visc_factor = config.viscosity * particles[j].mass * lap / particles[j].density;
84            f_viscosity = add3(f_viscosity, scale3(v_diff, visc_factor));
85        }
86        let f_gravity = scale3(config.gravity, particles[i].density);
87        forces[i] = add3(add3(f_pressure, f_viscosity), f_gravity);
88    }
89    for (i, p) in particles.iter_mut().enumerate() {
90        p.force = forces[i];
91    }
92}
93/// Integrate SPH particles using semi-implicit Euler.
94pub fn sph_integrate(particles: &mut [SphParticle], config: &SphConfig) {
95    for p in particles.iter_mut() {
96        let accel = scale3(p.force, 1.0 / p.density.max(1e-6));
97        p.velocity = add3(p.velocity, scale3(accel, config.dt));
98        p.position = add3(p.position, scale3(p.velocity, config.dt));
99    }
100}
101/// Full SPH step: density → pressure → forces → integrate.
102pub fn sph_step(particles: &mut [SphParticle], config: &SphConfig) {
103    sph_compute_density(particles, config);
104    sph_compute_pressure(particles, config);
105    sph_compute_forces(particles, config);
106    sph_integrate(particles, config);
107}
108/// LBM D3Q27 velocity set index count.
109pub const D3Q27_Q: usize = 27;
110/// LBM D2Q9 velocity set index count.
111pub const D2Q9_Q: usize = 9;
112/// D2Q9 x-component velocity vectors.
113pub const D2Q9_CX: [i32; 9] = [0, 1, 0, -1, 0, 1, -1, -1, 1];
114/// D2Q9 y-component velocity vectors.
115pub const D2Q9_CY: [i32; 9] = [0, 0, 1, 0, -1, 1, 1, -1, -1];
116/// D2Q9 equilibrium weights.
117pub const D2Q9_W: [f64; 9] = [
118    4.0 / 9.0,
119    1.0 / 9.0,
120    1.0 / 9.0,
121    1.0 / 9.0,
122    1.0 / 9.0,
123    1.0 / 36.0,
124    1.0 / 36.0,
125    1.0 / 36.0,
126    1.0 / 36.0,
127];
128/// D2Q9 opposite direction indices (for bounce-back).
129pub const D2Q9_OPP: [usize; 9] = [0, 3, 4, 1, 2, 7, 8, 5, 6];
130/// Semi-Lagrangian advection of a scalar field `phi` on a MAC grid.
131///
132/// Uses RK2 (midpoint method) for back-tracing.
133pub fn advect_scalar(phi: &[f64], grid: &MacGrid, dt: f64, gravity: [f64; 3]) -> Vec<f64> {
134    let nx = grid.nx;
135    let ny = grid.ny;
136    let nz = grid.nz;
137    let dx = grid.dx;
138    let mut phi_new = vec![0.0f64; nx * ny * nz];
139    for k in 0..nz {
140        for j in 0..ny {
141            for i in 0..nx {
142                let xc = (i as f64 + 0.5) * dx;
143                let yc = (j as f64 + 0.5) * dx;
144                let zc = (k as f64 + 0.5) * dx;
145                let u0 = grid.interp_u(xc, yc, zc);
146                let x_mid = xc - 0.5 * dt * (u0 + gravity[0]);
147                let y_mid = yc - 0.5 * dt * gravity[1];
148                let z_mid = zc - 0.5 * dt * gravity[2];
149                let u_mid = grid.interp_u(x_mid, y_mid, z_mid);
150                let x_back = xc - dt * (u_mid + gravity[0]);
151                let y_back = yc - dt * gravity[1];
152                let z_back = zc - dt * gravity[2];
153                let ix = (x_back / dx - 0.5).floor() as isize;
154                let iy = (y_back / dx - 0.5).floor() as isize;
155                let iz = (z_back / dx - 0.5).floor() as isize;
156                let ii = ix.clamp(0, nx as isize - 1) as usize;
157                let jj = iy.clamp(0, ny as isize - 1) as usize;
158                let kk = iz.clamp(0, nz as isize - 1) as usize;
159                phi_new[k * nx * ny + j * nx + i] = phi[kk * nx * ny + jj * nx + ii];
160            }
161        }
162    }
163    phi_new
164}
165/// Compute curl (vorticity) of the velocity field at cell centers.
166///
167/// Returns a flat buffer of 3-vectors (ωx, ωy, ωz).
168pub fn compute_vorticity(grid: &MacGrid) -> Vec<[f64; 3]> {
169    let nx = grid.nx;
170    let ny = grid.ny;
171    let nz = grid.nz;
172    let inv2dx = 1.0 / (2.0 * grid.dx);
173    let mut vorticity = vec![[0.0f64; 3]; nx * ny * nz];
174    for k in 1..nz - 1 {
175        for j in 1..ny - 1 {
176            for i in 1..nx - 1 {
177                let dwdy = (grid.get_w(i, j + 1, k) - grid.get_w(i, j - 1, k)) * inv2dx;
178                let dvdz = (grid.get_v(i, j, k + 1) - grid.get_v(i, j, k - 1)) * inv2dx;
179                let dudz = (grid.get_u(i, j, k + 1) - grid.get_u(i, j, k - 1)) * inv2dx;
180                let dwdx = (grid.get_w(i + 1, j, k) - grid.get_w(i - 1, j, k)) * inv2dx;
181                let dvdx = (grid.get_v(i + 1, j, k) - grid.get_v(i - 1, j, k)) * inv2dx;
182                let dudy = (grid.get_u(i, j + 1, k) - grid.get_u(i, j - 1, k)) * inv2dx;
183                let idx = k * nx * ny + j * nx + i;
184                vorticity[idx] = [dwdy - dvdz, dudz - dwdx, dvdx - dudy];
185            }
186        }
187    }
188    vorticity
189}
190/// Apply vorticity confinement forces to MAC grid velocities.
191///
192/// Vorticity confinement adds a force F = ε * (N × ω)
193/// where N is the normalized gradient of |ω|.
194pub fn vorticity_confinement(grid: &mut MacGrid, vorticity: &[[f64; 3]], epsilon: f64, dt: f64) {
195    let nx = grid.nx;
196    let ny = grid.ny;
197    let nz = grid.nz;
198    let inv2dx = 1.0 / (2.0 * grid.dx);
199    for k in 1..nz - 1 {
200        for j in 1..ny - 1 {
201            for i in 1..nx - 1 {
202                let mag = |x: usize, y: usize, z: usize| {
203                    let idx = z * nx * ny + y * nx + x;
204                    length3(vorticity[idx])
205                };
206                let gx = (mag(i + 1, j, k) - mag(i - 1, j, k)) * inv2dx;
207                let gy = (mag(i, j + 1, k) - mag(i, j - 1, k)) * inv2dx;
208                let gz = (mag(i, j, k + 1) - mag(i, j, k - 1)) * inv2dx;
209                let grad_len = (gx * gx + gy * gy + gz * gz).sqrt();
210                if grad_len < 1e-12 {
211                    continue;
212                }
213                let n_vec = [gx / grad_len, gy / grad_len, gz / grad_len];
214                let idx = k * nx * ny + j * nx + i;
215                let omega = vorticity[idx];
216                let force = cross3(n_vec, omega);
217                let force = scale3(force, epsilon);
218                if i < nx {
219                    let u = grid.get_u(i, j, k);
220                    grid.set_u(i, j, k, u + dt * force[0]);
221                }
222                if j < ny {
223                    let v = grid.get_v(i, j, k);
224                    grid.set_v(i, j, k, v + dt * force[1]);
225                }
226                if k < nz {
227                    let w = grid.get_w(i, j, k);
228                    grid.set_w(i, j, k, w + dt * force[2]);
229                }
230            }
231        }
232    }
233}
234/// Surface tension force via Continuum Surface Force (CSF) method.
235///
236/// Requires a level-set function `phi` (negative inside fluid, positive outside).
237/// Returns per-cell force vectors.
238pub fn surface_tension_csf(
239    phi: &[f64],
240    nx: usize,
241    ny: usize,
242    nz: usize,
243    dx: f64,
244    sigma: f64,
245) -> Vec<[f64; 3]> {
246    let inv_dx = 1.0 / dx;
247    let inv2dx = 0.5 * inv_dx;
248    let n = nx * ny * nz;
249    let mut forces = vec![[0.0f64; 3]; n];
250    let idx = |i: usize, j: usize, k: usize| k * nx * ny + j * nx + i;
251    for k in 1..nz - 1 {
252        for j in 1..ny - 1 {
253            for i in 1..nx - 1 {
254                let c = idx(i, j, k);
255                let gx = (phi[idx(i + 1, j, k)] - phi[idx(i - 1, j, k)]) * inv2dx;
256                let gy = (phi[idx(i, j + 1, k)] - phi[idx(i, j - 1, k)]) * inv2dx;
257                let gz = (phi[idx(i, j, k + 1)] - phi[idx(i, j, k - 1)]) * inv2dx;
258                let grad_mag = (gx * gx + gy * gy + gz * gz).sqrt();
259                if grad_mag < 1e-12 {
260                    continue;
261                }
262                let nx_n = gx / grad_mag;
263                let ny_n = gy / grad_mag;
264                let nz_n = gz / grad_mag;
265                let phi_xx = (phi[idx(i + 1, j, k)] - 2.0 * phi[c] + phi[idx(i - 1, j, k)])
266                    * inv_dx
267                    * inv_dx;
268                let phi_yy = (phi[idx(i, j + 1, k)] - 2.0 * phi[c] + phi[idx(i, j - 1, k)])
269                    * inv_dx
270                    * inv_dx;
271                let phi_zz = (phi[idx(i, j, k + 1)] - 2.0 * phi[c] + phi[idx(i, j, k - 1)])
272                    * inv_dx
273                    * inv_dx;
274                let kappa = -(phi_xx + phi_yy + phi_zz) / grad_mag;
275                let f_scale = sigma * kappa;
276                forces[c] = [f_scale * nx_n, f_scale * ny_n, f_scale * nz_n];
277            }
278        }
279    }
280    forces
281}
282/// Particle-to-Grid (P2G) transfer: splat particle velocities onto MAC grid.
283///
284/// Uses trilinear weighting.
285pub fn p2g_transfer(particles: &[FlipParticle], grid: &mut MacGrid) {
286    let dx = grid.dx;
287    let inv_dx = 1.0 / dx;
288    let nx = grid.nx;
289    let ny = grid.ny;
290    let nz = grid.nz;
291    let mut u_weight = vec![0.0f64; (nx + 1) * ny * nz];
292    let mut v_weight = vec![0.0f64; nx * (ny + 1) * nz];
293    let mut w_weight = vec![0.0f64; nx * ny * (nz + 1)];
294    let mut u_num = vec![0.0f64; (nx + 1) * ny * nz];
295    let mut v_num = vec![0.0f64; nx * (ny + 1) * nz];
296    let mut w_num = vec![0.0f64; nx * ny * (nz + 1)];
297    for p in particles {
298        let [px, py, pz] = p.position;
299        let iu = (px * inv_dx).floor() as isize;
300        let ju = (py * inv_dx - 0.5).floor() as isize;
301        let ku = (pz * inv_dx - 0.5).floor() as isize;
302        let fxu = px * inv_dx - iu as f64;
303        let fyu = py * inv_dx - 0.5 - ju as f64;
304        let fzu = pz * inv_dx - 0.5 - ku as f64;
305        for dz in 0..2 {
306            for dy in 0..2 {
307                for dx_off in 0..2 {
308                    let ni = (iu + dx_off as isize).clamp(0, nx as isize) as usize;
309                    let nj = (ju + dy as isize).clamp(0, ny as isize - 1) as usize;
310                    let nk = (ku + dz as isize).clamp(0, nz as isize - 1) as usize;
311                    let wx = if dx_off == 0 { 1.0 - fxu } else { fxu };
312                    let wy = if dy == 0 { 1.0 - fyu } else { fyu };
313                    let wz = if dz == 0 { 1.0 - fzu } else { fzu };
314                    let w = wx * wy * wz;
315                    let uidx = nk * (nx + 1) * ny + nj * (nx + 1) + ni;
316                    u_num[uidx] += w * p.velocity[0];
317                    u_weight[uidx] += w;
318                }
319            }
320        }
321        let iv = (px * inv_dx - 0.5).floor() as isize;
322        let jv = (py * inv_dx).floor() as isize;
323        let kv = (pz * inv_dx - 0.5).floor() as isize;
324        let fxv = px * inv_dx - 0.5 - iv as f64;
325        let fyv = py * inv_dx - jv as f64;
326        let fzv = pz * inv_dx - 0.5 - kv as f64;
327        for dz in 0..2 {
328            for dy in 0..2 {
329                for dx_off in 0..2 {
330                    let ni = (iv + dx_off as isize).clamp(0, nx as isize - 1) as usize;
331                    let nj = (jv + dy as isize).clamp(0, ny as isize) as usize;
332                    let nk = (kv + dz as isize).clamp(0, nz as isize - 1) as usize;
333                    let wx = if dx_off == 0 { 1.0 - fxv } else { fxv };
334                    let wy = if dy == 0 { 1.0 - fyv } else { fyv };
335                    let wz = if dz == 0 { 1.0 - fzv } else { fzv };
336                    let wt = wx * wy * wz;
337                    let vidx = nk * nx * (ny + 1) + nj * nx + ni;
338                    v_num[vidx] += wt * p.velocity[1];
339                    v_weight[vidx] += wt;
340                }
341            }
342        }
343        let iw = (px * inv_dx - 0.5).floor() as isize;
344        let jw = (py * inv_dx - 0.5).floor() as isize;
345        let kw = (pz * inv_dx).floor() as isize;
346        let fxw = px * inv_dx - 0.5 - iw as f64;
347        let fyw = py * inv_dx - 0.5 - jw as f64;
348        let fzw = pz * inv_dx - kw as f64;
349        for dz in 0..2 {
350            for dy in 0..2 {
351                for dx_off in 0..2 {
352                    let ni = (iw + dx_off as isize).clamp(0, nx as isize - 1) as usize;
353                    let nj = (jw + dy as isize).clamp(0, ny as isize - 1) as usize;
354                    let nk = (kw + dz as isize).clamp(0, nz as isize) as usize;
355                    let wx = if dx_off == 0 { 1.0 - fxw } else { fxw };
356                    let wy = if dy == 0 { 1.0 - fyw } else { fyw };
357                    let wz = if dz == 0 { 1.0 - fzw } else { fzw };
358                    let wt = wx * wy * wz;
359                    let widx = nk * nx * ny + nj * nx + ni;
360                    w_num[widx] += wt * p.velocity[2];
361                    w_weight[widx] += wt;
362                }
363            }
364        }
365    }
366    for i in 0..u_num.len() {
367        grid.u[i] = if u_weight[i] > 1e-15 {
368            u_num[i] / u_weight[i]
369        } else {
370            0.0
371        };
372    }
373    for i in 0..v_num.len() {
374        grid.v[i] = if v_weight[i] > 1e-15 {
375            v_num[i] / v_weight[i]
376        } else {
377            0.0
378        };
379    }
380    for i in 0..w_num.len() {
381        grid.w[i] = if w_weight[i] > 1e-15 {
382            w_num[i] / w_weight[i]
383        } else {
384            0.0
385        };
386    }
387}
388/// Grid-to-Particle (G2P) transfer: update particle velocities from MAC grid.
389///
390/// `flip_ratio` in \[0,1\]: 1.0 = pure FLIP, 0.0 = pure PIC.
391pub fn g2p_transfer(
392    particles: &mut [FlipParticle],
393    grid_new: &MacGrid,
394    grid_old: &MacGrid,
395    flip_ratio: f64,
396) {
397    for p in particles.iter_mut() {
398        let [px, py, pz] = p.position;
399        let dx = grid_new.dx;
400        let u_new = grid_new.interp_u(px, py, pz);
401        let u_old = grid_old.interp_u(px, py, pz);
402        let pic_vel = [u_new, 0.0, 0.0];
403        let flip_delta = [u_new - u_old, 0.0, 0.0];
404        let flip_vel = add3(p.velocity, flip_delta);
405        p.velocity = [
406            flip_ratio * flip_vel[0] + (1.0 - flip_ratio) * pic_vel[0],
407            p.velocity[1],
408            p.velocity[2],
409        ];
410        p.position = add3(p.position, scale3(p.velocity, dx));
411    }
412}
413/// GPU-accelerated SPH density summation.
414///
415/// Simulates a GPU kernel where each "thread" computes the density contribution
416/// from all neighbours within the smoothing radius h. In a real GPU
417/// implementation this would be launched as one thread per particle pair with
418/// atomic accumulation. Here we use a parallel-style double loop and write
419/// results into a temporary buffer first to keep the interface identical to a
420/// real GPU dispatch.
421pub fn gpu_sph_density_parallel(particles: &mut [SphParticle], config: &SphConfig) {
422    let n = particles.len();
423    let mut densities = vec![0.0f64; n];
424    for i in 0..n {
425        let mut rho = 0.0;
426        for j in 0..n {
427            let r_vec = sub3(particles[i].position, particles[j].position);
428            let r = length3(r_vec);
429            rho += particles[j].mass * SphKernels::poly6(r, config.h);
430        }
431        densities[i] = rho.max(1e-6);
432    }
433    for (i, p) in particles.iter_mut().enumerate() {
434        p.density = densities[i];
435    }
436}
437/// GPU Jacobi pressure solver on a MAC grid.
438///
439/// Identical algorithm to `MacGrid::jacobi_pressure_solve` but expressed as a
440/// function that could be launched as a compute shader (one thread per cell).
441/// The indirection through a separate function makes the GPU dispatch boundary
442/// explicit.
443pub fn gpu_jacobi_pressure_solve(grid: &mut MacGrid, rho: f64, dt: f64, iterations: usize) {
444    grid.compute_divergence();
445    let scale = rho * grid.dx * grid.dx / dt;
446    let nx = grid.nx;
447    let ny = grid.ny;
448    let nz = grid.nz;
449    let mut p_new = grid.p.clone();
450    for _ in 0..iterations {
451        for k in 0..nz {
452            for j in 0..ny {
453                for i in 0..nx {
454                    let idx = k * nx * ny + j * nx + i;
455                    if grid.flags[idx] != 1 {
456                        continue;
457                    }
458                    let mut nb_sum = 0.0;
459                    let mut nb_cnt = 0u32;
460                    macro_rules! nb {
461                        ($ii:expr, $jj:expr, $kk:expr) => {{
462                            nb_sum += grid.p[$kk * nx * ny + $jj * nx + $ii];
463                            nb_cnt += 1;
464                        }};
465                    }
466                    if i + 1 < nx {
467                        nb!(i + 1, j, k);
468                    }
469                    if i > 0 {
470                        nb!(i - 1, j, k);
471                    }
472                    if j + 1 < ny {
473                        nb!(i, j + 1, k);
474                    }
475                    if j > 0 {
476                        nb!(i, j - 1, k);
477                    }
478                    if k + 1 < nz {
479                        nb!(i, j, k + 1);
480                    }
481                    if k > 0 {
482                        nb!(i, j, k - 1);
483                    }
484                    if nb_cnt > 0 {
485                        p_new[idx] = (nb_sum - scale * grid.div[idx]) / nb_cnt as f64;
486                    }
487                }
488            }
489        }
490        grid.p.copy_from_slice(&p_new);
491    }
492}
493/// GPU-style BGK collision step for the D2Q9 LBM grid.
494///
495/// In a real GPU implementation each cell maps to one thread.  The function
496/// signature mirrors a compute-shader kernel: it takes the whole grid and
497/// applies BGK in-place.
498pub fn gpu_lbm_bgk_collide(lbm: &mut LbmD2Q9) {
499    let nx = lbm.nx;
500    let ny = lbm.ny;
501    let inv_tau = lbm.inv_tau;
502    let mut updates: Vec<(usize, usize, usize, f64)> = Vec::with_capacity(nx * ny * D2Q9_Q);
503    for y in 0..ny {
504        for x in 0..nx {
505            if lbm.cell_type[y * nx + x] == LbmCellType::Solid {
506                continue;
507            }
508            let rho = lbm.density(x, y);
509            let [ux, uy] = lbm.velocity(x, y);
510            for q in 0..D2Q9_Q {
511                let f_eq = LbmD2Q9::f_equilibrium(rho, ux, uy, q);
512                let f_old = lbm.get_f(x, y, q);
513                let f_new = f_old - inv_tau * (f_old - f_eq);
514                updates.push((x, y, q, f_new));
515            }
516        }
517    }
518    for (x, y, q, val) in updates {
519        lbm.set_f(x, y, q, val);
520    }
521}
522/// Expand a 10-bit integer into a 30-bit Morton code component (bit interleave).
523pub fn morton_expand_bits(mut v: u32) -> u32 {
524    v &= 0x000003ff;
525    v = (v | (v << 16)) & 0xff0000ff;
526    v = (v | (v << 8)) & 0x0300f00f;
527    v = (v | (v << 4)) & 0x030c30c3;
528    v = (v | (v << 2)) & 0x09249249;
529    v
530}
531/// Encode 3D integer coordinates into a 30-bit Morton (Z-order curve) code.
532pub fn morton_encode_3d(x: u32, y: u32, z: u32) -> u32 {
533    morton_expand_bits(x) | (morton_expand_bits(y) << 1) | (morton_expand_bits(z) << 2)
534}
535/// Sort SPH particles by their Morton code for cache-friendly GPU access.
536///
537/// Positions are normalised into the unit cube `[0, domain_size]` and then
538/// quantised to 10 bits per axis (1024 levels) before Morton encoding.
539pub fn morton_sort_particles(particles: &mut Vec<SphParticle>, domain_size: [f64; 3]) {
540    let bits = 1024u32;
541    let mut keyed: Vec<(u32, SphParticle)> = particles
542        .drain(..)
543        .map(|p| {
544            let xi = ((p.position[0] / domain_size[0]).clamp(0.0, 1.0) * (bits - 1) as f64) as u32;
545            let yi = ((p.position[1] / domain_size[1]).clamp(0.0, 1.0) * (bits - 1) as f64) as u32;
546            let zi = ((p.position[2] / domain_size[2]).clamp(0.0, 1.0) * (bits - 1) as f64) as u32;
547            (morton_encode_3d(xi, yi, zi), p)
548        })
549        .collect();
550    keyed.sort_by_key(|(code, _)| *code);
551    *particles = keyed.into_iter().map(|(_, p)| p).collect();
552}
553/// GPU Euler integration: v += a*dt, x += v*dt.
554///
555/// Maps to one GPU thread per particle.
556pub fn gpu_particle_integrate_euler(particles: &mut [SphParticle], dt: f64) {
557    for p in particles.iter_mut() {
558        let inv_rho = 1.0 / p.density.max(1e-6);
559        for d in 0..3 {
560            let accel = p.force[d] * inv_rho;
561            p.velocity[d] += accel * dt;
562            p.position[d] += p.velocity[d] * dt;
563        }
564    }
565}
566/// GPU Verlet integration using the previous time-step `dt_prev`.
567///
568/// x_new = x + v*dt + 0.5*a*dt^2 (Störmer–Verlet, velocity-explicit variant).
569pub fn gpu_particle_integrate_verlet(particles: &mut [SphParticle], dt: f64, _dt_prev: f64) {
570    for p in particles.iter_mut() {
571        let inv_rho = 1.0 / p.density.max(1e-6);
572        for d in 0..3 {
573            let accel = p.force[d] * inv_rho;
574            p.position[d] += p.velocity[d] * dt + 0.5 * accel * dt * dt;
575            p.velocity[d] += accel * dt;
576        }
577    }
578}
579/// GPU boundary condition: clamp particles inside an AABB and reflect velocities.
580///
581/// One GPU thread per particle; no inter-thread communication needed.
582pub fn gpu_apply_boundary_box(particles: &mut [SphParticle], bounds: &GpuBoundaryBox) {
583    for p in particles.iter_mut() {
584        for d in 0..3 {
585            if p.position[d] < bounds.min[d] {
586                p.position[d] = bounds.min[d];
587                p.velocity[d] = p.velocity[d].abs() * bounds.restitution;
588            }
589            if p.position[d] > bounds.max[d] {
590                p.position[d] = bounds.max[d];
591                p.velocity[d] = -p.velocity[d].abs() * bounds.restitution;
592            }
593        }
594    }
595}
596/// GPU parallel reduction for total kinetic energy.
597///
598/// In GPU hardware this would be a tree-reduction across threads in a work-group;
599/// here we use an iterator fold which has the same semantics.
600///
601/// KE = Σ 0.5 * mass * |v|²
602pub fn gpu_reduce_kinetic_energy(particles: &[SphParticle], mass: f64) -> f64 {
603    particles.iter().fold(0.0, |acc, p| {
604        let v2 = p.velocity[0] * p.velocity[0]
605            + p.velocity[1] * p.velocity[1]
606            + p.velocity[2] * p.velocity[2];
607        acc + 0.5 * mass * v2
608    })
609}
610/// GPU parallel reduction for total linear momentum.
611///
612/// Returns a 3-vector \[px, py, pz\] = Σ mass * v_i.
613pub fn gpu_reduce_momentum(particles: &[SphParticle], mass: f64) -> [f64; 3] {
614    particles.iter().fold([0.0f64; 3], |acc, p| {
615        [
616            acc[0] + mass * p.velocity[0],
617            acc[1] + mass * p.velocity[1],
618            acc[2] + mass * p.velocity[2],
619        ]
620    })
621}
622/// GPU semi-Lagrangian scalar field advection on a 2D regular grid.
623///
624/// Each cell is a separate GPU thread: back-traces the characteristic by dt
625/// through the supplied velocity field `vel` (one 2-vector per cell) and
626/// samples the field using bilinear interpolation.
627pub fn gpu_advect_2d(
628    field: &[f64],
629    vel: &[[f64; 2]],
630    nx: usize,
631    ny: usize,
632    dx: f64,
633    dt: f64,
634) -> Vec<f64> {
635    let mut out = vec![0.0f64; nx * ny];
636    let inv_dx = 1.0 / dx;
637    let sample = |x: f64, y: f64| -> f64 {
638        let ix = (x * inv_dx - 0.5).floor() as isize;
639        let iy = (y * inv_dx - 0.5).floor() as isize;
640        let fx = x * inv_dx - 0.5 - ix as f64;
641        let fy = y * inv_dx - 0.5 - iy as f64;
642        let ci = |v: isize| v.clamp(0, nx as isize - 1) as usize;
643        let cj = |v: isize| v.clamp(0, ny as isize - 1) as usize;
644        let f00 = field[cj(iy) * nx + ci(ix)];
645        let f10 = field[cj(iy) * nx + ci(ix + 1)];
646        let f01 = field[cj(iy + 1) * nx + ci(ix)];
647        let f11 = field[cj(iy + 1) * nx + ci(ix + 1)];
648        let f0 = f00 + fx * (f10 - f00);
649        let f1 = f01 + fx * (f11 - f01);
650        f0 + fy * (f1 - f0)
651    };
652    for j in 0..ny {
653        for i in 0..nx {
654            let xc = (i as f64 + 0.5) * dx;
655            let yc = (j as f64 + 0.5) * dx;
656            let idx = j * nx + i;
657            let [vx, vy] = vel[idx];
658            let xb = xc - vx * dt;
659            let yb = yc - vy * dt;
660            out[idx] = sample(xb, yb);
661        }
662    }
663    out
664}
665/// GPU Jacobi solver for the 2D pressure-Poisson equation on a staggered grid.
666///
667/// Solves ∇²p = rhs (passed as `div`) for `iterations` Jacobi sweeps.
668/// The staggered MAC discretisation gives the standard 5-point stencil:
669///   p\[i,j\] = (p\[i+1,j\] + p\[i-1,j\] + p\[i,j+1\] + p\[i,j-1\] - dx²*rhs\[i,j\]) / 4
670pub fn gpu_pressure_poisson_jacobi_2d(
671    pressure: &mut [f64],
672    div: &[f64],
673    nx: usize,
674    ny: usize,
675    dx: f64,
676    iterations: usize,
677) {
678    let dx2 = dx * dx;
679    let mut p_new = pressure.to_vec();
680    for _ in 0..iterations {
681        for j in 0..ny {
682            for i in 0..nx {
683                let idx = j * nx + i;
684                let mut nb = 0.0;
685                let mut cnt = 0u32;
686                if i + 1 < nx {
687                    nb += pressure[j * nx + i + 1];
688                    cnt += 1;
689                }
690                if i > 0 {
691                    nb += pressure[j * nx + i - 1];
692                    cnt += 1;
693                }
694                if j + 1 < ny {
695                    nb += pressure[(j + 1) * nx + i];
696                    cnt += 1;
697                }
698                if j > 0 {
699                    nb += pressure[(j - 1) * nx + i];
700                    cnt += 1;
701                }
702                if cnt > 0 {
703                    p_new[idx] = (nb - dx2 * div[idx]) / cnt as f64;
704                }
705            }
706        }
707        pressure.copy_from_slice(&p_new);
708    }
709}