- some bugfixes

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),
 }

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<Vec<f64>>),
@@ -452,7 +457,7 @@ where
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
 pub enum Sigma {
     None,
     Absolute(Vec<f64>),

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<D: Dimension> {
 }
 
 impl<D: Dimension> Registration<D> {
-    pub fn new<F: Into<FixedMu>>(fixed_mu: F) -> Self {
-        Self {
-            fixed_mu: fixed_mu.into(),
-            steps: None,
-            initial_guess: None,
-            optimizer: Optimizer::default(),
-            dimension: PhantomData,
+    pub fn new<F: Into<FixedMu>>(fixed_mu: F) -> Result<Self, Error> {
+        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<D: Dimension> Registration<D> {
         self
     }
 
-    pub fn with_initial_guess(mut self, initial_guess: Vec<f64>) -> Self {
-        self.initial_guess = Some(initial_guess);
-        self
+    pub fn with_initial_guess(mut self, initial_guess: Vec<f64>) -> Result<Self, 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(self)
+        }
     }
 
     pub fn with_optimizer(mut self, optimizer: Optimizer) -> Self {
@@ -198,8 +215,16 @@ impl<D: Dimension> Registration<D> {
         self.steps = Some(steps);
     }
 
-    pub fn set_initial_guess(&mut self, initial_guess: Vec<f64>) {
-        self.initial_guess = Some(initial_guess);
+    pub fn set_initial_guess(&mut self, initial_guess: Vec<f64>) -> 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<D: Dimension> Registration<D> {
         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<Transform<D>, Error>
     where
         F: AsArray<'a, T, D>,
@@ -261,7 +285,7 @@ impl<D: Dimension> Registration<D> {
             .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<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 p_var = metric.fixed_mu().extract_variable(&p);
-
             let optimization_result = Self::optimize_metric(
                 &metric,
                 &p_var,
@@ -305,10 +327,13 @@ 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);
+                p_var = result.optimal_point;
             }
         }
-        Ok(Transform::<D>::new(p, fixed.shape().to_vec()))
+        Ok(Transform::<D>::new(
+            self.fixed_mu.combine(p_var.as_slice()),
+            fixed.shape().to_vec(),
+        ))
     }
 
     fn optimize_metric<M: ObjectiveFunction<f64>>(
@@ -366,9 +391,10 @@ impl<D: Dimension> Registration<D> {
             .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<D: Dimension> Registration<D> {
                         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<D: Dimension> Registration<D> {
                         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<D: Dimension> Registration<D> {
                 .iter()
                 .all(|i| i.is_finite())
             {
+                p_var = optimization_result.optimal_point.clone();
                 p = metric
                     .fixed_mu()
                     .combine(&optimization_result.optimal_point);

src/transform.rs

@@ -100,9 +100,32 @@ fn matrix_inverse(matrix: Array2<f64>) -> Array2<f64> {
     inv
 }
 
+#[derive(Deserialize)]
+struct TransformS {
+    parameters: Vec<f64>,
+    dparameters: Vec<f64>,
+    center: Vec<f64>,
+    shape: Vec<usize>,
+}
+
+impl<D: Dimension> From<TransformS> for Transform<D> {
+    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<D: Dimension = Ix2> {
     /// flattened NxN rotation matrix + N translation parameters
     pub parameters: Vec<f64>,
@@ -112,7 +135,9 @@ pub struct Transform<D: Dimension = Ix2> {
     pub center: Vec<f64>,
     /// the shape of images for which this transform is meant
     pub shape: Vec<usize>,
+    #[serde(skip)]
     ndim: usize,
+    #[serde(skip)]
     dimension: PhantomData<D>,
 }
 
@@ -463,12 +488,12 @@ impl<D: Dimension> Transform<D> {
         T: 'a + Clone + AsPrimitive<f64>,
         G: Into<FixedMu>,
     {
-        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<D: Dimension> Transform<D> {
         T: 'a + Clone + AsPrimitive<f64>,
         G: Into<FixedMu>,
     {
-        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<D: Dimension> Transform<D> {
         Self::from_scaling(scaling) * self
     }
 
+    pub fn with_center(mut self, center: Vec<f64>) -> Self {
+        self.center = center;
+        self
+    }
+
     /// read a transform from a file
     pub fn from_file(path: PathBuf) -> Result<Self, Error> {
         let file = File::open(path)?;
@@ -677,6 +707,16 @@ impl<D: Dimension> Transform<D> {
         matrix
     }
 
+    pub fn with_parameters(mut self, parameters: Vec<f64>) -> Self {
+        self.parameters = parameters;
+        self
+    }
+
+    pub fn with_dparameters(mut self, dparameters: Vec<f64>) -> Self {
+        self.dparameters = dparameters;
+        self
+    }
+
     pub fn with_matrix(mut self, matrix: ArrayView2<f64>) -> Self {
         let mut parameters = matrix
             .slice(s![..self.ndim, ..self.ndim])