diff --git a/Cargo.toml b/Cargo.toml index ac27cf1..5450e5d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,6 +30,7 @@ thiserror = "2" tempfile = "3" tiff = "0.11" tiffwrite = "2026.6.0" +ndbioimage = { version = "0.2.0", default-features = false, features = ["czi"] } [profile.release] debug = true diff --git a/src/error.rs b/src/error.rs index 33f198c..d1f66f7 100644 --- a/src/error.rs +++ b/src/error.rs @@ -13,4 +13,8 @@ pub enum Error { /// the dimensionality of the data does not match #[error("dimensionality mismatch: {0} != {0}")] DimensionalityMismatch(usize, usize), + #[error("expected fixed_mu length of {0}, got {1}")] + FixedMuLength(usize, usize), + #[error("expected initial guess length of {0}, got {1}")] + InitialGuessLength(usize, usize), } diff --git a/src/metric.rs b/src/metric.rs index 6a86825..d12a9ea 100644 --- a/src/metric.rs +++ b/src/metric.rs @@ -11,6 +11,7 @@ use num::cast::AsPrimitive; use rand::RngExt; use rayon::iter::IntoParallelIterator; use rayon::iter::ParallelIterator; +use serde::{Deserialize, Serialize}; use std::cell::RefCell; use std::ops::Deref; @@ -71,6 +72,10 @@ impl FixedMu { .filter_map(|(i, j)| if i.is_none() { Some(*j) } else { None }) .collect() } + + pub fn len_variable(&self) -> usize { + self.len_variable + } } /// first_bin, last_bin: center of bin @@ -178,7 +183,7 @@ impl Sampling { } } -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] pub enum SamplingArg { Fixed(usize), FixedAt(Vec>), @@ -452,7 +457,7 @@ where } } -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] pub enum Sigma { None, Absolute(Vec), diff --git a/src/register.rs b/src/register.rs index 185e25d..2e29242 100644 --- a/src/register.rs +++ b/src/register.rs @@ -6,10 +6,11 @@ use crate::transform::Transform; use algos::{ObjectiveFunction, OptimizationConfig}; use ndarray::{AsArray, Dimension}; use num::cast::AsPrimitive; +use serde::{Deserialize, Serialize}; use std::marker::PhantomData; /// Optimizer type for registration -#[derive(Clone, Debug, PartialEq, Default)] +#[derive(Clone, Debug, PartialEq, Default, Serialize, Deserialize)] pub enum Optimizer { /// L-BFGS optimizer (fast, requires consistent gradients) #[default] @@ -18,7 +19,7 @@ pub enum Optimizer { ASGD, } -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] pub struct RegistrationStep { pub sigma: Sigma, pub samples: SamplingArg, @@ -32,6 +33,7 @@ pub struct RegistrationStep { } impl RegistrationStep { + #[expect(clippy::too_many_arguments)] pub fn new( sigma: Sigma, samples: SamplingArg, @@ -40,6 +42,7 @@ impl RegistrationStep { edge: f64, max_iterations: usize, learning_rate: f64, + optimizer: Optimizer, ) -> Self { Self { sigma, @@ -49,7 +52,7 @@ impl RegistrationStep { edge, max_iterations, learning_rate, - optimizer: Optimizer::default(), + optimizer, } } @@ -134,13 +137,20 @@ pub struct Registration { } impl Registration { - pub fn new>(fixed_mu: F) -> Self { - Self { - fixed_mu: fixed_mu.into(), - steps: None, - initial_guess: None, - optimizer: Optimizer::default(), - dimension: PhantomData, + pub fn new>(fixed_mu: F) -> Result { + let fixed_mu = fixed_mu.into(); + if let Some(ndim) = D::NDIM + && fixed_mu.len() != (ndim * ndim + ndim) + { + Err(Error::FixedMuLength(ndim * ndim + ndim, fixed_mu.len())) + } else { + Ok(Self { + fixed_mu, + steps: None, + initial_guess: None, + optimizer: Optimizer::default(), + dimension: PhantomData, + }) } } @@ -184,9 +194,16 @@ impl Registration { self } - pub fn with_initial_guess(mut self, initial_guess: Vec) -> Self { - self.initial_guess = Some(initial_guess); - self + pub fn with_initial_guess(mut self, initial_guess: Vec) -> Result { + if self.fixed_mu.len_variable() != initial_guess.len() { + Err(Error::InitialGuessLength( + self.fixed_mu.len(), + initial_guess.len(), + )) + } else { + self.initial_guess = Some(initial_guess); + Ok(self) + } } pub fn with_optimizer(mut self, optimizer: Optimizer) -> Self { @@ -198,8 +215,16 @@ impl Registration { self.steps = Some(steps); } - pub fn set_initial_guess(&mut self, initial_guess: Vec) { - self.initial_guess = Some(initial_guess); + pub fn set_initial_guess(&mut self, initial_guess: Vec) -> Result<(), Error> { + if self.fixed_mu.len_variable() != initial_guess.len() { + Err(Error::InitialGuessLength( + self.fixed_mu.len(), + initial_guess.len(), + )) + } else { + self.initial_guess = Some(initial_guess); + Ok(()) + } } pub fn set_optimizer(&mut self, optimizer: Optimizer) { @@ -243,8 +268,7 @@ impl Registration { scales } - /// find the transform which transforms moving into fixed and return the results of each - /// optimization step + /// find the transform which transforms moving into fixed pub fn register<'a, F, M, T>(&self, fixed: F, moving: M) -> Result, Error> where F: AsArray<'a, T, D>, @@ -261,7 +285,7 @@ impl Registration { .cloned() .unwrap_or_else(|| RegistrationStep::default_steps(ndim, n)); - let mut p = self.initial_guess.as_ref().cloned().unwrap_or_else(|| { + let mut p_var = self.initial_guess.as_ref().cloned().unwrap_or_else(|| { Self::compute_initial_transform(fixed.shape(), moving.shape(), ndim) }); @@ -290,8 +314,6 @@ impl Registration { 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 p_var = metric.fixed_mu().extract_variable(&p); - let optimization_result = Self::optimize_metric( &metric, &p_var, @@ -305,10 +327,13 @@ impl Registration { if let Some(result) = optimization_result && result.optimal_point.iter().all(|i| i.is_finite()) { - p = self.fixed_mu.combine(&result.optimal_point); + p_var = result.optimal_point; } } - Ok(Transform::::new(p, fixed.shape().to_vec())) + Ok(Transform::::new( + self.fixed_mu.combine(p_var.as_slice()), + fixed.shape().to_vec(), + )) } fn optimize_metric>( @@ -366,9 +391,10 @@ impl Registration { .cloned() .unwrap_or_else(|| RegistrationStep::default_steps(ndim, n)); - let mut p = self.initial_guess.as_ref().cloned().unwrap_or_else(|| { + let mut p_var = self.initial_guess.as_ref().cloned().unwrap_or_else(|| { Self::compute_initial_transform(fixed.shape(), moving.shape(), ndim) }); + let mut p = vec![0.0; n * n + n]; let mut registration_results = Vec::new(); for RegistrationStep { @@ -412,11 +438,7 @@ impl Registration { tolerance, learning_rate, }; - lbfgs_minimize( - &metric, - metric.fixed_mu().extract_variable(&p).as_slice(), - &config, - ) + lbfgs_minimize(&metric, &p_var, &config) } Optimizer::ASGD => { let config = AsgdConfig { @@ -428,11 +450,7 @@ impl Registration { scales: Some(scales), ..Default::default() }; - asgd_minimize( - &metric, - metric.fixed_mu().extract_variable(&p).as_slice(), - &config, - ) + asgd_minimize(&metric, &p_var, &config) } }; @@ -441,6 +459,7 @@ impl Registration { .iter() .all(|i| i.is_finite()) { + p_var = optimization_result.optimal_point.clone(); p = metric .fixed_mu() .combine(&optimization_result.optimal_point); diff --git a/src/transform.rs b/src/transform.rs index ebb9e74..7b58290 100644 --- a/src/transform.rs +++ b/src/transform.rs @@ -100,9 +100,32 @@ fn matrix_inverse(matrix: Array2) -> Array2 { inv } +#[derive(Deserialize)] +struct TransformS { + parameters: Vec, + dparameters: Vec, + center: Vec, + shape: Vec, +} + +impl From for Transform { + fn from(t: TransformS) -> Self { + let ndim = t.shape.len(); + Self { + parameters: t.parameters, + dparameters: t.dparameters, + center: t.center, + shape: t.shape, + ndim, + dimension: PhantomData::default(), + } + } +} + /// a struct describing the transform /// generic parameter N = # image dimensions #[derive(Clone, Debug, Deserialize, Serialize, PartialEq)] +#[serde(from = "TransformS")] pub struct Transform { /// flattened NxN rotation matrix + N translation parameters pub parameters: Vec, @@ -112,7 +135,9 @@ pub struct Transform { pub center: Vec, /// the shape of images for which this transform is meant pub shape: Vec, + #[serde(skip)] ndim: usize, + #[serde(skip)] dimension: PhantomData, } @@ -463,12 +488,12 @@ impl Transform { T: 'a + Clone + AsPrimitive, G: Into, { - let mut registration = Registration::new(fixed_mu); + let mut registration = Registration::new(fixed_mu)?; if let Some(steps) = steps { registration.set_steps(steps); } if let Some(initial_guess) = initial_guess { - registration.set_initial_guess(initial_guess); + registration.set_initial_guess(initial_guess)?; } registration.register(fixed, moving) } @@ -488,12 +513,12 @@ impl Transform { T: 'a + Clone + AsPrimitive, G: Into, { - let mut registration = Registration::new(fixed_mu); + let mut registration = Registration::new(fixed_mu)?; if let Some(steps) = steps { registration.set_steps(steps); } if let Some(initial_guess) = initial_guess { - registration.set_initial_guess(initial_guess); + registration.set_initial_guess(initial_guess)?; } registration.register_debug(fixed, moving) } @@ -571,6 +596,11 @@ impl Transform { Self::from_scaling(scaling) * self } + pub fn with_center(mut self, center: Vec) -> Self { + self.center = center; + self + } + /// read a transform from a file pub fn from_file(path: PathBuf) -> Result { let file = File::open(path)?; @@ -677,6 +707,16 @@ impl Transform { matrix } + pub fn with_parameters(mut self, parameters: Vec) -> Self { + self.parameters = parameters; + self + } + + pub fn with_dparameters(mut self, dparameters: Vec) -> Self { + self.dparameters = dparameters; + self + } + pub fn with_matrix(mut self, matrix: ArrayView2) -> Self { let mut parameters = matrix .slice(s![..self.ndim, ..self.ndim])