Mercurial > repos > alg_tools / changeset

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/src/discrete_gradient.rs	Mon Dec 30 09:49:08 2024 -0500
@@ -0,0 +1,489 @@
+/*!
+Simple disrete gradient operators
+ */
+use numeric_literals::replace_float_literals;
+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)]
+/// 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) -> 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);
+
+    }
+}

--- a/src/linops.rs	Sun Dec 22 17:42:14 2024 -0500
+++ b/src/linops.rs	Mon Dec 30 09:49:08 2024 -0500
@@ -50,11 +50,13 @@
     /// Computes  `y = αAx + βy`, where `A` is `Self`.
     fn gemv<I : Instance<X>>(&self, y : &mut Y, α : F, x : I, β : F);
 
+    #[inline]
     /// Computes `y = Ax`, where `A` is `Self`
     fn apply_mut<I : Instance<X>>(&self, y : &mut Y, x : I){
         self.gemv(y, 1.0, x, 0.0)
     }
 
+    #[inline]
     /// Computes `y += Ax`, where `A` is `Self`
     fn apply_add<I : Instance<X>>(&self, y : &mut Y, x : I){
         self.gemv(y, 1.0, x, 1.0)