- second commit

src/register.rs

@@ -0,0 +1,333 @@
+use crate::bspline::{BSpline, BSplineTrait};
+use crate::error::Error;
+use crate::metric::{FixedMu, MattesMetric, SamplingArg, Sigma};
+use crate::transform::Transform;
+use crate::optimize::lbfgs_minimize;
+use algos::OptimizationConfig;
+use ndarray::{AsArray, Dimension};
+use num::cast::AsPrimitive;
+use num::integer::Roots;
+use std::marker::PhantomData;
+
+#[derive(Clone, Debug)]
+pub struct RegistrationStep {
+    pub sigma: Sigma,
+    pub samples: SamplingArg,
+    pub n_bins: usize,
+    pub tolerance: f64,
+    pub edge: f64,
+    pub max_iterations: usize,
+    pub learning_rate: f64,
+}
+
+impl RegistrationStep {
+    pub fn new(
+        sigma: Sigma,
+        samples: SamplingArg,
+        n_bins: usize,
+        tolerance: f64,
+        edge: f64,
+        max_iterations: usize,
+        learning_rate: f64,
+    ) -> Self {
+        Self {
+            sigma,
+            samples,
+            n_bins,
+            tolerance,
+            edge,
+            max_iterations,
+            learning_rate,
+        }
+    }
+
+    pub fn default_steps(ndim: usize, n: usize) -> Vec<Self> {
+        vec![
+            RegistrationStep::new(
+                Sigma::Relative(vec![0.5; ndim]),
+                SamplingArg::Fixed(n.sqrt().max(n / 125).max(100)),
+                3,
+                1e-4,
+                0.05,
+                100,
+                100.0,
+            ),
+            RegistrationStep::new(
+                Sigma::Relative(vec![0.25; ndim]),
+                SamplingArg::Fixed(n.sqrt().max(n / 30).max(200)),
+                6,
+                1e-6,
+                0.04,
+                100,
+                1.0,
+            ),
+            RegistrationStep::new(
+                Sigma::Absolute(vec![8.0; ndim]),
+                SamplingArg::Fixed(n.sqrt().max(n / 20).max(400)),
+                24,
+                3e-7,
+                0.03,
+                100,
+                1e-2,
+            ),
+            RegistrationStep::new(
+                Sigma::Absolute(vec![2.0; ndim]),
+                SamplingArg::Fixed(n.sqrt().max(n / 10).max(800)),
+                96,
+                1e-7,
+                0.01,
+                100,
+                1e-2,
+            ),
+            RegistrationStep::new(
+                Sigma::None,
+                SamplingArg::Fixed(n.sqrt().max(n / 5).max(1600)),
+                240,
+                1e-8,
+                0.001,
+                100,
+                1e-3,
+            ),
+        ]
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct RegistrationResult {
+    pub sigma_fixed: Option<Vec<f64>>,
+    pub sigma_moving: Option<Vec<f64>>,
+    pub n_bins: usize,
+    pub n_samples: usize,
+    pub tolerance: f64,
+    pub edge: f64,
+    pub max_iterations: usize,
+    pub learning_rate: f64,
+    pub optimal_point: Vec<f64>,
+    pub optimal_value: f64,
+    pub iterations: usize,
+    pub converged: bool,
+}
+
+pub struct Registration<D: Dimension> {
+    fixed_mu: FixedMu,
+    steps: Option<Vec<RegistrationStep>>,
+    initial_guess: Option<Vec<f64>>,
+    dimension: PhantomData<D>,
+}
+
+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,
+            dimension: PhantomData,
+        }
+    }
+
+    pub fn new_affine() -> Result<Self, Error> {
+        let ndim = if let Some(ndim) = D::NDIM {
+            ndim
+        } else {
+            return Err(Error::NumberOfDimensionsNotDefined);
+        };
+        Ok(Self {
+            fixed_mu: FixedMu::new_none(ndim),
+            steps: None,
+            initial_guess: None,
+            dimension: PhantomData,
+        })
+    }
+
+    pub fn new_translation() -> Result<Self, Error> {
+        let ndim = if let Some(ndim) = D::NDIM {
+            ndim
+        } else {
+            return Err(Error::NumberOfDimensionsNotDefined);
+        };
+        let mut fixed_mu = vec![Some(0.0); ndim * ndim];
+        for i in 0..ndim {
+            fixed_mu[i * (ndim + 1)] = Some(1.0);
+            fixed_mu[ndim * ndim + i] = None;
+        }
+        Ok(Self {
+            fixed_mu: fixed_mu.into(),
+            steps: None,
+            initial_guess: None,
+            dimension: PhantomData,
+        })
+    }
+
+    pub fn with_steps(mut self, steps: Vec<RegistrationStep>) -> Self {
+        self.steps = Some(steps);
+        self
+    }
+
+    pub fn with_initial_guess(mut self, initial_guess: Vec<f64>) -> Self {
+        self.initial_guess = Some(initial_guess);
+        self
+    }
+
+    pub fn set_steps(&mut self, steps: Vec<RegistrationStep>) {
+        self.steps = Some(steps);
+    }
+
+    pub fn set_initial_guess(&mut self, initial_guess: Vec<f64>) {
+        self.initial_guess = Some(initial_guess);
+    }
+
+    /// find the transform which transforms moving into fixed and return the results of each
+    /// optimization step
+    pub fn register<'a, F, M, T>(&self, fixed: F, moving: M) -> Result<Transform<D>, Error>
+    where
+        F: AsArray<'a, T, D>,
+        M: AsArray<'a, T, D>,
+        T: 'a + Clone + AsPrimitive<f64>,
+    {
+        let fixed = fixed.into().mapv(|i| i.as_());
+        let moving = moving.into().mapv(|i| i.as_());
+        let ndim = fixed.ndim();
+        let n = fixed.len();
+        let steps = self
+            .steps
+            .as_ref()
+            .cloned()
+            .unwrap_or_else(|| RegistrationStep::default_steps(ndim, n));
+        let mut p = self
+            .initial_guess
+            .as_ref()
+            .cloned()
+            .unwrap_or_else(|| Transform::<D>::default().parameters);
+        for RegistrationStep {
+            sigma,
+            samples,
+            n_bins,
+            tolerance,
+            edge,
+            max_iterations,
+            learning_rate,
+        } 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;
+            }
+            let bf = BSpline::<0, _>::new(f.view());
+            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 optimization_config = OptimizationConfig {
+                max_iterations,
+                tolerance,
+                learning_rate,
+            };
+            let optimization_result = lbfgs_minimize(
+                &metric,
+                metric.fixed_mu().extract_variable(&p).as_slice(),
+                &optimization_config,
+            );
+            if optimization_result
+                .optimal_point
+                .iter()
+                .all(|i| i.is_finite())
+            {
+                p = metric
+                    .fixed_mu()
+                    .combine(&optimization_result.optimal_point);
+            }
+        }
+        Ok(Transform::<D>::new(p, fixed.shape().to_vec()))
+    }
+
+    /// find the transform which transforms moving into fixed and return the results of each
+    /// optimization step
+    pub fn register_debug<'a, F, M, T>(
+        &self,
+        fixed: F,
+        moving: M,
+    ) -> Result<(Transform<D>, Vec<RegistrationResult>), Error>
+    where
+        F: AsArray<'a, T, D>,
+        M: AsArray<'a, T, D>,
+        T: 'a + Clone + AsPrimitive<f64>,
+    {
+        let fixed = fixed.into().mapv(|i| i.as_());
+        let moving = moving.into().mapv(|i| i.as_());
+        let ndim = fixed.ndim();
+        let n = fixed.len();
+        let steps = self
+            .steps
+            .as_ref()
+            .cloned()
+            .unwrap_or_else(|| RegistrationStep::default_steps(ndim, n));
+        let mut p = self
+            .initial_guess
+            .as_ref()
+            .cloned()
+            .unwrap_or_else(|| Transform::<D>::default().parameters);
+        let mut registration_results = Vec::new();
+        for RegistrationStep {
+            sigma,
+            samples,
+            n_bins,
+            tolerance,
+            edge,
+            max_iterations,
+            learning_rate,
+        } 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;
+            }
+            let bf = BSpline::<0, _>::new(f.view());
+            let bm = BSpline::<3, _>::new(m.view());
+            let n_samples = match &samples {
+                SamplingArg::Fixed(n) => *n,
+                SamplingArg::Random(n) => *n,
+                SamplingArg::FixedAt(n) => n.len(),
+            };
+            let metric = MattesMetric::new(bf, bm, samples, n_bins, edge)?
+                .with_fixed_mu(self.fixed_mu.clone());
+            let optimization_config = OptimizationConfig {
+                max_iterations,
+                tolerance,
+                learning_rate,
+            };
+            let optimization_result = lbfgs_minimize(
+                &metric,
+                metric.fixed_mu().extract_variable(&p).as_slice(),
+                &optimization_config,
+            );
+            if optimization_result
+                .optimal_point
+                .iter()
+                .all(|i| i.is_finite())
+            {
+                p = metric
+                    .fixed_mu()
+                    .combine(&optimization_result.optimal_point);
+            }
+            registration_results.push(RegistrationResult {
+                sigma_fixed: sigma.sigma(f.shape()),
+                sigma_moving: sigma.sigma(m.shape()),
+                n_bins,
+                n_samples,
+                tolerance,
+                edge,
+                max_iterations,
+                learning_rate,
+                optimal_point: optimization_result.optimal_point,
+                optimal_value: optimization_result.optimal_value,
+                iterations: optimization_result.iterations,
+                converged: optimization_result.converged,
+            });
+        }
+        Ok((
+            Transform::<D>::new(p, fixed.shape().to_vec()),
+            registration_results,
+        ))
+    }
+}