- register2 passing

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<Self> {
         let sigma_schedule: Vec<f64> = 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<D>(array: ndarray::ArrayView<f64, D>, factor: usize) -> ndarray::Array<f64, D>
+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<usize> = 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::<D>()
+        .unwrap()
+}
+
 #[derive(Clone, Debug)]
 pub struct RegistrationResult {
     pub sigma_fixed: Option<Vec<f64>>,
@@ -255,10 +299,12 @@ impl<D: Dimension> Registration<D> {
             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<D: Dimension> Registration<D> {
             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<D: Dimension> Registration<D> {
                     )
                 }
                 Optimizer::ASGD => {
-                    let scales = Self::golden_standard_scales(fixed.shape(), ndim);
                     let config = AsgdConfig {
                         max_iterations,
                         tolerance,
@@ -347,13 +395,19 @@ impl<D: Dimension> Registration<D> {
             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<usize> = f.shape().to_vec();
+            let m_orig_shape: Vec<usize> = 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<D: Dimension> Registration<D> {
             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<D: Dimension> Registration<D> {
                     )
                 }
                 Optimizer::ASGD => {
-                    let scales = Self::golden_standard_scales(fixed.shape(), ndim);
                     let config = AsgdConfig {
                         max_iterations,
                         tolerance,
@@ -405,9 +461,10 @@ impl<D: Dimension> Registration<D> {
                     .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,