From e0cbb9934398d040a641885d90b5895fe1358d4e Mon Sep 17 00:00:00 2001 From: Omar Sobh Date: Thu, 17 Sep 2026 11:33:13 -0500 Subject: [PATCH] rtx-cfd three_d: the predictors split into piso_predictor.rs (file-size rule), the w predictor's index helpers cleaned; tests unchanged and green Co-Authored-By: Claude Fable 5.1 --- .../src/solvers/incompressible/three_d/mod.rs | 1 + .../incompressible/three_d/piso_host.rs | 704 +----------------- .../incompressible/three_d/piso_predictor.rs | 700 +++++++++++++++++ 3 files changed, 705 insertions(+), 700 deletions(-) create mode 100644 crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/piso_predictor.rs diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/mod.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/mod.rs index 1eb2ad3..e6fff1f 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/mod.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/mod.rs @@ -9,6 +9,7 @@ pub mod flow_field; pub mod piso_host; +mod piso_predictor; pub mod poisson; pub use flow_field::FlowField3D; diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/piso_host.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/piso_host.rs index 75a0375..6eb9f3e 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/piso_host.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/piso_host.rs @@ -41,7 +41,7 @@ impl Boundaries3 { .contains(&SideBoundary3::PressureOutlet) } - fn periodic_z(self) -> bool { + pub(super) fn periodic_z(self) -> bool { self.z0 == SideBoundary3::Periodic } } @@ -97,7 +97,7 @@ type Vec3Fn = Box (f64, f64, f64) + Send + Sync>; pub struct Piso3Solver { pub fluid: Fluid3, pub params: Piso3Parameters, - momentum_source: Option, + pub(super) momentum_source: Option, boundary_velocity: Option, pcg_cache: PcgCache3, time: f64, @@ -158,7 +158,7 @@ impl Piso3Solver { self.poisson_profile } - fn boundary(&self, x: f64, y: f64, z: f64, t: f64) -> (f64, f64, f64) { + pub(super) fn boundary(&self, x: f64, y: f64, z: f64, t: f64) -> (f64, f64, f64) { self.boundary_velocity .as_ref() .map_or((0.0, 0.0, 0.0), |f| f(x, y, z, t)) @@ -182,7 +182,7 @@ impl Piso3Solver { true } - fn upwind(face_velocity: f64, upstream: f64, downstream: f64) -> f64 { + pub(super) fn upwind(face_velocity: f64, upstream: f64, downstream: f64) -> f64 { if face_velocity >= 0.0 { upstream } else { @@ -233,702 +233,6 @@ impl Piso3Solver { } } - /// Plane above `k` (wrapping when periodic). - #[inline] - fn k_up(&self, k: usize, nz: usize) -> Option { - if k + 1 < nz { - Some(k + 1) - } else if self.params.boundaries.periodic_z() { - Some(0) - } else { - None - } - } - - #[inline] - fn k_down(&self, k: usize, nz: usize) -> Option { - if k > 0 { - Some(k - 1) - } else if self.params.boundaries.periodic_z() { - Some(nz - 1) - } else { - None - } - } - - /// The predictor's right-hand side on the u face `(k, j, i)`, `i = 1..nx`: - /// the 2D `u_rhs` expression for expression, then `− conv_z + diff_z`. - #[allow(clippy::too_many_lines)] - pub(crate) fn u_rhs( - &self, - field: &FlowField3D, - k: usize, - j: usize, - i: usize, - t_old: f64, - ) -> f64 { - let g = field.grid; - let (nx, ny, nz, dx, dy, dz) = (g.nx, g.ny, g.nz, g.dx, g.dy, g.dz); - let rho = self.fluid.density; - let nu = self.fluid.viscosity / rho; - let b = self.params.boundaries; - let velocity = SideBoundary3::Velocity; - let uo = &field.u_old; - let vo = &field.v_old; - let wo = &field.w_old; - let uf = |kk: usize, jj: usize, ii: usize| g.uface(kk, jj, ii); - let vf = |kk: usize, jj: usize, ii: usize| g.vface(kk, jj, ii); - let wf = |kk: usize, jj: usize, ii: usize| g.wface(kk, jj, ii); - let zc = (k as f64 + 0.5) * dz; - let u_p = uo[uf(k, j, i)]; - - let ue_face = 0.5 * (uo[uf(k, j, i)] + uo[uf(k, j, i + 1)]); - let uw_face = 0.5 * (uo[uf(k, j, i - 1)] + uo[uf(k, j, i)]); - - let south_is_wall = j == 0; - let north_is_wall = j + 1 == ny; - - let vn_face = 0.5 * (vo[vf(k, j + 1, i - 1)] + vo[vf(k, j + 1, i)]); - let vs_face = 0.5 * (vo[vf(k, j, i - 1)] + vo[vf(k, j, i)]); - - let beyond_north = if b.y1 == velocity { - self.boundary(i as f64 * dx, ny as f64 * dy, zc, t_old).0 - } else { - u_p - }; - let beyond_south = if b.y0 == velocity { - self.boundary(i as f64 * dx, 0.0, zc, t_old).0 - } else { - u_p - }; - - let conv_x = (ue_face * Self::upwind(ue_face, uo[uf(k, j, i)], uo[uf(k, j, i + 1)]) - - uw_face * Self::upwind(uw_face, uo[uf(k, j, i - 1)], uo[uf(k, j, i)])) - / dx; - let conv_y = (vn_face - * if north_is_wall { - Self::upwind(vn_face, u_p, beyond_north) - } else { - Self::upwind(vn_face, uo[uf(k, j, i)], uo[uf(k, j + 1, i)]) - } - - vs_face - * if south_is_wall { - Self::upwind(vs_face, beyond_south, u_p) - } else { - Self::upwind(vs_face, uo[uf(k, j - 1, i)], uo[uf(k, j, i)]) - }) - / dy; - - let scheme = self.params.convection_scheme; - let mut conv_x = conv_x; - let mut conv_y = conv_y; - if scheme != ConvectionScheme::Upwind { - let delta_e = if ue_face >= 0.0 { - scheme.face_correction( - Some(uo[uf(k, j, i - 1)]), - uo[uf(k, j, i)], - uo[uf(k, j, i + 1)], - ) - } else { - let far = (i + 2 <= nx).then(|| uo[uf(k, j, i + 2)]); - scheme.face_correction(far, uo[uf(k, j, i + 1)], uo[uf(k, j, i)]) - }; - let delta_w = if uw_face >= 0.0 { - let far = (i >= 2).then(|| uo[uf(k, j, i - 2)]); - scheme.face_correction(far, uo[uf(k, j, i - 1)], uo[uf(k, j, i)]) - } else { - scheme.face_correction( - Some(uo[uf(k, j, i + 1)]), - uo[uf(k, j, i)], - uo[uf(k, j, i - 1)], - ) - }; - let delta_n = if north_is_wall { - 0.0 - } else if vn_face >= 0.0 { - let far = (j >= 1).then(|| uo[uf(k, j - 1, i)]); - scheme.face_correction(far, uo[uf(k, j, i)], uo[uf(k, j + 1, i)]) - } else { - let far = (j + 2 < ny).then(|| uo[uf(k, j + 2, i)]); - scheme.face_correction(far, uo[uf(k, j + 1, i)], uo[uf(k, j, i)]) - }; - let delta_s = if south_is_wall { - 0.0 - } else if vs_face >= 0.0 { - let far = (j >= 2).then(|| uo[uf(k, j - 2, i)]); - scheme.face_correction(far, uo[uf(k, j - 1, i)], uo[uf(k, j, i)]) - } else { - let far = (j + 1 < ny).then(|| uo[uf(k, j + 1, i)]); - scheme.face_correction(far, uo[uf(k, j, i)], uo[uf(k, j - 1, i)]) - }; - conv_x += (ue_face * delta_e - uw_face * delta_w) / dx; - conv_y += (vn_face * delta_n - vs_face * delta_s) / dy; - } - - let diff_x = nu * (uo[uf(k, j, i + 1)] - 2.0 * u_p + uo[uf(k, j, i - 1)]) / (dx * dx); - - let flux_north = if north_is_wall { - if b.y1 == velocity { - let u_wall = self.boundary(i as f64 * dx, ny as f64 * dy, zc, t_old).0; - nu * (u_wall - u_p) / (0.5 * dy) - } else { - 0.0 - } - } else { - nu * (uo[uf(k, j + 1, i)] - u_p) / dy - }; - let flux_south = if south_is_wall { - if b.y0 == velocity { - let u_wall = self.boundary(i as f64 * dx, 0.0, zc, t_old).0; - nu * (u_p - u_wall) / (0.5 * dy) - } else { - 0.0 - } - } else { - nu * (u_p - uo[uf(k, j - 1, i)]) / dy - }; - let diff_y = (flux_north - flux_south) / dy; - - let pressure_gradient = - -(field.p[g.cell(k, j, i)] - field.p[g.cell(k, j, i - 1)]) / (rho * dx); - - let body_force = self.momentum_source.as_ref().map_or(0.0, |f| { - f(i as f64 * dx, (j as f64 + 0.5) * dy, zc, t_old).0 / rho - }); - - let rhs_2d = -conv_x - conv_y + diff_x + diff_y + pressure_gradient + body_force; - - // --- the z terms, the y pattern turned along k --- - let ku = self.k_up(k, nz); - let kd = self.k_down(k, nz); - let top_is_wall = ku.is_none(); - let bottom_is_wall = kd.is_none(); - // The w faces above/below the u face: on top of the cells west and - // east of it (face k + 1 of cell k is face index k + 1; periodic: - // the face at k = nz equals the face at 0). - let wt_face = 0.5 * (wo[wf(k + 1, j, i - 1)] + wo[wf(k + 1, j, i)]); - let wb_face = 0.5 * (wo[wf(k, j, i - 1)] + wo[wf(k, j, i)]); - let beyond_top = if b.z1 == velocity { - self.boundary(i as f64 * dx, (j as f64 + 0.5) * dy, nz as f64 * dz, t_old) - .0 - } else { - u_p - }; - let beyond_bottom = if b.z0 == velocity { - self.boundary(i as f64 * dx, (j as f64 + 0.5) * dy, 0.0, t_old) - .0 - } else { - u_p - }; - let u_up = ku.map(|kk| uo[uf(kk, j, i)]); - let u_dn = kd.map(|kk| uo[uf(kk, j, i)]); - let mut conv_z = (wt_face - * match u_up { - Some(un) => Self::upwind(wt_face, u_p, un), - None => Self::upwind(wt_face, u_p, beyond_top), - } - - wb_face - * match u_dn { - Some(ud) => Self::upwind(wb_face, ud, u_p), - None => Self::upwind(wb_face, beyond_bottom, u_p), - }) - / dz; - if scheme != ConvectionScheme::Upwind { - let far_up2 = ku - .and_then(|kk| self.k_up(kk, nz)) - .map(|kk| uo[uf(kk, j, i)]); - let far_dn2 = kd - .and_then(|kk| self.k_down(kk, nz)) - .map(|kk| uo[uf(kk, j, i)]); - let delta_t = if top_is_wall { - 0.0 - } else if wt_face >= 0.0 { - scheme.face_correction(u_dn, u_p, u_up.unwrap_or(u_p)) - } else { - scheme.face_correction(far_up2, u_up.unwrap_or(u_p), u_p) - }; - let delta_b = if bottom_is_wall { - 0.0 - } else if wb_face >= 0.0 { - scheme.face_correction(far_dn2, u_dn.unwrap_or(u_p), u_p) - } else { - scheme.face_correction(u_up, u_p, u_dn.unwrap_or(u_p)) - }; - conv_z += (wt_face * delta_t - wb_face * delta_b) / dz; - } - let flux_top = match u_up { - Some(un) => nu * (un - u_p) / dz, - None => { - if b.z1 == velocity { - nu * (beyond_top - u_p) / (0.5 * dz) - } else { - 0.0 - } - } - }; - let flux_bottom = match u_dn { - Some(ud) => nu * (u_p - ud) / dz, - None => { - if b.z0 == velocity { - nu * (u_p - beyond_bottom) / (0.5 * dz) - } else { - 0.0 - } - } - }; - let diff_z = (flux_top - flux_bottom) / dz; - - rhs_2d - conv_z + diff_z - } - - /// The v face `(k, j, i)`, `j = 1..ny`: the 2D `v_rhs` then the z terms. - #[allow(clippy::too_many_lines)] - pub(crate) fn v_rhs( - &self, - field: &FlowField3D, - k: usize, - j: usize, - i: usize, - t_old: f64, - ) -> f64 { - let g = field.grid; - let (nx, ny, nz, dx, dy, dz) = (g.nx, g.ny, g.nz, g.dx, g.dy, g.dz); - let rho = self.fluid.density; - let nu = self.fluid.viscosity / rho; - let b = self.params.boundaries; - let velocity = SideBoundary3::Velocity; - let uo = &field.u_old; - let vo = &field.v_old; - let wo = &field.w_old; - let uf = |kk: usize, jj: usize, ii: usize| g.uface(kk, jj, ii); - let vf = |kk: usize, jj: usize, ii: usize| g.vface(kk, jj, ii); - let wf = |kk: usize, jj: usize, ii: usize| g.wface(kk, jj, ii); - let zc = (k as f64 + 0.5) * dz; - let v_p = vo[vf(k, j, i)]; - - let vn_face = 0.5 * (vo[vf(k, j, i)] + vo[vf(k, j + 1, i)]); - let vs_face = 0.5 * (vo[vf(k, j - 1, i)] + vo[vf(k, j, i)]); - - let west_is_wall = i == 0; - let east_is_wall = i + 1 == nx; - - let ue_face = 0.5 * (uo[uf(k, j - 1, i + 1)] + uo[uf(k, j, i + 1)]); - let uw_face = 0.5 * (uo[uf(k, j - 1, i)] + uo[uf(k, j, i)]); - let beyond_east = if b.x1 == velocity { - self.boundary(nx as f64 * dx, j as f64 * dy, zc, t_old).1 - } else { - v_p - }; - let beyond_west = if b.x0 == velocity { - self.boundary(0.0, j as f64 * dy, zc, t_old).1 - } else { - v_p - }; - - let conv_y = (vn_face * Self::upwind(vn_face, vo[vf(k, j, i)], vo[vf(k, j + 1, i)]) - - vs_face * Self::upwind(vs_face, vo[vf(k, j - 1, i)], vo[vf(k, j, i)])) - / dy; - let conv_x = (ue_face - * if east_is_wall { - Self::upwind(ue_face, v_p, beyond_east) - } else { - Self::upwind(ue_face, vo[vf(k, j, i)], vo[vf(k, j, i + 1)]) - } - - uw_face - * if west_is_wall { - Self::upwind(uw_face, beyond_west, v_p) - } else { - Self::upwind(uw_face, vo[vf(k, j, i - 1)], vo[vf(k, j, i)]) - }) - / dx; - - let scheme = self.params.convection_scheme; - let mut conv_x = conv_x; - let mut conv_y = conv_y; - if scheme != ConvectionScheme::Upwind { - let delta_n = if vn_face >= 0.0 { - scheme.face_correction( - Some(vo[vf(k, j - 1, i)]), - vo[vf(k, j, i)], - vo[vf(k, j + 1, i)], - ) - } else { - let far = (j + 2 <= ny).then(|| vo[vf(k, j + 2, i)]); - scheme.face_correction(far, vo[vf(k, j + 1, i)], vo[vf(k, j, i)]) - }; - let delta_s = if vs_face >= 0.0 { - let far = (j >= 2).then(|| vo[vf(k, j - 2, i)]); - scheme.face_correction(far, vo[vf(k, j - 1, i)], vo[vf(k, j, i)]) - } else { - scheme.face_correction( - Some(vo[vf(k, j + 1, i)]), - vo[vf(k, j, i)], - vo[vf(k, j - 1, i)], - ) - }; - let delta_e = if east_is_wall { - 0.0 - } else if ue_face >= 0.0 { - let far = (i >= 1).then(|| vo[vf(k, j, i - 1)]); - scheme.face_correction(far, vo[vf(k, j, i)], vo[vf(k, j, i + 1)]) - } else { - let far = (i + 2 < nx).then(|| vo[vf(k, j, i + 2)]); - scheme.face_correction(far, vo[vf(k, j, i + 1)], vo[vf(k, j, i)]) - }; - let delta_w = if west_is_wall { - 0.0 - } else if uw_face >= 0.0 { - let far = (i >= 2).then(|| vo[vf(k, j, i - 2)]); - scheme.face_correction(far, vo[vf(k, j, i - 1)], vo[vf(k, j, i)]) - } else { - let far = (i + 1 < nx).then(|| vo[vf(k, j, i + 1)]); - scheme.face_correction(far, vo[vf(k, j, i)], vo[vf(k, j, i - 1)]) - }; - conv_y += (vn_face * delta_n - vs_face * delta_s) / dy; - conv_x += (ue_face * delta_e - uw_face * delta_w) / dx; - } - - let diff_y = nu * (vo[vf(k, j + 1, i)] - 2.0 * v_p + vo[vf(k, j - 1, i)]) / (dy * dy); - - let flux_east = if east_is_wall { - if b.x1 == velocity { - let v_wall = self.boundary(nx as f64 * dx, j as f64 * dy, zc, t_old).1; - nu * (v_wall - v_p) / (0.5 * dx) - } else { - 0.0 - } - } else { - nu * (vo[vf(k, j, i + 1)] - v_p) / dx - }; - let flux_west = if west_is_wall { - if b.x0 == velocity { - let v_wall = self.boundary(0.0, j as f64 * dy, zc, t_old).1; - nu * (v_p - v_wall) / (0.5 * dx) - } else { - 0.0 - } - } else { - nu * (v_p - vo[vf(k, j, i - 1)]) / dx - }; - let diff_x = (flux_east - flux_west) / dx; - - let pressure_gradient = - -(field.p[g.cell(k, j, i)] - field.p[g.cell(k, j - 1, i)]) / (rho * dy); - - let body_force = self.momentum_source.as_ref().map_or(0.0, |f| { - f((i as f64 + 0.5) * dx, j as f64 * dy, zc, t_old).1 / rho - }); - - let rhs_2d = -conv_x - conv_y + diff_x + diff_y + pressure_gradient + body_force; - - // --- z terms --- - let ku = self.k_up(k, nz); - let kd = self.k_down(k, nz); - let top_is_wall = ku.is_none(); - let bottom_is_wall = kd.is_none(); - let wt_face = 0.5 * (wo[wf(k + 1, j - 1, i)] + wo[wf(k + 1, j, i)]); - let wb_face = 0.5 * (wo[wf(k, j - 1, i)] + wo[wf(k, j, i)]); - let beyond_top = if b.z1 == velocity { - self.boundary((i as f64 + 0.5) * dx, j as f64 * dy, nz as f64 * dz, t_old) - .1 - } else { - v_p - }; - let beyond_bottom = if b.z0 == velocity { - self.boundary((i as f64 + 0.5) * dx, j as f64 * dy, 0.0, t_old) - .1 - } else { - v_p - }; - let v_up = ku.map(|kk| vo[vf(kk, j, i)]); - let v_dn = kd.map(|kk| vo[vf(kk, j, i)]); - let mut conv_z = (wt_face - * match v_up { - Some(vn) => Self::upwind(wt_face, v_p, vn), - None => Self::upwind(wt_face, v_p, beyond_top), - } - - wb_face - * match v_dn { - Some(vd) => Self::upwind(wb_face, vd, v_p), - None => Self::upwind(wb_face, beyond_bottom, v_p), - }) - / dz; - if scheme != ConvectionScheme::Upwind { - let far_up2 = ku - .and_then(|kk| self.k_up(kk, nz)) - .map(|kk| vo[vf(kk, j, i)]); - let far_dn2 = kd - .and_then(|kk| self.k_down(kk, nz)) - .map(|kk| vo[vf(kk, j, i)]); - let delta_t = if top_is_wall { - 0.0 - } else if wt_face >= 0.0 { - scheme.face_correction(v_dn, v_p, v_up.unwrap_or(v_p)) - } else { - scheme.face_correction(far_up2, v_up.unwrap_or(v_p), v_p) - }; - let delta_b = if bottom_is_wall { - 0.0 - } else if wb_face >= 0.0 { - scheme.face_correction(far_dn2, v_dn.unwrap_or(v_p), v_p) - } else { - scheme.face_correction(v_up, v_p, v_dn.unwrap_or(v_p)) - }; - conv_z += (wt_face * delta_t - wb_face * delta_b) / dz; - } - let flux_top = match v_up { - Some(vn) => nu * (vn - v_p) / dz, - None => { - if b.z1 == velocity { - nu * (beyond_top - v_p) / (0.5 * dz) - } else { - 0.0 - } - } - }; - let flux_bottom = match v_dn { - Some(vd) => nu * (v_p - vd) / dz, - None => { - if b.z0 == velocity { - nu * (v_p - beyond_bottom) / (0.5 * dz) - } else { - 0.0 - } - } - }; - let diff_z = (flux_top - flux_bottom) / dz; - - rhs_2d - conv_z + diff_z - } - - /// The w face `(k, j, i)` between cells `k − 1` (wrapping when periodic) - /// and `k`: the v pattern with z as its own direction and x, y transverse. - #[allow(clippy::too_many_lines)] - pub(crate) fn w_rhs( - &self, - field: &FlowField3D, - k: usize, - j: usize, - i: usize, - t_old: f64, - ) -> f64 { - let g = field.grid; - let (nx, ny, nz, dx, dy, dz) = (g.nx, g.ny, g.nz, g.dx, g.dy, g.dz); - let rho = self.fluid.density; - let nu = self.fluid.viscosity / rho; - let b = self.params.boundaries; - let velocity = SideBoundary3::Velocity; - let periodic = b.periodic_z(); - let uo = &field.u_old; - let vo = &field.v_old; - let wo = &field.w_old; - let uf = |kk: usize, jj: usize, ii: usize| g.uface(kk, jj, ii); - let vf = |kk: usize, jj: usize, ii: usize| g.vface(kk, jj, ii); - let wf = |kk: usize, jj: usize, ii: usize| g.wface(kk, jj, ii); - // The cells below and above this face, and the faces around them. - let k_below = if k > 0 { k - 1 } else { nz - 1 }; // k = 0 only when periodic - let k_above = k % nz; // k = nz only when periodic (the same face as 0) - let wface_of = |kk: usize| wf(kk % (nz + 1), j, i); - // Own-direction neighbours: the faces k − 1 and k + 1 (wrapping). - let w_dn_idx = if k > 0 { - wf(k - 1, j, i) - } else { - wf(nz - 1, j, i) - }; - let w_up_idx = if k + 1 <= nz { - if k + 1 == nz && periodic { - wf(0, j, i) - } else { - wf(k + 1, j, i) - } - } else { - wf(1, j, i) - }; - let _ = wface_of; - let zf = k as f64 * dz; - let w_p = wo[wf(k, j, i)]; - - let wt_face = 0.5 * (wo[wf(k, j, i)] + wo[w_up_idx]); - let wb_face = 0.5 * (wo[w_dn_idx] + wo[wf(k, j, i)]); - - let west_is_wall = i == 0; - let east_is_wall = i + 1 == nx; - let south_is_wall = j == 0; - let north_is_wall = j + 1 == ny; - - let ue_face = 0.5 * (uo[uf(k_below, j, i + 1)] + uo[uf(k_above, j, i + 1)]); - let uw_face = 0.5 * (uo[uf(k_below, j, i)] + uo[uf(k_above, j, i)]); - let vn_face = 0.5 * (vo[vf(k_below, j + 1, i)] + vo[vf(k_above, j + 1, i)]); - let vs_face = 0.5 * (vo[vf(k_below, j, i)] + vo[vf(k_above, j, i)]); - let beyond_east = if b.x1 == velocity { - self.boundary(nx as f64 * dx, (j as f64 + 0.5) * dy, zf, t_old) - .2 - } else { - w_p - }; - let beyond_west = if b.x0 == velocity { - self.boundary(0.0, (j as f64 + 0.5) * dy, zf, t_old).2 - } else { - w_p - }; - let beyond_north = if b.y1 == velocity { - self.boundary((i as f64 + 0.5) * dx, ny as f64 * dy, zf, t_old) - .2 - } else { - w_p - }; - let beyond_south = if b.y0 == velocity { - self.boundary((i as f64 + 0.5) * dx, 0.0, zf, t_old).2 - } else { - w_p - }; - - let conv_z = (wt_face * Self::upwind(wt_face, w_p, wo[w_up_idx]) - - wb_face * Self::upwind(wb_face, wo[w_dn_idx], w_p)) - / dz; - let conv_x = (ue_face - * if east_is_wall { - Self::upwind(ue_face, w_p, beyond_east) - } else { - Self::upwind(ue_face, w_p, wo[wf(k, j, i + 1)]) - } - - uw_face - * if west_is_wall { - Self::upwind(uw_face, beyond_west, w_p) - } else { - Self::upwind(uw_face, wo[wf(k, j, i - 1)], w_p) - }) - / dx; - let conv_y = (vn_face - * if north_is_wall { - Self::upwind(vn_face, w_p, beyond_north) - } else { - Self::upwind(vn_face, w_p, wo[wf(k, j + 1, i)]) - } - - vs_face - * if south_is_wall { - Self::upwind(vs_face, beyond_south, w_p) - } else { - Self::upwind(vs_face, wo[wf(k, j - 1, i)], w_p) - }) - / dy; - - let scheme = self.params.convection_scheme; - let mut conv_x = conv_x; - let mut conv_y = conv_y; - let mut conv_z = conv_z; - if scheme != ConvectionScheme::Upwind { - // Own direction: far nodes two faces away (wrapping when periodic). - let far_up2 = if periodic { - Some(wo[wf((k + 2) % nz, j, i)]) - } else { - (k + 2 <= nz).then(|| wo[wf(k + 2, j, i)]) - }; - let far_dn2 = if periodic { - Some(wo[wf((k + nz - 2) % nz, j, i)]) - } else { - (k >= 2).then(|| wo[wf(k - 2, j, i)]) - }; - let delta_t = if wt_face >= 0.0 { - scheme.face_correction(Some(wo[w_dn_idx]), w_p, wo[w_up_idx]) - } else { - scheme.face_correction(far_up2, wo[w_up_idx], w_p) - }; - let delta_b = if wb_face >= 0.0 { - scheme.face_correction(far_dn2, wo[w_dn_idx], w_p) - } else { - scheme.face_correction(Some(wo[w_up_idx]), w_p, wo[w_dn_idx]) - }; - let delta_e = if east_is_wall { - 0.0 - } else if ue_face >= 0.0 { - let far = (i >= 1).then(|| wo[wf(k, j, i - 1)]); - scheme.face_correction(far, w_p, wo[wf(k, j, i + 1)]) - } else { - let far = (i + 2 < nx).then(|| wo[wf(k, j, i + 2)]); - scheme.face_correction(far, wo[wf(k, j, i + 1)], w_p) - }; - let delta_w = if west_is_wall { - 0.0 - } else if uw_face >= 0.0 { - let far = (i >= 2).then(|| wo[wf(k, j, i - 2)]); - scheme.face_correction(far, wo[wf(k, j, i - 1)], w_p) - } else { - let far = (i + 1 < nx).then(|| wo[wf(k, j, i + 1)]); - scheme.face_correction(far, w_p, wo[wf(k, j, i - 1)]) - }; - let delta_n = if north_is_wall { - 0.0 - } else if vn_face >= 0.0 { - let far = (j >= 1).then(|| wo[wf(k, j - 1, i)]); - scheme.face_correction(far, w_p, wo[wf(k, j + 1, i)]) - } else { - let far = (j + 2 < ny).then(|| wo[wf(k, j + 2, i)]); - scheme.face_correction(far, wo[wf(k, j + 1, i)], w_p) - }; - let delta_s = if south_is_wall { - 0.0 - } else if vs_face >= 0.0 { - let far = (j >= 2).then(|| wo[wf(k, j - 2, i)]); - scheme.face_correction(far, wo[wf(k, j - 1, i)], w_p) - } else { - let far = (j + 1 < ny).then(|| wo[wf(k, j + 1, i)]); - scheme.face_correction(far, w_p, wo[wf(k, j - 1, i)]) - }; - conv_z += (wt_face * delta_t - wb_face * delta_b) / dz; - conv_x += (ue_face * delta_e - uw_face * delta_w) / dx; - conv_y += (vn_face * delta_n - vs_face * delta_s) / dy; - } - - let diff_z = nu * (wo[w_up_idx] - 2.0 * w_p + wo[w_dn_idx]) / (dz * dz); - let flux_east = if east_is_wall { - if b.x1 == velocity { - nu * (beyond_east - w_p) / (0.5 * dx) - } else { - 0.0 - } - } else { - nu * (wo[wf(k, j, i + 1)] - w_p) / dx - }; - let flux_west = if west_is_wall { - if b.x0 == velocity { - nu * (w_p - beyond_west) / (0.5 * dx) - } else { - 0.0 - } - } else { - nu * (w_p - wo[wf(k, j, i - 1)]) / dx - }; - let diff_x = (flux_east - flux_west) / dx; - let flux_north = if north_is_wall { - if b.y1 == velocity { - nu * (beyond_north - w_p) / (0.5 * dy) - } else { - 0.0 - } - } else { - nu * (wo[wf(k, j + 1, i)] - w_p) / dy - }; - let flux_south = if south_is_wall { - if b.y0 == velocity { - nu * (w_p - beyond_south) / (0.5 * dy) - } else { - 0.0 - } - } else { - nu * (w_p - wo[wf(k, j - 1, i)]) / dy - }; - let diff_y = (flux_north - flux_south) / dy; - - let pressure_gradient = - -(field.p[g.cell(k_above, j, i)] - field.p[g.cell(k_below, j, i)]) / (rho * dz); - let body_force = self.momentum_source.as_ref().map_or(0.0, |f| { - f((i as f64 + 0.5) * dx, (j as f64 + 0.5) * dy, zf, t_old).2 / rho - }); - - -conv_x - conv_y - conv_z + diff_x + diff_y + diff_z + pressure_gradient + body_force - } - /// The explicit predictor on the fluid faces; outlet faces zero-gradient. fn momentum_predictor(&self, field: &mut FlowField3D, dt: f64, t_old: f64) { let g = field.grid; diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/piso_predictor.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/piso_predictor.rs new file mode 100644 index 0000000..b3767f9 --- /dev/null +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/piso_predictor.rs @@ -0,0 +1,700 @@ +//! The three explicit predictors of the 3D PISO step (`Piso3Solver`): the +//! 2D `u_rhs`/`v_rhs` expression for expression with the z terms appended, +//! and `w_rhs` as the v pattern turned along z. Split from `piso_host.rs` +//! for the file-size rule; `impl Piso3Solver` continues here. + +use super::flow_field::FlowField3D; +use super::piso_host::{Piso3Solver, SideBoundary3}; +use crate::solvers::incompressible::simple::ConvectionScheme; + +impl Piso3Solver { + /// Plane above `k` (wrapping when periodic). + #[inline] + fn k_up(&self, k: usize, nz: usize) -> Option { + if k + 1 < nz { + Some(k + 1) + } else if self.params.boundaries.periodic_z() { + Some(0) + } else { + None + } + } + + #[inline] + fn k_down(&self, k: usize, nz: usize) -> Option { + if k > 0 { + Some(k - 1) + } else if self.params.boundaries.periodic_z() { + Some(nz - 1) + } else { + None + } + } + + /// The predictor's right-hand side on the u face `(k, j, i)`, `i = 1..nx`: + /// the 2D `u_rhs` expression for expression, then `− conv_z + diff_z`. + #[allow(clippy::too_many_lines)] + pub(crate) fn u_rhs( + &self, + field: &FlowField3D, + k: usize, + j: usize, + i: usize, + t_old: f64, + ) -> f64 { + let g = field.grid; + let (nx, ny, nz, dx, dy, dz) = (g.nx, g.ny, g.nz, g.dx, g.dy, g.dz); + let rho = self.fluid.density; + let nu = self.fluid.viscosity / rho; + let b = self.params.boundaries; + let velocity = SideBoundary3::Velocity; + let uo = &field.u_old; + let vo = &field.v_old; + let wo = &field.w_old; + let uf = |kk: usize, jj: usize, ii: usize| g.uface(kk, jj, ii); + let vf = |kk: usize, jj: usize, ii: usize| g.vface(kk, jj, ii); + let wf = |kk: usize, jj: usize, ii: usize| g.wface(kk, jj, ii); + let zc = (k as f64 + 0.5) * dz; + let u_p = uo[uf(k, j, i)]; + + let ue_face = 0.5 * (uo[uf(k, j, i)] + uo[uf(k, j, i + 1)]); + let uw_face = 0.5 * (uo[uf(k, j, i - 1)] + uo[uf(k, j, i)]); + + let south_is_wall = j == 0; + let north_is_wall = j + 1 == ny; + + let vn_face = 0.5 * (vo[vf(k, j + 1, i - 1)] + vo[vf(k, j + 1, i)]); + let vs_face = 0.5 * (vo[vf(k, j, i - 1)] + vo[vf(k, j, i)]); + + let beyond_north = if b.y1 == velocity { + self.boundary(i as f64 * dx, ny as f64 * dy, zc, t_old).0 + } else { + u_p + }; + let beyond_south = if b.y0 == velocity { + self.boundary(i as f64 * dx, 0.0, zc, t_old).0 + } else { + u_p + }; + + let conv_x = (ue_face * Self::upwind(ue_face, uo[uf(k, j, i)], uo[uf(k, j, i + 1)]) + - uw_face * Self::upwind(uw_face, uo[uf(k, j, i - 1)], uo[uf(k, j, i)])) + / dx; + let conv_y = (vn_face + * if north_is_wall { + Self::upwind(vn_face, u_p, beyond_north) + } else { + Self::upwind(vn_face, uo[uf(k, j, i)], uo[uf(k, j + 1, i)]) + } + - vs_face + * if south_is_wall { + Self::upwind(vs_face, beyond_south, u_p) + } else { + Self::upwind(vs_face, uo[uf(k, j - 1, i)], uo[uf(k, j, i)]) + }) + / dy; + + let scheme = self.params.convection_scheme; + let mut conv_x = conv_x; + let mut conv_y = conv_y; + if scheme != ConvectionScheme::Upwind { + let delta_e = if ue_face >= 0.0 { + scheme.face_correction( + Some(uo[uf(k, j, i - 1)]), + uo[uf(k, j, i)], + uo[uf(k, j, i + 1)], + ) + } else { + let far = (i + 2 <= nx).then(|| uo[uf(k, j, i + 2)]); + scheme.face_correction(far, uo[uf(k, j, i + 1)], uo[uf(k, j, i)]) + }; + let delta_w = if uw_face >= 0.0 { + let far = (i >= 2).then(|| uo[uf(k, j, i - 2)]); + scheme.face_correction(far, uo[uf(k, j, i - 1)], uo[uf(k, j, i)]) + } else { + scheme.face_correction( + Some(uo[uf(k, j, i + 1)]), + uo[uf(k, j, i)], + uo[uf(k, j, i - 1)], + ) + }; + let delta_n = if north_is_wall { + 0.0 + } else if vn_face >= 0.0 { + let far = (j >= 1).then(|| uo[uf(k, j - 1, i)]); + scheme.face_correction(far, uo[uf(k, j, i)], uo[uf(k, j + 1, i)]) + } else { + let far = (j + 2 < ny).then(|| uo[uf(k, j + 2, i)]); + scheme.face_correction(far, uo[uf(k, j + 1, i)], uo[uf(k, j, i)]) + }; + let delta_s = if south_is_wall { + 0.0 + } else if vs_face >= 0.0 { + let far = (j >= 2).then(|| uo[uf(k, j - 2, i)]); + scheme.face_correction(far, uo[uf(k, j - 1, i)], uo[uf(k, j, i)]) + } else { + let far = (j + 1 < ny).then(|| uo[uf(k, j + 1, i)]); + scheme.face_correction(far, uo[uf(k, j, i)], uo[uf(k, j - 1, i)]) + }; + conv_x += (ue_face * delta_e - uw_face * delta_w) / dx; + conv_y += (vn_face * delta_n - vs_face * delta_s) / dy; + } + + let diff_x = nu * (uo[uf(k, j, i + 1)] - 2.0 * u_p + uo[uf(k, j, i - 1)]) / (dx * dx); + + let flux_north = if north_is_wall { + if b.y1 == velocity { + let u_wall = self.boundary(i as f64 * dx, ny as f64 * dy, zc, t_old).0; + nu * (u_wall - u_p) / (0.5 * dy) + } else { + 0.0 + } + } else { + nu * (uo[uf(k, j + 1, i)] - u_p) / dy + }; + let flux_south = if south_is_wall { + if b.y0 == velocity { + let u_wall = self.boundary(i as f64 * dx, 0.0, zc, t_old).0; + nu * (u_p - u_wall) / (0.5 * dy) + } else { + 0.0 + } + } else { + nu * (u_p - uo[uf(k, j - 1, i)]) / dy + }; + let diff_y = (flux_north - flux_south) / dy; + + let pressure_gradient = + -(field.p[g.cell(k, j, i)] - field.p[g.cell(k, j, i - 1)]) / (rho * dx); + + let body_force = self.momentum_source.as_ref().map_or(0.0, |f| { + f(i as f64 * dx, (j as f64 + 0.5) * dy, zc, t_old).0 / rho + }); + + let rhs_2d = -conv_x - conv_y + diff_x + diff_y + pressure_gradient + body_force; + + // --- the z terms, the y pattern turned along k --- + let ku = self.k_up(k, nz); + let kd = self.k_down(k, nz); + let top_is_wall = ku.is_none(); + let bottom_is_wall = kd.is_none(); + // The w faces above/below the u face: on top of the cells west and + // east of it (face k + 1 of cell k is face index k + 1; periodic: + // the face at k = nz equals the face at 0). + let wt_face = 0.5 * (wo[wf(k + 1, j, i - 1)] + wo[wf(k + 1, j, i)]); + let wb_face = 0.5 * (wo[wf(k, j, i - 1)] + wo[wf(k, j, i)]); + let beyond_top = if b.z1 == velocity { + self.boundary(i as f64 * dx, (j as f64 + 0.5) * dy, nz as f64 * dz, t_old) + .0 + } else { + u_p + }; + let beyond_bottom = if b.z0 == velocity { + self.boundary(i as f64 * dx, (j as f64 + 0.5) * dy, 0.0, t_old) + .0 + } else { + u_p + }; + let u_up = ku.map(|kk| uo[uf(kk, j, i)]); + let u_dn = kd.map(|kk| uo[uf(kk, j, i)]); + let mut conv_z = (wt_face + * match u_up { + Some(un) => Self::upwind(wt_face, u_p, un), + None => Self::upwind(wt_face, u_p, beyond_top), + } + - wb_face + * match u_dn { + Some(ud) => Self::upwind(wb_face, ud, u_p), + None => Self::upwind(wb_face, beyond_bottom, u_p), + }) + / dz; + if scheme != ConvectionScheme::Upwind { + let far_up2 = ku + .and_then(|kk| self.k_up(kk, nz)) + .map(|kk| uo[uf(kk, j, i)]); + let far_dn2 = kd + .and_then(|kk| self.k_down(kk, nz)) + .map(|kk| uo[uf(kk, j, i)]); + let delta_t = if top_is_wall { + 0.0 + } else if wt_face >= 0.0 { + scheme.face_correction(u_dn, u_p, u_up.unwrap_or(u_p)) + } else { + scheme.face_correction(far_up2, u_up.unwrap_or(u_p), u_p) + }; + let delta_b = if bottom_is_wall { + 0.0 + } else if wb_face >= 0.0 { + scheme.face_correction(far_dn2, u_dn.unwrap_or(u_p), u_p) + } else { + scheme.face_correction(u_up, u_p, u_dn.unwrap_or(u_p)) + }; + conv_z += (wt_face * delta_t - wb_face * delta_b) / dz; + } + let flux_top = match u_up { + Some(un) => nu * (un - u_p) / dz, + None => { + if b.z1 == velocity { + nu * (beyond_top - u_p) / (0.5 * dz) + } else { + 0.0 + } + } + }; + let flux_bottom = match u_dn { + Some(ud) => nu * (u_p - ud) / dz, + None => { + if b.z0 == velocity { + nu * (u_p - beyond_bottom) / (0.5 * dz) + } else { + 0.0 + } + } + }; + let diff_z = (flux_top - flux_bottom) / dz; + + rhs_2d - conv_z + diff_z + } + + /// The v face `(k, j, i)`, `j = 1..ny`: the 2D `v_rhs` then the z terms. + #[allow(clippy::too_many_lines)] + pub(crate) fn v_rhs( + &self, + field: &FlowField3D, + k: usize, + j: usize, + i: usize, + t_old: f64, + ) -> f64 { + let g = field.grid; + let (nx, ny, nz, dx, dy, dz) = (g.nx, g.ny, g.nz, g.dx, g.dy, g.dz); + let rho = self.fluid.density; + let nu = self.fluid.viscosity / rho; + let b = self.params.boundaries; + let velocity = SideBoundary3::Velocity; + let uo = &field.u_old; + let vo = &field.v_old; + let wo = &field.w_old; + let uf = |kk: usize, jj: usize, ii: usize| g.uface(kk, jj, ii); + let vf = |kk: usize, jj: usize, ii: usize| g.vface(kk, jj, ii); + let wf = |kk: usize, jj: usize, ii: usize| g.wface(kk, jj, ii); + let zc = (k as f64 + 0.5) * dz; + let v_p = vo[vf(k, j, i)]; + + let vn_face = 0.5 * (vo[vf(k, j, i)] + vo[vf(k, j + 1, i)]); + let vs_face = 0.5 * (vo[vf(k, j - 1, i)] + vo[vf(k, j, i)]); + + let west_is_wall = i == 0; + let east_is_wall = i + 1 == nx; + + let ue_face = 0.5 * (uo[uf(k, j - 1, i + 1)] + uo[uf(k, j, i + 1)]); + let uw_face = 0.5 * (uo[uf(k, j - 1, i)] + uo[uf(k, j, i)]); + let beyond_east = if b.x1 == velocity { + self.boundary(nx as f64 * dx, j as f64 * dy, zc, t_old).1 + } else { + v_p + }; + let beyond_west = if b.x0 == velocity { + self.boundary(0.0, j as f64 * dy, zc, t_old).1 + } else { + v_p + }; + + let conv_y = (vn_face * Self::upwind(vn_face, vo[vf(k, j, i)], vo[vf(k, j + 1, i)]) + - vs_face * Self::upwind(vs_face, vo[vf(k, j - 1, i)], vo[vf(k, j, i)])) + / dy; + let conv_x = (ue_face + * if east_is_wall { + Self::upwind(ue_face, v_p, beyond_east) + } else { + Self::upwind(ue_face, vo[vf(k, j, i)], vo[vf(k, j, i + 1)]) + } + - uw_face + * if west_is_wall { + Self::upwind(uw_face, beyond_west, v_p) + } else { + Self::upwind(uw_face, vo[vf(k, j, i - 1)], vo[vf(k, j, i)]) + }) + / dx; + + let scheme = self.params.convection_scheme; + let mut conv_x = conv_x; + let mut conv_y = conv_y; + if scheme != ConvectionScheme::Upwind { + let delta_n = if vn_face >= 0.0 { + scheme.face_correction( + Some(vo[vf(k, j - 1, i)]), + vo[vf(k, j, i)], + vo[vf(k, j + 1, i)], + ) + } else { + let far = (j + 2 <= ny).then(|| vo[vf(k, j + 2, i)]); + scheme.face_correction(far, vo[vf(k, j + 1, i)], vo[vf(k, j, i)]) + }; + let delta_s = if vs_face >= 0.0 { + let far = (j >= 2).then(|| vo[vf(k, j - 2, i)]); + scheme.face_correction(far, vo[vf(k, j - 1, i)], vo[vf(k, j, i)]) + } else { + scheme.face_correction( + Some(vo[vf(k, j + 1, i)]), + vo[vf(k, j, i)], + vo[vf(k, j - 1, i)], + ) + }; + let delta_e = if east_is_wall { + 0.0 + } else if ue_face >= 0.0 { + let far = (i >= 1).then(|| vo[vf(k, j, i - 1)]); + scheme.face_correction(far, vo[vf(k, j, i)], vo[vf(k, j, i + 1)]) + } else { + let far = (i + 2 < nx).then(|| vo[vf(k, j, i + 2)]); + scheme.face_correction(far, vo[vf(k, j, i + 1)], vo[vf(k, j, i)]) + }; + let delta_w = if west_is_wall { + 0.0 + } else if uw_face >= 0.0 { + let far = (i >= 2).then(|| vo[vf(k, j, i - 2)]); + scheme.face_correction(far, vo[vf(k, j, i - 1)], vo[vf(k, j, i)]) + } else { + let far = (i + 1 < nx).then(|| vo[vf(k, j, i + 1)]); + scheme.face_correction(far, vo[vf(k, j, i)], vo[vf(k, j, i - 1)]) + }; + conv_y += (vn_face * delta_n - vs_face * delta_s) / dy; + conv_x += (ue_face * delta_e - uw_face * delta_w) / dx; + } + + let diff_y = nu * (vo[vf(k, j + 1, i)] - 2.0 * v_p + vo[vf(k, j - 1, i)]) / (dy * dy); + + let flux_east = if east_is_wall { + if b.x1 == velocity { + let v_wall = self.boundary(nx as f64 * dx, j as f64 * dy, zc, t_old).1; + nu * (v_wall - v_p) / (0.5 * dx) + } else { + 0.0 + } + } else { + nu * (vo[vf(k, j, i + 1)] - v_p) / dx + }; + let flux_west = if west_is_wall { + if b.x0 == velocity { + let v_wall = self.boundary(0.0, j as f64 * dy, zc, t_old).1; + nu * (v_p - v_wall) / (0.5 * dx) + } else { + 0.0 + } + } else { + nu * (v_p - vo[vf(k, j, i - 1)]) / dx + }; + let diff_x = (flux_east - flux_west) / dx; + + let pressure_gradient = + -(field.p[g.cell(k, j, i)] - field.p[g.cell(k, j - 1, i)]) / (rho * dy); + + let body_force = self.momentum_source.as_ref().map_or(0.0, |f| { + f((i as f64 + 0.5) * dx, j as f64 * dy, zc, t_old).1 / rho + }); + + let rhs_2d = -conv_x - conv_y + diff_x + diff_y + pressure_gradient + body_force; + + // --- z terms --- + let ku = self.k_up(k, nz); + let kd = self.k_down(k, nz); + let top_is_wall = ku.is_none(); + let bottom_is_wall = kd.is_none(); + let wt_face = 0.5 * (wo[wf(k + 1, j - 1, i)] + wo[wf(k + 1, j, i)]); + let wb_face = 0.5 * (wo[wf(k, j - 1, i)] + wo[wf(k, j, i)]); + let beyond_top = if b.z1 == velocity { + self.boundary((i as f64 + 0.5) * dx, j as f64 * dy, nz as f64 * dz, t_old) + .1 + } else { + v_p + }; + let beyond_bottom = if b.z0 == velocity { + self.boundary((i as f64 + 0.5) * dx, j as f64 * dy, 0.0, t_old) + .1 + } else { + v_p + }; + let v_up = ku.map(|kk| vo[vf(kk, j, i)]); + let v_dn = kd.map(|kk| vo[vf(kk, j, i)]); + let mut conv_z = (wt_face + * match v_up { + Some(vn) => Self::upwind(wt_face, v_p, vn), + None => Self::upwind(wt_face, v_p, beyond_top), + } + - wb_face + * match v_dn { + Some(vd) => Self::upwind(wb_face, vd, v_p), + None => Self::upwind(wb_face, beyond_bottom, v_p), + }) + / dz; + if scheme != ConvectionScheme::Upwind { + let far_up2 = ku + .and_then(|kk| self.k_up(kk, nz)) + .map(|kk| vo[vf(kk, j, i)]); + let far_dn2 = kd + .and_then(|kk| self.k_down(kk, nz)) + .map(|kk| vo[vf(kk, j, i)]); + let delta_t = if top_is_wall { + 0.0 + } else if wt_face >= 0.0 { + scheme.face_correction(v_dn, v_p, v_up.unwrap_or(v_p)) + } else { + scheme.face_correction(far_up2, v_up.unwrap_or(v_p), v_p) + }; + let delta_b = if bottom_is_wall { + 0.0 + } else if wb_face >= 0.0 { + scheme.face_correction(far_dn2, v_dn.unwrap_or(v_p), v_p) + } else { + scheme.face_correction(v_up, v_p, v_dn.unwrap_or(v_p)) + }; + conv_z += (wt_face * delta_t - wb_face * delta_b) / dz; + } + let flux_top = match v_up { + Some(vn) => nu * (vn - v_p) / dz, + None => { + if b.z1 == velocity { + nu * (beyond_top - v_p) / (0.5 * dz) + } else { + 0.0 + } + } + }; + let flux_bottom = match v_dn { + Some(vd) => nu * (v_p - vd) / dz, + None => { + if b.z0 == velocity { + nu * (v_p - beyond_bottom) / (0.5 * dz) + } else { + 0.0 + } + } + }; + let diff_z = (flux_top - flux_bottom) / dz; + + rhs_2d - conv_z + diff_z + } + + /// The w face `(k, j, i)` between cells `k − 1` (wrapping when periodic) + /// and `k`: the v pattern with z as its own direction and x, y transverse. + #[allow(clippy::too_many_lines)] + pub(crate) fn w_rhs( + &self, + field: &FlowField3D, + k: usize, + j: usize, + i: usize, + t_old: f64, + ) -> f64 { + let g = field.grid; + let (nx, ny, nz, dx, dy, dz) = (g.nx, g.ny, g.nz, g.dx, g.dy, g.dz); + let rho = self.fluid.density; + let nu = self.fluid.viscosity / rho; + let b = self.params.boundaries; + let velocity = SideBoundary3::Velocity; + let periodic = b.periodic_z(); + let uo = &field.u_old; + let vo = &field.v_old; + let wo = &field.w_old; + let uf = |kk: usize, jj: usize, ii: usize| g.uface(kk, jj, ii); + let vf = |kk: usize, jj: usize, ii: usize| g.vface(kk, jj, ii); + let wf = |kk: usize, jj: usize, ii: usize| g.wface(kk, jj, ii); + // The cells below and above this face (k = 0 only when periodic). + let k_below = if k > 0 { k - 1 } else { nz - 1 }; + let k_above = k % nz; + // Own-direction neighbours: the faces k − 1 and k + 1 (wrapping). + let w_dn_idx = if k > 0 { + wf(k - 1, j, i) + } else { + wf(nz - 1, j, i) + }; + let w_up_idx = if k + 1 == nz && periodic { + wf(0, j, i) + } else { + wf(k + 1, j, i) + }; + let zf = k as f64 * dz; + let w_p = wo[wf(k, j, i)]; + + let wt_face = 0.5 * (wo[wf(k, j, i)] + wo[w_up_idx]); + let wb_face = 0.5 * (wo[w_dn_idx] + wo[wf(k, j, i)]); + + let west_is_wall = i == 0; + let east_is_wall = i + 1 == nx; + let south_is_wall = j == 0; + let north_is_wall = j + 1 == ny; + + let ue_face = 0.5 * (uo[uf(k_below, j, i + 1)] + uo[uf(k_above, j, i + 1)]); + let uw_face = 0.5 * (uo[uf(k_below, j, i)] + uo[uf(k_above, j, i)]); + let vn_face = 0.5 * (vo[vf(k_below, j + 1, i)] + vo[vf(k_above, j + 1, i)]); + let vs_face = 0.5 * (vo[vf(k_below, j, i)] + vo[vf(k_above, j, i)]); + let beyond_east = if b.x1 == velocity { + self.boundary(nx as f64 * dx, (j as f64 + 0.5) * dy, zf, t_old) + .2 + } else { + w_p + }; + let beyond_west = if b.x0 == velocity { + self.boundary(0.0, (j as f64 + 0.5) * dy, zf, t_old).2 + } else { + w_p + }; + let beyond_north = if b.y1 == velocity { + self.boundary((i as f64 + 0.5) * dx, ny as f64 * dy, zf, t_old) + .2 + } else { + w_p + }; + let beyond_south = if b.y0 == velocity { + self.boundary((i as f64 + 0.5) * dx, 0.0, zf, t_old).2 + } else { + w_p + }; + + let conv_z = (wt_face * Self::upwind(wt_face, w_p, wo[w_up_idx]) + - wb_face * Self::upwind(wb_face, wo[w_dn_idx], w_p)) + / dz; + let conv_x = (ue_face + * if east_is_wall { + Self::upwind(ue_face, w_p, beyond_east) + } else { + Self::upwind(ue_face, w_p, wo[wf(k, j, i + 1)]) + } + - uw_face + * if west_is_wall { + Self::upwind(uw_face, beyond_west, w_p) + } else { + Self::upwind(uw_face, wo[wf(k, j, i - 1)], w_p) + }) + / dx; + let conv_y = (vn_face + * if north_is_wall { + Self::upwind(vn_face, w_p, beyond_north) + } else { + Self::upwind(vn_face, w_p, wo[wf(k, j + 1, i)]) + } + - vs_face + * if south_is_wall { + Self::upwind(vs_face, beyond_south, w_p) + } else { + Self::upwind(vs_face, wo[wf(k, j - 1, i)], w_p) + }) + / dy; + + let scheme = self.params.convection_scheme; + let mut conv_x = conv_x; + let mut conv_y = conv_y; + let mut conv_z = conv_z; + if scheme != ConvectionScheme::Upwind { + // Own direction: far nodes two faces away (wrapping when periodic). + let far_up2 = if periodic { + Some(wo[wf((k + 2) % nz, j, i)]) + } else { + (k + 2 <= nz).then(|| wo[wf(k + 2, j, i)]) + }; + let far_dn2 = if periodic { + Some(wo[wf((k + nz - 2) % nz, j, i)]) + } else { + (k >= 2).then(|| wo[wf(k - 2, j, i)]) + }; + let delta_t = if wt_face >= 0.0 { + scheme.face_correction(Some(wo[w_dn_idx]), w_p, wo[w_up_idx]) + } else { + scheme.face_correction(far_up2, wo[w_up_idx], w_p) + }; + let delta_b = if wb_face >= 0.0 { + scheme.face_correction(far_dn2, wo[w_dn_idx], w_p) + } else { + scheme.face_correction(Some(wo[w_up_idx]), w_p, wo[w_dn_idx]) + }; + let delta_e = if east_is_wall { + 0.0 + } else if ue_face >= 0.0 { + let far = (i >= 1).then(|| wo[wf(k, j, i - 1)]); + scheme.face_correction(far, w_p, wo[wf(k, j, i + 1)]) + } else { + let far = (i + 2 < nx).then(|| wo[wf(k, j, i + 2)]); + scheme.face_correction(far, wo[wf(k, j, i + 1)], w_p) + }; + let delta_w = if west_is_wall { + 0.0 + } else if uw_face >= 0.0 { + let far = (i >= 2).then(|| wo[wf(k, j, i - 2)]); + scheme.face_correction(far, wo[wf(k, j, i - 1)], w_p) + } else { + let far = (i + 1 < nx).then(|| wo[wf(k, j, i + 1)]); + scheme.face_correction(far, w_p, wo[wf(k, j, i - 1)]) + }; + let delta_n = if north_is_wall { + 0.0 + } else if vn_face >= 0.0 { + let far = (j >= 1).then(|| wo[wf(k, j - 1, i)]); + scheme.face_correction(far, w_p, wo[wf(k, j + 1, i)]) + } else { + let far = (j + 2 < ny).then(|| wo[wf(k, j + 2, i)]); + scheme.face_correction(far, wo[wf(k, j + 1, i)], w_p) + }; + let delta_s = if south_is_wall { + 0.0 + } else if vs_face >= 0.0 { + let far = (j >= 2).then(|| wo[wf(k, j - 2, i)]); + scheme.face_correction(far, wo[wf(k, j - 1, i)], w_p) + } else { + let far = (j + 1 < ny).then(|| wo[wf(k, j + 1, i)]); + scheme.face_correction(far, w_p, wo[wf(k, j - 1, i)]) + }; + conv_z += (wt_face * delta_t - wb_face * delta_b) / dz; + conv_x += (ue_face * delta_e - uw_face * delta_w) / dx; + conv_y += (vn_face * delta_n - vs_face * delta_s) / dy; + } + + let diff_z = nu * (wo[w_up_idx] - 2.0 * w_p + wo[w_dn_idx]) / (dz * dz); + let flux_east = if east_is_wall { + if b.x1 == velocity { + nu * (beyond_east - w_p) / (0.5 * dx) + } else { + 0.0 + } + } else { + nu * (wo[wf(k, j, i + 1)] - w_p) / dx + }; + let flux_west = if west_is_wall { + if b.x0 == velocity { + nu * (w_p - beyond_west) / (0.5 * dx) + } else { + 0.0 + } + } else { + nu * (w_p - wo[wf(k, j, i - 1)]) / dx + }; + let diff_x = (flux_east - flux_west) / dx; + let flux_north = if north_is_wall { + if b.y1 == velocity { + nu * (beyond_north - w_p) / (0.5 * dy) + } else { + 0.0 + } + } else { + nu * (wo[wf(k, j + 1, i)] - w_p) / dy + }; + let flux_south = if south_is_wall { + if b.y0 == velocity { + nu * (w_p - beyond_south) / (0.5 * dy) + } else { + 0.0 + } + } else { + nu * (w_p - wo[wf(k, j - 1, i)]) / dy + }; + let diff_y = (flux_north - flux_south) / dy; + + let pressure_gradient = + -(field.p[g.cell(k_above, j, i)] - field.p[g.cell(k_below, j, i)]) / (rho * dz); + let body_force = self.momentum_source.as_ref().map_or(0.0, |f| { + f((i as f64 + 0.5) * dx, (j as f64 + 0.5) * dy, zf, t_old).2 / rho + }); + + -conv_x - conv_y - conv_z + diff_x + diff_y + diff_z + pressure_gradient + body_force + } +}