fftw/
plan.rs

1//! Plan in FFTW
2//!
3//! See also [Using Plans] in the original document
4//! [Using Plans]: http://www.fftw.org/fftw3_doc/Using-Plans.html
5
6use crate::array::{alignment_of, AlignedAllocable, AlignedVec, Alignment};
7use crate::error::*;
8use crate::ffi::*;
9use crate::types::{c32, c64, Flag, R2RKind, Sign};
10
11use std::marker::PhantomData;
12
13pub type C2CPlan64 = Plan<c64, c64, Plan64>;
14pub type C2CPlan32 = Plan<c32, c32, Plan32>;
15pub type R2CPlan64 = Plan<f64, c64, Plan64>;
16pub type R2CPlan32 = Plan<f32, c32, Plan32>;
17pub type C2RPlan64 = Plan<c64, f64, Plan64>;
18pub type C2RPlan32 = Plan<c32, f32, Plan32>;
19pub type R2RPlan64 = Plan<f64, f64, Plan64>;
20pub type R2RPlan32 = Plan<f32, f32, Plan32>;
21
22/// Typed wrapper of `fftw_plan`
23///
24/// The plan in FFTW manages the contains all information necessary to compute the transform,
25/// including the pointers to the input and output arrays.
26/// However, this wrapper *does not modify* this pointer once after the plan is created
27/// since it should be *unsafe* in terms of Rust.
28/// Instead, this plan executes a transform for different arrays with [new-array execute functions]
29/// with related associated functions, e.g. `C2CPlan::c2c`.
30///
31/// [new-array execute functions]: http://www.fftw.org/fftw3_doc/New_002darray-Execute-Functions.html
32pub struct Plan<A, B, Plan: PlanSpec> {
33    plan: Plan,
34    input: (usize, Alignment),
35    output: (usize, Alignment),
36    phantom: PhantomData<(A, B)>,
37}
38
39unsafe impl<A: Send, B: Send> Send for Plan<A, B, Plan32> {}
40unsafe impl<A: Send, B: Send> Send for Plan<A, B, Plan64> {}
41
42impl<A, B, P: PlanSpec> Drop for Plan<A, B, P> {
43    fn drop(&mut self) {
44        self.plan.destroy();
45    }
46}
47
48impl<A, B, P: PlanSpec> Plan<A, B, P> {
49    fn check(&self, in_: &[A], out: &[B]) -> Result<()> {
50        if self.input != slice_info(in_) {
51            return Err(Error::InputArrayMismatch {
52                expect: self.input,
53                actual: slice_info(in_),
54            });
55        }
56        if self.output != slice_info(out) {
57            return Err(Error::OutputArrayMismatch {
58                expect: self.output,
59                actual: slice_info(out),
60            });
61        }
62        Ok(())
63    }
64}
65
66/// Trait for Plan makers
67pub trait PlanSpec: Clone + Copy {
68    fn validate(self) -> Result<Self>;
69    fn destroy(self);
70    fn print(self);
71}
72
73/// Marker for 64-bit floating point FFT
74pub type Plan64 = fftw_plan;
75/// Marker for 32-bit floating point FFT
76pub type Plan32 = fftwf_plan;
77
78/// Trait for the plan of Complex-to-Complex transformation
79pub trait C2CPlan: Sized {
80    type Complex: AlignedAllocable;
81
82    /// Create new plan with aligned vector
83    fn aligned(shape: &[usize], sign: Sign, flag: Flag) -> Result<Self> {
84        let n: usize = shape.iter().product();
85        let mut in_ = AlignedVec::new(n);
86        let mut out = AlignedVec::new(n);
87        Self::new(shape, &mut in_, &mut out, sign, flag)
88    }
89
90    /// Create new plan
91    fn new(
92        shape: &[usize],
93        in_: &mut [Self::Complex],
94        out: &mut [Self::Complex],
95        sign: Sign,
96        flag: Flag,
97    ) -> Result<Self>;
98
99    /// Execute complex-to-complex transform
100    fn c2c(&mut self, in_: &mut [Self::Complex], out: &mut [Self::Complex]) -> Result<()>;
101}
102
103/// Trait for the plan of Real-to-Complex transformation
104pub trait R2CPlan: Sized {
105    type Real: AlignedAllocable;
106    type Complex: AlignedAllocable;
107
108    /// Create new plan with aligned vector
109    fn aligned(shape: &[usize], flag: Flag) -> Result<Self> {
110        let n: usize = shape.iter().product();
111        let n_d = shape.last().unwrap();
112        let n_sub = (n / n_d) * (n_d / 2 + 1);
113        let mut in_ = AlignedVec::new(n);
114        let mut out = AlignedVec::new(n_sub);
115        Self::new(shape, &mut in_, &mut out, flag)
116    }
117
118    /// Create new plan
119    fn new(
120        shape: &[usize],
121        in_: &mut [Self::Real],
122        out: &mut [Self::Complex],
123        flag: Flag,
124    ) -> Result<Self>;
125
126    /// Execute real-to-complex transform
127    fn r2c(&mut self, in_: &mut [Self::Real], out: &mut [Self::Complex]) -> Result<()>;
128}
129
130/// Trait for the plan of Complex-to-Real transformation
131pub trait C2RPlan: Sized {
132    type Real: AlignedAllocable;
133    type Complex: AlignedAllocable;
134
135    /// Create new plan with aligned vector
136    fn aligned(shape: &[usize], flag: Flag) -> Result<Self> {
137        let n: usize = shape.iter().product();
138        let n_d = shape.last().unwrap();
139        let n_sub = (n / n_d) * (n_d / 2 + 1);
140        let mut in_ = AlignedVec::new(n_sub);
141        let mut out = AlignedVec::new(n);
142        Self::new(shape, &mut in_, &mut out, flag)
143    }
144
145    /// Create new plan
146    fn new(
147        shape: &[usize],
148        in_: &mut [Self::Complex],
149        out: &mut [Self::Real],
150        flag: Flag,
151    ) -> Result<Self>;
152
153    /// Execute complex-to-real transform
154    fn c2r(&mut self, in_: &mut [Self::Complex], out: &mut [Self::Real]) -> Result<()>;
155}
156
157pub trait R2RPlan: Sized {
158    type Real: AlignedAllocable;
159
160    /// Create new plan with aligned vector
161    fn aligned(shape: &[usize], kind: R2RKind, flag: Flag) -> Result<Self> {
162        let n: usize = shape.iter().product();
163        let mut in_ = AlignedVec::new(n);
164        let mut out = AlignedVec::new(n);
165        Self::new(shape, &mut in_, &mut out, kind, flag)
166    }
167
168    /// Create new plan
169    fn new(
170        shape: &[usize],
171        in_: &mut [Self::Real],
172        out: &mut [Self::Real],
173        kind: R2RKind,
174        flag: Flag,
175    ) -> Result<Self>;
176
177    /// Execute complex-to-complex transform
178    fn r2r(&mut self, in_: &mut [Self::Real], out: &mut [Self::Real]) -> Result<()>;
179}
180
181macro_rules! impl_c2c {
182    ($C:ty, $Plan:ty; $plan:ident, $exec:ident) => {
183        impl C2CPlan for Plan<$C, $C, $Plan> {
184            type Complex = $C;
185            fn new(
186                shape: &[usize],
187                in_: &mut [Self::Complex],
188                out: &mut [Self::Complex],
189                sign: Sign,
190                flag: Flag,
191            ) -> Result<Self> {
192                let plan = excall! { $plan(
193                    shape.len() as i32,
194                    shape.to_cint().as_mut_ptr() as *mut _,
195                    in_.as_mut_ptr(),
196                    out.as_mut_ptr(),
197                    sign as i32, flag.bits())
198                }
199                .validate()?;
200                Ok(Self {
201                    plan,
202                    input: slice_info(in_),
203                    output: slice_info(out),
204                    phantom: PhantomData,
205                })
206            }
207            fn c2c(&mut self, in_: &mut [Self::Complex], out: &mut [Self::Complex]) -> Result<()> {
208                self.check(in_, out)?;
209                unsafe { $exec(self.plan, in_.as_mut_ptr(), out.as_mut_ptr()) };
210                Ok(())
211            }
212        }
213    };
214} // impl_c2c!
215
216impl_c2c!(c64, Plan64; fftw_plan_dft, fftw_execute_dft);
217impl_c2c!(c32, Plan32; fftwf_plan_dft, fftwf_execute_dft);
218
219macro_rules! impl_r2c {
220    ($R:ty, $C:ty, $Plan:ty; $plan:ident, $exec:ident) => {
221        impl R2CPlan for Plan<$R, $C, $Plan> {
222            type Real = $R;
223            type Complex = $C;
224            fn new(
225                shape: &[usize],
226                in_: &mut [Self::Real],
227                out: &mut [Self::Complex],
228                flag: Flag,
229            ) -> Result<Self> {
230                let plan = excall! { $plan(
231                    shape.len() as i32,
232                    shape.to_cint().as_mut_ptr() as *mut _,
233                    in_.as_mut_ptr(),
234                    out.as_mut_ptr(),
235                    flag.bits())
236                }
237                .validate()?;
238                Ok(Self {
239                    plan,
240                    input: slice_info(in_),
241                    output: slice_info(out),
242                    phantom: PhantomData,
243                })
244            }
245            fn r2c(&mut self, in_: &mut [Self::Real], out: &mut [Self::Complex]) -> Result<()> {
246                self.check(in_, out)?;
247                unsafe { $exec(self.plan, in_.as_mut_ptr(), out.as_mut_ptr()) };
248                Ok(())
249            }
250        }
251    };
252} // impl_r2c!
253
254impl_r2c!(f64, c64, Plan64; fftw_plan_dft_r2c, fftw_execute_dft_r2c);
255impl_r2c!(f32, c32, Plan32; fftwf_plan_dft_r2c, fftwf_execute_dft_r2c);
256
257macro_rules! impl_c2r {
258    ($R:ty, $C:ty, $Plan:ty; $plan:ident, $exec:ident) => {
259        impl C2RPlan for Plan<$C, $R, $Plan> {
260            type Real = $R;
261            type Complex = $C;
262            fn new(
263                shape: &[usize],
264                in_: &mut [Self::Complex],
265                out: &mut [Self::Real],
266                flag: Flag,
267            ) -> Result<Self> {
268                let plan = excall! { $plan(
269                    shape.len() as i32,
270                    shape.to_cint().as_mut_ptr() as *mut _,
271                    in_.as_mut_ptr(),
272                    out.as_mut_ptr(),
273                    flag.bits())
274                }
275                .validate()?;
276                Ok(Self {
277                    plan,
278                    input: slice_info(in_),
279                    output: slice_info(out),
280                    phantom: PhantomData,
281                })
282            }
283            fn c2r(&mut self, in_: &mut [Self::Complex], out: &mut [Self::Real]) -> Result<()> {
284                self.check(in_, out)?;
285                unsafe { $exec(self.plan, in_.as_mut_ptr(), out.as_mut_ptr()) };
286                Ok(())
287            }
288        }
289    };
290} // impl_c2r!
291
292impl_c2r!(f64, c64, Plan64; fftw_plan_dft_c2r, fftw_execute_dft_c2r);
293impl_c2r!(f32, c32, Plan32; fftwf_plan_dft_c2r, fftwf_execute_dft_c2r);
294
295macro_rules! impl_r2r {
296    ($R:ty, $Plan:ty; $plan:ident, $exec:ident) => {
297        impl R2RPlan for Plan<$R, $R, $Plan> {
298            type Real = $R;
299            fn new(
300                shape: &[usize],
301                in_: &mut [Self::Real],
302                out: &mut [Self::Real],
303                kind: R2RKind,
304                flag: Flag,
305            ) -> Result<Self> {
306                let plan = excall! { $plan(
307                    shape.len() as i32,
308                    shape.to_cint().as_mut_ptr() as *mut _,
309                    in_.as_mut_ptr(),
310                    out.as_mut_ptr(),
311                    &kind as *const _, flag.bits())
312                }
313                .validate()?;
314                Ok(Self {
315                    plan,
316                    input: slice_info(in_),
317                    output: slice_info(out),
318                    phantom: PhantomData,
319                })
320            }
321            fn r2r(&mut self, in_: &mut [Self::Real], out: &mut [Self::Real]) -> Result<()> {
322                self.check(in_, out)?;
323                unsafe { $exec(self.plan, in_.as_mut_ptr(), out.as_mut_ptr()) };
324                Ok(())
325            }
326        }
327    };
328} // impl_r2r!
329
330impl_r2r!(f64, Plan64; fftw_plan_r2r, fftw_execute_r2r);
331impl_r2r!(f32, Plan32; fftwf_plan_r2r, fftwf_execute_r2r);
332
333macro_rules! impl_plan_spec {
334    ($Plan:ty; $destroy_plan:ident, $print_plan:ident) => {
335        impl PlanSpec for $Plan {
336            fn validate(self) -> Result<Self> {
337                if self.is_null() {
338                    Err(Error::InvalidPlanError {})
339                } else {
340                    Ok(self)
341                }
342            }
343            fn destroy(self) {
344                excall! { $destroy_plan(self) }
345            }
346            fn print(self) {
347                excall! { $print_plan(self) }
348            }
349        }
350    };
351} // impl_plan_spec!
352
353impl_plan_spec!(Plan64; fftw_destroy_plan, fftw_print_plan);
354impl_plan_spec!(Plan32; fftwf_destroy_plan, fftwf_print_plan);
355
356// Convert [usize] -> [i32]
357trait ToCInt {
358    fn to_cint(&self) -> Vec<i32>;
359}
360
361impl ToCInt for [usize] {
362    fn to_cint(&self) -> Vec<i32> {
363        self.iter().map(|&x| x as i32).collect()
364    }
365}
366
367fn slice_info<T>(a: &[T]) -> (usize, Alignment) {
368    (a.len(), alignment_of(a))
369}
fftw/plan.rs

fftw/
plan.rs
