- some bugfixes

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);