1use 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}
31pub 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}
50pub 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}
57pub 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}
93pub 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}
101pub 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}
108pub const D3Q27_Q: usize = 27;
110pub const D2Q9_Q: usize = 9;
112pub const D2Q9_CX: [i32; 9] = [0, 1, 0, -1, 0, 1, -1, -1, 1];
114pub const D2Q9_CY: [i32; 9] = [0, 0, 1, 0, -1, 1, 1, -1, -1];
116pub 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];
128pub const D2Q9_OPP: [usize; 9] = [0, 3, 4, 1, 2, 7, 8, 5, 6];
130pub 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}
165pub 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}
190pub 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}
234pub 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}
282pub 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}
388pub 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}
413pub 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}
437pub 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}
493pub 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}
522pub 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}
531pub 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}
535pub 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}
553pub 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}
566pub 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}
579pub 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}
596pub 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}
610pub 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}
622pub 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}
665pub 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}