- second commit

This commit is contained in:
w.pomp
2026-07-24 11:13:43 +02:00
parent 738b4fc8eb
commit 1cf672a275
11 changed files with 4275 additions and 0 deletions
+1105
View File
File diff suppressed because it is too large Load Diff
+17
View File
@@ -0,0 +1,17 @@
use thiserror::Error;
#[derive(Debug, Error)]
pub enum Error {
#[error(transparent)]
IO(#[from] std::io::Error),
#[error(transparent)]
SerdeYAML(#[from] serde_yaml::Error),
#[error(transparent)]
ShapeError(#[from] ndarray::ShapeError),
#[error(transparent)]
LinAlg(#[from] ndarray_linalg::error::LinalgError),
#[error(transparent)]
NpyError(#[from] ndarray_npy::WriteNpzError),
#[error("number of dimensions is not defined")]
NumberOfDimensionsNotDefined,
}
+277
View File
@@ -0,0 +1,277 @@
use crate::error::Error;
use itertools::Itertools;
use ndarray::{Array, Array1, ArrayD, AsArray, Axis, Dimension, RemoveAxis, concatenate, s};
use ndrustfft::{FftHandler, Normalization, ndfft_par, ndifft_par};
use num::Complex;
use num::traits::FloatConst;
use std::ops::MulAssign;
/// Fourier transform
pub fn fft<'a, A, D>(array: A) -> Result<Array<Complex<f64>, D>, Error>
where
A: AsArray<'a, f64, D>,
D: Dimension,
{
let mut input = array.into().mapv(|i| Complex::new(i, 0.0));
let shape = input.shape().to_vec();
let mut tmp = ArrayD::zeros(shape.clone()).into_dimensionality()?;
for (i, s) in shape.iter().enumerate() {
let handler = FftHandler::new(*s).normalization(Normalization::None);
ndfft_par(&input.view(), &mut tmp.view_mut(), &handler, i);
std::mem::swap(&mut input, &mut tmp);
}
Ok(input)
}
/// inverse Fourier transform
pub fn ifft<'a, A, D>(array: A) -> Result<Array<Complex<f64>, D>, Error>
where
A: AsArray<'a, Complex<f64>, D>,
D: Dimension,
{
let mut input = array.into().to_owned();
let shape = input.shape().to_vec();
let mut tmp = ArrayD::zeros(shape.clone()).into_dimensionality()?;
for (i, s) in shape.iter().enumerate() {
let handler = FftHandler::new(*s).normalization(Normalization::None);
ndifft_par(&input.view(), &mut tmp.view_mut(), &handler, i);
std::mem::swap(&mut input, &mut tmp);
}
Ok(input / Complex::from(shape.iter().product::<usize>() as f64))
}
pub fn fft_freq(size: usize) -> Array1<f64> {
// let s = size as isize;
// let h = s / 2;
// (-h..h).map(|i| (i % s - s) as f64 / s as f64).collect()
let val = 1.0 / size as f64;
let mut results = Array1::zeros(size);
let n = (size - 1) / 2 + 1;
let p1 = Array1::range(0.0, n as f64, 1.0);
results.slice_mut(s![..n]).assign(&p1);
let p2 = Array1::range(-((size / 2) as f64), 0.0, 1.0);
results.slice_mut(s![n..]).assign(&p2);
results * val
}
pub fn fft_shift<'a, A, T, D>(x: A) -> Array<T, D>
where
A: AsArray<'a, T, D>,
T: 'a + Clone,
D: Dimension + RemoveAxis,
{
let mut x = x.into().to_owned();
let shift = x.shape().iter().map(|i| i / 2).collect::<Vec<_>>();
for (i, s) in shift.iter().enumerate() {
let (a, b) = x.view().split_at(Axis(i), *s);
x = concatenate(Axis(i), &[b, a]).unwrap();
}
x
}
pub fn gaussian(x: &[f64], mu: f64, sigma: f64) -> Vec<f64> {
let a = 2.0 * sigma.powi(2);
let b = (a * f64::PI()).sqrt();
x.iter()
.map(|i| (-(i - mu).powi(2) / a).exp() / b)
.collect()
}
/// Gaussian kernel in frequency space
pub fn gaussian_kernel<D>(shape: &[usize], sigma: &[f64]) -> Result<Array<Complex<f64>, D>, Error>
where
D: Dimension,
{
let mut g = Array::ones(shape).into_dimensionality()?;
for (i, (s, t)) in shape.iter().zip_eq(sigma.iter()).enumerate() {
let a = 0.5 / t.powi(2);
let f = (-(f64::PI() * fft_freq(*s)).powi(2) / a)
.exp()
.mapv(Complex::from);
for mut lane in g.lanes_mut(Axis(i)) {
lane.mul_assign(&f);
}
}
Ok(g)
}
/// smooth an array using a Gaussian kernel
pub fn gaussian_smooth<'a, A, D>(array: A, sigma: &[f64]) -> Result<Array<f64, D>, Error>
where
A: AsArray<'a, f64, D>,
D: Dimension,
{
let array = array.into().to_owned();
let shape = array.shape();
let kernel = gaussian_kernel::<D>(shape, sigma)?;
Ok(ifft((kernel * fft(array.view())?).view())?.mapv(|i| i.re))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::julia_image;
use ndarray::{Ix1, Ix2, array};
use tiffwrite::IJTiffFile;
#[test]
fn smooth() -> Result<(), Box<dyn std::error::Error>> {
let 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 b = gaussian_smooth(a.view(), &[1.4])?;
let c = array![
0.13775053, 0.26443112, 0.45036765, 0.67622086, 0.93678676, 1.06888618, 0.93678676,
0.67622086, 0.45036765, 0.26443112, 0.13775053
];
debug_assert!(b.iter().zip_eq(c.iter()).all(|(x, y)| (x - y).abs() < 1e-8));
Ok(())
}
#[test]
fn smooth2() -> Result<(), Box<dyn std::error::Error>> {
let im_a = julia_image(
&[60, 80],
&[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
&[29.5, 39.5],
&[-0.8, 0.156],
);
let im_b = gaussian_smooth(im_a.mapv(|i| i as f64).view(), &[4.0, 4.0])?;
let mut t = IJTiffFile::new(std::env::home_dir().unwrap().join("tmp/julia.tif"))?;
t.save(im_a.mapv(|i| i as u32).view(), 0, 0, 0)?;
t.save(im_b.mapv(|i| i as u32).view(), 1, 0, 0)?;
Ok(())
}
#[test]
fn gaussian_kernel_test2() -> Result<(), Box<dyn std::error::Error>> {
let k = gaussian_kernel::<Ix2>(&[60, 80], &[16.0, 16.0])?;
ndarray_npy::write_npy("/home/wim/tmp/kernel.npy", &k)?;
Ok(())
}
#[test]
fn fft_test() -> Result<(), Box<dyn std::error::Error>> {
let x = Array1::linspace(0.0, 1.0, 11).to_vec();
let y = gaussian(x.as_slice(), 0.4, 0.1);
let j = fft(Array1::from_vec(y).view())?;
let k = array![
Complex::new(9.99998512536899, 0.0),
Complex::new(-5.562906924752932, -6.419930454175756),
Complex::new(-0.7410757333589753, 5.154238692693016),
Complex::new(1.9379781953109205, -1.2454762622580575),
Complex::new(-0.7086874511902606, -0.20810344788799306),
Complex::new(0.08206001372382088, 0.1796510861631586),
Complex::new(0.08206001372382088, -0.1796510861631586),
Complex::new(-0.7086874511902606, 0.20810344788799306),
Complex::new(1.9379781953109205, 1.2454762622580575),
Complex::new(-0.7410757333589753, -5.154238692693016),
Complex::new(-5.562906924752932, 6.419930454175756),
];
debug_assert!(
j.iter()
.zip_eq(k.iter())
.all(|(x, y)| (x.re - y.re) < 1e-8 && (x.im - y.im) < 1e-8)
);
Ok(())
}
#[test]
fn ifft_test() -> Result<(), Box<dyn std::error::Error>> {
let x = Array1::linspace(0.0, 1.0, 11).to_vec();
let y = gaussian(x.as_slice(), 0.4, 0.1)
.into_iter()
.map(Complex::from)
.collect::<Vec<_>>();
let j = ifft(Array1::from_vec(y).view())?;
let k = array![
Complex::new(0.9090895568517264, 0.0),
Complex::new(-0.5057188113411757, 0.5836300412887051),
Complex::new(-0.06737052121445229, -0.4685671538811833),
Complex::new(0.17617983593735642, 0.1132251147507325),
Complex::new(-0.06442613192638733, 0.018918495262544823),
Complex::new(0.00746000124762008, -0.01633191692392351),
Complex::new(0.00746000124762008, 0.01633191692392351),
Complex::new(-0.06442613192638733, -0.018918495262544823),
Complex::new(0.17617983593735642, -0.1132251147507325),
Complex::new(-0.06737052121445229, 0.4685671538811833),
Complex::new(-0.5057188113411757, -0.5836300412887051),
];
debug_assert!(
j.into_iter()
.zip_eq(k.into_iter())
.all(|(x, y)| ((x.re - y.re).abs() < 1e-8) && ((x.im - y.im).abs() < 1e-8))
);
Ok(())
}
#[test]
fn fft_freq_test() -> Result<(), Box<dyn std::error::Error>> {
let f = fft_freq(11);
let g = vec![
0.0,
0.09090909090909091,
0.18181818181818182,
0.2727272727272727,
0.36363636363636365,
0.4545454545454546,
-0.4545454545454546,
-0.36363636363636365,
-0.2727272727272727,
-0.18181818181818182,
-0.09090909090909091,
];
debug_assert!(f.iter().zip_eq(g.iter()).all(|(x, y)| (x - y).abs() < 1e-8));
Ok(())
}
#[test]
fn gaussian_kernel_test() -> Result<(), Box<dyn std::error::Error>> {
let x = gaussian_kernel::<Ix1>(&[11], &[1.5])?.mapv(|i| i.re);
println!("x = {:?}", x.to_vec());
let y = array![
1.00000000e+00,
6.92774033e-01,
2.30338431e-01,
3.67556755e-02,
2.81491714e-03,
1.03464181e-04,
1.03464181e-04,
2.81491714e-03,
3.67556755e-02,
2.30338431e-01,
6.92774033e-01
];
debug_assert!(y.iter().zip_eq(x.iter()).all(|(x, y)| (x - y).abs() < 1e-8));
Ok(())
}
#[test]
fn fft_shift_test() -> Result<(), Box<dyn std::error::Error>> {
let x = Array1::from_iter(0..10);
let y = fft_shift(x.view());
let z = array![5, 6, 7, 8, 9, 0, 1, 2, 3, 4];
debug_assert_eq!(y, z);
Ok(())
}
#[test]
fn gaussian_test() -> Result<(), Box<dyn std::error::Error>> {
let x = Array1::linspace(0.0, 1.0, 11).to_vec();
let y = gaussian(x.as_slice(), 0.4, 0.1);
let z = vec![
1.33830226e-03,
4.43184841e-02,
5.39909665e-01,
2.41970725e+00,
3.98942280e+00,
2.41970725e+00,
5.39909665e-01,
4.43184841e-02,
1.33830226e-03,
1.48671951e-05,
6.07588285e-08,
];
debug_assert!(y.iter().zip_eq(z.iter()).all(|(x, y)| (x - y).abs() < 1e-8));
Ok(())
}
}
+179
View File
@@ -0,0 +1,179 @@
use rayon::iter::ParallelIterator;
pub mod bspline;
pub mod error;
pub mod filter;
pub mod metric;
pub mod par_indexed_iter;
pub mod register;
pub mod transform;
mod optimize;
use ndarray::prelude::*;
use num::Complex;
use thiserror::Error;
use crate::par_indexed_iter::ParallelIndexedIterMut;
use crate::transform::transform_point;
use error::Error;
/// An example of generating julia fractals, for testing purposes.
/// parameters: 2x2 flattened rotation matrix + xy translation
/// center: center of rotation
pub fn julia_image(
shape: &[usize; 2],
parameters: &[f64; 6],
center: &[f64; 2],
c: &[f64; 2],
) -> Array2<u8> {
let c = Complex::new(c[0], c[1]);
let scaley = 3.0 / shape[0] as f64;
let scalex = 3.0 / shape[1] as f64;
let mut im = Array2::<u8>::zeros([shape[0], shape[1]]);
im.par_indexed_iter_mut().for_each(|(i, x)| {
let cv = transform_point(i.as_slice(), center, parameters);
let mut z = Complex::new(cv[1] * scalex - 1.5, cv[0] * scaley - 1.5);
let mut i = 0;
while i < 255 && z.norm() <= 2.0 {
z = z * z + c;
i += 1;
}
*x = i;
});
im
}
#[cfg(test)]
mod tests {
use super::*;
use crate::transform::Transform;
use tempfile::NamedTempFile;
use tiffwrite::IJTiffFile;
#[test]
fn serialization() -> Result<(), Error> {
let file = NamedTempFile::new()?;
let t = Transform::<Ix2>::new_with_center(
vec![1.2, 0.3, -0.4, 0.9, 10.2, -9.5],
vec![59.5, 49.5],
vec![120, 100],
);
t.to_file(file.path().to_path_buf())?;
let s = Transform::from_file(file.path().to_path_buf())?;
assert_eq!(s, t);
Ok(())
}
#[test]
fn transform_point() -> Result<(), Box<dyn std::error::Error>> {
let transform = Transform::<Ix2>::new_with_center(
vec![1.0, 0.0, 0.0, 1.0, -120.0, 10.0],
vec![299.5, 399.5],
vec![600, 800],
);
let p = transform.transform_point(&[300, 400]);
assert_eq!(p, [180.0, 410.0]);
Ok(())
}
#[test]
fn interpbs_f64() -> Result<(), Box<dyn std::error::Error>> {
let shape = [1200, 1600];
let transform_j = [1.0, 0.0, 0.0, 1.0, 0.0, 0.0];
let transform_k = [1.0, 0.0, 0.0, 1.0, 0.0, 0.0];
let center: [f64; 2] = shape
.iter()
.map(|&s| (s - 1) as f64 / 2.0)
.collect::<Vec<_>>()
.try_into()
.unwrap();
let transform =
Transform::<Ix2>::new(transform_j.to_vec(), shape.to_vec()).with_rotation(1.0);
let transform_j = transform.parameters.clone().try_into().unwrap();
let c = [-0.8, 0.156];
let j = julia_image(&shape, &transform_j, &center, &c).mapv(|x| x as f64);
let k = julia_image(&shape, &transform_k, &center, &c).mapv(|x| x as f64);
let n = transform.interpolate::<3, _, _>(k.view())?;
let sj = j.iter().sum::<f64>();
// let sk = k.iter().sum::<f64>();
let sn = n.iter().sum::<f64>();
// println!("sj: {}, sk: {}, sn: {}", sj, sk, sn);
// let mut tiff = tiffwrite::IJTiffFile::new("interpbs_f64_0.tif")?;
// tiff.save(j.view(), 0, 0, 0)?;
// tiff.save(k.view(), 1, 0, 0)?;
// tiff.save(n.view(), 2, 0, 0)?;
let s = (sj.ln() - sn.ln()).abs();
let d = (j.ln() - n.ln())
.powi(2)
.iter()
.filter(|i| i.is_finite())
.sum::<f64>()
.sqrt();
// println!("s: {}, d: {}", s, d);
assert!(s < 1e-2);
assert!(2000.0 * d <= (shape[0] * shape[1]) as f64);
Ok(())
}
#[test]
fn interpbs_f64_par() -> Result<(), Box<dyn std::error::Error>> {
let shape = [1200, 1600];
let transform_j = [1.0, 0.0, 0.0, 1.0, 0.0, 0.0];
let transform_k = [1.0, 0.0, 0.0, 1.0, 0.0, 0.0];
let center: [f64; 2] = shape
.iter()
.map(|&s| (s - 1) as f64 / 2.0)
.collect::<Vec<_>>()
.try_into()
.unwrap();
let transform =
Transform::<Ix2>::new(transform_j.to_vec(), shape.to_vec()).with_rotation(1.0);
let transform_j = transform.parameters.clone().try_into().unwrap();
let c = [-0.8, 0.156];
let j = julia_image(&shape, &transform_j, &center, &c).mapv(|x| x as f64);
let k = julia_image(&shape, &transform_k, &center, &c).mapv(|x| x as f64);
let n = transform.interpolate_par::<3, _, _>(k.view())?;
let sj = j.iter().sum::<f64>();
let sn = n.iter().sum::<f64>();
let s = (sj.ln() - sn.ln()).abs();
let d = (j.ln() - n.ln())
.powi(2)
.iter()
.filter(|i| i.is_finite())
.sum::<f64>()
.sqrt();
assert!(s < 1e-2);
assert!(2000.0 * d <= (shape[0] * shape[1]) as f64);
Ok(())
}
#[test]
fn interpolate_par_unsafe() -> Result<(), Box<dyn std::error::Error>> {
let image_a = julia_image(
&[8000, 6000],
&[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
&[3999.5, 2999.5],
&[-0.8, 0.156],
);
let transform = Transform::from_rotation(1.0, &[3999.5, 2999.5]);
let image_b = transform.interpolate_par::<3, _, _>(image_a.view())?;
assert_ne!(image_b[[4000, 3000]], 0.0);
Ok(())
}
#[test]
fn julia() -> Result<(), Box<dyn std::error::Error>> {
let image_a = julia_image(
&[60, 80],
&[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
&[29.5, 39.5],
&[-0.8, 0.156],
);
let mut t = IJTiffFile::new(std::env::home_dir().unwrap().join("tmp/julia.tif"))?;
t.save(image_a.view(), 0, 0, 0)?;
Ok(())
}
}
+1067
View File
File diff suppressed because it is too large Load Diff
+167
View File
@@ -0,0 +1,167 @@
use std::collections::VecDeque;
use std::fmt::Debug;
use algos::{ObjectiveFunction, OptimizationConfig, OptimizationResult};
use num::Float;
pub fn lbfgs_minimize<T, F>(
f: &F,
initial_point: &[T],
config: &OptimizationConfig<T>,
) -> OptimizationResult<T>
where
T: Float + Debug,
F: ObjectiveFunction<T>,
{
const M: usize = 10; // Number of corrections to store
let n = initial_point.len();
let mut current_point = initial_point.to_vec();
let mut iterations = 0;
let mut converged = false;
// Storage for the last M corrections
let mut s_list: VecDeque<Vec<T>> = VecDeque::with_capacity(M);
let mut y_list: VecDeque<Vec<T>> = VecDeque::with_capacity(M);
let mut rho_list: VecDeque<T> = VecDeque::with_capacity(M);
// Get initial gradient
let mut gradient = match f.gradient(&current_point) {
Some(g) => g,
None => {
return OptimizationResult {
optimal_point: current_point.clone(),
optimal_value: f.evaluate(&current_point),
iterations: 0,
converged: false,
};
}
};
while iterations < config.max_iterations {
// Check for convergence
let gradient_norm = gradient
.iter()
.fold(T::zero(), |acc, &x| acc + x * x)
.sqrt();
if gradient_norm < config.tolerance {
converged = true;
break;
}
// Compute search direction using L-BFGS two-loop recursion
let mut q = gradient.clone();
let mut alpha_list = Vec::with_capacity(s_list.len());
// First loop
for i in (0..s_list.len()).rev() {
let alpha = rho_list[i]
* s_list[i]
.iter()
.zip(q.iter())
.fold(T::zero(), |acc, (&s, &q)| acc + s * q);
alpha_list.push(alpha);
for (q_j, y_j) in q.iter_mut().zip(y_list[i].iter()) {
*q_j = *q_j - alpha * *y_j;
}
}
// Scale the initial Hessian approximation
let mut r = if !s_list.is_empty() {
let i = s_list.len() - 1;
let yy = y_list[i].iter().fold(T::zero(), |acc, &y| acc + y * y);
let ys = y_list[i]
.iter()
.zip(s_list[i].iter())
.fold(T::zero(), |acc, (&y, &s)| acc + y * s);
q.iter_mut().for_each(|r_j| *r_j = *r_j * (ys / yy));
q
} else {
q.iter_mut()
.for_each(|r_j| *r_j = *r_j * config.learning_rate);
q
};
// Second loop
for i in 0..s_list.len() {
let beta = rho_list[i]
* y_list[i]
.iter()
.zip(r.iter())
.fold(T::zero(), |acc, (&y, &r)| acc + y * r);
let alpha = alpha_list[s_list.len() - 1 - i];
for (r_j, s_j) in r.iter_mut().zip(s_list[i].iter()) {
*r_j = *r_j + (alpha - beta) * *s_j;
}
}
// r now contains the search direction
let direction: Vec<T> = r.iter().map(|&x| -x).collect();
// Line search to find step size
let mut alpha = T::one();
let mut new_point = vec![T::zero(); n];
let current_value = f.evaluate(&current_point);
// Simple backtracking line search
for _ in 0..20 {
for i in 0..n {
new_point[i] = current_point[i] + alpha * direction[i];
}
let new_value = f.evaluate(&new_point);
if new_value < current_value {
break;
}
alpha = alpha * T::from(0.5).unwrap();
}
// Get new gradient
let new_gradient = match f.gradient(&new_point) {
Some(g) => g,
None => break,
};
// Update the correction vectors
let s = new_point
.iter()
.zip(current_point.iter())
.map(|(&x_new, &x_old)| x_new - x_old)
.collect::<Vec<T>>();
let y = new_gradient
.iter()
.zip(gradient.iter())
.map(|(&g_new, &g_old)| g_new - g_old)
.collect::<Vec<T>>();
let ys = y
.iter()
.zip(s.iter())
.fold(T::zero(), |acc, (&y_i, &s_i)| acc + y_i * s_i);
if ys == T::zero() {
break;
}
let rho = T::one() / ys;
if s_list.len() == M {
s_list.pop_front();
y_list.pop_front();
rho_list.pop_front();
}
s_list.push_back(s);
y_list.push_back(y);
rho_list.push_back(rho);
// Update for next iteration
current_point = new_point;
gradient = new_gradient;
iterations += 1;
}
OptimizationResult {
optimal_point: current_point.clone(),
optimal_value: f.evaluate(&current_point),
iterations,
converged,
}
}
+407
View File
@@ -0,0 +1,407 @@
use ndarray::{Array, ArrayViewMut, Dimension};
use rayon::iter::plumbing::{
Consumer, Folder, Producer, ProducerCallback, UnindexedConsumer, UnindexedProducer, bridge,
bridge_unindexed,
};
use rayon::prelude::{IndexedParallelIterator, ParallelIterator};
use std::marker::PhantomData;
#[derive(Clone, Debug)]
struct BaseIndexIter {
start: usize,
end: usize,
cum_shape: Vec<usize>,
}
impl BaseIndexIter {
fn new(shape: Vec<usize>) -> Self {
let mut cum_shape = Vec::with_capacity(shape.len());
let mut c = 1;
for s in shape.iter().rev() {
cum_shape.push(c);
c *= s;
}
Self {
start: 0,
end: shape.iter().product::<usize>(),
cum_shape: cum_shape.into_iter().rev().collect(),
}
}
}
/// An iterator over array indices.
pub struct IndexIter(BaseIndexIter);
/// A parallel iterator over array indices.
pub struct ParIndexIter(BaseIndexIter);
impl IndexIter {
/// Create a new IndexIter using the shape of an array.
pub fn new(shape: Vec<usize>) -> Self {
Self(BaseIndexIter::new(shape))
}
}
impl ParIndexIter {
/// Create a new ParIndexIter using the shape of an array.
pub fn new(shape: Vec<usize>) -> Self {
Self(BaseIndexIter::new(shape))
}
}
impl ParallelIterator for ParIndexIter {
type Item = Vec<usize>;
fn drive_unindexed<C>(self, consumer: C) -> C::Result
where
C: UnindexedConsumer<Self::Item>,
{
bridge_unindexed(self, consumer)
}
}
impl IndexedParallelIterator for ParIndexIter {
fn len(&self) -> usize {
self.0.end - self.0.start
}
fn drive<C: Consumer<Self::Item>>(self, consumer: C) -> C::Result {
bridge(self, consumer)
}
fn with_producer<CB: ProducerCallback<Self::Item>>(self, callback: CB) -> CB::Output {
callback.callback(self)
}
}
impl UnindexedProducer for ParIndexIter {
type Item = Vec<usize>;
fn split(self) -> (Self, Option<Self>) {
let length = self.0.end - self.0.start;
if length > 1 {
let (a, b) = self.split_at(length / 2);
(a, Some(b))
} else {
(self, None)
}
}
fn fold_with<F>(self, folder: F) -> F
where
F: Folder<Self::Item>,
{
folder.consume_iter(IndexIter(self.0))
}
}
impl Producer for ParIndexIter {
type Item = Vec<usize>;
type IntoIter = IndexIter;
fn into_iter(self) -> Self::IntoIter {
IndexIter(self.0)
}
fn split_at(self, index: usize) -> (Self, Self) {
(
Self(BaseIndexIter {
start: self.0.start,
end: self.0.start + index,
cum_shape: self.0.cum_shape.clone(),
}),
Self(BaseIndexIter {
start: self.0.start + index,
end: self.0.end,
cum_shape: self.0.cum_shape.clone(),
}),
)
}
}
impl Iterator for IndexIter {
type Item = Vec<usize>;
fn next(&mut self) -> Option<Self::Item> {
if self.0.end <= self.0.start {
None
} else if self.0.cum_shape.len() == 1 {
let i = self.0.start;
self.0.start += 1;
Some(vec![i])
} else {
let mut n = self.0.start;
let mut i = Vec::new();
for c in self.0.cum_shape.iter() {
i.push(n / c);
n = n.saturating_sub(n - n % c);
}
self.0.start += 1;
Some(i)
}
}
}
impl DoubleEndedIterator for IndexIter {
fn next_back(&mut self) -> Option<Self::Item> {
if self.0.end <= self.0.start {
None
} else if self.0.cum_shape.len() == 1 {
self.0.end -= 1;
Some(vec![self.0.end])
} else {
self.0.end -= 1;
let mut n = self.0.end;
let mut i = Vec::new();
for c in self.0.cum_shape.iter() {
i.push(n / c);
n = n.saturating_sub(n - n % c);
}
Some(i)
}
}
}
impl ExactSizeIterator for IndexIter {}
/// A Sync and Send pointer, this should be safe because the pointer is only used with an offset,
/// and each offset is only used once, so no race conditions can occur.
#[derive(Debug, Clone)]
struct Ptr<'a, T> {
ptr: *mut T,
phantom_data: PhantomData<&'a T>,
}
unsafe impl<T> Send for Ptr<'_, T> {}
unsafe impl<T> Sync for Ptr<'_, T> {}
impl<'a, T> Ptr<'a, T> {
fn at(&self, i: usize) -> &'a mut T {
unsafe { &mut *self.ptr.add(i) }
}
}
struct BaseIndexedIterMut<'a, T> {
start: usize,
end: usize,
cum_shape: Vec<usize>,
ptr: Ptr<'a, T>,
}
impl<'a, T> BaseIndexedIterMut<'a, T> {
fn new<D: Dimension>(mut view: ArrayViewMut<'a, T, D>) -> Self {
let shape = view.shape();
let mut cum_shape = Vec::with_capacity(shape.len());
let mut c = 1;
for s in shape.iter().rev() {
cum_shape.push(c);
c *= s;
}
Self {
start: 0,
end: shape.iter().product::<usize>(),
cum_shape: cum_shape.into_iter().rev().collect(),
ptr: Ptr {
ptr: view.as_mut_ptr(),
phantom_data: PhantomData,
},
}
}
}
/// An parallel iterator over array indices and mutable values.
pub struct IndexedIterMut<'a, T>(BaseIndexedIterMut<'a, T>);
/// An iterator over array indices and mutable values.
pub struct ParIndexedIterMut<'a, T>(BaseIndexedIterMut<'a, T>);
impl<'a, T> IndexedIterMut<'a, T> {
/// Create a new IndexedIterMut for an array.
pub fn new<D: Dimension>(view: ArrayViewMut<'a, T, D>) -> Self {
Self(BaseIndexedIterMut::new(view))
}
}
impl<'a, T> ParIndexedIterMut<'a, T> {
/// Create a new ParIndexedIterMut for an array.
pub fn new<D: Dimension>(view: ArrayViewMut<'a, T, D>) -> Self {
Self(BaseIndexedIterMut::new(view))
}
}
impl<'a, T> ParallelIterator for ParIndexedIterMut<'a, T>
where
T: Clone + Send,
&'a T: Send,
{
type Item = (Vec<usize>, &'a mut T);
fn drive_unindexed<C>(self, consumer: C) -> C::Result
where
C: UnindexedConsumer<Self::Item>,
{
bridge_unindexed(self, consumer)
}
}
impl<'a, T> IndexedParallelIterator for ParIndexedIterMut<'a, T>
where
T: Clone + Send,
&'a T: Send,
{
fn len(&self) -> usize {
self.0.end - self.0.start
}
fn drive<C: Consumer<Self::Item>>(self, consumer: C) -> C::Result {
bridge(self, consumer)
}
fn with_producer<CB: ProducerCallback<Self::Item>>(self, callback: CB) -> CB::Output {
callback.callback(self)
}
}
impl<'a, T> Producer for ParIndexedIterMut<'a, T>
where
T: Clone,
{
type Item = (Vec<usize>, &'a mut T);
type IntoIter = IndexedIterMut<'a, T>;
fn into_iter(self) -> Self::IntoIter {
IndexedIterMut(self.0)
}
fn split_at(self, index: usize) -> (Self, Self) {
(
Self(BaseIndexedIterMut {
start: self.0.start,
end: self.0.start + index,
cum_shape: self.0.cum_shape.clone(),
ptr: self.0.ptr.clone(),
}),
Self(BaseIndexedIterMut {
start: self.0.start + index,
end: self.0.end,
cum_shape: self.0.cum_shape.clone(),
ptr: self.0.ptr.clone(),
}),
)
}
}
impl<'a, T> UnindexedProducer for ParIndexedIterMut<'a, T>
where
T: Clone + Send,
&'a T: Send,
{
type Item = (Vec<usize>, &'a mut T);
fn split(self) -> (Self, Option<Self>) {
let length = self.0.end - self.0.start;
if length > 1 {
let (a, b) = self.split_at(length / 2);
(a, Some(b))
} else {
(self, None)
}
}
fn fold_with<F>(self, folder: F) -> F
where
F: Folder<Self::Item>,
{
folder.consume_iter(IndexedIterMut(self.0))
}
}
impl<'a, T> Iterator for IndexedIterMut<'a, T> {
type Item = (Vec<usize>, &'a mut T);
fn next(&mut self) -> Option<Self::Item> {
if self.0.end <= self.0.start {
None
} else if self.0.cum_shape.len() == 1 {
self.0.start += 1;
let j = self.0.start;
Some((vec![j], self.0.ptr.at(j)))
} else {
let mut n = self.0.start;
let mut i = Vec::new();
for c in self.0.cum_shape.iter() {
i.push(n / c);
n = n.saturating_sub(n - n % c);
}
let j = self.0.start;
self.0.start += 1;
Some((i, self.0.ptr.at(j)))
}
}
}
impl<'a, T> DoubleEndedIterator for IndexedIterMut<'a, T> {
fn next_back(&mut self) -> Option<Self::Item> {
if self.0.end <= self.0.start {
None
} else if self.0.cum_shape.len() == 1 {
self.0.end -= 1;
Some((vec![self.0.end], self.0.ptr.at(self.0.end)))
} else {
self.0.end -= 1;
let mut n = self.0.end;
let mut i = Vec::new();
for c in self.0.cum_shape.iter() {
i.push(n / c);
n = n.saturating_sub(n - n % c);
}
Some((i, self.0.ptr.at(self.0.end)))
}
}
}
impl<'a, T> ExactSizeIterator for IndexedIterMut<'a, T> {}
/// A trait to use arrays as ParallelIndexedIterMut.
pub trait ParallelIndexedIterMut {
type IterMut;
fn par_indexed_iter_mut(self) -> Self::IterMut;
}
impl<'a, T, D> ParallelIndexedIterMut for ArrayViewMut<'a, T, D>
where
D: Dimension,
{
type IterMut = ParIndexedIterMut<'a, T>;
fn par_indexed_iter_mut(self) -> Self::IterMut {
Self::IterMut::new(self)
}
}
impl<'a, T, D> ParallelIndexedIterMut for &'a mut Array<T, D>
where
D: Dimension,
{
type IterMut = ParIndexedIterMut<'a, T>;
fn par_indexed_iter_mut(self) -> Self::IterMut {
Self::IterMut::new(self.view_mut())
}
}
#[cfg(test)]
mod tests {
use super::IndexIter;
use ndarray::Array3;
#[test]
fn par_indexed_iter() -> Result<(), Box<dyn std::error::Error>> {
let shape = [5, 7, 3];
let i = IndexIter::new(shape.to_vec()).collect::<Vec<_>>();
let j = Array3::<usize>::zeros(shape)
.indexed_iter()
.map(|(x, _)| vec![x.0, x.1, x.2])
.collect::<Vec<_>>();
assert_eq!(i.len(), shape.iter().product::<usize>());
assert_eq!(i, j);
Ok(())
}
}
+333
View File
@@ -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,
))
}
}
+618
View File
@@ -0,0 +1,618 @@
use crate::bspline::{BSpline, BSplineTrait};
use crate::error::Error;
use crate::metric::FixedMu;
use crate::register::{Registration, RegistrationResult, RegistrationStep};
use itertools::Itertools;
use ndarray::{Array, Array2, ArrayD, AsArray, Dimension, Ix2, s};
use ndarray_linalg::Inverse;
use num::cast::AsPrimitive;
use serde::{Deserialize, Serialize};
use serde_yaml::{from_reader, to_writer};
use std::fs::File;
use std::marker::PhantomData;
use std::ops::Mul;
use std::path::PathBuf;
/// get coordinates resulting from transforming input coordinates, coordinate must have N
/// columns: x, y, z, ...
#[inline]
pub fn transform_point<'a, T>(point: &'a [T], center: &'a [f64], parameters: &'a [f64]) -> Vec<f64>
where
T: AsPrimitive<f64>,
{
// [a, b, c] x = a*x + b*y + c
// [d, e, f] y = d*x + e*y + f
// [0, 0, 1] 1 = 1
debug_assert_eq!(center.len(), point.len());
debug_assert_eq!((center.len() + 1) * center.len(), parameters.len());
let m = center.len();
let n = m * m;
let point = point
.iter()
.zip_eq(center.iter())
.map(|(p, c)| p.as_() - c)
.collect::<Vec<_>>();
parameters
.iter()
.take(n)
.chunks(m)
.into_iter()
.zip_eq(parameters.iter().skip(n))
.zip_eq(center.iter())
.map(|((r, t), c)| r.zip_eq(point.iter()).map(|(i, p)| i * p).sum::<f64>() + t + c)
.collect()
}
/// a struct describing the transform
/// generic parameter N = # image dimensions
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
pub struct Transform<D: Dimension = Ix2> {
/// flattened NxN rotation matrix + N translation parameters
pub parameters: Vec<f64>,
/// error / significance on parameters
pub dparameters: Vec<f64>,
/// the point about which rotations are performed
pub center: Vec<f64>,
/// the shape of images for which this transform is meant
pub shape: Vec<usize>,
ndim: usize,
dimension: PhantomData<D>,
}
impl<D: Dimension> Mul for Transform<D> {
type Output = Transform<D>;
fn mul(self, rhs: Self) -> Self::Output {
&self * &rhs
}
}
impl<D: Dimension> Mul<&Transform<D>> for Transform<D> {
type Output = Transform<D>;
fn mul(self, rhs: &Transform<D>) -> Self::Output {
&self * rhs
}
}
impl<D: Dimension> Mul<Transform<D>> for &Transform<D> {
type Output = Transform<D>;
fn mul(self, rhs: Transform<D>) -> Self::Output {
self * &rhs
}
}
impl<D: Dimension> Mul<&Transform<D>> for &Transform<D> {
type Output = Transform<D>;
#[allow(clippy::suspicious_arithmetic_impl)]
fn mul(self, rhs: &Transform<D>) -> Self::Output {
let m = self.matrix().dot(&rhs.matrix());
let dm = self.dmatrix().dot(&rhs.matrix()) + self.matrix().dot(&rhs.dmatrix());
Self::Output {
parameters: m
.slice(s![..self.ndim, ..self.ndim])
.flatten()
.iter()
.chain(m.slice(s![..self.ndim, self.ndim]).iter())
.cloned()
.collect(),
dparameters: dm
.slice(s![..self.ndim, ..])
.flatten()
.iter()
.chain(dm.slice(s![..self.ndim, ..]).iter())
.cloned()
.collect(),
center: self.center.clone(),
shape: self.shape.clone(),
ndim: self.ndim,
dimension: self.dimension,
}
}
}
impl<D: Dimension> Eq for Transform<D> {}
impl<D: Dimension> Default for Transform<D> {
/// the unit transform
fn default() -> Self {
let ndim = D::NDIM.expect("number of dimensions must be known to initialise transform");
let mut parameters = vec![0.0; ndim * ndim + ndim];
for i in 0..ndim {
parameters[(ndim + 1) * i] = 1.0;
}
Self {
parameters,
dparameters: vec![0f64; ndim * ndim + ndim],
center: vec![0f64; ndim],
shape: vec![0usize; ndim],
ndim,
dimension: PhantomData,
}
}
}
impl<D: Dimension> Transform<D> {
/// parameters: flat NxN part of matrix + translation; center: center of rotation
pub fn new_with_center(parameters: Vec<f64>, center: Vec<f64>, shape: Vec<usize>) -> Self {
let ndim = if let Some(ndim) = D::NDIM {
ndim
} else {
center.len()
};
debug_assert_eq!(parameters.len(), ndim * ndim + ndim);
debug_assert_eq!(center.len(), ndim);
debug_assert_eq!(shape.len(), ndim);
Self {
parameters,
dparameters: vec![0.0; ndim * ndim + ndim],
center,
shape,
ndim,
dimension: PhantomData,
}
}
pub fn new(parameters: Vec<f64>, shape: Vec<usize>) -> Self {
let ndim = if let Some(ndim) = D::NDIM {
ndim
} else {
shape.len()
};
debug_assert_eq!(parameters.len(), ndim * ndim + ndim);
debug_assert_eq!(shape.len(), ndim);
let center = shape.iter().map(|s| ((*s as f64) - 1.0) / 2.0).collect();
Self {
parameters,
dparameters: vec![0.0; ndim * ndim + ndim],
center,
shape,
ndim,
dimension: PhantomData,
}
}
/// find the affine transform which transforms moving into fixed
pub fn register_affine<'a, F, M, T>(fixed: F, moving: M) -> Result<Self, Error>
where
F: AsArray<'a, T, D>,
M: AsArray<'a, T, D>,
T: 'a + Clone + AsPrimitive<f64>,
{
Registration::new_affine()?.register(fixed, moving)
}
/// find the translation which transforms moving into fixed
pub fn register_translation<'a, F, M, T>(fixed: F, moving: M) -> Result<Self, Error>
where
F: AsArray<'a, T, D>,
M: AsArray<'a, T, D>,
T: 'a + Clone + AsPrimitive<f64>,
{
Registration::new_translation()?.register(fixed, moving)
}
/// find the transform which transforms moving into fixed, using fixed_mu to specify which
/// parameters to keep fixed
pub fn register<'a, F, M, T, G>(fixed: F, moving: M, fixed_mu: G) -> Result<Self, Error>
where
F: AsArray<'a, T, D>,
M: AsArray<'a, T, D>,
T: 'a + Clone + AsPrimitive<f64>,
G: Into<FixedMu>,
{
Registration::new(fixed_mu).register(fixed, moving)
}
/// find the transform which transforms moving into fixed and return the results of each
/// optimization step
pub fn register_debug<'a, F, M, T, G>(
fixed: F,
moving: M,
fixed_mu: G,
steps: Option<Vec<RegistrationStep>>,
initial_guess: Option<Vec<f64>>,
) -> Result<(Self, Vec<RegistrationResult>), Error>
where
F: AsArray<'a, T, D>,
M: AsArray<'a, T, D>,
T: 'a + Clone + AsPrimitive<f64>,
G: Into<FixedMu>,
{
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.register_debug(fixed, moving)
}
/// create a transform from a translation
pub fn from_translation(translation: &[f64]) -> Self {
let ndim = if let Some(ndim) = D::NDIM {
ndim
} else {
translation.len()
};
debug_assert_eq!(translation.len(), ndim);
let mut parameters = vec![0.0; ndim * ndim + ndim];
for i in 0..ndim {
parameters[(ndim + 1) * i] = 1.0;
}
for (p, t) in parameters
.iter_mut()
.skip(ndim * ndim)
.zip_eq(translation.iter())
{
*p = *t;
}
Self {
parameters,
dparameters: vec![0f64; ndim * ndim + ndim],
center: vec![0f64; ndim],
shape: vec![0usize; ndim],
ndim,
dimension: PhantomData,
}
}
/// create a transform from a scaling
pub fn from_scaling(scaling: &[f64]) -> Self {
let ndim = if let Some(ndim) = D::NDIM {
ndim
} else {
scaling.len()
};
debug_assert_eq!(scaling.len(), ndim);
let mut parameters = vec![0.0; ndim * ndim + ndim];
parameters[ndim * ndim + ndim - 1] = 1.0;
for (p, s) in parameters
.iter_mut()
.step_by(ndim + 1)
.zip_eq(scaling.iter())
{
*p = *s;
}
Self {
parameters,
dparameters: vec![0f64; ndim * ndim + ndim],
center: vec![0f64; ndim],
shape: vec![0usize; ndim],
ndim,
dimension: PhantomData,
}
}
/// add a translation to self
pub fn with_translation(&self, translation: &[f64]) -> Self {
Self::from_translation(translation) * self
}
/// scale self
pub fn with_scaling(&self, scaling: &[f64]) -> Self {
Self::from_scaling(scaling) * self
}
/// read a transform from a file
pub fn from_file(path: PathBuf) -> Result<Self, Error> {
let file = File::open(path)?;
Ok(from_reader(file)?)
}
/// write a transform to a file
pub fn to_file(&self, path: PathBuf) -> Result<(), Error> {
let mut file = std::fs::OpenOptions::new()
.create(true)
.write(true)
.truncate(true)
.open(path)?;
to_writer(&mut file, self)?;
Ok(())
}
/// true if transform does nothing
pub fn is_unity(&self) -> bool {
let n = self.ndim * self.ndim;
let m = self.ndim + 1;
self.parameters
.iter()
.take(n)
.enumerate()
.all(|(i, &x)| if i % m == 0 { x == 1.0 } else { x == 0.0 })
&& self.parameters.iter().skip(n).all(|&x| x == 0.0)
}
/// get coordinates resulting from transforming input coordinates, coordinate must have N
/// columns: x, y, z, ...
pub fn transform_point<'a, T>(&'a self, point: &'a [T]) -> Vec<f64>
where
T: AsPrimitive<f64>,
{
transform_point(point, &self.center, &self.parameters)
}
/// get coordinates resulting from transforming input coordinates, coordinates must have N
/// columns: x, y, z, ...
pub fn transform_points<'a, A, T>(&'a self, points: A) -> Result<Array2<f64>, Error>
where
T: AsPrimitive<f64>,
A: AsArray<'a, T, Ix2>,
{
// [a, b, c] x = a*x + b*y + c
// [d, e, f] y = d*x + e*y + f
// [0, 0, 1] 1 = 1
let points = points.into();
let shape = points.shape();
let a = points
.rows()
.into_iter()
.flat_map(|p| {
if let Some(q) = p.as_slice() {
transform_point(q, &self.center, &self.parameters)
} else {
transform_point(&p.to_vec(), &self.center, &self.parameters)
}
})
.collect::<Vec<_>>();
Ok(ArrayD::from_shape_vec(shape, a)?.into_dimensionality()?)
}
/// get the matrix defining the transform
pub fn matrix(&self) -> Array2<f64> {
let n = self.ndim * self.ndim;
let mut matrix = Array2::eye(self.ndim + 1);
for (m, p) in matrix
.slice_mut(s![..self.ndim, ..self.ndim])
.iter_mut()
.zip_eq(self.parameters.iter().take(n))
{
*m = *p;
}
for (m, p) in matrix
.slice_mut(s![..self.ndim, self.ndim])
.iter_mut()
.zip_eq(self.parameters.iter().skip(n))
{
*m = *p;
}
matrix
}
/// get the matrix describing the error of the transform
pub fn dmatrix(&self) -> Array2<f64> {
let n = self.ndim * self.ndim;
let mut matrix = Array2::zeros([self.ndim + 1, self.ndim + 1]);
for (m, p) in matrix
.slice_mut(s![..self.ndim, ..self.ndim])
.iter_mut()
.zip_eq(self.dparameters.iter().take(n))
{
*m = *p;
}
for (m, p) in matrix
.slice_mut(s![..self.ndim, self.ndim])
.iter_mut()
.zip_eq(self.dparameters.iter().skip(n))
{
*m = *p;
}
matrix
}
/// get the inverse transform
pub fn inverse(&self) -> Result<Self, Error> {
let matrix = self.matrix();
let inverse = matrix.inv()?;
let parameters = inverse
.slice(s![..self.ndim, ..self.ndim])
.iter()
.chain(inverse.slice(s![..self.ndim, self.ndim]))
.cloned()
.collect();
Ok(Self {
parameters,
dparameters: vec![0f64; self.ndim * self.ndim + self.ndim],
center: self.center.clone(),
shape: self.shape.clone(),
ndim: self.ndim,
dimension: PhantomData,
})
}
/// adapt the transform to a new center and shape
pub fn adapt(&mut self, center: &[f64], shape: &[usize]) {
self.center = self
.shape
.iter()
.zip_eq(shape.iter())
.zip_eq(center)
.map(|((a, b), o)| o + (((a - b) as f64) / 2.0))
.collect();
self.shape = shape.to_vec();
}
pub fn interpolate<'a, const B: usize, T, A>(&self, image: A) -> Result<Array<f64, D>, Error>
where
A: AsArray<'a, T, D>,
D: Dimension,
T: 'a + Clone + AsPrimitive<f64>,
BSpline<B, D>: BSplineTrait<D>,
{
let image = image.into();
let bspline = BSpline::<B, _>::new(&image);
bspline.interpolate(self)
}
pub fn interpolate_par<'a, const B: usize, T, A>(
&self,
image: A,
) -> Result<Array<f64, D>, Error>
where
A: AsArray<'a, T, D>,
D: Dimension,
T: 'a + Clone + AsPrimitive<f64>,
BSpline<B, D>: BSplineTrait<D>,
{
let image = image.into();
let bspline = BSpline::<B, _>::new(&image);
bspline.interpolate_par(self)
}
}
impl Transform<Ix2> {
/// create a transform from a rotation in radians
pub fn from_rotation(theta: f64, center: &[f64]) -> Self {
Self {
parameters: vec![
theta.cos(),
-theta.sin(),
theta.sin(),
theta.cos(),
0.0,
0.0,
],
dparameters: vec![0f64; 6],
center: center.to_vec(),
shape: vec![0usize; 2],
ndim: 2,
dimension: PhantomData,
}
}
/// multiply self with a rotation matrix
pub fn with_rotation(&self, theta: f64) -> Self {
Self::from_rotation(theta, &self.center) * self
}
/// multiply self with a rotation matrix
pub fn with_rotation_around(&self, theta: f64, center: &[f64]) -> Self {
Self::from_rotation(theta, center) * self
}
}
#[cfg(test)]
mod tests {
use crate::julia_image;
use crate::transform::Transform;
use itertools::Itertools;
use ndarray::s;
use num::traits::FloatConst;
#[test]
fn interpolate() -> Result<(), Box<dyn std::error::Error>> {
let image_a = julia_image(
&[8000, 6000],
&[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
&[3999.5, 2999.5],
&[-0.8, 0.156],
);
let transform = Transform::from_rotation(1.0, &[3999.5, 2999.5]);
let image_b = transform.interpolate_par::<3, _, _>(image_a.view())?;
let image_c = transform.interpolate::<3, _, _>(image_a.view())?;
assert!(
image_b
.into_iter()
.zip_eq(image_c.into_iter())
.all(|(a, b)| a == b)
);
Ok(())
}
#[test]
fn interpolate_serial() -> Result<(), Box<dyn std::error::Error>> {
let image_a = julia_image(
&[8000, 6000],
&[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
&[3999.5, 2999.5],
&[-0.8, 0.156],
);
let transform = Transform::from_rotation(1.0, &[3999.5, 2999.5]);
let image_b = transform.interpolate::<3, _, _>(image_a.view())?;
assert_ne!(image_b[[4000, 3000]], 0.0);
Ok(())
}
#[test]
fn interpolate_par() -> Result<(), Box<dyn std::error::Error>> {
let image_a = julia_image(
&[8000, 6000],
&[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
&[3999.5, 2999.5],
&[-0.8, 0.156],
);
let transform = Transform::from_rotation(1.0, &[3999.5, 2999.5]);
let image_b = transform.interpolate_par::<3, _, _>(image_a.view())?;
assert_ne!(image_b[[4000, 3000]], 0.0);
Ok(())
}
#[test]
fn register1() -> Result<(), Box<dyn std::error::Error>> {
let im_a = julia_image(
&[100, 1],
&[1.0, 0.0, 0.0, 0.01, 0.0, 0.0],
&[99.5, 0.5],
&[-0.8, 0.156],
)
.slice(s![.., 0])
.mapv(|i| i as f64);
let im_b =
Transform::new(vec![0.85, 4.0], vec![im_a.shape()[0]]).interpolate::<1, _, _>(&im_a)?;
let (t, steps) =
Transform::register_debug(im_a.view(), im_b.view(), vec![None, None], None, None)?;
println!("steps:");
for step in steps {
println!(" {:?}", step);
}
println!("t: {:?}", t);
println!("i: {:?}", t.inverse()?);
Ok(())
}
#[test]
fn register2() -> Result<(), Box<dyn std::error::Error>> {
let t = Transform::from_rotation(f64::PI() / 4.0, &[299.5, 399.5]);
let s = Transform::new(vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0], vec![600, 800]);
let r = t * s;
let p = r.parameters;
println!("{:?}", p);
let im_a = julia_image(
&[600, 800],
&[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
&[299.5, 399.5],
&[-0.8, 0.156],
);
let im_b = julia_image(
&[600, 800],
&[p[0], p[1], p[2], p[3], p[4], p[5]],
&[299.5, 399.5],
&[-0.8, 0.156],
);
let (t, steps) = Transform::register_debug(
im_a.view(),
im_b.view(),
vec![None, None, None, None, Some(0.0), Some(0.0)],
None,
None,
)?;
println!("steps:");
for step in steps {
println!(" {:?}", step);
}
println!("t: {:?}", t);
println!("i: {:?}", t.inverse()?);
Ok(())
}
}