comparison: src/discrete_gradient.rs
src/discrete_gradient.rs
- branch
- dev
- changeset 67
- d7c0f431cbd6
- child 74
- 2c76df38d02b
- child 80
- f802ddbabcfc
equal
deleted
inserted
replaced
| 66:15bfebfbfa3e | 67:d7c0f431cbd6 |
|---|---|
| 1 /*! | |
| 2 Simple disrete gradient operators | |
| 3 */ | |
| 4 use numeric_literals::replace_float_literals; | |
| 5 use nalgebra::{ | |
| 6 DVector, Matrix, U1, Storage, StorageMut, Dyn | |
| 7 }; | |
| 8 use crate::types::Float; | |
| 9 use crate::instance::Instance; | |
| 10 use crate::linops::{Mapping, Linear, BoundedLinear, Adjointable, GEMV}; | |
| 11 use crate::norms::{Norm, L2}; | |
| 12 | |
| 13 #[derive(Copy, Clone, Debug)] | |
| 14 /// Forward differences with Neumann boundary conditions | |
| 15 pub struct ForwardNeumann; | |
| 16 | |
| 17 #[derive(Copy, Clone, Debug)] | |
| 18 /// Forward differences with Dirichlet boundary conditions | |
| 19 pub struct ForwardDirichlet; | |
| 20 | |
| 21 #[derive(Copy, Clone, Debug)] | |
| 22 /// Backward differences with Dirichlet boundary conditions | |
| 23 pub struct BackwardDirichlet; | |
| 24 | |
| 25 #[derive(Copy, Clone, Debug)] | |
| 26 /// Backward differences with Neumann boundary conditions | |
| 27 pub struct BackwardNeumann; | |
| 28 | |
| 29 /// Finite differences gradient | |
| 30 pub struct Grad< | |
| 31 F : Float + nalgebra::RealField, | |
| 32 B : Discretisation<F>, | |
| 33 const N : usize | |
| 34 > { | |
| 35 dims : [usize; N], | |
| 36 h : F, // may be negative to implement adjoints! | |
| 37 discretisation : B, | |
| 38 } | |
| 39 | |
| 40 | |
| 41 /// Finite differences divergence | |
| 42 pub struct Div< | |
| 43 F : Float + nalgebra::RealField, | |
| 44 B : Discretisation<F>, | |
| 45 const N : usize | |
| 46 > { | |
| 47 dims : [usize; N], | |
| 48 h : F, // may be negative to implement adjoints! | |
| 49 discretisation : B, | |
| 50 } | |
| 51 | |
| 52 /// Internal: classification of a point in a 1D discretisation | |
| 53 pub enum DiscretisationOrInterior { | |
| 54 /// center, forward | |
| 55 LeftBoundary(usize, usize), | |
| 56 /// center, backward | |
| 57 RightBoundary(usize, usize), | |
| 58 /// center, (backward, forward) | |
| 59 Interior(usize, (usize, usize)), | |
| 60 } | |
| 61 | |
| 62 use DiscretisationOrInterior::*; | |
| 63 | |
| 64 /// Trait for different discretisations | |
| 65 pub trait Discretisation<F : Float + nalgebra::RealField> : Copy { | |
| 66 /// Opposite discretisation, appropriate for adjoints with negated cell width. | |
| 67 type Opposite : Discretisation<F>; | |
| 68 | |
| 69 /// Add to appropiate index of `v` (as determined by `b`) the appropriate difference | |
| 70 /// of `x` with cell width `h`. | |
| 71 fn add_diff_mut<SMut, S>( | |
| 72 &self, | |
| 73 v : &mut Matrix<F, Dyn, U1, SMut>, | |
| 74 x : &Matrix<F, Dyn, U1, S>, | |
| 75 α : F, | |
| 76 b : DiscretisationOrInterior, | |
| 77 ) where | |
| 78 SMut : StorageMut<F, Dyn, U1>, | |
| 79 S : Storage<F, Dyn, U1>; | |
| 80 | |
| 81 /// Give the opposite discretisation, appropriate for adjoints with negated `h`. | |
| 82 fn opposite(&self) -> Self::Opposite; | |
| 83 | |
| 84 /// Bound for the corresponding operator norm. | |
| 85 #[replace_float_literals(F::cast_from(literal))] | |
| 86 fn opnorm_bound(&self, h : F) -> F { | |
| 87 // See: Chambolle, “An Algorithm for Total Variation Minimization and Applications”. | |
| 88 // Ok for forward and backward differences. | |
| 89 // | |
| 90 // Fuck nalgebra for polluting everything with its own shit. | |
| 91 num_traits::Float::sqrt(8.0) / h | |
| 92 } | |
| 93 } | |
| 94 | |
| 95 impl<F : Float + nalgebra::RealField> Discretisation<F> for ForwardNeumann { | |
| 96 type Opposite = BackwardDirichlet; | |
| 97 | |
| 98 #[inline] | |
| 99 fn add_diff_mut<SMut, S>( | |
| 100 &self, | |
| 101 v : &mut Matrix<F, Dyn, U1, SMut>, | |
| 102 x : &Matrix<F, Dyn, U1, S>, | |
| 103 α : F, | |
| 104 b : DiscretisationOrInterior, | |
| 105 ) where | |
| 106 SMut : StorageMut<F, Dyn, U1>, | |
| 107 S : Storage<F, Dyn, U1> | |
| 108 { | |
| 109 match b { | |
| 110 Interior(c, (_, f)) | LeftBoundary(c, f) => { v[c] += (x[f] - x[c]) * α }, | |
| 111 RightBoundary(_c, _b) => { }, | |
| 112 } | |
| 113 } | |
| 114 | |
| 115 #[inline] | |
| 116 fn opposite(&self) -> Self::Opposite { | |
| 117 BackwardDirichlet | |
| 118 } | |
| 119 } | |
| 120 | |
| 121 impl<F : Float + nalgebra::RealField> Discretisation<F> for BackwardNeumann { | |
| 122 type Opposite = ForwardDirichlet; | |
| 123 | |
| 124 #[inline] | |
| 125 fn add_diff_mut<SMut, S>( | |
| 126 &self, | |
| 127 v : &mut Matrix<F, Dyn, U1, SMut>, | |
| 128 x : &Matrix<F, Dyn, U1, S>, | |
| 129 α : F, | |
| 130 b : DiscretisationOrInterior, | |
| 131 ) where | |
| 132 SMut : StorageMut<F, Dyn, U1>, | |
| 133 S : Storage<F, Dyn, U1> | |
| 134 { | |
| 135 match b { | |
| 136 Interior(c, (b, _)) | RightBoundary(c, b) => { v[c] += (x[c] - x[b]) * α }, | |
| 137 LeftBoundary(_c, _f) => { }, | |
| 138 } | |
| 139 } | |
| 140 | |
| 141 #[inline] | |
| 142 fn opposite(&self) -> Self::Opposite { | |
| 143 ForwardDirichlet | |
| 144 } | |
| 145 } | |
| 146 | |
| 147 impl<F : Float + nalgebra::RealField> Discretisation<F> for BackwardDirichlet { | |
| 148 type Opposite = ForwardNeumann; | |
| 149 | |
| 150 #[inline] | |
| 151 fn add_diff_mut<SMut, S>( | |
| 152 &self, | |
| 153 v : &mut Matrix<F, Dyn, U1, SMut>, | |
| 154 x : &Matrix<F, Dyn, U1, S>, | |
| 155 α : F, | |
| 156 b : DiscretisationOrInterior, | |
| 157 ) where | |
| 158 SMut : StorageMut<F, Dyn, U1>, | |
| 159 S : Storage<F, Dyn, U1> | |
| 160 { | |
| 161 match b { | |
| 162 Interior(c, (b, _f)) => { v[c] += (x[c] - x[b]) * α }, | |
| 163 LeftBoundary(c, _f) => { v[c] += x[c] * α }, | |
| 164 RightBoundary(c, b) => { v[c] -= x[b] * α }, | |
| 165 } | |
| 166 } | |
| 167 | |
| 168 #[inline] | |
| 169 fn opposite(&self) -> Self::Opposite { | |
| 170 ForwardNeumann | |
| 171 } | |
| 172 } | |
| 173 | |
| 174 impl<F : Float + nalgebra::RealField> Discretisation<F> for ForwardDirichlet { | |
| 175 type Opposite = BackwardNeumann; | |
| 176 | |
| 177 #[inline] | |
| 178 fn add_diff_mut<SMut, S>( | |
| 179 &self, | |
| 180 v : &mut Matrix<F, Dyn, U1, SMut>, | |
| 181 x : &Matrix<F, Dyn, U1, S>, | |
| 182 α : F, | |
| 183 b : DiscretisationOrInterior, | |
| 184 ) where | |
| 185 SMut : StorageMut<F, Dyn, U1>, | |
| 186 S : Storage<F, Dyn, U1> | |
| 187 { | |
| 188 match b { | |
| 189 Interior(c, (_b, f)) => { v[c] += (x[f] - x[c]) * α }, | |
| 190 LeftBoundary(c, f) => { v[c] += x[f] * α }, | |
| 191 RightBoundary(c, _b) => { v[c] -= x[c] * α }, | |
| 192 } | |
| 193 } | |
| 194 | |
| 195 #[inline] | |
| 196 fn opposite(&self) -> Self::Opposite { | |
| 197 BackwardNeumann | |
| 198 } | |
| 199 } | |
| 200 | |
| 201 struct Iter<'a, const N : usize> { | |
| 202 /// Dimensions | |
| 203 dims : &'a [usize; N], | |
| 204 /// Dimension along which to calculate differences | |
| 205 d : usize, | |
| 206 /// Stride along coordinate d | |
| 207 d_stride : usize, | |
| 208 /// Cartesian indices | |
| 209 i : [usize; N], | |
| 210 /// Linear index | |
| 211 k : usize, | |
| 212 /// Maximal linear index | |
| 213 len : usize | |
| 214 } | |
| 215 | |
| 216 impl<'a, const N : usize> Iter<'a, N> { | |
| 217 fn new(dims : &'a [usize; N], d : usize) -> Self { | |
| 218 let d_stride = dims[0..d].iter().product::<usize>(); | |
| 219 let len = dims.iter().product::<usize>(); | |
| 220 Iter{ dims, d, d_stride, i : [0; N], k : 0, len } | |
| 221 } | |
| 222 } | |
| 223 | |
| 224 impl<'a, const N : usize> Iterator for Iter<'a, N> { | |
| 225 type Item = DiscretisationOrInterior; | |
| 226 fn next(&mut self) -> Option<Self::Item> { | |
| 227 let res = if self.k >= self.len { | |
| 228 None | |
| 229 } else { | |
| 230 let cartesian_idx = self.i[self.d]; | |
| 231 let cartesian_max = self.dims[self.d]; | |
| 232 let k = self.k; | |
| 233 | |
| 234 if cartesian_idx == 0 { | |
| 235 Some(LeftBoundary(k, k + self.d_stride)) | |
| 236 } else if cartesian_idx + 1 >= cartesian_max { | |
| 237 Some(RightBoundary(k, k - self.d_stride)) | |
| 238 } else { | |
| 239 Some(Interior(k, (k - self.d_stride, k + self.d_stride))) | |
| 240 } | |
| 241 }; | |
| 242 self.k += 1; | |
| 243 for j in 0..N { | |
| 244 if self.i[j] + 1 < self.dims[j] { | |
| 245 self.i[j] += 1; | |
| 246 break | |
| 247 } | |
| 248 self.i[j] = 0 | |
| 249 } | |
| 250 res | |
| 251 } | |
| 252 } | |
| 253 | |
| 254 impl<F, B, const N : usize> Mapping<DVector<F>> | |
| 255 for Grad<F, B, N> | |
| 256 where | |
| 257 B : Discretisation<F>, | |
| 258 F : Float + nalgebra::RealField, | |
| 259 { | |
| 260 type Codomain = DVector<F>; | |
| 261 fn apply<I : Instance<DVector<F>>>(&self, i : I) -> DVector<F> { | |
| 262 let mut y = DVector::zeros(N * self.len()); | |
| 263 self.apply_add(&mut y, i); | |
| 264 y | |
| 265 } | |
| 266 } | |
| 267 | |
| 268 #[replace_float_literals(F::cast_from(literal))] | |
| 269 impl<F, B, const N : usize> GEMV<F, DVector<F>> | |
| 270 for Grad<F, B, N> | |
| 271 where | |
| 272 B : Discretisation<F>, | |
| 273 F : Float + nalgebra::RealField, | |
| 274 { | |
| 275 fn gemv<I : Instance<DVector<F>>>( | |
| 276 &self, y : &mut DVector<F>, α : F, i : I, β : F | |
| 277 ) { | |
| 278 if β == 0.0 { | |
| 279 y.as_mut_slice().iter_mut().for_each(|x| *x = 0.0); | |
| 280 } else if β != 1.0 { | |
| 281 //*y *= β; | |
| 282 y.as_mut_slice().iter_mut().for_each(|x| *x *= β); | |
| 283 } | |
| 284 let h = self.h; | |
| 285 let m = self.len(); | |
| 286 i.eval(|x| { | |
| 287 assert_eq!(x.len(), m); | |
| 288 for d in 0..N { | |
| 289 let mut v = y.generic_view_mut((d*m, 0), (Dyn(m), U1)); | |
| 290 for b in Iter::new(&self.dims, d) { | |
| 291 self.discretisation.add_diff_mut(&mut v, x, α/h, b) | |
| 292 } | |
| 293 } | |
| 294 }) | |
| 295 } | |
| 296 } | |
| 297 | |
| 298 impl<F, B, const N : usize> Mapping<DVector<F>> | |
| 299 for Div<F, B, N> | |
| 300 where | |
| 301 B : Discretisation<F>, | |
| 302 F : Float + nalgebra::RealField, | |
| 303 { | |
| 304 type Codomain = DVector<F>; | |
| 305 fn apply<I : Instance<DVector<F>>>(&self, i : I) -> DVector<F> { | |
| 306 let mut y = DVector::zeros(self.len()); | |
| 307 self.apply_add(&mut y, i); | |
| 308 y | |
| 309 } | |
| 310 } | |
| 311 | |
| 312 #[replace_float_literals(F::cast_from(literal))] | |
| 313 impl<F, B, const N : usize> GEMV<F, DVector<F>> | |
| 314 for Div<F, B, N> | |
| 315 where | |
| 316 B : Discretisation<F>, | |
| 317 F : Float + nalgebra::RealField, | |
| 318 { | |
| 319 fn gemv<I : Instance<DVector<F>>>( | |
| 320 &self, y : &mut DVector<F>, α : F, i : I, β : F | |
| 321 ) { | |
| 322 if β == 0.0 { | |
| 323 y.as_mut_slice().iter_mut().for_each(|x| *x = 0.0); | |
| 324 } else if β != 1.0 { | |
| 325 //*y *= β; | |
| 326 y.as_mut_slice().iter_mut().for_each(|x| *x *= β); | |
| 327 } | |
| 328 let h = self.h; | |
| 329 let m = self.len(); | |
| 330 i.eval(|x| { | |
| 331 assert_eq!(x.len(), N * m); | |
| 332 for d in 0..N { | |
| 333 let v = x.generic_view((d*m, 0), (Dyn(m), U1)); | |
| 334 for b in Iter::new(&self.dims, d) { | |
| 335 self.discretisation.add_diff_mut(y, &v, α/h, b) | |
| 336 } | |
| 337 } | |
| 338 }) | |
| 339 } | |
| 340 } | |
| 341 | |
| 342 impl<F, B, const N : usize> Grad<F, B, N> | |
| 343 where | |
| 344 B : Discretisation<F>, | |
| 345 F : Float + nalgebra::RealField | |
| 346 { | |
| 347 /// Creates a new discrete gradient operator for the vector `u`, verifying dimensions. | |
| 348 pub fn new_for(u : &DVector<F>, h : F, dims : [usize; N], discretisation : B) | |
| 349 -> Option<Self> | |
| 350 { | |
| 351 if u.len() == dims.iter().product::<usize>() { | |
| 352 Some(Grad { dims, h, discretisation } ) | |
| 353 } else { | |
| 354 None | |
| 355 } | |
| 356 } | |
| 357 | |
| 358 fn len(&self) -> usize { | |
| 359 self.dims.iter().product::<usize>() | |
| 360 } | |
| 361 } | |
| 362 | |
| 363 | |
| 364 impl<F, B, const N : usize> Div<F, B, N> | |
| 365 where | |
| 366 B : Discretisation<F>, | |
| 367 F : Float + nalgebra::RealField | |
| 368 { | |
| 369 /// Creates a new discrete gradient operator for the vector `u`, verifying dimensions. | |
| 370 pub fn new_for(u : &DVector<F>, h : F, dims : [usize; N], discretisation : B) | |
| 371 -> Option<Self> | |
| 372 { | |
| 373 if u.len() == dims.iter().product::<usize>() * N { | |
| 374 Some(Div { dims, h, discretisation } ) | |
| 375 } else { | |
| 376 None | |
| 377 } | |
| 378 } | |
| 379 | |
| 380 fn len(&self) -> usize { | |
| 381 self.dims.iter().product::<usize>() | |
| 382 } | |
| 383 } | |
| 384 | |
| 385 impl<F, B, const N : usize> Linear<DVector<F>> | |
| 386 for Grad<F, B, N> | |
| 387 where | |
| 388 B : Discretisation<F>, | |
| 389 F : Float + nalgebra::RealField, | |
| 390 { | |
| 391 } | |
| 392 | |
| 393 impl<F, B, const N : usize> Linear<DVector<F>> | |
| 394 for Div<F, B, N> | |
| 395 where | |
| 396 B : Discretisation<F>, | |
| 397 F : Float + nalgebra::RealField, | |
| 398 { | |
| 399 } | |
| 400 | |
| 401 impl<F, B, const N : usize> BoundedLinear<DVector<F>, L2, L2, F> | |
| 402 for Grad<F, B, N> | |
| 403 where | |
| 404 B : Discretisation<F>, | |
| 405 F : Float + nalgebra::RealField, | |
| 406 DVector<F> : Norm<F, L2>, | |
| 407 { | |
| 408 fn opnorm_bound(&self, _ : L2, _ : L2) -> F { | |
| 409 // Fuck nalgebra. | |
| 410 self.discretisation.opnorm_bound(num_traits::Float::abs(self.h)) | |
| 411 } | |
| 412 } | |
| 413 | |
| 414 | |
| 415 impl<F, B, const N : usize> BoundedLinear<DVector<F>, L2, L2, F> | |
| 416 for Div<F, B, N> | |
| 417 where | |
| 418 B : Discretisation<F>, | |
| 419 F : Float + nalgebra::RealField, | |
| 420 DVector<F> : Norm<F, L2>, | |
| 421 { | |
| 422 fn opnorm_bound(&self, _ : L2, _ : L2) -> F { | |
| 423 // Fuck nalgebra. | |
| 424 self.discretisation.opnorm_bound(num_traits::Float::abs(self.h)) | |
| 425 } | |
| 426 } | |
| 427 | |
| 428 impl<F, B, const N : usize> | |
| 429 Adjointable<DVector<F>, DVector<F>> | |
| 430 for Grad<F, B, N> | |
| 431 where | |
| 432 B : Discretisation<F>, | |
| 433 F : Float + nalgebra::RealField, | |
| 434 { | |
| 435 type AdjointCodomain = DVector<F>; | |
| 436 type Adjoint<'a> = Div<F, B::Opposite, N> where Self : 'a; | |
| 437 | |
| 438 /// Form the adjoint operator of `self`. | |
| 439 fn adjoint(&self) -> Self::Adjoint<'_> { | |
| 440 Div { | |
| 441 dims : self.dims, | |
| 442 h : -self.h, | |
| 443 discretisation : self.discretisation.opposite(), | |
| 444 } | |
| 445 } | |
| 446 } | |
| 447 | |
| 448 | |
| 449 impl<F, B, const N : usize> | |
| 450 Adjointable<DVector<F>, DVector<F>> | |
| 451 for Div<F, B, N> | |
| 452 where | |
| 453 B : Discretisation<F>, | |
| 454 F : Float + nalgebra::RealField, | |
| 455 { | |
| 456 type AdjointCodomain = DVector<F>; | |
| 457 type Adjoint<'a> = Grad<F, B::Opposite, N> where Self : 'a; | |
| 458 | |
| 459 /// Form the adjoint operator of `self`. | |
| 460 fn adjoint(&self) -> Self::Adjoint<'_> { | |
| 461 Grad { | |
| 462 dims : self.dims, | |
| 463 h : -self.h, | |
| 464 discretisation : self.discretisation.opposite(), | |
| 465 } | |
| 466 } | |
| 467 } | |
| 468 | |
| 469 #[cfg(test)] | |
| 470 mod tests { | |
| 471 use super::*; | |
| 472 | |
| 473 #[test] | |
| 474 fn grad_adjoint() { | |
| 475 let im = DVector::from( (0..9).map(|t| t as f64).collect::<Vec<_>>()); | |
| 476 let v = DVector::from( (0..18).map(|t| t as f64).collect::<Vec<_>>()); | |
| 477 | |
| 478 let grad = Grad::new_for(&im, 1.0, [3, 3], ForwardNeumann).unwrap(); | |
| 479 let a = grad.apply(&im).dot(&v); | |
| 480 let b = grad.adjoint().apply(&v).dot(&im); | |
| 481 assert_eq!(a, b); | |
| 482 | |
| 483 let grad = Grad::new_for(&im, 1.0, [3, 3], ForwardDirichlet).unwrap(); | |
| 484 let a = grad.apply(&im).dot(&v); | |
| 485 let b = grad.adjoint().apply(&v).dot(&im); | |
| 486 assert_eq!(a, b); | |
| 487 | |
| 488 } | |
| 489 } |
