mlx_rs/fft/
rfftn.rs

1use mlx_internal_macros::generate_macro;
2
3use crate::{
4    error::Result,
5    utils::{guard::Guarded, IntoOption},
6    with_stream, Array, Stream,
7};
8
9use super::{
10    utils::{
11        legacy_fftn_options, require_real_axis, resolve_inverse_real_length,
12        resolve_lengths_and_axes, resolve_size_and_axis_unchecked,
13        resolve_sizes_and_axes_unchecked,
14    },
15    FftnOptions,
16};
17
18/// One dimensional discrete Fourier Transform on a real input.
19///
20/// The output has the same shape as the input except along `axis` in which case it has size `n // 2
21/// + 1`.
22///
23/// # Params
24///
25/// - `a`: The input array. If the array is complex it will be silently cast to a real type.
26/// - `n`: Size of the transformed axis. The corresponding axis in the input is truncated or padded
27///   with zeros to match `n`. The default value is `a.shape[axis]` if not specified.
28/// - `axis`: Axis along which to perform the FFT. The default is `-1` if not specified.
29pub fn rfft(
30    a: impl AsRef<Array>,
31    n: impl Into<Option<i32>>,
32    axis: impl Into<Option<i32>>,
33) -> Result<Array> {
34    let a = a.as_ref();
35    let (n, axis) = resolve_size_and_axis_unchecked(a, n.into(), axis.into());
36    let stream = Stream::thread_local_or_default();
37    Array::try_from_op(|res| unsafe {
38        mlx_sys::mlx_fft_rfft(
39            res,
40            a.as_ptr(),
41            n,
42            axis,
43            mlx_sys::mlx_fft_norm__MLX_FFT_NORM_BACKWARD,
44            stream.as_ref().as_ptr(),
45        )
46    })
47}
48
49/// Compatibility shim for [`rfft`].
50#[generate_macro(customize(forwarding_shim = true, root = "$crate::fft"))]
51#[deprecated(
52    since = "0.26.0",
53    note = "use `with_stream` or `with_device` around `rfft`"
54)]
55pub fn rfft_device(
56    a: impl AsRef<Array>,
57    #[optional] n: impl Into<Option<i32>>,
58    #[optional] axis: impl Into<Option<i32>>,
59    #[optional] stream: impl AsRef<Stream>,
60) -> Result<Array> {
61    with_stream(stream.as_ref(), || rfft(a, n, axis))
62}
63
64/// Two-dimensional real discrete Fourier Transform.
65///
66/// The output has the same shape as the input except along the dimensions in `axes` in which case
67/// it has sizes from `s`. The last axis in `axes` is treated as the real axis and will have size
68/// `s[s.len()-1] // 2 + 1`.
69///
70/// # Params
71///
72/// - `a`: The input array. If the array is complex it will be silently cast to a real type.
73/// - `s`: Sizes of the transformed axes. The corresponding axes in the input are truncated or
74///   padded with zeros to match `s`. The default value is the sizes of `a` along `axes`.
75/// - `axes`: Axes along which to perform the FFT. The default is `[-2, -1]`.
76pub fn rfft2<'a>(
77    a: impl AsRef<Array>,
78    s: impl IntoOption<&'a [i32]>,
79    axes: impl IntoOption<&'a [i32]>,
80) -> Result<Array> {
81    let a = a.as_ref();
82    let axes = axes.into_option().unwrap_or(&[-2, -1]);
83    let (s, axes) = resolve_sizes_and_axes_unchecked(a, s.into_option(), Some(axes));
84    require_real_axis(&axes)?;
85
86    let num_s = s.len();
87    let num_axes = axes.len();
88
89    let s_ptr = s.as_ptr();
90    let axes_ptr = axes.as_ptr();
91    let stream = Stream::thread_local_or_default();
92
93    Array::try_from_op(|res| unsafe {
94        mlx_sys::mlx_fft_rfft2(
95            res,
96            a.as_ptr(),
97            s_ptr,
98            num_s,
99            axes_ptr,
100            num_axes,
101            mlx_sys::mlx_fft_norm__MLX_FFT_NORM_BACKWARD,
102            stream.as_ref().as_ptr(),
103        )
104    })
105}
106
107/// Compatibility shim for [`rfft2`].
108#[generate_macro(customize(forwarding_shim = true, root = "$crate::fft"))]
109#[deprecated(
110    since = "0.26.0",
111    note = "use `with_stream` or `with_device` around `rfft2`"
112)]
113pub fn rfft2_device<'a>(
114    a: impl AsRef<Array>,
115    #[optional] s: impl IntoOption<&'a [i32]>,
116    #[optional] axes: impl IntoOption<&'a [i32]>,
117    #[optional] stream: impl AsRef<Stream>,
118) -> Result<Array> {
119    with_stream(stream.as_ref(), || rfft2(a, s, axes))
120}
121
122/// n-dimensional real discrete Fourier Transform.
123///
124/// The output has the same shape as the input except along the dimensions in `axes` in which case
125/// it has sizes from `s`. The last axis in `axes` is treated as the real axis and will have size
126/// `s[s.len()-1] // 2 + 1`.
127///
128/// # Params
129///
130/// - `a`: The input array. If the array is complex it will be silently cast to a real type.
131/// - `options`: Transform lengths and axes. The default transforms all axes at their input sizes.
132pub fn rfftn(a: impl AsRef<Array>, options: FftnOptions) -> Result<Array> {
133    let a = a.as_ref();
134    let (s, axes) = resolve_lengths_and_axes(a.shape(), options.lengths.as_deref(), &options.axes)?;
135    require_real_axis(&axes)?;
136
137    let num_s = s.len();
138    let num_axes = axes.len();
139
140    let s_ptr = s.as_ptr();
141    let axes_ptr = axes.as_ptr();
142    let stream = Stream::thread_local_or_default();
143
144    Array::try_from_op(|res| unsafe {
145        mlx_sys::mlx_fft_rfftn(
146            res,
147            a.as_ptr(),
148            s_ptr,
149            num_s,
150            axes_ptr,
151            num_axes,
152            mlx_sys::mlx_fft_norm__MLX_FFT_NORM_BACKWARD,
153            stream.as_ref().as_ptr(),
154        )
155    })
156}
157
158/// Compatibility shim for [`rfftn`].
159#[generate_macro(customize(forwarding_shim = true, root = "$crate::fft"))]
160#[deprecated(
161    since = "0.26.0",
162    note = "use `with_stream` or `with_device` around `rfftn` with `FftnOptions`"
163)]
164pub fn rfftn_device<'a>(
165    a: impl AsRef<Array>,
166    #[optional] s: impl IntoOption<&'a [i32]>,
167    #[optional] axes: impl IntoOption<&'a [i32]>,
168    #[optional] stream: impl AsRef<Stream>,
169) -> Result<Array> {
170    let options = legacy_fftn_options(s.into_option(), axes.into_option())?;
171    with_stream(stream.as_ref(), || rfftn(a, options))
172}
173
174/// The inverse of [`rfft()`].
175///
176/// The output has the same shape as the input except along axis in which case it has size n.
177///
178/// # Params
179///
180/// - `a`: The input array.
181/// - `n`: Size of the transformed axis. The corresponding axis in the input is truncated or padded
182///   with zeros to match `n // 2 + 1`. The default value is `a.shape[axis] // 2 + 1`.
183/// - `axis`: Axis along which to perform the FFT. The default is `-1`.
184pub fn irfft(
185    a: impl AsRef<Array>,
186    n: impl Into<Option<i32>>,
187    axis: impl Into<Option<i32>>,
188) -> Result<Array> {
189    let a = a.as_ref();
190    let n = n.into();
191    let axis = axis.into();
192    let modify_n = n.is_none();
193    let (mut n, axis) = resolve_size_and_axis_unchecked(a, n, axis);
194    if modify_n {
195        n = resolve_inverse_real_length(n)?;
196    }
197    let stream = Stream::thread_local_or_default();
198
199    Array::try_from_op(|res| unsafe {
200        mlx_sys::mlx_fft_irfft(
201            res,
202            a.as_ptr(),
203            n,
204            axis,
205            mlx_sys::mlx_fft_norm__MLX_FFT_NORM_BACKWARD,
206            stream.as_ref().as_ptr(),
207        )
208    })
209}
210
211/// Compatibility shim for [`irfft`].
212#[generate_macro(customize(forwarding_shim = true, root = "$crate::fft"))]
213#[deprecated(
214    since = "0.26.0",
215    note = "use `with_stream` or `with_device` around `irfft`"
216)]
217pub fn irfft_device(
218    a: impl AsRef<Array>,
219    #[optional] n: impl Into<Option<i32>>,
220    #[optional] axis: impl Into<Option<i32>>,
221    #[optional] stream: impl AsRef<Stream>,
222) -> Result<Array> {
223    with_stream(stream.as_ref(), || irfft(a, n, axis))
224}
225
226/// The inverse of [`rfft2()`].
227///
228/// Note the input is generally complex. The dimensions of the input specified in `axes` are padded
229/// or truncated to match the sizes from `s`. The last axis in `axes` is treated as the real axis
230/// and will have size `s[s.len()-1] // 2 + 1`.
231///
232/// # Params
233///
234/// - `a`: The input array.
235/// - `s`: Sizes of the transformed axes. The corresponding axes in the input are truncated or
236///   padded with zeros to match the sizes in `s` except for the last axis which has size
237///   `s[s.len()-1] // 2 + 1`. The default value is the sizes of `a` along `axes`.
238/// - `axes`: Axes along which to perform the FFT. The default is `[-2, -1]`.
239pub fn irfft2<'a>(
240    a: impl AsRef<Array>,
241    s: impl IntoOption<&'a [i32]>,
242    axes: impl IntoOption<&'a [i32]>,
243) -> Result<Array> {
244    let a = a.as_ref();
245    let s = s.into_option();
246    let axes = axes.into_option().unwrap_or(&[-2, -1]);
247    let modify_last_axis = s.is_none();
248
249    let (mut s, axes) = resolve_sizes_and_axes_unchecked(a, s, Some(axes));
250    require_real_axis(&axes)?;
251    if modify_last_axis {
252        let end = s.len() - 1;
253        s[end] = resolve_inverse_real_length(s[end])?;
254    }
255
256    let num_s = s.len();
257    let num_axes = axes.len();
258
259    let s_ptr = s.as_ptr();
260    let axes_ptr = axes.as_ptr();
261    let stream = Stream::thread_local_or_default();
262
263    Array::try_from_op(|res| unsafe {
264        mlx_sys::mlx_fft_irfft2(
265            res,
266            a.as_ptr(),
267            s_ptr,
268            num_s,
269            axes_ptr,
270            num_axes,
271            mlx_sys::mlx_fft_norm__MLX_FFT_NORM_BACKWARD,
272            stream.as_ref().as_ptr(),
273        )
274    })
275}
276
277/// Compatibility shim for [`irfft2`].
278#[generate_macro(customize(forwarding_shim = true, root = "$crate::fft"))]
279#[deprecated(
280    since = "0.26.0",
281    note = "use `with_stream` or `with_device` around `irfft2`"
282)]
283pub fn irfft2_device<'a>(
284    a: impl AsRef<Array>,
285    #[optional] s: impl IntoOption<&'a [i32]>,
286    #[optional] axes: impl IntoOption<&'a [i32]>,
287    #[optional] stream: impl AsRef<Stream>,
288) -> Result<Array> {
289    with_stream(stream.as_ref(), || irfft2(a, s, axes))
290}
291
292/// The inverse of [`rfftn()`].
293///
294/// Note the input is generally complex. The dimensions of the input specified in `axes` are padded
295/// or truncated to match the sizes from `s`. The last axis in `axes` is treated as the real axis
296/// and will have size `s[s.len()-1] // 2 + 1`.
297///
298/// # Params
299///
300/// - `a`: The input array.
301/// - `options`: Transform lengths and axes. The default transforms all axes at their input sizes.
302pub fn irfftn(a: impl AsRef<Array>, options: FftnOptions) -> Result<Array> {
303    let a = a.as_ref();
304    let modify_last_axis = options.lengths.is_none();
305
306    let (mut s, axes) =
307        resolve_lengths_and_axes(a.shape(), options.lengths.as_deref(), &options.axes)?;
308    require_real_axis(&axes)?;
309    if modify_last_axis {
310        let end = s.len() - 1;
311        s[end] = resolve_inverse_real_length(s[end])?;
312    }
313
314    let num_s = s.len();
315    let num_axes = axes.len();
316
317    let s_ptr = s.as_ptr();
318    let axes_ptr = axes.as_ptr();
319    let stream = Stream::thread_local_or_default();
320
321    Array::try_from_op(|res| unsafe {
322        mlx_sys::mlx_fft_irfftn(
323            res,
324            a.as_ptr(),
325            s_ptr,
326            num_s,
327            axes_ptr,
328            num_axes,
329            mlx_sys::mlx_fft_norm__MLX_FFT_NORM_BACKWARD,
330            stream.as_ref().as_ptr(),
331        )
332    })
333}
334
335/// Compatibility shim for [`irfftn`].
336#[generate_macro(customize(forwarding_shim = true, root = "$crate::fft"))]
337#[deprecated(
338    since = "0.26.0",
339    note = "use `with_stream` or `with_device` around `irfftn` with `FftnOptions`"
340)]
341pub fn irfftn_device<'a>(
342    a: impl AsRef<Array>,
343    #[optional] s: impl IntoOption<&'a [i32]>,
344    #[optional] axes: impl IntoOption<&'a [i32]>,
345    #[optional] stream: impl AsRef<Stream>,
346) -> Result<Array> {
347    let options = legacy_fftn_options(s.into_option(), axes.into_option())?;
348    with_stream(stream.as_ref(), || irfftn(a, options))
349}
350
351#[cfg(test)]
352mod tests {
353    use crate::{
354        complex64,
355        ops::indexing::TryIndexOp,
356        test_utils::{assert_array_eq, tolerances},
357        Array, Axes, Dtype,
358    };
359
360    #[test]
361    fn test_rfft() {
362        const RFFT_DATA: &[f32] = &[1.0, 2.0, 3.0, 4.0];
363        const RFFT_N: i32 = 4;
364        const RFFT_SHAPE: &[i32] = &[RFFT_N];
365        const RFFT_AXIS: i32 = -1;
366        const RFFT_EXPECTED: &[complex64] = &[
367            complex64::new(10.0, 0.0),
368            complex64::new(-2.0, 2.0),
369            complex64::new(-2.0, 0.0),
370        ];
371
372        let a = Array::from_slice(RFFT_DATA, RFFT_SHAPE);
373        let rfft = super::rfft(&a, RFFT_N, RFFT_AXIS).unwrap();
374        assert_eq!(rfft.dtype(), Dtype::Complex64);
375        assert_array_eq(
376            &rfft,
377            Array::from_slice(RFFT_EXPECTED, &[3]),
378            tolerances::EXACT.rtol,
379            tolerances::EXACT.atol,
380        );
381
382        let irfft = super::irfft(&rfft, RFFT_N, RFFT_AXIS).unwrap();
383        assert_eq!(irfft.dtype(), Dtype::Float32);
384        assert_array_eq(
385            irfft,
386            Array::from_slice(RFFT_DATA, RFFT_SHAPE),
387            tolerances::EXACT.rtol,
388            tolerances::EXACT.atol,
389        );
390    }
391
392    #[test]
393    fn test_rfft_shape_with_default_params() {
394        const IN_N: i32 = 8;
395        const OUT_N: i32 = IN_N / 2 + 1;
396
397        let a = Array::ones::<f32>(&[IN_N]).unwrap();
398        let rfft = super::rfft(&a, None, None).unwrap();
399        assert_eq!(rfft.shape(), &[OUT_N]);
400    }
401
402    #[test]
403    fn test_irfft_shape_with_default_params() {
404        const IN_N: i32 = 8;
405        const OUT_N: i32 = (IN_N - 1) * 2;
406
407        let a = Array::ones::<f32>(&[IN_N]).unwrap();
408        let irfft = super::irfft(&a, None, None).unwrap();
409        assert_eq!(irfft.shape(), &[OUT_N]);
410    }
411
412    #[test]
413    fn test_rfft2() {
414        const RFFT2_DATA: &[f32] = &[1.0; 4];
415        const RFFT2_SHAPE: &[i32] = &[2, 2];
416        const RFFT2_EXPECTED: &[complex64] = &[
417            complex64::new(4.0, 0.0),
418            complex64::new(0.0, 0.0),
419            complex64::new(0.0, 0.0),
420            complex64::new(0.0, 0.0),
421        ];
422
423        let a = Array::from_slice(RFFT2_DATA, RFFT2_SHAPE);
424        let rfft2 = super::rfft2(&a, None, None).unwrap();
425        assert_eq!(rfft2.dtype(), Dtype::Complex64);
426        assert_array_eq(
427            &rfft2,
428            Array::from_slice(RFFT2_EXPECTED, RFFT2_SHAPE),
429            tolerances::EXACT.rtol,
430            tolerances::EXACT.atol,
431        );
432
433        let irfft2 = super::irfft2(&rfft2, None, None).unwrap();
434        assert_eq!(irfft2.dtype(), Dtype::Float32);
435        assert_array_eq(
436            irfft2,
437            Array::from_slice(RFFT2_DATA, RFFT2_SHAPE),
438            tolerances::EXACT.rtol,
439            tolerances::EXACT.atol,
440        );
441    }
442
443    #[test]
444    fn test_rfft2_shape_with_default_params() {
445        const IN_SHAPE: &[i32] = &[6, 6];
446        const OUT_SHAPE: &[i32] = &[6, 6 / 2 + 1];
447
448        let a = Array::ones::<f32>(IN_SHAPE).unwrap();
449        let rfft2 = super::rfft2(&a, None, None).unwrap();
450        assert_eq!(rfft2.shape(), OUT_SHAPE);
451    }
452
453    #[test]
454    fn test_irfft2_shape_with_default_params() {
455        const IN_SHAPE: &[i32] = &[6, 6];
456        const OUT_SHAPE: &[i32] = &[6, (6 - 1) * 2];
457
458        let a = Array::ones::<f32>(IN_SHAPE).unwrap();
459        let irfft2 = super::irfft2(&a, None, None).unwrap();
460        assert_eq!(irfft2.shape(), OUT_SHAPE);
461    }
462
463    #[test]
464    fn test_rfftn() {
465        const RFFTN_DATA: &[f32] = &[1.0; 8];
466        const RFFTN_SHAPE: &[i32] = &[2, 2, 2];
467        const RFFTN_EXPECTED: &[complex64] = &[
468            complex64::new(8.0, 0.0),
469            complex64::new(0.0, 0.0),
470            complex64::new(0.0, 0.0),
471            complex64::new(0.0, 0.0),
472            complex64::new(0.0, 0.0),
473            complex64::new(0.0, 0.0),
474            complex64::new(0.0, 0.0),
475            complex64::new(0.0, 0.0),
476        ];
477
478        let a = Array::from_slice(RFFTN_DATA, RFFTN_SHAPE);
479        let rfftn = super::rfftn(&a, super::FftnOptions::default()).unwrap();
480        assert_eq!(rfftn.dtype(), Dtype::Complex64);
481        assert_array_eq(
482            &rfftn,
483            Array::from_slice(RFFTN_EXPECTED, RFFTN_SHAPE),
484            tolerances::EXACT.rtol,
485            tolerances::EXACT.atol,
486        );
487
488        let irfftn = super::irfftn(&rfftn, super::FftnOptions::default()).unwrap();
489        assert_eq!(irfftn.dtype(), Dtype::Float32);
490        assert_array_eq(
491            irfftn,
492            Array::from_slice(RFFTN_DATA, RFFTN_SHAPE),
493            tolerances::EXACT.rtol,
494            tolerances::EXACT.atol,
495        );
496    }
497
498    #[test]
499    fn asymmetric_real_multidimensional_transforms_use_backward_normalization() {
500        let input2 = Array::from_slice(&[1.0_f32, 2.0, 3.0, 5.0, 7.0, 11.0], &[2, 3]);
501        let spectrum2 = super::rfft2(&input2, None, None).unwrap();
502        assert_array_eq(
503            spectrum2.try_index((0, 0)).unwrap(),
504            Array::from(complex64::new(29.0, 0.0)),
505            tolerances::EXACT.rtol,
506            tolerances::EXACT.atol,
507        );
508        let roundtrip2 = super::irfft2(&spectrum2, &[2, 3], &[-2, -1]).unwrap();
509        assert_array_eq(
510            &roundtrip2,
511            &input2,
512            tolerances::STANDARD.rtol,
513            tolerances::STANDARD.atol,
514        );
515
516        let inputn = Array::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], &[2, 2, 2]);
517        let spectrumn = super::rfftn(&inputn, super::FftnOptions::default()).unwrap();
518        assert_array_eq(
519            spectrumn.try_index((0, 0, 0)).unwrap(),
520            Array::from(complex64::new(36.0, 0.0)),
521            tolerances::EXACT.rtol,
522            tolerances::EXACT.atol,
523        );
524        let roundtripn = super::irfftn(
525            &spectrumn,
526            super::FftnOptions {
527                lengths: Some(vec![2, 2, 2]),
528                axes: [0, 1, 2].into(),
529            },
530        )
531        .unwrap();
532        assert_array_eq(
533            &roundtripn,
534            &inputn,
535            tolerances::STANDARD.rtol,
536            tolerances::STANDARD.atol,
537        );
538    }
539
540    #[test]
541    fn test_fftn_shape_with_default_params() {
542        const IN_SHAPE: &[i32] = &[6, 6, 6];
543        const OUT_SHAPE: &[i32] = &[6, 6, 6 / 2 + 1];
544
545        let a = Array::ones::<f32>(IN_SHAPE).unwrap();
546        let rfftn = super::rfftn(&a, super::FftnOptions::default()).unwrap();
547        assert_eq!(rfftn.shape(), OUT_SHAPE);
548    }
549
550    #[test]
551    fn test_irfftn_shape_with_default_params() {
552        const IN_SHAPE: &[i32] = &[6, 6, 6];
553        const OUT_SHAPE: &[i32] = &[6, 6, (6 - 1) * 2];
554
555        let a = Array::ones::<f32>(IN_SHAPE).unwrap();
556        let irfftn = super::irfftn(&a, super::FftnOptions::default()).unwrap();
557        assert_eq!(irfftn.shape(), OUT_SHAPE);
558    }
559
560    #[test]
561    fn real_nd_transforms_reject_empty_axes_before_calling_mlx() {
562        let input = Array::ones::<f32>(&[2, 2]).unwrap();
563
564        let rfft2_error = super::rfft2(&input, None, &[]).unwrap_err();
565        assert!(rfft2_error.what().contains("requires at least one axis"));
566        let irfft2_error = super::irfft2(&input, None, &[]).unwrap_err();
567        assert!(irfft2_error.what().contains("requires at least one axis"));
568
569        for lengths in [None, Some(vec![])] {
570            let options = super::FftnOptions {
571                lengths,
572                axes: Axes::Axes(vec![]),
573            };
574            let rfft_error = super::rfftn(&input, options.clone()).unwrap_err();
575            assert!(rfft_error.what().contains("requires at least one axis"));
576
577            let irfft_error = super::irfftn(&input, options).unwrap_err();
578            assert!(irfft_error.what().contains("requires at least one axis"));
579        }
580    }
581}
mlx_rs/fft/rfftn.rs

mlx_rs/fft/
rfftn.rs
