- registration getting better

src/register.rs

@@ -3,7 +3,7 @@ use crate::error::Error;
 use crate::metric::{FixedMu, MattesMetric, SamplingArg, Sigma};
 use crate::optimize::{AsgdConfig, asgd_minimize, lbfgs_minimize};
 use crate::transform::Transform;
-use algos::OptimizationConfig;
+use algos::{ObjectiveFunction, OptimizationConfig};
 use ndarray::{AsArray, Dimension};
 use num::cast::AsPrimitive;
 use std::marker::PhantomData;
@@ -33,6 +33,8 @@ pub struct RegistrationStep {
     pub max_iterations: usize,
     pub learning_rate: f64,
     pub downsample: usize,
+    /// Optimizer to use for this level
+    pub optimizer: Optimizer,
 }
 
 impl RegistrationStep {
@@ -54,14 +56,14 @@ impl RegistrationStep {
             max_iterations,
             learning_rate,
             downsample: 1,
+            optimizer: Optimizer::default(),
         }
     }
 
     /// Default registration steps matching elastix `FixedSmoothingImagePyramid`.
     ///
-    /// 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.
-    /// Uses all pixels at all levels for deterministic, precise convergence.
+    /// Uses a multi-resolution approach with L-BFGS optimizer and all pixels on all levels
+    /// for deterministic, precise convergence.
     pub fn default_steps(ndim: usize, n: usize) -> Vec<Self> {
         if ndim == 1 {
             return vec![Self {
@@ -73,6 +75,7 @@ impl RegistrationStep {
                 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]
@@ -86,7 +89,7 @@ impl RegistrationStep {
             .map(|(level, (&s, &d))| {
                 let n_pixels = (n / (d * d)).max(4);
                 let is_finest = level == n_levels - 1;
-                // Use all pixels at all levels for deterministic results
+                // Use all pixels on all levels for deterministic results
                 let (max_iter, tol) = if is_finest { (2048, 1e-8) } else { (512, 1e-6) };
                 Self {
                     sigma: Sigma::Absolute(vec![s; ndim]),
@@ -97,6 +100,7 @@ impl RegistrationStep {
                     max_iterations: max_iter,
                     learning_rate: 1.0,
                     downsample: d,
+                    optimizer: Optimizer::LBFGS,
                 }
             })
             .collect()
@@ -313,6 +317,7 @@ impl<D: Dimension> Registration<D> {
             max_iterations,
             learning_rate,
             downsample: _,
+            optimizer,
         } in steps.into_iter()
         {
             let f = sigma.smooth(fixed.view())?;
@@ -321,58 +326,68 @@ impl<D: Dimension> Registration<D> {
             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 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 {
-                        max_iterations,
-                        tolerance,
-                        learning_rate,
-                    };
-                    lbfgs_minimize(
-                        &metric,
-                        metric.fixed_mu().extract_variable(&p).as_slice(),
-                        &config,
-                    )
-                }
-                Optimizer::ASGD => {
-                    let config = AsgdConfig {
-                        max_iterations,
-                        tolerance,
-                        maximum_step_length: 1.0,
-                        sp_a: 20.0,
-                        sp_alpha: 0.602,
-                        scales: Some(scales),
-                        ..Default::default()
-                    };
-                    asgd_minimize(
-                        &metric,
-                        metric.fixed_mu().extract_variable(&p).as_slice(),
-                        &config,
-                    )
-                }
-            };
+            let bf = BSpline::<0, _>::new(f.view());
+            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 p_var = metric.fixed_mu().extract_variable(&p);
 
-            if optimization_result
-                .optimal_point
-                .iter()
-                .all(|i| i.is_finite())
-            {
-                p = metric
-                    .fixed_mu()
-                    .combine(&optimization_result.optimal_point);
+            let optimization_result = Self::optimize_metric(
+                &metric,
+                &p_var,
+                &optimizer,
+                max_iterations,
+                tolerance,
+                learning_rate,
+                &scales,
+            );
+
+            if let Some(result) = optimization_result {
+                if result.optimal_point.iter().all(|i| i.is_finite()) {
+                    p = self.fixed_mu.combine(&result.optimal_point);
+                }
             }
         }
         Ok(Transform::<D>::new(p, fixed.shape().to_vec()))
     }
 
+    fn optimize_metric<M: ObjectiveFunction<f64>>(
+        metric: &M,
+        p_var: &[f64],
+        optimizer: &Optimizer,
+        max_iterations: usize,
+        tolerance: f64,
+        learning_rate: f64,
+        scales: &[f64],
+    ) -> Option<algos::OptimizationResult<f64>> {
+        match optimizer {
+            Optimizer::LBFGS => {
+                let config = OptimizationConfig {
+                    max_iterations,
+                    tolerance,
+                    learning_rate,
+                };
+                Some(lbfgs_minimize(metric, p_var, &config))
+            }
+            Optimizer::ASGD => {
+                let config = AsgdConfig {
+                    max_iterations,
+                    tolerance,
+                    maximum_step_length: 1.0,
+                    sp_a: 20.0,
+                    sp_alpha: 0.602,
+                    scales: Some(scales.to_vec()),
+                    ..Default::default()
+                };
+                Some(asgd_minimize(metric, p_var, &config))
+            }
+        }
+    }
+
     /// find the transform which transforms moving into fixed and return the results of each
     /// optimization step
     pub fn register_debug<'a, F, M, T>(
@@ -409,6 +424,7 @@ impl<D: Dimension> Registration<D> {
             max_iterations,
             learning_rate,
             downsample: _,
+            optimizer,
         } in steps.into_iter()
         {
             let f = sigma.smooth(fixed.view())?;
@@ -434,7 +450,7 @@ impl<D: Dimension> Registration<D> {
             // Golden standard scales from downsampled shape
             let scales = Self::golden_standard_scales(f.shape(), ndim);
 
-            let optimization_result = match &self.optimizer {
+            let optimization_result = match &optimizer {
                 Optimizer::LBFGS => {
                     let config = OptimizationConfig {
                         max_iterations,