nmr-schedule 0.2.1

//! Algorithms for reconstructing NMR signals.
//!
//! Supports FFT for uniform data, IST for nonuniform data, and quadrature.

use alloc::{boxed::Box, vec::Vec};
use ndarray::Ix1;
use rustfft::{
    FftPlanner,
    num_complex::{Complex, ComplexFloat},
};

use crate::{ComplexSequence, DisplayMode, Schedule};

/// Reconstructs signals from NMR data.
pub trait Reconstructor {
    /// Reconstruct data represented as a list of floating point values. Converts to complex numbers by setting the imaginary component to zero. Your data will have mirrored quadrature peaks.
    ///
    /// # Panics
    ///
    /// `out` must have the same length as `data`.
    fn reconstruct_from_reals(&mut self, data: &[f64], out: &mut [Complex<f64>]) {
        assert_eq!(data.len(), out.len());

        out.iter_mut()
            .zip(data)
            .for_each(|(complex, real)| *complex = Complex::new(*real, 0.));

        self.reconstruct(out);
    }

    /// Reconstruct a complex signal in place.
    ///
    /// This method should *not* be used for quadrature blindly. See [`Reconstructor::quadrature_reconstruct`].
    fn reconstruct(&mut self, data: &mut [Complex<f64>]);

    /// Reconstruct a spectrum from `cos` and `sin` signals using quadrature.
    ///
    /// This method adds `cos` and `-i·sin` componentwise, then calls `reconstruct` on the resulting list.
    ///
    /// # Panics
    ///
    /// `cos`, `sin`, and `out` must all have the same length.
    fn quadrature_reconstruct(&mut self, cos: &[f64], sin: &[f64], out: &mut [Complex<f64>]) {
        assert_eq!(cos.len(), sin.len());
        assert_eq!(sin.len(), out.len());

        out.iter_mut()
            .zip(cos.iter())
            .zip(sin.iter())
            .for_each(|((complex, cos), sin)| {
                *complex = Complex::new(*cos, -sin);
            });

        self.reconstruct(out);
    }
}

/// FFT reconstruction for uniformly sampled data.
pub struct Fft {
    fft_scratch: Vec<Complex<f64>>,
}

impl core::fmt::Debug for Fft {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("Fft").finish()
    }
}

impl Fft {
    /// Create a new `Fft` instance
    pub const fn new() -> Fft {
        Fft {
            fft_scratch: Vec::new(),
        }
    }
}

impl Default for Fft {
    fn default() -> Self {
        Self::new()
    }
}

impl Reconstructor for Fft {
    fn reconstruct(&mut self, data: &mut [Complex<f64>]) {
        let mut planner = FftPlanner::new();
        let fft = planner.plan_fft_forward(data.len());

        self.fft_scratch
            .resize(fft.get_inplace_scratch_len(), Complex::new(0., 0.));

        fft.process_with_scratch(data, &mut self.fft_scratch);

        let normalize = (data.len() as f64).recip().sqrt();

        for v in data.iter_mut() {
            *v *= normalize;
        }
    }
}

/// Iterative Soft Thresholding; an algorithm for reconstructing NUS data.
///
/// The [`Reconstructor::reconstruct`] implementation returns the fourier transform after interpolating unsampled data.
///
/// Interpolating is done by setting the unsampled positions in the data to zero, then iteratively thresholding the spectrum and setting the unsampled points to the IFT of what is above the threshold. On each iteration, the threshold is lowered until it reaches zero.
///
/// WARNING: IST is difficult to get right and it is unlikely that this does. These results should be trusted less than those of other, better tested, implementations. Feel free to contribute if you know any secret sauce.
///
/// A. Shchukina, P. Kasprzak, R. Dass et al. J Biomol NMR 68, 79–98 (2017). <https://doi.org/10.1007/s10858-016-0068-3>
/// (This implements IST-S with Virtual Echo)
pub struct Ist<'s> {
    schedule: &'s Schedule<Ix1>,
    threshold: Box<dyn Fn(f64) -> f64 + 's>,
    iterations: usize,
    fft_scratch: Vec<Complex<f64>>,
    ft_buffer: Vec<Complex<f64>>,
    mode: DisplayMode,
}

impl core::fmt::Debug for Ist<'_> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("Ist")
            .field("schedule", self.schedule)
            .field("iterations", &self.iterations)
            .field("mode", &self.mode)
            .finish()
    }
}

impl<'s> Ist<'s> {
    /// Create a new `Ist` instance
    ///
    /// `schedule` is the sampling schedule used.
    /// `threshold_function` is a function that returns the relative threshold for a given iteration. The function is called with values between zero and one, where each iteration i ∈ [0, iterations) is mapped to the range [0, 1]. The value returned must be between zero and one.
    /// `iterations` is the number of iterations to perform.
    /// `mode` is whether the data will be displayed in absolute value or real mode.
    pub fn new(
        schedule: &'s Schedule<Ix1>,
        threshold_function: impl Fn(f64) -> f64 + 's,
        iterations: usize,
        mode: DisplayMode,
    ) -> Ist<'s> {
        Ist {
            schedule,
            threshold: Box::new(threshold_function),
            iterations,
            fft_scratch: Vec::new(),
            ft_buffer: Vec::new(),
            mode,
        }
    }
}

impl Reconstructor for Ist<'_> {
    fn reconstruct(&mut self, data: &mut [Complex<f64>]) {
        assert_eq!(data.len(), self.schedule.len());

        if data.is_empty() {
            return;
        }

        self.ft_buffer
            .iter_mut()
            .for_each(|sample| *sample = Complex::new(0., 0.));

        let mode = self.mode;

        self.ft_buffer.resize(data.len() * 2, Complex::new(0., 0.));

        let ft_buffer = &mut *self.ft_buffer;

        let mut planner = FftPlanner::new();
        let fft = planner.plan_fft_forward(ft_buffer.len());
        let ift = planner.plan_fft_inverse(ft_buffer.len());

        self.fft_scratch.resize(
            fft.get_inplace_scratch_len()
                .max(ift.get_inplace_scratch_len()),
            Complex::new(0., 0.),
        );
        let fft_scratch = &mut self.fft_scratch;

        let normalize = (ft_buffer.len() as f64).recip().sqrt();

        for i in 0..self.iterations {
            // println!("{i}");

            ft_buffer
                .iter_mut()
                .zip(self.schedule.iter())
                .zip(data.iter())
                .for_each(|((sample, was_taken), data)| {
                    if *was_taken {
                        *sample = *data;
                    }
                });

            let (reconstruction, echo) = ft_buffer.split_at_mut(data.len());

            // Virtual echo
            reconstruction
                .iter()
                .zip(echo.iter_mut().rev())
                .for_each(|(sample, echo)| {
                    *echo = sample.conj();
                });

            fft.process_with_scratch(ft_buffer, fft_scratch);
            ft_buffer.apply(|v| v * normalize);

            let max = ft_buffer
                .iter()
                .map(|v| mode.magnitude(*v))
                .max_by(|a, b| a.total_cmp(b))
                .unwrap(); // The data is verified to not be empty

            let relative_threshold = (self.threshold)(i as f64 / (self.iterations as f64 - 1.));

            assert!(relative_threshold <= 1.);
            assert!(relative_threshold >= 0.);

            let threshold = relative_threshold * max;

            // println!("{}", ft_buffer.bar_chart(BarChartOptions::new()));

            mode.threshold(ft_buffer, threshold);

            // println!("{relative_threshold} → {threshold}");

            ift.process_with_scratch(ft_buffer, fft_scratch);
            ft_buffer.apply(|v| v * normalize);

            // println!(
            //     "{}",
            //     ft_buffer[0..data.len()].bar_chart(BarChartOptions::new())
            // );
        }

        let fft = planner.plan_fft_forward(data.len());
        let normalize = (data.len() as f64).recip().sqrt();

        ft_buffer
            .iter_mut()
            .zip(self.schedule.iter())
            .zip(data.iter())
            .for_each(|((sample, was_taken), data)| {
                if *was_taken {
                    *sample = *data;
                }
            });

        data.copy_from_slice(&ft_buffer[0..data.len()]);

        fft.process_with_scratch(data, fft_scratch);
        data.apply(|v| v * normalize);

        if let DisplayMode::RealPart = mode {
            data.apply(|v| Complex::new(v.re(), 0.))
        }
    }
}