src/discrete_gradient.rs

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

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

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

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

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

/// Finite differences gradient
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(bound(
    serialize = "[usize; N] : Serialize,
                 B : Serialize,
                 F : Serialize,",
    deserialize = "[usize; N] : for<'a> Deserialize<'a>,
                   B : for<'a> Deserialize<'a>,
                   F : for<'a> Deserialize<'a>,"
))]
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) -> 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.
        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(|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<F, L2>,
{
    fn opnorm_bound(&self, _ : L2, _ : L2) -> 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<F, L2>,
{
    fn opnorm_bound(&self, _ : L2, _ : L2) -> 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;

    /// Form the adjoint operator of `self`.
    fn adjoint(&self) -> Self::Adjoint<'_> {
        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;

    /// Form the adjoint operator of `self`.
    fn adjoint(&self) -> Self::Adjoint<'_> {
        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);

    }
}