1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

use num_complex::Complex;

/// Specifies a type of transform to perform.
#[derive(Copy, Clone, PartialEq, Eq)]
pub enum Transform {
    /// Fast Fourier transform
    Fft,
    /// Inverse fast Fourier transform
    Ifft,
    /// Unscaled IFFT (conventionally the IFFT is scaled by `1 / N`)
    UnscaledIfft,
    /// Square-root scaled FFT (a unitary transform)
    SqrtScaledFft,
    /// Square-root scaled IFFT (a unitary transform)
    SqrtScaledIfft,
}

impl Transform {
    /// Returns true if the transform is a forward transform.
    pub fn is_forward(&self) -> bool {
        match self {
            Self::Fft | Self::SqrtScaledFft => true,
            Self::Ifft | Self::UnscaledIfft | Self::SqrtScaledIfft => false,
        }
    }

    /// Returns the inverse transform, or `None` for `UnscaledIfft`.
    pub fn inverse(&self) -> Option<Self> {
        match self {
            Self::Fft => Some(Self::Ifft),
            Self::Ifft => Some(Self::Fft),
            Self::SqrtScaledFft => Some(Self::SqrtScaledIfft),
            Self::SqrtScaledIfft => Some(Self::SqrtScaledFft),
            Self::UnscaledIfft => None,
        }
    }
}

/// The interface for performing FFTs.
pub trait Fft {
    /// The real type used by the FFT.
    type Real: Copy;

    /// The size of the FFT.
    fn size(&self) -> usize;

    /// Apply an FFT or IFFT in-place.
    fn transform_in_place(&self, input: &mut [Complex<Self::Real>], transform: Transform);

    /// Apply an FFT or IFFT out-of-place.
    fn transform(
        &self,
        input: &[Complex<Self::Real>],
        output: &mut [Complex<Self::Real>],
        transform: Transform,
    ) {
        assert_eq!(input.len(), self.size());
        assert_eq!(output.len(), self.size());
        output.copy_from_slice(input);
        self.transform_in_place(output, transform);
    }

    /// Apply an FFT in-place.
    fn fft_in_place(&self, input: &mut [Complex<Self::Real>]) {
        self.transform_in_place(input, Transform::Fft);
    }

    /// Apply an IFFT in-place.
    fn ifft_in_place(&self, input: &mut [Complex<Self::Real>]) {
        self.transform_in_place(input, Transform::Ifft);
    }

    /// Apply an FFT out-of-place.
    fn fft(&self, input: &[Complex<Self::Real>], output: &mut [Complex<Self::Real>]) {
        self.transform(input, output, Transform::Fft);
    }

    /// Apply an IFFT out-of-place.
    fn ifft(&self, input: &[Complex<Self::Real>], output: &mut [Complex<Self::Real>]) {
        self.transform(input, output, Transform::Ifft);
    }
}