rlst 0.8.0

//! Definition of a complex-to-complex inplace FFT plan

use std::{
    ffi::{c_int, c_void},
    marker::PhantomData,
};

use crate::{
    BaseItem, RlstScalar, Shape, StableRawAccessMut, Stride, c32, c64,
    fftw::{FftwPlanFlags, FftwPlanPtrType, FftwPlanPtrTypeTrait, PlanInplace},
};

use itertools::{Itertools, izip};

use super::{FFTW_PLAN_INTERFACE, Fft, FftwPlanPtr, PlanInterfaceTrait};

/// The direction of the C2C FFT
#[derive(Copy, Clone)]
#[repr(i32)]
pub enum C2CDirection {
    /// Denotes the forward transform
    Forward = -1,
    /// Denotes the backward (inverse) transform
    Backward = 1,
}

/// Create a C2C FFT Plan
#[allow(private_bounds)]
pub trait C2CPlan<Item, ArrayImpl, const NDIM: usize>
where
    ArrayImpl: BaseItem<Item = Item>,
    Item: RlstScalar + 'static,
    FftwPlanPtrType<Item::Real>: FftwPlanPtrTypeTrait,
{
    /// Create an in-place complex-to-complex plan.
    ///
    /// Returns `None` for empty or noncontiguous arrays, ranks above three, or
    /// when FFTW cannot construct the requested plan. Planning flags other than
    /// [`FftwPlanFlags::Estimate`] may overwrite the array while benchmarking.
    /// FFTW's backward transform is unnormalized. For details see the
    /// [FFTW module description](crate::fftw).
    fn c2c_fft(
        &self,
        flags: FftwPlanFlags,
        direction: C2CDirection,
    ) -> Option<PlanInplace<Item, ArrayImpl, NDIM>>;
}

macro_rules! impl_c2c_inplace {
    ($ty: ty, $precision: ident) => {
        impl<ArrayImpl, const NDIM: usize> C2CPlan<$ty, ArrayImpl, NDIM> for Fft<ArrayImpl, NDIM>
        where
            ArrayImpl: Shape<NDIM> + Stride<NDIM> + StableRawAccessMut<Item = $ty>,
        {
            /// Create a new plan wrapped inside an array structure. Returns None if the Plan cannot be created.
            fn c2c_fft(
                &self,
                flags: FftwPlanFlags,
                direction: C2CDirection,
            ) -> Option<PlanInplace<$ty, ArrayImpl, NDIM>> {
                let mut guard = self.borrow_mut().unwrap();
                let arr = guard.array_mut();

                // Return None if NDIM > 3 as FFTW only supports up to three dimensions
                if NDIM > 3 {
                    return None;
                }
                // Return None if the array is not contiguous or is empty
                if !arr.is_contiguous() || arr.is_empty() {
                    return None;
                }

                let mut shape = arr
                    .shape()
                    .iter()
                    .map(|&elem| elem as i32)
                    .collect_array()
                    .unwrap();

                // If the array is column-major reverse the order of the shape as
                // FFTW uses row-major order

                if matches!(arr.memory_layout(), crate::MemoryLayout::ColumnMajor) {
                    shape = {
                        let mut rev_shape = [0_i32; NDIM];

                        for (&orig_index, rev_index) in
                            izip!(shape.iter().rev(), rev_shape.iter_mut())
                        {
                            *rev_index = orig_index;
                        }
                        rev_shape
                    };
                }

                let data = arr.data_mut()?.as_mut_ptr();

                let plan_ptr = unsafe {
                    FFTW_PLAN_INTERFACE.lock().unwrap().$precision.create_c2c(
                        NDIM as c_int,
                        shape.as_ptr(),
                        data,
                        data,
                        direction as i32,
                        flags as u32,
                    )
                };

                if plan_ptr.is_null() {
                    None
                } else {
                    Some(PlanInplace::new(
                        self.shared_ref.clone(),
                        FftwPlanPtr {
                            ptr: plan_ptr as *mut c_void,
                            _marker: PhantomData,
                        },
                    ))
                }
            }
        }
    };
}

impl_c2c_inplace!(c32, fftw_single);
impl_c2c_inplace!(c64, fftw_double);

#[cfg(test)]
mod test {

    use num::{Complex, FromPrimitive};
    use std::f64::consts::PI;

    use crate::fftw::c2c_inplace::C2CDirection;
    use crate::{EvaluateArray, c64};

    use crate::DynArray;
    use crate::fftw::{FftPlanInplace, FftwPlanFlags};

    use super::*;

    #[test]
    fn test_plan_keeps_storage_while_elements_are_mutated() {
        let input = [
            Complex::new(1.0, 0.5),
            Complex::new(-2.0, 1.0),
            Complex::new(0.25, -0.75),
            Complex::new(3.0, 2.0),
        ];
        let fft = DynArray::<c64, _>::from_shape([input.len()]).into_fft();
        let mut plan = fft
            .c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Forward)
            .unwrap();

        fft.borrow_mut()
            .unwrap()
            .data_mut()
            .unwrap()
            .copy_from_slice(&input);
        plan.execute();

        let actual = fft.borrow().unwrap().eval();
        for k in 0..input.len() {
            let expected = input
                .iter()
                .enumerate()
                .map(|(j, value)| {
                    let angle = -2.0 * PI * (j * k) as f64 / input.len() as f64;
                    value * Complex::new(angle.cos(), angle.sin())
                })
                .sum::<c64>();
            approx::assert_relative_eq!(actual[[k]], expected, epsilon = 1E-10);
        }

        // The same plan remains valid after another element-only mutation.
        fft.borrow_mut()
            .unwrap()
            .data_mut()
            .unwrap()
            .copy_from_slice(&input);
        plan.execute();
        let repeated = fft.borrow().unwrap().eval();
        crate::assert_array_relative_eq!(actual.r(), repeated.r(), 1E-10);
    }

    #[test]
    fn test_c2c_1d_f64() {
        let n = 10;
        let mut original = DynArray::<c64, _>::from_shape([10]);
        original.fill_from_seed_equally_distributed(0);

        let mut copied = original.eval();
        let sum = copied.iter_value().sum::<c64>();

        copied
            .r_mut()
            .into_fft()
            .c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Forward)
            .unwrap()
            .execute();

        approx::assert_relative_eq!(sum, copied[[0]], epsilon = 1E-10);

        let fft = original.eval().into_fft();

        fft.c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Forward)
            .unwrap()
            .execute();

        fft.c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Backward)
            .unwrap()
            .execute();

        let actual = fft.borrow().unwrap().eval();

        // Check that the zeroth frequency is the sum of all input elements.
        // Compute the forward and inverse FFT and check that the result is still the same.
        crate::assert_array_relative_eq!(
            c64::from_usize(n).unwrap() * original.r(),
            actual.r(),
            1E-10
        );
    }

    #[test]
    fn test_c2c_1d_f32() {
        let n = 10;
        let mut original = DynArray::<c32, _>::from_shape([10]);
        original.fill_from_seed_equally_distributed(0);

        let mut copied = original.eval();
        let sum = copied.iter_value().sum::<c32>();

        copied
            .r_mut()
            .into_fft()
            .c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Forward)
            .unwrap()
            .execute();

        approx::assert_relative_eq!(sum, copied[[0]], epsilon = 1E-10);

        let fft = original.eval().into_fft();

        fft.c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Forward)
            .unwrap()
            .execute();

        fft.c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Backward)
            .unwrap()
            .execute();

        let actual = fft.borrow().unwrap().eval();

        // Check that the zeroth frequency is the sum of all input elements.
        // Compute the forward and inverse FFT and check that the result is still the same.
        crate::assert_array_relative_eq!(
            c32::from_usize(n).unwrap() * original.r(),
            actual.r(),
            1E-5
        );
    }

    #[test]
    fn test_c2c_2d_f64() {
        let n = 15;
        let mut original = DynArray::<c64, _>::from_shape([3, 5]);
        original.fill_from_seed_equally_distributed(0);

        let mut copied = original.eval();
        let sum = copied.iter_value().sum::<c64>();

        copied
            .r_mut()
            .into_fft()
            .c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Forward)
            .unwrap()
            .execute();

        approx::assert_relative_eq!(sum, copied[[0, 0]], epsilon = 1E-10);

        let fft = original.eval().into_fft();

        fft.c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Forward)
            .unwrap()
            .execute();

        fft.c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Backward)
            .unwrap()
            .execute();

        let actual = fft.borrow().unwrap().eval();

        // Check that the zeroth frequency is the sum of all input elements.
        // Compute the forward and inverse FFT and check that the result is still the same.
        crate::assert_array_relative_eq!(
            c64::from_usize(n).unwrap() * original.r(),
            actual.r(),
            1E-10
        );
    }

    #[test]
    fn test_c2c_2d_f32() {
        let n = 15;
        let mut original = DynArray::<c32, _>::from_shape([3, 5]);
        original.fill_from_seed_equally_distributed(0);

        let mut copied = original.eval();
        let sum = copied.iter_value().sum::<c32>();

        copied
            .r_mut()
            .into_fft()
            .c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Forward)
            .unwrap()
            .execute();

        approx::assert_relative_eq!(sum, copied[[0, 0]], epsilon = 1E-10);

        let fft = original.eval().into_fft();

        fft.c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Forward)
            .unwrap()
            .execute();

        fft.c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Backward)
            .unwrap()
            .execute();

        let actual = fft.borrow().unwrap().eval();

        // Check that the zeroth frequency is the sum of all input elements.
        // Compute the forward and inverse FFT and check that the result is still the same.
        crate::assert_array_relative_eq!(
            c32::from_usize(n).unwrap() * original.r(),
            actual.r(),
            1E-5
        );
    }

    #[test]
    fn test_c2c_3d_f64() {
        let n = 30;
        let mut original = DynArray::<c64, _>::from_shape([3, 5, 2]);
        original.fill_from_seed_equally_distributed(0);

        let mut copied = original.eval();
        let sum = copied.iter_value().sum::<c64>();

        copied
            .r_mut()
            .into_fft()
            .c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Forward)
            .unwrap()
            .execute();

        approx::assert_relative_eq!(sum, copied[[0, 0, 0]], epsilon = 1E-10);

        let fft = original.eval().into_fft();

        fft.c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Forward)
            .unwrap()
            .execute();

        fft.c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Backward)
            .unwrap()
            .execute();

        let actual = fft.borrow().unwrap().eval();

        // Check that the zeroth frequency is the sum of all input elements.
        // Compute the forward and inverse FFT and check that the result is still the same.
        crate::assert_array_relative_eq!(
            c64::from_usize(n).unwrap() * original.r(),
            actual.r(),
            1E-10
        );
    }

    #[test]
    fn test_c2c_3d_f32() {
        let n = 30;
        let mut original = DynArray::<c32, _>::from_shape([3, 5, 2]);
        original.fill_from_seed_equally_distributed(0);

        let mut copied = original.eval();
        let sum = copied.iter_value().sum::<c32>();

        copied
            .r_mut()
            .into_fft()
            .c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Forward)
            .unwrap()
            .execute();

        approx::assert_relative_eq!(sum, copied[[0, 0, 0]], epsilon = 1E-10);

        let fft = original.eval().into_fft();

        fft.c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Forward)
            .unwrap()
            .execute();

        fft.c2c_fft(FftwPlanFlags::Estimate, C2CDirection::Backward)
            .unwrap()
            .execute();

        let actual = fft.borrow().unwrap().eval();

        // Check that the zeroth frequency is the sum of all input elements.
        // Compute the forward and inverse FFT and check that the result is still the same.
        crate::assert_array_relative_eq!(
            c32::from_usize(n).unwrap() * original.r(),
            actual.r(),
            1E-5
        );
    }
}