- tests register1 and register2 passing

src/lib.rs

@@ -3,10 +3,10 @@ pub mod bspline;
 pub mod error;
 pub mod filter;
 pub mod metric;
+pub mod optimize;
 pub mod par_indexed_iter;
 pub mod register;
 pub mod transform;
-mod optimize;
 
 use ndarray::prelude::*;
 use num::Complex;

src/metric.rs

@@ -152,7 +152,7 @@ fn parzen(
 
 enum Sampling {
     Fixed(Vec<Vec<f64>>),
-    Random((usize, Vec<f64>)),
+    Random((usize, Vec<f64>, RefCell<Vec<Vec<f64>>>)),
 }
 
 impl Sampling {
@@ -161,19 +161,21 @@ impl Sampling {
     }
 
     fn random(n_samples: usize, shape: Vec<f64>) -> Self {
-        Self::Random((n_samples, shape))
+        Self::Random((n_samples, shape, RefCell::new(Vec::new())))
     }
 
     fn index(&self) -> Vec<Vec<f64>> {
         match self {
             Self::Fixed(index) => index.clone(),
-            Self::Random((n_samples, shape)) => rand::rng()
+            Self::Random((n_samples, shape, _cached)) => {
-                .random_iter::<f64>()
+                rand::rng()
-                .take(n_samples * shape.len())
+                    .random_iter::<f64>()
-                .chunks(shape.len())
+                    .take(n_samples * shape.len())
-                .into_iter()
+                    .chunks(shape.len())
-                .map(|c| c.zip_eq(shape.iter()).map(|(i, s)| i * s - 0.5).collect())
+                    .into_iter()
-                .collect(),
+                    .map(|c| c.zip_eq(shape.iter()).map(|(i, s)| i * s - 0.5).collect())
+                    .collect()
+            }
         }
     }
 }
@@ -240,6 +242,7 @@ where
             .into_option()
             .expect("m and f cannot be empty");
         let minmax = [*min, *max];
+        let ndim = shape.len();
         let sampling = match sampling {
             SamplingArg::Fixed(n_samples) => {
                 Sampling::fixed(Sampling::random(n_samples, shape.clone()).index())
@@ -253,7 +256,7 @@ where
             center,
             fixed,
             moving,
-            fixed_mu: FixedMu::new_none(sampling.index().len()),
+            fixed_mu: FixedMu::new_none(ndim),
             minmax,
             sampling,
             n_bins,
@@ -319,7 +322,7 @@ where
                             let f = fixed.evaluate_at_continuous_index(&u, mem_f)?;
                             let (m, d) = moving
                                 .evaluate_value_and_derivative_at_continuous_index(&v, mem_m)?;
-                            let j = fixed_mu.extract_variable(&image_jacobian(&v, center, &d));
+                            let j = fixed_mu.extract_variable(&image_jacobian(&u, center, &d));
                             let mut k;
                             let mut w = 1.0;
                             let mut ws = Vec::with_capacity(shape.len());
@@ -347,7 +350,7 @@ where
                             for (dwi, wi) in dw.iter_mut().zip(ws) {
                                 *dwi *= w / wi;
                             }
-                            let dw = fixed_mu.extract_variable(&image_jacobian(&v, center, &dw));
+                            let dw = fixed_mu.extract_variable(&image_jacobian(&u, center, &dw));
                             Ok(Some((f, ((m, j), (w, dw)))))
                         } else {
                             Ok(None)
@@ -435,12 +438,8 @@ where
 
             self.metric.replace(IntMut {
                 mu: mu.to_vec(),
-                metric: metric * alpha,
+                metric: metric,
-                derivative: dmetric
+                derivative: dmetric.clone(),
-                    .iter()
-                    .zip_eq(dalpha)
-                    .map(|(dm, da)| *dm * alpha + da * metric)
-                    .collect(),
             });
             // self.metric.replace(IntMut {
             //     mu: mu.to_vec(),

src/optimize.rs

@@ -1,8 +1,9 @@
-use std::collections::VecDeque;
-use std::fmt::Debug;
 use algos::{ObjectiveFunction, OptimizationConfig, OptimizationResult};
 use num::Float;
+use std::collections::VecDeque;
+use std::fmt::Debug;
 
+/// L-BFGS optimizer with proper line search and initial Hessian scaling
 pub fn lbfgs_minimize<T, F>(
     f: &F,
     initial_point: &[T],
@@ -55,9 +56,9 @@ where
         for i in (0..s_list.len()).rev() {
             let alpha = rho_list[i]
                 * s_list[i]
-                .iter()
+                    .iter()
-                .zip(q.iter())
+                    .zip(q.iter())
-                .fold(T::zero(), |acc, (&s, &q)| acc + s * q);
+                    .fold(T::zero(), |acc, (&s, &q)| acc + s * q);
             alpha_list.push(alpha);
             for (q_j, y_j) in q.iter_mut().zip(y_list[i].iter()) {
                 *q_j = *q_j - alpha * *y_j;
@@ -84,9 +85,9 @@ where
         for i in 0..s_list.len() {
             let beta = rho_list[i]
                 * y_list[i]
-                .iter()
+                    .iter()
-                .zip(r.iter())
+                    .zip(r.iter())
-                .fold(T::zero(), |acc, (&y, &r)| acc + y * r);
+                    .fold(T::zero(), |acc, (&y, &r)| acc + y * r);
             let alpha = alpha_list[s_list.len() - 1 - i];
             for (r_j, s_j) in r.iter_mut().zip(s_list[i].iter()) {
                 *r_j = *r_j + (alpha - beta) * *s_j;
@@ -96,23 +97,36 @@ where
         // r now contains the search direction
         let direction: Vec<T> = r.iter().map(|&x| -x).collect();
 
-        // Line search to find step size
+        // Line search to find step size with Armijo sufficient decrease condition
         let mut alpha = T::one();
         let mut new_point = vec![T::zero(); n];
         let current_value = f.evaluate(&current_point);
+        let mut improved = false;
+        // directional derivative g^T * d (should be negative for descent)
+        let g_dot_d: T = gradient
+            .iter()
+            .zip(direction.iter())
+            .fold(T::zero(), |acc, (&g, &d)| acc + g * d);
+        let c1 = T::from(1e-4).unwrap(); // Armijo constant
 
-        // Simple backtracking line search
+        // Backtracking line search with Armijo condition
-        for _ in 0..20 {
+        for _ in 0..30 {
             for i in 0..n {
                 new_point[i] = current_point[i] + alpha * direction[i];
             }
             let new_value = f.evaluate(&new_point);
-            if new_value < current_value {
+            // Armijo: f(x + αd) ≤ f(x) + c₁ α (gᵀd)
+            if new_value <= current_value + c1 * alpha * g_dot_d {
+                improved = true;
                 break;
             }
             alpha = alpha * T::from(0.5).unwrap();
         }
 
+        // FIX: If no improvement found, don't update (would corrupt L-BFGS history)
+        if !improved {
+            break;
+        }
 
         // Get new gradient
         let new_gradient = match f.gradient(&new_point) {
@@ -137,7 +151,8 @@ where
             .zip(s.iter())
             .fold(T::zero(), |acc, (&y_i, &s_i)| acc + y_i * s_i);
 
-        if ys == T::zero() {
+        // FIX: Check for non-positive ys (corrupts Hessian approximation)
+        if ys <= T::zero() {
             break;
         }
 
@@ -151,7 +166,7 @@ where
         s_list.push_back(s);
         y_list.push_back(y);
         rho_list.push_back(rho);
-        
+
         // Update for next iteration
         current_point = new_point;
         gradient = new_gradient;
@@ -164,4 +179,400 @@ where
         iterations,
         converged,
     }
-}
+}
+
+/// Configuration for Adaptive Stochastic Gradient Descent (ASGD) optimizer
+#[derive(Clone, Debug)]
+pub struct AsgdConfig {
+    /// Maximum number of iterations per resolution
+    pub max_iterations: usize,
+    /// Convergence tolerance (gradient norm)
+    pub tolerance: f64,
+    /// Maximum step length (displacement in mm). Default: 1.0
+    pub maximum_step_length: f64,
+    /// Gain parameter A (denominator offset). Default: 20.0
+    pub sp_a: f64,
+    /// Learning rate parameter alpha. Default: 1.0
+    pub sp_alpha: f64,
+    /// Sigmoid maximum. Default: 1.0
+    pub sigmoid_max: f64,
+    /// Sigmoid minimum. Elastix: -0.99 + 0.98*noisefactor. For low noise: -0.01
+    pub sigmoid_min: f64,
+    /// Sigmoid scale (omega). Elastix: SigmoidScaleFactor * sigma3^2 * sqrt(TrCC).
+    /// For low noise (exact gradients), omega ≈ 0 (step function sigmoid).
+    pub sigmoid_scale: f64,
+    /// Per-parameter scales for gradient normalization (golden standard).
+    /// scales[i] = C[i][i] = E[|J_i|^2], the diagonal of Jacobian covariance.
+    pub scales: Option<Vec<f64>>,
+}
+
+impl Default for AsgdConfig {
+    fn default() -> Self {
+        Self {
+            max_iterations: 250,
+            tolerance: 1e-6,
+            maximum_step_length: 1.0,
+            sp_a: 20.0,
+            sp_alpha: 1.0,
+            sigmoid_max: 1.0,
+            sigmoid_min: -0.01,
+            sigmoid_scale: 1e-8,
+            scales: None,
+        }
+    }
+}
+
+/// Adaptive Stochastic Gradient Descent (ASGD) optimizer
+///
+/// This implements the ASGD optimizer from elastix, which is well-suited for
+/// stochastic gradient estimates from Mattes Mutual Information metric.
+///
+/// The learning rate at iteration k is:
+///   a(t_k) = a / (A + t_k + 1)^alpha
+///
+/// Where time t_k is updated via sigmoid:
+///   t_{k+1} = max(0, t_k + sigmoid(-g_k^T * g_{k-1}))
+pub fn asgd_minimize<T, F>(f: &F, initial_point: &[T], config: &AsgdConfig) -> OptimizationResult<T>
+where
+    T: Float + Debug,
+    F: ObjectiveFunction<T>,
+{
+    let n = initial_point.len();
+    let mut current_point = initial_point.to_vec();
+    let mut iterations = 0;
+    let mut converged = false;
+
+    // Get initial gradient
+    let mut gradient = match f.gradient(&current_point) {
+        Some(g) => g,
+        None => {
+            return OptimizationResult {
+                optimal_point: current_point.clone(),
+                optimal_value: f.evaluate(&current_point),
+                iterations: 0,
+                converged: false,
+            };
+        }
+    };
+
+    eprintln!(
+        "asgd_minimize: initial_point={:?} gradient={:?}",
+        initial_point
+            .iter()
+            .map(|x| num::cast::<_, f64>(*x).unwrap())
+            .collect::<Vec<_>>(),
+        gradient
+            .iter()
+            .map(|x| num::cast::<_, f64>(*x).unwrap())
+            .collect::<Vec<_>>(),
+    );
+
+    // Estimate ASGD parameters from gradient statistics
+    let (a, alpha, fmax, fmin, omega) =
+        estimate_asgd_parameters(f, &current_point, &gradient, config);
+
+    eprintln!(
+        "asgd_minimize: a={:.6} scales={:?}",
+        num::cast::<_, f64>(a).unwrap(),
+        config.scales.as_ref().map(|s| s
+            .iter()
+            .map(|x| num::cast::<_, f64>(*x).unwrap())
+            .collect::<Vec<_>>()),
+    );
+
+    // Initial time
+    let mut current_time = T::zero();
+    let mut previous_gradient: Option<Vec<T>> = None;
+
+    while iterations < config.max_iterations {
+        // Check for convergence
+        let gradient_norm = gradient
+            .iter()
+            .fold(T::zero(), |acc, &x| acc + x * x)
+            .sqrt();
+        let tol: T = num::cast(config.tolerance).unwrap();
+        if gradient_norm < tol {
+            converged = true;
+            break;
+        }
+
+        // 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);
+
+        // 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.
+        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²)
+                gradient
+                    .iter()
+                    .zip(prev_g.iter())
+                    .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)
+                    })
+            } else {
+                gradient
+                    .iter()
+                    .zip(prev_g.iter())
+                    .fold(T::zero(), |acc, (&g, &pg)| acc + g * pg)
+            };
+
+            // Elastix: sigmoid(-inprod) — time increases when gradients disagree
+            let sigmoid_val = sigmoid(-dot_product, fmax, fmin, omega);
+            current_time = T::max(T::zero(), current_time + sigmoid_val);
+            num::cast(dot_product).unwrap()
+        } else {
+            0.0
+        };
+
+        // Elastix gradient descent step in unscaled parameter space:
+        // Optimizer stores p_scaled = p * scales, updates p_scaled -= a_t * g_scaled
+        // where g_scaled = g / scales (chain rule from GetDerivative).
+        // Convert back: p_new = p_scaled_new / scales = p - a_t * 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)
+                })
+                .collect(),
+            None => gradient.iter().map(|&g| -g).collect(),
+        };
+
+        let mut new_point = current_point.clone();
+        for i in 0..n {
+            new_point[i] = current_point[i] + direction[i] * a_t;
+        }
+
+        if iterations < 5 || iterations % 20 == 0 {
+            eprintln!(
+                "iter={}: point={:?} grad={:?} dir={:?} a_t={:.4} t={:.4} dot={:.6}",
+                iterations,
+                current_point
+                    .iter()
+                    .map(|x| num::cast::<_, f64>(*x).unwrap())
+                    .collect::<Vec<_>>(),
+                gradient
+                    .iter()
+                    .map(|x| num::cast::<_, f64>(*x).unwrap())
+                    .collect::<Vec<_>>(),
+                direction
+                    .iter()
+                    .map(|x| num::cast::<_, f64>(*x).unwrap())
+                    .collect::<Vec<_>>(),
+                num::cast::<_, f64>(a_t).unwrap(),
+                num::cast::<_, f64>(current_time).unwrap(),
+                dot_product_val,
+            );
+        }
+
+        // Get new gradient
+        let new_gradient = match f.gradient(&new_point) {
+            Some(g) => g,
+            None => break,
+        };
+
+        // Update for next iteration
+        previous_gradient = Some(gradient.clone());
+        current_point = new_point;
+        gradient = new_gradient;
+        iterations += 1;
+    }
+
+    OptimizationResult {
+        optimal_point: current_point.clone(),
+        optimal_value: f.evaluate(&current_point),
+        iterations,
+        converged,
+    }
+}
+
+/// Sigmoid function for ASGD time update.
+/// Matches ITK SigmoidImageFilter: f(x) = (Max-Min) * 1/(1+exp(-(x-beta)/alpha)) + Min
+/// where alpha = omega, beta = omega * ln(-fmax/fmin), so sigmoid(0) = 0.
+fn sigmoid<T: Float + Debug>(x: T, fmax: T, fmin: T, omega: T) -> T {
+    let beta = omega * (-fmax / fmin).ln();
+    let z = (x - beta) / omega;
+    let sigmoid_raw = if z > T::from(20.0).unwrap() {
+        T::one()
+    } else if z < T::from(-20.0).unwrap() {
+        T::zero()
+    } else {
+        T::one() / (T::one() + (-z).exp())
+    };
+    (fmax - fmin) * sigmoid_raw + fmin
+}
+
+/// Estimate ASGD parameters from gradient statistics.
+///
+/// Matches elastix `AutomaticParameterEstimationOriginal()`:
+/// - `a_max = A * delta / sigma1 / sqrt(maxJCJ)`
+/// - `sigma1 = sqrt(gg / TrC)` where `gg = ||g_scaled||^2`
+/// - `TrC = sum(C_scaled[i][i]) = n` (when scales = sqrt(C[i][i]), C_scaled[i][i] = 1)
+/// - `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],
+    initial_gradient: &[T],
+    config: &AsgdConfig,
+) -> (T, T, T, T, T)
+where
+    T: Float + Debug,
+    F: ObjectiveFunction<T>,
+{
+    let a_param = T::from(config.sp_a).unwrap();
+    let alpha = T::from(config.sp_alpha).unwrap();
+    let fmax = T::from(config.sigmoid_max).unwrap();
+    let delta = T::from(config.maximum_step_length).unwrap();
+
+    // Compute scaled gradient: g_scaled = g / scales
+    // And its squared norm: gg = ||g/scales||^2
+    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;
+        }
+
+        // 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, n, T::one())
+    };
+
+    // Elastix: sigma1 = sqrt(gg / TrC), a_max = A * delta / sigma1 / sqrt(maxJCJ)
+    let sigma1 = if gg > T::from(1e-14).unwrap() && trc > T::from(1e-14).unwrap() {
+        (gg / trc).sqrt()
+    } else {
+        T::zero()
+    };
+
+    let a = if sigma1 > T::from(1e-14).unwrap() && max_jcj > T::from(1e-14).unwrap() {
+        a_param * delta / sigma1 / max_jcj.sqrt()
+    } else {
+        a_param
+    };
+
+    let fmin = T::from(config.sigmoid_min).unwrap();
+    let omega = T::from(config.sigmoid_scale).unwrap();
+
+    (a, alpha, fmax, fmin, omega)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use algos::OptimizationConfig;
+
+    // Test function: f(x, y) = x^2 + y^2
+    struct Quadratic;
+
+    impl ObjectiveFunction<f64> for Quadratic {
+        fn evaluate(&self, point: &[f64]) -> f64 {
+            point.iter().map(|x| x * x).sum()
+        }
+
+        fn gradient(&self, point: &[f64]) -> Option<Vec<f64>> {
+            Some(point.iter().map(|x| 2.0 * x).collect())
+        }
+    }
+
+    #[test]
+    fn test_lbfgs_quadratic() {
+        let f = Quadratic;
+        let initial_point = vec![1.0, 1.0];
+        let config = OptimizationConfig {
+            max_iterations: 100,
+            tolerance: 1e-6,
+            learning_rate: 1.0,
+        };
+
+        let result = lbfgs_minimize(&f, &initial_point, &config);
+
+        assert!(result.converged);
+        assert!(result.optimal_value < 1e-10);
+        for x in result.optimal_point {
+            assert!(x.abs() < 1e-5);
+        }
+    }
+
+    // Test function: f(x) = (x - 2)^2
+    struct QuadraticWithMinimum;
+
+    impl ObjectiveFunction<f64> for QuadraticWithMinimum {
+        fn evaluate(&self, point: &[f64]) -> f64 {
+            let x = point[0];
+            (x - 2.0).powi(2)
+        }
+
+        fn gradient(&self, point: &[f64]) -> Option<Vec<f64>> {
+            let x = point[0];
+            Some(vec![2.0 * (x - 2.0)])
+        }
+    }
+
+    #[test]
+    fn test_lbfgs_quadratic_with_minimum() {
+        let f = QuadraticWithMinimum;
+        let initial_point = vec![0.0];
+        let config = OptimizationConfig {
+            max_iterations: 100,
+            tolerance: 1e-6,
+            learning_rate: 1.0,
+        };
+
+        let result = lbfgs_minimize(&f, &initial_point, &config);
+
+        assert!(result.converged);
+        assert!((result.optimal_point[0] - 2.0).abs() < 1e-5);
+    }
+
+    #[test]
+    fn test_asgd_quadratic() {
+        let f = Quadratic;
+        let initial_point = vec![1.0, 1.0];
+        let config = AsgdConfig {
+            max_iterations: 1000,
+            tolerance: 1e-4,
+            sigmoid_min: -0.8,
+            ..Default::default()
+        };
+
+        let result = asgd_minimize(&f, &initial_point, &config);
+
+        // ASGD should converge close to minimum
+        assert!(result.optimal_value < 0.1);
+        for x in result.optimal_point {
+            assert!(x.abs() < 0.5);
+        }
+    }
+}

src/register.rs

@@ -1,14 +1,28 @@
 use crate::bspline::{BSpline, BSplineTrait};
 use crate::error::Error;
 use crate::metric::{FixedMu, MattesMetric, SamplingArg, Sigma};
+use crate::optimize::{AsgdConfig, asgd_minimize, lbfgs_minimize};
 use crate::transform::Transform;
-use crate::optimize::lbfgs_minimize;
 use algos::OptimizationConfig;
 use ndarray::{AsArray, Dimension};
 use num::cast::AsPrimitive;
-use num::integer::Roots;
 use std::marker::PhantomData;
 
+/// Optimizer type for registration
+#[derive(Clone, Debug, PartialEq)]
+pub enum Optimizer {
+    /// L-BFGS optimizer (fast, requires consistent gradients)
+    LBFGS,
+    /// Adaptive Stochastic Gradient Descent (robust, handles noisy gradients)
+    ASGD,
+}
+
+impl Default for Optimizer {
+    fn default() -> Self {
+        Self::ASGD
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct RegistrationStep {
     pub sigma: Sigma,
@@ -41,54 +55,38 @@ impl RegistrationStep {
         }
     }
 
+    /// Default registration steps matching SimpleElastix affine defaults.
+    ///
+    /// 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 defaults: 2048 samples, 32 histogram bins, 256 iterations.
     pub fn default_steps(ndim: usize, n: usize) -> Vec<Self> {
-        vec![
+        let sigma_schedule: Vec<f64> = vec![4.0, 2.0, 1.0, 0.5];
-            RegistrationStep::new(
+        let nlevels = sigma_schedule.len();
-                Sigma::Relative(vec![0.5; ndim]),
+
-                SamplingArg::Fixed(n.sqrt().max(n / 125).max(100)),
+        let mut steps = Vec::new();
-                3,
+
-                1e-4,
+        for (i, &sigma) in sigma_schedule.iter().enumerate() {
-                0.05,
+            let fraction = 1.0 / 2.0_f64.powi((nlevels - 1 - i) as i32);
-                100,
+            let samples = (n as f64 * fraction)
-                100.0,
+                .sqrt()
-            ),
+                .max(n as f64 * fraction / 10.0)
-            RegistrationStep::new(
+                .max(2048.0) as usize;
-                Sigma::Relative(vec![0.25; ndim]),
+            steps.push(RegistrationStep::new(
-                SamplingArg::Fixed(n.sqrt().max(n / 30).max(200)),
+                Sigma::Absolute(vec![sigma; ndim]),
-                6,
+                SamplingArg::Random(samples),
+                32,
                 1e-6,
-                0.04,
+                0.05,
-                100,
+                256,
                 1.0,
-            ),
+            ));
-            RegistrationStep::new(
+        }
-                Sigma::Absolute(vec![8.0; ndim]),
+
-                SamplingArg::Fixed(n.sqrt().max(n / 20).max(400)),
+        steps
-                24,
-                3e-7,
-                0.03,
-                100,
-                1e-2,
-            ),
-            RegistrationStep::new(
-                Sigma::Absolute(vec![2.0; ndim]),
-                SamplingArg::Fixed(n.sqrt().max(n / 10).max(800)),
-                96,
-                1e-7,
-                0.01,
-                100,
-                1e-2,
-            ),
-            RegistrationStep::new(
-                Sigma::None,
-                SamplingArg::Fixed(n.sqrt().max(n / 5).max(1600)),
-                240,
-                1e-8,
-                0.001,
-                100,
-                1e-3,
-            ),
-        ]
     }
 }
 
@@ -112,6 +110,7 @@ pub struct Registration<D: Dimension> {
     fixed_mu: FixedMu,
     steps: Option<Vec<RegistrationStep>>,
     initial_guess: Option<Vec<f64>>,
+    optimizer: Optimizer,
     dimension: PhantomData<D>,
 }
 
@@ -121,6 +120,7 @@ impl<D: Dimension> Registration<D> {
             fixed_mu: fixed_mu.into(),
             steps: None,
             initial_guess: None,
+            optimizer: Optimizer::default(),
             dimension: PhantomData,
         }
     }
@@ -135,6 +135,7 @@ impl<D: Dimension> Registration<D> {
             fixed_mu: FixedMu::new_none(ndim),
             steps: None,
             initial_guess: None,
+            optimizer: Optimizer::default(),
             dimension: PhantomData,
         })
     }
@@ -154,6 +155,7 @@ impl<D: Dimension> Registration<D> {
             fixed_mu: fixed_mu.into(),
             steps: None,
             initial_guess: None,
+            optimizer: Optimizer::default(),
             dimension: PhantomData,
         })
     }
@@ -168,6 +170,11 @@ impl<D: Dimension> Registration<D> {
         self
     }
 
+    pub fn with_optimizer(mut self, optimizer: Optimizer) -> Self {
+        self.optimizer = optimizer;
+        self
+    }
+
     pub fn set_steps(&mut self, steps: Vec<RegistrationStep>) {
         self.steps = Some(steps);
     }
@@ -176,6 +183,48 @@ impl<D: Dimension> Registration<D> {
         self.initial_guess = Some(initial_guess);
     }
 
+    pub fn set_optimizer(&mut self, optimizer: Optimizer) {
+        self.optimizer = optimizer;
+    }
+
+    /// Compute automatic transform initialization by aligning geometric centers
+    fn compute_initial_transform(
+        fixed_shape: &[usize],
+        moving_shape: &[usize],
+        ndim: usize,
+    ) -> Vec<f64> {
+        let mut params = vec![0.0; ndim * ndim + ndim];
+        for i in 0..ndim {
+            params[(ndim + 1) * i] = 1.0;
+        }
+        for i in 0..ndim {
+            let fixed_center = (fixed_shape[i] as f64 - 1.0) / 2.0;
+            let moving_center = (moving_shape[i] as f64 - 1.0) / 2.0;
+            params[ndim * ndim + i] = fixed_center - moving_center;
+        }
+        params
+    }
+
+    /// Compute golden standard scales for ASGD parameter normalization.
+    /// Matches elastix ScaledCostFunction convention: scales = sqrt(C[i][i]).
+    /// For affine transform v = R*(u-c) + t + c:
+    /// - Rotation/scale params: C[i][i] = Var(u-c) = (N^2-1)/12, scales[i] = sqrt(C[i][i])
+    /// - Translation params: C[i][i] = 1.0, scales[i] = 1.0
+    fn golden_standard_scales(fixed_shape: &[usize], ndim: usize) -> Vec<f64> {
+        let n_params = ndim * ndim + ndim;
+        let mut scales = Vec::with_capacity(n_params);
+        for _i in 0..ndim {
+            for j in 0..ndim {
+                let n = fixed_shape[j] as f64;
+                scales.push(((n * n - 1.0) / 12.0).sqrt());
+            }
+        }
+        for _ in 0..ndim {
+            scales.push(1.0);
+        }
+        scales
+    }
+
     /// find the transform which transforms moving into fixed and return the results of each
     /// optimization step
     pub fn register<'a, F, M, T>(&self, fixed: F, moving: M) -> Result<Transform<D>, Error>
@@ -193,11 +242,11 @@ impl<D: Dimension> Registration<D> {
             .as_ref()
             .cloned()
             .unwrap_or_else(|| RegistrationStep::default_steps(ndim, n));
-        let mut p = self
+
-            .initial_guess
+        let mut p = self.initial_guess.as_ref().cloned().unwrap_or_else(|| {
-            .as_ref()
+            Self::compute_initial_transform(fixed.shape(), moving.shape(), ndim)
-            .cloned()
+        });
-            .unwrap_or_else(|| Transform::<D>::default().parameters);
+
         for RegistrationStep {
             sigma,
             samples,
@@ -217,16 +266,39 @@ impl<D: Dimension> Registration<D> {
             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 optimization_config = OptimizationConfig {
+
-                max_iterations,
+            let optimization_result = match &self.optimizer {
-                tolerance,
+                Optimizer::LBFGS => {
-                learning_rate,
+                    let config = OptimizationConfig {
+                        max_iterations,
+                        tolerance,
+                        learning_rate,
+                    };
+                    lbfgs_minimize(
+                        &metric,
+                        metric.fixed_mu().extract_variable(&p).as_slice(),
+                        &config,
+                    )
+                }
+                Optimizer::ASGD => {
+                    let scales = Self::golden_standard_scales(fixed.shape(), ndim);
+                    let config = AsgdConfig {
+                        max_iterations,
+                        tolerance,
+                        maximum_step_length: 1.0,
+                        sp_a: 20.0,
+                        sp_alpha: 1.0,
+                        scales: Some(scales),
+                        ..Default::default()
+                    };
+                    asgd_minimize(
+                        &metric,
+                        metric.fixed_mu().extract_variable(&p).as_slice(),
+                        &config,
+                    )
+                }
             };
-            let optimization_result = lbfgs_minimize(
+
-                &metric,
-                metric.fixed_mu().extract_variable(&p).as_slice(),
-                &optimization_config,
-            );
             if optimization_result
                 .optimal_point
                 .iter()
@@ -261,11 +333,11 @@ impl<D: Dimension> Registration<D> {
             .as_ref()
             .cloned()
             .unwrap_or_else(|| RegistrationStep::default_steps(ndim, n));
-        let mut p = self
+
-            .initial_guess
+        let mut p = self.initial_guess.as_ref().cloned().unwrap_or_else(|| {
-            .as_ref()
+            Self::compute_initial_transform(fixed.shape(), moving.shape(), ndim)
-            .cloned()
+        });
-            .unwrap_or_else(|| Transform::<D>::default().parameters);
+
         let mut registration_results = Vec::new();
         for RegistrationStep {
             sigma,
@@ -291,16 +363,39 @@ impl<D: Dimension> Registration<D> {
             };
             let metric = MattesMetric::new(bf, bm, samples, n_bins, edge)?
                 .with_fixed_mu(self.fixed_mu.clone());
-            let optimization_config = OptimizationConfig {
+
-                max_iterations,
+            let optimization_result = match &self.optimizer {
-                tolerance,
+                Optimizer::LBFGS => {
-                learning_rate,
+                    let config = OptimizationConfig {
+                        max_iterations,
+                        tolerance,
+                        learning_rate,
+                    };
+                    lbfgs_minimize(
+                        &metric,
+                        metric.fixed_mu().extract_variable(&p).as_slice(),
+                        &config,
+                    )
+                }
+                Optimizer::ASGD => {
+                    let scales = Self::golden_standard_scales(fixed.shape(), ndim);
+                    let config = AsgdConfig {
+                        max_iterations,
+                        tolerance,
+                        maximum_step_length: 1.0,
+                        sp_a: 20.0,
+                        sp_alpha: 1.0,
+                        scales: Some(scales),
+                        ..Default::default()
+                    };
+                    asgd_minimize(
+                        &metric,
+                        metric.fixed_mu().extract_variable(&p).as_slice(),
+                        &config,
+                    )
+                }
             };
-            let optimization_result = lbfgs_minimize(
+
-                &metric,
-                metric.fixed_mu().extract_variable(&p).as_slice(),
-                &optimization_config,
-            );
             if optimization_result
                 .optimal_point
                 .iter()

src/transform.rs

@@ -504,7 +504,7 @@ mod tests {
     use crate::julia_image;
     use crate::transform::Transform;
     use itertools::Itertools;
-    use ndarray::s;
+    use ndarray::{Ix2, s};
     use num::traits::FloatConst;
 
     #[test]
@@ -556,6 +556,167 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn grad_check() -> Result<(), Box<dyn std::error::Error>> {
+        use crate::bspline::{BSpline, BSplineTrait};
+        use crate::metric::{MattesMetric, SamplingArg};
+        use algos::ObjectiveFunction;
+        let im_a = julia_image(
+            &[100, 1],
+            &[1.0, 0.0, 0.0, 0.01, 0.0, 0.0],
+            &[99.5, 0.5],
+            &[-0.8, 0.156],
+        )
+        .slice(s![.., 0])
+        .mapv(|i| i as f64);
+        let q = vec![0.85, 4.0];
+        let im_b =
+            Transform::new(q.clone(), vec![im_a.shape()[0]]).interpolate::<1, _, _>(&im_a)?;
+
+        // Test with ALL points (like grad_check uses)
+        let metric_all = MattesMetric::<ndarray::Ix1>::new(
+            BSpline::new(im_b.view()),
+            BSpline::new(im_a.view()),
+            SamplingArg::FixedAt((0..100).map(|i| vec![i as f64]).collect()),
+            3,
+            0.05,
+        )?
+        .with_fixed_mu(crate::metric::FixedMu::new_none(1));
+
+        let identity = vec![1.0, 0.0];
+        let val_i = metric_all.evaluate(&identity);
+        let grad_i = metric_all.gradient(&identity).unwrap();
+        println!("ALL POINTS - identity: val={}, grad={:?}", val_i, grad_i);
+        let truth = vec![0.85, 4.0];
+        let val_t = metric_all.evaluate(&truth);
+        let grad_t = metric_all.gradient(&truth).unwrap();
+        println!("ALL POINTS - truth: val={}, grad={:?}", val_t, grad_t);
+        let inv = vec![1.0 / 0.85, -4.0 / 0.85];
+        let val_inv = metric_all.evaluate(&inv);
+        println!("ALL POINTS - inverse: val={}", val_inv);
+
+        // Test with 100 random points (like register1 uses at finest level)
+        let metric_rand = MattesMetric::<ndarray::Ix1>::new(
+            BSpline::new(im_b.view()),
+            BSpline::new(im_a.view()),
+            SamplingArg::Fixed(100),
+            3,
+            0.05,
+        )?
+        .with_fixed_mu(crate::metric::FixedMu::new_none(1));
+
+        let val_i2 = metric_rand.evaluate(&identity);
+        let grad_i2 = metric_rand.gradient(&identity).unwrap();
+        println!("RAND 100 - identity: val={}, grad={:?}", val_i2, grad_i2);
+
+        // Test with 100 random points and 32 bins (like level 0)
+        let metric_32 = MattesMetric::<ndarray::Ix1>::new(
+            BSpline::new(im_b.view()),
+            BSpline::new(im_a.view()),
+            SamplingArg::Fixed(100),
+            32,
+            0.05,
+        )?
+        .with_fixed_mu(crate::metric::FixedMu::new_none(1));
+
+        let val_i3 = metric_32.evaluate(&identity);
+        let grad_i3 = metric_32.gradient(&identity).unwrap();
+        println!(
+            "RAND 100 bins=32 - identity: val={}, grad={:?}",
+            val_i3, grad_i3
+        );
+
+        let eps = 1e-5;
+        for i in 0..2 {
+            let mut p_plus = identity.clone();
+            let mut p_minus = identity.clone();
+            p_plus[i] += eps;
+            p_minus[i] -= eps;
+            let num_grad =
+                (metric_all.evaluate(&p_plus) - metric_all.evaluate(&p_minus)) / (2.0 * eps);
+            println!(
+                "numerical d/dmu[{}] = {} (analytical: {})",
+                i, num_grad, grad_i[i]
+            );
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn metric_landscape() -> Result<(), Box<dyn std::error::Error>> {
+        use crate::bspline::{BSpline, BSplineTrait};
+        use crate::metric::{MattesMetric, SamplingArg, Sigma};
+        use algos::ObjectiveFunction;
+        let im_a = julia_image(
+            &[100, 1],
+            &[1.0, 0.0, 0.0, 0.01, 0.0, 0.0],
+            &[99.5, 0.5],
+            &[-0.8, 0.156],
+        )
+        .slice(s![.., 0])
+        .mapv(|i| i as f64);
+        let q = vec![0.85, 4.0];
+        let im_b =
+            Transform::new(vec![0.85, 4.0], vec![im_a.shape()[0]]).interpolate::<1, _, _>(&im_a)?;
+
+        // Smooth both images with sigma=8 like level 0
+        let sigma = Sigma::Absolute(vec![8.0]);
+        let sf = sigma.smooth(im_a.view())?;
+        let sm = sigma.smooth(im_b.view())?;
+
+        let points: Vec<Vec<f64>> = (0..100).map(|i| vec![i as f64]).collect();
+
+        let metric = MattesMetric::<ndarray::Ix1>::new(
+            BSpline::new(sf.view()),
+            BSpline::new(sm.view()),
+            SamplingArg::FixedAt(points.clone()),
+            32,
+            0.05,
+        )?
+        .with_fixed_mu(crate::metric::FixedMu::new_none(1));
+
+        // Test several points
+        let test_points: Vec<(&str, Vec<f64>)> = vec![
+            ("identity".into(), vec![1.0, 0.0]),
+            ("truth".into(), vec![0.85, 4.0]),
+            ("neg_trans".into(), vec![1.0, -4.0]),
+            ("scale_0.9".into(), vec![0.9, 0.0]),
+            ("scale_1.1".into(), vec![1.1, 0.0]),
+        ];
+
+        println!("=== SMOOTHED ALL 100 integer points bins=32 ===");
+        for (name, p) in &test_points {
+            let val = metric.evaluate(p);
+            let grad = metric.gradient(p);
+            match grad {
+                Some(g) => println!("  {}: val={:.6}, grad={:?}", name, val, g),
+                None => println!("  {}: val={:.6}, grad=None", name, val),
+            }
+        }
+
+        // Now test with Fixed(100) random points
+        let metric_rand = MattesMetric::<ndarray::Ix1>::new(
+            BSpline::new(sf.view()),
+            BSpline::new(sm.view()),
+            SamplingArg::Fixed(100),
+            32,
+            0.05,
+        )?
+        .with_fixed_mu(crate::metric::FixedMu::new_none(1));
+
+        println!("=== SMOOTHED Random 100 points bins=32 ===");
+        for (name, p) in &test_points {
+            let val = metric_rand.evaluate(p);
+            let grad = metric_rand.gradient(p);
+            match grad {
+                Some(g) => println!("  {}: val={:.6}, grad={:?}", name, val, g),
+                None => println!("  {}: val={:.6}, grad=None", name, val),
+            }
+        }
+
+        Ok(())
+    }
+
     #[test]
     fn register1() -> Result<(), Box<dyn std::error::Error>> {
         let im_a = julia_image(
@@ -566,17 +727,111 @@ mod tests {
         )
         .slice(s![.., 0])
         .mapv(|i| i as f64);
+        let q = vec![0.85, 4.0];
         let im_b =
-            Transform::new(vec![0.85, 4.0], vec![im_a.shape()[0]]).interpolate::<1, _, _>(&im_a)?;
+            Transform::new(q.clone(), vec![im_a.shape()[0]]).interpolate::<1, _, _>(&im_a)?;
 
-        let (t, steps) =
+        // The registration finds T such that im_b(T(x)) = im_a(x), which is the inverse of q
-            Transform::register_debug(im_a.view(), im_b.view(), vec![None, None], None, None)?;
+        let q_inv = Transform::<ndarray::Ix1>::new(q.clone(), vec![im_a.shape()[0]])
+            .inverse()?
+            .parameters;
+
+        let all_points: Vec<Vec<f64>> = (0..100).map(|j| vec![j as f64]).collect();
+        let steps = vec![
+            crate::register::RegistrationStep::new(
+                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.clone()),
+                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(
+            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);
+        assert!(
+            t.parameters
+                .iter()
+                .zip(q_inv.iter())
+                .map(|(a, b)| (a - b).powi(2))
+                .sum::<f64>()
+                < 1.0
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn register2_interpolate() -> Result<(), Box<dyn std::error::Error>> {
+        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;
+
+        // Use default_steps which matches elastix FixedSmoothingImagePyramid:
+        // sigma = [4.0, 2.0, 1.0, 0.5] (schedule [8,4,2,1] with spacing=1)
+        let (t, steps) = Transform::register_debug(
+            im_a.view(),
+            im_b.view(),
+            vec![None, None, None, None, None, None],
+            None,
+            None,
+        )?;
+        println!("steps:");
+        for step in steps {
+            println!("  {:?}", step);
+        }
+        println!("t: {:?}", t);
+        println!("i: {:?}", t.inverse()?);
+        println!("q_inv: {:?}", q_inv);
+        assert!(
+            t.parameters
+                .iter()
+                .zip(q_inv.iter())
+                .map(|(a, b)| (a - b).powi(2))
+                .sum::<f64>()
+                < 1.0
+        );
         Ok(())
     }
 
@@ -603,7 +858,7 @@ mod tests {
         let (t, steps) = Transform::register_debug(
             im_a.view(),
             im_b.view(),
-            vec![None, None, None, None, Some(0.0), Some(0.0)],
+            vec![None, None, None, None, None, None],
             None,
             None,
         )?;
@@ -613,6 +868,19 @@ mod tests {
         }
         println!("t: {:?}", t);
         println!("i: {:?}", t.inverse()?);
+        println!("p: {:?}", p);
+        // julia_image applies transform to coordinates, so T maps im_b->im_a means T = p_inv
+        let p_inv = Transform::<ndarray::Ix2>::new(p.clone(), vec![600, 800])
+            .inverse()?
+            .parameters;
+        assert!(
+            t.parameters
+                .iter()
+                .zip(p_inv.iter())
+                .map(|(a, b)| (a - b).powi(2))
+                .sum::<f64>()
+                < 1.0
+        );
         Ok(())
     }
 }