register_real_images test working

Cargo.toml

@@ -31,6 +31,7 @@ tiffwrite = "2026.6.0"
 
 [dev-dependencies]
 tempfile = "3"
+tiff = "0.11"
 
 [profile.release]
 debug = true

src/metric.rs

@@ -167,13 +167,19 @@ impl Sampling {
     fn index(&self) -> Vec<Vec<f64>> {
         match self {
             Self::Fixed(index) => index.clone(),
-            Self::Random((n_samples, shape, _cached)) => rand::rng()
+            Self::Random((n_samples, shape, cached)) => {
+                let mut cache = cached.borrow_mut();
+                if cache.is_empty() {
+                    *cache = rand::rng()
                         .random_iter::<f64>()
                         .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(),
+                        .collect();
+                }
+                cache.clone()
+            }
         }
     }
 }

src/register.rs

@@ -59,35 +59,25 @@ impl RegistrationStep {
 
     /// Default registration steps matching elastix `FixedSmoothingImagePyramid`.
     ///
-    /// Elastix uses `MultiResolutionGaussianSmoothingPyramidImageFilter` which applies
+    /// Smooth-only multi-resolution pyramid (no downsampling).
-    /// Gaussian smoothing with σ = 0.5 × factor at **full resolution** — images are
+    /// Schedule [8, 4, 2, 1] → σ = [4.0, 2.0, 1.0, 0.5] at spacing=1.
-    /// **NOT downsampled**. See: `itkMultiResolutionGaussianSmoothingPyramidImageFilter.hxx`
+    /// Uses cached random sampling (NewSamplesEveryIteration=false per level).
-    ///
-    /// Schedule [8, 4, 2, 1] means σ = [4.0, 2.0, 1.0, 0.5] (spacing=1).
-    /// MaximumNumberOfIterations: 256 (matching SimpleElastix default affine).
     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 n_samples = (n / 16).clamp(2048, 8192);
-        let mut steps = Vec::new();
+        sigma_schedule
-
+            .iter()
-        for i in 0..nlevels {
+            .map(|&s| Self {
-            let fraction = 1.0 / 2.0_f64.powi((nlevels - 1 - i) as i32);
+                sigma: Sigma::Absolute(vec![s; ndim]),
-            let samples = (n as f64 * fraction)
+                samples: SamplingArg::Random(n_samples),
-                .sqrt()
+                n_bins: 32,
-                .max(n as f64 * fraction / 10.0)
+                tolerance: 1e-6,
-                .max(2048.0) as usize;
+                edge: 0.05,
-            steps.push(RegistrationStep::new(
+                max_iterations: 1500,
-                Sigma::Absolute(vec![sigma_schedule[i]; ndim]),
+                learning_rate: 4.0,
-                SamplingArg::Random(samples),
+                downsample: 1,
-                32,
+            })
-                1e-6,
+            .collect()
-                0.05,
-                256,
-                1.0,
-            ));
-        }
-
-        steps
     }
 }
 

src/transform.rs

@@ -198,14 +198,27 @@ impl<D: Dimension> Transform<D> {
 
     /// find the transform which transforms moving into fixed, using fixed_mu to specify which
     /// parameters to keep fixed
-    pub fn register<'a, F, M, T, G>(fixed: F, moving: M, fixed_mu: G) -> Result<Self, Error>
+    pub fn register<'a, F, M, T, G>(
+        fixed: F,
+        moving: M,
+        fixed_mu: G,
+        steps: Option<Vec<RegistrationStep>>,
+        initial_guess: Option<Vec<f64>>,
+    ) -> Result<Self, Error>
     where
         F: AsArray<'a, T, D>,
         M: AsArray<'a, T, D>,
         T: 'a + Clone + AsPrimitive<f64>,
         G: Into<FixedMu>,
     {
-        Registration::new(fixed_mu).register(fixed, moving)
+        let mut registration = Registration::new(fixed_mu);
+        if let Some(steps) = steps {
+            registration.set_steps(steps);
+        }
+        if let Some(initial_guess) = initial_guess {
+            registration.set_initial_guess(initial_guess);
+        }
+        registration.register(fixed, moving)
     }
 
     /// find the transform which transforms moving into fixed and return the results of each
@@ -504,8 +517,13 @@ mod tests {
     use crate::julia_image;
     use crate::transform::Transform;
     use itertools::Itertools;
-    use ndarray::{Ix2, s};
+    use ndarray::{Array2, Ix2, s};
     use num::traits::FloatConst;
+    use std::fs::File;
+    use std::path::Path;
+    use tiff::decoder::{Decoder, DecodingResult};
+    use tiff::tags::Tag;
+    use tiffwrite::IJTiffFile;
 
     #[test]
     fn interpolate() -> Result<(), Box<dyn std::error::Error>> {
@@ -736,39 +754,9 @@ mod tests {
             .inverse()?
             .parameters;
 
-        let all_points: Vec<Vec<f64>> = (0..100).map(|j| vec![j as f64]).collect();
+        // Use default steps with ASGD — the multi-resolution pyramid converges
-        let steps = vec![
+        // reliably for this 1D affine problem.
-            crate::register::RegistrationStep::new(
+        let t = Transform::register(im_a.view(), im_b.view(), vec![None; 2], None, None)?;
-                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),
-                64,
-                1e-6,
-                0.04,
-                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);
@@ -1063,4 +1051,73 @@ mod tests {
         );
         Ok(())
     }
+
+    fn read_tiff<P: AsRef<Path>>(path: P) -> Result<Array2<f64>, Box<dyn std::error::Error>> {
+        let mut reader = Decoder::new(File::open(path)?)?;
+        reader.seek_to_image(0)?;
+        let bytes = match reader.read_image()? {
+            DecodingResult::U8(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
+            DecodingResult::U16(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
+            DecodingResult::U32(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
+            DecodingResult::U64(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
+            DecodingResult::I8(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
+            DecodingResult::I16(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
+            DecodingResult::I32(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
+            DecodingResult::I64(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
+            DecodingResult::F16(data) => data.into_iter().map(|i| f64::from(i)).collect::<Vec<_>>(),
+            DecodingResult::F32(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
+            DecodingResult::F64(data) => data,
+        };
+        let width = reader.get_tag(Tag::ImageWidth)?.into_u32()? as usize;
+        let height = reader.get_tag(Tag::ImageLength)?.into_u32()? as usize;
+        Ok(Array2::from_shape_vec((width, height), bytes)?)
+    }
+
+    #[test]
+    fn register_real_images() -> Result<(), Box<dyn std::error::Error>> {
+        let fixed = read_tiff("test_files/fixed.tif")?;
+        let moving = read_tiff("test_files/moving.tif")?;
+
+        // SimpleElastix expected — the registration result is in TIFF-native coordinates.
+        let expected = [
+            0.9899559376493817,
+            0.011269992506480442,
+            0.017048860489651384,
+            0.8533772512806085,
+            -12.877028775909979,
+            4.143118928117275,
+        ];
+
+        let t = Transform::<Ix2>::register(
+            fixed.view(),
+            moving.view(),
+            vec![None, None, None, None, None, None],
+            None,
+            None,
+        )?;
+
+        let sse = t
+            .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);
+
+        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)?;
+
+        Ok(())
+    }
 }