diff --git a/README.md b/README.md index c77b00c..60a0b5c 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ Find an ND affine transformation that registers an ND test image to a fixed imag other images. `image-registration` is a pure-Rust, dependency-light implementation of intensity-based image registration in the spirit -of [SimpleElastix]([https://elastix.lumc.nl](https://simpleelastix.github.io/)). It finds an affine (rigid + scale + shear + translation) transform +of [SimpleElastix](https://simpleelastix.github.io/). It finds an affine (rigid + scale + shear + translation) transform that aligns a *moving* image to a *fixed* image by maximizing the [Mattes mutual-information](https://doi.org/10.1109/TMI.2003.809072) metric with a B-spline Parzen window, using a multi-resolution Gaussian image pyramid and a hybrid ASGD + L-BFGS optimization strategy. @@ -12,6 +12,10 @@ window, using a multi-resolution Gaussian image pyramid and a hybrid ASGD + L-BF Everything operates on [`ndarray`](https://docs.rs/ndarray) arrays, works in arbitrary dimensions (1D, 2D, 3D, …), and is parallelized with [rayon](https://docs.rs/rayon). +So far, 2D registration seems to work, but 3D and higher dimensions often do not. Also, this code is still a lot slower +than (Simple)Elastix. The biggest benefit of this crate is that it's Rust and does not require a lengthy compilation of +C++ code. + ## Features - **Affine registration** in any number of dimensions — rotation, scale, shear, and translation parameters @@ -36,13 +40,13 @@ Register a moving image to a fixed image and obtain the affine transform: ```rust use image_registration::error::Error; use image_registration::transform::Transform; -use ndarray::Array2; +use ndarray::{Array2, Ix2}; -fn register(fixed: &Array2, moving: &Array2) -> Result { +fn register(fixed: &Array2, moving: &Array2) -> Result, Error> { // All 6 parameters of the 2D affine transform are free (None = optimize). // steps=None uses the default multi-resolution pyramid; initial_guess=None // initializes by aligning the geometric centers of the two images. - Transform::::register( + Transform::::register( fixed.view(), moving.view(), vec![None; 6], @@ -57,7 +61,8 @@ The returned `Transform` maps coordinates from the *fixed* image to the *moving* ```rust let transform = register(&fixed, &moving)?; -let registered = transform.interpolate_par::<3, _, _>(moving.view())?; // cubic B-spline, parallel +// out-of-range pixels are filled with the default value (0.0 by default) +let registered = transform.interpolate_par::<3, _, _>(moving.view(), &Default::default())?; // cubic B-spline, parallel ``` ## Concepts @@ -71,15 +76,17 @@ in total: - 2D: `[m00, m01, m10, m11, tx, ty]` (row-major linear part, then translation) - 3D: `[m00, m01, m02, m10, m11, m12, m20, m21, m22, tx, ty, tz]` -The transformation is applied about a center saved in ```Transform```: +The transformation is applied about a center `c` that is supplied where the transform is used: ```text v = A·(p − c) + t + c ``` where `A` is the linear part, `t` the translation, `c` the center of the transformation, and `p` the input coordinates. -`Transform::new(parameters, shape)` sets the center to `(shape − 1) / 2` automatically; `Transform::new_with_center` -lets you override it. +`transform_point`/`transform_points` take the center as an argument, and resampling (`interpolate`, `interpolate_par`) +uses the geometric center of the image, `(shape − 1) / 2`. `Transform::new_around_center(parameters, center)` +re-expresses the parameters so that the transformation acts about a chosen center (this is how +`from_rotation`/`with_rotation` rotate about a given point). ### Registration convention @@ -101,19 +108,24 @@ fn main() -> Result<(), Error> { let rot = Transform::::from_rotation(std::f64::consts::FRAC_PI_4, &[99.5, 99.5]); // Compose: translate, then scale, then rotate - let t = rot.with_scaling(&[0.9, 1.1]).with_translation(&[10.0, -5.0]); + let t = rot.add_scaling(&[0.9, 1.1]).with_translation(&[10.0, -5.0]); // Transform a single point (the center is subtracted and re-added) - let p = t.transform_point(&[200, 300]); + let p = t.transform_point(&[200.0, 300.0], &[99.5, 99.5]); // Transform many points at once (rows of an N×2 array) let points = ndarray::array![[0.0, 0.0], [199.0, 199.0]]; - let pts = t.transform_points(points.view())?; + let pts = t.transform_points(points.view(), &[99.5, 99.5])?; // Matrix form (3×3 homogeneous) and exact inverse let m = t.matrix(); let inv = t.inverse()?; - assert!(m.dot(&inv.matrix()).iter().all(|x| (x - 1.0).abs() < 1e-9)); + let identity = ndarray::Array2::::eye(3); + assert!(m + .dot(&inv.matrix()) + .iter() + .zip(identity.iter()) + .all(|(x, y)| (x - y).abs() < 1e-9)); Ok(()) } ``` @@ -132,11 +144,11 @@ fn main() -> Result<(), Error> { let image = ndarray::Array2::::zeros((200, 150)); let transform = Transform::::from_rotation(0.3, &[99.5, 74.5]); - // Order-1 (linear) interpolation, single-threaded - let warped_linear = transform.interpolate::<1, _, _>(image.view())?; + // Order-1 (linear) interpolation, single-threaded (out-of-range pixels → 0.0) + let warped_linear = transform.interpolate::<1, _, _>(image.view(), &Default::default())?; // Order-3 (cubic) B-spline interpolation, parallel - let warped_cubic = transform.interpolate_par::<3, _, _>(image.view())?; + let warped_cubic = transform.interpolate_par::<3, _, _>(image.view(), &Default::default())?; Ok(()) } ``` @@ -163,11 +175,10 @@ fn main() -> Result<(), Error> { let theta = 10.0f64.to_radians(); let (s, c) = theta.sin_cos(); let params = vec![c * 1.05, -s, s, c * 0.95, 12.0, -8.0]; - let ground_truth = - Transform::::new_with_center(params.clone(), center.to_vec(), shape.to_vec()); + let ground_truth = Transform::::new_around_center(params.clone(), ¢er); // Moving image = ground truth applied to the fixed image - let moving = ground_truth.interpolate::<3, _, _>(fixed.view())?; + let moving = ground_truth.interpolate::<3, _, _>(fixed.view(), &Default::default())?; // Register (all 6 parameters free, default steps) let t = Transform::::register(fixed.view(), moving.view(), vec![None; 6], None, None)?; @@ -218,7 +229,7 @@ use image_registration::metric::{SamplingArg, Sigma}; use image_registration::register::{Optimizer, Registration, RegistrationStep}; fn main() -> Result<(), Error> { - let mut coarse = RegistrationStep::new( + let coarse = RegistrationStep::new( Sigma::Absolute(vec![4.0, 4.0]), // Gaussian sigma per dimension SamplingArg::Fixed(3000), // 3000 sampled points (cached) 32, // histogram bins @@ -226,10 +237,10 @@ fn main() -> Result<(), Error> { 0.05, // edge fraction 512, // max iterations 1.0, // learning rate + Optimizer::ASGD, // robust to noisy gradients ); - coarse.optimizer = Optimizer::ASGD; // robust to noisy gradients - let mut fine = RegistrationStep::new( + let fine = RegistrationStep::new( Sigma::Absolute(vec![0.5, 0.5]), SamplingArg::Fixed(3000), 32, @@ -237,9 +248,10 @@ fn main() -> Result<(), Error> { 0.05, 2048, 1.0, - ); // optimizer stays LBFGS (the default) + Optimizer::LBFGS, // the default optimizer + ); - let mut reg = Registration::new(vec![None; 6]); + let mut reg = Registration::new(vec![None; 6])?; reg.set_steps(vec![coarse, fine]); let t = reg.register(fixed.view(), moving.view())?; Ok(()) @@ -284,18 +296,23 @@ Build a `MattesMetric` and evaluate the (negative) mutual information and its gr - `FixedAt(points)` — an explicit list of sample positions (fully deterministic; recommended for reproducible results). ```rust -use algos::ObjectiveFunction; -use image_registration::bspline::BSpline; +use algos::ObjectiveFunction; // direct dependency of image-registration; add it to your crate to name the trait +use image_registration::bspline::{BSpline, BSplineTrait}; use image_registration::error::Error; use image_registration::metric::{MattesMetric, SamplingArg, Sigma}; +use image_registration::transform::Transform; fn main() -> Result<(), Error> { - // `fixed` and `moving` are Array2 - let fixed = ndarray::Array2::::zeros((64, 64)); - let moving = ndarray::Array2::::zeros((64, 64)); + // `fixed` and `moving` are Array2; moving is a rotated + shifted copy of fixed + let center = [31.5f64, 31.5]; + let fixed = julia_image(&[64, 64], &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0], ¢er, &[-0.8, 0.156]) + .mapv(|i| i as f64); + let moving = Transform::::from_rotation(10.0f64.to_radians(), ¢er) + .interpolate_par::<3, _, _>(fixed.view(), &Default::default())?; - let f = Sigma::Absolute(vec![2.0, 2.0]).smooth(fixed.view())?; - let m = Sigma::Absolute(vec![2.0, 2.0]).smooth(moving.view())?; + let sigma = Sigma::Absolute(vec![2.0, 2.0]); + let f = sigma.smooth(fixed.view())?; + let m = sigma.smooth(moving.view())?; let metric = MattesMetric::new( BSpline::<0, _>::new(f.view()), // fixed image B-spline (nearest-neighbor) @@ -336,11 +353,11 @@ fn main() -> Result<(), Error> { | Module | Contents | |----------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `image_registration::transform` | `Transform`, the affine transform type: construction (`new`, `new_around_center`, `from_translation`, `from_scaling`, `from_rotation`), algebra (`inverse`, `add_translation`, `add_scaling`, `with_rotation`, `with_rotation_around`, `Mul`), coordinate transforms (`transform_point`, `transform_points`, `matrix`, `dmatrix`), resampling (`interpolate`, `interpolate_par`), registration (`register`, `register_debug`, `register_affine`, `register_translation`), serialization (`to_file`, `from_file`), plus the free function `transform_point` and a blas-free `matrix_inverse`. | -| `image_registration::register` | `Registration`, `RegistrationStep`, `RegistrationResult`, `Optimizer` (LBFGS/ASGD). | -| `image_registration::metric` | `MattesMetric` (Mattes mutual information with analytic gradient), `SamplingArg` (Fixed/FixedAt/Random), `Sigma` (None/Absolute/Relative), `FixedMu`. | -| `image_registration::bspline` | `BSpline` B-spline interpolation of arbitrary order, `BSplineTrait`, `BSplineMem`, Parzen kernel helpers. | -| `image_registration::filter` | Gaussian smoothing and FFT helpers: `gaussian_smooth`, `gaussian_kernel`, `fft`, `ifft`, `fft_freq`, `fft_shift`. | +| `image_registration::transform` | `Transform`, the affine transform type: construction (`new`, `new_around_center`, `from_translation`, `from_scaling`, `from_rotation`), algebra (`inverse`, `add_translation`, `add_scaling`, `with_rotation`, `Mul`), coordinate transforms (`transform_point`, `transform_points`, `matrix`, `dmatrix`), resampling (`interpolate`, `interpolate_par`, `interpolate_order`, `interpolate_par_order`), registration (`register`, `register_debug`, `register_affine`, `register_translation`, `register_translation_fft`), serialization (`to_file`, `from_file`, and serde), plus the free function `transform_point`. | +| `image_registration::register` | `Registration` (`new_affine`, `new_translation`, `with_steps`, `with_optimizer`, …), `RegistrationStep`, `RegistrationResult`, `Optimizer` (LBFGS/ASGD). | +| `image_registration::metric` | `MattesMetric` (Mattes mutual information with analytic gradient; `evaluate`/`gradient` via the `algos::ObjectiveFunction` trait), `new_from_arrays`, `residual_fit`, `SamplingArg` (Fixed/FixedAt/Random), `Sigma` (None/Absolute/Relative), `FixedMu`. | +| `image_registration::bspline` | `BSpline` B-spline interpolation of arbitrary order, `BSplineTrait`, `BSplineMem`, `DefaultValue` (fill value for out-of-range pixels), Parzen kernel helpers. | +| `image_registration::filter` | Gaussian smoothing and FFT helpers: `gaussian_smooth`, `gaussian_kernel`, `gaussian`, `fft`, `ifft`, `fft_freq`, `fft_shift`. | | `image_registration::optimize` | Optimizers: `lbfgs_minimize`, `asgd_minimize`, `AsgdConfig`. | | `image_registration::par_indexed_iter` | Parallel iterators over array indices. | | `image_registration::error` | `Error` type. | diff --git a/src/bspline.rs b/src/bspline.rs index d7872fc..7c76f11 100644 --- a/src/bspline.rs +++ b/src/bspline.rs @@ -1,4 +1,3 @@ -use std::hash::Hash; use crate::Error; use crate::par_indexed_iter::ParallelIndexedIterMut; use crate::transform::{Transform, transform_point}; @@ -8,12 +7,13 @@ use ndarray::{ }; use num::traits::FloatConst; use num::{Complex, cast::AsPrimitive, complex::ComplexFloat}; +use rand::RngExt; +use rand::distr::StandardUniform; use rayon::iter::ParallelIterator; +use serde::{Deserialize, Serialize}; +use std::hash::Hash; use std::ops::{Deref, MulAssign}; use std::sync::{Arc, Mutex}; -use rand::distr::StandardUniform; -use rand::RngExt; -use serde::{Deserialize, Serialize}; fn in_shape(shape_f: &[f64], v: &[f64]) -> bool { shape_f @@ -72,9 +72,10 @@ impl DefaultValue { } }); let n = weights.iter().sum::(); - weights.iter_mut().fold(0.0, |w, x| { - *x += w / n; - *x + weights.iter_mut().fold(0.0, |mut acc, w| { + acc += *w / n; + *w = acc; + acc }); Self { rng: Self::make_rng(), @@ -96,7 +97,11 @@ impl DefaultValue { self.values[0] } else { let i = self.rng.lock().unwrap().sample::(StandardUniform); - let n = self.weights.iter().position(|&w| i >= w).expect("weights should sum to 1"); + let n = self + .weights + .iter() + .position(|&w| i <= w) + .unwrap_or(self.weights.len() - 1); self.values[n] } } @@ -190,7 +195,11 @@ where D: Dimension, BSpline: BSplineTrait, { - pub fn interpolate(&self, transform: &Transform, default_value: &DefaultValue) -> Result, Error> { + pub fn interpolate( + &self, + transform: &Transform, + default_value: &DefaultValue, + ) -> Result, Error> { let mut mem = self.get_mem(); let shape_f = self.shape_f(); let center = shape_f.iter().map(|i| i / 2.0).collect::>(); @@ -215,7 +224,11 @@ where Ok(new) } - pub fn interpolate_par(&self, transform: &Transform, default_value: &DefaultValue) -> Result, Error> { + pub fn interpolate_par( + &self, + transform: &Transform, + default_value: &DefaultValue, + ) -> Result, Error> { let mem = self.get_mem(); let shape_f = self.shape_f(); let center = shape_f.iter().map(|i| i / 2.0).collect::>(); diff --git a/src/error.rs b/src/error.rs index 6796bc4..5d19658 100644 --- a/src/error.rs +++ b/src/error.rs @@ -19,4 +19,6 @@ pub enum Error { InitialGuessLength(usize, usize), #[error("interpolation order {0} is not implemented")] InterpolationNotImplemented(usize), + #[error("matrix is singular")] + SingularMatrix, } diff --git a/src/lib.rs b/src/lib.rs index e9ef912..6869f09 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,5 @@ +#![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))] +#![doc = include_str!("../README.md")] use rayon::iter::ParallelIterator; pub mod bspline; pub mod error; @@ -148,7 +150,7 @@ mod tests { &[3999.5, 2999.5], &[-0.8, 0.156], ); - let transform = Transform::from_rotation(1.0, &[3999.5, 2999.5]); + let transform = Transform::from_rotation(1.0, &[0.0, 0.0]); let image_b = transform.interpolate_par::<3, _, _>(image_a.view(), &Default::default())?; assert_ne!(image_b[[4000, 3000]], 0.0); Ok(()) diff --git a/src/metric.rs b/src/metric.rs index 0a694ac..db10b1a 100644 --- a/src/metric.rs +++ b/src/metric.rs @@ -9,7 +9,7 @@ use itertools::Itertools; use ndarray::{Array, Array1, Array2, Array3, AsArray, Axis, Dimension, s}; use num::cast::AsPrimitive; use rand::RngExt; -use rayon::iter::{ParallelIterator, IntoParallelIterator}; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; use serde::{Deserialize, Serialize}; use std::cell::RefCell; use std::ops::Deref; @@ -64,6 +64,22 @@ impl FixedMu { .collect() } + /// combine the dparameters of the variable parameters into a complete + /// parameter set, with no uncertainty (0) on the fixed parameters + pub fn combine_dparameters(&self, dparameters: &[f64]) -> Vec { + debug_assert_eq!(dparameters.len(), self.len_variable); + let mut dparameters_iter = dparameters.iter(); + self.iter() + .map(|i| { + if i.is_some() { + 0.0 + } else { + *dparameters_iter.next().unwrap() + } + }) + .collect() + } + /// extract mu from complete parameter set pub fn extract_variable(&self, mu: &[f64]) -> Vec { self.iter() @@ -203,6 +219,24 @@ fn image_jacobian(point: &[f64], center: &[f64], image_derivative: &[f64]) -> Ve .collect() } +/// Local Gauss-Newton statistics of the intensity fit +/// `fixed(x) = moving(T(x)) + noise` at a parameter point. +/// +/// `hessian` is the weighted normal matrix `JᵀWJ` and `variance` is the +/// weighted residual variance `Σ w r² / (Σ w − q − 1)`, with `q` the number of +/// variable parameters. The extra degree of freedom accounts for the weighted +/// mean intensity offset that is removed from the residuals. The parameter +/// covariance of the fit is `variance * hessian⁻¹`. +#[derive(Clone, Debug)] +pub struct ResidualFit { + /// weighted normal matrix `JᵀWJ` + pub hessian: Vec>, + /// residual variance (mean squared residual, corrected for the dof) + pub variance: f64, + /// number of sample points that contributed + pub n_samples: usize, +} + pub struct MattesMetric where D: Dimension, @@ -439,6 +473,93 @@ where } Ok(()) } + + /// Compute the Gauss-Newton normal matrix and the residual variance of the + /// intensity fit at `point`, using the same sample points and edge + /// weighting as [`Self::eval`]. + /// + /// Residuals are `fixed(x) - moving(T(x))`, so a well-registered image pair + /// gives a small variance. Returns `None` when interpolation fails or when + /// too few samples contribute to estimate the residual variance. + pub fn residual_fit(&self, point: &[f64]) -> Option { + let complete_mu = self.fixed_mu.combine(point); + let idx = self.sampling.index(); + let shape = self.shape.as_slice(); + let shape_f = self.shape.iter().map(|s| s - 0.5).collect::>(); + let center = self.center.as_slice(); + let e = 1.0 / self.edge; + let n_params = point.len(); + let mut hessian = vec![vec![0.0; n_params]; n_params]; + let mut sum_w = 0.0; + let mut sum_wr = 0.0; + let mut sum_wr2 = 0.0; + let mut n_samples = 0; + let mut mem_f = self.fixed.get_mem(); + let mut mem_m = self.moving.get_mem(); + for u in idx { + let v = transform_point(&u, center, &complete_mu); + if !shape_f + .iter() + .zip_eq(v.iter()) + .all(|(s, x)| (-0.5 <= *x) && (x <= s)) + { + continue; + } + let f = self + .fixed + .evaluate_at_continuous_index(&u, &mut mem_f) + .ok()?; + let (m, d) = self + .moving + .evaluate_value_and_derivative_at_continuous_index(&v, &mut mem_m) + .ok()?; + // same smooth edge weighting as in `eval` + let mut w = 1.0; + for (s, x) in shape.iter().zip_eq(&v) { + if e * (x + 0.5) < *s { + let k = 3.0 * e * (x + 0.5) / s - 1.5; + w *= square_bspline_integral(k); + } else if e * (x + 0.5) > (e - 1.0) * s { + let k = 3.0 * e * (s - x - 0.5) / s - 1.5; + w *= square_bspline_integral(k); + } + } + if w <= 0.0 || !w.is_finite() { + continue; + } + let r = f - m; + sum_w += w; + sum_wr += w * r; + sum_wr2 += w * r * r; + n_samples += 1; + let j = self + .fixed_mu + .extract_variable(&image_jacobian(&u, center, &d)); + for (a, ja) in j.iter().enumerate() { + for (b, jb) in j.iter().enumerate() { + hessian[a][b] += w * ja * jb; + } + } + } + // the residual variance needs at least one degree of freedom + if n_samples == 0 || sum_w <= n_params as f64 + 1.0 { + return None; + } + // subtract the weighted mean intensity mismatch (e.g. a constant + // offset between the images) before estimating the variance + let mean_r = sum_wr / sum_w; + let sse = (sum_wr2 - sum_w * mean_r * mean_r).max(0.0); + // q parameters plus the mean intensity offset that was removed + let variance = sse / (sum_w - n_params as f64 - 1.0); + if !variance.is_finite() { + return None; + } + Some(ResidualFit { + hessian, + variance, + n_samples, + }) + } } impl ObjectiveFunction for MattesMetric @@ -454,6 +575,39 @@ where self.eval(point).unwrap(); Some(self.metric.borrow().derivative.clone()) } + + /// The hessian is computed numerically: central finite differences of + /// `gradient` with respect to each parameter (2 * n full evaluations). + /// `MattesMetric` only accumulates first-order terms, so an analytic hessian + /// is not available. The result is symmetrized, which also damps the small + /// asymmetry introduced by finite differencing. + fn hessian(&self, point: &[f64]) -> Option>> { + let n = point.len(); + let mut hessian = vec![vec![0.0; n]; n]; + let mut p = point.to_vec(); + for j in 0..n { + let h = 1e-5 * f64::max(1.0, point[j].abs()); + p[j] += h; + self.eval(&p).unwrap(); + let g_plus = self.metric.borrow().derivative.clone(); + p[j] -= 2.0 * h; + self.eval(&p).unwrap(); + let g_minus = self.metric.borrow().derivative.clone(); + p[j] += h; + for i in 0..n { + hessian[i][j] = (g_plus[i] - g_minus[i]) / (2.0 * h); + } + } + #[expect(clippy::needless_range_loop)] + for i in 0..n { + for j in (i + 1)..n { + let w = (hessian[i][j] + hessian[j][i]) / 2.0; + hessian[i][j] = w; + hessian[j][i] = w; + } + } + Some(hessian) + } } #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] @@ -587,6 +741,122 @@ mod tests { Ok(()) } + #[test] + fn residual_fit_1d() -> Result<(), Box> { + let a = array![1.0, 2.0, 3.0, 2.0, 1.0, 2.0, 3.0, 4.0, 2.0, 0.0, 1.0]; + // b is a shifted right by 3 (padded on the left with the first value) + let b = array![1.0, 1.0, 1.0, 1.0, 2.0, 3.0, 2.0, 1.0, 2.0, 3.0, 4.0]; + let points = (0..11).map(|i| vec![i as f64]).collect::>(); + let m = MattesMetric::new_from_arrays( + a.view(), + b.view(), + SamplingArg::FixedAt(points), + 16, + 0.05, + )? + .with_fixed_mu(vec![Some(1.0), None]); + + // fixed(x) = moving(x + 3): no residual + let fit = m.residual_fit(&[3.0]).expect("residual fit"); + assert!(fit.n_samples > 0); + assert!(fit.hessian[0][0] > 0.0); + assert!(fit.variance < 1e-12); + + // away from the optimum the residual variance must grow + let worse = m.residual_fit(&[0.0]).expect("residual fit"); + assert!(worse.variance > fit.variance); + Ok(()) + } + + #[test] + fn hessian_1d() -> 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 p = -3.49; + let h = m.hessian(&[p]).unwrap(); + assert_eq!(h.len(), 1); + assert_eq!(h[0].len(), 1); + + // cross-check against the second central difference of `evaluate` + let step = 1e-4; + let fd = (m.evaluate(&[p + step]) - 2.0 * m.evaluate(&[p]) + m.evaluate(&[p - step])) + / (step * step); + println!("hessian = {:?}\nfinite diff = {:?}", h[0][0], fd); + assert!((h[0][0] - fd).abs() < 1e-3 + 1e-2 * fd.abs()); + Ok(()) + } + + #[test] + fn hessian_2d() -> Result<(), Box> { + let a = Array2::from_shape_fn((12, 12), |(i, j)| i as f64 + 0.7 * j as f64); + let b = Array2::from_shape_fn((12, 12), |(i, j)| 0.5 * i as f64 + 24.0 - 0.7 * j as f64); + let a = gaussian_smooth(a.view(), &[2.0, 2.0])?; + let b = gaussian_smooth(b.view(), &[2.0, 2.0])?; + + let points = (0..12) + .flat_map(|i| (0..12).map(move |j| vec![i as f64, j as f64])) + .step_by(2) + .collect::>(); + + let m = MattesMetric::new_from_arrays( + a.view(), + b.view(), + SamplingArg::FixedAt(points), + 7, + 0.05, + )?; + let n = m.fixed_mu().len_variable(); + assert_eq!(n, 2 * 2 + 2); + + let p = vec![0.5, -0.3, 0.2, 0.7, -0.4, 0.3]; + let h = m.hessian(&p).unwrap(); + assert_eq!(h.len(), n); + assert!(h.iter().all(|row| row.len() == n)); + // symmetrized by construction + for (i, row) in h.iter().enumerate() { + for (j, v) in row.iter().enumerate() { + assert_eq!(*v, h[j][i]); + } + } + + // cross-check each diagonal entry against the 2nd central difference of `evaluate` + let step = 1e-4; + for j in 0..n { + let mut pp = p.clone(); + let mut pm = p.clone(); + pp[j] += step; + pm[j] -= step; + let fd = (m.evaluate(&pp) - 2.0 * m.evaluate(&p) + m.evaluate(&pm)) / (step * step); + println!("diagonal[{j}] = {:?}\nfinite diff = {:?}", h[j][j], fd); + assert!( + (h[j][j] - fd).abs() < 1e-3 + 1e-2 * fd.abs(), + "diagonal entry {j}" + ); + } + Ok(()) + } + #[test] fn test_square_bspline_integral() -> Result<(), Box> { let x = Array1::linspace(-2.0, 2.0, 500); diff --git a/src/register.rs b/src/register.rs index d6f5cb9..681cfb6 100644 --- a/src/register.rs +++ b/src/register.rs @@ -1,10 +1,10 @@ use crate::bspline::{BSpline, BSplineTrait}; use crate::error::Error; -use crate::metric::{FixedMu, MattesMetric, SamplingArg, Sigma}; +use crate::metric::{FixedMu, MattesMetric, ResidualFit, SamplingArg, Sigma}; use crate::optimize::{AsgdConfig, asgd_minimize, lbfgs_minimize}; -use crate::transform::Transform; +use crate::transform::{Transform, try_matrix_inverse}; use algos::{ObjectiveFunction, OptimizationConfig}; -use ndarray::{AsArray, Dimension}; +use ndarray::{Array2, AsArray, Dimension}; use num::cast::AsPrimitive; use serde::{Deserialize, Serialize}; use std::marker::PhantomData; @@ -286,8 +286,11 @@ impl Registration { .unwrap_or_else(|| RegistrationStep::default_steps(ndim, n)); let mut p_var = self.initial_guess.as_ref().cloned().unwrap_or_else(|| { - Self::compute_initial_transform(fixed.shape(), moving.shape(), ndim) + let full = Self::compute_initial_transform(fixed.shape(), moving.shape(), ndim); + self.fixed_mu.extract_variable(&full) }); + let mut last_metric: Option> = None; + let mut last_result: Option> = None; for RegistrationStep { sigma, @@ -324,13 +327,23 @@ impl Registration { &scales, ); - if let Some(result) = optimization_result + last_metric = Some(metric); + last_result = optimization_result; + if let Some(result) = last_result.as_ref() && result.optimal_point.iter().all(|i| i.is_finite()) { - p_var = result.optimal_point; + p_var = result.optimal_point.clone(); } } - Ok(Transform::::new(self.fixed_mu.combine(p_var.as_slice()))) + let mut transform = Transform::::new(self.fixed_mu.combine(p_var.as_slice())); + if let (Some(metric), Some(result)) = (&last_metric, last_result.as_ref()) + && result.converged + && let Some(fit) = metric.residual_fit(&p_var) + { + let dparameters = dparameters_from_residual_fit(&fit); + transform = transform.with_dparameters(self.fixed_mu.combine_dparameters(&dparameters)); + } + Ok(transform) } fn optimize_metric>( @@ -389,9 +402,12 @@ impl Registration { .unwrap_or_else(|| RegistrationStep::default_steps(ndim, n)); let mut p_var = self.initial_guess.as_ref().cloned().unwrap_or_else(|| { - Self::compute_initial_transform(fixed.shape(), moving.shape(), ndim) + let full = Self::compute_initial_transform(fixed.shape(), moving.shape(), ndim); + self.fixed_mu.extract_variable(&full) }); - let mut p = vec![0.0; n * n + n]; + let mut p = vec![0.0; ndim * ndim + ndim]; + let mut last_metric: Option> = None; + let mut last_result: Option> = None; let mut registration_results = Vec::new(); for RegistrationStep { @@ -461,6 +477,8 @@ impl Registration { .fixed_mu() .combine(&optimization_result.optimal_point); } + last_metric = Some(metric); + last_result = Some(optimization_result.clone()); registration_results.push(RegistrationResult { sigma_fixed: sigma.sigma(f_orig_shape.as_slice()), @@ -477,6 +495,265 @@ impl Registration { converged: optimization_result.converged, }); } - Ok((Transform::::new(p), registration_results)) + let mut transform = Transform::::new(p); + if let (Some(metric), Some(result)) = (&last_metric, last_result.as_ref()) + && result.converged + && let Some(fit) = metric.residual_fit(&p_var) + { + let dparameters = dparameters_from_residual_fit(&fit); + transform = transform.with_dparameters(self.fixed_mu.combine_dparameters(&dparameters)); + } + Ok((transform, registration_results)) + } +} + +/// Diagonal of the inverse of a symmetric positive-definite matrix. +/// +/// Non-finite entries are treated as zero. Returns `None` if the matrix is +/// singular, because a singular normal matrix means the corresponding +/// directions are not determined by the data. Diagonal entries of the inverse +/// that are not finite and positive are reported as 0. +fn variance_from_hessian(hessian: &[Vec]) -> Option> { + let n = hessian.len(); + if hessian.iter().any(|row| row.len() != n) { + return None; + } + let mut data: Vec = Vec::with_capacity(n * n); + data.extend( + hessian + .iter() + .flatten() + .map(|v| if v.is_finite() { *v } else { 0.0 }), + ); + let matrix = Array2::from_shape_vec((n, n), data).ok()?; + let inv = try_matrix_inverse(matrix)?; + Some( + (0..n) + .map(|i| { + let v = inv[[i, i]]; + if v.is_finite() && v > 0.0 { v } else { 0.0 } + }) + .collect(), + ) +} + +/// Compute parameter uncertainty (standard error, the `dparameters` field) from +/// the intensity fit at the optimum. +/// +/// The moving image is modelled as `moving(T(x)) = fixed(x) + noise` with +/// independent residuals. Around the optimum this is a linear least-squares +/// problem with Jacobian `J`, so the parameter covariance is +/// +/// ```text +/// Cov = σ² (JᵀWJ)⁻¹ +/// ``` +/// +/// with `σ²` the residual variance and `JᵀWJ` the weighted Gauss-Newton normal +/// matrix, both computed by [`MattesMetric::residual_fit`]. This is the usual +/// standard error of a nonlinear least-squares fit: it is calibrated in +/// parameter units and grows when the images fit poorly or when the image +/// gradients are small. +/// +/// This deliberately does *not* use the inverse Hessian of `-MI`: mutual +/// information is not a likelihood, so its curvature has an arbitrary scale +/// (it changes with the number of histogram bins and the Parzen width) and +/// `H⁻¹` is not a standard error. +/// +/// Each uncertainty is the square root of the corresponding diagonal. If the +/// normal matrix is singular or the residuals cannot be estimated, the result +/// is 0 for the affected parameters, meaning "not estimated". +pub(crate) fn dparameters_from_residual_fit(fit: &ResidualFit) -> Vec { + let n = fit.hessian.len(); + if !fit.variance.is_finite() || fit.variance <= 0.0 { + return vec![0.0; n]; + } + variance_from_hessian(&fit.hessian) + .map(|variance| { + variance + .into_iter() + .map(|v| { + let v = fit.variance * v; + if v.is_finite() && v > 0.0 { + v.sqrt() + } else { + 0.0 + } + }) + .collect() + }) + .unwrap_or_else(|| vec![0.0; n]) +} + +#[cfg(test)] +mod tests { + use super::*; + use ndarray::{Array1, Ix1}; + + #[test] + fn dparameters_from_residual_fit_1d() { + // Cov = variance / (JᵀWJ), so with JᵀWJ = 4 and variance = 1/4 the + // standard error is sqrt(1/16) = 1/4. + let fit = ResidualFit { + hessian: vec![vec![4.0]], + variance: 0.25, + n_samples: 10, + }; + let d = dparameters_from_residual_fit(&fit); + assert!((d[0] - 0.25).abs() < 1e-12); + + // the covariance scales linearly with the residual variance + let fit = ResidualFit { + variance: 1.0, + ..fit + }; + let d = dparameters_from_residual_fit(&fit); + assert!((d[0] - 0.5).abs() < 1e-12); + } + + #[test] + fn dparameters_from_residual_fit_2d() { + // inverse of [[2, 1], [1, 2]] has diagonal 2/3 + let fit = ResidualFit { + hessian: vec![vec![2.0, 1.0], vec![1.0, 2.0]], + variance: 0.5, + n_samples: 20, + }; + let d = dparameters_from_residual_fit(&fit); + let expected: f64 = (1.0f64 / 3.0).sqrt(); + assert!((d[0] - expected).abs() < 1e-9); + assert!((d[1] - expected).abs() < 1e-9); + } + + #[test] + fn dparameters_from_residual_fit_not_estimated() { + // a singular normal matrix leaves the parameters undetermined; there + // is no diagonal fallback + let singular = ResidualFit { + hessian: vec![vec![2.0, -2.0], vec![-2.0, 2.0]], + variance: 1.0, + n_samples: 20, + }; + assert_eq!(dparameters_from_residual_fit(&singular), vec![0.0, 0.0]); + + // a non-positive-definite or non-finite normal matrix is not + // informative either + for hessian in [vec![vec![0.0]], vec![vec![-4.0]], vec![vec![f64::NAN]]] { + let fit = ResidualFit { + hessian, + variance: 1.0, + n_samples: 10, + }; + assert_eq!(dparameters_from_residual_fit(&fit), vec![0.0]); + } + + // no residual variance means no standard error is available + for variance in [0.0, -1.0, f64::NAN, f64::INFINITY] { + let fit = ResidualFit { + hessian: vec![vec![4.0]], + variance, + n_samples: 10, + }; + assert_eq!(dparameters_from_residual_fit(&fit), vec![0.0]); + } + } + + #[test] + fn combine_dparameters() { + let fixed_mu: FixedMu = vec![Some(1.0), None, None].into(); + assert_eq!( + fixed_mu.combine_dparameters(&[0.1, 0.2]), + vec![0.0, 0.1, 0.2] + ); + } + + #[test] + fn register_deterministic_dparameters() -> Result<(), Box> { + // dense, non-periodic 1D signal; b is a, shifted right by 2: the + // optimal translation is +2 + // incommensurate frequencies keep the signal effectively aperiodic, + // so the mutual information peaks at the true shift + let n = 256; + let a: Array1 = (0..n) + .map(|i| { + let x = i as f64; + 1.5 + (0.131 * x).sin() + + 0.7 * (0.397 * x + 1.3).sin() + + 0.3 * (0.883 * x + 0.7).sin() + }) + .collect(); + let b: Array1 = (0..n) + .map(|i| if i < 2 { 0.0 } else { a[[i - 2]] }) + .collect(); + let points: Vec> = (0..n).map(|i| vec![i as f64]).collect(); + let steps = vec![ + RegistrationStep::new( + Sigma::Absolute(vec![2.0]), + SamplingArg::FixedAt(points.clone()), + 24, + 1e-6, + 0.05, + 1024, + 1.0, + Optimizer::LBFGS, + ), + RegistrationStep::new( + Sigma::None, + SamplingArg::FixedAt(points), + 48, + 1e-6, + 0.05, + 2048, + 1.0, + Optimizer::LBFGS, + ), + ]; + let registration = Registration::::new_translation()?.with_steps(steps); + let transform = registration.register(a.view(), b.view())?; + assert_eq!(transform.parameters.len(), 2); + assert_eq!(transform.dparameters.len(), 2); + // the fixed scale parameter has no uncertainty + assert_eq!(transform.dparameters[0], 0.0); + // the data are a perfect shifted copy, so the residual variance is + // (numerically) zero and so is the translation uncertainty + assert!(transform.dparameters[1].is_finite()); + assert!(transform.dparameters[1] < 0.05); + // the found translation should be close to the applied one + assert!((transform.parameters[1] - 2.0).abs() < 1.0); + Ok(()) + } + + #[test] + fn noisy_fit_has_positive_dparameters() -> Result<(), Box> { + let n = 128; + let a: Array1 = (0..n) + .map(|i| { + let x = i as f64; + 1.5 + (0.131 * x).sin() + + 0.7 * (0.397 * x + 1.3).sin() + + 0.3 * (0.883 * x + 0.7).sin() + }) + .collect(); + let b: Array1 = (0..n) + .map(|i| { + let base = if i < 2 { 0.0 } else { a[[i - 2]] }; + // deterministic broadband perturbation + base + 0.05 * ((i * 37 % 19) as f64 / 18.0 - 0.5) + }) + .collect(); + let points: Vec> = (0..n).map(|i| vec![i as f64]).collect(); + let metric = MattesMetric::new_from_arrays( + a.view(), + b.view(), + SamplingArg::FixedAt(points), + 16, + 0.05, + )? + .with_fixed_mu(vec![Some(1.0), None]); + + let fit = metric.residual_fit(&[2.0]).expect("residual fit"); + assert!(fit.variance > 0.0); + let d = dparameters_from_residual_fit(&fit); + assert!(d[0] > 0.0); + Ok(()) } } diff --git a/src/transform.rs b/src/transform.rs index 1a95633..ce03ef4 100644 --- a/src/transform.rs +++ b/src/transform.rs @@ -48,11 +48,18 @@ where .collect() } -/// Calculate the inverse of a square matrix using LU factorization with partial pivoting. -fn matrix_inverse(matrix: Array2) -> Array2 { +/// Calculate the inverse of a square matrix using LU factorization with partial +/// pivoting. +/// +/// Returns `None` if the matrix is singular, or numerically rank-deficient. +pub(crate) fn try_matrix_inverse(matrix: Array2) -> Option> { let n = matrix.nrows(); debug_assert_eq!(n, matrix.ncols(), "matrix must be square"); + // pivot magnitudes must exceed this relative threshold to count as significant + let max_abs = matrix.iter().fold(0.0f64, |m, &x| m.max(x.abs())); + let threshold = max_abs * (n as f64) * f64::EPSILON; + // PA = LU factorization (Doolittle with partial pivoting) let mut lu = matrix; let mut perm: Vec = (0..n).collect(); @@ -64,7 +71,9 @@ fn matrix_inverse(matrix: Array2) -> Array2 { pivot_row = i; } } - assert_ne!(lu[[pivot_row, k]], 0.0, "matrix is singular"); + if lu[[pivot_row, k]].abs() <= threshold || !lu[[pivot_row, k]].is_finite() { + return None; + } if pivot_row != k { for j in 0..n { lu.swap([k, j], [pivot_row, j]); @@ -102,7 +111,7 @@ fn matrix_inverse(matrix: Array2) -> Array2 { } } - inv + Some(inv) } /// a struct describing the transform @@ -403,8 +412,8 @@ impl Transform { } pub fn new_around_center(parameters: Vec, center: &[f64]) -> Self { - Self::new(parameters).with_translation(¢er.iter().map(|i| -i).collect::>()) - * Self::from_translation(center) + Self::new(parameters).with_translation(center) + * Self::from_translation(¢er.iter().map(|i| -i).collect::>()) } /// find the affine transform which transforms moving into fixed @@ -796,7 +805,7 @@ impl Transform { /// get the inverse transform pub fn inverse(&self) -> Result { let matrix = self.matrix(); - let inverse = matrix_inverse(matrix); + let inverse = try_matrix_inverse(matrix).ok_or(Error::SingularMatrix)?; let parameters = inverse .slice(s![..self.ndim(), ..self.ndim()]) .iter() @@ -811,7 +820,11 @@ impl Transform { }) } - pub fn interpolate<'a, const B: usize, T, A>(&self, image: A, default_value: &DefaultValue) -> Result, Error> + pub fn interpolate<'a, const B: usize, T, A>( + &self, + image: A, + default_value: &DefaultValue, + ) -> Result, Error> where A: AsArray<'a, T, D>, D: Dimension, @@ -994,7 +1007,7 @@ impl Transform { #[cfg(test)] mod tests { - use super::matrix_inverse; + use super::try_matrix_inverse; use crate::julia_image; use crate::transform::File; use crate::transform::Transform; @@ -1011,7 +1024,7 @@ mod tests { let a = array![[4.0, 7.0], [2.0, 6.0]]; let expected = array![[0.6, -0.7], [-0.2, 0.4]]; - let got = matrix_inverse(a.clone()); + let got = try_matrix_inverse(a.clone()).expect("matrix should not be singular"); assert!( got.iter() .zip_eq(expected.iter()) @@ -1020,7 +1033,7 @@ mod tests { let b = array![[0.0, 2.0], [3.0, 4.0]]; let expected = array![[-2.0 / 3.0, 1.0 / 3.0], [0.5, 0.0]]; - let got = matrix_inverse(b.clone()); + let got = try_matrix_inverse(b.clone()).expect("matrix should not be singular"); assert!( got.iter() .zip_eq(expected.iter()) @@ -1029,13 +1042,16 @@ mod tests { let c = array![[1.0, 2.0, 3.0], [0.0, 1.0, 4.0], [5.0, 6.0, 0.0]]; let reference = array![[-24.0, 18.0, 5.0], [20.0, -15.0, -4.0], [-5.0, 4.0, 1.0]]; - let got = matrix_inverse(c); + let got = try_matrix_inverse(c).expect("matrix should not be singular"); assert!( got.iter() .zip_eq(reference.iter()) .all(|(x, y)| (x - y).abs() < 1e-10) ); + let singular = array![[1.0, 2.0], [2.0, 4.0]]; + assert!(try_matrix_inverse(singular).is_none()); + Ok(()) } @@ -1186,7 +1202,8 @@ mod tests { ) .slice(s![.., 0]) .mapv(|i| i as f64); - let im_b = Transform::new(vec![0.85, 4.0]).interpolate::<1, _, _>(&im_a, &Default::default())?; + let im_b = + Transform::new(vec![0.85, 4.0]).interpolate::<1, _, _>(&im_a, &Default::default())?; // Smooth both images with sigma=8 like level 0 let sigma = Sigma::Absolute(vec![8.0]); @@ -1295,7 +1312,8 @@ mod tests { ]; let q = vec![0.85, 2.0]; - let moving = Transform::new(q.clone()).interpolate::<3, _, _>(&fixed, &Default::default())?; + let moving = + Transform::new(q.clone()).interpolate::<3, _, _>(&fixed, &Default::default())?; // The registration finds T such that im_b(T(x)) = im_a(x), which is the inverse of q let q_inv = Transform::::new(q.clone()).inverse()?; @@ -1520,11 +1538,11 @@ mod tests { let identity = vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0]; let edge = 0.01; - // Test at full resolution with different sigma levels (matching elastix approach) + // Test at full resolution with different sigma levels (matching Elastix approach) // Elastix does NOT downsample — it smooths at full resolution with // sigma = 0.5 * factor where factor is from the pyramid schedule [8, 4, 2, 1] println!( - "=== Full resolution (no downsampling, matching elastix FixedSmoothingImagePyramid) ===" + "=== Full resolution (no downsampling, matching Elastix FixedSmoothingImagePyramid) ===" ); for (level, sigma_val) in [4.0, 2.0, 1.0, 0.5].iter().enumerate() { let f = gaussian_smooth(im_a.view(), &[*sigma_val; 2])?;