src/discrete_gradient.rs

/*!
Simple disrete gradient operators
 */
use crate::error::DynResult;
use crate::instance::Instance;
use crate::linops::{Adjointable, BoundedLinear, Linear, Mapping, SimplyAdjointable, GEMV};
use crate::norms::{Norm, L2};
use crate::types::Float;
use nalgebra::{DVector, Dyn, Matrix, Storage, StorageMut, U1};
use numeric_literals::replace_float_literals;

#[derive(Copy, Clone, Debug)]
/// Forward differences with Neumann boundary conditions
pub struct ForwardNeumann;

#[derive(Copy, Clone, Debug)]
/// Forward differences with Dirichlet boundary conditions
pub struct ForwardDirichlet;

#[derive(Copy, Clone, Debug)]
/// Backward differences with Dirichlet boundary conditions
pub struct BackwardDirichlet;

#[derive(Copy, Clone, Debug)]
/// Backward differences with Neumann boundary conditions
pub struct BackwardNeumann;

/// Finite differences gradient
pub struct Grad<F: Float + nalgebra::RealField, B: Discretisation<F>, const N: usize> {
    dims: [usize; N],
    h: F, // may be negative to implement adjoints!
    discretisation: B,
}

/// Finite differences divergence
pub struct Div<F: Float + nalgebra::RealField, B: Discretisation<F>, const N: usize> {
    dims: [usize; N],
    h: F, // may be negative to implement adjoints!
    discretisation: B,
}

/// Internal: classification of a point in a 1D discretisation
pub enum DiscretisationOrInterior {
    /// center, forward
    LeftBoundary(usize, usize),
    /// center, backward
    RightBoundary(usize, usize),
    /// center, (backward, forward)
    Interior(usize, (usize, usize)),
}

use DiscretisationOrInterior::*;

/// Trait for different discretisations
pub trait Discretisation<F: Float + nalgebra::RealField>: Copy {
    /// Opposite discretisation, appropriate for adjoints with negated cell width.
    type Opposite: Discretisation<F>;

    /// Add to appropiate index of `v` (as determined by `b`) the appropriate difference
    /// of `x` with cell width `h`.
    fn add_diff_mut<SMut, S>(
        &self,
        v: &mut Matrix<F, Dyn, U1, SMut>,
        x: &Matrix<F, Dyn, U1, S>,
        α: F,
        b: DiscretisationOrInterior,
    ) where
        SMut: StorageMut<F, Dyn, U1>,
        S: Storage<F, Dyn, U1>;

    /// Give the opposite discretisation, appropriate for adjoints with negated `h`.
    fn opposite(&self) -> Self::Opposite;

    /// Bound for the corresponding operator norm.
    #[replace_float_literals(F::cast_from(literal))]
    fn opnorm_bound(&self, h: F) -> DynResult<F> {
        // See: Chambolle, “An Algorithm for Total Variation Minimization and Applications”.
        // Ok for forward and backward differences.
        //
        // Fuck nalgebra for polluting everything with its own shit.
        Ok(num_traits::Float::sqrt(8.0) / h)
    }
}

impl<F: Float + nalgebra::RealField> Discretisation<F> for ForwardNeumann {
    type Opposite = BackwardDirichlet;

    #[inline]
    fn add_diff_mut<SMut, S>(
        &self,
        v: &mut Matrix<F, Dyn, U1, SMut>,
        x: &Matrix<F, Dyn, U1, S>,
        α: F,
        b: DiscretisationOrInterior,
    ) where
        SMut: StorageMut<F, Dyn, U1>,
        S: Storage<F, Dyn, U1>,
    {
        match b {
            Interior(c, (_, f)) | LeftBoundary(c, f) => v[c] += (x[f] - x[c]) * α,
            RightBoundary(_c, _b) => {}
        }
    }

    #[inline]
    fn opposite(&self) -> Self::Opposite {
        BackwardDirichlet
    }
}

impl<F: Float + nalgebra::RealField> Discretisation<F> for BackwardNeumann {
    type Opposite = ForwardDirichlet;

    #[inline]
    fn add_diff_mut<SMut, S>(
        &self,
        v: &mut Matrix<F, Dyn, U1, SMut>,
        x: &Matrix<F, Dyn, U1, S>,
        α: F,
        b: DiscretisationOrInterior,
    ) where
        SMut: StorageMut<F, Dyn, U1>,
        S: Storage<F, Dyn, U1>,
    {
        match b {
            Interior(c, (b, _)) | RightBoundary(c, b) => v[c] += (x[c] - x[b]) * α,
            LeftBoundary(_c, _f) => {}
        }
    }

    #[inline]
    fn opposite(&self) -> Self::Opposite {
        ForwardDirichlet
    }
}

impl<F: Float + nalgebra::RealField> Discretisation<F> for BackwardDirichlet {
    type Opposite = ForwardNeumann;

    #[inline]
    fn add_diff_mut<SMut, S>(
        &self,
        v: &mut Matrix<F, Dyn, U1, SMut>,
        x: &Matrix<F, Dyn, U1, S>,
        α: F,
        b: DiscretisationOrInterior,
    ) where
        SMut: StorageMut<F, Dyn, U1>,
        S: Storage<F, Dyn, U1>,
    {
        match b {
            Interior(c, (b, _f)) => v[c] += (x[c] - x[b]) * α,
            LeftBoundary(c, _f) => v[c] += x[c] * α,
            RightBoundary(c, b) => v[c] -= x[b] * α,
        }
    }

    #[inline]
    fn opposite(&self) -> Self::Opposite {
        ForwardNeumann
    }
}

impl<F: Float + nalgebra::RealField> Discretisation<F> for ForwardDirichlet {
    type Opposite = BackwardNeumann;

    #[inline]
    fn add_diff_mut<SMut, S>(
        &self,
        v: &mut Matrix<F, Dyn, U1, SMut>,
        x: &Matrix<F, Dyn, U1, S>,
        α: F,
        b: DiscretisationOrInterior,
    ) where
        SMut: StorageMut<F, Dyn, U1>,
        S: Storage<F, Dyn, U1>,
    {
        match b {
            Interior(c, (_b, f)) => v[c] += (x[f] - x[c]) * α,
            LeftBoundary(c, f) => v[c] += x[f] * α,
            RightBoundary(c, _b) => v[c] -= x[c] * α,
        }
    }

    #[inline]
    fn opposite(&self) -> Self::Opposite {
        BackwardNeumann
    }
}

struct Iter<'a, const N: usize> {
    /// Dimensions
    dims: &'a [usize; N],
    /// Dimension along which to calculate differences
    d: usize,
    /// Stride along coordinate d
    d_stride: usize,
    /// Cartesian indices
    i: [usize; N],
    /// Linear index
    k: usize,
    /// Maximal linear index
    len: usize,
}

impl<'a, const N: usize> Iter<'a, N> {
    fn new(dims: &'a [usize; N], d: usize) -> Self {
        let d_stride = dims[0..d].iter().product::<usize>();
        let len = dims.iter().product::<usize>();
        Iter { dims, d, d_stride, i: [0; N], k: 0, len }
    }
}

impl<'a, const N: usize> Iterator for Iter<'a, N> {
    type Item = DiscretisationOrInterior;
    fn next(&mut self) -> Option<Self::Item> {
        let res = if self.k >= self.len {
            None
        } else {
            let cartesian_idx = self.i[self.d];
            let cartesian_max = self.dims[self.d];
            let k = self.k;

            if cartesian_idx == 0 {
                Some(LeftBoundary(k, k + self.d_stride))
            } else if cartesian_idx + 1 >= cartesian_max {
                Some(RightBoundary(k, k - self.d_stride))
            } else {
                Some(Interior(k, (k - self.d_stride, k + self.d_stride)))
            }
        };
        self.k += 1;
        for j in 0..N {
            if self.i[j] + 1 < self.dims[j] {
                self.i[j] += 1;
                break;
            }
            self.i[j] = 0
        }
        res
    }
}

impl<F, B, const N: usize> Mapping<DVector<F>> for Grad<F, B, N>
where
    B: Discretisation<F>,
    F: Float + nalgebra::RealField,
{
    type Codomain = DVector<F>;
    fn apply<I: Instance<DVector<F>>>(&self, i: I) -> DVector<F> {
        let mut y = DVector::zeros(N * self.len());
        self.apply_add(&mut y, i);
        y
    }
}

#[replace_float_literals(F::cast_from(literal))]
impl<F, B, const N: usize> GEMV<F, DVector<F>> for Grad<F, B, N>
where
    B: Discretisation<F>,
    F: Float + nalgebra::RealField,
{
    fn gemv<I: Instance<DVector<F>>>(&self, y: &mut DVector<F>, α: F, i: I, β: F) {
        if β == 0.0 {
            y.as_mut_slice().iter_mut().for_each(|x| *x = 0.0);
        } else if β != 1.0 {
            //*y *= β;
            y.as_mut_slice().iter_mut().for_each(|x| *x *= β);
        }
        let h = self.h;
        let m = self.len();
        i.eval(|x| {
            assert_eq!(x.len(), m);
            for d in 0..N {
                let mut v = y.generic_view_mut((d * m, 0), (Dyn(m), U1));
                for b in Iter::new(&self.dims, d) {
                    self.discretisation.add_diff_mut(&mut v, x, α / h, b)
                }
            }
        })
    }
}

impl<F, B, const N: usize> Mapping<DVector<F>> for Div<F, B, N>
where
    B: Discretisation<F>,
    F: Float + nalgebra::RealField,
{
    type Codomain = DVector<F>;
    fn apply<I: Instance<DVector<F>>>(&self, i: I) -> DVector<F> {
        let mut y = DVector::zeros(self.len());
        self.apply_add(&mut y, i);
        y
    }
}

#[replace_float_literals(F::cast_from(literal))]
impl<F, B, const N: usize> GEMV<F, DVector<F>> for Div<F, B, N>
where
    B: Discretisation<F>,
    F: Float + nalgebra::RealField,
{
    fn gemv<I: Instance<DVector<F>>>(&self, y: &mut DVector<F>, α: F, i: I, β: F) {
        if β == 0.0 {
            y.as_mut_slice().iter_mut().for_each(|x| *x = 0.0);
        } else if β != 1.0 {
            //*y *= β;
            y.as_mut_slice().iter_mut().for_each(|x| *x *= β);
        }
        let h = self.h;
        let m = self.len();
        i.eval_decompose(|x| {
            assert_eq!(x.len(), N * m);
            for d in 0..N {
                let v = x.generic_view((d * m, 0), (Dyn(m), U1));
                for b in Iter::new(&self.dims, d) {
                    self.discretisation.add_diff_mut(y, &v, α / h, b)
                }
            }
        })
    }
}

impl<F, B, const N: usize> Grad<F, B, N>
where
    B: Discretisation<F>,
    F: Float + nalgebra::RealField,
{
    /// Creates a new discrete gradient operator for the vector `u`, verifying dimensions.
    pub fn new_for(u: &DVector<F>, h: F, dims: [usize; N], discretisation: B) -> Option<Self> {
        if u.len() == dims.iter().product::<usize>() {
            Some(Grad { dims, h, discretisation })
        } else {
            None
        }
    }

    fn len(&self) -> usize {
        self.dims.iter().product::<usize>()
    }
}

impl<F, B, const N: usize> Div<F, B, N>
where
    B: Discretisation<F>,
    F: Float + nalgebra::RealField,
{
    /// Creates a new discrete gradient operator for the vector `u`, verifying dimensions.
    pub fn new_for(u: &DVector<F>, h: F, dims: [usize; N], discretisation: B) -> Option<Self> {
        if u.len() == dims.iter().product::<usize>() * N {
            Some(Div { dims, h, discretisation })
        } else {
            None
        }
    }

    fn len(&self) -> usize {
        self.dims.iter().product::<usize>()
    }
}

impl<F, B, const N: usize> Linear<DVector<F>> for Grad<F, B, N>
where
    B: Discretisation<F>,
    F: Float + nalgebra::RealField,
{
}

impl<F, B, const N: usize> Linear<DVector<F>> for Div<F, B, N>
where
    B: Discretisation<F>,
    F: Float + nalgebra::RealField,
{
}

impl<F, B, const N: usize> BoundedLinear<DVector<F>, L2, L2, F> for Grad<F, B, N>
where
    B: Discretisation<F>,
    F: Float + nalgebra::RealField,
    DVector<F>: Norm<L2, F>,
{
    fn opnorm_bound(&self, _: L2, _: L2) -> DynResult<F> {
        // Fuck nalgebra.
        self.discretisation
            .opnorm_bound(num_traits::Float::abs(self.h))
    }
}

impl<F, B, const N: usize> BoundedLinear<DVector<F>, L2, L2, F> for Div<F, B, N>
where
    B: Discretisation<F>,
    F: Float + nalgebra::RealField,
    DVector<F>: Norm<L2, F>,
{
    fn opnorm_bound(&self, _: L2, _: L2) -> DynResult<F> {
        // Fuck nalgebra.
        self.discretisation
            .opnorm_bound(num_traits::Float::abs(self.h))
    }
}

impl<F, B, const N: usize> Adjointable<DVector<F>, DVector<F>> for Grad<F, B, N>
where
    B: Discretisation<F>,
    F: Float + nalgebra::RealField,
{
    type AdjointCodomain = DVector<F>;
    type Adjoint<'a>
        = Div<F, B::Opposite, N>
    where
        Self: 'a;

    fn adjoint(&self) -> Self::Adjoint<'_> {
        Div { dims: self.dims, h: -self.h, discretisation: self.discretisation.opposite() }
    }
}

impl<F, B, const N: usize> SimplyAdjointable<DVector<F>, DVector<F>> for Grad<F, B, N>
where
    B: Discretisation<F>,
    F: Float + nalgebra::RealField,
{
    type AdjointCodomain = DVector<F>;
    type SimpleAdjoint = Div<F, B::Opposite, N>;

    fn adjoint(&self) -> Self::SimpleAdjoint {
        Div { dims: self.dims, h: -self.h, discretisation: self.discretisation.opposite() }
    }
}

impl<F, B, const N: usize> Adjointable<DVector<F>, DVector<F>> for Div<F, B, N>
where
    B: Discretisation<F>,
    F: Float + nalgebra::RealField,
{
    type AdjointCodomain = DVector<F>;
    type Adjoint<'a>
        = Grad<F, B::Opposite, N>
    where
        Self: 'a;

    fn adjoint(&self) -> Self::Adjoint<'_> {
        Grad { dims: self.dims, h: -self.h, discretisation: self.discretisation.opposite() }
    }
}

impl<F, B, const N: usize> SimplyAdjointable<DVector<F>, DVector<F>> for Div<F, B, N>
where
    B: Discretisation<F>,
    F: Float + nalgebra::RealField,
{
    type AdjointCodomain = DVector<F>;
    type SimpleAdjoint = Grad<F, B::Opposite, N>;

    fn adjoint(&self) -> Self::SimpleAdjoint {
        Grad { dims: self.dims, h: -self.h, discretisation: self.discretisation.opposite() }
    }
}

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

    #[test]
    fn grad_adjoint() {
        let im = DVector::from((0..9).map(|t| t as f64).collect::<Vec<_>>());
        let v = DVector::from((0..18).map(|t| t as f64).collect::<Vec<_>>());

        let grad = Grad::new_for(&im, 1.0, [3, 3], ForwardNeumann).unwrap();
        let a = grad.apply(&im).dot(&v);
        let b = grad.adjoint().apply(&v).dot(&im);
        assert_eq!(a, b);

        let grad = Grad::new_for(&im, 1.0, [3, 3], ForwardDirichlet).unwrap();
        let a = grad.apply(&im).dot(&v);
        let b = grad.adjoint().apply(&v).dot(&im);
        assert_eq!(a, b);
    }
}