Skip to main content

qsm_core/inversion/
tkd.rs

1//! Truncated k-space division (TKD) / TSVD for QSM
2//!
3//! TKD is the simplest dipole inversion method. It directly divides the
4//! field in k-space by the dipole kernel, with truncation to avoid
5//! division by small values near the magic angle.
6//!
7//! Reference:
8//! Shmueli, K., de Zwart, J.A., van Gelderen, P., Li, T.Q., Dodd, S.J., Duyn, J.H. (2009).
9//! "Magnetic susceptibility mapping of brain tissue in vivo using MRI phase data."
10//! Magnetic Resonance in Medicine, 62(6):1510-1522. https://doi.org/10.1002/mrm.22135
11//!
12//! Reference implementation: https://github.com/kamesy/QSM.jl
13
14use num_complex::Complex64;
15use crate::fft::{fft3d, ifft3d};
16use crate::kernels::dipole::dipole_kernel;
17
18/// TKD algorithm parameters
19#[derive(Clone, Debug)]
20pub struct TkdParams {
21    /// Truncation threshold (typically 0.1-0.2)
22    pub threshold: f64,
23}
24
25impl Default for TkdParams {
26    fn default() -> Self {
27        Self { threshold: 0.15 }
28    }
29}
30
31/// Truncated k-space division (TKD) for dipole inversion
32///
33/// Computes susceptibility map from local field using direct k-space division
34/// with threshold-based truncation.
35///
36/// # Arguments
37/// * `local_field` - Local field values (unwrapped phase / TE / gamma / B0)
38/// * `mask` - Binary mask (1 = inside ROI, 0 = outside)
39/// * `nx`, `ny`, `nz` - Array dimensions
40/// * `vsx`, `vsy`, `vsz` - Voxel sizes in mm
41/// * `bdir` - B0 field direction (bx, by, bz)
42/// * `threshold` - Truncation threshold (typically 0.1-0.2)
43///
44/// # Returns
45/// Susceptibility map (same size as input)
46pub fn tkd(
47    local_field: &[f64],
48    mask: &[u8],
49    nx: usize, ny: usize, nz: usize,
50    vsx: f64, vsy: f64, vsz: f64,
51    bdir: (f64, f64, f64),
52    threshold: f64,
53) -> Vec<f64> {
54    let n_total = nx * ny * nz;
55
56    // Generate dipole kernel
57    let d = dipole_kernel(nx, ny, nz, vsx, vsy, vsz, bdir);
58
59    // Compute inverse dipole kernel with truncation
60    // TKD: if |D| <= threshold, use sign(D)/threshold; else use 1/D
61    let inv_threshold = 1.0 / threshold;
62    let inv_d: Vec<f64> = d.iter().map(|&dval| {
63        if dval.abs() <= threshold {
64            // Truncate: use sign(D)/threshold
65            if dval >= 0.0 { inv_threshold } else { -inv_threshold }
66        } else {
67            // Normal inverse
68            1.0 / dval
69        }
70    }).collect();
71
72    // Convert local field to complex
73    let mut field_complex: Vec<Complex64> = local_field.iter()
74        .map(|&x| Complex64::new(x, 0.0))
75        .collect();
76
77    // FFT of local field
78    fft3d(&mut field_complex, nx, ny, nz);
79
80    // Multiply by inverse dipole kernel
81    for i in 0..n_total {
82        field_complex[i] *= inv_d[i];
83    }
84
85    // IFFT to get susceptibility
86    ifft3d(&mut field_complex, nx, ny, nz);
87
88    // Extract real part and apply mask
89    let mut chi: Vec<f64> = field_complex.iter()
90        .map(|c| c.re)
91        .collect();
92
93    // Apply mask
94    for i in 0..n_total {
95        if mask[i] == 0 {
96            chi[i] = 0.0;
97        }
98    }
99
100    chi
101}
102
103/// TKD with default B direction (0, 0, 1) and threshold 0.15
104pub fn tkd_default(
105    local_field: &[f64],
106    mask: &[u8],
107    nx: usize, ny: usize, nz: usize,
108    vsx: f64, vsy: f64, vsz: f64,
109) -> Vec<f64> {
110    let p = TkdParams::default();
111    tkd(local_field, mask, nx, ny, nz, vsx, vsy, vsz, (0.0, 0.0, 1.0), p.threshold)
112}
113
114/// Truncated singular value decomposition (TSVD) variant
115///
116/// Similar to TKD but zeros out values below threshold instead of truncating.
117/// This produces smoother results but may have more artifacts at the magic angle.
118pub fn tsvd(
119    local_field: &[f64],
120    mask: &[u8],
121    nx: usize, ny: usize, nz: usize,
122    vsx: f64, vsy: f64, vsz: f64,
123    bdir: (f64, f64, f64),
124    threshold: f64,
125) -> Vec<f64> {
126    let n_total = nx * ny * nz;
127
128    // Generate dipole kernel
129    let d = dipole_kernel(nx, ny, nz, vsx, vsy, vsz, bdir);
130
131    // Compute inverse dipole kernel with TSVD truncation
132    // TSVD: if |D| <= threshold, use 0; else use 1/D
133    let inv_d: Vec<f64> = d.iter().map(|&dval| {
134        if dval.abs() <= threshold {
135            0.0  // Zero out small values
136        } else {
137            1.0 / dval
138        }
139    }).collect();
140
141    // Convert local field to complex
142    let mut field_complex: Vec<Complex64> = local_field.iter()
143        .map(|&x| Complex64::new(x, 0.0))
144        .collect();
145
146    // FFT of local field
147    fft3d(&mut field_complex, nx, ny, nz);
148
149    // Multiply by inverse dipole kernel
150    for i in 0..n_total {
151        field_complex[i] *= inv_d[i];
152    }
153
154    // IFFT to get susceptibility
155    ifft3d(&mut field_complex, nx, ny, nz);
156
157    // Extract real part and apply mask
158    let mut chi: Vec<f64> = field_complex.iter()
159        .map(|c| c.re)
160        .collect();
161
162    // Apply mask
163    for i in 0..n_total {
164        if mask[i] == 0 {
165            chi[i] = 0.0;
166        }
167    }
168
169    chi
170}
171
172#[cfg(test)]
173mod tests {
174    use super::*;
175
176    #[test]
177    fn test_tkd_zero_field() {
178        // Zero field should give zero susceptibility
179        let n = 8;
180        let field = vec![0.0; n * n * n];
181        let mask = vec![1u8; n * n * n];
182
183        let chi = tkd_default(&field, &mask, n, n, n, 1.0, 1.0, 1.0);
184
185        for val in chi.iter() {
186            assert!(val.abs() < 1e-10, "Zero field should give zero chi");
187        }
188    }
189
190    #[test]
191    fn test_tkd_mask() {
192        // Values outside mask should be zero
193        let n = 8;
194        let field = vec![1.0; n * n * n];
195        let mut mask = vec![1u8; n * n * n];
196
197        // Set some voxels outside mask
198        mask[0] = 0;
199        mask[1] = 0;
200
201        let chi = tkd_default(&field, &mask, n, n, n, 1.0, 1.0, 1.0);
202
203        assert_eq!(chi[0], 0.0, "Outside mask should be 0");
204        assert_eq!(chi[1], 0.0, "Outside mask should be 0");
205    }
206
207    #[test]
208    fn test_tkd_finite() {
209        // Result should be finite (no NaN or Inf)
210        let n = 8;
211        let field: Vec<f64> = (0..n*n*n).map(|i| (i as f64) * 0.01).collect();
212        let mask = vec![1u8; n * n * n];
213
214        let chi = tkd_default(&field, &mask, n, n, n, 1.0, 1.0, 1.0);
215
216        for (i, val) in chi.iter().enumerate() {
217            assert!(val.is_finite(), "Chi should be finite at index {}", i);
218        }
219    }
220
221    #[test]
222    fn test_tsvd_vs_tkd() {
223        // TSVD and TKD should give different results near magic angle
224        let n = 16;
225        let field: Vec<f64> = (0..n*n*n).map(|i| ((i as f64) * 0.1).sin()).collect();
226        let mask = vec![1u8; n * n * n];
227
228        let chi_tkd = tkd(&field, &mask, n, n, n, 1.0, 1.0, 1.0, (0.0, 0.0, 1.0), 0.15);
229        let chi_tsvd = tsvd(&field, &mask, n, n, n, 1.0, 1.0, 1.0, (0.0, 0.0, 1.0), 0.15);
230
231        // They should be different
232        let diff: f64 = chi_tkd.iter().zip(chi_tsvd.iter())
233            .map(|(a, b)| (a - b).abs())
234            .sum();
235
236        assert!(diff > 1e-10, "TKD and TSVD should give different results");
237    }
238}