272 lines
9.1 KiB
Rust
272 lines
9.1 KiB
Rust
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, 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 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(())
|
|
}
|
|
}
|