nmr-schedule 0.2.1

//! This crate implements algorithms for Non-Uniform Sampling in NMR spectroscopy.
//!
//! This crate primarily implements algorithms for generating and filtering sampling schedules. It includes various base schedules along with various post-processing filters for schedules.
//!
//! Schedulers implemented:
//! - [Quantiles](`generators::Quantiles`)
//! - [Poisson Gap](`generators::SinWeightedPoissonGap`)
//! - [Random](`generators::RandomSampling`)
//! - [Averaging](`generators::Averaging`)
//!
//! Filters implemented:
//! - [Backfill](`modifiers::FillCorners`)
//! - [Seed searching](`modifiers::Iterate`)
//! - [Thue-Morse filter](`modifiers::TMFilter`)
//! - [PSF polisher](`modifiers::PSFPolisher`)
//!
//! This library currently only supports 1D schedules, however there are plans to support higher dimensional schedules in the future and the architecture is already generic over dimension.
//!
//! # Examples
//!
//! ```
//! # use nmr_schedule::{*, pdf::*, generators::*, modifiers::*, ndarray::Ix1};
//! // Generate a 64x256 schedule with Quantiles with QSin weighting, 8 points backfill, and TMPF filtering
//! let sched = Quantiles::new(|len| qsin(len, QSinBias::Low, 3.))
//!     .fill_corners(|_, _| [8, 1]) // Any function of the count and length of the schedule
//!     .tm_filter()
//!     .polish_psf(0.1, 0.32, DisplayMode::Abs)
//!     .generate(64, Ix1(256));
//!
//! println!("{sched}");
//! ```
//!
//! ```
//! # use nmr_schedule::{*, pdf::*, generators::*, modifiers::*, ndarray::Ix1};
//! // Apply TMPF filtering to an existing schedule
//! let sched_encoded = "0\n1\n2\n5\n7\n9\n20";
//! let sched =
//!     Schedule::decode(sched_encoded, EncodingType::ZeroBased, |dim| Ok(dim)).unwrap();
//!
//! let filtered =
//!     PSFPolisher::new(0.1, 0.32, DisplayMode::Abs).filter(TMFilter::new().filter(sched));
//!
//! let encoded = filtered.encode(EncodingType::ZeroBased);
//! println!("{encoded}");
//! ```

#![cfg_attr(not(test), no_std)]
#![cfg_attr(docsrs, feature(doc_cfg))]
#![warn(clippy::cargo)]
#![warn(clippy::complexity)]
#![warn(clippy::correctness)]
#![warn(clippy::perf)]
#![warn(clippy::style)]
#![warn(clippy::suspicious)]
#![warn(missing_docs)]
#![warn(missing_copy_implementations)]
#![warn(missing_debug_implementations)]
#![warn(clippy::missing_panics_doc)]
#![warn(clippy::missing_const_for_fn)]
// Remove allow after https://github.com/rust-lang/rust-clippy/pull/14609 goes into stable
#![allow(clippy::unused_unit)]

extern crate alloc;

use core::cmp::Ordering;

pub mod generators;
pub mod modifiers;
pub mod pdf;
pub mod reconstruction;
mod schedule;
#[cfg(feature = "terminal-viz")]
#[cfg_attr(docsrs, doc(cfg(feature = "terminal-viz")))]
pub mod terminal_viz;

use rand::Rng;
use rand_chacha::ChaCha12Rng;
use rustfft::num_complex::{Complex, ComplexFloat};
pub use schedule::*;

pub use rustfft;

pub use ndarray;

/// Represents whether a spectrum is meant to be displayed as the real part or as the absolute value
///
/// Note that this is very experimental and you should default to using [`DisplayMode::Abs`] in parameters. The PSF Polisher publication only uses [`DisplayMode::Abs`].
///
/// Future research can look into whether knowing the display mode of an experiment can inform schedule generation and reconstruction.
///
/// The PSF polisher, when passed in [`DisplayMode::RealPart`], will only "see" the real part of the PSF during filtering and ignore the imaginary part, potentially leaving artifacts in and moving PSF noise into the imaginary part. If it is known that all of the signal will be exclusively in the real part, then moving sampling noise from the real part to the imaginary part may be desirable.
///
/// The IST reconstructor will also only "see" the real part when passed in [`DisplayMode::RealPart`]. It will assume that all signals in the imaginary axis are noise and ignore them.
#[derive(Clone, Copy, Debug, Default)]
pub enum DisplayMode {
    /// The spectrum will be displayed in absolute value mode
    #[default]
    Abs,
    /// The spectrum will be displayed in real part mode; the complex part will be ignored.
    RealPart,
}

impl DisplayMode {
    /// Calculate the magnitude of the complex number given the display mode. In `Abs` mode, this will return `complex.abs()`, and in `RealPart` mode, this will return `complex.re().abs()`.
    pub fn magnitude(self, complex: Complex<f64>) -> f64 {
        match self {
            DisplayMode::Abs => complex.abs(),
            DisplayMode::RealPart => complex.re().abs(),
        }
    }

    /// Apply a soft threshold to a sequence, zeroing out all components with [`DisplayMode::magnitude`] less than the threshold.
    pub fn threshold<T: ComplexSequence + ?Sized>(self, sequence: &mut T, threshold: f64) {
        sequence.apply(|v| {
            let mag = self.magnitude(v);
            if mag > threshold {
                v * (mag - threshold) / mag
            } else {
                Complex::new(0., 0.)
            }
        })
    }

    /// Take the real part of the sequence if `self` is `RealPart`, otherwise leaving it alone.
    pub fn maybe_real_part<T: ComplexSequence + ?Sized>(self, sequence: &mut T) {
        if let Self::RealPart = self {
            sequence.apply(|v| Complex::new(v.re(), 0.));
        }
    }
}

fn partition<T>(rng: &mut ChaCha12Rng, slice: &mut [T], by: &impl Fn(&T, &T) -> Ordering) -> usize {
    slice.swap(0, rng.random_range(0..slice.len()));

    let mut i = 1;
    let mut j = slice.len() - 1;

    loop {
        while i < slice.len() && !matches!(by(&slice[i], &slice[0]), Ordering::Less) {
            i += 1;
        }

        while matches!(by(&slice[j], &slice[0]), Ordering::Less) {
            j -= 1;
        }

        // If the indices crossed, return
        if i > j {
            slice.swap(0, j);
            return j;
        }

        // Swap the elements at the left and right indices
        slice.swap(i, j);
        i += 1;
    }
}

/// Standard quickselect algorithm: https://en.wikipedia.org/wiki/Quickselect
/// Modified to sort in descending order (since I want maxima in all of my usecases)
///
/// After calling this function, the value at index `find_spot` is guaranteed to be at the correctly sorted position and all values at indices less than `find_spot` are guaranteed to be greater than the value at `find_spot` and vice versa for indices greater.
pub(crate) fn quickselect<T>(
    rng: &mut ChaCha12Rng,
    mut slice: &mut [T],
    by: impl Fn(&T, &T) -> Ordering,
    mut find_spot: usize,
) {
    loop {
        let len = slice.len();

        if len < 2 {
            return;
        }

        let spot_found = partition(rng, slice, &by);

        match find_spot.cmp(&spot_found) {
            Ordering::Less => slice = &mut slice[0..spot_found],
            Ordering::Equal => return,
            Ordering::Greater => {
                slice = &mut slice[spot_found + 1..len];
                find_spot = find_spot - spot_found - 1;
            }
        }
    }
}

/// An extension trait adding utility functions to lists of complex numbers.
pub trait ComplexSequence {
    /// Apply a function componentwise on each complex number.
    fn apply(&mut self, f: impl FnMut(Complex<f64>) -> Complex<f64>);

    /// Apply a function componentwise on each complex number and copy the output to the `out` array.
    ///
    /// # Panics
    ///
    /// The buffer of complex numbers must be the same length `out`.
    fn apply_into<T>(&self, out: &mut [T], f: impl FnMut(Complex<f64>) -> T);

    /// Multiply each value in the sequence by e^iθ, rotating the phase by θ radians.
    fn phase(&mut self, θ: f64) {
        let rotator = Complex::new(0., θ).exp();
        self.apply(|v| v * rotator);
    }

    /// Copy the real part of each value into `out`.
    ///
    /// # Panics
    ///
    /// The buffer of complex numbers must be the same length as `out`.
    fn re(&self, out: &mut [f64]) {
        self.apply_into(out, |v| v.re());
    }
}

impl ComplexSequence for [Complex<f64>] {
    fn apply(&mut self, mut f: impl FnMut(Complex<f64>) -> Complex<f64>) {
        for v in self {
            *v = f(*v);
        }
    }

    fn apply_into<T>(&self, out: &mut [T], mut f: impl FnMut(Complex<f64>) -> T) {
        assert_eq!(self.len(), out.len());

        out.iter_mut()
            .zip(self.iter())
            .for_each(|(out_v, complex)| *out_v = f(*complex));
    }
}

impl<V: AsMut<[Complex<f64>]> + AsRef<[Complex<f64>]>> ComplexSequence for V {
    fn apply(&mut self, f: impl FnMut(Complex<f64>) -> Complex<f64>) {
        self.as_mut().apply(f);
    }

    fn apply_into<T>(&self, out: &mut [T], f: impl FnMut(Complex<f64>) -> T) {
        self.as_ref().apply_into(out, f);
    }
}

/// Represents whether a function is monotonically increasing or decreasing.
#[derive(Clone, Copy, Debug)]
enum Monotonicity {
    /// The function is increasing
    Increasing,
    /// The function is decreasing
    Decreasing,
}

/// Represents the initial parameters of a binary search
#[derive(Clone, Copy, Debug)]
struct InitialState {
    start: f64,
    min: f64,
    max: f64,
}

impl InitialState {
    /// Create initial parameters for binary search
    ///
    /// `start` is the initial guess.
    pub const fn new(start: f64, min: f64, max: f64) -> Self {
        Self { start, min, max }
    }
}

type BinsearchPoint<T> = (f64, (f64, T));

/// Represents the precision required from a binary search
enum Precision<'a, T> {
    #[allow(dead_code)]
    // In case I make changes that need to use this, I don't want to get rid of it just yet
    Preimage(f64),
    Image {
        amount: f64,
        /// A function to be called if the required precision can't be acheived
        debug: &'a dyn Fn(BinsearchPoint<T>, BinsearchPoint<T>),
    },
}

/// Perform binary search over `f` to find where it is zero.
///
/// ```
/// ```
///
/// `monotonicity` tells whether `f` is monotonically increasing or decreasing.
/// `initial_state` defines the initial state for binary search.
/// `precision` defines the target precision. The absolute value of `f`'s return value is guaranteed to be less than `precision` except when convergence fails.
/// `f` is the function to perform binary search over. It is allowed to return an arbitrary value along with the value to binary search over.
/// `debug` will be called if the binary search doesn't converge and runs out of precision to differentiate the minimum and maximum values. It's called with tuples of (input, output) for the minimum and maximum states. In this situation, the function will return the value of the last guess regardless of whether it's correct.
fn find_zero<T>(
    monotonicity: Monotonicity,
    initial_state: InitialState,
    precision: Precision<'_, T>,
    f: impl Fn(f64) -> (f64, T),
) -> (f64, T) {
    let mut min = initial_state.min;
    let mut current = initial_state.start;
    let mut max = initial_state.max;

    loop {
        let (v, ret) = f(current);

        if let Precision::Image { amount, debug: _ } = precision {
            if v.abs() < amount {
                return (current, ret);
            }
        }

        match (v.total_cmp(&0.), monotonicity) {
            (Ordering::Less, Monotonicity::Increasing)
            | (Ordering::Greater, Monotonicity::Decreasing) => min = current,
            (Ordering::Equal, _) => return (current, ret),
            (Ordering::Greater, Monotonicity::Increasing)
            | (Ordering::Less, Monotonicity::Decreasing) => max = current,
        }

        if max != f64::INFINITY {
            let new_current = min / 2. + max / 2.;

            match precision {
                Precision::Preimage(amount) => {
                    if max - min < amount || new_current == current {
                        return (current, ret);
                    }
                }
                Precision::Image { amount: _, debug } => {
                    if new_current == current {
                        debug((min, f(min)), (max, f(max)));
                        return (current, ret);
                    }
                }
            }

            current = new_current;
        } else {
            current *= 2.;
        }
    }
}

#[cfg(test)]
mod tests {
    use alloc::borrow::ToOwned;
    use core::cell::OnceCell;
    use rand_chacha::ChaCha12Rng;

    use alloc::vec::Vec;
    use rand::{Rng, SeedableRng};
    use rustfft::num_complex::{Complex, ComplexFloat};

    use crate::{DisplayMode, InitialState, Monotonicity, find_zero, quickselect};

    #[test]
    fn test_display_mode() {
        assert_eq!(DisplayMode::Abs.magnitude(Complex::new(4., 3.)), 5.);
        assert_eq!(DisplayMode::RealPart.magnitude(Complex::new(4., 3.)), 4.);

        let mut seq = [Complex::new(6., 8.), Complex::new(1., 1.)];

        DisplayMode::Abs.threshold(&mut seq[..], 5.);
        DisplayMode::Abs.maybe_real_part(&mut seq[..]);
        assert!((seq[0] - Complex::new(3., 4.)).abs() < 0.000000001);
        assert_eq!(seq[1], Complex::ZERO);

        seq = [Complex::new(2., 9.), Complex::new(0.5, 1000.)];

        DisplayMode::RealPart.threshold(&mut seq[..], 1.);

        assert!((seq[0] - Complex::new(1., 4.5)).abs() < 0.000000001);
        assert_eq!(seq[1], Complex::ZERO);

        DisplayMode::RealPart.maybe_real_part(&mut seq[..]);
        assert!((seq[0] - Complex::new(1., 0.)).abs() < 0.000000001);

        let mut rng = ChaCha12Rng::from_seed(*b"Not all plants spread seeds, and");

        for mode in [DisplayMode::Abs, DisplayMode::RealPart] {
            let seq = (0..1000)
                .map(|_| Complex::new(rng.random::<f64>() - 0.5, rng.random::<f64>() - 0.5) * 128.)
                .collect::<Vec<_>>();

            let mut seq_new = seq.to_owned();

            mode.threshold(&mut *seq_new, 32.);

            seq.iter().zip(seq_new.iter()).for_each(|(before, after)| {
                if mode.magnitude(*before) < 32. {
                    assert_eq!(*after, Complex::ZERO);
                } else {
                    assert!(
                        (mode.magnitude(*before) - mode.magnitude(*after) - 32.).abs() < 0.00000001,
                        "{before} {after} {mode:?}"
                    );

                    assert!(
                        (before.arg() - after.arg()).abs() < 0.00000001,
                        "{before} {after} {mode:?}"
                    );

                    match mode {
                        DisplayMode::Abs => {}
                        DisplayMode::RealPart => {
                            assert_eq!(
                                before.re().signum(),
                                after.re().signum(),
                                "{before} {after} {mode:?}"
                            );
                        }
                    }
                }
            });

            let mut seq_maybe_re = seq_new.to_owned();

            mode.maybe_real_part(&mut *seq_maybe_re);

            seq_new
                .iter()
                .zip(seq_maybe_re.iter())
                .for_each(|(before, after)| match mode {
                    DisplayMode::Abs => assert_eq!(before, after),
                    DisplayMode::RealPart => {
                        assert_eq!(after.im(), 0.);
                    }
                });
        }
    }

    #[test]
    fn test_binsearch() {
        let v = find_zero(
            Monotonicity::Increasing,
            InitialState {
                start: 1.,
                min: 0.,
                max: f64::INFINITY,
            },
            crate::Precision::Preimage(0.1),
            |v| (-5. + v, ()),
        )
        .0;

        assert!(v > 5. - 0.1 && v < 5. + 0.1);

        let v = find_zero(
            Monotonicity::Decreasing,
            InitialState {
                start: 1.,
                min: 0.,
                max: f64::INFINITY,
            },
            crate::Precision::Preimage(0.1),
            |v| (4. - 0.1 * v, ()),
        )
        .0;

        assert!(v > 40. - 0.1 && v < 40. + 0.1, "v = {v}");

        let v = find_zero(
            Monotonicity::Increasing,
            InitialState {
                start: 1.,
                min: 0.,
                max: f64::INFINITY,
            },
            crate::Precision::Image {
                amount: 0.1,
                debug: (&|_, _| -> () { panic!("Debug should not have been called!") })
                    as &dyn Fn(_, _),
            },
            |v| (-3. + 10. * v, ()),
        )
        .0;

        assert!(v > 0.3 - 0.01 && v < 0.3 + 0.01, "v = {v}");

        let v = find_zero(
            Monotonicity::Decreasing,
            InitialState {
                start: 1.,
                min: 0.,
                max: f64::INFINITY,
            },
            crate::Precision::Image {
                amount: 0.1,
                debug: (&|_, _| -> () { panic!("Debug should not have been called!") })
                    as &dyn Fn(_, _),
            },
            |v| (5. - 5. * v, ()),
        )
        .0;

        assert!(v > 1. - 0.02 && v < 1. + 0.02, "v = {v}");

        let debug_called = OnceCell::new();

        find_zero(
            Monotonicity::Decreasing,
            InitialState {
                start: 1.,
                min: 0.,
                max: f64::INFINITY,
            },
            crate::Precision::Image {
                amount: 0.1,
                debug: (&|_, _| {
                    debug_called.set(true).unwrap();
                }) as &dyn Fn(_, _),
            },
            |v| (if v > 5. { 1. } else { -1. }, ()),
        );

        assert!(debug_called.get().is_some());
    }

    #[test]
    fn test_quickselect() {
        fn verify(rng: &mut ChaCha12Rng, pos: usize, slice: &[f64]) {
            let mut slice = slice
                .iter()
                .enumerate()
                .map(|(a, b)| (*b, a))
                .collect::<Vec<_>>();

            quickselect(rng, &mut slice, |a, b| a.0.total_cmp(&b.0), pos);

            for i in 0..pos {
                assert!(
                    slice[i].0 >= slice[pos].0,
                    "Pos: {pos}, Index: {i} - {slice:?}"
                );
            }

            for i in pos + 1..slice.len() {
                assert!(
                    slice[i].0 <= slice[pos].0,
                    "Pos: {pos}, Index: {i} - {slice:?}"
                );
            }

            let v = slice[pos];

            slice.sort_by(|a, b| b.0.total_cmp(&a.0));

            assert_eq!(slice[pos].0, v.0);
        }

        let mut rng = ChaCha12Rng::from_seed(*b"Not all seeds plant plants, some");

        verify(&mut rng, 2, &[5., 4., 3., 2., 1.]);
        verify(&mut rng, 2, &[1., 2., 3., 4., 5.]);
        verify(&mut rng, 3, &[1., 2., 1., 4., 3.]);

        for i in 0..100 {
            let pos = rng.random_range(0..i + 1);
            let data = (0..i + 1).map(|_| rng.random()).collect::<Vec<_>>();
            verify(&mut rng, pos, &data);
        }
    }
}