mlx_rs/fft/
mod.rs

1//! Fast Fourier Transform (FFT) and its inverse (IFFT) for one, two, and `N` dimensions.
2//!
3//! Unsuffixed functions are the canonical named `Result` paths and use the current thread's scoped
4//! stream or runtime default. Select an explicit stream or device with [`crate::with_stream`] or
5//! [`crate::with_device`]. The `_device` functions and generated operation macros are deprecated
6//! forwarding shims.
7//!
8//! N-dimensional transforms accept [`FftnOptions`]. Its default selects every axis at the input
9//! lengths, and correlated lengths and axes are validated before calling MLX.
10//!
11//! # Examples
12//!
13//! ## One dimension
14//!
15//! ```rust
16//! use mlx_rs::{
17//!     Dtype, Array, complex64, fft::*,
18//!     test_utils::{assert_array_eq, tolerances},
19//! };
20//!
21//! let src = [1.0f32, 2.0, 3.0, 4.0];
22//! let mut array = Array::from_slice(&src[..], &[4]);
23//!
24//! let mut fft_result = fft(&array, 4, 0).unwrap();
25//! assert_eq!(fft_result.dtype(), Dtype::Complex64);
26//!
27//! let expected = Array::from_slice(&[
28//!     complex64::new(10.0, 0.0),
29//!     complex64::new(-2.0, 2.0),
30//!     complex64::new(-2.0, 0.0),
31//!     complex64::new(-2.0, -2.0),
32//! ], &[4]);
33//! assert_array_eq(&fft_result, expected, tolerances::EXACT.rtol, tolerances::EXACT.atol);
34//!
35//! let mut ifft_result = ifft(&fft_result, 4, 0).unwrap();
36//! assert_eq!(ifft_result.dtype(), Dtype::Complex64);
37//!
38//! let expected = Array::from_slice(&[
39//!    complex64::new(1.0, 0.0),
40//!    complex64::new(2.0, 0.0),
41//!    complex64::new(3.0, 0.0),
42//!    complex64::new(4.0, 0.0),
43//! ], &[4]);
44//! assert_array_eq(ifft_result, expected, tolerances::EXACT.rtol, tolerances::EXACT.atol);
45//!
46//! let mut rfft_result = rfft(&array, 4, 0).unwrap();
47//! assert_eq!(rfft_result.dtype(), Dtype::Complex64);
48//!
49//! let expected = Array::from_slice(&[
50//!    complex64::new(10.0, 0.0),
51//!    complex64::new(-2.0, 2.0),
52//!    complex64::new(-2.0, 0.0),
53//! ], &[3]);
54//! assert_array_eq(&rfft_result, expected, tolerances::EXACT.rtol, tolerances::EXACT.atol);
55//!
56//! let mut irfft_result = irfft(&rfft_result, 4, 0).unwrap();
57//! assert_eq!(irfft_result.dtype(), Dtype::Float32);
58//! assert_array_eq(
59//!     irfft_result,
60//!     Array::from_slice(&src, &[4]),
61//!     tolerances::EXACT.rtol,
62//!     tolerances::EXACT.atol,
63//! );
64//!
65//! // The original array is not modified
66//! assert_array_eq(
67//!     array,
68//!     Array::from_slice(&src, &[4]),
69//!     tolerances::EXACT.rtol,
70//!     tolerances::EXACT.atol,
71//! );
72//! ```
73//!
74//! ## Two dimensions
75//!
76//! ```rust
77//! use mlx_rs::{
78//!     Dtype, Array, complex64, fft::*,
79//!     test_utils::{assert_array_eq, tolerances},
80//! };
81//!
82//! let src = [1.0f32, 1.0, 1.0, 1.0];
83//! let mut array = Array::from_slice(&src[..], &[2, 2]);
84//!
85//! let mut fft2_result = fft2(&array, None, None).unwrap();
86//! assert_eq!(fft2_result.dtype(), Dtype::Complex64);
87//! let expected = Array::from_slice(&[
88//!     complex64::new(4.0, 0.0),
89//!     complex64::new(0.0, 0.0),
90//!     complex64::new(0.0, 0.0),
91//!     complex64::new(0.0, 0.0),
92//! ], &[2, 2]);
93//! assert_array_eq(&fft2_result, expected, tolerances::EXACT.rtol, tolerances::EXACT.atol);
94//!
95//! let mut ifft2_result = ifft2(&fft2_result, None, None).unwrap();
96//! assert_eq!(ifft2_result.dtype(), Dtype::Complex64);
97//!
98//! let expected = Array::from_slice(&[
99//!    complex64::new(1.0, 0.0),
100//!    complex64::new(1.0, 0.0),
101//!    complex64::new(1.0, 0.0),
102//!    complex64::new(1.0, 0.0),
103//! ], &[2, 2]);
104//! assert_array_eq(ifft2_result, expected, tolerances::EXACT.rtol, tolerances::EXACT.atol);
105//!
106//! let mut rfft2_result = rfft2(&array, None, None).unwrap();
107//! assert_eq!(rfft2_result.dtype(), Dtype::Complex64);
108//!
109//! let expected = Array::from_slice(&[
110//!     complex64::new(4.0, 0.0),
111//!     complex64::new(0.0, 0.0),
112//!     complex64::new(0.0, 0.0),
113//!     complex64::new(0.0, 0.0),
114//! ], &[2, 2]);
115//! assert_array_eq(&rfft2_result, expected, tolerances::EXACT.rtol, tolerances::EXACT.atol);
116//!
117//! let mut irfft2_result = irfft2(&rfft2_result, None, None).unwrap();
118//! assert_eq!(irfft2_result.dtype(), Dtype::Float32);
119//! assert_array_eq(
120//!     irfft2_result,
121//!     Array::from_slice(&src, &[2, 2]),
122//!     tolerances::EXACT.rtol,
123//!     tolerances::EXACT.atol,
124//! );
125//!
126//! // The original array is not modified
127//! assert_array_eq(
128//!     array,
129//!     Array::from_slice(&src, &[2, 2]),
130//!     tolerances::EXACT.rtol,
131//!     tolerances::EXACT.atol,
132//! );
133//! ```
134//!
135//! ## `N` dimensions
136//!
137//! ```rust
138//! use mlx_rs::{
139//!     Dtype, Array, complex64, fft::*,
140//!     test_utils::{assert_array_eq, tolerances},
141//! };
142//!
143//! let mut array = Array::ones::<f32>(&[2, 2, 2]).unwrap();
144//! let mut fftn_result = fftn(&array, FftnOptions::default()).unwrap();
145//! assert_eq!(fftn_result.dtype(), Dtype::Complex64);
146//!
147//! let mut expected = [complex64::new(0.0, 0.0); 8];
148//! expected[0] = complex64::new(8.0, 0.0);
149//! assert_array_eq(
150//!     &fftn_result,
151//!     Array::from_slice(&expected, &[2, 2, 2]),
152//!     tolerances::EXACT.rtol,
153//!     tolerances::EXACT.atol,
154//! );
155//!
156//! let mut ifftn_result = ifftn(&fftn_result, FftnOptions::default()).unwrap();
157//! assert_eq!(ifftn_result.dtype(), Dtype::Complex64);
158//!
159//! let expected = [complex64::new(1.0, 0.0); 8];
160//! assert_array_eq(
161//!     ifftn_result,
162//!     Array::from_slice(&expected, &[2, 2, 2]),
163//!     tolerances::EXACT.rtol,
164//!     tolerances::EXACT.atol,
165//! );
166//!
167//! let mut rfftn_result = rfftn(&array, FftnOptions::default()).unwrap();
168//! assert_eq!(rfftn_result.dtype(), Dtype::Complex64);
169//!
170//! let mut expected = [complex64::new(0.0, 0.0); 8];
171//! expected[0] = complex64::new(8.0, 0.0);
172//! assert_array_eq(
173//!     &rfftn_result,
174//!     Array::from_slice(&expected, &[2, 2, 2]),
175//!     tolerances::EXACT.rtol,
176//!     tolerances::EXACT.atol,
177//! );
178//!
179//! let mut irfftn_result = irfftn(&rfftn_result, FftnOptions::default()).unwrap();
180//! assert_eq!(irfftn_result.dtype(), Dtype::Float32);
181//!
182//! let expected = [1.0; 8];
183//! assert_array_eq(
184//!     irfftn_result,
185//!     Array::from_slice(&expected, &[2, 2, 2]),
186//!     tolerances::EXACT.rtol,
187//!     tolerances::EXACT.atol,
188//! );
189//!
190//! // The original array is not modified
191//! assert_array_eq(
192//!     array,
193//!     Array::from_slice(&[1.0; 8], &[2, 2, 2]),
194//!     tolerances::EXACT.rtol,
195//!     tolerances::EXACT.atol,
196//! );
197//! ```
198
199mod fftn;
200mod frequencies;
201mod options;
202mod rfftn;
203mod shift;
204mod utils;
205
206pub use self::{fftn::*, frequencies::*, options::*, rfftn::*, shift::*};
207
208/* -------------------------------------------------------------------------- */
209/*                              Helper functions                              */
210/* -------------------------------------------------------------------------- */
mlx_rs/fft/mod.rs

mlx_rs/fft/
mod.rs
