wim/image-registration
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

- registration getting better

Browse Source
This commit is contained in:
Wim Pomp
2026-07-29 20:30:57 +02:00
parent e677e7400b
commit bb3d46cc9d
3 changed files with 136 additions and 304 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/optimize.rs
+59 -196
View File
@@ -13,18 +13,16 @@ where
T: Float + Debug,
F: ObjectiveFunction<T>,
{
const M: usize = 10; // Number of corrections to store
const M: usize = 10;
let n = initial_point.len();
let mut current_point = initial_point.to_vec();
let mut iterations = 0;
let mut converged = false;
// Storage for the last M corrections
let mut s_list: VecDeque<Vec<T>> = VecDeque::with_capacity(M);
let mut y_list: VecDeque<Vec<T>> = VecDeque::with_capacity(M);
let mut rho_list: VecDeque<T> = VecDeque::with_capacity(M);
// Get initial gradient
let mut gradient = match f.gradient(&current_point) {
Some(g) => g,
None => {
@@ -38,7 +36,6 @@ where
};
while iterations < config.max_iterations {
// Check for convergence
let gradient_norm = gradient
.iter()
.fold(T::zero(), |acc, &x| acc + x * x)
@@ -48,11 +45,9 @@ where
break;
}
// Compute search direction using L-BFGS two-loop recursion
let mut q = gradient.clone();
let mut alpha_list = Vec::with_capacity(s_list.len());
// First loop
for i in (0..s_list.len()).rev() {
let alpha = rho_list[i]
* s_list[i]
@@ -65,7 +60,6 @@ where
}
}
// Scale the initial Hessian approximation
let mut r = if !s_list.is_empty() {
let i = s_list.len() - 1;
let yy = y_list[i].iter().fold(T::zero(), |acc, &y| acc + y * y);
@@ -81,7 +75,6 @@ where
q
};
// Second loop
for i in 0..s_list.len() {
let beta = rho_list[i]
* y_list[i]
@@ -94,28 +87,23 @@ where
}
}
// r now contains the search direction
let direction: Vec<T> = r.iter().map(|&x| -x).collect();
// 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
let c1 = T::from(1e-4).unwrap();
// 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);
// Armijo: f(x + αd) ≤ f(x) + c₁ α (gᵀd)
if new_value <= current_value + c1 * alpha * g_dot_d {
improved = true;
break;
@@ -123,18 +111,15 @@ where
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) {
Some(g) => g,
None => break,
};
// Update the correction vectors
let s = new_point
.iter()
.zip(current_point.iter())
@@ -151,7 +136,6 @@ where
.zip(s.iter())
.fold(T::zero(), |acc, (&y_i, &s_i)| acc + y_i * s_i);
// FIX: Check for non-positive ys (corrupts Hessian approximation)
if ys <= T::zero() {
break;
}
@@ -167,7 +151,6 @@ where
y_list.push_back(y);
rho_list.push_back(rho);
// Update for next iteration
current_point = new_point;
gradient = new_gradient;
iterations += 1;
@@ -181,30 +164,16 @@ where
}
}
/// 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. 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,
/// 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>>,
}
@@ -215,7 +184,7 @@ impl Default for AsgdConfig {
tolerance: 1e-6,
maximum_step_length: 1.0,
sp_a: 20.0,
sp_alpha: 1.0,
sp_alpha: 0.602,
sigmoid_max: 1.0,
sigmoid_min: -0.01,
sigmoid_scale: 1e-8,
@@ -224,16 +193,6 @@ impl Default for AsgdConfig {
}
}
/// 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,
@@ -244,7 +203,6 @@ where
let mut iterations = 0;
let mut converged = false;
// Get initial gradient
let mut gradient = match f.gradient(&current_point) {
Some(g) => g,
None => {
@@ -257,37 +215,47 @@ where
}
};
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);
let alpha = T::from(config.sp_alpha).unwrap();
let a_param = T::from(config.sp_a).unwrap();
let delta = T::from(config.maximum_step_length).unwrap();
let fmax = T::from(config.sigmoid_max).unwrap();
let fmin = T::from(config.sigmoid_min).unwrap();
let omega = T::from(config.sigmoid_scale).unwrap();
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<_>>()),
);
// Compute max_j = max(|g_i / scales_i|)
let max_j = if let Some(ref scales) = config.scales {
let mut mj = T::zero();
for (&g, &s) in gradient.iter().zip(scales.iter()) {
let s_t = T::from(s.max(1e-10)).unwrap();
let gs = (g / s_t).abs();
if gs > mj {
mj = gs;
}
}
mj
} else {
let mut mj = T::zero();
for &g in gradient.iter() {
let ga = g.abs();
if ga > mj {
mj = ga;
}
}
mj
};
// a = delta * (A+1)^alpha / max_j (Elastix AutomaticParameterEstimation)
let a = if max_j > T::from(1e-14).unwrap() {
delta * (a_param + T::one()).powf(alpha) / max_j
} else {
delta * (a_param + T::one()).powf(alpha)
};
// 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)
@@ -298,19 +266,13 @@ where
break;
}
// Compute learning rate: a(t_k) = a / (A + t_k + 1)^alpha
let t_k: T = current_time;
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,
);
// a(k) = min(a / (A + t_k + 1)^alpha, delta)
let t_k = current_time + T::from(iterations).unwrap();
let a_t = T::min(a / (a_param + t_k + T::one()).powf(alpha), delta);
// 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/sqrt(C).
let dot_product_val: f64 = if let Some(ref prev_g) = previous_gradient {
// Update time: t_{k+1} = max(0, t_k + sigmoid(-g_k^T * g_{k-1}))
if let Some(ref prev_g) = previous_gradient {
let dot_product: T = if let Some(ref scales) = config.scales {
// Scaled dot product: sum((g/sqrt(C)) * (prev_g/sqrt(C))) = sum(g*pg/C)
gradient.iter().zip(prev_g.iter()).zip(scales.iter()).fold(
T::zero(),
|acc, ((&g, &pg), &s)| {
@@ -325,20 +287,12 @@ where
.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
};
}
previous_gradient = Some(gradient.clone());
// Elastix gradient descent step in unscaled parameter space:
// m_Position stores the unscaled position p.
// GetScaledDerivative computes g_scaled = g / scales.
// Update: p -= a_t * g_scaled = p - a_t * g / scales.
// So direction = -g / scales.
// direction = -g / scales
let direction: Vec<T> = match &config.scales {
Some(scales) => gradient
.iter()
@@ -356,36 +310,11 @@ where
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;
@@ -399,9 +328,6 @@ where
}
}
/// 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;
@@ -415,72 +341,11 @@ fn sigmoid<T: Float + Debug>(x: T, fmax: T, fmin: T, omega: T) -> T {
(fmax - fmin) * sigmoid_raw + fmin
}
/// Estimate ASGD parameters using DisplacementDistribution method.
///
/// Matches elastix `AutomaticParameterEstimationUsingDisplacementDistribution()`:
/// - `a = delta * (A+1)^alpha / (jacg + eps)` where `delta = MaximumStepLength`
/// - `jacg ≈ sqrt(gg) * 1.8` as analytical approximation
/// - `maximum_step_length` is also the per-iteration step clamp (same value as delta).
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 gg = ||g/scales||^2 = sum(g²/scales²)
let gg = 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
} else {
let mut gg = T::zero();
for &g in initial_gradient.iter() {
gg = gg + g * g;
}
gg
};
// Compute jacg = displacement distribution estimate.
// For golden standard scales: E[||d||^2] = gg, so jacg ≈ 1.8 * sqrt(gg)
let sqrt_gg = if gg > T::from(1e-20).unwrap() {
gg.sqrt()
} else {
T::zero()
};
let jacg = sqrt_gg * T::from(1.8).unwrap();
// Elastix DisplacementDistribution: a = delta * (A+1)^alpha / (jacg + eps)
let a = if jacg > T::from(1e-14).unwrap() {
delta * a_param.powf(alpha) / jacg
} else {
delta * a_param.powf(alpha)
};
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 {
@@ -512,21 +377,6 @@ mod tests {
}
}
// 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;
@@ -556,10 +406,23 @@ mod tests {
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);
}
}
}
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)])
}
}