diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..9b13019 --- /dev/null +++ b/.gitignore @@ -0,0 +1,79 @@ +/target +/Cargo.lock + +# Byte-compiled / optimized / DLL files +__pycache__/ +.pytest_cache/ +*.py[cod] + +# C extensions +*.so + +# Distribution / packaging +.Python +.venv/ +env/ +bin/ +build/ +develop-eggs/ +dist/ +eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +include/ +man/ +venv/ +*.egg-info/ +.installed.cfg +*.egg + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt +pip-selfcheck.json + +# Unit test / coverage reports +htmlcov/ +.tox/ +.coverage +.cache +nosetests.xml +coverage.xml + +# Translations +*.mo + +# Mr Developer +.mr.developer.cfg +.project +.pydevproject + +# Rope +.ropeproject + +# Django stuff: +*.log +*.pot + +.DS_Store + +# Sphinx documentation +docs/_build/ + +# PyCharm +.idea/ + +# VSCode +.vscode/ + +# Pyenv +.python-version + +AGENTS.md +.agentbridge + +*.tif +*.svg diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..4b593cc --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "image-registration" +version = "0.1.0" +edition = "2024" + +[dependencies] +algos = "0.6" +itertools = "0.14" +ndarray = { version = "0.17", features = ["rayon"] } +ndarray-linalg = { version = "0.18", features = ["openblas-static"] } +ndarray-npy = { version = "0.10.0", features = ["npz"] } +ndrustfft = "0.6" +num = "0.4" +rayon = "1" +serde = { version = "1", features = ["derive"] } +serde_yaml = "0.9" +rand = "0.9" +thiserror = "2" +tiffwrite = "2025.12.0" + +[dev-dependencies] +tempfile = "3" +tiffwrite = "2025.12.0" + +[profile.release] +debug = true diff --git a/src/bspline.rs b/src/bspline.rs index e69de29..1a38996 100644 --- a/src/bspline.rs +++ b/src/bspline.rs @@ -0,0 +1,1105 @@ +use crate::Error; +use crate::filter::{fft, ifft}; +use crate::par_indexed_iter::ParallelIndexedIterMut; +use crate::transform::{Transform, transform_point}; +use itertools::Itertools; +use ndarray::{ + Array, Array1, Array2, ArrayD, ArrayViewMut1, AsArray, Axis, Dimension, IntoDimension, IxDyn, + SliceInfoElem, +}; +use num::traits::FloatConst; +use num::{Complex, cast::AsPrimitive, complex::ComplexFloat}; +use rayon::iter::ParallelIterator; +use std::ops::{Deref, MulAssign}; +use tiffwrite::IJTiffFile; + +fn in_shape(shape_f: &[f64], v: &[f64]) -> bool { + shape_f + .iter() + .zip_eq(v.iter()) + .all(|(s, x)| (-0.5 <= *x) && (x <= s)) +} + +#[derive(Clone, Debug)] +pub struct BSplineMem { + ndim: usize, + evaluate_index: Array2, + weights: Array2, + derivative_weights: Array2, +} + +pub trait BSplineTrait { + fn new<'a, A, T>(array: A) -> Self + where + A: AsArray<'a, T, D>, + T: 'a + Clone + AsPrimitive, + Self: Sized; + + fn set_interpolation_weights(index: &[f64], mem: &mut BSplineMem); + + fn set_derivative_weights(index: &[f64], mem: &mut BSplineMem); + + /// reciprocal of table I B-Spline Signal Processing: Part II-Efficient Design and Applications, + /// Unser et al. 1993 + fn b_transfer<'a, A, B>(z: A) -> Array, B> + where + A: AsArray<'a, Complex, B>, + B: Dimension; +} + +#[derive(Clone, Debug)] +pub struct BSpline { + /// coefficients describing the bspline + pub coefficients: Array, + /// number of neighborhood points used for interpolation + pub max_number_interpolation_points: usize, + points_to_index: Array2, +} + +impl BSpline { + pub fn shape_f(&self) -> Vec { + self.shape().iter().map(|&i| i as f64 - 0.5).collect() + } + + pub fn get_mem(&self) -> BSplineMem { + let ndim = self.ndim(); + BSplineMem { + ndim, + evaluate_index: Array2::zeros((ndim, N + 1)), + weights: Array2::zeros((ndim, N + 1)), + derivative_weights: Array2::zeros((ndim, N + 1)), + } + } +} + +impl BSpline +where + D: Dimension, + BSpline: BSplineTrait, +{ + pub fn interpolate(&self, transform: &Transform) -> Result, Error> { + let mut mem = self.get_mem(); + let shape_f = self.shape_f(); + let mut new: Array = + Array::::zeros(self.shape()).into_dimensionality()?; + for (i, x) in new.indexed_iter_mut() { + let d = i.into_dimension(); + let j = d.as_array_view(); + let v = if let Some(k) = j.as_slice() { + transform_point(k, &transform.center, &transform.parameters) + } else { + let k = j.to_vec(); + transform_point(&k, &transform.center, &transform.parameters) + }; + if in_shape(&shape_f, &v) { + *x = self.evaluate_at_continuous_index(&v, &mut mem)?; + } + } + + Ok(new) + } + + pub fn interpolate_par(&self, transform: &Transform) -> Result, Error> { + let mem = self.get_mem(); + let shape_f = self.shape_f(); + let mut new: Array = + Array::::zeros(self.shape()).into_dimensionality()?; + new.par_indexed_iter_mut() + .try_fold( + || mem.clone(), + |mut mem, (i, x)| { + let v = transform_point(&i, &transform.center, &transform.parameters); + if shape_f + .iter() + .zip_eq(v.iter()) + .all(|(s, x)| (-0.5 <= *x) && (x <= s)) + { + *x = self.evaluate_at_continuous_index(&v, &mut mem)?; + } + Ok(mem) + }, + ) + .collect::, Error>>()?; + Ok(new) + } +} + +impl Deref for BSpline { + type Target = Array; + + fn deref(&self) -> &Self::Target { + &self.coefficients + } +} + +impl BSplineTrait for BSpline<0, D> { + fn new<'a, A, T>(array: A) -> Self + where + A: AsArray<'a, T, D>, + T: 'a + Clone + AsPrimitive, + { + // See Unser, 1997. Part II, Table I for Pole values. + // See also, Handbook of Medical Imaging, Processing and Analysis, Ed. Isaac + // N. Bankman, 2000, pg. 416. + let array = array.into(); + let max_number_interpolation_points = 1usize.pow(array.ndim() as u32); + let coefficients = array.mapv(|i| i.as_()); + let points_to_index = + Self::generate_points_to_index(max_number_interpolation_points, coefficients.ndim()); + Self { + coefficients, + max_number_interpolation_points, + points_to_index, + } + } + + fn set_interpolation_weights(_index: &[f64], mem: &mut BSplineMem) { + for n in 0..mem.ndim { + mem.weights[[n, 0]] = 1.0; + } + } + + fn set_derivative_weights(_index: &[f64], mem: &mut BSplineMem) { + for n in 0..mem.ndim { + mem.derivative_weights[[n, 0]] = 0.0; + } + } + + fn b_transfer<'a, A, B>(z: A) -> Array, B> + where + A: AsArray<'a, Complex, B>, + B: Dimension, + { + z.into().to_owned() + } +} + +impl BSplineTrait for BSpline<1, D> { + fn new<'a, A, T>(array: A) -> Self + where + A: AsArray<'a, T, D>, + T: 'a + Clone + AsPrimitive, + { + // See Unser, 1997. Part II, Table I for Pole values. + // See also, Handbook of Medical Imaging, Processing and Analysis, Ed. Isaac + // N. Bankman, 2000, pg. 416. + let array = array.into(); + let max_number_interpolation_points = 2usize.pow(array.ndim() as u32); + let coefficients = array.mapv(|i| i.as_()); + let points_to_index = + Self::generate_points_to_index(max_number_interpolation_points, coefficients.ndim()); + Self { + coefficients, + max_number_interpolation_points, + points_to_index, + } + } + + fn set_interpolation_weights(index: &[f64], mem: &mut BSplineMem) { + for ((idx, evaluate_index), mut weights) in index + .iter() + .zip_eq(mem.evaluate_index.rows_mut()) + .zip_eq(mem.weights.rows_mut()) + { + let w = idx - evaluate_index[0] as f64; + weights[1] = w; + weights[0] = 1.0 - w; + } + } + + fn set_derivative_weights(_index: &[f64], mem: &mut BSplineMem) { + for mut derivative_weights in mem.derivative_weights.rows_mut() { + derivative_weights[0] = -1.0; + derivative_weights[1] = 1.0; + } + } + + fn b_transfer<'a, A, B>(z: A) -> Array, B> + where + A: AsArray<'a, Complex, B>, + B: Dimension, + { + z.into().to_owned() + } +} + +impl BSplineTrait for BSpline<2, D> { + fn new<'a, A, T>(array: A) -> Self + where + A: AsArray<'a, T, D>, + T: 'a + Clone + AsPrimitive, + { + // See Unser, 1997. Part II, Table I for Pole values. + // See also, Handbook of Medical Imaging, Processing and Analysis, Ed. Isaac + // N. Bankman, 2000, pg. 416. + let array = array.into(); + let poles = vec![8f64.sqrt() - 3f64]; + let tolerance = 1e-10; + let max_number_interpolation_points = 3usize.pow(array.ndim() as u32); + let coefficients = Self::data_to_coefficients(tolerance, &poles, array); + let points_to_index = + Self::generate_points_to_index(max_number_interpolation_points, coefficients.ndim()); + Self { + coefficients, + max_number_interpolation_points, + points_to_index, + } + } + + fn set_interpolation_weights(index: &[f64], mem: &mut BSplineMem) { + for ((idx, evaluate_index), mut weights) in index + .iter() + .zip_eq(mem.evaluate_index.rows_mut()) + .zip_eq(mem.weights.rows_mut()) + { + let w = idx - evaluate_index[1] as f64; + weights[1] = 0.75 - w.powi(2); + weights[2] = 0.5 * (w - weights[1] + 1.0); + weights[0] = 1.0 - weights[1] - weights[2]; + } + } + + fn set_derivative_weights(index: &[f64], mem: &mut BSplineMem) { + for ((idx, evaluate_index), mut derivative_weights) in index + .iter() + .zip_eq(mem.evaluate_index.rows_mut()) + .zip_eq(mem.derivative_weights.rows_mut()) + { + let w = idx + 0.5 - evaluate_index[1] as f64; + let w1 = 1.0 - w; + derivative_weights[0] = 0.0 - w1; + derivative_weights[1] = w1 - w; + derivative_weights[2] = w; + } + } + + fn b_transfer<'a, A, B>(z: A) -> Array, B> + where + A: AsArray<'a, Complex, B>, + B: Dimension, + { + z.into().mapv(|i| (i + 6.0 + 1.0 / i) / 8.0) + } +} + +impl BSplineTrait for BSpline<3, D> { + fn new<'a, A, T>(array: A) -> Self + where + A: AsArray<'a, T, D>, + T: 'a + Clone + AsPrimitive, + { + // See Unser, 1997. Part II, Table I for Pole values. + // See also, Handbook of Medical Imaging, Processing and Analysis, Ed. Isaac + // N. Bankman, 2000, pg. 416. + let array = array.into(); + let poles = vec![3f64.sqrt() - 2f64]; + let tolerance = 1e-10; + let max_number_interpolation_points = 4usize.pow(array.ndim() as u32); + let coefficients = Self::data_to_coefficients(tolerance, &poles, array); + let points_to_index = + Self::generate_points_to_index(max_number_interpolation_points, coefficients.ndim()); + Self { + coefficients, + max_number_interpolation_points, + points_to_index, + } + } + + fn set_interpolation_weights(index: &[f64], mem: &mut BSplineMem) { + for ((idx, evaluate_index), mut weights) in index + .iter() + .zip_eq(mem.evaluate_index.rows_mut()) + .zip_eq(mem.weights.rows_mut()) + { + let w = idx - evaluate_index[1] as f64; + weights[3] = w.powi(3) / 6.0; + weights[0] = (1.0 / 6.0) + 0.5 * w * (w - 1.0) - weights[3]; + weights[2] = w + weights[0] - 2.0 * weights[3]; + weights[1] = 1.0 - weights[0] - weights[2] - weights[3]; + } + } + + fn set_derivative_weights(index: &[f64], mem: &mut BSplineMem) { + for ((idx, evaluate_index), mut derivative_weights) in index + .iter() + .zip_eq(mem.evaluate_index.rows_mut()) + .zip_eq(mem.derivative_weights.rows_mut()) + { + let w = idx + 0.5 - evaluate_index[2] as f64; + let w2 = 0.75 - w.powi(2); + let w3 = 0.5 * (w - w2 + 1.0); + let w1 = 1.0 - w2 - w3; + derivative_weights[0] = 0.0 - w1; + derivative_weights[1] = w1 - w2; + derivative_weights[2] = w2 - w3; + derivative_weights[3] = w3; + } + } + + fn b_transfer<'a, A, B>(z: A) -> Array, B> + where + A: AsArray<'a, Complex, B>, + B: Dimension, + { + z.into().mapv(|i| (i + 4.0 + 1.0 / i) / 6.0) + } +} + +impl BSplineTrait for BSpline<4, D> { + fn new<'a, A, T>(array: A) -> Self + where + A: AsArray<'a, T, D>, + T: 'a + Clone + AsPrimitive, + { + // See Unser, 1997. Part II, Table I for Pole values. + // See also, Handbook of Medical Imaging, Processing and Analysis, Ed. Isaac + // N. Bankman, 2000, pg. 416. + let array = array.into(); + let poles = vec![ + (664f64 - 438976f64.sqrt()).sqrt() + 304f64.sqrt() - 19f64, + (664f64 + 438976f64.sqrt()).sqrt() - 304f64.sqrt() - 19f64, + ]; + let tolerance = 1e-10; + let max_number_interpolation_points = 5usize.pow(array.ndim() as u32); + let coefficients = Self::data_to_coefficients(tolerance, &poles, array); + let points_to_index = + Self::generate_points_to_index(max_number_interpolation_points, coefficients.ndim()); + Self { + coefficients, + max_number_interpolation_points, + points_to_index, + } + } + + fn set_interpolation_weights(index: &[f64], mem: &mut BSplineMem) { + for ((idx, evaluate_index), mut weights) in index + .iter() + .zip_eq(mem.evaluate_index.rows_mut()) + .zip_eq(mem.weights.rows_mut()) + { + let w = idx - evaluate_index[2] as f64; + let w2 = w.powi(2); + let t = w2 / 6.0; + let t0 = w * (t - 11.0 / 24.0); + let t1 = 19.0 / 96.0 + w2 * (0.25 - t); + weights[0] = (0.5 - w).powi(2) / 24.0; + weights[1] = t1 + t0; + weights[3] = t1 - t0; + weights[4] = weights[0] + t0 + 0.5 * w; + weights[2] = 1.0 - weights[0] - weights[1] - weights[3] - weights[4]; + } + } + + fn set_derivative_weights(index: &[f64], mem: &mut BSplineMem) { + for ((idx, evaluate_index), mut derivative_weights) in index + .iter() + .zip_eq(mem.evaluate_index.rows_mut()) + .zip_eq(mem.derivative_weights.rows_mut()) + { + let w = idx + 0.5 - evaluate_index[2] as f64; + let w4 = w.powi(3) / 6.0; + let w1 = (1.0 / 6.0) + 0.5 * w * (w - 1.0) - w4; + let w3 = w + w1 - 2.0 * w4; + let w2 = 1.0 - w1 - w3 - w4; + derivative_weights[0] = 0.0 - w1; + derivative_weights[1] = w1 - w2; + derivative_weights[2] = w2 - w3; + derivative_weights[3] = w3 - w4; + derivative_weights[4] = w4; + } + } + + fn b_transfer<'a, A, B>(z: A) -> Array, B> + where + A: AsArray<'a, Complex, B>, + B: Dimension, + { + z.into() + .mapv(|i| (i.powi(2) + 76.0 * i + 230.0 + 76.0 / i + 1.0 / i.powi(2)) / 384.0) + } +} + +impl BSplineTrait for BSpline<5, D> { + fn new<'a, A, T>(array: A) -> Self + where + A: AsArray<'a, T, D>, + T: 'a + Clone + AsPrimitive, + { + // See Unser, 1997. Part II, Table I for Pole values. + // See also, Handbook of Medical Imaging, Processing and Analysis, Ed. Isaac + // N. Bankman, 2000, pg. 416. + let array = array.into(); + let poles = vec![ + (135f64 / 2f64 - (17745f64 / 4f64).sqrt()).sqrt() + (105f64 / 4f64).sqrt() + - 13f64 / 2f64, + (135f64 / 2f64 + (17745f64 / 4f64).sqrt()).sqrt() + - (105f64 / 4f64).sqrt() + - 13f64 / 2f64, + ]; + let tolerance = 1e-10; + let max_number_interpolation_points = 6usize.pow(array.ndim() as u32); + let coefficients = Self::data_to_coefficients(tolerance, &poles, array); + let points_to_index = + Self::generate_points_to_index(max_number_interpolation_points, coefficients.ndim()); + Self { + coefficients, + max_number_interpolation_points, + points_to_index, + } + } + + fn set_interpolation_weights(index: &[f64], mem: &mut BSplineMem) { + for ((idx, evaluate_index), mut weights) in index + .iter() + .zip_eq(mem.evaluate_index.rows_mut()) + .zip_eq(mem.weights.rows_mut()) + { + let mut w = idx - evaluate_index[2] as f64; + let mut w2 = w.powi(2); + weights[5] = (1.0 / 120.0) * w.powi(5); + w2 -= w; + let w4 = w2.powi(2); + w -= 0.5; + let t = w2 * (w2 - 3.0); + weights[0] = (1.0 / 24.0) * (1.0 / 5.0 + w2 + w4) - weights[5]; + let mut t0 = (1.0 / 24.0) * (w2 * (w2 - 5.0) + 46.0 / 5.0); + let mut t1 = (-1.0 / 12.0) * w * (t + 4.0); + weights[2] = t0 + t1; + weights[3] = t0 - t1; + t0 = (1.0 / 16.0) * (9.0 / 5.0 - t); + t1 = (1.0 / 24.0) * w * (w4 - w2 - 5.0); + weights[1] = t0 + t1; + weights[4] = t0 - t1; + } + } + + fn set_derivative_weights(index: &[f64], mem: &mut BSplineMem) { + for ((idx, evaluate_index), mut derivative_weights) in index + .iter() + .zip_eq(mem.evaluate_index.rows_mut()) + .zip_eq(mem.derivative_weights.rows_mut()) + { + let w = idx + 0.5 - evaluate_index[3] as f64; + let t2 = w.powi(2); + let t = t2 / 6.0; + let w1 = (0.5 - w).powi(4) / 24.0; + let t0 = w * (t - 11.0 / 24.0); + let t1 = 19.0 / 96.0 + t2 * (0.25 - t); + let w2 = t1 + t0; + let w4 = t1 - t0; + let w5 = w1 + t0 + 0.5 * w; + let w3 = 1.0 - w1 - w2 - w4 - w5; + derivative_weights[0] = 0.0 - w1; + derivative_weights[1] = w1 - w2; + derivative_weights[2] = w2 - w3; + derivative_weights[3] = w3 - w4; + derivative_weights[4] = w4 - w5; + derivative_weights[5] = w5; + } + } + + fn b_transfer<'a, A, B>(z: A) -> Array, B> + where + A: AsArray<'a, Complex, B>, + B: Dimension, + { + z.into() + .mapv(|i| (i.powi(2) + 26.0 * i + 66.0 + 26.0 / i + 1.0 / i.powi(2)) / 120.0) + } +} + +impl BSpline +where + D: Dimension, + BSpline: BSplineTrait, +{ + /// See Unser, 1993, Part II, Equation 2.5, + /// or Unser, 1999, Box 2. for an explanation. + pub fn data_to_coefficients<'a, A, T>(tolerance: f64, poles: &[f64], data: A) -> Array + where + A: AsArray<'a, T, D>, + T: 'a + Clone + AsPrimitive, + { + let mut coefficients = data.into().mapv(|i| i.as_()); + for axis in 0..coefficients.ndim() { + for mut line in coefficients.lanes_mut(Axis(axis)) { + let mut c0 = 1.0; + + // compute overall gain + for pole in poles.iter() { + c0 *= (1.0 - pole) * (1.0 - 1.0 / pole); + } + + // apply the gain + line *= c0; + + // loop over all poles + for pole in poles.iter() { + // causal initialization + Self::set_initial_causal_coefficient(tolerance, *pole, &mut line); + + // causal recursion + for n in 1..line.len() { + line[n] += pole * line[n - 1]; + } + + // anticausal initialization + Self::set_initial_anti_causal_coefficient(*pole, &mut line); + + // anticausal recursion + for n in (0..line.len() - 1).rev() { + line[n] = pole * (line[n + 1] - line[n]); + } + } + } + } + coefficients + } + + /// See Unser, 1999, Box 2 for explanation + fn set_initial_causal_coefficient(tolerance: f64, z: f64, line: &mut ArrayViewMut1) { + let mut zn = z; + if tolerance > 0.0 { + let horizon = (tolerance.ln() / z.abs().ln()).ceil() as usize; + if horizon < line.len() { + // accelerated loop + let mut sum = line[0]; + for n in 1..horizon { + sum += zn * line[n]; + zn *= z; + } + line[0] = sum; + return; // return early + } + } + + // full loop + let iz = 1.0 / z; + let mut z2n = z.powi((line.len() - 1) as i32); + let mut sum = line[0] + z2n * line[line.len() - 1]; + z2n = z2n * z2n * iz; + for n in 1..=line.len() - 2 { + sum += (zn + z2n) * line[n]; + zn *= z; + z2n *= iz; + } + line[0] = sum / (1.0 - zn * zn); + } + + /// This initialization corresponds to mirror boundaries. + /// See Unser, 1999, Box 2 for explanation. + /// Also see erratum at http://bigwww.epfl.ch/publications/unser9902.html + fn set_initial_anti_causal_coefficient(z: f64, line: &mut ArrayViewMut1) { + let n = line.len(); + if n >= 2 { + line[n - 1] = (z / (z * z - 1.0)) * (z * line[n - 2] + line[n - 1]); + } else if n == 1 { + line[n - 1] *= z / (z * z - 1.0); + } + } + + pub fn evaluate_at_continuous_index( + &self, + index: &[f64], + mem: &mut BSplineMem, + ) -> Result { + let coefficients = self.coefficients.view().into_dyn(); + + // compute the interpolation indexes + Self::determine_region_of_support(index, mem); + + // determine weights + Self::set_interpolation_weights(index, mem); + + // modify evaluate_index at the boundaries using mirror boundary conditions + Self::apply_mirror_boundary_conditions(coefficients.shape(), mem); + + // perform interpolation + let mut interpolated = 0.0; + let mut coefficient_index = vec![0; mem.ndim]; + + // step through each point in the n-dimensional interpolation cube. + for p in 0..self.max_number_interpolation_points { + let mut w = 1.0; + for (n, c) in coefficient_index.iter_mut().enumerate() { + let indx = self.points_to_index[[p, n]]; + w *= mem.weights[[n, indx]]; + *c = mem.evaluate_index[[n, indx]] as usize; + } + interpolated += w * coefficients[IxDyn(&coefficient_index)]; + } + Ok(interpolated) + } + + pub fn evaluate_derivative_at_continuous_index( + &self, + index: &[f64], + mem: &mut BSplineMem, + ) -> Result, Error> { + let ndim = self.coefficients.ndim(); + let coefficients = self.coefficients.view().into_dyn(); + + // compute the interpolation indexes + Self::determine_region_of_support(index, mem); + + // determine weights + Self::set_interpolation_weights(index, mem); + + // determine derivative weights + Self::set_derivative_weights(index, mem); + + // modify evaluate_index at the boundaries using mirror boundary conditions + Self::apply_mirror_boundary_conditions(coefficients.shape(), mem); + + // perform interpolation + let mut derivative = vec![0.0; ndim]; + let mut coefficient_index = vec![0; ndim]; + for (n, d) in derivative.iter_mut().enumerate() { + // step through each point in the n-dimensional interpolation cube. + for p in 0..self.max_number_interpolation_points { + let mut w1 = 1.0; + for (m, c) in coefficient_index.iter_mut().enumerate() { + let indx = self.points_to_index[[p, m]]; + *c = mem.evaluate_index[[m, indx]] as usize; + if n == m { + w1 *= mem.derivative_weights[[m, indx]]; + } else { + w1 *= mem.weights[[m, indx]]; + } + } + *d += w1 * coefficients[IxDyn(&coefficient_index)]; + } + } + Ok(derivative) + } + + pub fn evaluate_value_and_derivative_at_continuous_index( + &self, + index: &[f64], + mem: &mut BSplineMem, + ) -> Result<(f64, Vec), Error> { + let ndim = self.coefficients.ndim(); + let coefficients = self.coefficients.view().into_dyn(); + + // compute the interpolation indexes + Self::determine_region_of_support(index, mem); + + // determine weights + Self::set_interpolation_weights(index, mem); + + // determine derivative weights + Self::set_derivative_weights(index, mem); + + // modify evaluate_index at the boundaries using mirror boundary conditions + Self::apply_mirror_boundary_conditions(coefficients.shape(), mem); + + // perform interpolation + let mut interpolated = 0.0; + let mut coefficient_index = vec![0; ndim]; + + // step through each point in the n-dimensional interpolation cube. + for p in 0..self.max_number_interpolation_points { + let mut w = 1.0; + for (n, c) in coefficient_index.iter_mut().enumerate() { + let indx = self.points_to_index[[p, n]]; + w *= mem.weights[[n, indx]]; + *c = mem.evaluate_index[[n, indx]] as usize; + } + interpolated += w * coefficients[IxDyn(&coefficient_index)]; + } + + // perform interpolation + let mut derivative = vec![0.0; ndim]; + let mut coefficient_index = vec![0; ndim]; + for (n, d) in derivative.iter_mut().enumerate() { + // step through each point in the n-dimensional interpolation cube. + for p in 0..self.max_number_interpolation_points { + let mut w1 = 1.0; + for (m, c) in coefficient_index.iter_mut().enumerate() { + let indx = self.points_to_index[[p, m]]; + *c = mem.evaluate_index[[m, indx]] as usize; + if n == m { + w1 *= mem.derivative_weights[[m, indx]]; + } else { + w1 *= mem.weights[[m, indx]]; + } + } + *d += w1 * coefficients[IxDyn(&coefficient_index)]; + } + } + Ok((interpolated, derivative)) + } + + /// points_to_index is used to convert a sequential location to an N-dimension + /// index vector. This is precomputed to save time during the interpolation + /// routine. + fn generate_points_to_index( + max_number_interpolation_points: usize, + ndim: usize, + ) -> Array2 { + let mut points_to_index = Array2::::zeros([max_number_interpolation_points, ndim]); + for p in 0..max_number_interpolation_points { + let mut pp = p; + let mut index_factor = vec![0; ndim]; + index_factor[0] = 1; + for j in 1..ndim { + index_factor[j] = index_factor[j - 1] * (N + 1); + } + for j in (0..ndim).rev() { + points_to_index[[p, j]] = pp / index_factor[j]; + pp %= index_factor[j]; + } + } + points_to_index + } + + fn determine_region_of_support(index: &[f64], mem: &mut BSplineMem) { + let order = N as isize; + let half_offset = if order & 1 == 1 { 0.0 } else { 0.5 }; + for (i, mut r) in index.iter().zip_eq(mem.evaluate_index.rows_mut()) { + let indx = ((i + half_offset).floor() as isize) - order / 2; + for (j, k) in r.iter_mut().take(N + 1).enumerate() { + *k = (j as isize) + indx; + } + } + } + + fn apply_mirror_boundary_conditions(shape: &[usize], mem: &mut BSplineMem) { + // apply the mirror boundary conditions + // TODO: We could implement other boundary options beside mirror + for (s, mut r) in shape.iter().zip_eq(mem.evaluate_index.rows_mut()) { + if *s == 0 { + for k in r.iter_mut() { + *k = 0; + } + } else { + let end = (*s as isize) - 1; + for k in r.iter_mut() { + if *k < 0 { + *k = -*k; + } + if *k >= end { + *k = end - (*k - end); + } + } + } + } + } + + pub fn subsample2(&self) -> Result { + let shape = self.coefficients.shape(); + let v = Self::reduction_filter2(shape)?; + let c = fft(self.coefficients.view())?; + let mut t = IJTiffFile::new(std::env::home_dir().unwrap().join("tmp/fft.tif")).unwrap(); + let v0 = v.mapv(|i| i.re).into_dimensionality()?; + let c0 = c.mapv(|i| i.re).into_dimensionality()?; + let filtered = ifft((v * c).view())?.into_dyn().mapv(|i| i.re); + t.save(v0.view(), 0, 0, 0).unwrap(); + t.save(c0.view(), 1, 0, 0).unwrap(); + t.save(filtered.view().into_dimensionality()?.view(), 2, 0, 0) + .unwrap(); + + let slice = vec![ + SliceInfoElem::Slice { + start: 0, + end: None, + step: 2 + }; + filtered.ndim() + ]; + let coefficients = filtered + .slice(slice.as_slice()) + .to_owned() + .into_dimensionality()?; + Ok(BSpline { + coefficients, + max_number_interpolation_points: self.max_number_interpolation_points, + points_to_index: self.points_to_index.clone(), + }) + } + + pub fn evaluate(&self, index: Array) -> Result, Error> + where + D: Dimension, + { + let shape = index.shape()[1..].to_vec(); + let mut output = ArrayD::zeros(shape); + let mut mem = self.get_mem(); + for (i, o) in index.lanes(Axis(0)).into_iter().zip_eq(output.iter_mut()) { + let j = i.to_vec(); + *o = self.evaluate_at_continuous_index(&j, &mut mem)?; + } + Ok(output.into_dimensionality()?) + } + + /// eq 2.8 The L2 Polynomial Spline Pyramid, Unser et al. 1993 + fn u<'a, A, B>(m: usize, z: A) -> Array, B> + where + A: AsArray<'a, Complex, B>, + B: Dimension, + { + let n = N as u32; + let m = m as u32; + let k0 = (n + 1) * (m - 1) / 2; + z.into().mapv(|i| { + i.powu(k0) + * (0..m as i32) + .map(|j| i.powi(-j)) + .sum::>() + .powu(n + 1) + }) / (m as f64).powi(n as i32) + } + + /// table III The L2 Polynomial Spline Pyramid, Unser et al. 1993 + pub(crate) fn reduction_filter2(shape: &[usize]) -> Result, D>, Error> { + let mut filter: Array, D> = ArrayD::ones(shape).into_dimensionality()?; + for (i, &s) in shape.iter().enumerate() { + let si = s as isize; + let sf = s as f64; + let h = (si + 1) / 2; + let f = (h..h + si) + .map(|k| (((k % si) - h) as f64) / sf) + .collect::>(); + let z = f.mapv(|i| { + let theta = 2.0 * i * f64::PI(); + Complex::new(theta.cos(), theta.sin()) + }); + let z2 = f.mapv(|i| { + let theta = 4.0 * i * f64::PI(); + Complex::new(theta.cos(), theta.sin()) + }); + let v = Self::u(2, z.view()) + * (Self::b_transfer(z.view()) / Self::b_transfer(z2.view())) + .mapv(|i| i.powi(2 * N as i32 + 2) / 2.0); + for mut lane in filter.lanes_mut(Axis(i)) { + lane.mul_assign(&v); + } + } + Ok(filter) + } +} + +#[inline] +pub fn cubic_bspline(x: f64) -> f64 { + (x + 2.0).max(0.0).powi(3) / 6.0 - (x + 1.0).max(0.0).powi(3) / 1.5 + x.max(0.0).powi(3) + - (x - 1.0).max(0.0).powi(3) / 1.5 + + (x - 2.0).max(0.0).powi(3) / 6.0 +} + +#[inline] +pub fn square_bspline(x: f64) -> f64 { + (x + 1.5).max(0.0).powi(2) / 2.0 - (x + 0.5).max(0.0).powi(2) * 1.5 + + (x - 0.5).max(0.0).powi(2) * 1.5 + - (x - 1.5).max(0.0).powi(2) / 2.0 +} + +#[inline] +pub fn square_bspline_integral(x: f64) -> f64 { + (x + 1.5).max(0.0).powi(3) / 6.0 - (x + 0.5).max(0.0).powi(3) / 2.0 + + (x - 0.5).max(0.0).powi(3) / 2.0 + - (x - 1.5).max(0.0).powi(3) / 6.0 +} + +pub fn cubic_bspline_parzen_1d( + data: &[f64], + first_bin: f64, + last_bin: f64, + n_bins: usize, + e: f64, +) -> (Array1, Array1, Array1) { + let n_bins_i = n_bins as isize; + let mut hist = Array1::zeros(n_bins); + let mut grad = Array1::zeros(n_bins); + let d = ((n_bins - 1) as f64) / (last_bin - first_bin); + let m = (2.0 * d / e) as isize; + let bins = Array1::linspace(first_bin, last_bin, n_bins); + for x in data { + let n = (d * (x - first_bin)).round() as isize; + for i in (n - m)..=(n + m) { + if (0 <= i) && (i < n_bins_i) { + let j = i as usize; + let k = (x - bins[j]) / e; + hist[j] += cubic_bspline(k); + grad[j] += square_bspline(k - 0.5) - square_bspline(k + 0.5); + } + } + } + let alpha = (data.len() as f64) * e * d; + (bins, hist / alpha, grad / alpha) +} + +pub fn cubic_bspline_joint_parzen( + data_a: &[f64], + data_b: &[f64], + first_bin: f64, + last_bin: f64, + n_bins: usize, + e: f64, +) -> (Array1, Array2, Array2, Array2) { + let n_bins_i = n_bins as isize; + let mut hist = Array2::zeros([n_bins, n_bins]); + let mut grad_a = Array2::zeros([n_bins, n_bins]); + let mut grad_b = Array2::zeros([n_bins, n_bins]); + let d = ((n_bins - 1) as f64) / (last_bin - first_bin); + let m = (2.0 * d / e) as isize; + let bins = Array1::linspace(first_bin, last_bin, n_bins); + for xa in data_a { + for xb in data_b { + let na = (d * (xa - first_bin)).round() as isize; + let nb = (d * (xb - first_bin)).round() as isize; + for ia in (na - m)..=(na + m) { + for ib in (nb - m)..=(nb + m) { + if (0 <= ia) && (ia < n_bins_i) && (0 <= ib) && (ib < n_bins_i) { + let ja = ia as usize; + let jb = ib as usize; + let ka = (xa - bins[ja]) / e; + let kb = (xb - bins[jb]) / e; + let csa = cubic_bspline(ka); + let csb = cubic_bspline(kb); + hist[[ja, jb]] += csa * csb; + grad_a[[ja, jb]] += + (square_bspline(ka - 0.5) - square_bspline(ka + 0.5)) * csb; + grad_b[[ja, jb]] += + (square_bspline(kb - 0.5) - square_bspline(kb + 0.5)) * csa; + } + } + } + } + } + let alpha = (data_a.len() as f64) * e * d * (data_b.len() as f64) * e * d; + (bins, hist / alpha, grad_a / alpha, grad_b / alpha) +} + +pub fn cubic_bspline_joint_parzen_grad_b( + data_a: &[f64], + data_b: &[f64], + first_bin: f64, + last_bin: f64, + n_bins: usize, + e: f64, +) -> (Array1, Array2, Array2) { + let n_bins_i = n_bins as isize; + let mut hist = Array2::zeros([n_bins, n_bins]); + let mut grad_b = Array2::zeros([n_bins, n_bins]); + let d = ((n_bins - 1) as f64) / (last_bin - first_bin); + let m = (2.0 * d / e) as isize; + let bins = Array1::linspace(first_bin, last_bin, n_bins); + for (xa, xb) in data_a.iter().zip_eq(data_b.iter()) { + let na = (d * (xa - first_bin)).round() as isize; + let nb = (d * (xb - first_bin)).round() as isize; + for ia in (na - m)..=(na + m) { + if (0 <= ia) && (ia < n_bins_i) { + let ja = ia as usize; + let ka = (xa - bins[ja]) / e; + let csa = cubic_bspline(ka); + for ib in (nb - m)..=(nb + m) { + if (0 <= ib) && (ib < n_bins_i) { + let jb = ib as usize; + let kb = (xb - bins[jb]) / e; + hist[[ja, jb]] += csa * cubic_bspline(kb); + grad_b[[ja, jb]] += + csa * (square_bspline(kb - 0.5) - square_bspline(kb + 0.5)); + } + } + } + } + } + let alpha = (data_a.len() as f64) * e * d; + // println!("alpha: {}", alpha); + (bins, hist / alpha, grad_b / alpha) +} + +#[cfg(test)] +mod test { + use super::{BSpline, BSplineTrait, cubic_bspline_joint_parzen_grad_b}; + use crate::error::Error; + use crate::julia_image; + use ndarray::{Array1, Array3, Ix1, array, s}; + use ndarray_npy::NpzWriter; + use std::fs::File; + + #[test] + fn bspline_grad() -> Result<(), Box> { + let a = vec![0.0, 0.0, 1.0, 0.0, 1.0, 5.0, 1.0, 0.0, 1.0, 0.0, 0.0]; + let first_bin = 0.0; + let last_bin = 5.0; + let n_bins = 3; + let e = (last_bin - first_bin) / (n_bins - 1) as f64; + let (bins, jpdf, _d_jpdf_m) = + cubic_bspline_joint_parzen_grad_b(&a, &a, first_bin, last_bin, n_bins, e); + println!("bins: {:?}, jpdf: {:?}", bins, jpdf); + println!("sum: {}", jpdf.sum()); + Ok(()) + } + + #[test] + fn reduction_filter() -> Result<(), Error> { + let _v = BSpline::<3, Ix1>::reduction_filter2(&[10]); + Ok(()) + } + + #[test] + fn subsample2() -> Result<(), Box> { + let shape = [600, 800]; + let transform = [1.0, 0.0, 0.0, 1.0, 0.0, 0.0]; + let center = shape + .iter() + .map(|&s| (s - 1) as f64 / 2.0) + .collect::>(); + let center = [center[0], center[1]]; + let k = julia_image(&shape, &transform, ¢er, &[-0.4, 0.6]); + let b = BSpline::<3, _>::new(k.view()); + let b2 = b.subsample2()?; + let shape = k.shape().iter().map(|i| i / 2).collect::>(); + let mut xy = Array3::zeros((2, shape[0], shape[1])); + for (i, mut x) in xy + .slice_mut(s![1, .., ..]) + .columns_mut() + .into_iter() + .enumerate() + { + x.fill(i as f64); + } + for (i, mut y) in xy + .slice_mut(s![0, .., ..]) + .rows_mut() + .into_iter() + .enumerate() + { + y.fill(i as f64); + } + let k2 = b2.evaluate(xy)?.mapv(|i| i.clamp(0.0, 255.0) as u8); + let mut t1 = + tiffwrite::IJTiffFile::new(std::env::home_dir().unwrap().join("tmp/subsample1.tif"))?; + // t1.save(b.coefficients.view(), 0, 0, 0)?; + t1.save(k.view(), 0, 0, 0)?; + + let mut t2 = + tiffwrite::IJTiffFile::new(std::env::home_dir().unwrap().join("tmp/subsample2.tif"))?; + // t2.save(b2.coefficients.view(), 0, 0, 0)?; + t2.save(k2.view(), 0, 0, 0)?; + Ok(()) + } + + #[test] + fn derivative() -> Result<(), Box> { + let a = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0]; + let b = BSpline::<3, _>::new(a.view()); + let mut mem = b.get_mem(); + let x = Array1::linspace(-0.5, 10.5, 500); + let mut v = Vec::new(); + let mut d = Vec::new(); + for i in &x { + let (vi, di) = b.evaluate_value_and_derivative_at_continuous_index(&[*i], &mut mem)?; + v.push(vi); + d.push(di[0]); + } + let mut npz = NpzWriter::new(File::create( + std::env::home_dir().unwrap().join("tmp/metric.npz"), + )?); + npz.add_array("x", &x)?; + npz.add_array("v", &Array1::from(v))?; + npz.add_array("d", &Array1::from(d))?; + Ok(()) + } +} diff --git a/src/error.rs b/src/error.rs index e69de29..c013d2c 100644 --- a/src/error.rs +++ b/src/error.rs @@ -0,0 +1,17 @@ +use thiserror::Error; + +#[derive(Debug, Error)] +pub enum Error { + #[error(transparent)] + IO(#[from] std::io::Error), + #[error(transparent)] + SerdeYAML(#[from] serde_yaml::Error), + #[error(transparent)] + ShapeError(#[from] ndarray::ShapeError), + #[error(transparent)] + LinAlg(#[from] ndarray_linalg::error::LinalgError), + #[error(transparent)] + NpyError(#[from] ndarray_npy::WriteNpzError), + #[error("number of dimensions is not defined")] + NumberOfDimensionsNotDefined, +} diff --git a/src/filter.rs b/src/filter.rs index e69de29..8f4a5d8 100644 --- a/src/filter.rs +++ b/src/filter.rs @@ -0,0 +1,277 @@ +use crate::error::Error; +use itertools::Itertools; +use ndarray::{Array, Array1, ArrayD, AsArray, Axis, Dimension, RemoveAxis, concatenate, s}; +use ndrustfft::{FftHandler, Normalization, ndfft_par, ndifft_par}; +use num::Complex; +use num::traits::FloatConst; +use std::ops::MulAssign; + +/// Fourier transform +pub fn fft<'a, A, D>(array: A) -> Result, D>, Error> +where + A: AsArray<'a, f64, D>, + D: Dimension, +{ + let mut input = array.into().mapv(|i| Complex::new(i, 0.0)); + let shape = input.shape().to_vec(); + let mut tmp = ArrayD::zeros(shape.clone()).into_dimensionality()?; + for (i, s) in shape.iter().enumerate() { + let handler = FftHandler::new(*s).normalization(Normalization::None); + ndfft_par(&input.view(), &mut tmp.view_mut(), &handler, i); + std::mem::swap(&mut input, &mut tmp); + } + Ok(input) +} + +/// inverse Fourier transform +pub fn ifft<'a, A, D>(array: A) -> Result, D>, Error> +where + A: AsArray<'a, Complex, D>, + D: Dimension, +{ + let mut input = array.into().to_owned(); + let shape = input.shape().to_vec(); + let mut tmp = ArrayD::zeros(shape.clone()).into_dimensionality()?; + for (i, s) in shape.iter().enumerate() { + let handler = FftHandler::new(*s).normalization(Normalization::None); + ndifft_par(&input.view(), &mut tmp.view_mut(), &handler, i); + std::mem::swap(&mut input, &mut tmp); + } + Ok(input / Complex::from(shape.iter().product::() as f64)) +} + +pub fn fft_freq(size: usize) -> Array1 { + // let s = size as isize; + // let h = s / 2; + // (-h..h).map(|i| (i % s - s) as f64 / s as f64).collect() + + let val = 1.0 / size as f64; + let mut results = Array1::zeros(size); + let n = (size - 1) / 2 + 1; + let p1 = Array1::range(0.0, n as f64, 1.0); + results.slice_mut(s![..n]).assign(&p1); + let p2 = Array1::range(-((size / 2) as f64), 0.0, 1.0); + results.slice_mut(s![n..]).assign(&p2); + results * val +} + +pub fn fft_shift<'a, A, T, D>(x: A) -> Array +where + A: AsArray<'a, T, D>, + T: 'a + Clone, + D: Dimension + RemoveAxis, +{ + let mut x = x.into().to_owned(); + let shift = x.shape().iter().map(|i| i / 2).collect::>(); + for (i, s) in shift.iter().enumerate() { + let (a, b) = x.view().split_at(Axis(i), *s); + x = concatenate(Axis(i), &[b, a]).unwrap(); + } + x +} + +pub fn gaussian(x: &[f64], mu: f64, sigma: f64) -> Vec { + let a = 2.0 * sigma.powi(2); + let b = (a * f64::PI()).sqrt(); + x.iter() + .map(|i| (-(i - mu).powi(2) / a).exp() / b) + .collect() +} + +/// Gaussian kernel in frequency space +pub fn gaussian_kernel(shape: &[usize], sigma: &[f64]) -> Result, D>, Error> +where + D: Dimension, +{ + let mut g = Array::ones(shape).into_dimensionality()?; + for (i, (s, t)) in shape.iter().zip_eq(sigma.iter()).enumerate() { + let a = 0.5 / t.powi(2); + let f = (-(f64::PI() * fft_freq(*s)).powi(2) / a) + .exp() + .mapv(Complex::from); + for mut lane in g.lanes_mut(Axis(i)) { + lane.mul_assign(&f); + } + } + Ok(g) +} + +/// smooth an array using a Gaussian kernel +pub fn gaussian_smooth<'a, A, D>(array: A, sigma: &[f64]) -> Result, Error> +where + A: AsArray<'a, f64, D>, + D: Dimension, +{ + let array = array.into().to_owned(); + let shape = array.shape(); + let kernel = gaussian_kernel::(shape, sigma)?; + Ok(ifft((kernel * fft(array.view())?).view())?.mapv(|i| i.re)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::julia_image; + use ndarray::{Ix1, Ix2, array}; + use tiffwrite::IJTiffFile; + + #[test] + fn smooth() -> Result<(), Box> { + let a = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0]; + let b = gaussian_smooth(a.view(), &[1.4])?; + let c = array![ + 0.13775053, 0.26443112, 0.45036765, 0.67622086, 0.93678676, 1.06888618, 0.93678676, + 0.67622086, 0.45036765, 0.26443112, 0.13775053 + ]; + debug_assert!(b.iter().zip_eq(c.iter()).all(|(x, y)| (x - y).abs() < 1e-8)); + Ok(()) + } + + #[test] + fn smooth2() -> Result<(), Box> { + let im_a = julia_image( + &[60, 80], + &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0], + &[29.5, 39.5], + &[-0.8, 0.156], + ); + let im_b = gaussian_smooth(im_a.mapv(|i| i as f64).view(), &[4.0, 4.0])?; + let mut t = IJTiffFile::new(std::env::home_dir().unwrap().join("tmp/julia.tif"))?; + t.save(im_a.mapv(|i| i as u32).view(), 0, 0, 0)?; + t.save(im_b.mapv(|i| i as u32).view(), 1, 0, 0)?; + Ok(()) + } + + #[test] + fn gaussian_kernel_test2() -> Result<(), Box> { + let k = gaussian_kernel::(&[60, 80], &[16.0, 16.0])?; + ndarray_npy::write_npy("/home/wim/tmp/kernel.npy", &k)?; + Ok(()) + } + + #[test] + fn fft_test() -> Result<(), Box> { + let x = Array1::linspace(0.0, 1.0, 11).to_vec(); + let y = gaussian(x.as_slice(), 0.4, 0.1); + let j = fft(Array1::from_vec(y).view())?; + let k = array![ + Complex::new(9.99998512536899, 0.0), + Complex::new(-5.562906924752932, -6.419930454175756), + Complex::new(-0.7410757333589753, 5.154238692693016), + Complex::new(1.9379781953109205, -1.2454762622580575), + Complex::new(-0.7086874511902606, -0.20810344788799306), + Complex::new(0.08206001372382088, 0.1796510861631586), + Complex::new(0.08206001372382088, -0.1796510861631586), + Complex::new(-0.7086874511902606, 0.20810344788799306), + Complex::new(1.9379781953109205, 1.2454762622580575), + Complex::new(-0.7410757333589753, -5.154238692693016), + Complex::new(-5.562906924752932, 6.419930454175756), + ]; + debug_assert!( + j.iter() + .zip_eq(k.iter()) + .all(|(x, y)| (x.re - y.re) < 1e-8 && (x.im - y.im) < 1e-8) + ); + Ok(()) + } + + #[test] + fn ifft_test() -> Result<(), Box> { + let x = Array1::linspace(0.0, 1.0, 11).to_vec(); + let y = gaussian(x.as_slice(), 0.4, 0.1) + .into_iter() + .map(Complex::from) + .collect::>(); + let j = ifft(Array1::from_vec(y).view())?; + let k = array![ + Complex::new(0.9090895568517264, 0.0), + Complex::new(-0.5057188113411757, 0.5836300412887051), + Complex::new(-0.06737052121445229, -0.4685671538811833), + Complex::new(0.17617983593735642, 0.1132251147507325), + Complex::new(-0.06442613192638733, 0.018918495262544823), + Complex::new(0.00746000124762008, -0.01633191692392351), + Complex::new(0.00746000124762008, 0.01633191692392351), + Complex::new(-0.06442613192638733, -0.018918495262544823), + Complex::new(0.17617983593735642, -0.1132251147507325), + Complex::new(-0.06737052121445229, 0.4685671538811833), + Complex::new(-0.5057188113411757, -0.5836300412887051), + ]; + debug_assert!( + j.into_iter() + .zip_eq(k.into_iter()) + .all(|(x, y)| ((x.re - y.re).abs() < 1e-8) && ((x.im - y.im).abs() < 1e-8)) + ); + Ok(()) + } + + #[test] + fn fft_freq_test() -> Result<(), Box> { + let f = fft_freq(11); + let g = vec![ + 0.0, + 0.09090909090909091, + 0.18181818181818182, + 0.2727272727272727, + 0.36363636363636365, + 0.4545454545454546, + -0.4545454545454546, + -0.36363636363636365, + -0.2727272727272727, + -0.18181818181818182, + -0.09090909090909091, + ]; + debug_assert!(f.iter().zip_eq(g.iter()).all(|(x, y)| (x - y).abs() < 1e-8)); + Ok(()) + } + + #[test] + fn gaussian_kernel_test() -> Result<(), Box> { + let x = gaussian_kernel::(&[11], &[1.5])?.mapv(|i| i.re); + println!("x = {:?}", x.to_vec()); + let y = array![ + 1.00000000e+00, + 6.92774033e-01, + 2.30338431e-01, + 3.67556755e-02, + 2.81491714e-03, + 1.03464181e-04, + 1.03464181e-04, + 2.81491714e-03, + 3.67556755e-02, + 2.30338431e-01, + 6.92774033e-01 + ]; + debug_assert!(y.iter().zip_eq(x.iter()).all(|(x, y)| (x - y).abs() < 1e-8)); + Ok(()) + } + + #[test] + fn fft_shift_test() -> Result<(), Box> { + let x = Array1::from_iter(0..10); + let y = fft_shift(x.view()); + let z = array![5, 6, 7, 8, 9, 0, 1, 2, 3, 4]; + debug_assert_eq!(y, z); + Ok(()) + } + + #[test] + fn gaussian_test() -> Result<(), Box> { + let x = Array1::linspace(0.0, 1.0, 11).to_vec(); + let y = gaussian(x.as_slice(), 0.4, 0.1); + let z = vec![ + 1.33830226e-03, + 4.43184841e-02, + 5.39909665e-01, + 2.41970725e+00, + 3.98942280e+00, + 2.41970725e+00, + 5.39909665e-01, + 4.43184841e-02, + 1.33830226e-03, + 1.48671951e-05, + 6.07588285e-08, + ]; + debug_assert!(y.iter().zip_eq(z.iter()).all(|(x, y)| (x - y).abs() < 1e-8)); + Ok(()) + } +} diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..140efb1 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,179 @@ +use rayon::iter::ParallelIterator; +pub mod bspline; +pub mod error; +pub mod filter; +pub mod metric; +pub mod par_indexed_iter; +pub mod register; +pub mod transform; +mod optimize; + +use ndarray::prelude::*; +use num::Complex; +use thiserror::Error; + +use crate::par_indexed_iter::ParallelIndexedIterMut; +use crate::transform::transform_point; +use error::Error; + +/// An example of generating julia fractals, for testing purposes. +/// parameters: 2x2 flattened rotation matrix + xy translation +/// center: center of rotation +pub fn julia_image( + shape: &[usize; 2], + parameters: &[f64; 6], + center: &[f64; 2], + c: &[f64; 2], +) -> Array2 { + let c = Complex::new(c[0], c[1]); + let scaley = 3.0 / shape[0] as f64; + let scalex = 3.0 / shape[1] as f64; + let mut im = Array2::::zeros([shape[0], shape[1]]); + im.par_indexed_iter_mut().for_each(|(i, x)| { + let cv = transform_point(i.as_slice(), center, parameters); + let mut z = Complex::new(cv[1] * scalex - 1.5, cv[0] * scaley - 1.5); + let mut i = 0; + while i < 255 && z.norm() <= 2.0 { + z = z * z + c; + i += 1; + } + *x = i; + }); + + im +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::transform::Transform; + use tempfile::NamedTempFile; + use tiffwrite::IJTiffFile; + + #[test] + fn serialization() -> Result<(), Error> { + let file = NamedTempFile::new()?; + let t = Transform::::new_with_center( + vec![1.2, 0.3, -0.4, 0.9, 10.2, -9.5], + vec![59.5, 49.5], + vec![120, 100], + ); + t.to_file(file.path().to_path_buf())?; + let s = Transform::from_file(file.path().to_path_buf())?; + assert_eq!(s, t); + Ok(()) + } + + #[test] + fn transform_point() -> Result<(), Box> { + let transform = Transform::::new_with_center( + vec![1.0, 0.0, 0.0, 1.0, -120.0, 10.0], + vec![299.5, 399.5], + vec![600, 800], + ); + let p = transform.transform_point(&[300, 400]); + assert_eq!(p, [180.0, 410.0]); + Ok(()) + } + + #[test] + fn interpbs_f64() -> Result<(), Box> { + let shape = [1200, 1600]; + let transform_j = [1.0, 0.0, 0.0, 1.0, 0.0, 0.0]; + let transform_k = [1.0, 0.0, 0.0, 1.0, 0.0, 0.0]; + let center: [f64; 2] = shape + .iter() + .map(|&s| (s - 1) as f64 / 2.0) + .collect::>() + .try_into() + .unwrap(); + let transform = + Transform::::new(transform_j.to_vec(), shape.to_vec()).with_rotation(1.0); + let transform_j = transform.parameters.clone().try_into().unwrap(); + 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 sj = j.iter().sum::(); + // let sk = k.iter().sum::(); + let sn = n.iter().sum::(); + // println!("sj: {}, sk: {}, sn: {}", sj, sk, sn); + + // let mut tiff = tiffwrite::IJTiffFile::new("interpbs_f64_0.tif")?; + // tiff.save(j.view(), 0, 0, 0)?; + // tiff.save(k.view(), 1, 0, 0)?; + // tiff.save(n.view(), 2, 0, 0)?; + + let s = (sj.ln() - sn.ln()).abs(); + let d = (j.ln() - n.ln()) + .powi(2) + .iter() + .filter(|i| i.is_finite()) + .sum::() + .sqrt(); + // println!("s: {}, d: {}", s, d); + assert!(s < 1e-2); + assert!(2000.0 * d <= (shape[0] * shape[1]) as f64); + Ok(()) + } + + #[test] + fn interpbs_f64_par() -> Result<(), Box> { + let shape = [1200, 1600]; + let transform_j = [1.0, 0.0, 0.0, 1.0, 0.0, 0.0]; + let transform_k = [1.0, 0.0, 0.0, 1.0, 0.0, 0.0]; + let center: [f64; 2] = shape + .iter() + .map(|&s| (s - 1) as f64 / 2.0) + .collect::>() + .try_into() + .unwrap(); + let transform = + Transform::::new(transform_j.to_vec(), shape.to_vec()).with_rotation(1.0); + let transform_j = transform.parameters.clone().try_into().unwrap(); + 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 sj = j.iter().sum::(); + let sn = n.iter().sum::(); + + let s = (sj.ln() - sn.ln()).abs(); + let d = (j.ln() - n.ln()) + .powi(2) + .iter() + .filter(|i| i.is_finite()) + .sum::() + .sqrt(); + assert!(s < 1e-2); + assert!(2000.0 * d <= (shape[0] * shape[1]) as f64); + Ok(()) + } + + #[test] + fn interpolate_par_unsafe() -> Result<(), Box> { + let image_a = julia_image( + &[8000, 6000], + &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0], + &[3999.5, 2999.5], + &[-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())?; + assert_ne!(image_b[[4000, 3000]], 0.0); + Ok(()) + } + + #[test] + fn julia() -> Result<(), Box> { + let image_a = julia_image( + &[60, 80], + &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0], + &[29.5, 39.5], + &[-0.8, 0.156], + ); + let mut t = IJTiffFile::new(std::env::home_dir().unwrap().join("tmp/julia.tif"))?; + t.save(image_a.view(), 0, 0, 0)?; + Ok(()) + } +} diff --git a/src/metric.rs b/src/metric.rs index e69de29..47d1b04 100644 --- a/src/metric.rs +++ b/src/metric.rs @@ -0,0 +1,1067 @@ +use crate::bspline::{ + BSpline, BSplineTrait, cubic_bspline, square_bspline, square_bspline_integral, +}; +use crate::error::Error; +use crate::filter::gaussian_smooth; +use crate::transform::transform_point; +use algos::ObjectiveFunction; +use itertools::Itertools; +use ndarray::{Array, Array1, Array2, Array3, AsArray, Axis, Dimension, s}; +use num::cast::AsPrimitive; +use rand::Rng; +use rayon::iter::IntoParallelIterator; +use rayon::iter::ParallelIterator; +use std::cell::RefCell; +use std::ops::Deref; + +#[derive(Default)] +struct IntMut { + mu: Vec, + metric: f64, + derivative: Vec, +} + +pub struct FixedMu { + mu: Vec>, + len_variable: usize, +} + +impl Deref for FixedMu { + type Target = Vec>; + + fn deref(&self) -> &Self::Target { + &self.mu + } +} + +impl From>> for FixedMu { + fn from(mu: Vec>) -> Self { + let len_variable = mu.iter().filter(|v| v.is_none()).count(); + FixedMu { mu, len_variable } + } +} + +impl FixedMu { + pub fn new_none(n: usize) -> Self { + Self { + mu: vec![None; n * n + n], + len_variable: n * n + n, + } + } + + /// combine mu and fixed_mu into complete parameter set + pub fn combine(&self, mu: &[f64]) -> Vec { + debug_assert_eq!(mu.len(), self.len_variable); + let mut mu_iter = mu.iter(); + self.iter() + .map(|i| { + if let Some(i) = i { + *i + } else { + *mu_iter.next().unwrap() + } + }) + .collect() + } + + /// extract mu from complete parameter set + pub fn extract_variable(&self, mu: &[f64]) -> Vec { + self.iter() + .zip_eq(mu) + .filter_map(|(i, j)| if i.is_none() { Some(*j) } else { None }) + .collect() + } +} + +/// first_bin, last_bin: center of bin +/// n_bins >= 2 +#[allow(clippy::too_many_arguments)] +fn parzen( + alpha: f64, + dalpha: &[f64], + fixed: &[f64], + moving: &[f64], + dmoving: &[Vec], + weight: &[f64], + dweight: &[Vec], + first_bin: f64, + last_bin: f64, + n_bins: usize, +) -> (Array2, Array3) { + let ndim = dalpha.len(); + let width = (last_bin - first_bin) / ((n_bins - 1) as f64); + let v = 1.0 / width; + // add two bins left and right + let first_bin = first_bin - 2.0 * width; + let last_bin = last_bin + 2.0 * width; + let n_bins = n_bins + 4; + let n_bins_i = n_bins as isize; + let mut hist = Array2::zeros([n_bins, n_bins]); + let mut grad_b = Array3::zeros([n_bins, n_bins, ndim]); + let left = first_bin - width / 2.0; + let bins = Array1::linspace(first_bin, last_bin, n_bins); + let mut nf; + let mut nm; + let mut uf; + let mut um; + let mut kf; + let mut km; + let mut csf; + for ((f, m), (dm, (w, dw))) in fixed + .iter() + .zip_eq(moving) + .zip_eq(dmoving.iter().zip_eq(weight.iter().zip_eq(dweight))) + { + nf = (v * (f - left)) as isize; + nm = (v * (m - left)) as isize; + for if_ in (nf - 2)..=(nf + 2) { + if (0 <= if_) && (if_ < n_bins_i) { + uf = if_ as usize; + kf = (f - bins[uf]) * v; + csf = cubic_bspline(kf); + for im in (nm - 2)..=(nm + 2) { + if (0 <= im) && (im < n_bins_i) { + um = im as usize; + km = (m - bins[um]) * v; + hist[[uf, um]] += w * csf * cubic_bspline(km); + for ((dmi, dwi), d) in + dm.iter() + .zip_eq(dw) + .zip_eq(grad_b.slice_mut(s![uf, um, ..])) + { + *d += w + * csf + * (square_bspline(km + 0.5) - square_bspline(km - 0.5)) + * v + * dmi + + dwi * csf * cubic_bspline(km); + } + } + } + } + } + } + let hist = hist / alpha; + let mut hist_bc = hist + .broadcast([dalpha.len(), hist.shape()[0], hist.shape()[1]]) + .unwrap(); + hist_bc.permute_axes([1, 2, 0]); + let dhist = (grad_b - (&hist_bc * &Array1::from_iter(dalpha))) / alpha; + (hist, dhist) +} + +enum Sampling { + Fixed(Vec>), + Random((usize, Vec)), +} + +impl Sampling { + fn fixed(index: Vec>) -> Self { + Self::Fixed(index) + } + + fn random(n_samples: usize, shape: Vec) -> Self { + Self::Random((n_samples, shape)) + } + + fn index(&self) -> Vec> { + match self { + Self::Fixed(index) => index.clone(), + Self::Random((n_samples, shape)) => rand::rng() + .random_iter::() + .take(n_samples * shape.len()) + .chunks(shape.len()) + .into_iter() + .map(|c| c.zip_eq(shape.iter()).map(|(i, s)| i * s - 0.5).collect()) + .collect(), + } + } +} + +#[derive(Clone, Debug)] +pub enum SamplingArg { + Fixed(usize), + FixedAt(Vec>), + Random(usize), +} + +/// (dI / dT(x)) * (dT(x) / dmu) +fn image_jacobian(point: &[f64], center: &[f64], image_derivative: &[f64]) -> Vec { + image_derivative + .iter() + .flat_map(|d| { + point + .iter() + .zip_eq(center.iter()) + .map(|(p, c)| *d * (p - c)) + }) + .chain(image_derivative.iter().cloned()) + .collect() +} + +pub struct MattesMetric +where + D: Dimension, +{ + shape: Vec, + center: Vec, + fixed: BSpline<0, D>, + moving: BSpline<3, D>, + fixed_mu: FixedMu, + minmax: [f64; 2], + sampling: Sampling, + n_bins: usize, + /// memory for (point, value, gradient) combination, refcell gives interior mutability + metric: RefCell, + edge: f64, + // mem_f: RefCell, + // mem_m: RefCell, +} + +impl MattesMetric +where + D: Dimension, +{ + pub fn new( + fixed: BSpline<0, D>, + moving: BSpline<3, D>, + sampling: SamplingArg, + n_bins: usize, + edge: f64, + ) -> Result { + let shape = fixed.shape().iter().map(|&i| i as f64).collect::>(); + let center = shape.iter().map(|i| (i - 1.0) / 2.0).collect(); + // let mem_f = fixed.get_mem(); + // let mem_m = moving.get_mem(); + let (min, max) = fixed + .iter() + .chain(moving.iter()) + .minmax() + .into_option() + .expect("m and f cannot be empty"); + let minmax = [*min, *max]; + let sampling = match sampling { + SamplingArg::Fixed(n_samples) => { + Sampling::fixed(Sampling::random(n_samples, shape.clone()).index()) + } + SamplingArg::FixedAt(points) => Sampling::fixed(points), + SamplingArg::Random(n_samples) => Sampling::random(n_samples, shape.clone()), + }; + + Ok(Self { + shape, + center, + fixed, + moving, + fixed_mu: FixedMu::new_none(sampling.index().len()), + minmax, + sampling, + n_bins, + metric: RefCell::default(), + edge, + }) + } + + pub fn new_from_arrays<'a, T, A>( + fixed: A, + moving: A, + sampling: SamplingArg, + n_bins: usize, + edge: f64, + ) -> Result + where + A: AsArray<'a, T, D>, + T: 'a + Clone + AsPrimitive, + { + Self::new( + BSpline::new(fixed), + BSpline::new(moving), + sampling, + n_bins, + edge, + ) + } + + pub fn with_fixed_mu>(mut self, fixed_mu: F) -> Self { + self.fixed_mu = fixed_mu.into(); + self + } + + pub fn fixed_mu(&self) -> &FixedMu { + &self.fixed_mu + } + + pub fn eval(&self, mu: &[f64]) -> Result<(), Error> { + assert!(mu.iter().all(|i| i.is_finite())); + if self.metric.borrow().mu != mu { + let complete_mu = self.fixed_mu.combine(mu); + let idx = self.sampling.index(); + let shape = self.shape.clone(); + let shape_f = self.shape.iter().map(|s| s - 0.5).collect::>(); + let center = self.center.as_slice(); + let fixed = &self.fixed; + let mem_f = fixed.get_mem(); + let moving = &self.moving; + let mem_m = moving.get_mem(); + let fixed_mu = &self.fixed_mu; + let e = 1.0 / self.edge; + let (f, ((m, j), (w, dw))) = idx + .into_par_iter() + .map_init( + || (mem_f.clone(), mem_m.clone()), + |(mem_f, mem_m), u| { + let v = transform_point(&u, center, &complete_mu); + if shape_f + .iter() + .zip_eq(v.iter()) + .all(|(s, x)| (-0.5 <= *x) && (x <= s)) + { + let f = fixed.evaluate_at_continuous_index(&u, mem_f)?; + let (m, d) = moving + .evaluate_value_and_derivative_at_continuous_index(&v, mem_m)?; + let j = fixed_mu.extract_variable(&image_jacobian(&v, center, &d)); + let mut k; + let mut w = 1.0; + let mut ws = Vec::with_capacity(shape.len()); + let mut dw = Vec::with_capacity(shape.len()); + // the metric wouldn't be smooth if we'd drop a point abruptly, + // so apply a 0 <= weight <= 1 to points within 5% of the edge + for (s, x) in shape.iter().zip_eq(&v) { + if e * (x + 0.5) < *s { + k = 3.0 * e * (x + 0.5) / s - 1.5; + let b = square_bspline_integral(k); + w *= b; + ws.push(b); + dw.push(square_bspline(k) * 3.0 * e / s); + } else if e * (x + 0.5) > (e - 1.0) * s { + k = 3.0 * e * (s - x - 0.5) / s - 1.5; + let b = square_bspline_integral(k); + w *= b; + ws.push(b); + dw.push(-square_bspline(k) * 3.0 * e / s); + } else { + ws.push(1.0); + dw.push(0.0); + } + } + for (dwi, wi) in dw.iter_mut().zip(ws) { + *dwi *= w / wi; + } + let dw = fixed_mu.extract_variable(&image_jacobian(&v, center, &dw)); + Ok(Some((f, ((m, j), (w, dw))))) + } else { + Ok(None) + } + }, + ) + .filter_map(|i| i.transpose()) + .collect::, + ((Vec, Vec>), (Vec, Vec>)), + ), + Error, + >>()?; + let alpha = w.iter().sum::(); + let mut dalpha = vec![0.0; mu.len()]; + for i in dw.iter() { + for (j, a) in i.iter().zip_eq(dalpha.iter_mut()) { + *a += j; + } + } + + let (jpdf, d_jpdf_m) = parzen( + alpha, + &dalpha, + &f, + &m, + j.as_slice(), + &w, + &dw, + self.minmax[0], + self.minmax[1], + self.n_bins, + ); + + let pdf_f = jpdf.axis_iter(Axis(0)).map(|i| i.sum()).collect::>(); + let pdf_m = jpdf.axis_iter(Axis(1)).map(|i| i.sum()).collect::>(); + let d_pdf_m = d_jpdf_m + .axis_iter(Axis(1)) + .zip(pdf_m.iter()) + .map(|(di, i)| di.sum() / i) + .collect::>(); + + let mut metric = 0.0; + let mut dmetric = vec![0.0; self.fixed_mu.len_variable]; + let mut n; + let mut m; + let mut ik; + let mut ln; + + for ((row, d_row), k) in jpdf + .axis_iter(Axis(0)) + .zip_eq(d_jpdf_m.axis_iter(Axis(0))) + .zip_eq(pdf_f) + { + debug_assert_eq!(row.len(), d_row.shape()[0]); + debug_assert_eq!(row.len(), pdf_m.len()); + + for ((&j, dj), (i, di)) in row + .into_iter() + .zip_eq(d_row.axis_iter(Axis(0))) + .zip_eq(pdf_m.iter().zip(d_pdf_m.iter())) + { + ik = i * k; + // check for non-zero bin contribution + if (j > 1e-16) && (ik > 1e-16) { + ln = (j / (ik)).ln(); + n = 1.0 + ln; + m = di * j; + metric -= j * ln; + for (d, g) in dmetric.iter_mut().zip_eq(dj) { + *d += m - n * g; + } + } + } + } + + debug_assert!(mu.iter().all(|i| i.is_finite())); + debug_assert!(alpha.is_finite()); + debug_assert!(metric.is_finite()); + debug_assert!(dmetric.iter().all(|i| i.is_finite())); + + // let alpha = 1.0; + // let dalpha = vec![0.0; dalpha.len()]; + + self.metric.replace(IntMut { + mu: mu.to_vec(), + metric: metric * alpha, + derivative: dmetric + .iter() + .zip_eq(dalpha) + .map(|(dm, da)| *dm * alpha + da * metric) + .collect(), + }); + // self.metric.replace(IntMut { + // mu: mu.to_vec(), + // metric: metric / alpha, + // derivative: dmetric + // .iter() + // .zip_eq(dalpha) + // .map(|(dm, da)| *dm / alpha + da * metric / (alpha * alpha)) + // .collect(), + // }); + } + Ok(()) + } +} + +impl ObjectiveFunction for MattesMetric +where + D: Dimension, +{ + fn evaluate(&self, point: &[f64]) -> f64 { + self.eval(point).unwrap(); + self.metric.borrow().metric + } + + fn gradient(&self, point: &[f64]) -> Option> { + self.eval(point).unwrap(); + Some(self.metric.borrow().derivative.clone()) + } +} + +#[derive(Clone, Debug)] +pub enum Sigma { + None, + Absolute(Vec), + Relative(Vec), +} + +impl Sigma { + pub fn smooth<'a, A, D>(&self, array: A) -> Result, Error> + where + A: AsArray<'a, f64, D>, + D: Dimension, + { + match self { + Self::None => Ok(array.into().to_owned()), + Self::Absolute(sigma) => gaussian_smooth(array, sigma.as_slice()), + Self::Relative(sigma) => { + let array = array.into(); + let sigma = sigma + .iter() + .zip(array.shape()) + .map(|(i, j)| i * (*j as f64)) + .collect::>(); + gaussian_smooth(array, sigma.as_slice()) + } + } + } + + pub fn sigma(&self, shape: &[usize]) -> Option> { + match self { + Self::None => None, + Self::Absolute(sigma) => Some(sigma.to_vec()), + Self::Relative(sigma) => Some( + sigma + .iter() + .zip(shape) + .map(|(i, j)| i * (*j as f64)) + .collect::>(), + ), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::filter::gaussian_smooth; + use crate::julia_image; + use crate::transform::Transform; + use algos::OptimizationConfig; + use algos::optimization::{bfgs_minimize, gradient_descent_minimize}; + use ndarray::{MeshIndex, array, meshgrid, stack}; + use ndarray_npy::NpzWriter; + use num::integer::Roots; + use num::traits::FloatConst; + use std::fs::File; + + #[test] + fn derivative() -> Result<(), Box> { + let im = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0]; + let b = BSpline::<3, _>::new(im.view()); + let mut mem = b.get_mem(); + let mut value = Vec::new(); + let mut derivative = Vec::new(); + let x = (0..100 * im.len()) + .map(|i| i as f64 / 100.0) + .collect::>(); + for i in x.iter() { + let (v, d) = b.evaluate_value_and_derivative_at_continuous_index(&[*i], &mut mem)?; + value.push(v); + derivative.push(d[0]); + } + println!("x = {:?}\nv = {:?}\nd = {:?}", x, value, derivative); + Ok(()) + } + + #[test] + fn grad() -> Result<(), Box> { + let mut rng = rand::rng(); + let u1 = (&mut rng) + .random_iter::() + .take(1000) + .collect::>(); + let u2 = (&mut rng) + .random_iter::() + .take(1000) + .collect::>(); + let f = u1 + .iter() + .zip_eq(u2.iter()) + .map(|(i, j)| (-2.0 * i.ln()).sqrt() * (2.0 * f64::PI() * j).cos() - 1.0) + .collect::>(); + let m = u1 + .iter() + .zip_eq(u2.iter()) + .map(|(i, j)| (-2.0 * i.ln()).sqrt() * (2.0 * f64::PI() * j).sin() + 1.0) + .collect::>(); + + let s = f.len().sqrt(); + let n_bins = s.max(11); + let e = 2.0 / (n_bins - 3) as f64; + let first_bin = -5.0 - e; + let last_bin = 5.0 + e; + // println!("f: {:?}, m: {:?}", f, m); + // println!("first_bin: {}, last_bin: {}, n_bins: {}, e: {}", first_bin, last_bin, n_bins, e); + let w = vec![1.0; f.len()]; + let dw = vec![vec![0.0]; f.len()]; + let (jpdf, d_jpdf_m) = parzen( + 1.0, + &[1.0], + &f, + &m, + &[], + &w, + &dw, + first_bin, + last_bin, + n_bins, + ); + let pdf_f = jpdf.rows().into_iter().map(|i| i.sum()).collect::>(); + let pdf_m = jpdf + .columns() + .into_iter() + .map(|i| i.sum()) + .collect::>(); + let d_pdf_m = d_jpdf_m + .columns() + .into_iter() + .map(|i| i.sum()) + .collect::>(); + println!( + "pdf_f = {:?}\npdf_m = {:?}\nd_pdf_m = {:?}", + pdf_f.to_vec(), + pdf_m.to_vec(), + d_pdf_m.to_vec() + ); + Ok(()) + } + + #[test] + fn metric() -> Result<(), Box> { + let im_a = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0]; + let im_b = array![1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0]; + // let a = gaussian_smooth(im_a.view(), &[4.0])?; + // let b = gaussian_smooth(im_b.view(), &[4.0])?; + let a = im_a; + let b = im_b; + + // let points = vec![vec![3.0], vec![4.0], vec![5.0], vec![6.0], vec![7.0], vec![8.0]]; + // let m = MattesMetric::new_from_arrays(a.view(), b.view(), SamplingArg::FixedAt(points))? + // .with_fixed_mu(vec![Some(1.0), None]); + let m = MattesMetric::new_from_arrays(a.view(), b.view(), SamplingArg::Fixed(20), 5, 0.05)? + .with_fixed_mu(vec![Some(1.0), None]); + // let m = MattesMetric::new_all_from_arrays(im_a.view(), im_b.view())?.with_fixed_mu(vec![Some(1.0), None]); + let mu = Array1::linspace(-4.0, 2.0, 601); + let mut v = Vec::new(); + let mut d = Vec::new(); + for x in &mu { + v.push(m.evaluate(&[*x])); + d.push(m.gradient(&[*x]).unwrap()[0]); + } + let mut npz = NpzWriter::new(File::create( + std::env::home_dir().unwrap().join("tmp/metric.npz"), + )?); + npz.add_array("mu", &mu)?; + npz.add_array("v", &Array1::from_vec(v))?; + npz.add_array("d", &Array1::from_vec(d))?; + npz.add_array("a", &a)?; + npz.add_array("b", &b)?; + Ok(()) + } + + #[test] + fn metric2() -> Result<(), Box> { + let im_a = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0]; + let im_b = array![1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0]; + let a = gaussian_smooth(im_a.view(), &[2.0])?; + let b = gaussian_smooth(im_b.view(), &[2.0])?; + + let points = vec![ + vec![3.0], + vec![4.0], + vec![5.0], + vec![6.0], + vec![7.0], + vec![8.0], + ]; + + let m = MattesMetric::new_from_arrays( + a.view(), + b.view(), + SamplingArg::FixedAt(points), + 5, + 0.05, + )? + .with_fixed_mu(vec![Some(1.0), None]); + // let mu = vec![-3.52, -3.51, -3.50, -3.49, -3.48]; + let mu = vec![-3.51, -3.49]; + // let mu = vec![-2.0, -3.0]; + let mut v = Vec::new(); + let mut d = Vec::new(); + for x in mu.iter() { + v.push(m.evaluate(&[*x])); + d.push(m.gradient(&[*x]).unwrap()[0]); + } + println!("a = {:?}\nb = {:?}", a.to_vec(), b.to_vec()); + println!("mu = {:?}\nv = {:?}\nd = {:?}", mu, v.to_vec(), d.to_vec()); + Ok(()) + } + + #[test] + fn metric1() -> Result<(), Box> { + let im_a = julia_image( + &[60, 80], + &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0], + &[29.5, 39.5], + &[-0.8, 0.156], + ) + .mapv(|i| i as f64); + let im_b = julia_image( + &[60, 80], + &[1.0, 0.0, 0.0, 1.0, 10.0, -20.0], + &[29.5, 39.5], + &[-0.8, 0.156], + ) + .mapv(|i| i as f64); + let a = gaussian_smooth(im_a.view(), &[30.0, 40.0])?; + let b = gaussian_smooth(im_b.view(), &[30.0, 40.0])?; + let m = MattesMetric::new_from_arrays(a.view(), b.view(), SamplingArg::Fixed(250), 3, 0.2)? + .with_fixed_mu(vec![Some(1.0), Some(0.0), Some(0.0), Some(1.0), None, None]); + let mut npz = NpzWriter::new(File::create( + std::env::home_dir().unwrap().join("tmp/metric.npz"), + )?); + + let s = 200; + let mu = Array1::linspace(-50.0, 50.0, s); + let (mux, muy) = meshgrid((&mu, &mu), MeshIndex::XY); + let mu = stack(Axis(0), &[mux, muy])?; + let mut v = Array2::zeros([s, s]); + let mut d = Array3::zeros([2, s, s]); + for (mui, (vi, mut di)) in mu + .lanes(Axis(0)) + .into_iter() + .zip_eq(v.iter_mut().zip_eq(d.lanes_mut(Axis(0)))) + { + *vi = m.evaluate(mui.to_vec().as_slice()); + di.assign(&Array1::from_vec( + m.gradient(mui.to_vec().as_slice()).unwrap(), + )); + } + npz.add_array("a", &a)?; + npz.add_array("b", &b)?; + npz.add_array("mu", &mu)?; + npz.add_array("v", &v)?; + npz.add_array("d", &d)?; + Ok(()) + } + + #[test] + fn metric3() -> Result<(), Box> { + let im_a = julia_image( + &[100, 1], + &[1.0, 0.0, 0.0, 0.01, 0.0, 0.0], + &[99.5, 0.5], + &[-0.8, 0.156], + ) + .slice(s![.., 0]) + .mapv(|i| i as f64); + let im_b = + Transform::new(vec![1.0, 0.0], vec![im_a.shape()[0]]).interpolate::<1, _, _>(&im_a)?; + let a = gaussian_smooth(im_a.view(), &[25.0])?; + let b = gaussian_smooth(im_b.view(), &[25.0])?; + let m = MattesMetric::new_from_arrays(a.view(), b.view(), SamplingArg::Fixed(100), 6, 0.1)? + .with_fixed_mu(vec![None, None]); + let mut npz = NpzWriter::new(File::create( + std::env::home_dir().unwrap().join("tmp/metric.npz"), + )?); + + let s = 200; + let translate = Array1::linspace(-50.0, 50.0, s); + let scale = Array1::logspace(10.0, -1.0, 1.0, s); + let (mux, muy) = meshgrid((&scale, &translate), MeshIndex::XY); + let mu = stack(Axis(0), &[mux, muy])?; + let mut v = Array2::zeros([s, s]); + let mut d = Array3::zeros([2, s, s]); + for (mui, (vi, mut di)) in mu + .lanes(Axis(0)) + .into_iter() + .zip_eq(v.iter_mut().zip_eq(d.lanes_mut(Axis(0)))) + { + *vi = m.evaluate(mui.to_vec().as_slice()); + di.assign(&Array1::from_vec( + m.gradient(mui.to_vec().as_slice()).unwrap(), + )); + } + npz.add_array("a", &a)?; + npz.add_array("b", &b)?; + npz.add_array("mu", &mu)?; + npz.add_array("v", &v)?; + npz.add_array("d", &d)?; + Ok(()) + } + + #[test] + fn metric4() -> Result<(), Box> { + let im_a = julia_image( + &[100, 1], + &[1.0, 0.0, 0.0, 0.01, 0.0, 0.0], + &[99.5, 0.5], + &[-0.8, 0.156], + ) + .slice(s![.., 0]) + .mapv(|i| i as f64); + let im_b = + Transform::new(vec![1.0, 0.0], vec![im_a.shape()[0]]).interpolate::<1, _, _>(&im_a)?; + let a = gaussian_smooth(im_a.view(), &[50.0])?; + let b = gaussian_smooth(im_b.view(), &[50.0])?; + let m = MattesMetric::new_from_arrays(a.view(), b.view(), SamplingArg::Fixed(1000), 3, 0.05)? + .with_fixed_mu(vec![None, Some(0.0)]); + let mut npz = NpzWriter::new(File::create( + std::env::home_dir().unwrap().join("tmp/metric.npz"), + )?); + + let mu = Array1::logspace(2.0, -1.0, 1.0, 1000); + let mut v = Vec::new(); + let mut d = Vec::new(); + for mui in &mu { + v.push(m.evaluate(&[*mui])); + d.push(m.gradient(&[*mui]).unwrap()[0]); + } + npz.add_array("a", &a)?; + npz.add_array("b", &b)?; + npz.add_array("mu", &mu)?; + npz.add_array("v", &Array1::from(v))?; + npz.add_array("d", &Array1::from(d))?; + Ok(()) + } + + #[test] + fn optimize() -> Result<(), Box> { + let im_a = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0]; + let im_b = array![1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0]; + let a = gaussian_smooth(im_a.view(), &[4.0])?; + let b = gaussian_smooth(im_b.view(), &[4.0])?; + + let m = + MattesMetric::new_from_arrays(a.view(), b.view(), SamplingArg::Random(100), 5, 0.05)? + .with_fixed_mu(vec![Some(1.0), None]); + let optimization_config = OptimizationConfig { + max_iterations: 1000, + tolerance: 1e-6, + learning_rate: 0.001, + }; + + let r_grad = gradient_descent_minimize(&m, &[-0.0], &optimization_config); + let r_bfgs = bfgs_minimize(&m, &[-0.0], &optimization_config); + println!("r_grad = {:?}", r_grad); + println!("r_bfgs = {:?}", r_bfgs); + Ok(()) + } + + #[test] + fn parzen_test() -> Result<(), Box> { + let a = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0]; + let b = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0]; + let fixed = BSpline::<0, _>::new(a.view()); + let moving = BSpline::<3, _>::new(b.view()); + + let mus = vec![0.0, 0.001]; + let mut npz = NpzWriter::new(File::create( + std::env::home_dir().unwrap().join("tmp/metric.npz"), + )?); + let s = a.len() as f64; + // let center = (s - 1.0) / 2.0; + let mut rng = rand::rng(); + let idx = (&mut rng) + .random_iter::() + .take(25) + .map(|i| i * s - 0.5) + .collect::>(); + + for (index, mu) in mus.iter().enumerate() { + let mut mem_f = fixed.get_mem(); + let mut mem_m = moving.get_mem(); + let mut w = Vec::new(); + let mut dw = Vec::new(); + let mut a = Vec::new(); + let mut b = Vec::new(); + let mut db = Vec::new(); + + for u in idx.clone() { + let v = u + mu; + if (-0.5 < v) && (v <= s - 0.5) { + let f = fixed.evaluate_at_continuous_index(&[u], &mut mem_f)?; + let (m, d) = moving + .evaluate_value_and_derivative_at_continuous_index(&[v], &mut mem_m)?; + // let j = image_jacobian(&[v], &[center], &d); + a.push(f); + b.push(m); + db.push(d); + // the metric wouldn't be smooth if we'd drop a point abruptly, + // so apply a 0 <= weight <= 1 to points within 5% of the edge + if 20.0 * (v + 0.5) < s { + let k = 60.0 * (v + 0.5) / s - 1.5; + w.push(square_bspline_integral(k)); + dw.push(vec![square_bspline(k) * 60.0 / s]); + } else if 20.0 * (v + 0.5) > 19.0 * s { + let k = 60.0 * (s - v - 0.5) / s - 1.5; + w.push(square_bspline_integral(k)); + dw.push(vec![-square_bspline(k) * 60.0 / s]); + } else { + w.push(1.0); + dw.push(vec![0.0]); + } + // let dw = image_jacobian(&[v], &[center], &dw); + } + } + + let alpha = w.iter().sum::(); + let mut dalpha = [0.0]; + for i in dw.iter() { + for (j, a) in i.iter().zip_eq(dalpha.iter_mut()) { + *a += j; + } + } + + let (j, dj) = parzen(alpha, &dalpha, &a, &b, &db, &w, &dw, 0.0, 2.0, 5); + let db = db.iter().map(|i| i[0]).collect::>(); + let dw = dw.iter().map(|i| i[0]).collect::>(); + + npz.add_array(format!("idx{}", index), &Array1::from_vec(idx.clone()))?; + npz.add_array(format!("a{}", index), &Array1::from_vec(a))?; + npz.add_array(format!("b{}", index), &Array1::from_vec(b))?; + npz.add_array(format!("db{}", index), &Array1::from_vec(db))?; + npz.add_array(format!("w{}", index), &Array1::from_vec(w))?; + npz.add_array(format!("dw{}", index), &Array1::from_vec(dw))?; + npz.add_array(format!("j{}", index), &j)?; + npz.add_array(format!("dj{}", index), &dj)?; + } + Ok(()) + } + + #[test] + fn metric_test() -> Result<(), Box> { + let a = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0]; + let b = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0]; + let fixed = BSpline::<0, _>::new(a.view()); + let moving = BSpline::<3, _>::new(b.view()); + + let mus = Array1::linspace(-10.0, 10.0, 500); + let s = a.len() as f64; + let mut rng = rand::rng(); + let idx = (&mut rng) + .random_iter::() + .take(25) + .map(|i| i * s - 0.5) + .collect::>(); + + let mut metrics = Vec::new(); + let mut dmetrics = Vec::new(); + + for mu in &mus { + let mut mem_f = fixed.get_mem(); + let mut mem_m = moving.get_mem(); + let mut w = Vec::new(); + let mut dw = Vec::new(); + let mut a = Vec::new(); + let mut b = Vec::new(); + let mut db = Vec::new(); + + for u in idx.clone() { + let v = u + mu; + if (-0.5 < v) && (v <= s - 0.5) { + let f = fixed.evaluate_at_continuous_index(&[u], &mut mem_f)?; + let (m, d) = moving + .evaluate_value_and_derivative_at_continuous_index(&[v], &mut mem_m)?; + // let j = image_jacobian(&[v], &[center], &d); + a.push(f); + b.push(m); + db.push(d); + // the metric wouldn't be smooth if we'd drop a point abruptly, + // so apply a 0 <= weight <= 1 to points within 5% of the edge + if 20.0 * (v + 0.5) < s { + let k = 60.0 * (v + 0.5) / s - 1.5; + w.push(square_bspline_integral(k)); + dw.push(vec![square_bspline(k) * 60.0 / s]); + } else if 20.0 * (v + 0.5) > 19.0 * s { + let k = 60.0 * (s - v - 0.5) / s - 1.5; + w.push(square_bspline_integral(k)); + dw.push(vec![-square_bspline(k) * 60.0 / s]); + } else { + w.push(1.0); + dw.push(vec![0.0]); + } + // let dw = image_jacobian(&[v], &[center], &dw); + } + } + + let alpha = w.iter().sum::(); + let mut dalpha = [0.0]; + for i in dw.iter() { + for (j, a) in i.iter().zip_eq(dalpha.iter_mut()) { + *a += j; + } + } + + let (jpdf, d_jpdf_m) = parzen(alpha, &dalpha, &a, &b, &db, &w, &dw, 0.0, 2.0, 5); + + let pdf_f = jpdf.axis_iter(Axis(0)).map(|i| i.sum()).collect::>(); + let pdf_m = jpdf.axis_iter(Axis(1)).map(|i| i.sum()).collect::>(); + let d_pdf_m = d_jpdf_m + .axis_iter(Axis(1)) + .map(|i| i.sum()) + .collect::>(); + + let mut metric = 0.0; + let mut dmetric = vec![0.0]; + let mut n; + let mut m; + + for ((row, d_row), k) in jpdf + .axis_iter(Axis(0)) + .zip_eq(d_jpdf_m.axis_iter(Axis(0))) + .zip_eq(pdf_f) + { + debug_assert_eq!(row.len(), d_row.shape()[0]); + debug_assert_eq!(row.len(), pdf_m.len()); + + for ((&j, dj), (i, di)) in row + .into_iter() + .zip_eq(d_row.axis_iter(Axis(0))) + .zip_eq(pdf_m.iter().zip(d_pdf_m.iter())) + { + // check for non-zero bin contribution + if (j > 1e-16) && (i * k > 1e-16) { + metric -= j * (j / (i * k)).ln(); + n = (j / (i * k)).ln() + 1.0; + m = di * j / i; + for (d, g) in dmetric.iter_mut().zip_eq(dj) { + *d += m - n * g; // eq 23 of Thevenaz & Unser paper [3] + } + } + } + } + + metrics.push(metric); + dmetrics.push(dmetric); + } + let mut npz = NpzWriter::new(File::create( + std::env::home_dir().unwrap().join("tmp/metric.npz"), + )?); + npz.add_array("mu", &mus)?; + npz.add_array("metric", &Array1::from_vec(metrics))?; + npz.add_array( + "dmetric", + &Array1::from_vec(dmetrics.iter().map(|d| d[0]).collect()), + )?; + Ok(()) + } + + #[test] + fn test_square_bspline_integral() -> Result<(), Box> { + let x = Array1::linspace(-2.0, 2.0, 500); + let y = x.mapv(square_bspline); + let z = x.mapv(square_bspline_integral); + let mut w = vec![0.0]; + let mut j = 0.0; + let dx = (x[x.len() - 1] - x[0]) / (x.len() - 1) as f64; + for i in y.windows([2]) { + j += dx * (i[0] + i[1]) / 2.0; + w.push(j); + } + let w = Array1::from(w); + assert!((w - z).abs().iter().all(|i| *i < 1e-5)); + Ok(()) + } + + #[test] + fn test_cubic_bspline_derivative() -> Result<(), Box> { + let x = Array1::linspace(-2.0, 2.0, 500); + let y = x.mapv(cubic_bspline); + let z = x.mapv(|i| square_bspline(i + 0.5) - square_bspline(i - 0.5)); + let mut w = vec![0.0]; + let mut j = 0.0; + let dx = (x[x.len() - 1] - x[0]) / (x.len() - 1) as f64; + for i in z.windows([2]) { + j += dx * (i[0] + i[1]) / 2.0; + w.push(j); + } + let w = Array1::from(w); + assert!((w - y).abs().iter().all(|i| *i < 1e-4)); + Ok(()) + } +} diff --git a/src/optimize.rs b/src/optimize.rs index e69de29..fc59a72 100644 --- a/src/optimize.rs +++ b/src/optimize.rs @@ -0,0 +1,167 @@ +use std::collections::VecDeque; +use std::fmt::Debug; +use algos::{ObjectiveFunction, OptimizationConfig, OptimizationResult}; +use num::Float; + +pub fn lbfgs_minimize( + f: &F, + initial_point: &[T], + config: &OptimizationConfig, +) -> OptimizationResult +where + T: Float + Debug, + F: ObjectiveFunction, +{ + const M: usize = 10; // Number of corrections to store + let n = initial_point.len(); + let mut current_point = initial_point.to_vec(); + let mut iterations = 0; + let mut converged = false; + + // Storage for the last M corrections + let mut s_list: VecDeque> = VecDeque::with_capacity(M); + let mut y_list: VecDeque> = VecDeque::with_capacity(M); + let mut rho_list: VecDeque = VecDeque::with_capacity(M); + + // Get initial gradient + let mut gradient = match f.gradient(¤t_point) { + Some(g) => g, + None => { + return OptimizationResult { + optimal_point: current_point.clone(), + optimal_value: f.evaluate(¤t_point), + iterations: 0, + converged: false, + }; + } + }; + + while iterations < config.max_iterations { + // Check for convergence + let gradient_norm = gradient + .iter() + .fold(T::zero(), |acc, &x| acc + x * x) + .sqrt(); + if gradient_norm < config.tolerance { + converged = true; + break; + } + + // Compute search direction using L-BFGS two-loop recursion + let mut q = gradient.clone(); + let mut alpha_list = Vec::with_capacity(s_list.len()); + + // First loop + for i in (0..s_list.len()).rev() { + let alpha = rho_list[i] + * s_list[i] + .iter() + .zip(q.iter()) + .fold(T::zero(), |acc, (&s, &q)| acc + s * q); + alpha_list.push(alpha); + for (q_j, y_j) in q.iter_mut().zip(y_list[i].iter()) { + *q_j = *q_j - alpha * *y_j; + } + } + + // Scale the initial Hessian approximation + let mut r = if !s_list.is_empty() { + let i = s_list.len() - 1; + let yy = y_list[i].iter().fold(T::zero(), |acc, &y| acc + y * y); + let ys = y_list[i] + .iter() + .zip(s_list[i].iter()) + .fold(T::zero(), |acc, (&y, &s)| acc + y * s); + q.iter_mut().for_each(|r_j| *r_j = *r_j * (ys / yy)); + q + } else { + q.iter_mut() + .for_each(|r_j| *r_j = *r_j * config.learning_rate); + q + }; + + // Second loop + for i in 0..s_list.len() { + let beta = rho_list[i] + * y_list[i] + .iter() + .zip(r.iter()) + .fold(T::zero(), |acc, (&y, &r)| acc + y * r); + let alpha = alpha_list[s_list.len() - 1 - i]; + for (r_j, s_j) in r.iter_mut().zip(s_list[i].iter()) { + *r_j = *r_j + (alpha - beta) * *s_j; + } + } + + // r now contains the search direction + let direction: Vec = r.iter().map(|&x| -x).collect(); + + // Line search to find step size + let mut alpha = T::one(); + let mut new_point = vec![T::zero(); n]; + let current_value = f.evaluate(¤t_point); + + // Simple backtracking line search + for _ in 0..20 { + for i in 0..n { + new_point[i] = current_point[i] + alpha * direction[i]; + } + let new_value = f.evaluate(&new_point); + if new_value < current_value { + break; + } + alpha = alpha * T::from(0.5).unwrap(); + } + + + // Get new gradient + let new_gradient = match f.gradient(&new_point) { + Some(g) => g, + None => break, + }; + + // Update the correction vectors + let s = new_point + .iter() + .zip(current_point.iter()) + .map(|(&x_new, &x_old)| x_new - x_old) + .collect::>(); + let y = new_gradient + .iter() + .zip(gradient.iter()) + .map(|(&g_new, &g_old)| g_new - g_old) + .collect::>(); + + let ys = y + .iter() + .zip(s.iter()) + .fold(T::zero(), |acc, (&y_i, &s_i)| acc + y_i * s_i); + + if ys == T::zero() { + break; + } + + let rho = T::one() / ys; + + if s_list.len() == M { + s_list.pop_front(); + y_list.pop_front(); + rho_list.pop_front(); + } + s_list.push_back(s); + y_list.push_back(y); + rho_list.push_back(rho); + + // Update for next iteration + current_point = new_point; + gradient = new_gradient; + iterations += 1; + } + + OptimizationResult { + optimal_point: current_point.clone(), + optimal_value: f.evaluate(¤t_point), + iterations, + converged, + } +} \ No newline at end of file diff --git a/src/par_indexed_iter.rs b/src/par_indexed_iter.rs index e69de29..cfacf28 100644 --- a/src/par_indexed_iter.rs +++ b/src/par_indexed_iter.rs @@ -0,0 +1,407 @@ +use ndarray::{Array, ArrayViewMut, Dimension}; +use rayon::iter::plumbing::{ + Consumer, Folder, Producer, ProducerCallback, UnindexedConsumer, UnindexedProducer, bridge, + bridge_unindexed, +}; +use rayon::prelude::{IndexedParallelIterator, ParallelIterator}; +use std::marker::PhantomData; + +#[derive(Clone, Debug)] +struct BaseIndexIter { + start: usize, + end: usize, + cum_shape: Vec, +} + +impl BaseIndexIter { + fn new(shape: Vec) -> Self { + let mut cum_shape = Vec::with_capacity(shape.len()); + let mut c = 1; + for s in shape.iter().rev() { + cum_shape.push(c); + c *= s; + } + Self { + start: 0, + end: shape.iter().product::(), + cum_shape: cum_shape.into_iter().rev().collect(), + } + } +} + +/// An iterator over array indices. +pub struct IndexIter(BaseIndexIter); + +/// A parallel iterator over array indices. +pub struct ParIndexIter(BaseIndexIter); + +impl IndexIter { + /// Create a new IndexIter using the shape of an array. + pub fn new(shape: Vec) -> Self { + Self(BaseIndexIter::new(shape)) + } +} + +impl ParIndexIter { + /// Create a new ParIndexIter using the shape of an array. + pub fn new(shape: Vec) -> Self { + Self(BaseIndexIter::new(shape)) + } +} + +impl ParallelIterator for ParIndexIter { + type Item = Vec; + + fn drive_unindexed(self, consumer: C) -> C::Result + where + C: UnindexedConsumer, + { + bridge_unindexed(self, consumer) + } +} + +impl IndexedParallelIterator for ParIndexIter { + fn len(&self) -> usize { + self.0.end - self.0.start + } + + fn drive>(self, consumer: C) -> C::Result { + bridge(self, consumer) + } + + fn with_producer>(self, callback: CB) -> CB::Output { + callback.callback(self) + } +} + +impl UnindexedProducer for ParIndexIter { + type Item = Vec; + + fn split(self) -> (Self, Option) { + let length = self.0.end - self.0.start; + if length > 1 { + let (a, b) = self.split_at(length / 2); + (a, Some(b)) + } else { + (self, None) + } + } + + fn fold_with(self, folder: F) -> F + where + F: Folder, + { + folder.consume_iter(IndexIter(self.0)) + } +} + +impl Producer for ParIndexIter { + type Item = Vec; + type IntoIter = IndexIter; + + fn into_iter(self) -> Self::IntoIter { + IndexIter(self.0) + } + + fn split_at(self, index: usize) -> (Self, Self) { + ( + Self(BaseIndexIter { + start: self.0.start, + end: self.0.start + index, + cum_shape: self.0.cum_shape.clone(), + }), + Self(BaseIndexIter { + start: self.0.start + index, + end: self.0.end, + cum_shape: self.0.cum_shape.clone(), + }), + ) + } +} + +impl Iterator for IndexIter { + type Item = Vec; + + fn next(&mut self) -> Option { + if self.0.end <= self.0.start { + None + } else if self.0.cum_shape.len() == 1 { + let i = self.0.start; + self.0.start += 1; + Some(vec![i]) + } else { + let mut n = self.0.start; + let mut i = Vec::new(); + for c in self.0.cum_shape.iter() { + i.push(n / c); + n = n.saturating_sub(n - n % c); + } + self.0.start += 1; + Some(i) + } + } +} + +impl DoubleEndedIterator for IndexIter { + fn next_back(&mut self) -> Option { + if self.0.end <= self.0.start { + None + } else if self.0.cum_shape.len() == 1 { + self.0.end -= 1; + Some(vec![self.0.end]) + } else { + self.0.end -= 1; + let mut n = self.0.end; + let mut i = Vec::new(); + for c in self.0.cum_shape.iter() { + i.push(n / c); + n = n.saturating_sub(n - n % c); + } + Some(i) + } + } +} + +impl ExactSizeIterator for IndexIter {} + +/// A Sync and Send pointer, this should be safe because the pointer is only used with an offset, +/// and each offset is only used once, so no race conditions can occur. +#[derive(Debug, Clone)] +struct Ptr<'a, T> { + ptr: *mut T, + phantom_data: PhantomData<&'a T>, +} +unsafe impl Send for Ptr<'_, T> {} +unsafe impl Sync for Ptr<'_, T> {} + +impl<'a, T> Ptr<'a, T> { + fn at(&self, i: usize) -> &'a mut T { + unsafe { &mut *self.ptr.add(i) } + } +} + +struct BaseIndexedIterMut<'a, T> { + start: usize, + end: usize, + cum_shape: Vec, + ptr: Ptr<'a, T>, +} + +impl<'a, T> BaseIndexedIterMut<'a, T> { + fn new(mut view: ArrayViewMut<'a, T, D>) -> Self { + let shape = view.shape(); + let mut cum_shape = Vec::with_capacity(shape.len()); + let mut c = 1; + for s in shape.iter().rev() { + cum_shape.push(c); + c *= s; + } + Self { + start: 0, + end: shape.iter().product::(), + cum_shape: cum_shape.into_iter().rev().collect(), + ptr: Ptr { + ptr: view.as_mut_ptr(), + phantom_data: PhantomData, + }, + } + } +} + +/// An parallel iterator over array indices and mutable values. +pub struct IndexedIterMut<'a, T>(BaseIndexedIterMut<'a, T>); + +/// An iterator over array indices and mutable values. +pub struct ParIndexedIterMut<'a, T>(BaseIndexedIterMut<'a, T>); + +impl<'a, T> IndexedIterMut<'a, T> { + /// Create a new IndexedIterMut for an array. + pub fn new(view: ArrayViewMut<'a, T, D>) -> Self { + Self(BaseIndexedIterMut::new(view)) + } +} + +impl<'a, T> ParIndexedIterMut<'a, T> { + /// Create a new ParIndexedIterMut for an array. + pub fn new(view: ArrayViewMut<'a, T, D>) -> Self { + Self(BaseIndexedIterMut::new(view)) + } +} + +impl<'a, T> ParallelIterator for ParIndexedIterMut<'a, T> +where + T: Clone + Send, + &'a T: Send, +{ + type Item = (Vec, &'a mut T); + + fn drive_unindexed(self, consumer: C) -> C::Result + where + C: UnindexedConsumer, + { + bridge_unindexed(self, consumer) + } +} + +impl<'a, T> IndexedParallelIterator for ParIndexedIterMut<'a, T> +where + T: Clone + Send, + &'a T: Send, +{ + fn len(&self) -> usize { + self.0.end - self.0.start + } + + fn drive>(self, consumer: C) -> C::Result { + bridge(self, consumer) + } + + fn with_producer>(self, callback: CB) -> CB::Output { + callback.callback(self) + } +} + +impl<'a, T> Producer for ParIndexedIterMut<'a, T> +where + T: Clone, +{ + type Item = (Vec, &'a mut T); + type IntoIter = IndexedIterMut<'a, T>; + + fn into_iter(self) -> Self::IntoIter { + IndexedIterMut(self.0) + } + + fn split_at(self, index: usize) -> (Self, Self) { + ( + Self(BaseIndexedIterMut { + start: self.0.start, + end: self.0.start + index, + cum_shape: self.0.cum_shape.clone(), + ptr: self.0.ptr.clone(), + }), + Self(BaseIndexedIterMut { + start: self.0.start + index, + end: self.0.end, + cum_shape: self.0.cum_shape.clone(), + ptr: self.0.ptr.clone(), + }), + ) + } +} + +impl<'a, T> UnindexedProducer for ParIndexedIterMut<'a, T> +where + T: Clone + Send, + &'a T: Send, +{ + type Item = (Vec, &'a mut T); + + fn split(self) -> (Self, Option) { + let length = self.0.end - self.0.start; + if length > 1 { + let (a, b) = self.split_at(length / 2); + (a, Some(b)) + } else { + (self, None) + } + } + + fn fold_with(self, folder: F) -> F + where + F: Folder, + { + folder.consume_iter(IndexedIterMut(self.0)) + } +} + +impl<'a, T> Iterator for IndexedIterMut<'a, T> { + type Item = (Vec, &'a mut T); + + fn next(&mut self) -> Option { + if self.0.end <= self.0.start { + None + } else if self.0.cum_shape.len() == 1 { + self.0.start += 1; + let j = self.0.start; + Some((vec![j], self.0.ptr.at(j))) + } else { + let mut n = self.0.start; + let mut i = Vec::new(); + for c in self.0.cum_shape.iter() { + i.push(n / c); + n = n.saturating_sub(n - n % c); + } + let j = self.0.start; + self.0.start += 1; + Some((i, self.0.ptr.at(j))) + } + } +} + +impl<'a, T> DoubleEndedIterator for IndexedIterMut<'a, T> { + fn next_back(&mut self) -> Option { + if self.0.end <= self.0.start { + None + } else if self.0.cum_shape.len() == 1 { + self.0.end -= 1; + Some((vec![self.0.end], self.0.ptr.at(self.0.end))) + } else { + self.0.end -= 1; + let mut n = self.0.end; + let mut i = Vec::new(); + for c in self.0.cum_shape.iter() { + i.push(n / c); + n = n.saturating_sub(n - n % c); + } + Some((i, self.0.ptr.at(self.0.end))) + } + } +} + +impl<'a, T> ExactSizeIterator for IndexedIterMut<'a, T> {} + +/// A trait to use arrays as ParallelIndexedIterMut. +pub trait ParallelIndexedIterMut { + type IterMut; + fn par_indexed_iter_mut(self) -> Self::IterMut; +} + +impl<'a, T, D> ParallelIndexedIterMut for ArrayViewMut<'a, T, D> +where + D: Dimension, +{ + type IterMut = ParIndexedIterMut<'a, T>; + fn par_indexed_iter_mut(self) -> Self::IterMut { + Self::IterMut::new(self) + } +} + +impl<'a, T, D> ParallelIndexedIterMut for &'a mut Array +where + D: Dimension, +{ + type IterMut = ParIndexedIterMut<'a, T>; + fn par_indexed_iter_mut(self) -> Self::IterMut { + Self::IterMut::new(self.view_mut()) + } +} + +#[cfg(test)] +mod tests { + use super::IndexIter; + use ndarray::Array3; + + #[test] + fn par_indexed_iter() -> Result<(), Box> { + let shape = [5, 7, 3]; + let i = IndexIter::new(shape.to_vec()).collect::>(); + let j = Array3::::zeros(shape) + .indexed_iter() + .map(|(x, _)| vec![x.0, x.1, x.2]) + .collect::>(); + assert_eq!(i.len(), shape.iter().product::()); + assert_eq!(i, j); + Ok(()) + } +} diff --git a/src/register.rs b/src/register.rs index e69de29..8e0f9e1 100644 --- a/src/register.rs +++ b/src/register.rs @@ -0,0 +1,333 @@ +use crate::bspline::{BSpline, BSplineTrait}; +use crate::error::Error; +use crate::metric::{FixedMu, MattesMetric, SamplingArg, Sigma}; +use crate::transform::Transform; +use crate::optimize::lbfgs_minimize; +use algos::OptimizationConfig; +use ndarray::{AsArray, Dimension}; +use num::cast::AsPrimitive; +use num::integer::Roots; +use std::marker::PhantomData; + +#[derive(Clone, Debug)] +pub struct RegistrationStep { + pub sigma: Sigma, + pub samples: SamplingArg, + pub n_bins: usize, + pub tolerance: f64, + pub edge: f64, + pub max_iterations: usize, + pub learning_rate: f64, +} + +impl RegistrationStep { + pub fn new( + sigma: Sigma, + samples: SamplingArg, + n_bins: usize, + tolerance: f64, + edge: f64, + max_iterations: usize, + learning_rate: f64, + ) -> Self { + Self { + sigma, + samples, + n_bins, + tolerance, + edge, + max_iterations, + learning_rate, + } + } + + pub fn default_steps(ndim: usize, n: usize) -> Vec { + vec![ + RegistrationStep::new( + Sigma::Relative(vec![0.5; ndim]), + SamplingArg::Fixed(n.sqrt().max(n / 125).max(100)), + 3, + 1e-4, + 0.05, + 100, + 100.0, + ), + RegistrationStep::new( + Sigma::Relative(vec![0.25; ndim]), + SamplingArg::Fixed(n.sqrt().max(n / 30).max(200)), + 6, + 1e-6, + 0.04, + 100, + 1.0, + ), + RegistrationStep::new( + Sigma::Absolute(vec![8.0; ndim]), + SamplingArg::Fixed(n.sqrt().max(n / 20).max(400)), + 24, + 3e-7, + 0.03, + 100, + 1e-2, + ), + RegistrationStep::new( + Sigma::Absolute(vec![2.0; ndim]), + SamplingArg::Fixed(n.sqrt().max(n / 10).max(800)), + 96, + 1e-7, + 0.01, + 100, + 1e-2, + ), + RegistrationStep::new( + Sigma::None, + SamplingArg::Fixed(n.sqrt().max(n / 5).max(1600)), + 240, + 1e-8, + 0.001, + 100, + 1e-3, + ), + ] + } +} + +#[derive(Clone, Debug)] +pub struct RegistrationResult { + pub sigma_fixed: Option>, + pub sigma_moving: Option>, + pub n_bins: usize, + pub n_samples: usize, + pub tolerance: f64, + pub edge: f64, + pub max_iterations: usize, + pub learning_rate: f64, + pub optimal_point: Vec, + pub optimal_value: f64, + pub iterations: usize, + pub converged: bool, +} + +pub struct Registration { + fixed_mu: FixedMu, + steps: Option>, + initial_guess: Option>, + dimension: PhantomData, +} + +impl Registration { + pub fn new>(fixed_mu: F) -> Self { + Self { + fixed_mu: fixed_mu.into(), + steps: None, + initial_guess: None, + dimension: PhantomData, + } + } + + pub fn new_affine() -> Result { + let ndim = if let Some(ndim) = D::NDIM { + ndim + } else { + return Err(Error::NumberOfDimensionsNotDefined); + }; + Ok(Self { + fixed_mu: FixedMu::new_none(ndim), + steps: None, + initial_guess: None, + dimension: PhantomData, + }) + } + + pub fn new_translation() -> Result { + let ndim = if let Some(ndim) = D::NDIM { + ndim + } else { + return Err(Error::NumberOfDimensionsNotDefined); + }; + let mut fixed_mu = vec![Some(0.0); ndim * ndim]; + for i in 0..ndim { + fixed_mu[i * (ndim + 1)] = Some(1.0); + fixed_mu[ndim * ndim + i] = None; + } + Ok(Self { + fixed_mu: fixed_mu.into(), + steps: None, + initial_guess: None, + dimension: PhantomData, + }) + } + + pub fn with_steps(mut self, steps: Vec) -> Self { + self.steps = Some(steps); + self + } + + pub fn with_initial_guess(mut self, initial_guess: Vec) -> Self { + self.initial_guess = Some(initial_guess); + self + } + + pub fn set_steps(&mut self, steps: Vec) { + self.steps = Some(steps); + } + + pub fn set_initial_guess(&mut self, initial_guess: Vec) { + self.initial_guess = Some(initial_guess); + } + + /// find the transform which transforms moving into fixed and return the results of each + /// optimization step + pub fn register<'a, F, M, T>(&self, fixed: F, moving: M) -> Result, Error> + where + F: AsArray<'a, T, D>, + M: AsArray<'a, T, D>, + T: 'a + Clone + AsPrimitive, + { + let fixed = fixed.into().mapv(|i| i.as_()); + let moving = moving.into().mapv(|i| i.as_()); + let ndim = fixed.ndim(); + let n = fixed.len(); + let steps = self + .steps + .as_ref() + .cloned() + .unwrap_or_else(|| RegistrationStep::default_steps(ndim, n)); + let mut p = self + .initial_guess + .as_ref() + .cloned() + .unwrap_or_else(|| Transform::::default().parameters); + for RegistrationStep { + sigma, + samples, + n_bins, + tolerance, + edge, + max_iterations, + learning_rate, + } in steps.into_iter() + { + let f = sigma.smooth(fixed.view())?; + let m = sigma.smooth(moving.view())?; + if (f.std(0.0) == 0.0) || (m.std(0.0) == 0.0) { + continue; + } + let bf = BSpline::<0, _>::new(f.view()); + let bm = BSpline::<3, _>::new(m.view()); + let metric = MattesMetric::new(bf, bm, samples, n_bins, edge)? + .with_fixed_mu(self.fixed_mu.clone()); + let optimization_config = OptimizationConfig { + max_iterations, + tolerance, + learning_rate, + }; + let optimization_result = lbfgs_minimize( + &metric, + metric.fixed_mu().extract_variable(&p).as_slice(), + &optimization_config, + ); + if optimization_result + .optimal_point + .iter() + .all(|i| i.is_finite()) + { + p = metric + .fixed_mu() + .combine(&optimization_result.optimal_point); + } + } + Ok(Transform::::new(p, fixed.shape().to_vec())) + } + + /// find the transform which transforms moving into fixed and return the results of each + /// optimization step + pub fn register_debug<'a, F, M, T>( + &self, + fixed: F, + moving: M, + ) -> Result<(Transform, Vec), Error> + where + F: AsArray<'a, T, D>, + M: AsArray<'a, T, D>, + T: 'a + Clone + AsPrimitive, + { + let fixed = fixed.into().mapv(|i| i.as_()); + let moving = moving.into().mapv(|i| i.as_()); + let ndim = fixed.ndim(); + let n = fixed.len(); + let steps = self + .steps + .as_ref() + .cloned() + .unwrap_or_else(|| RegistrationStep::default_steps(ndim, n)); + let mut p = self + .initial_guess + .as_ref() + .cloned() + .unwrap_or_else(|| Transform::::default().parameters); + let mut registration_results = Vec::new(); + for RegistrationStep { + sigma, + samples, + n_bins, + tolerance, + edge, + max_iterations, + learning_rate, + } in steps.into_iter() + { + let f = sigma.smooth(fixed.view())?; + let m = sigma.smooth(moving.view())?; + if (f.std(0.0) == 0.0) || (m.std(0.0) == 0.0) { + continue; + } + let bf = BSpline::<0, _>::new(f.view()); + let bm = BSpline::<3, _>::new(m.view()); + let n_samples = match &samples { + SamplingArg::Fixed(n) => *n, + SamplingArg::Random(n) => *n, + SamplingArg::FixedAt(n) => n.len(), + }; + let metric = MattesMetric::new(bf, bm, samples, n_bins, edge)? + .with_fixed_mu(self.fixed_mu.clone()); + let optimization_config = OptimizationConfig { + max_iterations, + tolerance, + learning_rate, + }; + let optimization_result = lbfgs_minimize( + &metric, + metric.fixed_mu().extract_variable(&p).as_slice(), + &optimization_config, + ); + if optimization_result + .optimal_point + .iter() + .all(|i| i.is_finite()) + { + p = metric + .fixed_mu() + .combine(&optimization_result.optimal_point); + } + registration_results.push(RegistrationResult { + sigma_fixed: sigma.sigma(f.shape()), + sigma_moving: sigma.sigma(m.shape()), + n_bins, + n_samples, + tolerance, + edge, + max_iterations, + learning_rate, + optimal_point: optimization_result.optimal_point, + optimal_value: optimization_result.optimal_value, + iterations: optimization_result.iterations, + converged: optimization_result.converged, + }); + } + Ok(( + Transform::::new(p, fixed.shape().to_vec()), + registration_results, + )) + } +} diff --git a/src/transform.rs b/src/transform.rs index e69de29..0191924 100644 --- a/src/transform.rs +++ b/src/transform.rs @@ -0,0 +1,618 @@ +use crate::bspline::{BSpline, BSplineTrait}; +use crate::error::Error; +use crate::metric::FixedMu; +use crate::register::{Registration, RegistrationResult, RegistrationStep}; +use itertools::Itertools; +use ndarray::{Array, Array2, ArrayD, AsArray, Dimension, Ix2, s}; +use ndarray_linalg::Inverse; +use num::cast::AsPrimitive; +use serde::{Deserialize, Serialize}; +use serde_yaml::{from_reader, to_writer}; +use std::fs::File; +use std::marker::PhantomData; +use std::ops::Mul; +use std::path::PathBuf; + +/// get coordinates resulting from transforming input coordinates, coordinate must have N +/// columns: x, y, z, ... +#[inline] +pub fn transform_point<'a, T>(point: &'a [T], center: &'a [f64], parameters: &'a [f64]) -> Vec +where + T: AsPrimitive, +{ + // [a, b, c] x = a*x + b*y + c + // [d, e, f] y = d*x + e*y + f + // [0, 0, 1] 1 = 1 + debug_assert_eq!(center.len(), point.len()); + debug_assert_eq!((center.len() + 1) * center.len(), parameters.len()); + + let m = center.len(); + let n = m * m; + let point = point + .iter() + .zip_eq(center.iter()) + .map(|(p, c)| p.as_() - c) + .collect::>(); + parameters + .iter() + .take(n) + .chunks(m) + .into_iter() + .zip_eq(parameters.iter().skip(n)) + .zip_eq(center.iter()) + .map(|((r, t), c)| r.zip_eq(point.iter()).map(|(i, p)| i * p).sum::() + t + c) + .collect() +} + +/// a struct describing the transform +/// generic parameter N = # image dimensions +#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)] +pub struct Transform { + /// flattened NxN rotation matrix + N translation parameters + pub parameters: Vec, + /// error / significance on parameters + pub dparameters: Vec, + /// the point about which rotations are performed + pub center: Vec, + /// the shape of images for which this transform is meant + pub shape: Vec, + ndim: usize, + dimension: PhantomData, +} + +impl Mul for Transform { + type Output = Transform; + + fn mul(self, rhs: Self) -> Self::Output { + &self * &rhs + } +} + +impl Mul<&Transform> for Transform { + type Output = Transform; + + fn mul(self, rhs: &Transform) -> Self::Output { + &self * rhs + } +} + +impl Mul> for &Transform { + type Output = Transform; + + fn mul(self, rhs: Transform) -> Self::Output { + self * &rhs + } +} + +impl Mul<&Transform> for &Transform { + type Output = Transform; + + #[allow(clippy::suspicious_arithmetic_impl)] + fn mul(self, rhs: &Transform) -> Self::Output { + let m = self.matrix().dot(&rhs.matrix()); + let dm = self.dmatrix().dot(&rhs.matrix()) + self.matrix().dot(&rhs.dmatrix()); + Self::Output { + parameters: m + .slice(s![..self.ndim, ..self.ndim]) + .flatten() + .iter() + .chain(m.slice(s![..self.ndim, self.ndim]).iter()) + .cloned() + .collect(), + dparameters: dm + .slice(s![..self.ndim, ..]) + .flatten() + .iter() + .chain(dm.slice(s![..self.ndim, ..]).iter()) + .cloned() + .collect(), + center: self.center.clone(), + shape: self.shape.clone(), + ndim: self.ndim, + dimension: self.dimension, + } + } +} + +impl Eq for Transform {} + +impl Default for Transform { + /// the unit transform + fn default() -> Self { + let ndim = D::NDIM.expect("number of dimensions must be known to initialise transform"); + let mut parameters = vec![0.0; ndim * ndim + ndim]; + for i in 0..ndim { + parameters[(ndim + 1) * i] = 1.0; + } + + Self { + parameters, + dparameters: vec![0f64; ndim * ndim + ndim], + center: vec![0f64; ndim], + shape: vec![0usize; ndim], + ndim, + dimension: PhantomData, + } + } +} + +impl Transform { + /// parameters: flat NxN part of matrix + translation; center: center of rotation + pub fn new_with_center(parameters: Vec, center: Vec, shape: Vec) -> Self { + let ndim = if let Some(ndim) = D::NDIM { + ndim + } else { + center.len() + }; + debug_assert_eq!(parameters.len(), ndim * ndim + ndim); + debug_assert_eq!(center.len(), ndim); + debug_assert_eq!(shape.len(), ndim); + Self { + parameters, + dparameters: vec![0.0; ndim * ndim + ndim], + center, + shape, + ndim, + dimension: PhantomData, + } + } + + pub fn new(parameters: Vec, shape: Vec) -> Self { + let ndim = if let Some(ndim) = D::NDIM { + ndim + } else { + shape.len() + }; + debug_assert_eq!(parameters.len(), ndim * ndim + ndim); + debug_assert_eq!(shape.len(), ndim); + let center = shape.iter().map(|s| ((*s as f64) - 1.0) / 2.0).collect(); + Self { + parameters, + dparameters: vec![0.0; ndim * ndim + ndim], + center, + shape, + ndim, + dimension: PhantomData, + } + } + + /// find the affine transform which transforms moving into fixed + pub fn register_affine<'a, F, M, T>(fixed: F, moving: M) -> Result + where + F: AsArray<'a, T, D>, + M: AsArray<'a, T, D>, + T: 'a + Clone + AsPrimitive, + { + Registration::new_affine()?.register(fixed, moving) + } + + /// find the translation which transforms moving into fixed + pub fn register_translation<'a, F, M, T>(fixed: F, moving: M) -> Result + where + F: AsArray<'a, T, D>, + M: AsArray<'a, T, D>, + T: 'a + Clone + AsPrimitive, + { + Registration::new_translation()?.register(fixed, moving) + } + + /// find the transform which transforms moving into fixed, using fixed_mu to specify which + /// parameters to keep fixed + pub fn register<'a, F, M, T, G>(fixed: F, moving: M, fixed_mu: G) -> Result + where + F: AsArray<'a, T, D>, + M: AsArray<'a, T, D>, + T: 'a + Clone + AsPrimitive, + G: Into, + { + Registration::new(fixed_mu).register(fixed, moving) + } + + /// find the transform which transforms moving into fixed and return the results of each + /// optimization step + pub fn register_debug<'a, F, M, T, G>( + fixed: F, + moving: M, + fixed_mu: G, + steps: Option>, + initial_guess: Option>, + ) -> Result<(Self, Vec), Error> + where + F: AsArray<'a, T, D>, + M: AsArray<'a, T, D>, + T: 'a + Clone + AsPrimitive, + G: Into, + { + let mut registration = Registration::new(fixed_mu); + if let Some(steps) = steps { + registration.set_steps(steps); + } + if let Some(initial_guess) = initial_guess { + registration.set_initial_guess(initial_guess); + } + registration.register_debug(fixed, moving) + } + + /// create a transform from a translation + pub fn from_translation(translation: &[f64]) -> Self { + let ndim = if let Some(ndim) = D::NDIM { + ndim + } else { + translation.len() + }; + debug_assert_eq!(translation.len(), ndim); + let mut parameters = vec![0.0; ndim * ndim + ndim]; + for i in 0..ndim { + parameters[(ndim + 1) * i] = 1.0; + } + for (p, t) in parameters + .iter_mut() + .skip(ndim * ndim) + .zip_eq(translation.iter()) + { + *p = *t; + } + + Self { + parameters, + dparameters: vec![0f64; ndim * ndim + ndim], + center: vec![0f64; ndim], + shape: vec![0usize; ndim], + ndim, + dimension: PhantomData, + } + } + + /// create a transform from a scaling + pub fn from_scaling(scaling: &[f64]) -> Self { + let ndim = if let Some(ndim) = D::NDIM { + ndim + } else { + scaling.len() + }; + debug_assert_eq!(scaling.len(), ndim); + let mut parameters = vec![0.0; ndim * ndim + ndim]; + parameters[ndim * ndim + ndim - 1] = 1.0; + for (p, s) in parameters + .iter_mut() + .step_by(ndim + 1) + .zip_eq(scaling.iter()) + { + *p = *s; + } + + Self { + parameters, + dparameters: vec![0f64; ndim * ndim + ndim], + center: vec![0f64; ndim], + shape: vec![0usize; ndim], + ndim, + dimension: PhantomData, + } + } + + /// add a translation to self + pub fn with_translation(&self, translation: &[f64]) -> Self { + Self::from_translation(translation) * self + } + + /// scale self + pub fn with_scaling(&self, scaling: &[f64]) -> Self { + Self::from_scaling(scaling) * self + } + + /// read a transform from a file + pub fn from_file(path: PathBuf) -> Result { + let file = File::open(path)?; + Ok(from_reader(file)?) + } + + /// write a transform to a file + pub fn to_file(&self, path: PathBuf) -> Result<(), Error> { + let mut file = std::fs::OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(path)?; + to_writer(&mut file, self)?; + Ok(()) + } + + /// true if transform does nothing + pub fn is_unity(&self) -> bool { + let n = self.ndim * self.ndim; + let m = self.ndim + 1; + self.parameters + .iter() + .take(n) + .enumerate() + .all(|(i, &x)| if i % m == 0 { x == 1.0 } else { x == 0.0 }) + && self.parameters.iter().skip(n).all(|&x| x == 0.0) + } + + /// get coordinates resulting from transforming input coordinates, coordinate must have N + /// columns: x, y, z, ... + pub fn transform_point<'a, T>(&'a self, point: &'a [T]) -> Vec + where + T: AsPrimitive, + { + transform_point(point, &self.center, &self.parameters) + } + + /// get coordinates resulting from transforming input coordinates, coordinates must have N + /// columns: x, y, z, ... + pub fn transform_points<'a, A, T>(&'a self, points: A) -> Result, Error> + where + T: AsPrimitive, + A: AsArray<'a, T, Ix2>, + { + // [a, b, c] x = a*x + b*y + c + // [d, e, f] y = d*x + e*y + f + // [0, 0, 1] 1 = 1 + let points = points.into(); + let shape = points.shape(); + let a = points + .rows() + .into_iter() + .flat_map(|p| { + if let Some(q) = p.as_slice() { + transform_point(q, &self.center, &self.parameters) + } else { + transform_point(&p.to_vec(), &self.center, &self.parameters) + } + }) + .collect::>(); + Ok(ArrayD::from_shape_vec(shape, a)?.into_dimensionality()?) + } + + /// get the matrix defining the transform + pub fn matrix(&self) -> Array2 { + let n = self.ndim * self.ndim; + let mut matrix = Array2::eye(self.ndim + 1); + for (m, p) in matrix + .slice_mut(s![..self.ndim, ..self.ndim]) + .iter_mut() + .zip_eq(self.parameters.iter().take(n)) + { + *m = *p; + } + for (m, p) in matrix + .slice_mut(s![..self.ndim, self.ndim]) + .iter_mut() + .zip_eq(self.parameters.iter().skip(n)) + { + *m = *p; + } + matrix + } + + /// get the matrix describing the error of the transform + pub fn dmatrix(&self) -> Array2 { + let n = self.ndim * self.ndim; + let mut matrix = Array2::zeros([self.ndim + 1, self.ndim + 1]); + for (m, p) in matrix + .slice_mut(s![..self.ndim, ..self.ndim]) + .iter_mut() + .zip_eq(self.dparameters.iter().take(n)) + { + *m = *p; + } + for (m, p) in matrix + .slice_mut(s![..self.ndim, self.ndim]) + .iter_mut() + .zip_eq(self.dparameters.iter().skip(n)) + { + *m = *p; + } + matrix + } + + /// get the inverse transform + pub fn inverse(&self) -> Result { + let matrix = self.matrix(); + let inverse = matrix.inv()?; + let parameters = inverse + .slice(s![..self.ndim, ..self.ndim]) + .iter() + .chain(inverse.slice(s![..self.ndim, self.ndim])) + .cloned() + .collect(); + + Ok(Self { + parameters, + dparameters: vec![0f64; self.ndim * self.ndim + self.ndim], + center: self.center.clone(), + shape: self.shape.clone(), + ndim: self.ndim, + dimension: PhantomData, + }) + } + + /// adapt the transform to a new center and shape + pub fn adapt(&mut self, center: &[f64], shape: &[usize]) { + self.center = self + .shape + .iter() + .zip_eq(shape.iter()) + .zip_eq(center) + .map(|((a, b), o)| o + (((a - b) as f64) / 2.0)) + .collect(); + self.shape = shape.to_vec(); + } + + pub fn interpolate<'a, const B: usize, T, A>(&self, image: A) -> Result, Error> + where + A: AsArray<'a, T, D>, + D: Dimension, + T: 'a + Clone + AsPrimitive, + BSpline: BSplineTrait, + { + let image = image.into(); + let bspline = BSpline::::new(&image); + bspline.interpolate(self) + } + + pub fn interpolate_par<'a, const B: usize, T, A>( + &self, + image: A, + ) -> Result, Error> + where + A: AsArray<'a, T, D>, + D: Dimension, + T: 'a + Clone + AsPrimitive, + BSpline: BSplineTrait, + { + let image = image.into(); + let bspline = BSpline::::new(&image); + bspline.interpolate_par(self) + } +} + +impl Transform { + /// create a transform from a rotation in radians + pub fn from_rotation(theta: f64, center: &[f64]) -> Self { + Self { + parameters: vec![ + theta.cos(), + -theta.sin(), + theta.sin(), + theta.cos(), + 0.0, + 0.0, + ], + dparameters: vec![0f64; 6], + center: center.to_vec(), + shape: vec![0usize; 2], + ndim: 2, + dimension: PhantomData, + } + } + + /// multiply self with a rotation matrix + pub fn with_rotation(&self, theta: f64) -> Self { + Self::from_rotation(theta, &self.center) * self + } + + /// multiply self with a rotation matrix + pub fn with_rotation_around(&self, theta: f64, center: &[f64]) -> Self { + Self::from_rotation(theta, center) * self + } +} + +#[cfg(test)] +mod tests { + use crate::julia_image; + use crate::transform::Transform; + use itertools::Itertools; + use ndarray::s; + use num::traits::FloatConst; + + #[test] + fn interpolate() -> Result<(), Box> { + let image_a = julia_image( + &[8000, 6000], + &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0], + &[3999.5, 2999.5], + &[-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())?; + + assert!( + image_b + .into_iter() + .zip_eq(image_c.into_iter()) + .all(|(a, b)| a == b) + ); + Ok(()) + } + + #[test] + fn interpolate_serial() -> Result<(), Box> { + let image_a = julia_image( + &[8000, 6000], + &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0], + &[3999.5, 2999.5], + &[-0.8, 0.156], + ); + let transform = Transform::from_rotation(1.0, &[3999.5, 2999.5]); + let image_b = transform.interpolate::<3, _, _>(image_a.view())?; + assert_ne!(image_b[[4000, 3000]], 0.0); + Ok(()) + } + + #[test] + fn interpolate_par() -> Result<(), Box> { + let image_a = julia_image( + &[8000, 6000], + &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0], + &[3999.5, 2999.5], + &[-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())?; + assert_ne!(image_b[[4000, 3000]], 0.0); + Ok(()) + } + + #[test] + fn register1() -> Result<(), Box> { + let im_a = julia_image( + &[100, 1], + &[1.0, 0.0, 0.0, 0.01, 0.0, 0.0], + &[99.5, 0.5], + &[-0.8, 0.156], + ) + .slice(s![.., 0]) + .mapv(|i| i as f64); + let im_b = + Transform::new(vec![0.85, 4.0], vec![im_a.shape()[0]]).interpolate::<1, _, _>(&im_a)?; + + let (t, steps) = + Transform::register_debug(im_a.view(), im_b.view(), vec![None, None], None, None)?; + println!("steps:"); + for step in steps { + println!(" {:?}", step); + } + println!("t: {:?}", t); + println!("i: {:?}", t.inverse()?); + Ok(()) + } + + #[test] + fn register2() -> Result<(), Box> { + let t = Transform::from_rotation(f64::PI() / 4.0, &[299.5, 399.5]); + let s = Transform::new(vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0], vec![600, 800]); + let r = t * s; + let p = r.parameters; + println!("{:?}", p); + + let im_a = julia_image( + &[600, 800], + &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0], + &[299.5, 399.5], + &[-0.8, 0.156], + ); + let im_b = julia_image( + &[600, 800], + &[p[0], p[1], p[2], p[3], p[4], p[5]], + &[299.5, 399.5], + &[-0.8, 0.156], + ); + let (t, steps) = Transform::register_debug( + im_a.view(), + im_b.view(), + vec![None, None, None, None, Some(0.0), Some(0.0)], + None, + None, + )?; + println!("steps:"); + for step in steps { + println!(" {:?}", step); + } + println!("t: {:?}", t); + println!("i: {:?}", t.inverse()?); + Ok(()) + } +}