From cc12a1a9b98309a1cecd818cdf26a679182903df Mon Sep 17 00:00:00 2001 From: "w.pomp" Date: Fri, 31 Jul 2026 17:49:11 +0200 Subject: [PATCH] - cleanup --- src/filter.rs | 1 - src/metric.rs | 38 +++------- src/optimize.rs | 28 +++---- src/register.rs | 167 ++++++++++++++-------------------------- src/transform.rs | 194 +++++++++++++++++++++++++++++++++++++++++------ 5 files changed, 253 insertions(+), 175 deletions(-) diff --git a/src/filter.rs b/src/filter.rs index 292301c..8ba8c23 100644 --- a/src/filter.rs +++ b/src/filter.rs @@ -142,7 +142,6 @@ mod tests { Ok(()) } - #[test] fn fft_test() -> Result<(), Box> { let x = Array1::linspace(0.0, 1.0, 11).to_vec(); diff --git a/src/metric.rs b/src/metric.rs index 4fe4d02..0d1ef6e 100644 --- a/src/metric.rs +++ b/src/metric.rs @@ -152,7 +152,7 @@ fn parzen( enum Sampling { Fixed(Vec>), - Random((usize, Vec, RefCell>>)), + Random((usize, Vec)), } impl Sampling { @@ -161,25 +161,19 @@ impl Sampling { } fn random(n_samples: usize, shape: Vec) -> Self { - Self::Random((n_samples, shape, RefCell::new(Vec::new()))) + Self::Random((n_samples, shape)) } fn index(&self) -> Vec> { match self { Self::Fixed(index) => index.clone(), - Self::Random((n_samples, shape, cached)) => { - let mut cache = cached.borrow_mut(); - if cache.is_empty() { - *cache = 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(); - } - cache.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(), } } } @@ -437,23 +431,11 @@ where 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, derivative: dmetric.clone(), }); - // 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(()) } @@ -778,7 +760,7 @@ mod tests { let fixed = BSpline::<0, _>::new(a.view()); let moving = BSpline::<3, _>::new(b.view()); - let mus = vec![0.0, 0.001]; + let mus = [0.0, 0.001]; let mut npz = NpzWriter::new(File::create( std::env::home_dir().unwrap().join("tmp/metric.npz"), )?); diff --git a/src/optimize.rs b/src/optimize.rs index f5ac083..7c20af7 100644 --- a/src/optimize.rs +++ b/src/optimize.rs @@ -358,6 +358,20 @@ mod tests { } } + struct QuadraticWithMinimum; + + impl ObjectiveFunction for QuadraticWithMinimum { + fn evaluate(&self, point: &[f64]) -> f64 { + let x = point[0]; + (x - 2.0).powi(2) + } + + fn gradient(&self, point: &[f64]) -> Option> { + let x = point[0]; + Some(vec![2.0 * (x - 2.0)]) + } + } + #[test] fn test_lbfgs_quadratic() { let f = Quadratic; @@ -379,20 +393,6 @@ mod tests { #[test] fn test_lbfgs_quadratic_with_minimum() { - struct QuadraticWithMinimum; - - impl ObjectiveFunction for QuadraticWithMinimum { - fn evaluate(&self, point: &[f64]) -> f64 { - let x = point[0]; - (x - 2.0).powi(2) - } - - fn gradient(&self, point: &[f64]) -> Option> { - let x = point[0]; - Some(vec![2.0 * (x - 2.0)]) - } - } - let f = QuadraticWithMinimum; let initial_point = vec![0.0]; let config = OptimizationConfig { diff --git a/src/register.rs b/src/register.rs index 630a613..185e25d 100644 --- a/src/register.rs +++ b/src/register.rs @@ -9,8 +9,7 @@ use num::cast::AsPrimitive; use std::marker::PhantomData; /// Optimizer type for registration -#[derive(Clone, Debug, PartialEq)] -#[derive(Default)] +#[derive(Clone, Debug, PartialEq, Default)] pub enum Optimizer { /// L-BFGS optimizer (fast, requires consistent gradients) #[default] @@ -19,7 +18,6 @@ pub enum Optimizer { ASGD, } - #[derive(Clone, Debug)] pub struct RegistrationStep { pub sigma: Sigma, @@ -29,7 +27,6 @@ pub struct RegistrationStep { pub edge: f64, pub max_iterations: usize, pub learning_rate: f64, - pub downsample: usize, /// Optimizer to use for this level pub optimizer: Optimizer, } @@ -52,114 +49,66 @@ impl RegistrationStep { edge, max_iterations, learning_rate, - downsample: 1, optimizer: Optimizer::default(), } } - /// Default registration steps matching elastix `FixedSmoothingImagePyramid`. - /// - /// Uses a hybrid multi-fidelity approach: - /// - Coarse levels (sigma=4,2): ASGD optimizer with all pixels for speed - /// - Fine levels (sigma=1,0.5): L-BFGS optimizer with all pixels for accuracy - /// - /// Note: Samples are cached per level (not randomized per iteration like SimpleElastix), - /// but ASGD's sigmoid momentum handles the stochastic gradients effectively. pub fn default_steps(ndim: usize, n: usize) -> Vec { - if ndim == 1 { - return vec![Self { - sigma: Sigma::Absolute(vec![1.0; 1]), - samples: SamplingArg::Random(n), + vec![ + Self { + sigma: Sigma::Absolute(vec![8.0; ndim]), + samples: SamplingArg::Random((n / 256).max(1024)), + n_bins: 16, + tolerance: 1e-1, + edge: 0.05, + max_iterations: 128, + learning_rate: 1.0, + optimizer: Optimizer::ASGD, + }, + Self { + sigma: Sigma::Absolute(vec![4.0; ndim]), + samples: SamplingArg::Random((n / 64).max(4096)), + n_bins: 24, + tolerance: 1e-2, + edge: 0.05, + max_iterations: 256, + learning_rate: 1.0, + optimizer: Optimizer::ASGD, + }, + Self { + sigma: Sigma::Absolute(vec![2.0; ndim]), + samples: SamplingArg::Random((n / 16).max(4096)), n_bins: 32, + tolerance: 1e-4, + edge: 0.05, + max_iterations: 512, + learning_rate: 1.0, + optimizer: Optimizer::ASGD, + }, + Self { + sigma: Sigma::Absolute(vec![1.0; ndim]), + samples: SamplingArg::Random((n / 4).max(8192)), + n_bins: 48, + tolerance: 1e-6, + edge: 0.05, + max_iterations: 1024, + learning_rate: 1.0, + optimizer: Optimizer::ASGD, + }, + Self { + sigma: Sigma::None, + samples: SamplingArg::Fixed(n.max(16384)), + n_bins: 64, tolerance: 1e-8, edge: 0.05, max_iterations: 2048, learning_rate: 1.0, - downsample: 1, optimizer: Optimizer::LBFGS, - }]; - } - // Pyramid with downsampling: schedule [8,4,2,1], sigma [4,2,1,0.5] - let sigma_schedule: Vec = vec![4.0, 2.0, 1.0, 0.5]; - let downsample_schedule: Vec = vec![8, 4, 2, 1]; - let n_levels = sigma_schedule.len(); - sigma_schedule - .iter() - .zip(downsample_schedule.iter()) - .enumerate() - .map(|(level, (&s, &d))| { - let n_pixels = (n / (d * d)).max(4); - let is_finest = level == n_levels - 1; - let is_coarse = level < 2; - // Hybrid: ASGD on coarse levels for speed, L-BFGS on fine levels for accuracy - // Use all pixels on all levels for consistent gradient estimates - let (optimizer, max_iter, tol) = if is_coarse { - (Optimizer::ASGD, 512, 1e-4) - } else if is_finest { - (Optimizer::LBFGS, 2048, 1e-8) - } else { - (Optimizer::LBFGS, 512, 1e-6) - }; - Self { - sigma: Sigma::Absolute(vec![s; ndim]), - samples: SamplingArg::Random(n_pixels), - n_bins: 32, - tolerance: tol, - edge: 0.05, - max_iterations: max_iter, - learning_rate: 1.0, - downsample: d, - optimizer, - } - }) - .collect() + }, + ] } } -/// Downsample an n-dimensional array by integer factor along all axes. -/// Takes every `factor`-th element along each axis. -/// Only supports 1D and 2D arrays (the only types used in this crate). -pub fn downsample_nd(array: ndarray::ArrayView, factor: usize) -> ndarray::Array -where - D: ndarray::Dimension, -{ - if factor <= 1 { - return array.to_owned(); - } - let ndim = array.ndim(); - let src = array.to_owned(); - let src_shape = src.shape().to_vec(); - let src_slice = src.as_slice().unwrap(); - - let new_shape: Vec = src_shape.iter().map(|&s| (s - 1) / factor + 1).collect(); - - // Compute strides for row-major layout - let mut strides = vec![1usize; ndim]; - for i in (0..ndim - 1).rev() { - strides[i] = strides[i + 1] * src_shape[i + 1]; - } - - let total: usize = new_shape.iter().product(); - let mut data = Vec::with_capacity(total); - for out_flat in 0..total { - let mut out_coords = vec![0usize; ndim]; - let mut tmp = out_flat; - for i in 0..ndim { - out_coords[i] = tmp % new_shape[i]; - tmp /= new_shape[i]; - } - let mut src_flat = 0; - for i in 0..ndim { - src_flat += out_coords[i] * factor * strides[i]; - } - data.push(src_slice[src_flat]); - } - ndarray::Array::from_shape_vec(new_shape, data) - .unwrap() - .into_dimensionality::() - .unwrap() -} - #[derive(Clone, Debug)] pub struct RegistrationResult { pub sigma_fixed: Option>, @@ -283,15 +232,14 @@ impl Registration { fn golden_standard_scales(fixed_shape: &[usize], ndim: usize) -> Vec { let n_params = ndim * ndim + ndim; let mut scales = Vec::with_capacity(n_params); - for _i in 0..ndim { - for j in 0..ndim { - let n = fixed_shape[j] as f64; - scales.push(((n * n - 1.0) / 12.0).sqrt()); - } - } for _ in 0..ndim { - scales.push(1.0); + scales.extend( + fixed_shape + .iter() + .map(|&i| ((i * i - 1) as f64 / 12.0).sqrt()), + ); } + scales.extend(vec![1.0; ndim]); scales } @@ -325,7 +273,6 @@ impl Registration { edge, max_iterations, learning_rate, - downsample: _, optimizer, } in steps.into_iter() { @@ -356,9 +303,10 @@ impl Registration { ); if let Some(result) = optimization_result - && result.optimal_point.iter().all(|i| i.is_finite()) { - p = self.fixed_mu.combine(&result.optimal_point); - } + && result.optimal_point.iter().all(|i| i.is_finite()) + { + p = self.fixed_mu.combine(&result.optimal_point); + } } Ok(Transform::::new(p, fixed.shape().to_vec())) } @@ -431,7 +379,6 @@ impl Registration { edge, max_iterations, learning_rate, - downsample: _, optimizer, } in steps.into_iter() { diff --git a/src/transform.rs b/src/transform.rs index e4895ae..fde35d8 100644 --- a/src/transform.rs +++ b/src/transform.rs @@ -579,7 +579,7 @@ mod tests { use crate::julia_image; use crate::transform::Transform; use itertools::Itertools; - use ndarray::{Array2, Ix2, s}; + use ndarray::{Array2, Ix2, array, s}; use num::traits::FloatConst; use std::fs::File; use std::path::Path; @@ -789,11 +789,11 @@ mod tests { // Test several points let test_points: Vec<(&str, Vec)> = vec![ - ("identity".into(), vec![1.0, 0.0]), - ("truth".into(), vec![0.85, 4.0]), - ("neg_trans".into(), vec![1.0, -4.0]), - ("scale_0.9".into(), vec![0.9, 0.0]), - ("scale_1.1".into(), vec![1.1, 0.0]), + ("identity", vec![1.0, 0.0]), + ("truth", vec![0.85, 4.0]), + ("neg_trans", vec![1.0, -4.0]), + ("scale_0.9", vec![0.9, 0.0]), + ("scale_1.1", vec![1.1, 0.0]), ]; println!("=== SMOOTHED ALL 100 integer points bins=32 ==="); @@ -831,6 +831,96 @@ mod tests { #[test] fn register1() -> Result<(), Box> { + let fixed = array![ + 0.00133107, 0.00138621, 0.00141719, 0.00137673, 0.00133192, 0.00131586, 0.00129737, + 0.00131649, 0.00132131, 0.0013368, 0.0013557, 0.00135557, 0.00135903, 0.00136225, + 0.00136999, 0.0013847, 0.00137785, 0.00138041, 0.00140829, 0.00142085, 0.0014344, + 0.00149646, 0.00156956, 0.00172352, 0.0019473, 0.00226109, 0.00269055, 0.00312467, + 0.00343227, 0.00352425, 0.00335073, 0.00305988, 0.00276482, 0.0024797, 0.00223358, + 0.00204837, 0.00191672, 0.00183362, 0.00179435, 0.00177568, 0.00177995, 0.00177444, + 0.00178801, 0.00180613, 0.00181906, 0.00183161, 0.00185975, 0.00185648, 0.00187759, + 0.00188436, 0.00185824, 0.00185461, 0.00188928, 0.00190146, 0.00185707, 0.00185682, + 0.00185372, 0.00184588, 0.00183992, 0.0018354, 0.0018307, 0.00181903, 0.00182027, + 0.00181646, 0.00180596, 0.00180369, 0.00180493, 0.0018042, 0.00180887, 0.00180575, + 0.00180144, 0.00180303, 0.00180743, 0.00180482, 0.0017817, 0.00175849, 0.00174724, + 0.00174177, 0.00174569, 0.00176198, 0.00178241, 0.00180581, 0.00184414, 0.00188772, + 0.00192021, 0.00195737, 0.00201045, 0.00207885, 0.00216532, 0.00227495, 0.00240301, + 0.00254983, 0.00273925, 0.00298603, 0.00329465, 0.00367626, 0.00415677, 0.0047952, + 0.00565066, 0.006755, 0.00808893, 0.00958421, 0.0111741, 0.01279892, 0.01434347, + 0.0156247, 0.01649142, 0.01690067, 0.01689863, 0.0165402, 0.0158715, 0.01497264, + 0.0139421, 0.01286888, 0.01182453, 0.01084795, 0.00994079, 0.00908256, 0.00826134, + 0.00748309, 0.00676281, 0.00611948, 0.00556657, 0.00510174, 0.00470918, 0.0043702, + 0.00407155, 0.00381196, 0.00359745, 0.00342831, 0.00329565, 0.00318778, 0.00309373, + 0.00300738, 0.00292954, 0.00286483, 0.00281634, 0.00278162, 0.00275487, 0.00273168, + 0.0027102, 0.0026901, 0.00267163, 0.00265643, 0.00264636, 0.00264188, 0.00264216, + 0.00264419, 0.00264316, 0.00263564, 0.00262231, 0.00260767, 0.00259683, 0.00259233, + 0.00259301, 0.00259521, 0.0025956, 0.00259316, 0.00258932, 0.00258678, 0.00258754, + 0.00259157, 0.00259742, 0.00260376, 0.00261015, 0.00261699, 0.00262534, 0.00263624, + 0.00265, 0.00266603, 0.00268315, 0.00269997, 0.00271538, 0.0027293, 0.00274301, + 0.00275855, 0.00277762, 0.00280069, 0.00282693, 0.00285474, 0.00288248, 0.00290876, + 0.00293257, 0.00295365, 0.0029727, 0.00299094, 0.00300944, 0.0030288, 0.00304906, + 0.00306975, 0.00309011, 0.00310917, 0.00312581, 0.00313895, 0.00314813, 0.00315378, + 0.00315703, 0.0031592, 0.00316134, 0.00316407, 0.00316761, 0.00317182, 0.00317628, + 0.00318021, 0.00318272, 0.00318301, 0.00318065, 0.00317563, 0.00316838, 0.00315962, + 0.00315024, 0.00314115, 0.00313327, 0.00312741, 0.0031241, 0.00312351, 0.00312549, + 0.00312965, 0.00313552, 0.00314259, 0.00315038, 0.00315846, 0.00316645, 0.00317408, + 0.00318122, 0.00318795, 0.0031945, 0.00320115, 0.00320818, 0.0032158, 0.00322408, + 0.00323293, 0.00324213, 0.00325147, 0.00326083, 0.00327028, 0.00328004, 0.00329046, + 0.00330194, 0.00331489, 0.0033297, 0.00334668, 0.00336601, 0.00338771, 0.00341162, + 0.0034374, 0.0034646, 0.00349265, 0.003521, 0.00354909, 0.00357647, 0.00360277, + 0.00362773, 0.00365124, 0.00367333, 0.0036942, 0.00371413, 0.00373349, 0.00375265, + 0.0037719, 0.00379146, 0.00381142, 0.00383177, 0.00385234, 0.0038729, 0.0038931, + 0.00391258, 0.00393098, 0.00394801, 0.00396348, 0.00397735, 0.0039897, 0.00400073, + 0.00401069, 0.0040199, 0.00402867, 0.00403726, 0.0040459, 0.00405473, 0.00406382, + 0.00407317, 0.00408273, 0.00409241, 0.00410213, 0.00411181, 0.00412144, 0.00413098, + 0.00414044, 0.00414982, 0.00415912, 0.00416834, 0.00417744, 0.00418636, 0.00419505, + 0.00420339, 0.0042113, 0.00421866, 0.00422539, 0.00423144, 0.00423681, 0.00424153 + ]; + + let q = vec![0.85, 2.0]; + let moving = + Transform::new(q.clone(), vec![fixed.shape()[0]]).interpolate::<3, _, _>(&fixed)?; + + // 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(), vec![fixed.shape()[0]]).inverse()?; + + // Use default steps with ASGD — the multi-resolution pyramid converges + // reliably for this 1D affine problem. + let t = Transform::register(fixed.view(), moving.view(), vec![None; 2], None, None)?; + + // Compare transformed coordinates of all pixels + let mut sum_diff = 0.0; + let mut count = 0; + for point in 0..fixed.len() { + let t_point = t.transform_point(&[point]); + let e_point = q_inv.transform_point(&[point]); + let diff_sq = (t_point[0] - e_point[0]).powi(2); + sum_diff += diff_sq.sqrt(); + count += 1; + } + let mean_diff = sum_diff / count as f64; + println!("Our: {:?}\nmean_coord_diff: {:.6}", t, mean_diff); + + println!("t: {:?}", t); + println!("i: {:?}", t.inverse()?); + println!("q_inv: {:?}", q_inv.parameters); + assert!( + t.parameters + .iter() + .zip(q_inv.parameters.iter()) + .map(|(a, b)| (a - b).powi(2)) + .sum::() + < 1.0 + ); + assert!(mean_diff < 0.1); + Ok(()) + } + #[test] + fn diag_register1_metric() -> Result<(), Box> { + use crate::bspline::{BSpline, BSplineTrait}; + use crate::metric::{FixedMu, MattesMetric, SamplingArg, Sigma}; + use algos::ObjectiveFunction; + let im_a = julia_image( &[100, 1], &[1.0, 0.0, 0.0, 0.01, 0.0, 0.0], @@ -843,32 +933,92 @@ mod tests { let im_b = Transform::new(q.clone(), vec![im_a.shape()[0]]).interpolate::<1, _, _>(&im_a)?; - // The registration finds T such that im_b(T(x)) = im_a(x), which is the inverse of q + let f = Sigma::Absolute(vec![1.0]).smooth(im_a.view())?; + let m = Sigma::Absolute(vec![1.0]).smooth(im_b.view())?; + + let metric = MattesMetric::::new( + BSpline::new(f.view()), + BSpline::new(m.view()), + SamplingArg::Random(100), + 32, + 0.05, + )? + .with_fixed_mu(FixedMu::new_none(1)); + let q_inv = Transform::::new(q.clone(), vec![im_a.shape()[0]]) .inverse()? .parameters; - // Use default steps with ASGD — the multi-resolution pyramid converges - // reliably for this 1D affine problem. - let t = Transform::register(im_a.view(), im_b.view(), vec![None; 2], None, None)?; - println!("t: {:?}", t); - println!("i: {:?}", t.inverse()?); - println!("q_inv: {:?}", q_inv); - assert!( - t.parameters - .iter() - .zip(q_inv.iter()) - .map(|(a, b)| (a - b).powi(2)) - .sum::() - < 1.0 - ); + // sweep scale at the optimal translation, and translation at scale 1.0 + let mut best = (f64::MAX, vec![0.0, 0.0]); + let mut rows = Vec::new(); + for s in (60..=360).map(|i| i as f64 / 100.0) { + for t in (-120..=80).map(|i| i as f64 / 40.0) { + let p = vec![s, t]; + let v = metric.evaluate(&p); + if v < best.0 { + best = (v, p); + } + } + let v0 = metric.evaluate(&[s, -2.353]); + rows.push(format!("s={:.2} t=-2.353 val={:.4}", s, v0)); + } + println!("BEST on grid: val={:.4} params={:?}", best.0, best.1); + println!("rows:"); + + // line from identity to q_inv and beyond + let mut prev = f64::NAN; + for (i, alpha) in (0..=40).map(|i| i as f64 / 10.0).enumerate() { + let p = vec![1.0 + (q_inv[0] - 1.0) * alpha, 0.0 + q_inv[1] * alpha]; + let v = metric.evaluate(&p); + println!("line[{}] alpha={:.1} p={:?} val={:.4}", i, alpha, p, v); + if !prev.is_nan() && v > prev { + println!(" ^ INCREASING here — minimum is before this point"); + } + prev = v; + } + + // metric value along diagonal from q_inv to the registered optimum + for p in [q_inv.clone(), vec![2.0, -1.0], vec![2.5, -0.3]] { + println!("eval {:?}: val={:.4}", p, metric.evaluate(&p)); + } + + // numerical vs analytical gradient at identity and q_inv + let eps = 1e-5; + for p in [vec![1.0, 0.0], q_inv.clone()] { + let grad = metric.gradient(&p).unwrap(); + let mut num = Vec::new(); + for i in 0..2 { + let mut pp = p.clone(); + let mut pm = p.clone(); + pp[i] += eps; + pm[i] -= eps; + num.push((metric.evaluate(&pp) - metric.evaluate(&pm)) / (2.0 * eps)); + } + println!("grad at {:?}: analytical={:?} numerical={:?}", p, grad, num); + } + + let (t, results) = Transform::::register_debug( + im_a.view(), + im_b.view(), + vec![None, None], + None, + None, + )?; + println!("register_debug t: {:?}", t); + for r in results { + println!( + " step: opt={:?} iters={} converged={} val={}", + r.optimal_point, r.iterations, r.converged, r.optimal_value + ); + } Ok(()) } #[test] fn register2_random_affine() -> Result<(), Box> { use rand::prelude::*; - let mut rng = rand::rngs::StdRng::seed_from_u64(1337); + let mut rng = StdRng::seed_from_u64(1337); let shape = [200, 200]; let center = [99.5, 99.5];