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

//! Flow field data structures and operations
//!
//! This module defines the core data structures for storing and manipulating
//! flow field variables (velocity, pressure) on structured grids.

use crate::{CfdError, CfdResult};
use nalgebra::DMatrix;

/// Flow field containing all flow variables on a structured grid
#[derive(Debug, Clone)]
pub struct FlowField {
    /// Grid dimensions
    pub nx: usize,
    pub ny: usize,

    /// Grid spacing
    pub dx: f64,
    pub dy: f64,

    /// x-component of velocity (u) - located at cell faces (i+1/2, j)
    pub u: DMatrix<f64>,
    /// y-component of velocity (v) - located at cell faces (i, j+1/2)
    pub v: DMatrix<f64>,
    /// Pressure (p) - located at cell centers (i, j)
    pub p: DMatrix<f64>,

    /// Previous time step values for time integration
    pub u_old: DMatrix<f64>,
    pub v_old: DMatrix<f64>,
    pub p_old: DMatrix<f64>,

    /// Auxiliary fields for solver algorithms
    pub u_star: DMatrix<f64>, // Predicted velocity (SIMPLE/PISO)
    pub v_star: DMatrix<f64>,  // Predicted velocity (SIMPLE/PISO)
    pub p_prime: DMatrix<f64>, // Pressure correction (SIMPLE/PISO)

    /// Source terms
    pub su: DMatrix<f64>, // u-momentum source
    pub sv: DMatrix<f64>, // v-momentum source
    pub sp: DMatrix<f64>, // Pressure source (mass source)
}

impl FlowField {
    /// Create new flow field with given dimensions
    pub fn new(nx: usize, ny: usize, dx: f64, dy: f64) -> CfdResult<Self> {
        if nx < 3 || ny < 3 {
            return Err(CfdError::invalid_parameter("Grid must be at least 3x3"));
        }

        if dx <= 0.0 || dy <= 0.0 {
            return Err(CfdError::invalid_parameter("Grid spacing must be positive"));
        }

        // For staggered grid:
        // u: (nx+1, ny) - face-centered in x-direction
        // v: (nx, ny+1) - face-centered in y-direction
        // p: (nx, ny) - cell-centered

        let zeros_u = DMatrix::zeros(ny, nx + 1); // Note: nalgebra is (rows, cols)
        let zeros_v = DMatrix::zeros(ny + 1, nx);
        let zeros_p = DMatrix::zeros(ny, nx);

        Ok(Self {
            nx,
            ny,
            dx,
            dy,

            u: zeros_u.clone(),
            v: zeros_v.clone(),
            p: zeros_p.clone(),

            u_old: zeros_u.clone(),
            v_old: zeros_v.clone(),
            p_old: zeros_p.clone(),

            u_star: zeros_u.clone(),
            v_star: zeros_v.clone(),
            p_prime: zeros_p.clone(),

            su: zeros_u,
            sv: zeros_v,
            sp: zeros_p,
        })
    }

    /// Set velocity at a given grid point
    pub fn set_velocity(&mut self, i: usize, j: usize, u_val: f64, v_val: f64) -> CfdResult<()> {
        if i >= self.nx || j >= self.ny {
            return Err(CfdError::invalid_parameter("Grid indices out of bounds"));
        }

        // For staggered grid, velocity components are at different locations
        // u is stored at (j, i) for face (i+1/2, j)
        if i < self.nx {
            self.u[(j, i)] = u_val;
        }

        // v is stored at (j, i) for face (i, j+1/2)
        if j < self.ny {
            self.v[(j, i)] = v_val;
        }

        Ok(())
    }

    /// Get velocity at a given grid point (interpolated to cell center)
    pub fn get_velocity_at(&self, i: usize, j: usize) -> CfdResult<(f64, f64)> {
        if i >= self.nx || j >= self.ny {
            return Err(CfdError::invalid_parameter("Grid indices out of bounds"));
        }

        // Interpolate to the cell centre from the cell's own two faces.
        //
        // On this staggered layout `u` is `(ny, nx + 1)` and `v` is
        // `(ny + 1, nx)`, and cell `i` is bounded by u-faces `i` and `i + 1` —
        // which is the convention `compute_mass_source` uses to form the
        // divergence, and therefore the one that defines the grid.
        //
        // This previously averaged faces `i - 1` and `i`, half a cell to the
        // west of the cell it claimed to be reporting, with the first and last
        // cells special-cased to a single face and the outermost face never
        // read at all. Every profile taken through this function was shifted
        // by half a cell against the field the solver actually computed.
        let u_center = 0.5 * (self.u[(j, i)] + self.u[(j, i + 1)]);
        let v_center = 0.5 * (self.v[(j, i)] + self.v[(j + 1, i)]);

        Ok((u_center, v_center))
    }

    /// Set pressure at a given grid point
    pub fn set_pressure(&mut self, i: usize, j: usize, p_val: f64) -> CfdResult<()> {
        if i >= self.nx || j >= self.ny {
            return Err(CfdError::invalid_parameter("Grid indices out of bounds"));
        }

        self.p[(j, i)] = p_val;
        Ok(())
    }

    /// Get pressure at a given grid point
    pub fn get_pressure_at(&self, i: usize, j: usize) -> CfdResult<f64> {
        if i >= self.nx || j >= self.ny {
            return Err(CfdError::invalid_parameter("Grid indices out of bounds"));
        }

        Ok(self.p[(j, i)])
    }

    /// Apply boundary conditions to the flow field
    pub fn apply_boundary_conditions(&mut self, bcs: &super::BoundaryConditions) -> CfdResult<()> {
        bcs.apply_to_flow_field(self)
    }

    /// Compute divergence of velocity field (mass conservation check)
    pub fn compute_divergence(&self) -> CfdResult<DMatrix<f64>> {
        let mut divergence = DMatrix::zeros(self.ny, self.nx);

        for j in 0..self.ny {
            for i in 0..self.nx {
                // ∇·u = ∂u/∂x + ∂v/∂y
                let du_dx = if i == self.nx - 1 {
                    (self.u[(j, i)] - self.u[(j, i - 1)]) / self.dx
                } else {
                    (self.u[(j, i + 1)] - self.u[(j, i)]) / self.dx
                };

                let dv_dy = if j == self.ny - 1 {
                    (self.v[(j, i)] - self.v[(j - 1, i)]) / self.dy
                } else {
                    (self.v[(j + 1, i)] - self.v[(j, i)]) / self.dy
                };

                divergence[(j, i)] = du_dx + dv_dy;
            }
        }

        Ok(divergence)
    }

    /// Compute maximum divergence (for mass conservation check)
    pub fn compute_max_divergence(&self) -> CfdResult<f64> {
        let divergence = self.compute_divergence()?;
        Ok(divergence.iter().map(|&x| x.abs()).fold(0.0, f64::max))
    }

    /// Find maximum u-velocity and its location
    pub fn find_max_u_velocity(&self) -> CfdResult<(f64, (usize, usize))> {
        let mut max_u = f64::NEG_INFINITY;
        let mut max_loc = (0, 0);

        for j in 0..self.ny {
            for i in 0..=self.nx {
                if self.u[(j, i)] > max_u {
                    max_u = self.u[(j, i)];
                    max_loc = (i, j);
                }
            }
        }

        Ok((max_u, max_loc))
    }

    /// Compute total kinetic energy
    pub fn compute_kinetic_energy(&self) -> CfdResult<f64> {
        let mut ke = 0.0;

        for j in 0..self.ny {
            for i in 0..self.nx {
                let (u_center, v_center) = self.get_velocity_at(i, j)?;
                ke += 0.5 * (u_center * u_center + v_center * v_center) * self.dx * self.dy;
            }
        }

        Ok(ke)
    }

    /// Update old values (for time stepping)
    pub fn update_old_values(&mut self) {
        self.u_old.copy_from(&self.u);
        self.v_old.copy_from(&self.v);
        self.p_old.copy_from(&self.p);
    }

    /// Copy current values to starred values (for predictor step)
    pub fn copy_to_starred(&mut self) {
        self.u_star.copy_from(&self.u);
        self.v_star.copy_from(&self.v);
    }

    /// Apply under-relaxation to velocity field
    pub fn apply_velocity_relaxation(&mut self, relaxation_factor: f64) -> CfdResult<()> {
        if relaxation_factor <= 0.0 || relaxation_factor > 1.0 {
            return Err(CfdError::invalid_parameter(
                "Relaxation factor must be in (0, 1]",
            ));
        }

        // u = α * u_new + (1 - α) * u_old
        for j in 0..self.ny {
            for i in 0..=self.nx {
                self.u[(j, i)] = relaxation_factor * self.u[(j, i)]
                    + (1.0 - relaxation_factor) * self.u_old[(j, i)];
            }
        }

        for j in 0..=self.ny {
            for i in 0..self.nx {
                self.v[(j, i)] = relaxation_factor * self.v[(j, i)]
                    + (1.0 - relaxation_factor) * self.v_old[(j, i)];
            }
        }

        Ok(())
    }

    /// Apply under-relaxation to pressure field
    pub fn apply_pressure_relaxation(&mut self, relaxation_factor: f64) -> CfdResult<()> {
        if relaxation_factor <= 0.0 || relaxation_factor > 1.0 {
            return Err(CfdError::invalid_parameter(
                "Relaxation factor must be in (0, 1]",
            ));
        }

        for j in 0..self.ny {
            for i in 0..self.nx {
                self.p[(j, i)] = relaxation_factor * self.p[(j, i)]
                    + (1.0 - relaxation_factor) * self.p_old[(j, i)];
            }
        }

        Ok(())
    }

    /// Compute L2 norm of residual
    #[must_use]
    pub fn compute_velocity_residual(&self) -> f64 {
        let mut residual = 0.0;

        // u-momentum residual
        for j in 0..self.ny {
            for i in 0..=self.nx {
                let diff = self.u[(j, i)] - self.u_old[(j, i)];
                residual += diff * diff;
            }
        }

        // v-momentum residual
        for j in 0..=self.ny {
            for i in 0..self.nx {
                let diff = self.v[(j, i)] - self.v_old[(j, i)];
                residual += diff * diff;
            }
        }

        residual.sqrt()
    }

    /// Compute pressure residual
    #[must_use]
    pub fn compute_pressure_residual(&self) -> f64 {
        let mut residual = 0.0;

        for j in 0..self.ny {
            for i in 0..self.nx {
                let diff = self.p[(j, i)] - self.p_old[(j, i)];
                residual += diff * diff;
            }
        }

        residual.sqrt()
    }

    /// Get grid information
    #[must_use]
    pub fn grid_info(&self) -> (usize, usize, f64, f64) {
        (self.nx, self.ny, self.dx, self.dy)
    }

    /// Initialize with analytical solution (for testing)
    pub fn initialize_with_analytical(
        &mut self,
        solution_type: AnalyticalSolution,
    ) -> CfdResult<()> {
        match solution_type {
            AnalyticalSolution::PoiseuillePlane { u_max } => {
                // Plane Poiseuille flow: u(y) = u_max * 4 * y * (1-y)
                for j in 0..self.ny {
                    let y = (j as f64 + 0.5) * self.dy; // Cell center y-coordinate
                    let y_normalized = y / (self.ny as f64 * self.dy);
                    let u_analytical = u_max * 4.0 * y_normalized * (1.0 - y_normalized);

                    for i in 0..=self.nx {
                        self.u[(j, i)] = u_analytical;
                    }
                }

                // v = 0 everywhere
                self.v.fill(0.0);

                // Pressure gradient to drive the flow
                for j in 0..self.ny {
                    for i in 0..self.nx {
                        self.p[(j, i)] = -(i as f64) * self.dx; // Linear pressure drop
                    }
                }
            }
            AnalyticalSolution::TaylorGreenVortex { amplitude } => {
                // Taylor-Green vortex: analytical solution for 2D Navier-Stokes
                for j in 0..self.ny {
                    for i in 0..=self.nx {
                        let x = i as f64 * self.dx;
                        let y = (j as f64 + 0.5) * self.dy;
                        self.u[(j, i)] = amplitude
                            * (2.0 * std::f64::consts::PI * x).sin()
                            * (2.0 * std::f64::consts::PI * y).cos();
                    }
                }

                for j in 0..=self.ny {
                    for i in 0..self.nx {
                        let x = (i as f64 + 0.5) * self.dx;
                        let y = j as f64 * self.dy;
                        self.v[(j, i)] = -amplitude
                            * (2.0 * std::f64::consts::PI * x).cos()
                            * (2.0 * std::f64::consts::PI * y).sin();
                    }
                }

                // Pressure field for Taylor-Green vortex
                for j in 0..self.ny {
                    for i in 0..self.nx {
                        let x = (i as f64 + 0.5) * self.dx;
                        let y = (j as f64 + 0.5) * self.dy;
                        self.p[(j, i)] = -0.25
                            * amplitude
                            * amplitude
                            * ((4.0 * std::f64::consts::PI * x).cos()
                                + (4.0 * std::f64::consts::PI * y).cos());
                    }
                }
            }
        }

        Ok(())
    }
}

impl FlowField {
    const SAVE_MAGIC: [u8; 4] = *b"RTXF";
    const SAVE_VERSION: u32 = 1;

    /// Serialize the complete field state to a file, bit-exact.
    ///
    /// Every matrix is written (including `*_old`, the starred
    /// predictors and the sources), so a [`Self::load`] of the file is
    /// a true restart state, not a view: a solver resumed from it sees
    /// exactly the arrays the saved solver held. Layout: magic `RTXF`,
    /// version, `nx`/`ny` (u64 LE), `dx`/`dy` (f64 LE), then each
    /// matrix as `nrows`/`ncols` (u64 LE) + column-major f64 LE data,
    /// in declaration order.
    pub fn save(&self, path: &std::path::Path) -> CfdResult<()> {
        use std::io::Write as _;
        let file = std::fs::File::create(path)
            .map_err(|e| CfdError::invalid_parameter(format!("save {}: {e}", path.display())))?;
        let mut w = std::io::BufWriter::new(file);
        let mut write = |bytes: &[u8]| -> CfdResult<()> {
            w.write_all(bytes)
                .map_err(|e| CfdError::invalid_parameter(format!("save write: {e}")))
        };
        write(&Self::SAVE_MAGIC)?;
        write(&Self::SAVE_VERSION.to_le_bytes())?;
        write(&(self.nx as u64).to_le_bytes())?;
        write(&(self.ny as u64).to_le_bytes())?;
        write(&self.dx.to_le_bytes())?;
        write(&self.dy.to_le_bytes())?;
        for m in self.matrices() {
            write(&(m.nrows() as u64).to_le_bytes())?;
            write(&(m.ncols() as u64).to_le_bytes())?;
            for v in m.iter() {
                write(&v.to_le_bytes())?;
            }
        }
        w.flush()
            .map_err(|e| CfdError::invalid_parameter(format!("save flush: {e}")))
    }

    /// Deserialize a field saved by [`Self::save`], validating magic,
    /// version and every matrix shape against a fresh field of the
    /// stored dimensions.
    pub fn load(path: &std::path::Path) -> CfdResult<Self> {
        use std::io::Read as _;
        let mut data = Vec::new();
        std::fs::File::open(path)
            .and_then(|mut f| f.read_to_end(&mut data))
            .map_err(|e| CfdError::invalid_parameter(format!("load {}: {e}", path.display())))?;
        let mut off = 0usize;
        let take = |off: &mut usize, n: usize| -> CfdResult<&[u8]> {
            let s = data
                .get(*off..*off + n)
                .ok_or_else(|| CfdError::invalid_parameter("load: truncated file"))?;
            *off += n;
            Ok(s)
        };
        if take(&mut off, 4)? != Self::SAVE_MAGIC {
            return Err(CfdError::invalid_parameter("load: bad magic"));
        }
        let version = u32::from_le_bytes(take(&mut off, 4)?.try_into().unwrap());
        if version != Self::SAVE_VERSION {
            return Err(CfdError::invalid_parameter(format!(
                "load: unsupported version {version}"
            )));
        }
        let nx = u64::from_le_bytes(take(&mut off, 8)?.try_into().unwrap()) as usize;
        let ny = u64::from_le_bytes(take(&mut off, 8)?.try_into().unwrap()) as usize;
        let dx = f64::from_le_bytes(take(&mut off, 8)?.try_into().unwrap());
        let dy = f64::from_le_bytes(take(&mut off, 8)?.try_into().unwrap());
        let mut field = Self::new(nx, ny, dx, dy)?;
        for m in field.matrices_mut() {
            let nrows = u64::from_le_bytes(take(&mut off, 8)?.try_into().unwrap()) as usize;
            let ncols = u64::from_le_bytes(take(&mut off, 8)?.try_into().unwrap()) as usize;
            if nrows != m.nrows() || ncols != m.ncols() {
                return Err(CfdError::invalid_parameter(format!(
                    "load: matrix shape {nrows}x{ncols} does not match field {}x{}",
                    m.nrows(),
                    m.ncols()
                )));
            }
            for v in m.iter_mut() {
                *v = f64::from_le_bytes(take(&mut off, 8)?.try_into().unwrap());
            }
        }
        if off != data.len() {
            return Err(CfdError::invalid_parameter("load: trailing bytes"));
        }
        Ok(field)
    }

    fn matrices(&self) -> [&DMatrix<f64>; 12] {
        [
            &self.u,
            &self.v,
            &self.p,
            &self.u_old,
            &self.v_old,
            &self.p_old,
            &self.u_star,
            &self.v_star,
            &self.p_prime,
            &self.su,
            &self.sv,
            &self.sp,
        ]
    }

    fn matrices_mut(&mut self) -> [&mut DMatrix<f64>; 12] {
        [
            &mut self.u,
            &mut self.v,
            &mut self.p,
            &mut self.u_old,
            &mut self.v_old,
            &mut self.p_old,
            &mut self.u_star,
            &mut self.v_star,
            &mut self.p_prime,
            &mut self.su,
            &mut self.sv,
            &mut self.sp,
        ]
    }
}

/// Analytical solutions for testing and validation
#[derive(Debug, Clone, Copy)]
pub enum AnalyticalSolution {
    /// Plane Poiseuille flow between parallel plates
    PoiseuillePlane { u_max: f64 },
    /// Taylor-Green vortex (decaying vortex solution)
    TaylorGreenVortex { amplitude: f64 },
}

#[cfg(test)]
mod save_load_tests {
    use super::*;

    fn scratch(name: &str) -> std::path::PathBuf {
        let dir = std::env::temp_dir().join("rtx_cfd_flow_field_tests");
        std::fs::create_dir_all(&dir).unwrap();
        dir.join(name)
    }

    #[test]
    fn save_load_roundtrip_is_bit_exact_across_every_matrix() {
        let mut field = FlowField::new(7, 5, 0.125, 0.25).unwrap();
        // Fill every matrix with distinct full-mantissa values so a
        // field mix-up or truncation cannot roundtrip by accident.
        for (k, m) in field.matrices_mut().into_iter().enumerate() {
            for (i, v) in m.iter_mut().enumerate() {
                *v = ((k * 1000 + i) as f64 * 0.7391 + 0.001).sin() * 3.7e3;
            }
        }
        let path = scratch("roundtrip.rtxf");
        field.save(&path).unwrap();
        let loaded = FlowField::load(&path).unwrap();

        assert_eq!(loaded.nx, field.nx);
        assert_eq!(loaded.ny, field.ny);
        assert_eq!(loaded.dx.to_bits(), field.dx.to_bits());
        assert_eq!(loaded.dy.to_bits(), field.dy.to_bits());
        for (a, b) in field.matrices().iter().zip(loaded.matrices().iter()) {
            assert_eq!(a.nrows(), b.nrows());
            assert_eq!(a.ncols(), b.ncols());
            for (x, y) in a.iter().zip(b.iter()) {
                assert_eq!(x.to_bits(), y.to_bits(), "field value changed in roundtrip");
            }
        }
    }

    #[test]
    fn load_rejects_truncated_and_corrupt_files() {
        let field = FlowField::new(5, 4, 0.1, 0.1).unwrap();
        let path = scratch("truncate.rtxf");
        field.save(&path).unwrap();
        let full = std::fs::read(&path).unwrap();

        let cut = scratch("truncate_cut.rtxf");
        std::fs::write(&cut, &full[..full.len() / 2]).unwrap();
        assert!(FlowField::load(&cut).is_err(), "truncated file must fail");

        let bad = scratch("bad_magic.rtxf");
        let mut corrupted = full.clone();
        corrupted[0] = b'X';
        std::fs::write(&bad, &corrupted).unwrap();
        assert!(FlowField::load(&bad).is_err(), "bad magic must fail");
    }
}