Ivan Ukhov
10 years ago
4 changed files with 285 additions and 200 deletions
-
42src/complex.rs
-
183src/lib.rs
-
256tests/fixtures.rs
-
2tests/lib.rs
@ -0,0 +1,42 @@ |
|||||
|
use std::ops::{Add, Mul, Sub};
|
||||
|
|
||||
|
/// A complex number.
|
||||
|
#[allow(non_camel_case_types)]
|
||||
|
#[derive(Clone, Copy, Debug)]
|
||||
|
pub struct c64(pub f64, pub f64);
|
||||
|
|
||||
|
impl Add for c64 {
|
||||
|
type Output = Self;
|
||||
|
|
||||
|
#[inline(always)]
|
||||
|
fn add(self, rhs: c64) -> c64 {
|
||||
|
c64(self.0 + rhs.0, self.1 + rhs.1)
|
||||
|
}
|
||||
|
}
|
||||
|
|
||||
|
impl Mul for c64 {
|
||||
|
type Output = Self;
|
||||
|
|
||||
|
#[inline(always)]
|
||||
|
fn mul(self, rhs: c64) -> c64 {
|
||||
|
c64(self.0 * rhs.0 - self.1 * rhs.1, self.0 * rhs.1 + self.1 * rhs.0)
|
||||
|
}
|
||||
|
}
|
||||
|
|
||||
|
impl Mul<f64> for c64 {
|
||||
|
type Output = Self;
|
||||
|
|
||||
|
#[inline(always)]
|
||||
|
fn mul(self, rhs: f64) -> c64 {
|
||||
|
c64(self.0 * rhs, self.1 * rhs)
|
||||
|
}
|
||||
|
}
|
||||
|
|
||||
|
impl Sub for c64 {
|
||||
|
type Output = Self;
|
||||
|
|
||||
|
#[inline(always)]
|
||||
|
fn sub(self, rhs: c64) -> c64 {
|
||||
|
c64(self.0 - rhs.0, self.1 - rhs.1)
|
||||
|
}
|
||||
|
}
|
||||
@ -1,96 +1,139 @@ |
|||||
#![feature(step_by)]
|
|
||||
|
use std::slice;
|
||||
|
|
||||
|
mod complex;
|
||||
|
pub use complex::c64;
|
||||
|
|
||||
|
/// A means of obtaining a slice of mutable complex numbers.
|
||||
|
pub trait AsMutComplex<'l> {
|
||||
|
fn as_mut_complex(self) -> &'l mut [c64];
|
||||
|
}
|
||||
|
|
||||
|
impl<'l> AsMutComplex<'l> for &'l mut [c64] {
|
||||
|
#[inline(always)]
|
||||
|
fn as_mut_complex(self) -> &'l mut [c64] {
|
||||
|
self
|
||||
|
}
|
||||
|
}
|
||||
|
|
||||
|
impl<'l> AsMutComplex<'l> for &'l mut [f64] {
|
||||
|
/// Treat the slice as a collection of pairs of real and imaginary parts and
|
||||
|
/// reinterpret it as a slice of complex numbers.
|
||||
|
///
|
||||
|
/// ## Panics
|
||||
|
///
|
||||
|
/// The function panics if the number of elements is not even.
|
||||
|
#[inline]
|
||||
|
fn as_mut_complex(self) -> &'l mut [c64] {
|
||||
|
unsafe {
|
||||
|
let length = self.len();
|
||||
|
assert!(length % 2 == 0, "the number of elements should be even");
|
||||
|
slice::from_raw_parts_mut(self.as_mut_ptr() as *mut _, length / 2)
|
||||
|
}
|
||||
|
}
|
||||
|
}
|
||||
|
|
||||
|
impl<'l> AsMutComplex<'l> for &'l mut Vec<f64> {
|
||||
|
#[inline]
|
||||
|
fn as_mut_complex(self) -> &'l mut [c64] {
|
||||
|
<&mut [f64]>::as_mut_complex(&mut *self)
|
||||
|
}
|
||||
|
}
|
||||
|
|
||||
/// Perform the Fourier transform.
|
/// Perform the Fourier transform.
|
||||
///
|
///
|
||||
/// `data` should contain `n` complex numbers where `n` is a power of two. Each
|
|
||||
/// complex number is stored as a pair of `f64`s so that the first is the real
|
|
||||
/// part and the second is the corresponding imaginary part. Hence, the total
|
|
||||
/// length of `data` should be `2 × n`.
|
|
||||
///
|
|
||||
/// # Panics
|
|
||||
///
|
|
||||
/// The function panics if `data.len()` is not even or `data.len() / 2` is not a
|
|
||||
/// power of two.
|
|
||||
|
/// The number of points should be a power of two.
|
||||
#[inline(always)]
|
#[inline(always)]
|
||||
pub fn forward(data: &mut [f64]) {
|
|
||||
transform(data, 1.0);
|
|
||||
|
pub fn forward<'l, T: AsMutComplex<'l>>(data: T) {
|
||||
|
let data = data.as_mut_complex();
|
||||
|
|
||||
|
let n = data.len();
|
||||
|
if n < 1 || n & (n - 1) != 0 {
|
||||
|
panic!("expected the number of points to be a power of two");
|
||||
|
}
|
||||
|
|
||||
|
rearrange(data, n);
|
||||
|
perform(data, n, false);
|
||||
}
|
}
|
||||
|
|
||||
/// Perform the inverse Fourier transform.
|
/// Perform the inverse Fourier transform.
|
||||
///
|
///
|
||||
/// `data` should contain `n` complex numbers where `n` is a power of two. Each
|
|
||||
/// complex number is stored as a pair of `f64`s so that the first is the real
|
|
||||
/// part and the second is the corresponding imaginary part. Hence, the total
|
|
||||
/// length of `data` should be `2 × n`.
|
|
||||
///
|
|
||||
/// # Panics
|
|
||||
///
|
|
||||
/// The function panics if `data.len()` is not even or `data.len() / 2` is not a
|
|
||||
/// power of two.
|
|
||||
|
/// The number of points should be a power of two.
|
||||
#[inline(always)]
|
#[inline(always)]
|
||||
pub fn inverse(data: &mut [f64]) {
|
|
||||
transform(data, -1.0);
|
|
||||
}
|
|
||||
|
|
||||
fn transform(data: &mut [f64], isign: f64) {
|
|
||||
use std::f64::consts::PI;
|
|
||||
|
pub fn inverse<'l, T: AsMutComplex<'l>>(data: T, scaling: bool) {
|
||||
|
let data = data.as_mut_complex();
|
||||
|
|
||||
let l = data.len();
|
|
||||
if l % 2 != 0 {
|
|
||||
panic!("expected the length of the data to be even");
|
|
||||
}
|
|
||||
|
|
||||
let n = l / 2;
|
|
||||
|
let n = data.len();
|
||||
if n < 1 || n & (n - 1) != 0 {
|
if n < 1 || n & (n - 1) != 0 {
|
||||
panic!("expected the number of points to be a power of two");
|
panic!("expected the number of points to be a power of two");
|
||||
}
|
}
|
||||
|
|
||||
let nn = n << 1;
|
|
||||
|
rearrange(data, n);
|
||||
|
perform(data, n, true);
|
||||
|
if scaling {
|
||||
|
scale(data, n);
|
||||
|
}
|
||||
|
}
|
||||
|
|
||||
let mut j = 1;
|
|
||||
for i in (1..nn).step_by(2) {
|
|
||||
if j > i {
|
|
||||
data.swap(j - 1, i - 1);
|
|
||||
data.swap(j, i);
|
|
||||
|
#[inline(always)]
|
||||
|
fn rearrange(data: &mut [c64], n: usize) {
|
||||
|
let mut target = 0;
|
||||
|
for position in 0..n {
|
||||
|
if target > position {
|
||||
|
data.swap(position, target);
|
||||
}
|
}
|
||||
let mut m = n;
|
|
||||
while m >= 2 && j > m {
|
|
||||
j -= m;
|
|
||||
m >>= 1;
|
|
||||
|
let mut mask = n >> 1;
|
||||
|
while target & mask != 0 {
|
||||
|
target &= !mask;
|
||||
|
mask >>= 1;
|
||||
}
|
}
|
||||
j += m;
|
|
||||
|
target |= mask;
|
||||
}
|
}
|
||||
|
}
|
||||
|
|
||||
|
#[inline(always)]
|
||||
|
fn perform(data: &mut [c64], n: usize, inverse: bool) {
|
||||
|
use std::f64::consts::PI;
|
||||
|
|
||||
|
let pi = if inverse { PI } else { -PI };
|
||||
|
|
||||
let mut mmax = 2;
|
|
||||
while nn > mmax {
|
|
||||
let istep = mmax << 1;
|
|
||||
let theta = isign * (2.0 * PI / mmax as f64);
|
|
||||
let wtemp = (0.5 * theta).sin();
|
|
||||
let wpr = -2.0 * wtemp * wtemp;
|
|
||||
let wpi = theta.sin();
|
|
||||
let mut wr = 1.0;
|
|
||||
let mut wi = 0.0;
|
|
||||
for m in (1..mmax).step_by(2) {
|
|
||||
for i in (m..(nn + 1)).step_by(istep) {
|
|
||||
let j = i + mmax;
|
|
||||
let tempr = wr * data[j - 1] - wi * data[j];
|
|
||||
let tempi = wr * data[j] + wi * data[j - 1];
|
|
||||
data[j - 1] = data[i - 1] - tempr;
|
|
||||
data[j] = data[i] - tempi;
|
|
||||
data[i - 1] += tempr;
|
|
||||
data[i] += tempi;
|
|
||||
|
let mut step = 1;
|
||||
|
while step < n {
|
||||
|
let jump = step << 1;
|
||||
|
let delta = pi / step as f64;
|
||||
|
let sine = (0.5 * delta).sin();
|
||||
|
let multiplier = c64(-2.0 * sine * sine, delta.sin());
|
||||
|
let mut factor = c64(1.0, 0.0);
|
||||
|
for group in 0..step {
|
||||
|
let mut pair = group;
|
||||
|
while pair < n {
|
||||
|
let match_pair = pair + step;
|
||||
|
let product = factor * data[match_pair];
|
||||
|
data[match_pair] = data[pair] - product;
|
||||
|
data[pair] = data[pair] + product;
|
||||
|
pair += jump;
|
||||
}
|
}
|
||||
let wtemp = wr;
|
|
||||
wr = wr * wpr - wi * wpi + wr;
|
|
||||
wi = wi * wpr + wtemp * wpi + wi;
|
|
||||
|
factor = multiplier * factor + factor;
|
||||
}
|
}
|
||||
mmax = istep;
|
|
||||
|
step <<= 1;
|
||||
}
|
}
|
||||
|
}
|
||||
|
|
||||
if isign == -1.0 {
|
|
||||
let scale = 1.0 / n as f64;
|
|
||||
for i in 0..l {
|
|
||||
data[i] *= scale;
|
|
||||
|
#[inline(always)]
|
||||
|
fn scale(data: &mut [c64], n: usize) {
|
||||
|
let factor = 1.0 / n as f64;
|
||||
|
for position in 0..n {
|
||||
|
data[position] = data[position] * factor;
|
||||
}
|
}
|
||||
|
}
|
||||
|
|
||||
|
#[cfg(test)]
|
||||
|
mod tests {
|
||||
|
use c64;
|
||||
|
use std::mem;
|
||||
|
|
||||
|
#[test]
|
||||
|
fn size_of() {
|
||||
|
assert_eq!(mem::size_of::<c64>(), 2 * mem::size_of::<f64>());
|
||||
|
assert_eq!(mem::size_of::<[c64; 42]>(), 2 * mem::size_of::<[f64; 42]>());
|
||||
}
|
}
|
||||
}
|
}
|
||||
Write
Preview
Loading…
Cancel
Save
Reference in new issue