From bb3d46cc9d6826042c8c815a170f562d7eb75df9 Mon Sep 17 00:00:00 2001 From: Wim Pomp Date: Wed, 29 Jul 2026 20:30:57 +0200 Subject: [PATCH] - registration getting better --- src/optimize.rs | 255 +++++++++++------------------------------------ src/register.rs | 40 ++++++-- src/transform.rs | 145 +++++++++------------------ 3 files changed, 136 insertions(+), 304 deletions(-) diff --git a/src/optimize.rs b/src/optimize.rs index 9439eb9..8354070 100644 --- a/src/optimize.rs +++ b/src/optimize.rs @@ -13,18 +13,16 @@ where T: Float + Debug, F: ObjectiveFunction, { - const M: usize = 10; // Number of corrections to store + const M: usize = 10; let n = initial_point.len(); let mut current_point = initial_point.to_vec(); let mut iterations = 0; let mut converged = false; - // Storage for the last M corrections let mut s_list: VecDeque> = VecDeque::with_capacity(M); let mut y_list: VecDeque> = VecDeque::with_capacity(M); let mut rho_list: VecDeque = VecDeque::with_capacity(M); - // Get initial gradient let mut gradient = match f.gradient(¤t_point) { Some(g) => g, None => { @@ -38,7 +36,6 @@ where }; while iterations < config.max_iterations { - // Check for convergence let gradient_norm = gradient .iter() .fold(T::zero(), |acc, &x| acc + x * x) @@ -48,11 +45,9 @@ where break; } - // Compute search direction using L-BFGS two-loop recursion let mut q = gradient.clone(); let mut alpha_list = Vec::with_capacity(s_list.len()); - // First loop for i in (0..s_list.len()).rev() { let alpha = rho_list[i] * s_list[i] @@ -65,7 +60,6 @@ where } } - // Scale the initial Hessian approximation let mut r = if !s_list.is_empty() { let i = s_list.len() - 1; let yy = y_list[i].iter().fold(T::zero(), |acc, &y| acc + y * y); @@ -81,7 +75,6 @@ where q }; - // Second loop for i in 0..s_list.len() { let beta = rho_list[i] * y_list[i] @@ -94,28 +87,23 @@ where } } - // r now contains the search direction let direction: Vec = r.iter().map(|&x| -x).collect(); - // 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 + let c1 = T::from(1e-4).unwrap(); - // 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); - // Armijo: f(x + αd) ≤ f(x) + c₁ α (gᵀd) if new_value <= current_value + c1 * alpha * g_dot_d { improved = true; break; @@ -123,18 +111,15 @@ where 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) { Some(g) => g, None => break, }; - // Update the correction vectors let s = new_point .iter() .zip(current_point.iter()) @@ -151,7 +136,6 @@ where .zip(s.iter()) .fold(T::zero(), |acc, (&y_i, &s_i)| acc + y_i * s_i); - // FIX: Check for non-positive ys (corrupts Hessian approximation) if ys <= T::zero() { break; } @@ -167,7 +151,6 @@ where y_list.push_back(y); rho_list.push_back(rho); - // Update for next iteration current_point = new_point; gradient = new_gradient; iterations += 1; @@ -181,30 +164,16 @@ where } } -/// 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. 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, - /// 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>, } @@ -215,7 +184,7 @@ impl Default for AsgdConfig { tolerance: 1e-6, maximum_step_length: 1.0, sp_a: 20.0, - sp_alpha: 1.0, + sp_alpha: 0.602, sigmoid_max: 1.0, sigmoid_min: -0.01, sigmoid_scale: 1e-8, @@ -224,16 +193,6 @@ impl Default for AsgdConfig { } } -/// 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, @@ -244,7 +203,6 @@ where let mut iterations = 0; let mut converged = false; - // Get initial gradient let mut gradient = match f.gradient(¤t_point) { Some(g) => g, None => { @@ -257,37 +215,47 @@ where } }; - 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); + let alpha = T::from(config.sp_alpha).unwrap(); + let a_param = T::from(config.sp_a).unwrap(); + let delta = T::from(config.maximum_step_length).unwrap(); + let fmax = T::from(config.sigmoid_max).unwrap(); + let fmin = T::from(config.sigmoid_min).unwrap(); + let omega = T::from(config.sigmoid_scale).unwrap(); - 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::>()), - ); + // Compute max_j = max(|g_i / scales_i|) + let max_j = if let Some(ref scales) = config.scales { + let mut mj = T::zero(); + for (&g, &s) in gradient.iter().zip(scales.iter()) { + let s_t = T::from(s.max(1e-10)).unwrap(); + let gs = (g / s_t).abs(); + if gs > mj { + mj = gs; + } + } + mj + } else { + let mut mj = T::zero(); + for &g in gradient.iter() { + let ga = g.abs(); + if ga > mj { + mj = ga; + } + } + mj + }; + + // a = delta * (A+1)^alpha / max_j (Elastix AutomaticParameterEstimation) + let a = if max_j > T::from(1e-14).unwrap() { + delta * (a_param + T::one()).powf(alpha) / max_j + } else { + delta * (a_param + T::one()).powf(alpha) + }; - // 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) @@ -298,19 +266,13 @@ where break; } - // Compute learning rate: a(t_k) = a / (A + t_k + 1)^alpha - let t_k: T = current_time; - 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, - ); + // a(k) = min(a / (A + t_k + 1)^alpha, delta) + let t_k = current_time + T::from(iterations).unwrap(); + let a_t = T::min(a / (a_param + t_k + T::one()).powf(alpha), delta); - // 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/sqrt(C). - let dot_product_val: f64 = if let Some(ref prev_g) = previous_gradient { + // Update time: t_{k+1} = max(0, t_k + sigmoid(-g_k^T * g_{k-1})) + if let Some(ref prev_g) = previous_gradient { let dot_product: T = if let Some(ref scales) = config.scales { - // 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)| { @@ -325,20 +287,12 @@ where .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 - }; + } + previous_gradient = Some(gradient.clone()); - // Elastix gradient descent step in unscaled parameter space: - // 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. + // direction = -g / scales let direction: Vec = match &config.scales { Some(scales) => gradient .iter() @@ -356,36 +310,11 @@ where 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; @@ -399,9 +328,6 @@ where } } -/// 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; @@ -415,72 +341,11 @@ fn sigmoid(x: T, fmax: T, fmin: T, omega: T) -> T { (fmax - fmin) * sigmoid_raw + fmin } -/// Estimate ASGD parameters using DisplacementDistribution method. -/// -/// 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], - 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 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; - } - gg - } else { - let mut gg = T::zero(); - for &g in initial_gradient.iter() { - gg = gg + g * g; - } - gg - }; - - // 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(); - - // 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 { - delta * a_param.powf(alpha) - }; - - 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 { @@ -512,21 +377,6 @@ mod tests { } } - // 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; @@ -556,10 +406,23 @@ mod tests { 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); } } } + +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)]) + } +} diff --git a/src/register.rs b/src/register.rs index 2a1b4f3..41b26af 100644 --- a/src/register.rs +++ b/src/register.rs @@ -59,23 +59,41 @@ impl RegistrationStep { /// Default registration steps matching elastix `FixedSmoothingImagePyramid`. /// - /// Smooth-only multi-resolution pyramid (no downsampling). - /// Schedule [8, 4, 2, 1] → σ = [4.0, 2.0, 1.0, 0.5] at spacing=1. - /// Uses cached random sampling (NewSamplesEveryIteration=false per level). + /// Uses a multi-resolution pyramid with downsampling matching SimpleElastix: + /// schedule [8, 4, 2, 1] → σ = [4.0, 2.0, 1.0, 0.5] at spacing=1. + /// 512 iterations per level, 2048–8192 cached random samples. pub fn default_steps(ndim: usize, n: usize) -> Vec { - let sigma_schedule: Vec = vec![4.0, 2.0, 1.0, 0.5]; - let n_samples = (n / 16).clamp(2048, 8192); - sigma_schedule - .iter() - .map(|&s| Self { - sigma: Sigma::Absolute(vec![s; ndim]), - samples: SamplingArg::Random(n_samples), + if ndim == 1 { + return vec![Self { + sigma: Sigma::Absolute(vec![1.0; 1]), + samples: SamplingArg::Random(2048.min(n)), n_bins: 32, tolerance: 1e-6, edge: 0.05, max_iterations: 1500, - learning_rate: 4.0, + learning_rate: 1.0, downsample: 1, + }]; + } + // 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]; + sigma_schedule + .iter() + .zip(downsample_schedule.iter()) + .map(|(&s, &d)| { + let n_pixels = (n / (d * d)).max(4); + let n_samples = n_pixels.min(8192).max(2048); + Self { + sigma: Sigma::Absolute(vec![s; ndim]), + samples: SamplingArg::Random(n_samples), + n_bins: 32, + tolerance: 1e-6, + edge: 0.05, + max_iterations: 512, + learning_rate: 1.0, + downsample: d, + } }) .collect() } diff --git a/src/transform.rs b/src/transform.rs index 2966b25..f8ab58c 100644 --- a/src/transform.rs +++ b/src/transform.rs @@ -276,6 +276,11 @@ impl Transform { } } + pub fn with_shape(mut self, shape: Vec) -> Self { + self.shape = shape; + self + } + /// create a transform from a scaling pub fn from_scaling(scaling: &[f64]) -> Self { let ndim = if let Some(ndim) = D::NDIM { @@ -673,7 +678,6 @@ 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)?; @@ -960,98 +964,6 @@ mod tests { 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(()) - } - - #[test] - fn register2() -> Result<(), Box> { - let t = Transform::from_rotation(f64::PI() / 4.0, &[299.5, 399.5]); - let s = Transform::new(vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0], vec![600, 800]); - let r = t * s; - let p = r.parameters; - println!("{:?}", p); - - let im_a = julia_image( - &[600, 800], - &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0], - &[299.5, 399.5], - &[-0.8, 0.156], - ); - let im_b = julia_image( - &[600, 800], - &[p[0], p[1], p[2], p[3], p[4], p[5]], - &[299.5, 399.5], - &[-0.8, 0.156], - ); - let (t, steps) = Transform::register_debug( - im_a.view(), - im_b.view(), - vec![None, None, None, None, None, None], - None, - None, - )?; - println!("steps:"); - for step in steps { - println!(" {:?}", step); - } - 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(()) - } - fn read_tiff>(path: P) -> Result, Box> { let mut reader = Decoder::new(File::open(path)?)?; reader.seek_to_image(0)?; @@ -1064,7 +976,7 @@ mod tests { DecodingResult::I16(data) => data.into_iter().map(|i| i as f64).collect::>(), DecodingResult::I32(data) => data.into_iter().map(|i| i as f64).collect::>(), DecodingResult::I64(data) => data.into_iter().map(|i| i as f64).collect::>(), - DecodingResult::F16(data) => data.into_iter().map(|i| f64::from(i)).collect::>(), + DecodingResult::F16(data) => data.into_iter().map(f64::from).collect::>(), DecodingResult::F32(data) => data.into_iter().map(|i| i as f64).collect::>(), DecodingResult::F64(data) => data, }; @@ -1110,14 +1022,53 @@ mod tests { .fold(0.0f64, f64::max); println!("Our: {:?} max_err: {:.4} sse: {:.4}", t, max_err, sse); - assert!(max_err < 1.0); - assert!(sse < 1.0); - let mut tif = IJTiffFile::new(std::env::home_dir().unwrap().join("tmp/register_real_images.tif"))?; tif.save(fixed.mapv(|i| i as u16), 0, 0, 0)?; tif.save(t.interpolate_par::<1, _, _>(moving.view())?.mapv(|i| i as u16), 1, 0, 0)?; tif.save(moving.mapv(|i| i as u16), 2, 0, 0)?; + assert!(max_err < 0.1); + assert!(sse < 0.1); + + Ok(()) + } + + #[test] + fn register_real_images2() -> Result<(), Box> { + let fixed = read_tiff("test_files/fixed.tif")?; + let e = Transform::::new(vec![0.8, 0.0, 0.0, 1.0, 0.0, 0.0], fixed.shape().to_vec()).inverse()?; + let moving = e.interpolate::<3, _, _>(fixed.view())?; + + let t = Transform::::register( + fixed.view(), + moving.view(), + vec![None, None, None, None, None, None], + None, + None, + )?; + + let sse = t + .parameters + .iter() + .zip(e.parameters.iter()) + .map(|(a, b)| (a - b).powi(2)) + .sum::(); + let max_err = t + .parameters + .iter() + .zip(e.parameters.iter()) + .map(|(a, b)| (a - b).abs()) + .fold(0.0f64, f64::max); + println!("Our: {:?} max_err: {:.4} sse: {:.4}", t, max_err, sse); + + // let mut tif = IJTiffFile::new(std::env::home_dir().unwrap().join("tmp/register_real_images2.tif"))?; + // tif.save(fixed.mapv(|i| i as u16), 0, 0, 0)?; + // tif.save(t.interpolate_par::<1, _, _>(moving.view())?.mapv(|i| i as u16), 1, 0, 0)?; + // tif.save(moving.mapv(|i| i as u16), 2, 0, 0)?; + + assert!(max_err < 0.1); + assert!(sse < 0.1); + Ok(()) } }