Skip to content

Instantly share code, notes, and snippets.

@saethlin
Created April 15, 2018 04:03
Show Gist options
  • Save saethlin/8035ebf79f576db715f2ce08b5a54c95 to your computer and use it in GitHub Desktop.
Save saethlin/8035ebf79f576db715f2ce08b5a54c95 to your computer and use it in GitHub Desktop.
FFT convolution bug somewhere
extern crate rustfft;
extern crate num_complex;
extern crate gnuplot;
use num_complex::Complex;
use gnuplot::*;
use rustfft::FFT;
fn main() {
// create an rv and ccf to convolve
let rv : Vec<_> = (0..1_000).map(|n| ((n as f64) / 1000. - 0.5) * 200.).collect();
let sqrt = f64::sqrt;
let ln = f64::ln;
let exp = f64::exp;
let c = 299792458.0;
let resolution = 115000.0;
let profile_fwhm = c / resolution / 1000.;
let profile_sigma = profile_fwhm / (2.0 * sqrt(2.0 * ln(2.0)));
let kernel: Vec<Complex<f64>> = rv.iter()
.map(|r| exp(-r.powi(2) / (2.0 * profile_sigma.powi(2))))
.map(|n| Complex::new(n, 0.0))
.collect();
let mut ccf_complex = kernel.clone();
let mut fg = Figure::new();
let ccf : Vec<_> = ccf_complex.iter().map(|n| n.re).collect();
fg.axes2d().lines(&rv, &ccf, &[]);
fg.show();
// FFT convolution
let fft = rustfft::FFTplanner::new(false).plan_fft(rv.len());
let inverse_fft = rustfft::FFTplanner::new(true).plan_fft(rv.len());
// FFT of the kernel
let mut fft_kernel = vec![Complex::new(0.0, 0.0); rv.len()];
fft.process(&mut kernel.clone(), fft_kernel.as_mut_slice());
// FFT of the CCF
let mut fft_ccf = vec![Complex::new(0.0, 0.0); rv.len()];
fft.process(&mut ccf_complex.clone(), fft_ccf.as_mut_slice());
// Multiply both frequencies and compute inverse transform
let mut convolved_fft: Vec<Complex<f64>> = fft_kernel
.iter()
.rev()
.zip(fft_ccf.iter())
.map(|(k, c)| k * c)
.collect();
let mut output = vec![Complex::new(0.0, 0.0); rv.len()];
inverse_fft.process(&mut convolved_fft, output.as_mut_slice());
let result = output.iter().map(|n| n.re).collect::<Vec<f64>>();
let mut fg = Figure::new();
fg.axes2d().lines(&rv, &result, &[]);
fg.show();
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment