lenia_ca/
fft.rs

1//! Contains the required functionality for performing n-dimensional fast-fourier-transforms.
2
3#![allow(dead_code)]
4#![allow(unused_variables)]
5
6use std::{fmt, sync::Arc};
7use rustfft::{Fft, FftPlanner, FftDirection};
8use rustfft::num_complex::Complex;
9use rayon::prelude::*;
10
11/// Holds all the relevant data for a pre-planned FFT, which is to say, once initialized, 
12/// it can perform efficient FFT-s on data of the initially specified length.
13#[derive(Clone)]
14pub struct PlannedFFT {
15    fft: Arc<dyn Fft<f64>>,
16    scratch_space: Vec<Complex<f64>>,
17}
18
19impl fmt::Debug for PlannedFFT {
20    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
21        f.debug_struct("PreplannedFFT")
22         .field("scratch_space", &format!("Vec<Complex<f64>>, len: {}", self.scratch_space.len()))
23         .field("fft", &format!("Arc<dyn rustfft::Fft<f64>> => len: {}, direction: {}", self.fft.len(), self.fft.fft_direction()))
24         .finish()
25    }
26}
27
28impl PlannedFFT {
29    pub fn new(length: usize, inverse: bool) -> Self {
30        if length == 0 { panic!("PlannedFFT::new() - Provided length was 0. Length must be at least 1!"); }
31        let mut planner = FftPlanner::new();
32        let direction: FftDirection;
33        match inverse {
34            true => { direction = FftDirection::Inverse; }
35            false => { direction = FftDirection::Forward; }
36        }
37        let fft = planner.plan_fft(
38            length, 
39            direction,
40        );
41        let scratch_space: Vec<Complex<f64>> = Vec::from_iter(std::iter::repeat(Complex::new(0.0, 0.0)).take(fft.get_inplace_scratch_len()));
42
43        PlannedFFT {
44            fft: fft,
45            scratch_space: scratch_space,
46        }
47    }
48
49    pub fn inverse(&self) -> bool {
50        match self.fft.fft_direction() {
51            FftDirection::Forward => { return false; }
52            FftDirection::Inverse => { return true; }
53        }
54    }
55
56    pub fn length(&self) -> usize {
57        self.fft.len()
58    }
59
60    pub fn transform(&mut self, data: &mut [Complex<f64>]) {
61        self.fft.process_with_scratch(data, &mut self.scratch_space);
62        if self.inverse() {     // I fekin' forgot this AGAIN...
63            let inverse_len = 1.0 / data.len() as f64;
64            for v in data.iter_mut() {
65                v.re *= inverse_len;
66                v.im *= inverse_len;
67            }
68        }
69    }
70}
71
72/// Holds all the relevant data for a pre-planned N-dimensional fast-fourier-transform. Operates only
73/// on data of the initially specified length.
74#[derive(Debug)]
75pub struct PlannedFFTND {
76    shape: Vec<usize>,
77    fft_instances: Vec<PlannedFFT>,
78    inverse: bool
79}
80
81impl PlannedFFTND {
82    pub fn new(shape: &[usize], inverse: bool) -> Self {
83        if shape.is_empty() { panic!("PlannedFFTND::new() - Provided shape was empty! Needs at least 1 dimension!"); }
84        let mut ffts: Vec<PlannedFFT> = Vec::with_capacity(shape.len());
85        for dim in shape {
86            ffts.push(PlannedFFT::new(*dim, inverse));
87        }
88        PlannedFFTND {
89            shape: shape.to_vec(),
90            fft_instances: ffts,
91            inverse: inverse,
92        }
93    }
94
95    pub fn shape(&self) -> &[usize] {
96        &self.shape
97    }
98
99    pub fn inverse(&self) -> bool {
100        self.inverse
101    }
102
103    pub fn transform(&mut self, data: &mut ndarray::ArrayD<Complex<f64>>) {
104        if data.shape() != self.shape { panic!("PlannedFFTND::transform() - shape of the data to be transformed does not agree with the shape that the fft can work on!"); }
105        let mut axis_iterator: Vec<usize> = Vec::with_capacity(self.shape.len());
106        if self.inverse() {
107            for i in (0..self.shape.len()).rev() {
108                axis_iterator.push(i);
109            }
110        }
111        else {
112            for i in 0..self.shape.len() {
113                axis_iterator.push(i);
114            }
115        }
116        for axis in axis_iterator {
117            for mut lane in data.lanes_mut(ndarray::Axis(axis)) {
118                let mut buf = lane.to_vec();
119                self.fft_instances[axis].transform(&mut buf);
120                for i in 0..lane.len() {
121                    lane[i] = buf[i];
122                }
123            }
124        }
125    }
126}
127
128
129/// Parallel version (multithreaded) of the PlannedFFTND.
130#[derive(Debug)]
131pub struct ParPlannedFFTND {
132    shape: Vec<usize>,
133    fft_instances: Vec<PlannedFFT>,
134    inverse: bool
135}
136
137impl ParPlannedFFTND {
138    pub fn new(shape: &[usize], inverse: bool) -> Self {
139        if shape.is_empty() { panic!("ParPlannedFFTND::new() - Provided shape was empty! Needs at least 1 dimension!"); }
140        let mut ffts: Vec<PlannedFFT> = Vec::with_capacity(shape.len());
141        for dim in shape {
142            ffts.push(PlannedFFT::new(*dim, inverse));
143        }
144        ParPlannedFFTND {
145            shape: shape.to_vec(),
146            fft_instances: ffts,
147            inverse: inverse,
148        }
149    }
150
151    pub fn shape(&self) -> &[usize] {
152        &self.shape
153    }
154
155    pub fn inverse(&self) -> bool {
156        self.inverse
157    }
158
159    pub fn transform(&mut self, data: &mut ndarray::ArrayD<Complex<f64>>) {
160        if data.shape() != self.shape { panic!("ParPlannedFFTND::transform() - shape of the data to be transformed does not agree with the shape that the fft can work on!"); }
161        let mut axis_iterator: Vec<usize> = Vec::with_capacity(self.shape.len());
162        if self.inverse() {
163            for i in (0..self.shape.len()).rev() {
164                axis_iterator.push(i);
165            }
166        }
167        else {
168            for i in 0..self.shape.len() {
169                axis_iterator.push(i);
170            }
171        }
172        for axis in axis_iterator {
173            let mut data_lane = data.lanes_mut(ndarray::Axis(axis));
174            let mut fft_lane = &mut self.fft_instances[axis];
175            ndarray::Zip::from(data_lane)
176                .into_par_iter()
177                .for_each_with(self.fft_instances[axis].clone(), |mut fft, mut row| {
178                    let mut lane = row.0.to_vec();
179                    fft.transform(&mut lane);
180                    for i in 0..row.0.len() {
181                        row.0[i] = lane[i];
182                    }
183                });
184        }
185    }
186}
lenia_ca/fft.rs

lenia_ca/
fft.rs
