nmr_schedule/reconstruction/
mod.rs

1//! Algorithms for reconstructing NMR signals.
2//!
3//! Supports FFT for uniform data, IST for nonuniform data, and quadrature.
4
5use alloc::{boxed::Box, vec::Vec};
6use ndarray::Ix1;
7use rustfft::{
8    num_complex::{Complex, ComplexFloat},
9    FftPlanner,
10};
11
12use crate::{ComplexSequence, DisplayMode, Schedule};
13
14/// Reconstructs signals from NMR data.
15pub trait Reconstructor {
16    /// 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.
17    ///
18    /// # Panics
19    ///
20    /// `out` must have the same length as `data`.
21    fn reconstruct_from_reals(&mut self, data: &[f64], out: &mut [Complex<f64>]) {
22        assert_eq!(data.len(), out.len());
23
24        out.iter_mut()
25            .zip(data)
26            .for_each(|(complex, real)| *complex = Complex::new(*real, 0.));
27
28        self.reconstruct(out);
29    }
30
31    /// Reconstruct a complex signal in place.
32    ///
33    /// This method should *not* be used for quadrature blindly. See [`Reconstructor::quadrature_reconstruct`].
34    fn reconstruct(&mut self, data: &mut [Complex<f64>]);
35
36    /// Reconstruct a spectrum from `cos` and `sin` signals using quadrature.
37    ///
38    /// This method adds `cos` and `-i·sin` componentwise, then calls `reconstruct` on the resulting list.
39    ///
40    /// # Panics
41    ///
42    /// `cos`, `sin`, and `out` must all have the same length.
43    fn quadrature_reconstruct(&mut self, cos: &[f64], sin: &[f64], out: &mut [Complex<f64>]) {
44        assert_eq!(cos.len(), sin.len());
45        assert_eq!(sin.len(), out.len());
46
47        out.iter_mut()
48            .zip(cos.iter())
49            .zip(sin.iter())
50            .for_each(|((complex, cos), sin)| {
51                *complex = Complex::new(*cos, -sin);
52            });
53
54        self.reconstruct(out);
55    }
56}
57
58/// FFT reconstruction for uniformly sampled data.
59pub struct Fft {
60    fft_scratch: Vec<Complex<f64>>,
61}
62
63impl core::fmt::Debug for Fft {
64    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
65        f.debug_struct("Fft").finish()
66    }
67}
68
69impl Fft {
70    /// Create a new `Fft` instance
71    pub const fn new() -> Fft {
72        Fft {
73            fft_scratch: Vec::new(),
74        }
75    }
76}
77
78impl Default for Fft {
79    fn default() -> Self {
80        Self::new()
81    }
82}
83
84impl Reconstructor for Fft {
85    fn reconstruct(&mut self, data: &mut [Complex<f64>]) {
86        let mut planner = FftPlanner::new();
87        let fft = planner.plan_fft_forward(data.len());
88
89        self.fft_scratch.resize(data.len(), Complex::new(0., 0.));
90
91        fft.process_with_scratch(data, &mut self.fft_scratch);
92
93        let normalize = (data.len() as f64).recip().sqrt();
94
95        for v in data.iter_mut() {
96            *v *= normalize;
97        }
98    }
99}
100
101/// Iterative Soft Thresholding; an algorithm for reconstructing NUS data.
102///
103/// The [`Reconstructor::reconstruct`] implementation returns the fourier transform after interpolating unsampled data.
104///
105/// 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.
106///
107/// 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.
108///
109/// A. Shchukina, P. Kasprzak, R. Dass et al. J Biomol NMR 68, 79–98 (2017). <https://doi.org/10.1007/s10858-016-0068-3>
110/// (This implements IST-S with Virtual Echo)
111pub struct Ist<'s> {
112    schedule: &'s Schedule<Ix1>,
113    threshold: Box<dyn Fn(f64) -> f64 + 's>,
114    iterations: usize,
115    fft_scratch: Vec<Complex<f64>>,
116    ft_buffer: Vec<Complex<f64>>,
117    mode: DisplayMode,
118}
119
120impl<'s> core::fmt::Debug for Ist<'s> {
121    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
122        f.debug_struct("Ist")
123            .field("schedule", self.schedule)
124            .field("iterations", &self.iterations)
125            .field("mode", &self.mode)
126            .finish()
127    }
128}
129
130impl<'s> Ist<'s> {
131    /// Create a new `Ist` instance
132    ///
133    /// `schedule` is the sampling schedule used.
134    /// `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.
135    /// `iterations` is the number of iterations to perform.
136    /// `mode` is whether the data will be displayed in absolute value or real mode.
137    pub fn new(
138        schedule: &'s Schedule<Ix1>,
139        threshold_function: impl Fn(f64) -> f64 + 's,
140        iterations: usize,
141        mode: DisplayMode,
142    ) -> Ist<'s> {
143        Ist {
144            schedule,
145            threshold: Box::new(threshold_function),
146            iterations,
147            fft_scratch: Vec::new(),
148            ft_buffer: Vec::new(),
149            mode,
150        }
151    }
152}
153
154impl<'s> Reconstructor for Ist<'s> {
155    fn reconstruct(&mut self, data: &mut [Complex<f64>]) {
156        assert_eq!(data.len(), self.schedule.len());
157
158        if data.is_empty() {
159            return;
160        }
161
162        self.ft_buffer
163            .iter_mut()
164            .for_each(|sample| *sample = Complex::new(0., 0.));
165
166        let mode = self.mode;
167
168        self.ft_buffer.resize(data.len() * 2, Complex::new(0., 0.));
169        self.fft_scratch
170            .resize(data.len() * 2, Complex::new(0., 0.));
171
172        let ft_buffer = &mut *self.ft_buffer;
173        let fft_scratch = &mut self.fft_scratch;
174
175        let mut planner = FftPlanner::new();
176        let fft = planner.plan_fft_forward(ft_buffer.len());
177        let ift = planner.plan_fft_inverse(ft_buffer.len());
178
179        let normalize = (ft_buffer.len() as f64).recip().sqrt();
180
181        for i in 0..self.iterations {
182            // println!("{i}");
183
184            ft_buffer
185                .iter_mut()
186                .zip(self.schedule.iter())
187                .zip(data.iter())
188                .for_each(|((sample, was_taken), data)| {
189                    if *was_taken {
190                        *sample = *data;
191                    }
192                });
193
194            let (reconstruction, echo) = ft_buffer.split_at_mut(data.len());
195
196            // Virtual echo
197            reconstruction
198                .iter()
199                .zip(echo.iter_mut().rev())
200                .for_each(|(sample, echo)| {
201                    *echo = sample.conj();
202                });
203
204            fft.process_with_scratch(ft_buffer, fft_scratch);
205            ft_buffer.apply(|v| v * normalize);
206
207            let max = ft_buffer
208                .iter()
209                .map(|v| mode.magnitude(*v))
210                .max_by(|a, b| a.total_cmp(b))
211                .unwrap(); // The data is verified to not be empty
212
213            let relative_threshold = (self.threshold)(i as f64 / (self.iterations as f64 - 1.));
214
215            assert!(relative_threshold <= 1.);
216            assert!(relative_threshold >= 0.);
217
218            let threshold = relative_threshold * max;
219
220            // println!("{}", ft_buffer.bar_chart(BarChartOptions::new()));
221
222            mode.threshold(ft_buffer, threshold);
223
224            // println!("{relative_threshold} → {threshold}");
225
226            ift.process_with_scratch(ft_buffer, fft_scratch);
227            ft_buffer.apply(|v| v * normalize);
228
229            // println!(
230            //     "{}",
231            //     ft_buffer[0..data.len()].bar_chart(BarChartOptions::new())
232            // );
233        }
234
235        let fft = planner.plan_fft_forward(data.len());
236        let normalize = (data.len() as f64).recip().sqrt();
237
238        ft_buffer
239            .iter_mut()
240            .zip(self.schedule.iter())
241            .zip(data.iter())
242            .for_each(|((sample, was_taken), data)| {
243                if *was_taken {
244                    *sample = *data;
245                }
246            });
247
248        data.copy_from_slice(&ft_buffer[0..data.len()]);
249
250        fft.process_with_scratch(data, fft_scratch);
251        data.apply(|v| v * normalize);
252
253        if let DisplayMode::RealPart = mode {
254            data.apply(|v| Complex::new(v.re(), 0.))
255        }
256    }
257}
nmr_schedule/reconstruction/mod.rs

nmr_schedule/reconstruction/
mod.rs
