163 lines
4.2 KiB
Rust
163 lines
4.2 KiB
Rust
// Implementation of raders's fft for prime sized ffts
|
|
|
|
use std::f32::consts::PI;
|
|
|
|
use crate::{
|
|
complex::Complex32,
|
|
fft::{
|
|
create_fft, is_prime, DFTAlgorithm, FFTDirection
|
|
},
|
|
};
|
|
|
|
pub struct Rader2FFT {
|
|
permutations: Box<[usize]>,
|
|
convolution_operand: Box<[Complex32]>,
|
|
|
|
convolution_fft_input: Box<[Complex32]>,
|
|
convolution_ifft: Box<dyn DFTAlgorithm>,
|
|
convolution_fft: Box<dyn DFTAlgorithm>,
|
|
|
|
output: Box<[Complex32]>,
|
|
|
|
sub_size: usize,
|
|
size: usize,
|
|
}
|
|
|
|
impl DFTAlgorithm for Rader2FFT {
|
|
fn create(size: usize, direction: FFTDirection) -> Self
|
|
where
|
|
Self: Sized,
|
|
{
|
|
assert!(is_prime(size));
|
|
|
|
// Primitive root and its powers
|
|
let g = compute_prime_primitive_root(size);
|
|
let permutations: Box<[usize]> = (0..(size - 1)).map(|i| exp_mod(g, i + 1, size)).collect();
|
|
let sub_size = next_pow2(2 * size - 3);
|
|
|
|
// Compute fourrier transform of twiddle factors
|
|
let twiddle_factors = (0..sub_size)
|
|
.map(|i| {
|
|
Complex32::cexp(
|
|
-2. * PI * direction.sign() * (permutations[i % (size - 1)] as f32)
|
|
/ (size as f32),
|
|
)
|
|
})
|
|
.collect::<Vec<Complex32>>();
|
|
|
|
let mut convolution_fft = create_fft(sub_size, FFTDirection::Forward);
|
|
convolution_fft.execute(&twiddle_factors);
|
|
Rader2FFT {
|
|
permutations,
|
|
convolution_operand: convolution_fft.get_output().iter().copied().collect(),
|
|
|
|
convolution_fft,
|
|
convolution_ifft: create_fft(sub_size, FFTDirection::Inverse),
|
|
convolution_fft_input: vec![Complex32::zero(); sub_size].into(),
|
|
|
|
output: vec![Complex32::zero(); size].into(),
|
|
size,
|
|
sub_size,
|
|
}
|
|
}
|
|
|
|
fn execute(&mut self, input: &[Complex32]) {
|
|
// Compute fft of input signal
|
|
|
|
self.convolution_fft_input[0] = input[self.permutations[self.size - 2]];
|
|
for i in 0..(self.sub_size - self.size + 1) {
|
|
self.convolution_fft_input[i + 1] = Complex32::zero();
|
|
}
|
|
for i in 1..(self.size - 1) {
|
|
// reverse sequence
|
|
let k = self.permutations[self.size - 1 - i - 1];
|
|
self.convolution_fft_input[i + self.sub_size - self.size + 1] = input[k];
|
|
}
|
|
|
|
self.convolution_fft.execute(&self.convolution_fft_input);
|
|
|
|
// Compute convolution by multiplying in freq domain
|
|
for i in 0..self.sub_size {
|
|
// Using output as staging buffer
|
|
self.convolution_fft_input[i] =
|
|
self.convolution_fft.get_output()[i] * self.convolution_operand[i];
|
|
}
|
|
|
|
self.convolution_ifft.execute(&self.convolution_fft_input);
|
|
|
|
self.output[0] = Complex32::zero();
|
|
for x in input {
|
|
self.output[0] = self.output[0] + *x;
|
|
}
|
|
|
|
for i in 0..(self.size - 1) {
|
|
// Actually compute the output
|
|
let k = self.permutations[i];
|
|
self.output[k] =
|
|
(self.convolution_ifft.get_output()[i] / (self.sub_size) as f32) + input[0];
|
|
}
|
|
}
|
|
|
|
fn get_output(&self) -> &[Complex32] {
|
|
&self.output
|
|
}
|
|
}
|
|
|
|
pub fn compute_prime_primitive_root(n: usize) -> usize {
|
|
assert!(is_prime(n));
|
|
|
|
let phi = n - 1; // Euler's totient for n prime
|
|
|
|
// Test all candidates
|
|
for i in 1..(n + 1) {
|
|
// Find multiplicative order of i
|
|
let mut val = i;
|
|
let mut order = 1;
|
|
for _ in 0..n {
|
|
if val == 1 {
|
|
break;
|
|
}
|
|
val = (val * i) % n;
|
|
order += 1;
|
|
}
|
|
|
|
if order == phi {
|
|
return i;
|
|
}
|
|
}
|
|
|
|
unreachable!("Prime must have primitive root");
|
|
}
|
|
|
|
pub fn exp_mod(mut n: usize, mut exp: usize, m: usize) -> usize {
|
|
if m == 1 {
|
|
return 0;
|
|
}
|
|
|
|
n %= m;
|
|
let mut r = 1;
|
|
while exp > 0 {
|
|
if exp % 2 == 1 {
|
|
r = (r * n) % m;
|
|
}
|
|
|
|
n = (n * n) % m;
|
|
exp >>= 1;
|
|
}
|
|
|
|
r
|
|
}
|
|
|
|
pub fn next_pow2(mut n: usize) -> usize {
|
|
if n.count_ones() == 1 {
|
|
n
|
|
} else {
|
|
let mut p = 0;
|
|
while n > 0 {
|
|
n >>= 1;
|
|
p += 1;
|
|
}
|
|
1 << p
|
|
}
|
|
}
|