diff -r 943c6b3b9414 -r a1278320be26 src/discrete_gradient.rs --- a/src/discrete_gradient.rs Wed Apr 30 01:06:25 2025 -0500 +++ b/src/discrete_gradient.rs Wed Apr 30 16:39:01 2025 -0500 @@ -1,14 +1,13 @@ /*! Simple disrete gradient operators */ +use crate::error::DynResult; +use crate::instance::Instance; +use crate::linops::{Adjointable, BoundedLinear, Linear, Mapping, GEMV}; +use crate::norms::{Norm, L2}; +use crate::types::Float; +use nalgebra::{DVector, Dyn, Matrix, Storage, StorageMut, U1}; 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 @@ -27,26 +26,17 @@ pub struct BackwardNeumann; /// Finite differences gradient -pub struct Grad< - F : Float + nalgebra::RealField, - B : Discretisation, - const N : usize -> { - dims : [usize; N], - h : F, // may be negative to implement adjoints! - discretisation : B, +pub struct Grad, 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, - const N : usize -> { - dims : [usize; N], - h : F, // may be negative to implement adjoints! - discretisation : B, +pub struct Div, 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 @@ -62,53 +52,53 @@ use DiscretisationOrInterior::*; /// Trait for different discretisations -pub trait Discretisation : Copy { +pub trait Discretisation: Copy { /// Opposite discretisation, appropriate for adjoints with negated cell width. - type Opposite : Discretisation; + type Opposite: Discretisation; /// Add to appropiate index of `v` (as determined by `b`) the appropriate difference /// of `x` with cell width `h`. fn add_diff_mut( &self, - v : &mut Matrix, - x : &Matrix, - α : F, - b : DiscretisationOrInterior, + v: &mut Matrix, + x: &Matrix, + α: F, + b: DiscretisationOrInterior, ) where - SMut : StorageMut, - S : Storage; + SMut: StorageMut, + S: Storage; /// 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 { + fn opnorm_bound(&self, h: F) -> DynResult { // 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 + Ok(num_traits::Float::sqrt(8.0) / h) } } -impl Discretisation for ForwardNeumann { +impl Discretisation for ForwardNeumann { type Opposite = BackwardDirichlet; #[inline] fn add_diff_mut( &self, - v : &mut Matrix, - x : &Matrix, - α : F, - b : DiscretisationOrInterior, + v: &mut Matrix, + x: &Matrix, + α: F, + b: DiscretisationOrInterior, ) where - SMut : StorageMut, - S : Storage + SMut: StorageMut, + S: Storage, { match b { - Interior(c, (_, f)) | LeftBoundary(c, f) => { v[c] += (x[f] - x[c]) * α }, - RightBoundary(_c, _b) => { }, + Interior(c, (_, f)) | LeftBoundary(c, f) => v[c] += (x[f] - x[c]) * α, + RightBoundary(_c, _b) => {} } } @@ -118,23 +108,23 @@ } } -impl Discretisation for BackwardNeumann { +impl Discretisation for BackwardNeumann { type Opposite = ForwardDirichlet; #[inline] fn add_diff_mut( &self, - v : &mut Matrix, - x : &Matrix, - α : F, - b : DiscretisationOrInterior, + v: &mut Matrix, + x: &Matrix, + α: F, + b: DiscretisationOrInterior, ) where - SMut : StorageMut, - S : Storage + SMut: StorageMut, + S: Storage, { match b { - Interior(c, (b, _)) | RightBoundary(c, b) => { v[c] += (x[c] - x[b]) * α }, - LeftBoundary(_c, _f) => { }, + Interior(c, (b, _)) | RightBoundary(c, b) => v[c] += (x[c] - x[b]) * α, + LeftBoundary(_c, _f) => {} } } @@ -144,24 +134,24 @@ } } -impl Discretisation for BackwardDirichlet { +impl Discretisation for BackwardDirichlet { type Opposite = ForwardNeumann; #[inline] fn add_diff_mut( &self, - v : &mut Matrix, - x : &Matrix, - α : F, - b : DiscretisationOrInterior, + v: &mut Matrix, + x: &Matrix, + α: F, + b: DiscretisationOrInterior, ) where - SMut : StorageMut, - S : Storage + SMut: StorageMut, + S: Storage, { 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] * α }, + Interior(c, (b, _f)) => v[c] += (x[c] - x[b]) * α, + LeftBoundary(c, _f) => v[c] += x[c] * α, + RightBoundary(c, b) => v[c] -= x[b] * α, } } @@ -171,24 +161,24 @@ } } -impl Discretisation for ForwardDirichlet { +impl Discretisation for ForwardDirichlet { type Opposite = BackwardNeumann; #[inline] fn add_diff_mut( &self, - v : &mut Matrix, - x : &Matrix, - α : F, - b : DiscretisationOrInterior, + v: &mut Matrix, + x: &Matrix, + α: F, + b: DiscretisationOrInterior, ) where - SMut : StorageMut, - S : Storage + SMut: StorageMut, + S: Storage, { 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] * α }, + Interior(c, (_b, f)) => v[c] += (x[f] - x[c]) * α, + LeftBoundary(c, f) => v[c] += x[f] * α, + RightBoundary(c, _b) => v[c] -= x[c] * α, } } @@ -198,30 +188,37 @@ } } -struct Iter<'a, const N : usize> { +struct Iter<'a, const N: usize> { /// Dimensions - dims : &'a [usize; N], + dims: &'a [usize; N], /// Dimension along which to calculate differences - d : usize, + d: usize, /// Stride along coordinate d - d_stride : usize, + d_stride: usize, /// Cartesian indices - i : [usize; N], + i: [usize; N], /// Linear index - k : usize, + k: usize, /// Maximal linear index - len : usize + len: usize, } -impl<'a, const N : usize> Iter<'a, N> { - fn new(dims : &'a [usize; N], d : usize) -> Self { +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::(); let len = dims.iter().product::(); - Iter{ dims, d, d_stride, i : [0; N], k : 0, len } + Iter { + dims, + d, + d_stride, + i: [0; N], + k: 0, + len, + } } } -impl<'a, const N : usize> Iterator for Iter<'a, N> { +impl<'a, const N: usize> Iterator for Iter<'a, N> { type Item = DiscretisationOrInterior; fn next(&mut self) -> Option { let res = if self.k >= self.len { @@ -243,7 +240,7 @@ for j in 0..N { if self.i[j] + 1 < self.dims[j] { self.i[j] += 1; - break + break; } self.i[j] = 0 } @@ -251,14 +248,13 @@ } } -impl Mapping> -for Grad +impl Mapping> for Grad where - B : Discretisation, - F : Float + nalgebra::RealField, + B: Discretisation, + F: Float + nalgebra::RealField, { type Codomain = DVector; - fn apply>>(&self, i : I) -> DVector { + fn apply>>(&self, i: I) -> DVector { let mut y = DVector::zeros(N * self.len()); self.apply_add(&mut y, i); y @@ -266,15 +262,12 @@ } #[replace_float_literals(F::cast_from(literal))] -impl GEMV> -for Grad +impl GEMV> for Grad where - B : Discretisation, - F : Float + nalgebra::RealField, + B: Discretisation, + F: Float + nalgebra::RealField, { - fn gemv>>( - &self, y : &mut DVector, α : F, i : I, β : F - ) { + fn gemv>>(&self, y: &mut DVector, α: F, i: I, β: F) { if β == 0.0 { y.as_mut_slice().iter_mut().for_each(|x| *x = 0.0); } else if β != 1.0 { @@ -286,23 +279,22 @@ 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)); + 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) + self.discretisation.add_diff_mut(&mut v, x, α / h, b) } } }) } } -impl Mapping> -for Div +impl Mapping> for Div where - B : Discretisation, - F : Float + nalgebra::RealField, + B: Discretisation, + F: Float + nalgebra::RealField, { type Codomain = DVector; - fn apply>>(&self, i : I) -> DVector { + fn apply>>(&self, i: I) -> DVector { let mut y = DVector::zeros(self.len()); self.apply_add(&mut y, i); y @@ -310,15 +302,12 @@ } #[replace_float_literals(F::cast_from(literal))] -impl GEMV> -for Div +impl GEMV> for Div where - B : Discretisation, - F : Float + nalgebra::RealField, + B: Discretisation, + F: Float + nalgebra::RealField, { - fn gemv>>( - &self, y : &mut DVector, α : F, i : I, β : F - ) { + fn gemv>>(&self, y: &mut DVector, α: F, i: I, β: F) { if β == 0.0 { y.as_mut_slice().iter_mut().for_each(|x| *x = 0.0); } else if β != 1.0 { @@ -330,26 +319,51 @@ 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)); + 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) + self.discretisation.add_diff_mut(y, &v, α / h, b) } } }) } } -impl Grad +impl Grad where - B : Discretisation, - F : Float + nalgebra::RealField + B: Discretisation, + F: Float + nalgebra::RealField, { /// Creates a new discrete gradient operator for the vector `u`, verifying dimensions. - pub fn new_for(u : &DVector, h : F, dims : [usize; N], discretisation : B) - -> Option - { + pub fn new_for(u: &DVector, h: F, dims: [usize; N], discretisation: B) -> Option { if u.len() == dims.iter().product::() { - Some(Grad { dims, h, discretisation } ) + Some(Grad { + dims, + h, + discretisation, + }) + } else { + None + } + } + + fn len(&self) -> usize { + self.dims.iter().product::() + } +} + +impl Div +where + B: Discretisation, + F: Float + nalgebra::RealField, +{ + /// Creates a new discrete gradient operator for the vector `u`, verifying dimensions. + pub fn new_for(u: &DVector, h: F, dims: [usize; N], discretisation: B) -> Option { + if u.len() == dims.iter().product::() * N { + Some(Div { + dims, + h, + discretisation, + }) } else { None } @@ -360,108 +374,84 @@ } } - -impl Div +impl Linear> for Grad where - B : Discretisation, - F : Float + nalgebra::RealField + B: Discretisation, + F: Float + nalgebra::RealField, { - /// Creates a new discrete gradient operator for the vector `u`, verifying dimensions. - pub fn new_for(u : &DVector, h : F, dims : [usize; N], discretisation : B) - -> Option - { - if u.len() == dims.iter().product::() * N { - Some(Div { dims, h, discretisation } ) - } else { - None - } - } - - fn len(&self) -> usize { - self.dims.iter().product::() - } } -impl Linear> -for Grad +impl Linear> for Div where - B : Discretisation, - F : Float + nalgebra::RealField, + B: Discretisation, + F: Float + nalgebra::RealField, { } -impl Linear> -for Div +impl BoundedLinear, L2, L2, F> for Grad where - B : Discretisation, - F : Float + nalgebra::RealField, + B: Discretisation, + F: Float + nalgebra::RealField, + DVector: Norm, { -} - -impl BoundedLinear, L2, L2, F> -for Grad -where - B : Discretisation, - F : Float + nalgebra::RealField, - DVector : Norm, -{ - fn opnorm_bound(&self, _ : L2, _ : L2) -> F { + fn opnorm_bound(&self, _: L2, _: L2) -> DynResult { // Fuck nalgebra. - self.discretisation.opnorm_bound(num_traits::Float::abs(self.h)) + self.discretisation + .opnorm_bound(num_traits::Float::abs(self.h)) } } - -impl BoundedLinear, L2, L2, F> -for Div +impl BoundedLinear, L2, L2, F> for Div where - B : Discretisation, - F : Float + nalgebra::RealField, - DVector : Norm, + B: Discretisation, + F: Float + nalgebra::RealField, + DVector: Norm, { - fn opnorm_bound(&self, _ : L2, _ : L2) -> F { + fn opnorm_bound(&self, _: L2, _: L2) -> DynResult { // Fuck nalgebra. - self.discretisation.opnorm_bound(num_traits::Float::abs(self.h)) + self.discretisation + .opnorm_bound(num_traits::Float::abs(self.h)) } } -impl -Adjointable, DVector> -for Grad +impl Adjointable, DVector> for Grad where - B : Discretisation, - F : Float + nalgebra::RealField, + B: Discretisation, + F: Float + nalgebra::RealField, { type AdjointCodomain = DVector; - type Adjoint<'a> = Div where Self : 'a; + type Adjoint<'a> + = Div + 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(), + dims: self.dims, + h: -self.h, + discretisation: self.discretisation.opposite(), } } } - -impl -Adjointable, DVector> -for Div +impl Adjointable, DVector> for Div where - B : Discretisation, - F : Float + nalgebra::RealField, + B: Discretisation, + F: Float + nalgebra::RealField, { type AdjointCodomain = DVector; - type Adjoint<'a> = Grad where Self : 'a; + type Adjoint<'a> + = Grad + 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(), + dims: self.dims, + h: -self.h, + discretisation: self.discretisation.opposite(), } } } @@ -472,8 +462,8 @@ #[test] fn grad_adjoint() { - let im = DVector::from( (0..9).map(|t| t as f64).collect::>()); - let v = DVector::from( (0..18).map(|t| t as f64).collect::>()); + let im = DVector::from((0..9).map(|t| t as f64).collect::>()); + let v = DVector::from((0..18).map(|t| t as f64).collect::>()); let grad = Grad::new_for(&im, 1.0, [3, 3], ForwardNeumann).unwrap(); let a = grad.apply(&im).dot(&v); @@ -484,6 +474,5 @@ let a = grad.apply(&im).dot(&v); let b = grad.adjoint().apply(&v).dot(&im); assert_eq!(a, b); - } }