qsm_core/inversion/
rts.rs

1//! Rapid Two-Step (RTS) dipole inversion
2//!
3//! Two-step approach that combines:
4//! 1. LSMR for well-conditioned k-space regions
5//! 2. TV regularization for ill-conditioned regions
6//!
7//! Reference:
8//! Kames, C., Wiggermann, V., Rauscher, A. (2018).
9//! "Rapid two-step dipole inversion for susceptibility mapping with sparsity priors."
10//! NeuroImage, 167:276-283. https://doi.org/10.1016/j.neuroimage.2017.11.018
11//!
12//! Reference implementation: https://github.com/kamesy/QSM.jl
13
14use num_complex::Complex64;
15use crate::fft::Fft3dWorkspace;
16use crate::kernels::dipole::dipole_kernel;
17use crate::kernels::laplacian::laplacian_kernel;
18use crate::utils::{shrink, apply_mask_zero};
19use crate::Grid;
20use super::admm::{AdmmBuffers, admm_step};
21
22/// RTS algorithm parameters
23#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
24#[derive(Clone, Debug)]
25pub struct RtsParams {
26    /// Threshold for ill-conditioned region (typically 0.15)
27    pub delta: f64,
28    /// Regularization parameter for well-conditioned region (typically 1e5)
29    pub mu: f64,
30    /// ADMM penalty parameter (typically 10)
31    pub rho: f64,
32    /// Convergence tolerance
33    pub tol: f64,
34    /// Maximum ADMM iterations
35    pub max_iter: usize,
36    /// LSMR iterations for step 1 (typically 4)
37    pub lsmr_iter: usize,
38}
39
40impl Default for RtsParams {
41    fn default() -> Self {
42        Self {
43            delta: 0.15,
44            mu: 1e5,
45            rho: 10.0,
46            tol: 1e-2,
47            max_iter: 20,
48            lsmr_iter: 4,
49        }
50    }
51}
52
53/// RTS dipole inversion
54///
55/// Optimized implementation with:
56/// - Pre-allocated buffers (zero allocations per iteration)
57/// - In-place gradient/divergence operations
58/// - Buffer swapping instead of cloning
59/// - Fused z-subproblem and u-update
60///
61/// # Arguments
62/// * `local_field` - Local field values (nx * ny * nz)
63/// * `mask` - Binary mask (nx * ny * nz), 1 = inside ROI
64/// * `grid` - Volume grid (dimensions and voxel sizes)
65/// * `bdir` - B0 field direction
66/// * `params` - RTS parameters
67/// * `progress` - Progress callback `(iteration, max_iter)`
68///
69/// # Returns
70/// Susceptibility map
71pub fn rts(
72    local_field: &[f64],
73    mask: &[u8],
74    grid: &Grid,
75    bdir: (f64, f64, f64),
76    params: &RtsParams,
77    mut progress: impl FnMut(usize, usize),
78) -> Vec<f64> {
79    let (nx, ny, nz) = grid.dims;
80    let n_total = grid.n_total();
81
82    // ========================================================================
83    // Pre-compute kernels (done once)
84    // ========================================================================
85
86    // Create FFT workspace (caches plans and scratch buffers for reuse)
87    let mut fft_ws = Fft3dWorkspace::new(nx, ny, nz);
88
89    // Generate dipole kernel D
90    let d_kernel = dipole_kernel(grid, bdir);
91
92    // Generate negative Laplacian kernel
93    let l_kernel = laplacian_kernel(grid, true);
94
95    // FFT of Laplacian kernel (reuse buffer for other purposes later)
96    let mut work_complex: Vec<Complex64> = l_kernel.iter()
97        .map(|&x| Complex64::new(x, 0.0))
98        .collect();
99    fft_ws.fft3d(&mut work_complex);
100
101    // Compute well-conditioned mask M and inverse operator iA
102    let mut m_mask: Vec<f64> = vec![0.0; n_total];
103    let mut inv_a: Vec<f64> = vec![0.0; n_total];
104
105    for i in 0..n_total {
106        let l_fft_i = work_complex[i].re;
107        if d_kernel[i].abs() > params.delta {
108            m_mask[i] = params.mu;
109        }
110        let a = m_mask[i] + params.rho * l_fft_i;
111        if a.abs() > 1e-20 {
112            inv_a[i] = params.rho / a;
113        }
114    }
115
116    // ========================================================================
117    // Step 1: Well-conditioned k-space (simplified LSMR)
118    // ========================================================================
119
120    // FFT of field (reuse work_complex)
121    for i in 0..n_total {
122        work_complex[i] = Complex64::new(local_field[i], 0.0);
123    }
124    fft_ws.fft3d(&mut work_complex);
125
126    // Store field_fft for LSMR iterations
127    let field_fft: Vec<Complex64> = work_complex.clone();
128
129    // Initial estimate: chi = D * f / (D^2 + epsilon) for well-conditioned
130    // Stored in work_complex
131    for i in 0..n_total {
132        let d = d_kernel[i];
133        if d.abs() > params.delta {
134            work_complex[i] = field_fft[i] * d / (d * d + 1e-6);
135        } else {
136            work_complex[i] = Complex64::new(0.0, 0.0);
137        }
138    }
139
140    // Simple iterative refinement for well-conditioned region
141    // Use a temporary buffer for residual
142    let mut residual = vec![Complex64::new(0.0, 0.0); n_total];
143    for _ in 0..params.lsmr_iter {
144        // residual = f - D * chi
145        for i in 0..n_total {
146            residual[i] = field_fft[i] - work_complex[i] * d_kernel[i];
147        }
148
149        // update chi for well-conditioned region
150        for i in 0..n_total {
151            let d = d_kernel[i];
152            if d.abs() > params.delta {
153                work_complex[i] += residual[i] * d / (d * d + 1e-6);
154            }
155        }
156    }
157
158    // Transform to spatial domain
159    fft_ws.ifft3d(&mut work_complex);
160
161    // Initialize x and apply mask
162    let mut x = vec![0.0; n_total];
163    for i in 0..n_total {
164        x[i] = if mask[i] != 0 { work_complex[i].re } else { 0.0 };
165    }
166
167    // ========================================================================
168    // Pre-compute constant part of RHS for ADMM
169    // ========================================================================
170
171    // F_hat = inv_a * M * FFT(x) / rho
172    for i in 0..n_total {
173        work_complex[i] = Complex64::new(x[i], 0.0);
174    }
175    fft_ws.fft3d(&mut work_complex);
176
177    let mut f_hat: Vec<Complex64> = vec![Complex64::new(0.0, 0.0); n_total];
178    for i in 0..n_total {
179        if m_mask[i].abs() > 1e-20 && inv_a[i].abs() > 1e-20 {
180            f_hat[i] = work_complex[i] * (m_mask[i] / params.rho) * inv_a[i];
181        }
182    }
183
184    // ========================================================================
185    // Pre-allocate working buffers and run ADMM iterations
186    // ========================================================================
187
188    let mut buf = AdmmBuffers::new(n_total);
189    // Copy LSMR result into buf.x
190    buf.x.copy_from_slice(&x);
191
192    let inv_rho = 1.0 / params.rho;
193
194    for iter in 0..params.max_iter {
195        progress(iter + 1, params.max_iter);
196
197        // RTS uses rho=1.0 in admm_step because inv_a already incorporates rho
198        let converged = admm_step(
199            &mut buf, &mut fft_ws, &f_hat, &inv_a, 1.0, grid, params.tol,
200            |vx, vy, vz, _| (shrink(vx, inv_rho), shrink(vy, inv_rho), shrink(vz, inv_rho)),
201        );
202
203        if converged {
204            progress(iter + 1, iter + 1);
205            break;
206        }
207    }
208
209    apply_mask_zero(&mut buf.x, mask);
210
211    buf.x
212}
213
214#[cfg(test)]
215mod tests {
216    use super::*;
217
218    #[test]
219    fn test_rts_zero_field() {
220        let n = 8;
221        let field = vec![0.0; n * n * n];
222        let mask = vec![1u8; n * n * n];
223        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
224        let params = RtsParams { delta: 0.15, mu: 1e5, rho: 10.0, tol: 1e-2, max_iter: 5, lsmr_iter: 2 };
225
226        let chi = rts(&field, &mask, &grid, (0.0, 0.0, 1.0), &params, |_, _| {});
227
228        for &val in chi.iter() {
229            assert!(val.abs() < 1e-6, "Zero field should give near-zero chi");
230        }
231    }
232
233    #[test]
234    fn test_rts_finite() {
235        let n = 8;
236        let field: Vec<f64> = (0..n*n*n).map(|i| (i as f64) * 0.001).collect();
237        let mask = vec![1u8; n * n * n];
238        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
239        let params = RtsParams { delta: 0.15, mu: 1e5, rho: 10.0, tol: 1e-2, max_iter: 5, lsmr_iter: 2 };
240
241        let chi = rts(&field, &mask, &grid, (0.0, 0.0, 1.0), &params, |_, _| {});
242
243        for (i, &val) in chi.iter().enumerate() {
244            assert!(val.is_finite(), "Chi should be finite at index {}", i);
245        }
246    }
247
248    #[test]
249    fn test_rts_mask() {
250        let n = 8;
251        let field: Vec<f64> = (0..n*n*n).map(|i| (i as f64) * 0.001).collect();
252        let mut mask = vec![1u8; n * n * n];
253        // Zero out some mask values
254        mask[0] = 0;
255        mask[10] = 0;
256        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
257        let params = RtsParams { delta: 0.15, mu: 1e5, rho: 10.0, tol: 1e-2, max_iter: 5, lsmr_iter: 2 };
258
259        let chi = rts(&field, &mask, &grid, (0.0, 0.0, 1.0), &params, |_, _| {});
260
261        assert_eq!(chi[0], 0.0, "Masked voxel should be zero");
262        assert_eq!(chi[10], 0.0, "Masked voxel should be zero");
263    }
264
265    /// Verify parallel and sequential RTS produce identical results.
266    #[cfg(feature = "parallel")]
267    #[test]
268    fn test_rts_parallel_matches_sequential() {
269        let n = 16;
270        let field: Vec<f64> = (0..n*n*n).map(|i| ((i as f64) * 0.7).sin() * 0.01).collect();
271        let mask = vec![1u8; n * n * n];
272        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
273        let params = RtsParams { delta: 0.15, mu: 1e5, rho: 10.0, tol: 1e-4, max_iter: 20, lsmr_iter: 4 };
274
275        // Sequential (1 thread)
276        let pool_1 = rayon::ThreadPoolBuilder::new().num_threads(1).build().unwrap();
277        let chi_seq = pool_1.install(|| {
278            rts(&field, &mask, &grid, (0.0, 0.0, 1.0), &params, |_, _| {})
279        });
280
281        // Parallel (default threads)
282        let chi_par = rts(&field, &mask, &grid, (0.0, 0.0, 1.0), &params, |_, _| {});
283
284        // Compare
285        for (i, (s, p)) in chi_seq.iter().zip(chi_par.iter()).enumerate() {
286            assert!(
287                (s - p).abs() < 1e-10,
288                "RTS mismatch at voxel {}: seq={} par={} diff={}",
289                i, s, p, (s - p).abs()
290            );
291        }
292    }
293}
qsm_core/inversion/rts.rs

qsm_core/inversion/
rts.rs
