- tests register1 and register2 passing

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()
-                .zip(q.iter())
-                .fold(T::zero(), |acc, (&s, &q)| acc + s * q);
+                    .iter()
+                    .zip(q.iter())
+                    .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()
-                .zip(r.iter())
-                .fold(T::zero(), |acc, (&y, &r)| acc + y * r);
+                    .iter()
+                    .zip(r.iter())
+                    .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
-        for _ in 0..20 {
+        // Backtracking line search with Armijo condition
+        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);
+        }
+    }
+}