rustyverse/rustytorch
2
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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
CI / Test (macos-latest) (push) Blocked by required conditions
Details
CI / Test (ubuntu-latest) (push) Blocked by required conditions
Details
CI / Python Bindings (maturin) (macos-latest) (push) Blocked by required conditions
Details
CI / Build (macos-latest) (push) Waiting to run
Details
CI / Python Bindings (maturin) (ubuntu-latest) (push) Blocked by required conditions
Details
CI / WASM Build + Size Check (push) Blocked by required conditions
Details
CI / Distributed Training Tests (push) Blocked by required conditions
Details
CI / CI Success (push) Blocked by required conditions
Details
Documentation / Build User Guide (push) Successful in 6s
Details
CI / Build (ubuntu-latest) (push) Failing after 3s
Details
CI / Format Check (push) Failing after 4s
Details
CI / Clippy Check (push) Failing after 4s
Details
Performance Benchmarks / Run Benchmarks (push) Failing after 15s
Details
CI / Build CPU-Only (Explicit) (push) Failing after 1m30s
Details
Documentation / Build API Documentation (push) Failing after 1m33s
Details

Browse Source 
Co-Authored-By: Claude Fable 5.1 <[email protected]>
This commit is contained in:
Omar Sobh
2026-09-17 11:33:13 -05:00
co-authored by Claude Fable 5.1
parent 616d2a3394
commit e0cbb99343
3 changed files with 705 additions and 700 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/mod.rs
+1
View File
@@ -9,6 +9,7 @@
pub mod flow_field;
pub mod piso_host;
mod piso_predictor;
pub mod poisson;
pub use flow_field::FlowField3D;
crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/piso_host.rs
+4 -700
View File
@@ -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<dyn Fn(f64, f64, f64, f64) -> (f64, f64, f64) + Send + Sync>;
pub struct Piso3Solver {
pub fluid: Fluid3,
pub params: Piso3Parameters,
momentum_source: Option<Vec3Fn>,
pub(super) momentum_source: Option<Vec3Fn>,
boundary_velocity: Option<Vec3Fn>,
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<usize> {
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<usize> {
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;
crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/piso_predictor.rs
+700
View File
@@ -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<usize> {
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<usize> {
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
}
}