rustytorch/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/impose.rs

//! The wall's imposition on the velocity field (`impl Mask` continued
//! from `wall.rs`, split for the file-size rule): prescribed faces take
//! the surface velocity, ghost faces their reconstruction from the source
//! field minus the shared flux compatibility correction.

use super::body::Body;
use super::wall::{FaceKind, Mask};

impl Mask {
    /// Impose the wall on `(u, v, w)` from the same field.
    pub fn impose(&self, body: &Body, u: &mut [f64], v: &mut [f64], w: &mut [f64], t: f64) -> f64 {
        let (us, vs, ws) = (u.to_vec(), v.to_vec(), w.to_vec());
        self.impose_from(body, &us, &vs, &ws, u, v, w, t)
    }

    /// Solid faces: the surface velocity; ghost faces: the reconstruction
    /// from the SOURCE field, minus the shared flux compatibility
    /// correction over the flux-carrying ghosts. Returns the correction.
    #[allow(clippy::too_many_arguments)]
    pub fn impose_from(
        &self,
        body: &Body,
        u_src: &[f64],
        v_src: &[f64],
        w_src: &[f64],
        u: &mut [f64],
        v: &mut [f64],
        w: &mut [f64],
        t: f64,
    ) -> f64 {
        let g = self.grid;
        let (nx, ny, nz, dx, dy, dz) = (g.nx, g.ny, g.nz, g.dx, g.dy, g.dz);
        for k in 0..nz {
            for j in 0..ny {
                for i in 1..nx {
                    let idx = g.uface(k, j, i);
                    if self.u_kind[idx] == FaceKind::Solid {
                        u[idx] = body
                            .surface_velocity(
                                i as f64 * dx,
                                (j as f64 + 0.5) * dy,
                                (k as f64 + 0.5) * dz,
                                t,
                            )
                            .0;
                    }
                }
            }
            for j in 1..ny {
                for i in 0..nx {
                    let idx = g.vface(k, j, i);
                    if self.v_kind[idx] == FaceKind::Solid {
                        v[idx] = body
                            .surface_velocity(
                                (i as f64 + 0.5) * dx,
                                j as f64 * dy,
                                (k as f64 + 0.5) * dz,
                                t,
                            )
                            .1;
                    }
                }
            }
        }
        for k in 0..=nz {
            for j in 0..ny {
                for i in 0..nx {
                    let idx = g.wface(k, j, i);
                    if self.w_kind[idx] == FaceKind::Solid {
                        w[idx] = body
                            .surface_velocity(
                                (i as f64 + 0.5) * dx,
                                (j as f64 + 0.5) * dy,
                                k as f64 * dz,
                                t,
                            )
                            .2;
                    }
                }
            }
        }
        let u_vals: Vec<f64> = self
            .u_ghosts
            .iter()
            .map(|gh| gh.reconstruct(u_src))
            .collect();
        let v_vals: Vec<f64> = self
            .v_ghosts
            .iter()
            .map(|gh| gh.reconstruct(v_src))
            .collect();
        let w_vals: Vec<f64> = self
            .w_ghosts
            .iter()
            .map(|gh| gh.reconstruct(w_src))
            .collect();
        let (au, av, aw) = (dy * dz, dx * dz, dx * dy);
        let mut net = 0.0;
        let mut area = 0.0;
        for (gh, &val) in self.u_ghosts.iter().zip(&u_vals) {
            if gh.flux_sign != 0.0 {
                net += gh.flux_sign * val * au;
                area += au;
            }
        }
        for (gh, &val) in self.v_ghosts.iter().zip(&v_vals) {
            if gh.flux_sign != 0.0 {
                net += gh.flux_sign * val * av;
                area += av;
            }
        }
        for (gh, &val) in self.w_ghosts.iter().zip(&w_vals) {
            if gh.flux_sign != 0.0 {
                net += gh.flux_sign * val * aw;
                area += aw;
            }
        }
        let correction = if area > 0.0 { net / area } else { 0.0 };
        for (gh, &val) in self.u_ghosts.iter().zip(&u_vals) {
            u[gh.idx] = val - gh.flux_sign * correction;
        }
        for (gh, &val) in self.v_ghosts.iter().zip(&v_vals) {
            v[gh.idx] = val - gh.flux_sign * correction;
        }
        for (gh, &val) in self.w_ghosts.iter().zip(&w_vals) {
            w[gh.idx] = val - gh.flux_sign * correction;
        }
        // The periodic seam: the w face at k = nz is the face at k = 0.
        for j in 0..ny {
            for i in 0..nx {
                let (f0, fn_) = (g.wface(0, j, i), g.wface(nz, j, i));
                if self.w_kind[f0] != FaceKind::Fluid && self.w_kind[fn_] == self.w_kind[f0] {
                    w[fn_] = w[f0];
                }
            }
        }
        correction
    }
}