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

//! The uniform Cartesian grid and its index conventions.

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Grid {
    pub nx: usize,
    pub ny: usize,
    pub nz: usize,
    pub dx: f64,
    pub dy: f64,
    pub dz: f64,
}

impl Grid {
    /// Cubic cells of spacing `h`.
    #[must_use]
    pub fn cubic(nx: usize, ny: usize, nz: usize, h: f64) -> Self {
        Self {
            nx,
            ny,
            nz,
            dx: h,
            dy: h,
            dz: h,
        }
    }

    #[inline]
    #[must_use]
    pub fn cells(&self) -> usize {
        self.nx * self.ny * self.nz
    }

    #[inline]
    #[must_use]
    pub fn cell(&self, k: usize, j: usize, i: usize) -> usize {
        (k * self.ny + j) * self.nx + i
    }

    /// `(k, j, i)` of a cell index.
    #[inline]
    #[must_use]
    pub fn kji(&self, idx: usize) -> (usize, usize, usize) {
        let nxy = self.nx * self.ny;
        (idx / nxy, (idx % nxy) / self.nx, idx % self.nx)
    }

    /// The u face west of cell `(k, j, i)`; `i = nx` is the east face of the last cell.
    #[inline]
    #[must_use]
    pub fn uface(&self, k: usize, j: usize, i: usize) -> usize {
        (k * self.ny + j) * (self.nx + 1) + i
    }

    #[inline]
    #[must_use]
    pub fn vface(&self, k: usize, j: usize, i: usize) -> usize {
        (k * (self.ny + 1) + j) * self.nx + i
    }

    #[inline]
    #[must_use]
    pub fn wface(&self, k: usize, j: usize, i: usize) -> usize {
        (k * self.ny + j) * self.nx + i
    }

    #[inline]
    #[must_use]
    pub fn n_ufaces(&self) -> usize {
        (self.nx + 1) * self.ny * self.nz
    }

    #[inline]
    #[must_use]
    pub fn n_vfaces(&self) -> usize {
        self.nx * (self.ny + 1) * self.nz
    }

    #[inline]
    #[must_use]
    pub fn n_wfaces(&self) -> usize {
        self.nx * self.ny * (self.nz + 1)
    }
}