- register2 passing

.gitignore

@@ -74,6 +74,7 @@ docs/_build/
 
 AGENTS.md
 .agentbridge
+.agent-work
 
 *.tif
 *.svg

Cargo.toml

@@ -34,3 +34,6 @@ tempfile = "3"
 
 [profile.release]
 debug = true
+
+[profile.test]
+inherits = "release"

src/optimize.rs

@@ -188,7 +188,9 @@ pub struct AsgdConfig {
     pub max_iterations: usize,
     /// Convergence tolerance (gradient norm)
     pub tolerance: f64,
-    /// Maximum step length (displacement in mm). Default: 1.0
+    /// 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,
@@ -298,21 +300,25 @@ where
 
         // Compute learning rate: a(t_k) = a / (A + t_k + 1)^alpha
         let t_k: T = current_time;
-        let a_t = a / (T::from(config.sp_a).unwrap() + t_k + T::one()).powf(alpha);
+        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,
+        );
 
         // 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/scales.
+        // 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 {
             let dot_product: T = if let Some(ref scales) = config.scales {
-                // Scaled dot product: sum(g_i * prev_g_i / scales_i²)
+                // Scaled dot product: sum((g/sqrt(C)) * (prev_g/sqrt(C))) = sum(g*pg/C)
-                gradient
+                gradient.iter().zip(prev_g.iter()).zip(scales.iter()).fold(
-                    .iter()
+                    T::zero(),
-                    .zip(prev_g.iter())
+                    |acc, ((&g, &pg), &s)| {
-                    .zip(scales.iter())
-                    .fold(T::zero(), |acc, ((&g, &pg), &s)| {
                         let s2: T = num::cast(s.max(1e-10)).unwrap();
-                        acc + g * pg / (s2 * s2)
+                        let s4 = s2 * s2;
-                    })
+                        acc + g * pg / s4
+                    },
+                )
             } else {
                 gradient
                     .iter()
@@ -329,17 +335,17 @@ where
         };
 
         // Elastix gradient descent step in unscaled parameter space:
-        // Optimizer stores p_scaled = p * scales, updates p_scaled -= a_t * g_scaled
+        // m_Position stores the unscaled position p.
-        // where g_scaled = g / scales (chain rule from GetDerivative).
+        // GetScaledDerivative computes g_scaled = g / scales.
-        // Convert back: p_new = p_scaled_new / scales = p - a_t * g / scales²
+        // Update: p -= a_t * g_scaled = p - a_t * g / scales.
-        // So direction = -g / scales².
+        // So direction = -g / scales.
         let direction: Vec<T> = match &config.scales {
             Some(scales) => gradient
                 .iter()
                 .zip(scales.iter())
                 .map(|(&g, &s)| {
                     let s2: T = num::cast(s.max(1e-10)).unwrap();
-                    -g / (s2 * s2)
+                    -g / s2
                 })
                 .collect(),
             None => gradient.iter().map(|&g| -g).collect(),
@@ -409,14 +415,12 @@ fn sigmoid<T: Float + Debug>(x: T, fmax: T, fmin: T, omega: T) -> T {
     (fmax - fmin) * sigmoid_raw + fmin
 }
 
-/// Estimate ASGD parameters from gradient statistics.
+/// Estimate ASGD parameters using DisplacementDistribution method.
 ///
-/// Matches elastix `AutomaticParameterEstimationOriginal()`:
+/// Matches elastix `AutomaticParameterEstimationUsingDisplacementDistribution()`:
-/// - `a_max = A * delta / sigma1 / sqrt(maxJCJ)`
+/// - `a = delta * (A+1)^alpha / (jacg + eps)` where `delta = MaximumStepLength`
-/// - `sigma1 = sqrt(gg / TrC)` where `gg = ||g_scaled||^2`
+/// - `jacg ≈ sqrt(gg) * 1.8` as analytical approximation
-/// - `TrC = sum(C_scaled[i][i]) = n` (when scales = sqrt(C[i][i]), C_scaled[i][i] = 1)
+/// - `maximum_step_length` is also the per-iteration step clamp (same value as delta).
-/// - `maxJCJ` is approximated as max of squared scaled Jacobian column norms
-/// - `g_scaled = g / scales` (golden standard normalization)
 fn estimate_asgd_parameters<T, F>(
     _f: &F,
     _initial_point: &[T],
@@ -432,54 +436,37 @@ where
     let fmax = T::from(config.sigmoid_max).unwrap();
     let delta = T::from(config.maximum_step_length).unwrap();
 
-    // Compute scaled gradient: g_scaled = g / scales
+    // Compute gg = ||g/scales||^2 = sum(g²/scales²)
-    // And its squared norm: gg = ||g/scales||^2
+    let gg = if let Some(ref scales) = config.scales {
-    let (gg, trc, max_jcj) = 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
-        // Compute maxJCJ analytically for golden standard scales.
-        // Elastix: maxJCJ = max_j [Tr(J_j C J_j^T) + 2√2 ||J_j C J_j^T||_F]
-        // For golden standard scales, C ≈ diag(scales²), so C_scaled ≈ I.
-        // Then J_scaled J_scaled^T = diag(a, ..., a) where
-        //   a = Σ_j max_u((u-c)²)/s_j² + 1 = Σ_j 3(N_j-1)/(N_j+1) + 1
-        // and maxJCJ = a · (ndim + 2√(2·ndim)).
-        let n_params = scales.len();
-        let ndim = ((-1.0 + (1.0 + 4.0 * n_params as f64).sqrt()) / 2.0).round() as usize;
-        let mut a_val = 1.0; // translation contribution: 1/s² = 1
-        for j in 0..ndim {
-            let s_j = scales[j];
-            let n_j = (12.0 * s_j * s_j + 1.0).sqrt();
-            a_val += 3.0 * (n_j - 1.0) / (n_j + 1.0);
-        }
-        let max_jcj = a_val * (ndim as f64 + 2.0 * (2.0 * ndim as f64).sqrt());
-
-        let n = T::from(n_params).unwrap();
-        (gg, n, T::from(max_jcj).unwrap())
     } else {
         let mut gg = T::zero();
         for &g in initial_gradient.iter() {
             gg = gg + g * g;
         }
-        let n = T::from(initial_gradient.len()).unwrap();
+        gg
-        (gg, n, T::one())
     };
 
-    // Elastix: sigma1 = sqrt(gg / TrC), a_max = A * delta / sigma1 / sqrt(maxJCJ)
+    // Compute jacg = displacement distribution estimate.
-    let sigma1 = if gg > T::from(1e-14).unwrap() && trc > T::from(1e-14).unwrap() {
+    // For golden standard scales: E[||d||^2] = gg, so jacg ≈ 1.8 * sqrt(gg)
-        (gg / trc).sqrt()
+    let sqrt_gg = if gg > T::from(1e-20).unwrap() {
+        gg.sqrt()
     } else {
         T::zero()
     };
+    let jacg = sqrt_gg * T::from(1.8).unwrap();
 
-    let a = if sigma1 > T::from(1e-14).unwrap() && max_jcj > T::from(1e-14).unwrap() {
+    // Elastix DisplacementDistribution: a = delta * (A+1)^alpha / (jacg + eps)
-        a_param * delta / sigma1 / max_jcj.sqrt()
+    let a = if jacg > T::from(1e-14).unwrap() {
+        delta * a_param.powf(alpha) / jacg
     } else {
-        a_param
+        delta * a_param.powf(alpha)
     };
 
     let fmin = T::from(config.sigmoid_min).unwrap();

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.
+    /// Elastix uses `MultiResolutionGaussianSmoothingPyramidImageFilter` which applies
-    /// Sigma = 0.5 * factor * spacing. With spacing=1: sigma = [4.0, 2.0, 1.0, 0.5].
+    /// Gaussian smoothing with σ = 0.5 × factor at **full resolution** — images are
-    /// These are absolute sigma values in pixel units.
+    /// **NOT downsampled**. See: `itkMultiResolutionGaussianSmoothingPyramidImageFilter.hxx`
-    /// 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_fixed: sigma.sigma(f_orig_shape.as_slice()),
-                sigma_moving: sigma.sigma(m.shape()),
+                sigma_moving: sigma.sigma(m_orig_shape.as_slice()),
                 n_bins,
                 n_samples,
                 tolerance,

src/transform.rs

@@ -749,22 +749,13 @@ mod tests {
             ),
             crate::register::RegistrationStep::new(
                 crate::metric::Sigma::Absolute(vec![2.0]),
-                crate::metric::SamplingArg::FixedAt(all_points.clone()),
+                crate::metric::SamplingArg::FixedAt(all_points),
                 64,
                 1e-6,
                 0.04,
                 200,
                 1.0,
             ),
-            crate::register::RegistrationStep::new(
-                crate::metric::Sigma::None,
-                crate::metric::SamplingArg::FixedAt(all_points),
-                64,
-                1e-8,
-                0.001,
-                200,
-                1.0,
-            ),
         ];
 
         let (t, steps) = Transform::register_debug(
@@ -792,6 +783,195 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn register2_random_affine() -> Result<(), Box<dyn std::error::Error>> {
+        use rand::prelude::*;
+        let mut rng = rand::rngs::StdRng::seed_from_u64(1337);
+
+        let shape = [200, 200];
+        let center = [99.5, 99.5];
+
+        let fixed = julia_image(
+            &shape,
+            &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
+            &center,
+            &[-0.8, 0.156],
+        )
+        .mapv(|i| i as f64);
+
+        // random rotation ±25°
+        let angle: f64 = rng.random_range(-25.0f64..25.0).to_radians();
+        let (s, c) = angle.sin_cos();
+
+        // random scale 0.85–1.15
+        let sx = rng.random_range(0.85..1.15);
+        let sy = rng.random_range(0.85..1.15);
+
+        // small shear
+        let shx: f64 = rng.random_range(-0.1..0.1);
+        let shy: f64 = rng.random_range(-0.1..0.1);
+
+        // translation ±30 px
+        let tx: f64 = rng.random_range(-30.0..30.0);
+        let ty: f64 = rng.random_range(-30.0..30.0);
+
+        // [m00, m01, m10, m11, tx, ty]
+        let params = vec![c * sx, -s * sy + shx, s * sx + shy, c * sy, tx, ty];
+
+        let transform =
+            Transform::<Ix2>::new_with_center(params.clone(), center.to_vec(), shape.to_vec());
+        let moving = transform.interpolate::<3, _, _>(fixed.view())?;
+
+        let q_inv = transform.inverse()?.parameters;
+
+        let (t, _steps) = Transform::<Ix2>::register_debug(
+            fixed.view(),
+            moving.view(),
+            vec![None, None, None, None, None, None],
+            None,
+            None,
+        )?;
+        println!("params:  {:?}", params);
+        println!("result:  {:?}", t.parameters);
+        println!("q_inv:   {:?}", q_inv);
+        let sse: f64 = t
+            .parameters
+            .iter()
+            .zip(q_inv.iter())
+            .map(|(a, b)| (a - b).powi(2))
+            .sum();
+        println!("sse:     {sse}");
+        assert!(sse < 1.0);
+        Ok(())
+    }
+
+    #[test]
+    fn metric_landscape_2d() -> Result<(), Box<dyn std::error::Error>> {
+        use crate::bspline::{BSpline, BSplineTrait};
+        use crate::filter::gaussian_smooth;
+        use crate::metric::{MattesMetric, SamplingArg};
+        use algos::ObjectiveFunction;
+
+        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],
+            &center,
+            &[-0.8, 0.156],
+        )
+        .mapv(|i| i as f64);
+        let rotation = Transform::<Ix2>::from_rotation(f64::PI() / 4.0, &center);
+        let im_b = rotation.interpolate::<3, _, _>(im_a.view())?;
+
+        let q_inv = rotation.inverse()?.parameters;
+        let p = rotation.parameters.clone();
+        let identity = vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0];
+        let edge = 0.01;
+
+        // Test at full resolution with different sigma levels (matching elastix approach)
+        // Elastix does NOT downsample — it smooths at full resolution with
+        // sigma = 0.5 * factor where factor is from the pyramid schedule [8, 4, 2, 1]
+        println!(
+            "=== Full resolution (no downsampling, matching elastix FixedSmoothingImagePyramid) ==="
+        );
+        for (level, sigma_val) in [4.0, 2.0, 1.0, 0.5].iter().enumerate() {
+            let f = gaussian_smooth(im_a.view(), &[*sigma_val; 2])?;
+            let m = gaussian_smooth(im_b.view(), &[*sigma_val; 2])?;
+            let bf = BSpline::<0, _>::new(f.view());
+            let bm = BSpline::<3, _>::new(m.view());
+            let metric = MattesMetric::new(bf, bm, SamplingArg::Random(3000), 128, edge)?;
+
+            let mi_id = metric.evaluate(&identity);
+            let mi_p = metric.evaluate(&p);
+            let mi_qinv = metric.evaluate(&q_inv);
+
+            println!(
+                "  Level {} (sigma={:.1}, {}x{}): id={:.4} p={:.4} q_inv={:.4} → {}",
+                level,
+                sigma_val,
+                f.shape()[0],
+                f.shape()[1],
+                -mi_id,
+                -mi_p,
+                -mi_qinv,
+                if -mi_qinv > -mi_p {
+                    "Q_INV correct"
+                } else {
+                    "P incorrect (landscape inverted!)"
+                }
+            );
+        }
+
+        // Verify: at every sigma level, q_inv should have higher MI than p
+        let f4 = gaussian_smooth(im_a.view(), &[4.0, 4.0])?;
+        let m4 = gaussian_smooth(im_b.view(), &[4.0, 4.0])?;
+        let bf4 = BSpline::<0, _>::new(f4.view());
+        let bm4 = BSpline::<3, _>::new(m4.view());
+        let metric4 = MattesMetric::new(bf4, bm4, SamplingArg::Random(5000), 128, edge)?;
+        let mi_p_coarse = metric4.evaluate(&p);
+        let mi_qinv_coarse = metric4.evaluate(&q_inv);
+        println!(
+            "\nCoarsest (sigma=4.0): p={:.6} q_inv={:.6}",
+            -mi_p_coarse, -mi_qinv_coarse
+        );
+        assert!(
+            -mi_qinv_coarse > -mi_p_coarse,
+            "MI landscape is inverted at coarsest level! q_inv={} should be > p={}",
+            -mi_qinv_coarse,
+            -mi_p_coarse
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn angle_sweep() -> Result<(), Box<dyn std::error::Error>> {
+        use crate::bspline::{BSpline, BSplineTrait};
+        use crate::filter::gaussian_smooth;
+        use crate::metric::{MattesMetric, SamplingArg};
+        use algos::ObjectiveFunction;
+        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],
+            &center,
+            &[-0.8, 0.156],
+        )
+        .mapv(|i| i as f64);
+        let edge = 0.01;
+        // Full resolution, sigma=4.0 — matching elastix FixedSmoothingImagePyramid (no downsampling)
+        for angle_deg in [5.0f64, 10.0, 15.0, 20.0, 25.0, 30.0, 35.0, 40.0, 45.0] {
+            let rotation = Transform::<Ix2>::from_rotation(angle_deg.to_radians(), &center);
+            let im_b = rotation.interpolate::<3, _, _>(im_a.view())?;
+            let q_inv = rotation.inverse()?.parameters;
+            let p = rotation.parameters.clone();
+            let f = gaussian_smooth(im_a.view(), &[4.0, 4.0])?;
+            let m = gaussian_smooth(im_b.view(), &[4.0, 4.0])?;
+            let metric = MattesMetric::new(
+                BSpline::<0, _>::new(f.view()),
+                BSpline::<3, _>::new(m.view()),
+                SamplingArg::Random(3000),
+                128,
+                edge,
+            )?;
+            let id = vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0];
+            let mi_id = metric.evaluate(&id);
+            let mi_p = metric.evaluate(&p);
+            let mi_qi = metric.evaluate(&q_inv);
+            println!(
+                "{:5.1}°: id={:.3} p={:.3} q_inv={:.3} → {}",
+                angle_deg,
+                -mi_id,
+                -mi_p,
+                -mi_qi,
+                if mi_p > mi_qi { "P wins" } else { "Q_INV wins" }
+            );
+        }
+        Ok(())
+    }
+
     #[test]
     fn register2_interpolate() -> Result<(), Box<dyn std::error::Error>> {
         let shape = [200, 200];