- registration getting better

src/metric.rs

@@ -281,8 +281,8 @@ where
         T: 'a + Clone + AsPrimitive<f64>,
     {
         Self::new(
-            BSpline::new(fixed),
+            BSpline::<0, _>::new(fixed),
-            BSpline::new(moving),
+            BSpline::<3, _>::new(moving),
             sampling,
             n_bins,
             edge,

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:
+    /// Uses a multi-resolution approach with L-BFGS optimizer and all pixels on all levels
-    /// schedule [8, 4, 2, 1] → σ = [4.0, 2.0, 1.0, 0.5] at spacing=1.
+    /// for deterministic, precise convergence.
-    /// Uses all pixels at 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 {
+            let bf = BSpline::<0, _>::new(f.view());
-                Optimizer::LBFGS => {
+            let bm = BSpline::<3, _>::new(m.view());
-                    let config = OptimizationConfig {
+            let metric = MattesMetric::new(bf, bm, samples, n_bins, edge)?
-                        max_iterations,
+                .with_fixed_mu(self.fixed_mu.clone());
-                        tolerance,
+            let p_var = metric.fixed_mu().extract_variable(&p);
-                        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,
-                    )
-                }
-            };
 
-            if optimization_result
+            let optimization_result = Self::optimize_metric(
-                .optimal_point
+                &metric,
-                .iter()
+                &p_var,
-                .all(|i| i.is_finite())
+                &optimizer,
-            {
+                max_iterations,
-                p = metric
+                tolerance,
-                    .fixed_mu()
+                learning_rate,
-                    .combine(&optimization_result.optimal_point);
+                &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,

src/transform.rs

@@ -1008,20 +1008,24 @@ mod tests {
             None,
         )?;
 
-        let sse = t
+        let expected_transform = Transform::<Ix2>::new(expected.to_vec(), fixed.shape().to_vec());
-            .parameters
-            .iter()
-            .zip(expected.iter())
-            .map(|(a, b)| (a - b).powi(2))
-            .sum::<f64>();
-        let max_err = t
-            .parameters
-            .iter()
-            .zip(expected.iter())
-            .map(|(a, b)| (a - b).abs())
-            .fold(0.0f64, f64::max);
-        println!("Our: {:?}  max_err: {:.4} sse: {:.4}", t, max_err, sse);
 
+        // Compare transformed coordinates of all pixels
+        let shape = fixed.shape();
+        let mut sum_diff = 0.0;
+        let mut count = 0;
+        for row in 0..shape[0] {
+            for col in 0..shape[1] {
+                let point = [row as f64, col as f64];
+                let t_point = t.transform_point(&point);
+                let e_point = expected_transform.transform_point(&point);
+                let diff_sq = (t_point[0] - e_point[0]).powi(2) + (t_point[1] - e_point[1]).powi(2);
+                sum_diff += diff_sq.sqrt();
+                count += 1;
+            }
+        }
+        let mean_diff = sum_diff / count as f64;
+        println!("Our: {:?}  mean_coord_diff: {:.6}", t, mean_diff);
 
         let mut tif = IJTiffFile::new(
             std::env::home_dir()
@@ -1038,8 +1042,7 @@ mod tests {
         )?;
         tif.save(moving.mapv(|i| i as u16), 2, 0, 0)?;
 
-        assert!(max_err < 0.02);
+        assert!(mean_diff < 0.1);
-        assert!(sse < 0.02);
 
         Ok(())
     }
@@ -1060,27 +1063,40 @@ mod tests {
         )?;
 
         let e_inv = e.inverse()?;
-        let sse = t
-            .parameters
-            .iter()
-            .zip(e_inv.parameters.iter())
-            .map(|(a, b)| (a - b).powi(2))
-            .sum::<f64>();
-        let max_err = t
-            .parameters
-            .iter()
-            .zip(e_inv.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"))?;
+        // Compare transformed coordinates of all pixels
+        let shape = fixed.shape();
+        let mut sum_diff = 0.0;
+        let mut count = 0;
+        for row in 0..shape[0] {
+            for col in 0..shape[1] {
+                let point = [row as f64, col as f64];
+                let t_point = t.transform_point(&point);
+                let e_point = e_inv.transform_point(&point);
+                let diff_sq = (t_point[0] - e_point[0]).powi(2) + (t_point[1] - e_point[1]).powi(2);
+                sum_diff += diff_sq.sqrt();
+                count += 1;
+            }
+        }
+        let mean_diff = sum_diff / count as f64;
+        println!("Our: {:?}  mean_coord_diff: {:.6}", t, mean_diff);
+
+        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(
+            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.02);
+        assert!(mean_diff < 0.1);
-        assert!(sse < 0.02);
 
         Ok(())
     }