diff --git a/.gitignore b/.gitignore index 9b13019..d0749ff 100644 --- a/.gitignore +++ b/.gitignore @@ -74,6 +74,7 @@ docs/_build/ AGENTS.md .agentbridge +.agent-work *.tif *.svg diff --git a/Cargo.toml b/Cargo.toml index eea0d76..9e44e70 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,3 +34,6 @@ tempfile = "3" [profile.release] debug = true + +[profile.test] +inherits = "release" \ No newline at end of file diff --git a/src/optimize.rs b/src/optimize.rs index 2b0b9cd..9439eb9 100644 --- a/src/optimize.rs +++ b/src/optimize.rs @@ -188,7 +188,9 @@ pub struct AsgdConfig { pub max_iterations: usize, /// Convergence tolerance (gradient norm) pub tolerance: f64, - /// Maximum step length (displacement in mm). Default: 1.0 + /// Maximum step length. In elastix this is `delta` — used both for initial + /// learning rate estimation (`a = delta * (A+1)^alpha / (jacg + eps)`) and + /// as the per-iteration step clamp. Default: 1.0 pub maximum_step_length: f64, /// Gain parameter A (denominator offset). Default: 20.0 pub sp_a: f64, @@ -298,21 +300,25 @@ where // Compute learning rate: a(t_k) = a / (A + t_k + 1)^alpha let t_k: T = current_time; - let a_t = a / (T::from(config.sp_a).unwrap() + t_k + T::one()).powf(alpha); + let max_step: T = num::cast(config.maximum_step_length).unwrap(); + let a_t = T::min( + a / (T::from(config.sp_a).unwrap() + t_k + T::one()).powf(alpha), + max_step, + ); // Update time using sigmoid: t_{k+1} = max(0, t_k + sigmoid(-g_k^T * g_{k-1})) - // In elastix, the dot product is computed in SCALED space: g_scaled = g/scales. + // In elastix, the dot product is computed in SCALED space: g_scaled = g/sqrt(C). let dot_product_val: f64 = if let Some(ref prev_g) = previous_gradient { let dot_product: T = if let Some(ref scales) = config.scales { - // Scaled dot product: sum(g_i * prev_g_i / scales_i²) - gradient - .iter() - .zip(prev_g.iter()) - .zip(scales.iter()) - .fold(T::zero(), |acc, ((&g, &pg), &s)| { + // Scaled dot product: sum((g/sqrt(C)) * (prev_g/sqrt(C))) = sum(g*pg/C) + gradient.iter().zip(prev_g.iter()).zip(scales.iter()).fold( + T::zero(), + |acc, ((&g, &pg), &s)| { let s2: T = num::cast(s.max(1e-10)).unwrap(); - acc + g * pg / (s2 * s2) - }) + let s4 = s2 * s2; + acc + g * pg / s4 + }, + ) } else { gradient .iter() @@ -329,17 +335,17 @@ where }; // Elastix gradient descent step in unscaled parameter space: - // Optimizer stores p_scaled = p * scales, updates p_scaled -= a_t * g_scaled - // where g_scaled = g / scales (chain rule from GetDerivative). - // Convert back: p_new = p_scaled_new / scales = p - a_t * g / scales² - // So direction = -g / scales². + // m_Position stores the unscaled position p. + // GetScaledDerivative computes g_scaled = g / scales. + // Update: p -= a_t * g_scaled = p - a_t * g / scales. + // So direction = -g / scales. let direction: Vec = match &config.scales { Some(scales) => gradient .iter() .zip(scales.iter()) .map(|(&g, &s)| { let s2: T = num::cast(s.max(1e-10)).unwrap(); - -g / (s2 * s2) + -g / s2 }) .collect(), None => gradient.iter().map(|&g| -g).collect(), @@ -409,14 +415,12 @@ fn sigmoid(x: T, fmax: T, fmin: T, omega: T) -> T { (fmax - fmin) * sigmoid_raw + fmin } -/// Estimate ASGD parameters from gradient statistics. +/// Estimate ASGD parameters using DisplacementDistribution method. /// -/// Matches elastix `AutomaticParameterEstimationOriginal()`: -/// - `a_max = A * delta / sigma1 / sqrt(maxJCJ)` -/// - `sigma1 = sqrt(gg / TrC)` where `gg = ||g_scaled||^2` -/// - `TrC = sum(C_scaled[i][i]) = n` (when scales = sqrt(C[i][i]), C_scaled[i][i] = 1) -/// - `maxJCJ` is approximated as max of squared scaled Jacobian column norms -/// - `g_scaled = g / scales` (golden standard normalization) +/// Matches elastix `AutomaticParameterEstimationUsingDisplacementDistribution()`: +/// - `a = delta * (A+1)^alpha / (jacg + eps)` where `delta = MaximumStepLength` +/// - `jacg ≈ sqrt(gg) * 1.8` as analytical approximation +/// - `maximum_step_length` is also the per-iteration step clamp (same value as delta). fn estimate_asgd_parameters( _f: &F, _initial_point: &[T], @@ -432,54 +436,37 @@ where let fmax = T::from(config.sigmoid_max).unwrap(); let delta = T::from(config.maximum_step_length).unwrap(); - // Compute scaled gradient: g_scaled = g / scales - // And its squared norm: gg = ||g/scales||^2 - let (gg, trc, max_jcj) = if let Some(ref scales) = config.scales { + // Compute gg = ||g/scales||^2 = sum(g²/scales²) + let gg = if let Some(ref scales) = config.scales { let mut gg = T::zero(); for (&g, &s) in initial_gradient.iter().zip(scales.iter()) { let s_t = T::from(s.max(1e-10)).unwrap(); let g_scaled = g / s_t; gg = gg + g_scaled * g_scaled; } - - // Compute maxJCJ analytically for golden standard scales. - // Elastix: maxJCJ = max_j [Tr(J_j C J_j^T) + 2√2 ||J_j C J_j^T||_F] - // For golden standard scales, C ≈ diag(scales²), so C_scaled ≈ I. - // Then J_scaled J_scaled^T = diag(a, ..., a) where - // a = Σ_j max_u((u-c)²)/s_j² + 1 = Σ_j 3(N_j-1)/(N_j+1) + 1 - // and maxJCJ = a · (ndim + 2√(2·ndim)). - let n_params = scales.len(); - let ndim = ((-1.0 + (1.0 + 4.0 * n_params as f64).sqrt()) / 2.0).round() as usize; - let mut a_val = 1.0; // translation contribution: 1/s² = 1 - for j in 0..ndim { - let s_j = scales[j]; - let n_j = (12.0 * s_j * s_j + 1.0).sqrt(); - a_val += 3.0 * (n_j - 1.0) / (n_j + 1.0); - } - let max_jcj = a_val * (ndim as f64 + 2.0 * (2.0 * ndim as f64).sqrt()); - - let n = T::from(n_params).unwrap(); - (gg, n, T::from(max_jcj).unwrap()) + gg } else { let mut gg = T::zero(); for &g in initial_gradient.iter() { gg = gg + g * g; } - let n = T::from(initial_gradient.len()).unwrap(); - (gg, n, T::one()) + gg }; - // Elastix: sigma1 = sqrt(gg / TrC), a_max = A * delta / sigma1 / sqrt(maxJCJ) - let sigma1 = if gg > T::from(1e-14).unwrap() && trc > T::from(1e-14).unwrap() { - (gg / trc).sqrt() + // Compute jacg = displacement distribution estimate. + // For golden standard scales: E[||d||^2] = gg, so jacg ≈ 1.8 * sqrt(gg) + let sqrt_gg = if gg > T::from(1e-20).unwrap() { + gg.sqrt() } else { T::zero() }; + let jacg = sqrt_gg * T::from(1.8).unwrap(); - let a = if sigma1 > T::from(1e-14).unwrap() && max_jcj > T::from(1e-14).unwrap() { - a_param * delta / sigma1 / max_jcj.sqrt() + // Elastix DisplacementDistribution: a = delta * (A+1)^alpha / (jacg + eps) + let a = if jacg > T::from(1e-14).unwrap() { + delta * a_param.powf(alpha) / jacg } else { - a_param + delta * a_param.powf(alpha) }; let fmin = T::from(config.sigmoid_min).unwrap(); diff --git a/src/register.rs b/src/register.rs index e3f4fa0..3d7b7d4 100644 --- a/src/register.rs +++ b/src/register.rs @@ -32,6 +32,7 @@ pub struct RegistrationStep { pub edge: f64, pub max_iterations: usize, pub learning_rate: f64, + pub downsample: usize, } impl RegistrationStep { @@ -52,36 +53,35 @@ impl RegistrationStep { edge, max_iterations, learning_rate, + downsample: 1, } } - /// Default registration steps matching SimpleElastix affine defaults. + /// Default registration steps matching elastix `FixedSmoothingImagePyramid`. /// - /// Elastix uses `FixedSmoothingImagePyramid` with schedule [8, 4, 2, 1] for 4 levels. - /// Sigma = 0.5 * factor * spacing. With spacing=1: sigma = [4.0, 2.0, 1.0, 0.5]. - /// These are absolute sigma values in pixel units. - /// See: `itkMultiResolutionGaussianSmoothingPyramidImageFilter.hxx` + /// Elastix uses `MultiResolutionGaussianSmoothingPyramidImageFilter` which applies + /// Gaussian smoothing with σ = 0.5 × factor at **full resolution** — images are + /// **NOT downsampled**. See: `itkMultiResolutionGaussianSmoothingPyramidImageFilter.hxx` /// - /// Elastix defaults: 2048 samples, 32 histogram bins, 256 iterations. + /// Schedule [8, 4, 2, 1] means σ = [4.0, 2.0, 1.0, 0.5] (spacing=1). pub fn default_steps(ndim: usize, n: usize) -> Vec { let sigma_schedule: Vec = vec![4.0, 2.0, 1.0, 0.5]; let nlevels = sigma_schedule.len(); - let mut steps = Vec::new(); - for (i, &sigma) in sigma_schedule.iter().enumerate() { + for i in 0..nlevels { let fraction = 1.0 / 2.0_f64.powi((nlevels - 1 - i) as i32); let samples = (n as f64 * fraction) .sqrt() .max(n as f64 * fraction / 10.0) .max(2048.0) as usize; steps.push(RegistrationStep::new( - Sigma::Absolute(vec![sigma; ndim]), + Sigma::Absolute(vec![sigma_schedule[i]; ndim]), SamplingArg::Random(samples), 32, 1e-6, 0.05, - 256, + 512, 1.0, )); } @@ -90,6 +90,50 @@ impl RegistrationStep { } } +/// 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>, @@ -255,10 +299,12 @@ impl Registration { edge, max_iterations, learning_rate, + downsample: _, } 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; } @@ -267,6 +313,9 @@ impl Registration { let metric = MattesMetric::new(bf, bm, samples, n_bins, edge)? .with_fixed_mu(self.fixed_mu.clone()); + // Golden standard scales from downsampled shape + let scales = Self::golden_standard_scales(f.shape(), ndim); + let optimization_result = match &self.optimizer { Optimizer::LBFGS => { let config = OptimizationConfig { @@ -281,7 +330,6 @@ impl Registration { ) } Optimizer::ASGD => { - let scales = Self::golden_standard_scales(fixed.shape(), ndim); let config = AsgdConfig { max_iterations, tolerance, @@ -347,13 +395,19 @@ impl Registration { edge, max_iterations, learning_rate, + downsample: _, } in steps.into_iter() { let f = sigma.smooth(fixed.view())?; let m = sigma.smooth(moving.view())?; + + let f_orig_shape: Vec = f.shape().to_vec(); + let m_orig_shape: Vec = m.shape().to_vec(); + 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 { @@ -364,6 +418,9 @@ impl Registration { let metric = MattesMetric::new(bf, bm, samples, n_bins, edge)? .with_fixed_mu(self.fixed_mu.clone()); + // Golden standard scales from downsampled shape + let scales = Self::golden_standard_scales(f.shape(), ndim); + let optimization_result = match &self.optimizer { Optimizer::LBFGS => { let config = OptimizationConfig { @@ -378,7 +435,6 @@ impl Registration { ) } Optimizer::ASGD => { - let scales = Self::golden_standard_scales(fixed.shape(), ndim); let config = AsgdConfig { max_iterations, tolerance, @@ -405,9 +461,10 @@ impl Registration { .fixed_mu() .combine(&optimization_result.optimal_point); } + registration_results.push(RegistrationResult { - sigma_fixed: sigma.sigma(f.shape()), - sigma_moving: sigma.sigma(m.shape()), + sigma_fixed: sigma.sigma(f_orig_shape.as_slice()), + sigma_moving: sigma.sigma(m_orig_shape.as_slice()), n_bins, n_samples, tolerance, diff --git a/src/transform.rs b/src/transform.rs index 4946ad1..6d7f76f 100644 --- a/src/transform.rs +++ b/src/transform.rs @@ -749,22 +749,13 @@ mod tests { ), crate::register::RegistrationStep::new( crate::metric::Sigma::Absolute(vec![2.0]), - crate::metric::SamplingArg::FixedAt(all_points.clone()), + crate::metric::SamplingArg::FixedAt(all_points), 64, 1e-6, 0.04, 200, 1.0, ), - crate::register::RegistrationStep::new( - crate::metric::Sigma::None, - crate::metric::SamplingArg::FixedAt(all_points), - 64, - 1e-8, - 0.001, - 200, - 1.0, - ), ]; let (t, steps) = Transform::register_debug( @@ -792,6 +783,195 @@ mod tests { Ok(()) } + #[test] + fn register2_random_affine() -> Result<(), Box> { + use rand::prelude::*; + let mut rng = rand::rngs::StdRng::seed_from_u64(1337); + + let shape = [200, 200]; + let center = [99.5, 99.5]; + + let fixed = julia_image( + &shape, + &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0], + ¢er, + &[-0.8, 0.156], + ) + .mapv(|i| i as f64); + + // random rotation ±25° + let angle: f64 = rng.random_range(-25.0f64..25.0).to_radians(); + let (s, c) = angle.sin_cos(); + + // random scale 0.85–1.15 + let sx = rng.random_range(0.85..1.15); + let sy = rng.random_range(0.85..1.15); + + // small shear + let shx: f64 = rng.random_range(-0.1..0.1); + let shy: f64 = rng.random_range(-0.1..0.1); + + // translation ±30 px + let tx: f64 = rng.random_range(-30.0..30.0); + let ty: f64 = rng.random_range(-30.0..30.0); + + // [m00, m01, m10, m11, tx, ty] + let params = vec![c * sx, -s * sy + shx, s * sx + shy, c * sy, tx, ty]; + + let transform = + Transform::::new_with_center(params.clone(), center.to_vec(), shape.to_vec()); + let moving = transform.interpolate::<3, _, _>(fixed.view())?; + + let q_inv = transform.inverse()?.parameters; + + let (t, _steps) = Transform::::register_debug( + fixed.view(), + moving.view(), + vec![None, None, None, None, None, None], + None, + None, + )?; + println!("params: {:?}", params); + println!("result: {:?}", t.parameters); + println!("q_inv: {:?}", q_inv); + let sse: f64 = t + .parameters + .iter() + .zip(q_inv.iter()) + .map(|(a, b)| (a - b).powi(2)) + .sum(); + println!("sse: {sse}"); + assert!(sse < 1.0); + Ok(()) + } + + #[test] + fn metric_landscape_2d() -> Result<(), Box> { + use crate::bspline::{BSpline, BSplineTrait}; + use crate::filter::gaussian_smooth; + use crate::metric::{MattesMetric, SamplingArg}; + use algos::ObjectiveFunction; + + let shape = [200, 200]; + let center = [99.5, 99.5]; + let im_a = julia_image( + &shape, + &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0], + ¢er, + &[-0.8, 0.156], + ) + .mapv(|i| i as f64); + let rotation = Transform::::from_rotation(f64::PI() / 4.0, ¢er); + let im_b = rotation.interpolate::<3, _, _>(im_a.view())?; + + let q_inv = rotation.inverse()?.parameters; + let p = rotation.parameters.clone(); + 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) + // 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) ===" + ); + 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])?; + let m = gaussian_smooth(im_b.view(), &[*sigma_val; 2])?; + let bf = BSpline::<0, _>::new(f.view()); + let bm = BSpline::<3, _>::new(m.view()); + let metric = MattesMetric::new(bf, bm, SamplingArg::Random(3000), 128, edge)?; + + let mi_id = metric.evaluate(&identity); + let mi_p = metric.evaluate(&p); + let mi_qinv = metric.evaluate(&q_inv); + + println!( + " Level {} (sigma={:.1}, {}x{}): id={:.4} p={:.4} q_inv={:.4} → {}", + level, + sigma_val, + f.shape()[0], + f.shape()[1], + -mi_id, + -mi_p, + -mi_qinv, + if -mi_qinv > -mi_p { + "Q_INV correct" + } else { + "P incorrect (landscape inverted!)" + } + ); + } + + // Verify: at every sigma level, q_inv should have higher MI than p + let f4 = gaussian_smooth(im_a.view(), &[4.0, 4.0])?; + let m4 = gaussian_smooth(im_b.view(), &[4.0, 4.0])?; + let bf4 = BSpline::<0, _>::new(f4.view()); + let bm4 = BSpline::<3, _>::new(m4.view()); + let metric4 = MattesMetric::new(bf4, bm4, SamplingArg::Random(5000), 128, edge)?; + let mi_p_coarse = metric4.evaluate(&p); + let mi_qinv_coarse = metric4.evaluate(&q_inv); + println!( + "\nCoarsest (sigma=4.0): p={:.6} q_inv={:.6}", + -mi_p_coarse, -mi_qinv_coarse + ); + assert!( + -mi_qinv_coarse > -mi_p_coarse, + "MI landscape is inverted at coarsest level! q_inv={} should be > p={}", + -mi_qinv_coarse, + -mi_p_coarse + ); + + Ok(()) + } + + #[test] + fn angle_sweep() -> Result<(), Box> { + use crate::bspline::{BSpline, BSplineTrait}; + use crate::filter::gaussian_smooth; + use crate::metric::{MattesMetric, SamplingArg}; + use algos::ObjectiveFunction; + let shape = [200, 200]; + let center = [99.5, 99.5]; + let im_a = julia_image( + &shape, + &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0], + ¢er, + &[-0.8, 0.156], + ) + .mapv(|i| i as f64); + let edge = 0.01; + // Full resolution, sigma=4.0 — matching elastix FixedSmoothingImagePyramid (no downsampling) + for angle_deg in [5.0f64, 10.0, 15.0, 20.0, 25.0, 30.0, 35.0, 40.0, 45.0] { + let rotation = Transform::::from_rotation(angle_deg.to_radians(), ¢er); + let im_b = rotation.interpolate::<3, _, _>(im_a.view())?; + let q_inv = rotation.inverse()?.parameters; + let p = rotation.parameters.clone(); + let f = gaussian_smooth(im_a.view(), &[4.0, 4.0])?; + let m = gaussian_smooth(im_b.view(), &[4.0, 4.0])?; + let metric = MattesMetric::new( + BSpline::<0, _>::new(f.view()), + BSpline::<3, _>::new(m.view()), + SamplingArg::Random(3000), + 128, + edge, + )?; + let id = vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0]; + let mi_id = metric.evaluate(&id); + let mi_p = metric.evaluate(&p); + let mi_qi = metric.evaluate(&q_inv); + println!( + "{:5.1}°: id={:.3} p={:.3} q_inv={:.3} → {}", + angle_deg, + -mi_id, + -mi_p, + -mi_qi, + if mi_p > mi_qi { "P wins" } else { "Q_INV wins" } + ); + } + Ok(()) + } + #[test] fn register2_interpolate() -> Result<(), Box> { let shape = [200, 200];