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    FftPlanner,
9    num_complex::{Complex, ComplexFloat},
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
90            .resize(fft.get_inplace_scratch_len(), Complex::new(0., 0.));
91
92        fft.process_with_scratch(data, &mut self.fft_scratch);
93
94        let normalize = (data.len() as f64).recip().sqrt();
95
96        for v in data.iter_mut() {
97            *v *= normalize;
98        }
99    }
100}
101
102/// Iterative Soft Thresholding; an algorithm for reconstructing NUS data.
103///
104/// The [`Reconstructor::reconstruct`] implementation returns the fourier transform after interpolating unsampled data.
105///
106/// 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.
107///
108/// 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.
109///
110/// A. Shchukina, P. Kasprzak, R. Dass et al. J Biomol NMR 68, 79–98 (2017). <https://doi.org/10.1007/s10858-016-0068-3>
111/// (This implements IST-S with Virtual Echo)
112pub struct Ist<'s> {
113    schedule: &'s Schedule<Ix1>,
114    threshold: Box<dyn Fn(f64) -> f64 + 's>,
115    iterations: usize,
116    fft_scratch: Vec<Complex<f64>>,
117    ft_buffer: Vec<Complex<f64>>,
118    mode: DisplayMode,
119}
120
121impl core::fmt::Debug for Ist<'_> {
122    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
123        f.debug_struct("Ist")
124            .field("schedule", self.schedule)
125            .field("iterations", &self.iterations)
126            .field("mode", &self.mode)
127            .finish()
128    }
129}
130
131impl<'s> Ist<'s> {
132    /// Create a new `Ist` instance
133    ///
134    /// `schedule` is the sampling schedule used.
135    /// `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.
136    /// `iterations` is the number of iterations to perform.
137    /// `mode` is whether the data will be displayed in absolute value or real mode.
138    pub fn new(
139        schedule: &'s Schedule<Ix1>,
140        threshold_function: impl Fn(f64) -> f64 + 's,
141        iterations: usize,
142        mode: DisplayMode,
143    ) -> Ist<'s> {
144        Ist {
145            schedule,
146            threshold: Box::new(threshold_function),
147            iterations,
148            fft_scratch: Vec::new(),
149            ft_buffer: Vec::new(),
150            mode,
151        }
152    }
153}
154
155impl Reconstructor for Ist<'_> {
156    fn reconstruct(&mut self, data: &mut [Complex<f64>]) {
157        assert_eq!(data.len(), self.schedule.len());
158
159        if data.is_empty() {
160            return;
161        }
162
163        self.ft_buffer
164            .iter_mut()
165            .for_each(|sample| *sample = Complex::new(0., 0.));
166
167        let mode = self.mode;
168
169        self.ft_buffer.resize(data.len() * 2, Complex::new(0., 0.));
170
171        let ft_buffer = &mut *self.ft_buffer;
172
173        let mut planner = FftPlanner::new();
174        let fft = planner.plan_fft_forward(ft_buffer.len());
175        let ift = planner.plan_fft_inverse(ft_buffer.len());
176
177        self.fft_scratch.resize(
178            fft.get_inplace_scratch_len()
179                .max(ift.get_inplace_scratch_len()),
180            Complex::new(0., 0.),
181        );
182        let fft_scratch = &mut self.fft_scratch;
183
184        let normalize = (ft_buffer.len() as f64).recip().sqrt();
185
186        for i in 0..self.iterations {
187            // println!("{i}");
188
189            ft_buffer
190                .iter_mut()
191                .zip(self.schedule.iter())
192                .zip(data.iter())
193                .for_each(|((sample, was_taken), data)| {
194                    if *was_taken {
195                        *sample = *data;
196                    }
197                });
198
199            let (reconstruction, echo) = ft_buffer.split_at_mut(data.len());
200
201            // Virtual echo
202            reconstruction
203                .iter()
204                .zip(echo.iter_mut().rev())
205                .for_each(|(sample, echo)| {
206                    *echo = sample.conj();
207                });
208
209            fft.process_with_scratch(ft_buffer, fft_scratch);
210            ft_buffer.apply(|v| v * normalize);
211
212            let max = ft_buffer
213                .iter()
214                .map(|v| mode.magnitude(*v))
215                .max_by(|a, b| a.total_cmp(b))
216                .unwrap(); // The data is verified to not be empty
217
218            let relative_threshold = (self.threshold)(i as f64 / (self.iterations as f64 - 1.));
219
220            assert!(relative_threshold <= 1.);
221            assert!(relative_threshold >= 0.);
222
223            let threshold = relative_threshold * max;
224
225            // println!("{}", ft_buffer.bar_chart(BarChartOptions::new()));
226
227            mode.threshold(ft_buffer, threshold);
228
229            // println!("{relative_threshold} → {threshold}");
230
231            ift.process_with_scratch(ft_buffer, fft_scratch);
232            ft_buffer.apply(|v| v * normalize);
233
234            // println!(
235            //     "{}",
236            //     ft_buffer[0..data.len()].bar_chart(BarChartOptions::new())
237            // );
238        }
239
240        let fft = planner.plan_fft_forward(data.len());
241        let normalize = (data.len() as f64).recip().sqrt();
242
243        ft_buffer
244            .iter_mut()
245            .zip(self.schedule.iter())
246            .zip(data.iter())
247            .for_each(|((sample, was_taken), data)| {
248                if *was_taken {
249                    *sample = *data;
250                }
251            });
252
253        data.copy_from_slice(&ft_buffer[0..data.len()]);
254
255        fft.process_with_scratch(data, fft_scratch);
256        data.apply(|v| v * normalize);
257
258        if let DisplayMode::RealPart = mode {
259            data.apply(|v| Complex::new(v.re(), 0.))
260        }
261    }
262}
nmr_schedule/reconstruction/mod.rs

nmr_schedule/reconstruction/
mod.rs
