diff --git a/src/bspline.rs b/src/bspline.rs index 7d333e2..d7872fc 100644 --- a/src/bspline.rs +++ b/src/bspline.rs @@ -1,3 +1,4 @@ +use std::hash::Hash; use crate::Error; use crate::par_indexed_iter::ParallelIndexedIterMut; use crate::transform::{Transform, transform_point}; @@ -9,6 +10,10 @@ use num::traits::FloatConst; use num::{Complex, cast::AsPrimitive, complex::ComplexFloat}; use rayon::iter::ParallelIterator; use std::ops::{Deref, MulAssign}; +use std::sync::{Arc, Mutex}; +use rand::distr::StandardUniform; +use rand::RngExt; +use serde::{Deserialize, Serialize}; fn in_shape(shape_f: &[f64], v: &[f64]) -> bool { shape_f @@ -17,6 +22,117 @@ fn in_shape(shape_f: &[f64], v: &[f64]) -> bool { .all(|(s, x)| (-0.5 <= *x) && (x <= s)) } +#[derive(Debug, Serialize, Deserialize)] +pub struct DefaultValue { + #[serde(skip, default = "DefaultValue::make_rng")] + rng: Arc>, + values: Vec, + weights: Vec, +} + +impl Clone for DefaultValue { + fn clone(&self) -> Self { + Self { + rng: Self::make_rng(), + values: self.values.clone(), + weights: self.weights.clone(), + } + } +} + +impl Default for DefaultValue { + fn default() -> Self { + Self { + rng: Self::make_rng(), + values: vec![0.0], + weights: vec![1.0], + } + } +} + +impl PartialEq for DefaultValue { + fn eq(&self, other: &Self) -> bool { + (self.values == other.values) && (self.weights == other.weights) + } +} + +impl DefaultValue { + pub fn new(mut values: Vec, mut weights: Vec) -> Self { + if values.is_empty() { + values.push(f64::NAN); + } + if weights.len() != values.len() { + weights.resize(values.len(), 1.0); + } + weights.iter_mut().for_each(|w| { + if !w.is_finite() { + *w = 1.0; + } else if *w < 0.0 { + *w = 0.0 + } + }); + let n = weights.iter().sum::(); + weights.iter_mut().fold(0.0, |w, x| { + *x += w / n; + *x + }); + Self { + rng: Self::make_rng(), + values, + weights, + } + } + + pub fn values(&self) -> &[f64] { + &self.values + } + + pub fn weights(&self) -> &[f64] { + &self.weights + } + + pub fn sample(&self) -> f64 { + if self.values.len() == 1 { + self.values[0] + } else { + let i = self.rng.lock().unwrap().sample::(StandardUniform); + let n = self.weights.iter().position(|&w| i >= w).expect("weights should sum to 1"); + self.values[n] + } + } + + fn make_rng() -> Arc> { + Arc::new(Mutex::new(rand::make_rng())) + } +} + +impl Eq for DefaultValue {} + +impl Hash for DefaultValue { + fn hash(&self, state: &mut H) { + self.values.iter().for_each(|v| v.to_bits().hash(state)); + self.weights.iter().for_each(|v| v.to_bits().hash(state)); + } +} + +impl From for DefaultValue { + fn from(value: f64) -> Self { + Self::new(vec![value], vec![1.0]) + } +} + +impl From> for DefaultValue { + fn from(value: Vec) -> Self { + Self::new(value, Vec::new()) + } +} + +impl From<&[f64]> for DefaultValue { + fn from(value: &[f64]) -> Self { + Self::new(value.to_vec(), Vec::new()) + } +} + #[derive(Clone, Debug)] pub struct BSplineMem { ndim: usize, @@ -74,7 +190,7 @@ where D: Dimension, BSpline: BSplineTrait, { - pub fn interpolate(&self, transform: &Transform) -> Result, Error> { + pub fn interpolate(&self, transform: &Transform, default_value: &DefaultValue) -> Result, Error> { let mut mem = self.get_mem(); let shape_f = self.shape_f(); let center = shape_f.iter().map(|i| i / 2.0).collect::>(); @@ -91,13 +207,15 @@ where }; if in_shape(&shape_f, &v) { *x = self.evaluate_at_continuous_index(&v, &mut mem)?; + } else { + *x = default_value.sample(); } } Ok(new) } - pub fn interpolate_par(&self, transform: &Transform) -> Result, Error> { + pub fn interpolate_par(&self, transform: &Transform, default_value: &DefaultValue) -> Result, Error> { let mem = self.get_mem(); let shape_f = self.shape_f(); let center = shape_f.iter().map(|i| i / 2.0).collect::>(); @@ -114,6 +232,8 @@ where .all(|(s, x)| (-0.5 <= *x) && (x <= s)) { *x = self.evaluate_at_continuous_index(&v, &mut mem)?; + } else { + *x = default_value.sample(); } Ok(mem) }, diff --git a/src/lib.rs b/src/lib.rs index 453fe9d..e9ef912 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -84,7 +84,7 @@ mod tests { let c = [-0.8, 0.156]; let j = julia_image(&shape, &transform_j, ¢er, &c).mapv(|x| x as f64); let k = julia_image(&shape, &transform_k, ¢er, &c).mapv(|x| x as f64); - let n = transform.interpolate::<3, _, _>(k.view())?; + let n = transform.interpolate::<3, _, _>(k.view(), &Default::default())?; let sj = j.iter().sum::(); // let sk = k.iter().sum::(); let sn = n.iter().sum::(); @@ -124,7 +124,7 @@ mod tests { let c = [-0.8, 0.156]; let j = julia_image(&shape, &transform_j, ¢er, &c).mapv(|x| x as f64); let k = julia_image(&shape, &transform_k, ¢er, &c).mapv(|x| x as f64); - let n = transform.interpolate_par::<3, _, _>(k.view())?; + let n = transform.interpolate_par::<3, _, _>(k.view(), &Default::default())?; let sj = j.iter().sum::(); let sn = n.iter().sum::(); @@ -149,7 +149,7 @@ mod tests { &[-0.8, 0.156], ); let transform = Transform::from_rotation(1.0, &[3999.5, 2999.5]); - let image_b = transform.interpolate_par::<3, _, _>(image_a.view())?; + let image_b = transform.interpolate_par::<3, _, _>(image_a.view(), &Default::default())?; assert_ne!(image_b[[4000, 3000]], 0.0); Ok(()) } diff --git a/src/metric.rs b/src/metric.rs index d12a9ea..0a694ac 100644 --- a/src/metric.rs +++ b/src/metric.rs @@ -9,8 +9,7 @@ use itertools::Itertools; use ndarray::{Array, Array1, Array2, Array3, AsArray, Axis, Dimension, s}; use num::cast::AsPrimitive; use rand::RngExt; -use rayon::iter::IntoParallelIterator; -use rayon::iter::ParallelIterator; +use rayon::iter::{ParallelIterator, IntoParallelIterator}; use serde::{Deserialize, Serialize}; use std::cell::RefCell; use std::ops::Deref; diff --git a/src/transform.rs b/src/transform.rs index aa6add8..1a95633 100644 --- a/src/transform.rs +++ b/src/transform.rs @@ -1,4 +1,4 @@ -use crate::bspline::{BSpline, BSplineTrait}; +use crate::bspline::{BSpline, BSplineTrait, DefaultValue}; use crate::error::Error; use crate::filter::{fft, ifft}; use crate::metric::FixedMu; @@ -811,7 +811,7 @@ impl Transform { }) } - pub fn interpolate<'a, const B: usize, T, A>(&self, image: A) -> Result, Error> + pub fn interpolate<'a, const B: usize, T, A>(&self, image: A, default_value: &DefaultValue) -> Result, Error> where A: AsArray<'a, T, D>, D: Dimension, @@ -820,12 +820,13 @@ impl Transform { { let image = image.into(); let bspline = BSpline::::new(&image); - bspline.interpolate(self) + bspline.interpolate(self, default_value) } pub fn interpolate_par<'a, const B: usize, T, A>( &self, image: A, + default_value: &DefaultValue, ) -> Result, Error> where A: AsArray<'a, T, D>, @@ -835,13 +836,14 @@ impl Transform { { let image = image.into(); let bspline = BSpline::::new(&image); - bspline.interpolate_par(self) + bspline.interpolate_par(self, default_value) } pub fn interpolate_order<'a, T, A>( &self, image: A, order: usize, + default_value: &DefaultValue, ) -> Result, Error> where A: AsArray<'a, T, D>, @@ -849,12 +851,12 @@ impl Transform { T: 'a + Clone + AsPrimitive, { match order { - 0 => Ok(self.interpolate::<0, _, _>(image)?), - 1 => Ok(self.interpolate::<1, _, _>(image)?), - 2 => Ok(self.interpolate::<2, _, _>(image)?), - 3 => Ok(self.interpolate::<3, _, _>(image)?), - 4 => Ok(self.interpolate::<4, _, _>(image)?), - 5 => Ok(self.interpolate::<5, _, _>(image)?), + 0 => Ok(self.interpolate::<0, _, _>(image, default_value)?), + 1 => Ok(self.interpolate::<1, _, _>(image, default_value)?), + 2 => Ok(self.interpolate::<2, _, _>(image, default_value)?), + 3 => Ok(self.interpolate::<3, _, _>(image, default_value)?), + 4 => Ok(self.interpolate::<4, _, _>(image, default_value)?), + 5 => Ok(self.interpolate::<5, _, _>(image, default_value)?), _ => Err(Error::InterpolationNotImplemented(order)), } } @@ -863,6 +865,7 @@ impl Transform { &self, image: A, order: usize, + default_value: &DefaultValue, ) -> Result, Error> where A: AsArray<'a, T, D>, @@ -870,12 +873,12 @@ impl Transform { T: 'a + Clone + AsPrimitive, { match order { - 0 => Ok(self.interpolate_par::<0, _, _>(image)?), - 1 => Ok(self.interpolate_par::<1, _, _>(image)?), - 2 => Ok(self.interpolate_par::<2, _, _>(image)?), - 3 => Ok(self.interpolate_par::<3, _, _>(image)?), - 4 => Ok(self.interpolate_par::<4, _, _>(image)?), - 5 => Ok(self.interpolate_par::<5, _, _>(image)?), + 0 => Ok(self.interpolate_par::<0, _, _>(image, default_value)?), + 1 => Ok(self.interpolate_par::<1, _, _>(image, default_value)?), + 2 => Ok(self.interpolate_par::<2, _, _>(image, default_value)?), + 3 => Ok(self.interpolate_par::<3, _, _>(image, default_value)?), + 4 => Ok(self.interpolate_par::<4, _, _>(image, default_value)?), + 5 => Ok(self.interpolate_par::<5, _, _>(image, default_value)?), _ => Err(Error::InterpolationNotImplemented(order)), } } @@ -1045,8 +1048,8 @@ mod tests { &[-0.8, 0.156], ); let transform = Transform::from_rotation(1.0, &[3999.5, 2999.5]); - let image_b = transform.interpolate_par::<3, _, _>(image_a.view())?; - let image_c = transform.interpolate::<3, _, _>(image_a.view())?; + let image_b = transform.interpolate_par::<3, _, _>(image_a.view(), &Default::default())?; + let image_c = transform.interpolate::<3, _, _>(image_a.view(), &Default::default())?; assert!( image_b @@ -1066,7 +1069,7 @@ mod tests { &[-0.8, 0.156], ); let transform = Transform::from_rotation(1.0, &[3999.5, 2999.5]); - let image_b = transform.interpolate::<3, _, _>(image_a.view())?; + let image_b = transform.interpolate::<3, _, _>(image_a.view(), &Default::default())?; assert_ne!(image_b[[4000, 3000]], 0.0); Ok(()) } @@ -1080,7 +1083,7 @@ mod tests { &[-0.8, 0.156], ); let transform = Transform::from_rotation(1.0, &[3999.5, 2999.5]); - let image_b = transform.interpolate_par::<3, _, _>(image_a.view())?; + let image_b = transform.interpolate_par::<3, _, _>(image_a.view(), &Default::default())?; assert_ne!(image_b[[4000, 3000]], 0.0); Ok(()) } @@ -1099,7 +1102,7 @@ mod tests { .slice(s![.., 0]) .mapv(|i| i as f64); let q = vec![0.85, 4.0]; - let im_b = Transform::new(q).interpolate::<1, _, _>(&im_a)?; + let im_b = Transform::new(q).interpolate::<1, _, _>(&im_a, &Default::default())?; // Test with ALL points (like grad_check uses) let metric_all = MattesMetric::::new( @@ -1183,7 +1186,7 @@ mod tests { ) .slice(s![.., 0]) .mapv(|i| i as f64); - let im_b = Transform::new(vec![0.85, 4.0]).interpolate::<1, _, _>(&im_a)?; + let im_b = Transform::new(vec![0.85, 4.0]).interpolate::<1, _, _>(&im_a, &Default::default())?; // Smooth both images with sigma=8 like level 0 let sigma = Sigma::Absolute(vec![8.0]); @@ -1292,7 +1295,7 @@ mod tests { ]; let q = vec![0.85, 2.0]; - let moving = Transform::new(q.clone()).interpolate::<3, _, _>(&fixed)?; + let moving = Transform::new(q.clone()).interpolate::<3, _, _>(&fixed, &Default::default())?; // The registration finds T such that im_b(T(x)) = im_a(x), which is the inverse of q let q_inv = Transform::::new(q.clone()).inverse()?; @@ -1348,7 +1351,7 @@ mod tests { .slice(s![.., 0]) .mapv(|i| i as f64); let q = vec![0.85, 2.0]; - let im_b = Transform::new(q.clone()).interpolate::<1, _, _>(&im_a)?; + let im_b = Transform::new(q.clone()).interpolate::<1, _, _>(&im_a, &Default::default())?; let f = Sigma::Absolute(vec![1.0]).smooth(im_a.view())?; let m = Sigma::Absolute(vec![1.0]).smooth(im_b.view())?; @@ -1468,7 +1471,7 @@ mod tests { let params = vec![c * sx, -s * sy + shx, s * sx + shy, c * sy, tx, ty]; let transform = Transform::::new_around_center(params.clone(), ¢er); - let moving = transform.interpolate::<3, _, _>(fixed.view())?; + let moving = transform.interpolate::<3, _, _>(fixed.view(), &Default::default())?; let q_inv = transform.inverse()?.parameters; @@ -1510,7 +1513,7 @@ mod tests { ) .mapv(|i| i as f64); let rotation = Transform::::from_rotation(f64::PI() / 4.0, ¢er); - let im_b = rotation.interpolate::<3, _, _>(im_a.view())?; + let im_b = rotation.interpolate::<3, _, _>(im_a.view(), &Default::default())?; let q_inv = rotation.inverse()?.parameters; let p = rotation.parameters.clone(); @@ -1592,7 +1595,7 @@ mod tests { // Full resolution, sigma=4.0 — matching elastix FixedSmoothingImagePyramid (no downsampling) for angle_deg in [5.0f64, 10.0, 15.0, 20.0, 25.0, 30.0, 35.0, 40.0, 45.0] { let rotation = Transform::::from_rotation(angle_deg.to_radians(), ¢er); - let im_b = rotation.interpolate::<3, _, _>(im_a.view())?; + let im_b = rotation.interpolate::<3, _, _>(im_a.view(), &Default::default())?; let q_inv = rotation.inverse()?.parameters; let p = rotation.parameters.clone(); let f = gaussian_smooth(im_a.view(), &[4.0, 4.0])?; @@ -1727,7 +1730,7 @@ mod tests { fn register_real_images2() -> Result<(), Box> { let fixed = read_tiff("test_files/fixed.tif")?; let e = Transform::::new(vec![0.8, 0.0, 0.0, 1.0, 0.0, 0.0]).inverse()?; - let moving = e.interpolate::<3, _, _>(fixed.view())?; + let moving = e.interpolate::<3, _, _>(fixed.view(), &Default::default())?; let t = Transform::::register( fixed.view(),