diff --git a/src/lib.rs b/src/lib.rs index 140efb1..986f0e9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,10 +3,10 @@ pub mod bspline; pub mod error; pub mod filter; pub mod metric; +pub mod optimize; pub mod par_indexed_iter; pub mod register; pub mod transform; -mod optimize; use ndarray::prelude::*; use num::Complex; diff --git a/src/metric.rs b/src/metric.rs index ec9400c..1e98e91 100644 --- a/src/metric.rs +++ b/src/metric.rs @@ -152,7 +152,7 @@ fn parzen( enum Sampling { Fixed(Vec>), - Random((usize, Vec)), + Random((usize, Vec, RefCell>>)), } impl Sampling { @@ -161,19 +161,21 @@ impl Sampling { } fn random(n_samples: usize, shape: Vec) -> Self { - Self::Random((n_samples, shape)) + Self::Random((n_samples, shape, RefCell::new(Vec::new()))) } 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(), + Self::Random((n_samples, shape, _cached)) => { + 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() + } } } } @@ -240,6 +242,7 @@ where .into_option() .expect("m and f cannot be empty"); let minmax = [*min, *max]; + let ndim = shape.len(); let sampling = match sampling { SamplingArg::Fixed(n_samples) => { Sampling::fixed(Sampling::random(n_samples, shape.clone()).index()) @@ -253,7 +256,7 @@ where center, fixed, moving, - fixed_mu: FixedMu::new_none(sampling.index().len()), + fixed_mu: FixedMu::new_none(ndim), minmax, sampling, n_bins, @@ -319,7 +322,7 @@ where 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 j = fixed_mu.extract_variable(&image_jacobian(&u, center, &d)); let mut k; let mut w = 1.0; let mut ws = Vec::with_capacity(shape.len()); @@ -347,7 +350,7 @@ where for (dwi, wi) in dw.iter_mut().zip(ws) { *dwi *= w / wi; } - let dw = fixed_mu.extract_variable(&image_jacobian(&v, center, &dw)); + let dw = fixed_mu.extract_variable(&image_jacobian(&u, center, &dw)); Ok(Some((f, ((m, j), (w, dw))))) } else { Ok(None) @@ -435,12 +438,8 @@ where 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(), + metric: metric, + derivative: dmetric.clone(), }); // self.metric.replace(IntMut { // mu: mu.to_vec(), diff --git a/src/optimize.rs b/src/optimize.rs index fc59a72..2b0b9cd 100644 --- a/src/optimize.rs +++ b/src/optimize.rs @@ -1,8 +1,9 @@ -use std::collections::VecDeque; -use std::fmt::Debug; use algos::{ObjectiveFunction, OptimizationConfig, OptimizationResult}; use num::Float; +use std::collections::VecDeque; +use std::fmt::Debug; +/// L-BFGS optimizer with proper line search and initial Hessian scaling pub fn lbfgs_minimize( f: &F, initial_point: &[T], @@ -55,9 +56,9 @@ where 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); + .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; @@ -84,9 +85,9 @@ where 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); + .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; @@ -96,23 +97,36 @@ where // r now contains the search direction let direction: Vec = r.iter().map(|&x| -x).collect(); - // Line search to find step size + // Line search to find step size with Armijo sufficient decrease condition let mut alpha = T::one(); let mut new_point = vec![T::zero(); n]; let current_value = f.evaluate(¤t_point); + let mut improved = false; + // directional derivative g^T * d (should be negative for descent) + let g_dot_d: T = gradient + .iter() + .zip(direction.iter()) + .fold(T::zero(), |acc, (&g, &d)| acc + g * d); + let c1 = T::from(1e-4).unwrap(); // Armijo constant - // Simple backtracking line search - for _ in 0..20 { + // Backtracking line search with Armijo condition + for _ in 0..30 { 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 { + // Armijo: f(x + αd) ≤ f(x) + c₁ α (gᵀd) + if new_value <= current_value + c1 * alpha * g_dot_d { + improved = true; break; } alpha = alpha * T::from(0.5).unwrap(); } + // FIX: If no improvement found, don't update (would corrupt L-BFGS history) + if !improved { + break; + } // Get new gradient let new_gradient = match f.gradient(&new_point) { @@ -137,7 +151,8 @@ where .zip(s.iter()) .fold(T::zero(), |acc, (&y_i, &s_i)| acc + y_i * s_i); - if ys == T::zero() { + // FIX: Check for non-positive ys (corrupts Hessian approximation) + if ys <= T::zero() { break; } @@ -151,7 +166,7 @@ where 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; @@ -164,4 +179,400 @@ where iterations, converged, } -} \ No newline at end of file +} + +/// Configuration for Adaptive Stochastic Gradient Descent (ASGD) optimizer +#[derive(Clone, Debug)] +pub struct AsgdConfig { + /// Maximum number of iterations per resolution + pub max_iterations: usize, + /// Convergence tolerance (gradient norm) + pub tolerance: f64, + /// Maximum step length (displacement in mm). Default: 1.0 + pub maximum_step_length: f64, + /// Gain parameter A (denominator offset). Default: 20.0 + pub sp_a: f64, + /// Learning rate parameter alpha. Default: 1.0 + pub sp_alpha: f64, + /// Sigmoid maximum. Default: 1.0 + pub sigmoid_max: f64, + /// Sigmoid minimum. Elastix: -0.99 + 0.98*noisefactor. For low noise: -0.01 + pub sigmoid_min: f64, + /// Sigmoid scale (omega). Elastix: SigmoidScaleFactor * sigma3^2 * sqrt(TrCC). + /// For low noise (exact gradients), omega ≈ 0 (step function sigmoid). + pub sigmoid_scale: f64, + /// Per-parameter scales for gradient normalization (golden standard). + /// scales[i] = C[i][i] = E[|J_i|^2], the diagonal of Jacobian covariance. + pub scales: Option>, +} + +impl Default for AsgdConfig { + fn default() -> Self { + Self { + max_iterations: 250, + tolerance: 1e-6, + maximum_step_length: 1.0, + sp_a: 20.0, + sp_alpha: 1.0, + sigmoid_max: 1.0, + sigmoid_min: -0.01, + sigmoid_scale: 1e-8, + scales: None, + } + } +} + +/// Adaptive Stochastic Gradient Descent (ASGD) optimizer +/// +/// This implements the ASGD optimizer from elastix, which is well-suited for +/// stochastic gradient estimates from Mattes Mutual Information metric. +/// +/// The learning rate at iteration k is: +/// a(t_k) = a / (A + t_k + 1)^alpha +/// +/// Where time t_k is updated via sigmoid: +/// t_{k+1} = max(0, t_k + sigmoid(-g_k^T * g_{k-1})) +pub fn asgd_minimize(f: &F, initial_point: &[T], config: &AsgdConfig) -> OptimizationResult +where + T: Float + Debug, + F: ObjectiveFunction, +{ + let n = initial_point.len(); + let mut current_point = initial_point.to_vec(); + let mut iterations = 0; + let mut converged = false; + + // 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, + }; + } + }; + + eprintln!( + "asgd_minimize: initial_point={:?} gradient={:?}", + initial_point + .iter() + .map(|x| num::cast::<_, f64>(*x).unwrap()) + .collect::>(), + gradient + .iter() + .map(|x| num::cast::<_, f64>(*x).unwrap()) + .collect::>(), + ); + + // Estimate ASGD parameters from gradient statistics + let (a, alpha, fmax, fmin, omega) = + estimate_asgd_parameters(f, ¤t_point, &gradient, config); + + eprintln!( + "asgd_minimize: a={:.6} scales={:?}", + num::cast::<_, f64>(a).unwrap(), + config.scales.as_ref().map(|s| s + .iter() + .map(|x| num::cast::<_, f64>(*x).unwrap()) + .collect::>()), + ); + + // Initial time + let mut current_time = T::zero(); + let mut previous_gradient: Option> = None; + + while iterations < config.max_iterations { + // Check for convergence + let gradient_norm = gradient + .iter() + .fold(T::zero(), |acc, &x| acc + x * x) + .sqrt(); + let tol: T = num::cast(config.tolerance).unwrap(); + if gradient_norm < tol { + converged = true; + break; + } + + // 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); + + // 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. + 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)| { + let s2: T = num::cast(s.max(1e-10)).unwrap(); + acc + g * pg / (s2 * s2) + }) + } else { + gradient + .iter() + .zip(prev_g.iter()) + .fold(T::zero(), |acc, (&g, &pg)| acc + g * pg) + }; + + // Elastix: sigmoid(-inprod) — time increases when gradients disagree + let sigmoid_val = sigmoid(-dot_product, fmax, fmin, omega); + current_time = T::max(T::zero(), current_time + sigmoid_val); + num::cast(dot_product).unwrap() + } else { + 0.0 + }; + + // 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². + 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) + }) + .collect(), + None => gradient.iter().map(|&g| -g).collect(), + }; + + let mut new_point = current_point.clone(); + for i in 0..n { + new_point[i] = current_point[i] + direction[i] * a_t; + } + + if iterations < 5 || iterations % 20 == 0 { + eprintln!( + "iter={}: point={:?} grad={:?} dir={:?} a_t={:.4} t={:.4} dot={:.6}", + iterations, + current_point + .iter() + .map(|x| num::cast::<_, f64>(*x).unwrap()) + .collect::>(), + gradient + .iter() + .map(|x| num::cast::<_, f64>(*x).unwrap()) + .collect::>(), + direction + .iter() + .map(|x| num::cast::<_, f64>(*x).unwrap()) + .collect::>(), + num::cast::<_, f64>(a_t).unwrap(), + num::cast::<_, f64>(current_time).unwrap(), + dot_product_val, + ); + } + + // Get new gradient + let new_gradient = match f.gradient(&new_point) { + Some(g) => g, + None => break, + }; + + // Update for next iteration + previous_gradient = Some(gradient.clone()); + current_point = new_point; + gradient = new_gradient; + iterations += 1; + } + + OptimizationResult { + optimal_point: current_point.clone(), + optimal_value: f.evaluate(¤t_point), + iterations, + converged, + } +} + +/// Sigmoid function for ASGD time update. +/// Matches ITK SigmoidImageFilter: f(x) = (Max-Min) * 1/(1+exp(-(x-beta)/alpha)) + Min +/// where alpha = omega, beta = omega * ln(-fmax/fmin), so sigmoid(0) = 0. +fn sigmoid(x: T, fmax: T, fmin: T, omega: T) -> T { + let beta = omega * (-fmax / fmin).ln(); + let z = (x - beta) / omega; + let sigmoid_raw = if z > T::from(20.0).unwrap() { + T::one() + } else if z < T::from(-20.0).unwrap() { + T::zero() + } else { + T::one() / (T::one() + (-z).exp()) + }; + (fmax - fmin) * sigmoid_raw + fmin +} + +/// Estimate ASGD parameters from gradient statistics. +/// +/// 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) +fn estimate_asgd_parameters( + _f: &F, + _initial_point: &[T], + initial_gradient: &[T], + config: &AsgdConfig, +) -> (T, T, T, T, T) +where + T: Float + Debug, + F: ObjectiveFunction, +{ + let a_param = T::from(config.sp_a).unwrap(); + let alpha = T::from(config.sp_alpha).unwrap(); + 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 { + 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()) + } 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()) + }; + + // 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() + } else { + T::zero() + }; + + let a = if sigma1 > T::from(1e-14).unwrap() && max_jcj > T::from(1e-14).unwrap() { + a_param * delta / sigma1 / max_jcj.sqrt() + } else { + a_param + }; + + let fmin = T::from(config.sigmoid_min).unwrap(); + let omega = T::from(config.sigmoid_scale).unwrap(); + + (a, alpha, fmax, fmin, omega) +} + +#[cfg(test)] +mod tests { + use super::*; + use algos::OptimizationConfig; + + // Test function: f(x, y) = x^2 + y^2 + struct Quadratic; + + impl ObjectiveFunction for Quadratic { + fn evaluate(&self, point: &[f64]) -> f64 { + point.iter().map(|x| x * x).sum() + } + + fn gradient(&self, point: &[f64]) -> Option> { + Some(point.iter().map(|x| 2.0 * x).collect()) + } + } + + #[test] + fn test_lbfgs_quadratic() { + let f = Quadratic; + let initial_point = vec![1.0, 1.0]; + let config = OptimizationConfig { + max_iterations: 100, + tolerance: 1e-6, + learning_rate: 1.0, + }; + + let result = lbfgs_minimize(&f, &initial_point, &config); + + assert!(result.converged); + assert!(result.optimal_value < 1e-10); + for x in result.optimal_point { + assert!(x.abs() < 1e-5); + } + } + + // Test function: f(x) = (x - 2)^2 + 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_with_minimum() { + let f = QuadraticWithMinimum; + let initial_point = vec![0.0]; + let config = OptimizationConfig { + max_iterations: 100, + tolerance: 1e-6, + learning_rate: 1.0, + }; + + let result = lbfgs_minimize(&f, &initial_point, &config); + + assert!(result.converged); + assert!((result.optimal_point[0] - 2.0).abs() < 1e-5); + } + + #[test] + fn test_asgd_quadratic() { + let f = Quadratic; + let initial_point = vec![1.0, 1.0]; + let config = AsgdConfig { + max_iterations: 1000, + tolerance: 1e-4, + sigmoid_min: -0.8, + ..Default::default() + }; + + let result = asgd_minimize(&f, &initial_point, &config); + + // ASGD should converge close to minimum + assert!(result.optimal_value < 0.1); + for x in result.optimal_point { + assert!(x.abs() < 0.5); + } + } +} diff --git a/src/register.rs b/src/register.rs index 8e0f9e1..e3f4fa0 100644 --- a/src/register.rs +++ b/src/register.rs @@ -1,14 +1,28 @@ use crate::bspline::{BSpline, BSplineTrait}; use crate::error::Error; use crate::metric::{FixedMu, MattesMetric, SamplingArg, Sigma}; +use crate::optimize::{AsgdConfig, asgd_minimize, lbfgs_minimize}; 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; +/// Optimizer type for registration +#[derive(Clone, Debug, PartialEq)] +pub enum Optimizer { + /// L-BFGS optimizer (fast, requires consistent gradients) + LBFGS, + /// Adaptive Stochastic Gradient Descent (robust, handles noisy gradients) + ASGD, +} + +impl Default for Optimizer { + fn default() -> Self { + Self::ASGD + } +} + #[derive(Clone, Debug)] pub struct RegistrationStep { pub sigma: Sigma, @@ -41,54 +55,38 @@ impl RegistrationStep { } } + /// Default registration steps matching SimpleElastix affine defaults. + /// + /// 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 defaults: 2048 samples, 32 histogram bins, 256 iterations. 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, + 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() { + 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]), + SamplingArg::Random(samples), + 32, 1e-6, - 0.04, - 100, + 0.05, + 256, 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, - ), - ] + )); + } + + steps } } @@ -112,6 +110,7 @@ pub struct Registration { fixed_mu: FixedMu, steps: Option>, initial_guess: Option>, + optimizer: Optimizer, dimension: PhantomData, } @@ -121,6 +120,7 @@ impl Registration { fixed_mu: fixed_mu.into(), steps: None, initial_guess: None, + optimizer: Optimizer::default(), dimension: PhantomData, } } @@ -135,6 +135,7 @@ impl Registration { fixed_mu: FixedMu::new_none(ndim), steps: None, initial_guess: None, + optimizer: Optimizer::default(), dimension: PhantomData, }) } @@ -154,6 +155,7 @@ impl Registration { fixed_mu: fixed_mu.into(), steps: None, initial_guess: None, + optimizer: Optimizer::default(), dimension: PhantomData, }) } @@ -168,6 +170,11 @@ impl Registration { self } + pub fn with_optimizer(mut self, optimizer: Optimizer) -> Self { + self.optimizer = optimizer; + self + } + pub fn set_steps(&mut self, steps: Vec) { self.steps = Some(steps); } @@ -176,6 +183,48 @@ impl Registration { self.initial_guess = Some(initial_guess); } + pub fn set_optimizer(&mut self, optimizer: Optimizer) { + self.optimizer = optimizer; + } + + /// Compute automatic transform initialization by aligning geometric centers + fn compute_initial_transform( + fixed_shape: &[usize], + moving_shape: &[usize], + ndim: usize, + ) -> Vec { + let mut params = vec![0.0; ndim * ndim + ndim]; + for i in 0..ndim { + params[(ndim + 1) * i] = 1.0; + } + for i in 0..ndim { + let fixed_center = (fixed_shape[i] as f64 - 1.0) / 2.0; + let moving_center = (moving_shape[i] as f64 - 1.0) / 2.0; + params[ndim * ndim + i] = fixed_center - moving_center; + } + params + } + + /// Compute golden standard scales for ASGD parameter normalization. + /// Matches elastix ScaledCostFunction convention: scales = sqrt(C[i][i]). + /// For affine transform v = R*(u-c) + t + c: + /// - Rotation/scale params: C[i][i] = Var(u-c) = (N^2-1)/12, scales[i] = sqrt(C[i][i]) + /// - Translation params: C[i][i] = 1.0, scales[i] = 1.0 + 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 + } + /// 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> @@ -193,11 +242,11 @@ impl Registration { .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 p = self.initial_guess.as_ref().cloned().unwrap_or_else(|| { + Self::compute_initial_transform(fixed.shape(), moving.shape(), ndim) + }); + for RegistrationStep { sigma, samples, @@ -217,16 +266,39 @@ impl Registration { 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 = match &self.optimizer { + Optimizer::LBFGS => { + let config = OptimizationConfig { + max_iterations, + tolerance, + learning_rate, + }; + lbfgs_minimize( + &metric, + metric.fixed_mu().extract_variable(&p).as_slice(), + &config, + ) + } + Optimizer::ASGD => { + let scales = Self::golden_standard_scales(fixed.shape(), ndim); + let config = AsgdConfig { + max_iterations, + tolerance, + maximum_step_length: 1.0, + sp_a: 20.0, + sp_alpha: 1.0, + scales: Some(scales), + ..Default::default() + }; + asgd_minimize( + &metric, + metric.fixed_mu().extract_variable(&p).as_slice(), + &config, + ) + } }; - let optimization_result = lbfgs_minimize( - &metric, - metric.fixed_mu().extract_variable(&p).as_slice(), - &optimization_config, - ); + if optimization_result .optimal_point .iter() @@ -261,11 +333,11 @@ impl Registration { .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 p = self.initial_guess.as_ref().cloned().unwrap_or_else(|| { + Self::compute_initial_transform(fixed.shape(), moving.shape(), ndim) + }); + let mut registration_results = Vec::new(); for RegistrationStep { sigma, @@ -291,16 +363,39 @@ impl Registration { }; 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 = match &self.optimizer { + Optimizer::LBFGS => { + let config = OptimizationConfig { + max_iterations, + tolerance, + learning_rate, + }; + lbfgs_minimize( + &metric, + metric.fixed_mu().extract_variable(&p).as_slice(), + &config, + ) + } + Optimizer::ASGD => { + let scales = Self::golden_standard_scales(fixed.shape(), ndim); + let config = AsgdConfig { + max_iterations, + tolerance, + maximum_step_length: 1.0, + sp_a: 20.0, + sp_alpha: 1.0, + scales: Some(scales), + ..Default::default() + }; + asgd_minimize( + &metric, + metric.fixed_mu().extract_variable(&p).as_slice(), + &config, + ) + } }; - let optimization_result = lbfgs_minimize( - &metric, - metric.fixed_mu().extract_variable(&p).as_slice(), - &optimization_config, - ); + if optimization_result .optimal_point .iter() diff --git a/src/transform.rs b/src/transform.rs index 0191924..4946ad1 100644 --- a/src/transform.rs +++ b/src/transform.rs @@ -504,7 +504,7 @@ mod tests { use crate::julia_image; use crate::transform::Transform; use itertools::Itertools; - use ndarray::s; + use ndarray::{Ix2, s}; use num::traits::FloatConst; #[test] @@ -556,6 +556,167 @@ mod tests { Ok(()) } + #[test] + fn grad_check() -> Result<(), Box> { + use crate::bspline::{BSpline, BSplineTrait}; + use crate::metric::{MattesMetric, SamplingArg}; + use algos::ObjectiveFunction; + 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 q = vec![0.85, 4.0]; + let im_b = + Transform::new(q.clone(), vec![im_a.shape()[0]]).interpolate::<1, _, _>(&im_a)?; + + // Test with ALL points (like grad_check uses) + let metric_all = MattesMetric::::new( + BSpline::new(im_b.view()), + BSpline::new(im_a.view()), + SamplingArg::FixedAt((0..100).map(|i| vec![i as f64]).collect()), + 3, + 0.05, + )? + .with_fixed_mu(crate::metric::FixedMu::new_none(1)); + + let identity = vec![1.0, 0.0]; + let val_i = metric_all.evaluate(&identity); + let grad_i = metric_all.gradient(&identity).unwrap(); + println!("ALL POINTS - identity: val={}, grad={:?}", val_i, grad_i); + let truth = vec![0.85, 4.0]; + let val_t = metric_all.evaluate(&truth); + let grad_t = metric_all.gradient(&truth).unwrap(); + println!("ALL POINTS - truth: val={}, grad={:?}", val_t, grad_t); + let inv = vec![1.0 / 0.85, -4.0 / 0.85]; + let val_inv = metric_all.evaluate(&inv); + println!("ALL POINTS - inverse: val={}", val_inv); + + // Test with 100 random points (like register1 uses at finest level) + let metric_rand = MattesMetric::::new( + BSpline::new(im_b.view()), + BSpline::new(im_a.view()), + SamplingArg::Fixed(100), + 3, + 0.05, + )? + .with_fixed_mu(crate::metric::FixedMu::new_none(1)); + + let val_i2 = metric_rand.evaluate(&identity); + let grad_i2 = metric_rand.gradient(&identity).unwrap(); + println!("RAND 100 - identity: val={}, grad={:?}", val_i2, grad_i2); + + // Test with 100 random points and 32 bins (like level 0) + let metric_32 = MattesMetric::::new( + BSpline::new(im_b.view()), + BSpline::new(im_a.view()), + SamplingArg::Fixed(100), + 32, + 0.05, + )? + .with_fixed_mu(crate::metric::FixedMu::new_none(1)); + + let val_i3 = metric_32.evaluate(&identity); + let grad_i3 = metric_32.gradient(&identity).unwrap(); + println!( + "RAND 100 bins=32 - identity: val={}, grad={:?}", + val_i3, grad_i3 + ); + + let eps = 1e-5; + for i in 0..2 { + let mut p_plus = identity.clone(); + let mut p_minus = identity.clone(); + p_plus[i] += eps; + p_minus[i] -= eps; + let num_grad = + (metric_all.evaluate(&p_plus) - metric_all.evaluate(&p_minus)) / (2.0 * eps); + println!( + "numerical d/dmu[{}] = {} (analytical: {})", + i, num_grad, grad_i[i] + ); + } + Ok(()) + } + + #[test] + fn metric_landscape() -> Result<(), Box> { + use crate::bspline::{BSpline, BSplineTrait}; + use crate::metric::{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], + &[99.5, 0.5], + &[-0.8, 0.156], + ) + .slice(s![.., 0]) + .mapv(|i| i as f64); + let q = vec![0.85, 4.0]; + let im_b = + Transform::new(vec![0.85, 4.0], vec![im_a.shape()[0]]).interpolate::<1, _, _>(&im_a)?; + + // Smooth both images with sigma=8 like level 0 + let sigma = Sigma::Absolute(vec![8.0]); + let sf = sigma.smooth(im_a.view())?; + let sm = sigma.smooth(im_b.view())?; + + let points: Vec> = (0..100).map(|i| vec![i as f64]).collect(); + + let metric = MattesMetric::::new( + BSpline::new(sf.view()), + BSpline::new(sm.view()), + SamplingArg::FixedAt(points.clone()), + 32, + 0.05, + )? + .with_fixed_mu(crate::metric::FixedMu::new_none(1)); + + // 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]), + ]; + + println!("=== SMOOTHED ALL 100 integer points bins=32 ==="); + for (name, p) in &test_points { + let val = metric.evaluate(p); + let grad = metric.gradient(p); + match grad { + Some(g) => println!(" {}: val={:.6}, grad={:?}", name, val, g), + None => println!(" {}: val={:.6}, grad=None", name, val), + } + } + + // Now test with Fixed(100) random points + let metric_rand = MattesMetric::::new( + BSpline::new(sf.view()), + BSpline::new(sm.view()), + SamplingArg::Fixed(100), + 32, + 0.05, + )? + .with_fixed_mu(crate::metric::FixedMu::new_none(1)); + + println!("=== SMOOTHED Random 100 points bins=32 ==="); + for (name, p) in &test_points { + let val = metric_rand.evaluate(p); + let grad = metric_rand.gradient(p); + match grad { + Some(g) => println!(" {}: val={:.6}, grad={:?}", name, val, g), + None => println!(" {}: val={:.6}, grad=None", name, val), + } + } + + Ok(()) + } + #[test] fn register1() -> Result<(), Box> { let im_a = julia_image( @@ -566,17 +727,111 @@ mod tests { ) .slice(s![.., 0]) .mapv(|i| i as f64); + let q = vec![0.85, 4.0]; let im_b = - Transform::new(vec![0.85, 4.0], vec![im_a.shape()[0]]).interpolate::<1, _, _>(&im_a)?; + Transform::new(q.clone(), 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)?; + // 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![im_a.shape()[0]]) + .inverse()? + .parameters; + + let all_points: Vec> = (0..100).map(|j| vec![j as f64]).collect(); + let steps = vec![ + crate::register::RegistrationStep::new( + crate::metric::Sigma::Absolute(vec![8.0]), + crate::metric::SamplingArg::FixedAt(all_points.clone()), + 32, + 1e-4, + 0.05, + 200, + 1.0, + ), + crate::register::RegistrationStep::new( + crate::metric::Sigma::Absolute(vec![2.0]), + crate::metric::SamplingArg::FixedAt(all_points.clone()), + 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( + im_a.view(), + im_b.view(), + vec![None, None], + Some(steps), + None, + )?; println!("steps:"); for step in steps { println!(" {:?}", step); } 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 + ); + Ok(()) + } + + #[test] + fn register2_interpolate() -> Result<(), Box> { + 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; + + // Use default_steps which matches elastix FixedSmoothingImagePyramid: + // sigma = [4.0, 2.0, 1.0, 0.5] (schedule [8,4,2,1] with spacing=1) + let (t, steps) = Transform::register_debug( + im_a.view(), + im_b.view(), + vec![None, None, None, None, None, None], + None, + None, + )?; + println!("steps:"); + for step in steps { + println!(" {:?}", step); + } + 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 + ); Ok(()) } @@ -603,7 +858,7 @@ mod tests { let (t, steps) = Transform::register_debug( im_a.view(), im_b.view(), - vec![None, None, None, None, Some(0.0), Some(0.0)], + vec![None, None, None, None, None, None], None, None, )?; @@ -613,6 +868,19 @@ mod tests { } println!("t: {:?}", t); println!("i: {:?}", t.inverse()?); + println!("p: {:?}", p); + // julia_image applies transform to coordinates, so T maps im_b->im_a means T = p_inv + let p_inv = Transform::::new(p.clone(), vec![600, 800]) + .inverse()? + .parameters; + assert!( + t.parameters + .iter() + .zip(p_inv.iter()) + .map(|(a, b)| (a - b).powi(2)) + .sum::() + < 1.0 + ); Ok(()) } }