- cleanup
This commit is contained in:
w.pomp committed
1 parent
f75f12db16
commit
cc12a1a9b9
5 files changed
+253
-175
No files matched your search
+57
-110
@@ -9,8 +9,7 @@ use num::cast::AsPrimitive;
|
||||
use std::marker::PhantomData;
|
||||
|
||||
/// Optimizer type for registration
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
#[derive(Default)]
|
||||
#[derive(Clone, Debug, PartialEq, Default)]
|
||||
pub enum Optimizer {
|
||||
/// L-BFGS optimizer (fast, requires consistent gradients)
|
||||
#[default]
|
||||
@@ -19,7 +18,6 @@ pub enum Optimizer {
|
||||
ASGD,
|
||||
}
|
||||
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct RegistrationStep {
|
||||
pub sigma: Sigma,
|
||||
@@ -29,7 +27,6 @@ pub struct RegistrationStep {
|
||||
pub edge: f64,
|
||||
pub max_iterations: usize,
|
||||
pub learning_rate: f64,
|
||||
pub downsample: usize,
|
||||
/// Optimizer to use for this level
|
||||
pub optimizer: Optimizer,
|
||||
}
|
||||
@@ -52,114 +49,66 @@ impl RegistrationStep {
|
||||
edge,
|
||||
max_iterations,
|
||||
learning_rate,
|
||||
downsample: 1,
|
||||
optimizer: Optimizer::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Default registration steps matching elastix `FixedSmoothingImagePyramid`.
|
||||
///
|
||||
/// Uses a hybrid multi-fidelity approach:
|
||||
/// - Coarse levels (sigma=4,2): ASGD optimizer with all pixels for speed
|
||||
/// - Fine levels (sigma=1,0.5): L-BFGS optimizer with all pixels for accuracy
|
||||
///
|
||||
/// Note: Samples are cached per level (not randomized per iteration like SimpleElastix),
|
||||
/// but ASGD's sigmoid momentum handles the stochastic gradients effectively.
|
||||
pub fn default_steps(ndim: usize, n: usize) -> Vec<Self> {
|
||||
if ndim == 1 {
|
||||
return vec![Self {
|
||||
sigma: Sigma::Absolute(vec![1.0; 1]),
|
||||
samples: SamplingArg::Random(n),
|
||||
vec![
|
||||
Self {
|
||||
sigma: Sigma::Absolute(vec![8.0; ndim]),
|
||||
samples: SamplingArg::Random((n / 256).max(1024)),
|
||||
n_bins: 16,
|
||||
tolerance: 1e-1,
|
||||
edge: 0.05,
|
||||
max_iterations: 128,
|
||||
learning_rate: 1.0,
|
||||
optimizer: Optimizer::ASGD,
|
||||
},
|
||||
Self {
|
||||
sigma: Sigma::Absolute(vec![4.0; ndim]),
|
||||
samples: SamplingArg::Random((n / 64).max(4096)),
|
||||
n_bins: 24,
|
||||
tolerance: 1e-2,
|
||||
edge: 0.05,
|
||||
max_iterations: 256,
|
||||
learning_rate: 1.0,
|
||||
optimizer: Optimizer::ASGD,
|
||||
},
|
||||
Self {
|
||||
sigma: Sigma::Absolute(vec![2.0; ndim]),
|
||||
samples: SamplingArg::Random((n / 16).max(4096)),
|
||||
n_bins: 32,
|
||||
tolerance: 1e-4,
|
||||
edge: 0.05,
|
||||
max_iterations: 512,
|
||||
learning_rate: 1.0,
|
||||
optimizer: Optimizer::ASGD,
|
||||
},
|
||||
Self {
|
||||
sigma: Sigma::Absolute(vec![1.0; ndim]),
|
||||
samples: SamplingArg::Random((n / 4).max(8192)),
|
||||
n_bins: 48,
|
||||
tolerance: 1e-6,
|
||||
edge: 0.05,
|
||||
max_iterations: 1024,
|
||||
learning_rate: 1.0,
|
||||
optimizer: Optimizer::ASGD,
|
||||
},
|
||||
Self {
|
||||
sigma: Sigma::None,
|
||||
samples: SamplingArg::Fixed(n.max(16384)),
|
||||
n_bins: 64,
|
||||
tolerance: 1e-8,
|
||||
edge: 0.05,
|
||||
max_iterations: 2048,
|
||||
learning_rate: 1.0,
|
||||
downsample: 1,
|
||||
optimizer: Optimizer::LBFGS,
|
||||
}];
|
||||
}
|
||||
// Pyramid with downsampling: schedule [8,4,2,1], sigma [4,2,1,0.5]
|
||||
let sigma_schedule: Vec<f64> = vec![4.0, 2.0, 1.0, 0.5];
|
||||
let downsample_schedule: Vec<usize> = vec![8, 4, 2, 1];
|
||||
let n_levels = sigma_schedule.len();
|
||||
sigma_schedule
|
||||
.iter()
|
||||
.zip(downsample_schedule.iter())
|
||||
.enumerate()
|
||||
.map(|(level, (&s, &d))| {
|
||||
let n_pixels = (n / (d * d)).max(4);
|
||||
let is_finest = level == n_levels - 1;
|
||||
let is_coarse = level < 2;
|
||||
// Hybrid: ASGD on coarse levels for speed, L-BFGS on fine levels for accuracy
|
||||
// Use all pixels on all levels for consistent gradient estimates
|
||||
let (optimizer, max_iter, tol) = if is_coarse {
|
||||
(Optimizer::ASGD, 512, 1e-4)
|
||||
} else if is_finest {
|
||||
(Optimizer::LBFGS, 2048, 1e-8)
|
||||
} else {
|
||||
(Optimizer::LBFGS, 512, 1e-6)
|
||||
};
|
||||
Self {
|
||||
sigma: Sigma::Absolute(vec![s; ndim]),
|
||||
samples: SamplingArg::Random(n_pixels),
|
||||
n_bins: 32,
|
||||
tolerance: tol,
|
||||
edge: 0.05,
|
||||
max_iterations: max_iter,
|
||||
learning_rate: 1.0,
|
||||
downsample: d,
|
||||
optimizer,
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
},
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
/// 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>>,
|
||||
@@ -283,15 +232,14 @@ impl<D: Dimension> Registration<D> {
|
||||
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.extend(
|
||||
fixed_shape
|
||||
.iter()
|
||||
.map(|&i| ((i * i - 1) as f64 / 12.0).sqrt()),
|
||||
);
|
||||
}
|
||||
scales.extend(vec![1.0; ndim]);
|
||||
scales
|
||||
}
|
||||
|
||||
@@ -325,7 +273,6 @@ impl<D: Dimension> Registration<D> {
|
||||
edge,
|
||||
max_iterations,
|
||||
learning_rate,
|
||||
downsample: _,
|
||||
optimizer,
|
||||
} in steps.into_iter()
|
||||
{
|
||||
@@ -356,9 +303,10 @@ impl<D: Dimension> Registration<D> {
|
||||
);
|
||||
|
||||
if let Some(result) = optimization_result
|
||||
&& result.optimal_point.iter().all(|i| i.is_finite()) {
|
||||
p = self.fixed_mu.combine(&result.optimal_point);
|
||||
}
|
||||
&& result.optimal_point.iter().all(|i| i.is_finite())
|
||||
{
|
||||
p = self.fixed_mu.combine(&result.optimal_point);
|
||||
}
|
||||
}
|
||||
Ok(Transform::<D>::new(p, fixed.shape().to_vec()))
|
||||
}
|
||||
@@ -431,7 +379,6 @@ impl<D: Dimension> Registration<D> {
|
||||
edge,
|
||||
max_iterations,
|
||||
learning_rate,
|
||||
downsample: _,
|
||||
optimizer,
|
||||
} in steps.into_iter()
|
||||
{
|
||||
|
||||
Reference in new issue
Block a user