894 lines
30 KiB
Rust
894 lines
30 KiB
Rust
use crate::bspline::{
|
||
BSpline, BSplineTrait, cubic_bspline, square_bspline, square_bspline_integral,
|
||
};
|
||
use crate::error::Error;
|
||
use crate::filter::gaussian_smooth;
|
||
use crate::transform::transform_point;
|
||
use algos::ObjectiveFunction;
|
||
use itertools::Itertools;
|
||
use ndarray::{Array, Array1, Array2, Array3, AsArray, Axis, Dimension, s};
|
||
use num::cast::AsPrimitive;
|
||
use rand::RngExt;
|
||
use rayon::iter::{IntoParallelIterator, ParallelIterator};
|
||
use serde::{Deserialize, Serialize};
|
||
use std::cell::RefCell;
|
||
use std::ops::Deref;
|
||
|
||
#[derive(Default)]
|
||
struct IntMut {
|
||
mu: Vec<f64>,
|
||
metric: f64,
|
||
derivative: Vec<f64>,
|
||
}
|
||
|
||
pub struct FixedMu {
|
||
mu: Vec<Option<f64>>,
|
||
len_variable: usize,
|
||
}
|
||
|
||
impl Deref for FixedMu {
|
||
type Target = Vec<Option<f64>>;
|
||
|
||
fn deref(&self) -> &Self::Target {
|
||
&self.mu
|
||
}
|
||
}
|
||
|
||
impl From<Vec<Option<f64>>> for FixedMu {
|
||
fn from(mu: Vec<Option<f64>>) -> Self {
|
||
let len_variable = mu.iter().filter(|v| v.is_none()).count();
|
||
FixedMu { mu, len_variable }
|
||
}
|
||
}
|
||
|
||
impl FixedMu {
|
||
pub fn new_none(n: usize) -> Self {
|
||
Self {
|
||
mu: vec![None; n * n + n],
|
||
len_variable: n * n + n,
|
||
}
|
||
}
|
||
|
||
/// combine mu and fixed_mu into complete parameter set
|
||
pub fn combine(&self, mu: &[f64]) -> Vec<f64> {
|
||
debug_assert_eq!(mu.len(), self.len_variable);
|
||
let mut mu_iter = mu.iter();
|
||
self.iter()
|
||
.map(|i| {
|
||
if let Some(i) = i {
|
||
*i
|
||
} else {
|
||
*mu_iter.next().unwrap()
|
||
}
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
/// combine the dparameters of the variable parameters into a complete
|
||
/// parameter set, with no uncertainty (0) on the fixed parameters
|
||
pub fn combine_dparameters(&self, dparameters: &[f64]) -> Vec<f64> {
|
||
debug_assert_eq!(dparameters.len(), self.len_variable);
|
||
let mut dparameters_iter = dparameters.iter();
|
||
self.iter()
|
||
.map(|i| {
|
||
if i.is_some() {
|
||
0.0
|
||
} else {
|
||
*dparameters_iter.next().unwrap()
|
||
}
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
/// extract mu from complete parameter set
|
||
pub fn extract_variable(&self, mu: &[f64]) -> Vec<f64> {
|
||
self.iter()
|
||
.zip_eq(mu)
|
||
.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
|
||
/// n_bins >= 2
|
||
#[expect(clippy::too_many_arguments)]
|
||
fn parzen(
|
||
alpha: f64,
|
||
dalpha: &[f64],
|
||
fixed: &[f64],
|
||
moving: &[f64],
|
||
dmoving: &[Vec<f64>],
|
||
weight: &[f64],
|
||
dweight: &[Vec<f64>],
|
||
first_bin: f64,
|
||
last_bin: f64,
|
||
n_bins: usize,
|
||
) -> (Array2<f64>, Array3<f64>) {
|
||
let ndim = dalpha.len();
|
||
let width = (last_bin - first_bin) / ((n_bins - 1) as f64);
|
||
let v = 1.0 / width;
|
||
// add two bins left and right
|
||
let first_bin = first_bin - 2.0 * width;
|
||
let last_bin = last_bin + 2.0 * width;
|
||
let n_bins = n_bins + 4;
|
||
let n_bins_i = n_bins as isize;
|
||
let mut hist = Array2::zeros([n_bins, n_bins]);
|
||
let mut grad_b = Array3::zeros([n_bins, n_bins, ndim]);
|
||
let left = first_bin - width / 2.0;
|
||
let bins = Array1::linspace(first_bin, last_bin, n_bins);
|
||
let mut nf;
|
||
let mut nm;
|
||
let mut uf;
|
||
let mut um;
|
||
let mut kf;
|
||
let mut km;
|
||
let mut csf;
|
||
for ((f, m), (dm, (w, dw))) in fixed
|
||
.iter()
|
||
.zip_eq(moving)
|
||
.zip_eq(dmoving.iter().zip_eq(weight.iter().zip_eq(dweight)))
|
||
{
|
||
nf = (v * (f - left)) as isize;
|
||
nm = (v * (m - left)) as isize;
|
||
for if_ in (nf - 2)..=(nf + 2) {
|
||
if (0 <= if_) && (if_ < n_bins_i) {
|
||
uf = if_ as usize;
|
||
kf = (f - bins[uf]) * v;
|
||
csf = cubic_bspline(kf);
|
||
for im in (nm - 2)..=(nm + 2) {
|
||
if (0 <= im) && (im < n_bins_i) {
|
||
um = im as usize;
|
||
km = (m - bins[um]) * v;
|
||
hist[[uf, um]] += w * csf * cubic_bspline(km);
|
||
for ((dmi, dwi), d) in
|
||
dm.iter()
|
||
.zip_eq(dw)
|
||
.zip_eq(grad_b.slice_mut(s![uf, um, ..]))
|
||
{
|
||
*d += w
|
||
* csf
|
||
* (square_bspline(km + 0.5) - square_bspline(km - 0.5))
|
||
* v
|
||
* dmi
|
||
+ dwi * csf * cubic_bspline(km);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
let hist = hist / alpha;
|
||
let mut hist_bc = hist
|
||
.broadcast([dalpha.len(), hist.shape()[0], hist.shape()[1]])
|
||
.unwrap();
|
||
hist_bc.permute_axes([1, 2, 0]);
|
||
let dhist = (grad_b - (&hist_bc * &Array1::from_iter(dalpha))) / alpha;
|
||
(hist, dhist)
|
||
}
|
||
|
||
enum Sampling {
|
||
Fixed(Vec<Vec<f64>>),
|
||
Random((usize, Vec<f64>)),
|
||
}
|
||
|
||
impl Sampling {
|
||
fn fixed(index: Vec<Vec<f64>>) -> Self {
|
||
Self::Fixed(index)
|
||
}
|
||
|
||
fn random(n_samples: usize, shape: Vec<f64>) -> Self {
|
||
Self::Random((n_samples, shape))
|
||
}
|
||
|
||
fn index(&self) -> Vec<Vec<f64>> {
|
||
match self {
|
||
Self::Fixed(index) => index.clone(),
|
||
Self::Random((n_samples, shape)) => rand::rng()
|
||
.random_iter::<f64>()
|
||
.take(n_samples * shape.len())
|
||
.chunks(shape.len())
|
||
.into_iter()
|
||
.map(|c| c.zip_eq(shape.iter()).map(|(i, s)| i * s - 0.5).collect())
|
||
.collect(),
|
||
}
|
||
}
|
||
}
|
||
|
||
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
|
||
pub enum SamplingArg {
|
||
Fixed(usize),
|
||
FixedAt(Vec<Vec<f64>>),
|
||
Random(usize),
|
||
}
|
||
|
||
/// (dI / dT(x)) * (dT(x) / dmu)
|
||
fn image_jacobian(point: &[f64], center: &[f64], image_derivative: &[f64]) -> Vec<f64> {
|
||
image_derivative
|
||
.iter()
|
||
.flat_map(|d| {
|
||
point
|
||
.iter()
|
||
.zip_eq(center.iter())
|
||
.map(|(p, c)| *d * (p - c))
|
||
})
|
||
.chain(image_derivative.iter().cloned())
|
||
.collect()
|
||
}
|
||
|
||
/// Local Gauss-Newton statistics of the intensity fit
|
||
/// `fixed(x) = moving(T(x)) + noise` at a parameter point.
|
||
///
|
||
/// `hessian` is the weighted normal matrix `JᵀWJ` and `variance` is the
|
||
/// weighted residual variance `Σ w r² / (Σ w − q − 1)`, with `q` the number of
|
||
/// variable parameters. The extra degree of freedom accounts for the weighted
|
||
/// mean intensity offset that is removed from the residuals. The parameter
|
||
/// covariance of the fit is `variance * hessian⁻¹`.
|
||
#[derive(Clone, Debug)]
|
||
pub struct ResidualFit {
|
||
/// weighted normal matrix `JᵀWJ`
|
||
pub hessian: Vec<Vec<f64>>,
|
||
/// residual variance (mean squared residual, corrected for the dof)
|
||
pub variance: f64,
|
||
/// number of sample points that contributed
|
||
pub n_samples: usize,
|
||
}
|
||
|
||
pub struct MattesMetric<D>
|
||
where
|
||
D: Dimension,
|
||
{
|
||
shape: Vec<f64>,
|
||
center: Vec<f64>,
|
||
fixed: BSpline<0, D>,
|
||
moving: BSpline<3, D>,
|
||
fixed_mu: FixedMu,
|
||
minmax: [f64; 2],
|
||
sampling: Sampling,
|
||
n_bins: usize,
|
||
/// memory for (point, value, gradient) combination, refcell gives interior mutability
|
||
metric: RefCell<IntMut>,
|
||
edge: f64,
|
||
}
|
||
|
||
impl<D> MattesMetric<D>
|
||
where
|
||
D: Dimension,
|
||
{
|
||
pub fn new(
|
||
fixed: BSpline<0, D>,
|
||
moving: BSpline<3, D>,
|
||
sampling: SamplingArg,
|
||
n_bins: usize,
|
||
edge: f64,
|
||
) -> Result<Self, Error> {
|
||
let shape = fixed.shape().iter().map(|&i| i as f64).collect::<Vec<_>>();
|
||
let center = shape.iter().map(|i| (i - 1.0) / 2.0).collect();
|
||
let (min, max) = fixed
|
||
.iter()
|
||
.chain(moving.iter())
|
||
.minmax()
|
||
.into_option()
|
||
.expect("m and f cannot be empty");
|
||
let minmax = [*min, *max];
|
||
let ndim = shape.len();
|
||
let sampling = match sampling {
|
||
SamplingArg::Fixed(n_samples) => {
|
||
Sampling::fixed(Sampling::random(n_samples, shape.clone()).index())
|
||
}
|
||
SamplingArg::FixedAt(points) => Sampling::fixed(points),
|
||
SamplingArg::Random(n_samples) => Sampling::random(n_samples, shape.clone()),
|
||
};
|
||
|
||
Ok(Self {
|
||
shape,
|
||
center,
|
||
fixed,
|
||
moving,
|
||
fixed_mu: FixedMu::new_none(ndim),
|
||
minmax,
|
||
sampling,
|
||
n_bins,
|
||
metric: RefCell::default(),
|
||
edge,
|
||
})
|
||
}
|
||
|
||
pub fn new_from_arrays<'a, T, A>(
|
||
fixed: A,
|
||
moving: A,
|
||
sampling: SamplingArg,
|
||
n_bins: usize,
|
||
edge: f64,
|
||
) -> Result<Self, Error>
|
||
where
|
||
A: AsArray<'a, T, D>,
|
||
T: 'a + Clone + AsPrimitive<f64>,
|
||
{
|
||
Self::new(
|
||
BSpline::<0, _>::new(fixed),
|
||
BSpline::<3, _>::new(moving),
|
||
sampling,
|
||
n_bins,
|
||
edge,
|
||
)
|
||
}
|
||
|
||
pub fn with_fixed_mu<F: Into<FixedMu>>(mut self, fixed_mu: F) -> Self {
|
||
self.fixed_mu = fixed_mu.into();
|
||
self
|
||
}
|
||
|
||
pub fn fixed_mu(&self) -> &FixedMu {
|
||
&self.fixed_mu
|
||
}
|
||
|
||
pub fn eval(&self, mu: &[f64]) -> Result<(), Error> {
|
||
assert!(mu.iter().all(|i| i.is_finite()));
|
||
if self.metric.borrow().mu != mu {
|
||
let complete_mu = self.fixed_mu.combine(mu);
|
||
let idx = self.sampling.index();
|
||
let shape = self.shape.clone();
|
||
let shape_f = self.shape.iter().map(|s| s - 0.5).collect::<Vec<_>>();
|
||
let center = self.center.as_slice();
|
||
let fixed = &self.fixed;
|
||
let mem_f = fixed.get_mem();
|
||
let moving = &self.moving;
|
||
let mem_m = moving.get_mem();
|
||
let fixed_mu = &self.fixed_mu;
|
||
let e = 1.0 / self.edge;
|
||
let (f, ((m, j), (w, dw))) = idx
|
||
.into_par_iter()
|
||
.map_init(
|
||
|| (mem_f.clone(), mem_m.clone()),
|
||
|(mem_f, mem_m), u| {
|
||
let v = transform_point(&u, center, &complete_mu);
|
||
if shape_f
|
||
.iter()
|
||
.zip_eq(v.iter())
|
||
.all(|(s, x)| (-0.5 <= *x) && (x <= s))
|
||
{
|
||
let f = fixed.evaluate_at_continuous_index(&u, mem_f)?;
|
||
let (m, d) = moving
|
||
.evaluate_value_and_derivative_at_continuous_index(&v, mem_m)?;
|
||
let j = fixed_mu.extract_variable(&image_jacobian(&u, center, &d));
|
||
let mut k;
|
||
let mut w = 1.0;
|
||
let mut ws = Vec::with_capacity(shape.len());
|
||
let mut dw = Vec::with_capacity(shape.len());
|
||
// the metric wouldn't be smooth if we'd drop a point abruptly,
|
||
// so apply a 0 <= weight <= 1 to points within 5% of the edge
|
||
for (s, x) in shape.iter().zip_eq(&v) {
|
||
if e * (x + 0.5) < *s {
|
||
k = 3.0 * e * (x + 0.5) / s - 1.5;
|
||
let b = square_bspline_integral(k);
|
||
w *= b;
|
||
ws.push(b);
|
||
dw.push(square_bspline(k) * 3.0 * e / s);
|
||
} else if e * (x + 0.5) > (e - 1.0) * s {
|
||
k = 3.0 * e * (s - x - 0.5) / s - 1.5;
|
||
let b = square_bspline_integral(k);
|
||
w *= b;
|
||
ws.push(b);
|
||
dw.push(-square_bspline(k) * 3.0 * e / s);
|
||
} else {
|
||
ws.push(1.0);
|
||
dw.push(0.0);
|
||
}
|
||
}
|
||
for (dwi, wi) in dw.iter_mut().zip(ws) {
|
||
*dwi *= w / wi;
|
||
}
|
||
let dw = fixed_mu.extract_variable(&image_jacobian(&u, center, &dw));
|
||
Ok(Some((f, ((m, j), (w, dw)))))
|
||
} else {
|
||
Ok(None)
|
||
}
|
||
},
|
||
)
|
||
.filter_map(|i| i.transpose())
|
||
.collect::<Result<
|
||
(
|
||
Vec<f64>,
|
||
((Vec<f64>, Vec<Vec<f64>>), (Vec<f64>, Vec<Vec<f64>>)),
|
||
),
|
||
Error,
|
||
>>()?;
|
||
let alpha = w.iter().sum::<f64>();
|
||
let mut dalpha = vec![0.0; mu.len()];
|
||
for i in dw.iter() {
|
||
for (j, a) in i.iter().zip_eq(dalpha.iter_mut()) {
|
||
*a += j;
|
||
}
|
||
}
|
||
|
||
let (jpdf, d_jpdf_m) = parzen(
|
||
alpha,
|
||
&dalpha,
|
||
&f,
|
||
&m,
|
||
j.as_slice(),
|
||
&w,
|
||
&dw,
|
||
self.minmax[0],
|
||
self.minmax[1],
|
||
self.n_bins,
|
||
);
|
||
|
||
let pdf_f = jpdf.axis_iter(Axis(0)).map(|i| i.sum()).collect::<Vec<_>>();
|
||
let pdf_m = jpdf.axis_iter(Axis(1)).map(|i| i.sum()).collect::<Vec<_>>();
|
||
let d_pdf_m = d_jpdf_m
|
||
.axis_iter(Axis(1))
|
||
.zip(pdf_m.iter())
|
||
.map(|(di, i)| di.sum() / i)
|
||
.collect::<Vec<_>>();
|
||
|
||
let mut metric = 0.0;
|
||
let mut dmetric = vec![0.0; self.fixed_mu.len_variable];
|
||
let mut n;
|
||
let mut m;
|
||
let mut ik;
|
||
let mut ln;
|
||
|
||
for ((row, d_row), k) in jpdf
|
||
.axis_iter(Axis(0))
|
||
.zip_eq(d_jpdf_m.axis_iter(Axis(0)))
|
||
.zip_eq(pdf_f)
|
||
{
|
||
debug_assert_eq!(row.len(), d_row.shape()[0]);
|
||
debug_assert_eq!(row.len(), pdf_m.len());
|
||
|
||
for ((&j, dj), (i, di)) in row
|
||
.into_iter()
|
||
.zip_eq(d_row.axis_iter(Axis(0)))
|
||
.zip_eq(pdf_m.iter().zip(d_pdf_m.iter()))
|
||
{
|
||
ik = i * k;
|
||
// check for non-zero bin contribution
|
||
if (j > 1e-16) && (ik > 1e-16) {
|
||
ln = (j / (ik)).ln();
|
||
n = 1.0 + ln;
|
||
m = di * j;
|
||
metric -= j * ln;
|
||
for (d, g) in dmetric.iter_mut().zip_eq(dj) {
|
||
*d += m - n * g;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
debug_assert!(mu.iter().all(|i| i.is_finite()));
|
||
debug_assert!(alpha.is_finite());
|
||
debug_assert!(metric.is_finite());
|
||
debug_assert!(dmetric.iter().all(|i| i.is_finite()));
|
||
|
||
self.metric.replace(IntMut {
|
||
mu: mu.to_vec(),
|
||
metric,
|
||
derivative: dmetric.clone(),
|
||
});
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Compute the Gauss-Newton normal matrix and the residual variance of the
|
||
/// intensity fit at `point`, using the same sample points and edge
|
||
/// weighting as [`Self::eval`].
|
||
///
|
||
/// Residuals are `fixed(x) - moving(T(x))`, so a well-registered image pair
|
||
/// gives a small variance. Returns `None` when interpolation fails or when
|
||
/// too few samples contribute to estimate the residual variance.
|
||
pub fn residual_fit(&self, point: &[f64]) -> Option<ResidualFit> {
|
||
let complete_mu = self.fixed_mu.combine(point);
|
||
let idx = self.sampling.index();
|
||
let shape = self.shape.as_slice();
|
||
let shape_f = self.shape.iter().map(|s| s - 0.5).collect::<Vec<_>>();
|
||
let center = self.center.as_slice();
|
||
let e = 1.0 / self.edge;
|
||
let n_params = point.len();
|
||
let mut hessian = vec![vec![0.0; n_params]; n_params];
|
||
let mut sum_w = 0.0;
|
||
let mut sum_wr = 0.0;
|
||
let mut sum_wr2 = 0.0;
|
||
let mut n_samples = 0;
|
||
let mut mem_f = self.fixed.get_mem();
|
||
let mut mem_m = self.moving.get_mem();
|
||
for u in idx {
|
||
let v = transform_point(&u, center, &complete_mu);
|
||
if !shape_f
|
||
.iter()
|
||
.zip_eq(v.iter())
|
||
.all(|(s, x)| (-0.5 <= *x) && (x <= s))
|
||
{
|
||
continue;
|
||
}
|
||
let f = self
|
||
.fixed
|
||
.evaluate_at_continuous_index(&u, &mut mem_f)
|
||
.ok()?;
|
||
let (m, d) = self
|
||
.moving
|
||
.evaluate_value_and_derivative_at_continuous_index(&v, &mut mem_m)
|
||
.ok()?;
|
||
// same smooth edge weighting as in `eval`
|
||
let mut w = 1.0;
|
||
for (s, x) in shape.iter().zip_eq(&v) {
|
||
if e * (x + 0.5) < *s {
|
||
let k = 3.0 * e * (x + 0.5) / s - 1.5;
|
||
w *= square_bspline_integral(k);
|
||
} else if e * (x + 0.5) > (e - 1.0) * s {
|
||
let k = 3.0 * e * (s - x - 0.5) / s - 1.5;
|
||
w *= square_bspline_integral(k);
|
||
}
|
||
}
|
||
if w <= 0.0 || !w.is_finite() {
|
||
continue;
|
||
}
|
||
let r = f - m;
|
||
sum_w += w;
|
||
sum_wr += w * r;
|
||
sum_wr2 += w * r * r;
|
||
n_samples += 1;
|
||
let j = self
|
||
.fixed_mu
|
||
.extract_variable(&image_jacobian(&u, center, &d));
|
||
for (a, ja) in j.iter().enumerate() {
|
||
for (b, jb) in j.iter().enumerate() {
|
||
hessian[a][b] += w * ja * jb;
|
||
}
|
||
}
|
||
}
|
||
// the residual variance needs at least one degree of freedom
|
||
if n_samples == 0 || sum_w <= n_params as f64 + 1.0 {
|
||
return None;
|
||
}
|
||
// subtract the weighted mean intensity mismatch (e.g. a constant
|
||
// offset between the images) before estimating the variance
|
||
let mean_r = sum_wr / sum_w;
|
||
let sse = (sum_wr2 - sum_w * mean_r * mean_r).max(0.0);
|
||
// q parameters plus the mean intensity offset that was removed
|
||
let variance = sse / (sum_w - n_params as f64 - 1.0);
|
||
if !variance.is_finite() {
|
||
return None;
|
||
}
|
||
Some(ResidualFit {
|
||
hessian,
|
||
variance,
|
||
n_samples,
|
||
})
|
||
}
|
||
}
|
||
|
||
impl<D> ObjectiveFunction<f64> for MattesMetric<D>
|
||
where
|
||
D: Dimension,
|
||
{
|
||
fn evaluate(&self, point: &[f64]) -> f64 {
|
||
self.eval(point).unwrap();
|
||
self.metric.borrow().metric
|
||
}
|
||
|
||
fn gradient(&self, point: &[f64]) -> Option<Vec<f64>> {
|
||
self.eval(point).unwrap();
|
||
Some(self.metric.borrow().derivative.clone())
|
||
}
|
||
|
||
/// The hessian is computed numerically: central finite differences of
|
||
/// `gradient` with respect to each parameter (2 * n full evaluations).
|
||
/// `MattesMetric` only accumulates first-order terms, so an analytic hessian
|
||
/// is not available. The result is symmetrized, which also damps the small
|
||
/// asymmetry introduced by finite differencing.
|
||
fn hessian(&self, point: &[f64]) -> Option<Vec<Vec<f64>>> {
|
||
let n = point.len();
|
||
let mut hessian = vec![vec![0.0; n]; n];
|
||
let mut p = point.to_vec();
|
||
for j in 0..n {
|
||
let h = 1e-5 * f64::max(1.0, point[j].abs());
|
||
p[j] += h;
|
||
self.eval(&p).unwrap();
|
||
let g_plus = self.metric.borrow().derivative.clone();
|
||
p[j] -= 2.0 * h;
|
||
self.eval(&p).unwrap();
|
||
let g_minus = self.metric.borrow().derivative.clone();
|
||
p[j] += h;
|
||
for i in 0..n {
|
||
hessian[i][j] = (g_plus[i] - g_minus[i]) / (2.0 * h);
|
||
}
|
||
}
|
||
#[expect(clippy::needless_range_loop)]
|
||
for i in 0..n {
|
||
for j in (i + 1)..n {
|
||
let w = (hessian[i][j] + hessian[j][i]) / 2.0;
|
||
hessian[i][j] = w;
|
||
hessian[j][i] = w;
|
||
}
|
||
}
|
||
Some(hessian)
|
||
}
|
||
}
|
||
|
||
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
|
||
pub enum Sigma {
|
||
None,
|
||
Absolute(Vec<f64>),
|
||
Relative(Vec<f64>),
|
||
}
|
||
|
||
impl Sigma {
|
||
pub fn smooth<'a, A, D>(&self, array: A) -> Result<Array<f64, D>, Error>
|
||
where
|
||
A: AsArray<'a, f64, D>,
|
||
D: Dimension,
|
||
{
|
||
match self {
|
||
Self::None => Ok(array.into().to_owned()),
|
||
Self::Absolute(sigma) => gaussian_smooth(array, sigma.as_slice()),
|
||
Self::Relative(sigma) => {
|
||
let array = array.into();
|
||
let sigma = sigma
|
||
.iter()
|
||
.zip(array.shape())
|
||
.map(|(i, j)| i * (*j as f64))
|
||
.collect::<Vec<_>>();
|
||
gaussian_smooth(array, sigma.as_slice())
|
||
}
|
||
}
|
||
}
|
||
|
||
pub fn sigma(&self, shape: &[usize]) -> Option<Vec<f64>> {
|
||
match self {
|
||
Self::None => None,
|
||
Self::Absolute(sigma) => Some(sigma.to_vec()),
|
||
Self::Relative(sigma) => Some(
|
||
sigma
|
||
.iter()
|
||
.zip(shape)
|
||
.map(|(i, j)| i * (*j as f64))
|
||
.collect::<Vec<_>>(),
|
||
),
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use crate::filter::gaussian_smooth;
|
||
use algos::OptimizationConfig;
|
||
use algos::optimization::{bfgs_minimize, gradient_descent_minimize};
|
||
use ndarray::array;
|
||
|
||
#[test]
|
||
fn derivative() -> Result<(), Box<dyn std::error::Error>> {
|
||
let im = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0];
|
||
let b = BSpline::<3, _>::new(im.view());
|
||
let mut mem = b.get_mem();
|
||
let mut value = Vec::new();
|
||
let mut derivative = Vec::new();
|
||
let x = (0..100 * im.len())
|
||
.map(|i| i as f64 / 100.0)
|
||
.collect::<Vec<_>>();
|
||
for i in x.iter() {
|
||
let (v, d) = b.evaluate_value_and_derivative_at_continuous_index(&[*i], &mut mem)?;
|
||
value.push(v);
|
||
derivative.push(d[0]);
|
||
}
|
||
println!("x = {:?}\nv = {:?}\nd = {:?}", x, value, derivative);
|
||
Ok(())
|
||
}
|
||
|
||
#[test]
|
||
fn metric2() -> Result<(), Box<dyn std::error::Error>> {
|
||
let im_a = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0];
|
||
let im_b = array![1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0];
|
||
let a = gaussian_smooth(im_a.view(), &[2.0])?;
|
||
let b = gaussian_smooth(im_b.view(), &[2.0])?;
|
||
|
||
let points = vec![
|
||
vec![3.0],
|
||
vec![4.0],
|
||
vec![5.0],
|
||
vec![6.0],
|
||
vec![7.0],
|
||
vec![8.0],
|
||
];
|
||
|
||
let m = MattesMetric::new_from_arrays(
|
||
a.view(),
|
||
b.view(),
|
||
SamplingArg::FixedAt(points),
|
||
5,
|
||
0.05,
|
||
)?
|
||
.with_fixed_mu(vec![Some(1.0), None]);
|
||
// let mu = vec![-3.52, -3.51, -3.50, -3.49, -3.48];
|
||
let mu = vec![-3.51, -3.49];
|
||
// let mu = vec![-2.0, -3.0];
|
||
let mut v = Vec::new();
|
||
let mut d = Vec::new();
|
||
for x in mu.iter() {
|
||
v.push(m.evaluate(&[*x]));
|
||
d.push(m.gradient(&[*x]).unwrap()[0]);
|
||
}
|
||
println!("a = {:?}\nb = {:?}", a.to_vec(), b.to_vec());
|
||
println!("mu = {:?}\nv = {:?}\nd = {:?}", mu, v.to_vec(), d.to_vec());
|
||
Ok(())
|
||
}
|
||
|
||
#[test]
|
||
fn optimize() -> Result<(), Box<dyn std::error::Error>> {
|
||
let im_a = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0];
|
||
let im_b = array![1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0];
|
||
let a = gaussian_smooth(im_a.view(), &[4.0])?;
|
||
let b = gaussian_smooth(im_b.view(), &[4.0])?;
|
||
|
||
let m =
|
||
MattesMetric::new_from_arrays(a.view(), b.view(), SamplingArg::Random(100), 5, 0.05)?
|
||
.with_fixed_mu(vec![Some(1.0), None]);
|
||
let optimization_config = OptimizationConfig {
|
||
max_iterations: 1000,
|
||
tolerance: 1e-6,
|
||
learning_rate: 0.001,
|
||
};
|
||
|
||
let r_grad = gradient_descent_minimize(&m, &[-0.0], &optimization_config);
|
||
let r_bfgs = bfgs_minimize(&m, &[-0.0], &optimization_config);
|
||
println!("r_grad = {:?}", r_grad);
|
||
println!("r_bfgs = {:?}", r_bfgs);
|
||
Ok(())
|
||
}
|
||
|
||
#[test]
|
||
fn residual_fit_1d() -> Result<(), Box<dyn std::error::Error>> {
|
||
let a = array![1.0, 2.0, 3.0, 2.0, 1.0, 2.0, 3.0, 4.0, 2.0, 0.0, 1.0];
|
||
// b is a shifted right by 3 (padded on the left with the first value)
|
||
let b = array![1.0, 1.0, 1.0, 1.0, 2.0, 3.0, 2.0, 1.0, 2.0, 3.0, 4.0];
|
||
let points = (0..11).map(|i| vec![i as f64]).collect::<Vec<_>>();
|
||
let m = MattesMetric::new_from_arrays(
|
||
a.view(),
|
||
b.view(),
|
||
SamplingArg::FixedAt(points),
|
||
16,
|
||
0.05,
|
||
)?
|
||
.with_fixed_mu(vec![Some(1.0), None]);
|
||
|
||
// fixed(x) = moving(x + 3): no residual
|
||
let fit = m.residual_fit(&[3.0]).expect("residual fit");
|
||
assert!(fit.n_samples > 0);
|
||
assert!(fit.hessian[0][0] > 0.0);
|
||
assert!(fit.variance < 1e-12);
|
||
|
||
// away from the optimum the residual variance must grow
|
||
let worse = m.residual_fit(&[0.0]).expect("residual fit");
|
||
assert!(worse.variance > fit.variance);
|
||
Ok(())
|
||
}
|
||
|
||
#[test]
|
||
fn hessian_1d() -> Result<(), Box<dyn std::error::Error>> {
|
||
let im_a = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0];
|
||
let im_b = array![1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0];
|
||
let a = gaussian_smooth(im_a.view(), &[2.0])?;
|
||
let b = gaussian_smooth(im_b.view(), &[2.0])?;
|
||
|
||
let points = vec![
|
||
vec![3.0],
|
||
vec![4.0],
|
||
vec![5.0],
|
||
vec![6.0],
|
||
vec![7.0],
|
||
vec![8.0],
|
||
];
|
||
|
||
let m = MattesMetric::new_from_arrays(
|
||
a.view(),
|
||
b.view(),
|
||
SamplingArg::FixedAt(points),
|
||
5,
|
||
0.05,
|
||
)?
|
||
.with_fixed_mu(vec![Some(1.0), None]);
|
||
|
||
let p = -3.49;
|
||
let h = m.hessian(&[p]).unwrap();
|
||
assert_eq!(h.len(), 1);
|
||
assert_eq!(h[0].len(), 1);
|
||
|
||
// cross-check against the second central difference of `evaluate`
|
||
let step = 1e-4;
|
||
let fd = (m.evaluate(&[p + step]) - 2.0 * m.evaluate(&[p]) + m.evaluate(&[p - step]))
|
||
/ (step * step);
|
||
println!("hessian = {:?}\nfinite diff = {:?}", h[0][0], fd);
|
||
assert!((h[0][0] - fd).abs() < 1e-3 + 1e-2 * fd.abs());
|
||
Ok(())
|
||
}
|
||
|
||
#[test]
|
||
fn hessian_2d() -> Result<(), Box<dyn std::error::Error>> {
|
||
let a = Array2::from_shape_fn((12, 12), |(i, j)| i as f64 + 0.7 * j as f64);
|
||
let b = Array2::from_shape_fn((12, 12), |(i, j)| 0.5 * i as f64 + 24.0 - 0.7 * j as f64);
|
||
let a = gaussian_smooth(a.view(), &[2.0, 2.0])?;
|
||
let b = gaussian_smooth(b.view(), &[2.0, 2.0])?;
|
||
|
||
let points = (0..12)
|
||
.flat_map(|i| (0..12).map(move |j| vec![i as f64, j as f64]))
|
||
.step_by(2)
|
||
.collect::<Vec<_>>();
|
||
|
||
let m = MattesMetric::new_from_arrays(
|
||
a.view(),
|
||
b.view(),
|
||
SamplingArg::FixedAt(points),
|
||
7,
|
||
0.05,
|
||
)?;
|
||
let n = m.fixed_mu().len_variable();
|
||
assert_eq!(n, 2 * 2 + 2);
|
||
|
||
let p = vec![0.5, -0.3, 0.2, 0.7, -0.4, 0.3];
|
||
let h = m.hessian(&p).unwrap();
|
||
assert_eq!(h.len(), n);
|
||
assert!(h.iter().all(|row| row.len() == n));
|
||
// symmetrized by construction
|
||
for (i, row) in h.iter().enumerate() {
|
||
for (j, v) in row.iter().enumerate() {
|
||
assert_eq!(*v, h[j][i]);
|
||
}
|
||
}
|
||
|
||
// cross-check each diagonal entry against the 2nd central difference of `evaluate`
|
||
let step = 1e-4;
|
||
for j in 0..n {
|
||
let mut pp = p.clone();
|
||
let mut pm = p.clone();
|
||
pp[j] += step;
|
||
pm[j] -= step;
|
||
let fd = (m.evaluate(&pp) - 2.0 * m.evaluate(&p) + m.evaluate(&pm)) / (step * step);
|
||
println!("diagonal[{j}] = {:?}\nfinite diff = {:?}", h[j][j], fd);
|
||
assert!(
|
||
(h[j][j] - fd).abs() < 1e-3 + 1e-2 * fd.abs(),
|
||
"diagonal entry {j}"
|
||
);
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
#[test]
|
||
fn test_square_bspline_integral() -> Result<(), Box<dyn std::error::Error>> {
|
||
let x = Array1::linspace(-2.0, 2.0, 500);
|
||
let y = x.mapv(square_bspline);
|
||
let z = x.mapv(square_bspline_integral);
|
||
let mut w = vec![0.0];
|
||
let mut j = 0.0;
|
||
let dx = (x[x.len() - 1] - x[0]) / (x.len() - 1) as f64;
|
||
for i in y.windows([2]) {
|
||
j += dx * (i[0] + i[1]) / 2.0;
|
||
w.push(j);
|
||
}
|
||
let w = Array1::from(w);
|
||
assert!((w - z).abs().iter().all(|i| *i < 1e-5));
|
||
Ok(())
|
||
}
|
||
|
||
#[test]
|
||
fn test_cubic_bspline_derivative() -> Result<(), Box<dyn std::error::Error>> {
|
||
let x = Array1::linspace(-2.0, 2.0, 500);
|
||
let y = x.mapv(cubic_bspline);
|
||
let z = x.mapv(|i| square_bspline(i + 0.5) - square_bspline(i - 0.5));
|
||
let mut w = vec![0.0];
|
||
let mut j = 0.0;
|
||
let dx = (x[x.len() - 1] - x[0]) / (x.len() - 1) as f64;
|
||
for i in z.windows([2]) {
|
||
j += dx * (i[0] + i[1]) / 2.0;
|
||
w.push(j);
|
||
}
|
||
let w = Array1::from(w);
|
||
assert!((w - y).abs().iter().all(|i| *i < 1e-4));
|
||
Ok(())
|
||
}
|
||
}
|