leto 0.46.1

//! Typed contracts shared by CPU and accelerator finite-difference stencils.

use aequitas::systems::si::{quantities::Length, units::Meter};
use eunomia::{FloatElement, NumericElement, UnitScalar};
use thiserror::Error;

/// Boundary condition applied by a Cartesian finite-difference stencil.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BoundaryCondition {
    /// Homogeneous Dirichlet: `u = 0` on boundaries.
    Dirichlet,
    /// Homogeneous Neumann: `∂u/∂n = 0` on boundaries.
    Neumann,
    /// Periodic wrapping across domain boundaries.
    Periodic,
}

impl From<BoundaryCondition> for u32 {
    #[inline]
    fn from(value: BoundaryCondition) -> Self {
        match value {
            BoundaryCondition::Dirichlet => 0,
            BoundaryCondition::Neumann => 1,
            BoundaryCondition::Periodic => 2,
        }
    }
}

/// Sign convention for a Laplacian operator.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LaplacianPolarity {
    /// The differential operator `∇²`.
    Laplacian,
    /// The positive-semidefinite operator `-∇²`.
    NegativeLaplacian,
}

/// Typed Laplacian contract failure.
#[derive(Debug, Clone, PartialEq, Eq, Error)]
#[non_exhaustive]
pub enum LaplacianError {
    /// An axis is too short for the boundary stencil.
    #[error("Laplacian grid axes must contain at least two points: nx={nx}, ny={ny}")]
    GridTooSmall {
        /// Number of x-axis points.
        nx: usize,
        /// Number of y-axis points.
        ny: usize,
    },
    /// The flattened grid size does not fit in `usize`.
    #[error("Laplacian grid size overflows usize: nx={nx}, ny={ny}")]
    GridSizeOverflow {
        /// Number of x-axis points.
        nx: usize,
        /// Number of y-axis points.
        ny: usize,
    },
    /// A physical grid spacing is non-finite or non-positive.
    #[error("Laplacian spacing must be finite and positive: dx={dx}, dy={dy}")]
    InvalidSpacing {
        /// Debug representation of x spacing in meters.
        dx: String,
        /// Debug representation of y spacing in meters.
        dy: String,
    },
    /// Signed inverse squared spacings are non-finite, zero, or of mixed sign,
    /// so they describe neither polarity of a valid grid.
    #[error(
        "Laplacian signed inverse squared spacings must be finite, non-zero, and share one sign: x={x}, y={y}"
    )]
    InvalidInverseSpacing {
        /// Debug representation of the signed x value in m⁻².
        x: String,
        /// Debug representation of the signed y value in m⁻².
        y: String,
    },
    /// The input array length does not match the flattened grid.
    #[error("Laplacian input length mismatch: expected {expected}, actual {actual}")]
    InputLength {
        /// Required flattened length.
        expected: usize,
        /// Supplied input length.
        actual: usize,
    },
    /// The output array length does not match the flattened grid.
    #[error("Laplacian output length mismatch: expected {expected}, actual {actual}")]
    OutputLength {
        /// Required flattened length.
        expected: usize,
        /// Supplied output length.
        actual: usize,
    },
}

/// Validated two-dimensional Cartesian Laplacian contract.
///
/// Spacing crosses the boundary as [`Length`] and is reduced once to inverse
/// squared SI spacing in the native precision of `T`. The same contract drives
/// Leto CPU evaluation and Hephaestus accelerator dispatch.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Laplacian2D<T> {
    nx: usize,
    ny: usize,
    len: usize,
    inverse_spacing_squared: [T; 2],
    boundary: BoundaryCondition,
    polarity: LaplacianPolarity,
}

impl<T> Laplacian2D<T>
where
    T: FloatElement + UnitScalar,
{
    /// Validate a uniform endpoint-inclusive Cartesian grid.
    ///
    /// # Errors
    ///
    /// Returns [`LaplacianError`] when an axis has fewer than two points, the
    /// flattened size overflows, or either spacing is non-finite or non-positive.
    pub fn new(
        nx: usize,
        ny: usize,
        dx: Length<T>,
        dy: Length<T>,
        boundary: BoundaryCondition,
    ) -> Result<Self, LaplacianError> {
        let len = Self::validated_len(nx, ny)?;
        let dx_m = dx.in_unit::<Meter>();
        let dy_m = dy.in_unit::<Meter>();
        if !dx_m.is_finite()
            || dx_m <= <T as NumericElement>::ZERO
            || !dy_m.is_finite()
            || dy_m <= <T as NumericElement>::ZERO
        {
            return Err(LaplacianError::InvalidSpacing {
                dx: format!("{dx_m:?} m"),
                dy: format!("{dy_m:?} m"),
            });
        }

        Ok(Self {
            nx,
            ny,
            len,
            inverse_spacing_squared: [dx_m.recip().powi(2), dy_m.recip().powi(2)],
            boundary,
            polarity: LaplacianPolarity::Laplacian,
        })
    }

    /// Rebuild the contract from the value [`Self::signed_inverse_spacing_squared`]
    /// returns, as carried by an accelerator parameter block.
    ///
    /// The polarity is the shared sign of the two values, so for every contract
    /// `c`, rebuilding from `c.signed_inverse_spacing_squared()` with
    /// `c.boundary()` yields `c` exactly: no spacing is recomputed through a
    /// square root.
    ///
    /// # Errors
    ///
    /// Returns [`LaplacianError::GridTooSmall`] or
    /// [`LaplacianError::GridSizeOverflow`] as [`Self::new`] does, and
    /// [`LaplacianError::InvalidInverseSpacing`] when a value is non-finite or
    /// zero, or the two differ in sign.
    pub fn from_signed_inverse_spacing_squared(
        nx: usize,
        ny: usize,
        signed_inverse_spacing_squared: [T; 2],
        boundary: BoundaryCondition,
    ) -> Result<Self, LaplacianError> {
        let len = Self::validated_len(nx, ny)?;
        let [x, y] = signed_inverse_spacing_squared;
        let zero = <T as NumericElement>::ZERO;
        let valid = |value: T| value.is_finite() && value != zero;
        let polarity = if valid(x) && valid(y) && (x > zero) == (y > zero) {
            if x > zero {
                LaplacianPolarity::Laplacian
            } else {
                LaplacianPolarity::NegativeLaplacian
            }
        } else {
            return Err(LaplacianError::InvalidInverseSpacing {
                x: format!("{x:?}"),
                y: format!("{y:?}"),
            });
        };
        let magnitude = |value: T| if value > zero { value } else { zero - value };
        Ok(Self {
            nx,
            ny,
            len,
            inverse_spacing_squared: [magnitude(x), magnitude(y)],
            boundary,
            polarity,
        })
    }

    /// Validate the grid axes and return the flattened length.
    fn validated_len(nx: usize, ny: usize) -> Result<usize, LaplacianError> {
        if nx < 2 || ny < 2 {
            return Err(LaplacianError::GridTooSmall { nx, ny });
        }
        nx.checked_mul(ny)
            .ok_or(LaplacianError::GridSizeOverflow { nx, ny })
    }

    /// Select the operator sign convention.
    #[must_use]
    pub const fn with_polarity(mut self, polarity: LaplacianPolarity) -> Self {
        self.polarity = polarity;
        self
    }

    /// Number of x-axis points.
    #[must_use]
    pub const fn nx(&self) -> usize {
        self.nx
    }

    /// Number of y-axis points.
    #[must_use]
    pub const fn ny(&self) -> usize {
        self.ny
    }

    /// Required flattened array length.
    #[must_use]
    pub const fn len(&self) -> usize {
        self.len
    }

    /// Whether the flattened grid is empty.
    #[must_use]
    pub const fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// Boundary condition encoded by this stencil.
    #[must_use]
    pub const fn boundary(&self) -> BoundaryCondition {
        self.boundary
    }

    /// Operator sign convention.
    #[must_use]
    pub const fn polarity(&self) -> LaplacianPolarity {
        self.polarity
    }

    /// Signed inverse squared spacing for x and y.
    #[must_use]
    pub fn signed_inverse_spacing_squared(&self) -> [T; 2] {
        match self.polarity {
            LaplacianPolarity::Laplacian => self.inverse_spacing_squared,
            LaplacianPolarity::NegativeLaplacian => [
                <T as NumericElement>::ZERO - self.inverse_spacing_squared[0],
                <T as NumericElement>::ZERO - self.inverse_spacing_squared[1],
            ],
        }
    }
}

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

    #[test]
    fn rejects_non_positive_spacing() {
        let error = Laplacian2D::new(
            4,
            4,
            Length::from_unit::<Meter>(0.0f32),
            Length::from_unit::<Meter>(1.0f32),
            BoundaryCondition::Dirichlet,
        )
        .expect_err("zero spacing must be rejected");

        assert!(matches!(error, LaplacianError::InvalidSpacing { .. }));
    }

    /// Rebuilding from the signed inverse squared spacings reproduces the
    /// contract exactly, for both polarities and a spacing (0.3 m) whose
    /// inverse square is inexact in binary.
    #[test]
    fn signed_inverse_spacing_round_trips_exactly() {
        for polarity in [
            LaplacianPolarity::Laplacian,
            LaplacianPolarity::NegativeLaplacian,
        ] {
            let contract = Laplacian2D::new(
                5,
                3,
                Length::from_unit::<Meter>(0.3f32),
                Length::from_unit::<Meter>(0.7f32),
                BoundaryCondition::Neumann,
            )
            .expect("valid grid")
            .with_polarity(polarity);
            let rebuilt = Laplacian2D::from_signed_inverse_spacing_squared(
                5,
                3,
                contract.signed_inverse_spacing_squared(),
                BoundaryCondition::Neumann,
            )
            .expect("round trip");
            assert_eq!(rebuilt, contract);
        }
    }

    #[test]
    fn rejects_zero_non_finite_or_mixed_sign_inverse_spacing() {
        for signed in [
            [0.0f32, 1.0],
            [1.0, f32::INFINITY],
            [f32::NAN, 1.0],
            [1.0, -1.0],
            [-4.0, 4.0],
        ] {
            let error = Laplacian2D::from_signed_inverse_spacing_squared(
                4,
                4,
                signed,
                BoundaryCondition::Periodic,
            )
            .expect_err("must be rejected");
            assert!(
                matches!(error, LaplacianError::InvalidInverseSpacing { .. }),
                "{signed:?}: {error}"
            );
        }
    }

    #[test]
    fn inverse_spacing_constructor_validates_the_grid() {
        let error = Laplacian2D::from_signed_inverse_spacing_squared(
            1,
            4,
            [1.0f32, 1.0],
            BoundaryCondition::Dirichlet,
        )
        .expect_err("a one-point axis must be rejected");
        assert_eq!(error, LaplacianError::GridTooSmall { nx: 1, ny: 4 });
    }
}