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

- register2 passing

Browse Source
This commit is contained in:
w.pomp
2026-07-28 11:27:08 +02:00
parent 5b7dc18a4d
commit 9f8764ee00
5 changed files with 305 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -74,6 +74,7 @@ docs/_build/
AGENTS.md
.agentbridge
.agent-work
*.tif
*.svg
Cargo.toml
+3
View File
@@ -34,3 +34,6 @@ tempfile = "3"
[profile.release]
debug = true
[profile.test]
inherits = "release"
src/optimize.rs
+40 -53
View File
@@ -188,7 +188,9 @@ pub struct AsgdConfig {
pub max_iterations: usize,
/// Convergence tolerance (gradient norm)
pub tolerance: f64,
/// Maximum step length (displacement in mm). Default: 1.0
/// Maximum step length. In elastix this is `delta` — used both for initial
/// learning rate estimation (`a = delta * (A+1)^alpha / (jacg + eps)`) and
/// as the per-iteration step clamp. Default: 1.0
pub maximum_step_length: f64,
/// Gain parameter A (denominator offset). Default: 20.0
pub sp_a: f64,
@@ -298,21 +300,25 @@ where
// Compute learning rate: a(t_k) = a / (A + t_k + 1)^alpha
let t_k: T = current_time;
let a_t = a / (T::from(config.sp_a).unwrap() + t_k + T::one()).powf(alpha);
let max_step: T = num::cast(config.maximum_step_length).unwrap();
let a_t = T::min(
a / (T::from(config.sp_a).unwrap() + t_k + T::one()).powf(alpha),
max_step,
);
// Update time using sigmoid: t_{k+1} = max(0, t_k + sigmoid(-g_k^T * g_{k-1}))
// In elastix, the dot product is computed in SCALED space: g_scaled = g/scales.
// In elastix, the dot product is computed in SCALED space: g_scaled = g/sqrt(C).
let dot_product_val: f64 = if let Some(ref prev_g) = previous_gradient {
let dot_product: T = if let Some(ref scales) = config.scales {
// Scaled dot product: sum(g_i * prev_g_i / scales_i²)
gradient
.iter()
.zip(prev_g.iter())
.zip(scales.iter())
.fold(T::zero(), |acc, ((&g, &pg), &s)| {
// 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)| {
let s2: T = num::cast(s.max(1e-10)).unwrap();
acc + g * pg / (s2 * s2)
})
let s4 = s2 * s2;
acc + g * pg / s4
},
)
} else {
gradient
.iter()
@@ -329,17 +335,17 @@ where
};
// Elastix gradient descent step in unscaled parameter space:
// Optimizer stores p_scaled = p * scales, updates p_scaled -= a_t * g_scaled
// 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².
// 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.
let direction: Vec<T> = match &config.scales {
Some(scales) => gradient
.iter()
.zip(scales.iter())
.map(|(&g, &s)| {
let s2: T = num::cast(s.max(1e-10)).unwrap();
-g / (s2 * s2)
-g / s2
})
.collect(),
None => gradient.iter().map(|&g| -g).collect(),
@@ -409,14 +415,12 @@ fn sigmoid<T: Float + Debug>(x: T, fmax: T, fmin: T, omega: T) -> T {
(fmax - fmin) * sigmoid_raw + fmin
}
/// Estimate ASGD parameters from gradient statistics.
/// Estimate ASGD parameters using DisplacementDistribution method.
///
/// Matches elastix `AutomaticParameterEstimationOriginal()`:
/// - `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)
/// 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],
@@ -432,54 +436,37 @@ where
let fmax = T::from(config.sigmoid_max).unwrap();
let delta = T::from(config.maximum_step_length).unwrap();
// Compute scaled gradient: g_scaled = g / scales
// And its squared norm: gg = ||g/scales||^2
let (gg, trc, max_jcj) = if let Some(ref scales) = config.scales {
// 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;
}
// 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())
gg
} 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())
gg
};
// 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()
// 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();
let a = if sigma1 > T::from(1e-14).unwrap() && max_jcj > T::from(1e-14).unwrap() {
a_param * delta / sigma1 / max_jcj.sqrt()
// 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 {
a_param
delta * a_param.powf(alpha)
};
let fmin = T::from(config.sigmoid_min).unwrap();
src/register.rs
+71 -14
View File
@@ -32,6 +32,7 @@ pub struct RegistrationStep {
pub edge: f64,
pub max_iterations: usize,
pub learning_rate: f64,
pub downsample: usize,
}
impl RegistrationStep {
@@ -52,36 +53,35 @@ impl RegistrationStep {
edge,
max_iterations,
learning_rate,
downsample: 1,
}
}
/// Default registration steps matching SimpleElastix affine defaults.
/// Default registration steps matching elastix `FixedSmoothingImagePyramid`.
///
/// Elastix uses `FixedSmoothingImagePyramid` with schedule [8, 4, 2, 1] for 4 levels.
/// 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 uses `MultiResolutionGaussianSmoothingPyramidImageFilter` which applies
/// Gaussian smoothing with σ = 0.5 × factor at **full resolution** — images are
/// **NOT downsampled**. See: `itkMultiResolutionGaussianSmoothingPyramidImageFilter.hxx`
///
/// Elastix defaults: 2048 samples, 32 histogram bins, 256 iterations.
/// Schedule [8, 4, 2, 1] means σ = [4.0, 2.0, 1.0, 0.5] (spacing=1).
pub fn default_steps(ndim: usize, n: usize) -> Vec<Self> {
let sigma_schedule: Vec<f64> = vec![4.0, 2.0, 1.0, 0.5];
let nlevels = sigma_schedule.len();
let mut steps = Vec::new();
for (i, &sigma) in sigma_schedule.iter().enumerate() {
for i in 0..nlevels {
let fraction = 1.0 / 2.0_f64.powi((nlevels - 1 - i) as i32);
let samples = (n as f64 * fraction)
.sqrt()
.max(n as f64 * fraction / 10.0)
.max(2048.0) as usize;
steps.push(RegistrationStep::new(
Sigma::Absolute(vec![sigma; ndim]),
Sigma::Absolute(vec![sigma_schedule[i]; ndim]),
SamplingArg::Random(samples),
32,
1e-6,
0.05,
256,
512,
1.0,
));
}
@@ -90,6 +90,50 @@ impl RegistrationStep {
}
}
/// Downsample an n-dimensional array by integer factor along all axes.
/// Takes every `factor`-th element along each axis.
/// Only supports 1D and 2D arrays (the only types used in this crate).
pub fn downsample_nd<D>(array: ndarray::ArrayView<f64, D>, factor: usize) -> ndarray::Array<f64, D>
where
D: ndarray::Dimension,
{
if factor <= 1 {
return array.to_owned();
}
let ndim = array.ndim();
let src = array.to_owned();
let src_shape = src.shape().to_vec();
let src_slice = src.as_slice().unwrap();
let new_shape: Vec<usize> = src_shape.iter().map(|&s| (s - 1) / factor + 1).collect();
// Compute strides for row-major layout
let mut strides = vec![1usize; ndim];
for i in (0..ndim - 1).rev() {
strides[i] = strides[i + 1] * src_shape[i + 1];
}
let total: usize = new_shape.iter().product();
let mut data = Vec::with_capacity(total);
for out_flat in 0..total {
let mut out_coords = vec![0usize; ndim];
let mut tmp = out_flat;
for i in 0..ndim {
out_coords[i] = tmp % new_shape[i];
tmp /= new_shape[i];
}
let mut src_flat = 0;
for i in 0..ndim {
src_flat += out_coords[i] * factor * strides[i];
}
data.push(src_slice[src_flat]);
}
ndarray::Array::from_shape_vec(new_shape, data)
.unwrap()
.into_dimensionality::<D>()
.unwrap()
}
#[derive(Clone, Debug)]
pub struct RegistrationResult {
pub sigma_fixed: Option<Vec<f64>>,
@@ -255,10 +299,12 @@ impl<D: Dimension> Registration<D> {
edge,
max_iterations,
learning_rate,
downsample: _,
} in steps.into_iter()
{
let f = sigma.smooth(fixed.view())?;
let m = sigma.smooth(moving.view())?;
if (f.std(0.0) == 0.0) || (m.std(0.0) == 0.0) {
continue;
}
@@ -267,6 +313,9 @@ impl<D: Dimension> Registration<D> {
let metric = MattesMetric::new(bf, bm, samples, n_bins, edge)?
.with_fixed_mu(self.fixed_mu.clone());
// Golden standard scales from downsampled shape
let scales = Self::golden_standard_scales(f.shape(), ndim);
let optimization_result = match &self.optimizer {
Optimizer::LBFGS => {
let config = OptimizationConfig {
@@ -281,7 +330,6 @@ impl<D: Dimension> Registration<D> {
)
}
Optimizer::ASGD => {
let scales = Self::golden_standard_scales(fixed.shape(), ndim);
let config = AsgdConfig {
max_iterations,
tolerance,
@@ -347,13 +395,19 @@ impl<D: Dimension> Registration<D> {
edge,
max_iterations,
learning_rate,
downsample: _,
} in steps.into_iter()
{
let f = sigma.smooth(fixed.view())?;
let m = sigma.smooth(moving.view())?;
let f_orig_shape: Vec<usize> = f.shape().to_vec();
let m_orig_shape: Vec<usize> = m.shape().to_vec();
if (f.std(0.0) == 0.0) || (m.std(0.0) == 0.0) {
continue;
}
let bf = BSpline::<0, _>::new(f.view());
let bm = BSpline::<3, _>::new(m.view());
let n_samples = match &samples {
@@ -364,6 +418,9 @@ impl<D: Dimension> Registration<D> {
let metric = MattesMetric::new(bf, bm, samples, n_bins, edge)?
.with_fixed_mu(self.fixed_mu.clone());
// Golden standard scales from downsampled shape
let scales = Self::golden_standard_scales(f.shape(), ndim);
let optimization_result = match &self.optimizer {
Optimizer::LBFGS => {
let config = OptimizationConfig {
@@ -378,7 +435,6 @@ impl<D: Dimension> Registration<D> {
)
}
Optimizer::ASGD => {
let scales = Self::golden_standard_scales(fixed.shape(), ndim);
let config = AsgdConfig {
max_iterations,
tolerance,
@@ -405,9 +461,10 @@ impl<D: Dimension> Registration<D> {
.fixed_mu()
.combine(&optimization_result.optimal_point);
}
registration_results.push(RegistrationResult {
sigma_fixed: sigma.sigma(f.shape()),
sigma_moving: sigma.sigma(m.shape()),
sigma_fixed: sigma.sigma(f_orig_shape.as_slice()),
sigma_moving: sigma.sigma(m_orig_shape.as_slice()),
n_bins,
n_samples,
tolerance,
src/transform.rs
+190 -10
View File
@@ -749,22 +749,13 @@ mod tests {
),
crate::register::RegistrationStep::new(
crate::metric::Sigma::Absolute(vec![2.0]),
crate::metric::SamplingArg::FixedAt(all_points.clone()),
crate::metric::SamplingArg::FixedAt(all_points),
64,
1e-6,
0.04,
200,
1.0,
),
crate::register::RegistrationStep::new(
crate::metric::Sigma::None,
crate::metric::SamplingArg::FixedAt(all_points),
64,
1e-8,
0.001,
200,
1.0,
),
];
let (t, steps) = Transform::register_debug(
@@ -792,6 +783,195 @@ mod tests {
Ok(())
}
#[test]
fn register2_random_affine() -> Result<(), Box<dyn std::error::Error>> {
use rand::prelude::*;
let mut rng = rand::rngs::StdRng::seed_from_u64(1337);
let shape = [200, 200];
let center = [99.5, 99.5];
let fixed = julia_image(
&shape,
&[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
&center,
&[-0.8, 0.156],
)
.mapv(|i| i as f64);
// random rotation ±25°
let angle: f64 = rng.random_range(-25.0f64..25.0).to_radians();
let (s, c) = angle.sin_cos();
// random scale 0.85–1.15
let sx = rng.random_range(0.85..1.15);
let sy = rng.random_range(0.85..1.15);
// small shear
let shx: f64 = rng.random_range(-0.1..0.1);
let shy: f64 = rng.random_range(-0.1..0.1);
// translation ±30 px
let tx: f64 = rng.random_range(-30.0..30.0);
let ty: f64 = rng.random_range(-30.0..30.0);
// [m00, m01, m10, m11, tx, ty]
let params = vec![c * sx, -s * sy + shx, s * sx + shy, c * sy, tx, ty];
let transform =
Transform::<Ix2>::new_with_center(params.clone(), center.to_vec(), shape.to_vec());
let moving = transform.interpolate::<3, _, _>(fixed.view())?;
let q_inv = transform.inverse()?.parameters;
let (t, _steps) = Transform::<Ix2>::register_debug(
fixed.view(),
moving.view(),
vec![None, None, None, None, None, None],
None,
None,
)?;
println!("params: {:?}", params);
println!("result: {:?}", t.parameters);
println!("q_inv: {:?}", q_inv);
let sse: f64 = t
.parameters
.iter()
.zip(q_inv.iter())
.map(|(a, b)| (a - b).powi(2))
.sum();
println!("sse: {sse}");
assert!(sse < 1.0);
Ok(())
}
#[test]
fn metric_landscape_2d() -> Result<(), Box<dyn std::error::Error>> {
use crate::bspline::{BSpline, BSplineTrait};
use crate::filter::gaussian_smooth;
use crate::metric::{MattesMetric, SamplingArg};
use algos::ObjectiveFunction;
let shape = [200, 200];
let center = [99.5, 99.5];
let im_a = julia_image(
&shape,
&[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
&center,
&[-0.8, 0.156],
)
.mapv(|i| i as f64);
let rotation = Transform::<Ix2>::from_rotation(f64::PI() / 4.0, &center);
let im_b = rotation.interpolate::<3, _, _>(im_a.view())?;
let q_inv = rotation.inverse()?.parameters;
let p = rotation.parameters.clone();
let identity = vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0];
let edge = 0.01;
// Test at full resolution with different sigma levels (matching elastix approach)
// Elastix does NOT downsample — it smooths at full resolution with
// sigma = 0.5 * factor where factor is from the pyramid schedule [8, 4, 2, 1]
println!(
"=== Full resolution (no downsampling, matching elastix FixedSmoothingImagePyramid) ==="
);
for (level, sigma_val) in [4.0, 2.0, 1.0, 0.5].iter().enumerate() {
let f = gaussian_smooth(im_a.view(), &[*sigma_val; 2])?;
let m = gaussian_smooth(im_b.view(), &[*sigma_val; 2])?;
let bf = BSpline::<0, _>::new(f.view());
let bm = BSpline::<3, _>::new(m.view());
let metric = MattesMetric::new(bf, bm, SamplingArg::Random(3000), 128, edge)?;
let mi_id = metric.evaluate(&identity);
let mi_p = metric.evaluate(&p);
let mi_qinv = metric.evaluate(&q_inv);
println!(
" Level {} (sigma={:.1}, {}x{}): id={:.4} p={:.4} q_inv={:.4} → {}",
level,
sigma_val,
f.shape()[0],
f.shape()[1],
-mi_id,
-mi_p,
-mi_qinv,
if -mi_qinv > -mi_p {
"Q_INV correct"
} else {
"P incorrect (landscape inverted!)"
}
);
}
// Verify: at every sigma level, q_inv should have higher MI than p
let f4 = gaussian_smooth(im_a.view(), &[4.0, 4.0])?;
let m4 = gaussian_smooth(im_b.view(), &[4.0, 4.0])?;
let bf4 = BSpline::<0, _>::new(f4.view());
let bm4 = BSpline::<3, _>::new(m4.view());
let metric4 = MattesMetric::new(bf4, bm4, SamplingArg::Random(5000), 128, edge)?;
let mi_p_coarse = metric4.evaluate(&p);
let mi_qinv_coarse = metric4.evaluate(&q_inv);
println!(
"\nCoarsest (sigma=4.0): p={:.6} q_inv={:.6}",
-mi_p_coarse, -mi_qinv_coarse
);
assert!(
-mi_qinv_coarse > -mi_p_coarse,
"MI landscape is inverted at coarsest level! q_inv={} should be > p={}",
-mi_qinv_coarse,
-mi_p_coarse
);
Ok(())
}
#[test]
fn angle_sweep() -> Result<(), Box<dyn std::error::Error>> {
use crate::bspline::{BSpline, BSplineTrait};
use crate::filter::gaussian_smooth;
use crate::metric::{MattesMetric, SamplingArg};
use algos::ObjectiveFunction;
let shape = [200, 200];
let center = [99.5, 99.5];
let im_a = julia_image(
&shape,
&[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
&center,
&[-0.8, 0.156],
)
.mapv(|i| i as f64);
let edge = 0.01;
// Full resolution, sigma=4.0 — matching elastix FixedSmoothingImagePyramid (no downsampling)
for angle_deg in [5.0f64, 10.0, 15.0, 20.0, 25.0, 30.0, 35.0, 40.0, 45.0] {
let rotation = Transform::<Ix2>::from_rotation(angle_deg.to_radians(), &center);
let im_b = rotation.interpolate::<3, _, _>(im_a.view())?;
let q_inv = rotation.inverse()?.parameters;
let p = rotation.parameters.clone();
let f = gaussian_smooth(im_a.view(), &[4.0, 4.0])?;
let m = gaussian_smooth(im_b.view(), &[4.0, 4.0])?;
let metric = MattesMetric::new(
BSpline::<0, _>::new(f.view()),
BSpline::<3, _>::new(m.view()),
SamplingArg::Random(3000),
128,
edge,
)?;
let id = vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0];
let mi_id = metric.evaluate(&id);
let mi_p = metric.evaluate(&p);
let mi_qi = metric.evaluate(&q_inv);
println!(
"{:5.1}°: id={:.3} p={:.3} q_inv={:.3} → {}",
angle_deg,
-mi_id,
-mi_p,
-mi_qi,
if mi_p > mi_qi { "P wins" } else { "Q_INV wins" }
);
}
Ok(())
}
#[test]
fn register2_interpolate() -> Result<(), Box<dyn std::error::Error>> {
let shape = [200, 200];