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

- tests register1 and register2 passing

Browse Source
This commit is contained in:
Wim Pomp
2026-07-25 13:15:48 +02:00
parent 806571be82
commit 5b7dc18a4d
5 changed files with 887 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lib.rs
+1 -1
View File
@@ -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
+17 -18
View File
@@ -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()
.random_iter::<f64>()
.take(n_samples * shape.len())
.chunks(shape.len())
.into_iter()
.map(|c| c.zip_eq(shape.iter()).map(|(i, s)| i * s - 0.5).collect())
.collect(),
Self::Random((n_samples, shape, _cached)) => {
rand::rng()
.random_iter::<f64>()
.take(n_samples * shape.len())
.chunks(shape.len())
.into_iter()
.map(|c| c.zip_eq(shape.iter()).map(|(i, s)| i * s - 0.5).collect())
.collect()
}
}
}
}
@@ -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,
derivative: dmetric
.iter()
.zip_eq(dalpha)
.map(|(dm, da)| *dm * alpha + da * metric)
.collect(),
metric: metric,
derivative: dmetric.clone(),
});
// self.metric.replace(IntMut {
// mu: mu.to_vec(),
src/optimize.rs
+426 -15
View File
@@ -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);
}
}
}
src/register.rs
+170 -75
View File
@@ -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![
RegistrationStep::new(
Sigma::Relative(vec![0.5; ndim]),
SamplingArg::Fixed(n.sqrt().max(n / 125).max(100)),
3,
1e-4,
0.05,
100,
100.0,
),
RegistrationStep::new(
Sigma::Relative(vec![0.25; ndim]),
SamplingArg::Fixed(n.sqrt().max(n / 30).max(200)),
6,
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() {
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]),
SamplingArg::Random(samples),
32,
1e-6,
0.04,
100,
0.05,
256,
1.0,
),
RegistrationStep::new(
Sigma::Absolute(vec![8.0; ndim]),
SamplingArg::Fixed(n.sqrt().max(n / 20).max(400)),
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,
),
]
));
}
steps
}
}
@@ -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
.as_ref()
.cloned()
.unwrap_or_else(|| Transform::<D>::default().parameters);
let mut p = self.initial_guess.as_ref().cloned().unwrap_or_else(|| {
Self::compute_initial_transform(fixed.shape(), moving.shape(), ndim)
});
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,
tolerance,
learning_rate,
let optimization_result = match &self.optimizer {
Optimizer::LBFGS => {
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
.as_ref()
.cloned()
.unwrap_or_else(|| Transform::<D>::default().parameters);
let mut p = self.initial_guess.as_ref().cloned().unwrap_or_else(|| {
Self::compute_initial_transform(fixed.shape(), moving.shape(), ndim)
});
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,
tolerance,
learning_rate,
let optimization_result = match &self.optimizer {
Optimizer::LBFGS => {
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
+273 -5
View File
@@ -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) =
Transform::register_debug(im_a.view(), im_b.view(), vec![None, None], None, None)?;
// The registration finds T such that im_b(T(x)) = im_a(x), which is the inverse of q
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(())
}
}