qsm_core/inversion/
admm.rs

1//! Shared ADMM iteration infrastructure for TV-based inversion methods.
2//!
3//! Provides pre-allocated buffers and a generic iteration step used by
4//! TV-ADMM, NLTV, and RTS algorithms.
5
6use num_complex::Complex64;
7use crate::fft::Fft3dWorkspace;
8use crate::Grid;
9use crate::kernels::dipole::dipole_kernel;
10use crate::kernels::laplacian::laplacian_kernel;
11use crate::utils::gradient::{bdiv_inplace, fgrad_inplace};
12use crate::utils::relative_change;
13
14/// Pre-allocated buffers for ADMM-based inversion algorithms.
15///
16/// All buffers are sized to `n_total` voxels. Zero allocations occur
17/// during iteration when using these buffers.
18pub struct AdmmBuffers {
19    /// Current solution
20    pub x: Vec<f64>,
21    /// Previous solution (for convergence)
22    pub x_prev: Vec<f64>,
23    /// Dual variables (x/y/z components)
24    pub ux: Vec<f64>,
25    pub uy: Vec<f64>,
26    pub uz: Vec<f64>,
27    /// Gradient / (z-u) buffers (dual-purpose)
28    pub gx: Vec<f64>,
29    pub gy: Vec<f64>,
30    pub gz: Vec<f64>,
31    /// Divergence buffer
32    pub div_buf: Vec<f64>,
33    /// Complex FFT work buffer
34    pub work_complex: Vec<Complex64>,
35}
36
37impl AdmmBuffers {
38    /// Allocate all ADMM buffers for a given volume size.
39    pub fn new(n_total: usize) -> Self {
40        Self {
41            x: vec![0.0; n_total],
42            x_prev: vec![0.0; n_total],
43            ux: vec![0.0; n_total],
44            uy: vec![0.0; n_total],
45            uz: vec![0.0; n_total],
46            gx: vec![0.0; n_total],
47            gy: vec![0.0; n_total],
48            gz: vec![0.0; n_total],
49            div_buf: vec![0.0; n_total],
50            work_complex: vec![Complex64::new(0.0, 0.0); n_total],
51        }
52    }
53}
54
55/// Run one ADMM iteration step.
56///
57/// Performs the x-subproblem (spectral solve), convergence check,
58/// and fused z-subproblem + dual variable update.
59///
60/// The `shrink_fn(vx, vy, vz, index) -> (zx, zy, zz)` closure captures
61/// per-algorithm differences in the proximal operator.
62///
63/// # Arguments
64/// * `buf` - Pre-allocated buffers (modified in place)
65/// * `fft_ws` - FFT workspace
66/// * `f_hat` - Pre-computed constant RHS in frequency domain
67/// * `inv_a` - Pre-computed inverse operator
68/// * `rho` - ADMM penalty multiplier for divergence term
69/// * `grid` - Volume grid
70/// * `tol` - Convergence tolerance
71/// * `shrink_fn` - Proximal operator for z-subproblem
72///
73/// # Returns
74/// `true` if converged (relative change < tol)
75#[inline]
76pub fn admm_step<S>(
77    buf: &mut AdmmBuffers,
78    fft_ws: &mut Fft3dWorkspace,
79    f_hat: &[Complex64],
80    inv_a: &[f64],
81    rho: f64,
82    grid: &Grid,
83    tol: f64,
84    shrink_fn: S,
85) -> bool
86where
87    S: Fn(f64, f64, f64, usize) -> (f64, f64, f64),
88{
89    let n_total = grid.n_total();
90
91    // Swap x and x_prev (no allocation, just pointer swap)
92    std::mem::swap(&mut buf.x, &mut buf.x_prev);
93
94    // === x-subproblem: solve in frequency domain ===
95
96    // Compute div(z - u) — gx/gy/gz hold (z-u) from previous step
97    bdiv_inplace(&mut buf.div_buf, &buf.gx, &buf.gy, &buf.gz, grid);
98
99    // FFT of divergence
100    for i in 0..n_total {
101        buf.work_complex[i] = Complex64::new(buf.div_buf[i], 0.0);
102    }
103    fft_ws.fft3d(&mut buf.work_complex);
104
105    // x_hat = f_hat - rho * FFT(div) * inv_a
106    for i in 0..n_total {
107        buf.work_complex[i] = f_hat[i] - rho * buf.work_complex[i] * inv_a[i];
108    }
109
110    // IFFT to spatial domain
111    fft_ws.ifft3d(&mut buf.work_complex);
112    for i in 0..n_total {
113        buf.x[i] = buf.work_complex[i].re;
114    }
115
116    // === Convergence check ===
117    if relative_change(&buf.x, &buf.x_prev) < tol {
118        return true;
119    }
120
121    // === Fused z-subproblem + u-update ===
122
123    // Compute gradient of x
124    fgrad_inplace(&mut buf.gx, &mut buf.gy, &mut buf.gz, &buf.x, grid);
125
126    // Apply proximal operator and update duals
127    for i in 0..n_total {
128        let vx = buf.gx[i] + buf.ux[i];
129        let vy = buf.gy[i] + buf.uy[i];
130        let vz = buf.gz[i] + buf.uz[i];
131
132        let (zx, zy, zz) = shrink_fn(vx, vy, vz, i);
133
134        buf.ux[i] = vx - zx;
135        buf.uy[i] = vy - zy;
136        buf.uz[i] = vz - zz;
137
138        // Store (2z - v) = (z - u_new) for next iteration's divergence
139        buf.gx[i] = 2.0 * zx - vx;
140        buf.gy[i] = 2.0 * zy - vy;
141        buf.gz[i] = 2.0 * zz - vz;
142    }
143
144    false
145}
146
147/// Pre-compute ADMM spectral operators (dipole kernel, inverse operator, and RHS).
148///
149/// Shared by TV-ADMM and NLTV. Returns (fft_workspace, inv_a, f_hat).
150pub fn prepare_admm_spectral(
151    local_field: &[f64],
152    grid: &Grid,
153    bdir: (f64, f64, f64),
154    rho: f64,
155) -> (Fft3dWorkspace, Vec<f64>, Vec<Complex64>) {
156    let n_total = grid.n_total();
157    let mut fft_ws = Fft3dWorkspace::new(grid.nx(), grid.ny(), grid.nz());
158    let d_kernel = dipole_kernel(grid, bdir);
159    let l_kernel = laplacian_kernel(grid, true);
160    let mut l_complex: Vec<Complex64> = l_kernel.iter()
161        .map(|&x| Complex64::new(x, 0.0))
162        .collect();
163    fft_ws.fft3d(&mut l_complex);
164    let mut inv_a: Vec<f64> = vec![0.0; n_total];
165    for i in 0..n_total {
166        let a = d_kernel[i] * d_kernel[i] + rho * l_complex[i].re;
167        inv_a[i] = if a.abs() > 1e-20 { 1.0 / a } else { 0.0 };
168    }
169    let f_hat = &mut l_complex;
170    for i in 0..n_total {
171        f_hat[i] = Complex64::new(local_field[i], 0.0);
172    }
173    fft_ws.fft3d(f_hat);
174    for i in 0..n_total {
175        f_hat[i] = f_hat[i] * d_kernel[i] * inv_a[i];
176    }
177    (fft_ws, inv_a, l_complex)
178}
179
180/// Pre-compute spectral operators shared by the FANSI-family nonlinear solvers
181/// (NDI, nlTV, nlTGV, L1-QSM, WH-QSM, HD-QSM).
182///
183/// Unlike [`prepare_admm_spectral`], this does not bake the data term into an RHS,
184/// because those solvers carry a separately-updated data auxiliary variable. It
185/// returns the reusable primitives instead:
186///
187/// * `Fft3dWorkspace` — reusable FFT plans (unnormalized forward `fft3d`,
188///   normalized inverse `ifft3d`, matching MATLAB `fftn`/`ifftn`).
189/// * `k_kernel` — the real k-space dipole kernel `D(k)` (same convention as the
190///   rest of the crate: continuous kernel, includes voxel-size scaling, DC = 0).
191/// * `ee2` — the real spectral Laplacian `EE2(k) = |E1|² + |E2|² + |E3|²`, i.e.
192///   the frequency response of `bdiv ∘ fgrad`. Use as the regularization term in
193///   the x-subproblem denominator so it stays consistent with real-space
194///   `fgrad_inplace`/`bdiv_inplace` (which share the same voxel-size scaling).
195///
196/// The x-subproblem denominator for these solvers is therefore
197/// `mu2 * k_kernel² + mu * ee2` (all real), matching FANSI's
198/// `mu2*abs(Kernel).^2 + mu*EE2`.
199pub fn prepare_fansi_spectral(
200    grid: &Grid,
201    bdir: (f64, f64, f64),
202) -> (Fft3dWorkspace, Vec<f64>, Vec<f64>) {
203    let mut fft_ws = Fft3dWorkspace::new(grid.nx(), grid.ny(), grid.nz());
204    let k_kernel = dipole_kernel(grid, bdir);
205    let l_kernel = laplacian_kernel(grid, true);
206    let mut l_complex: Vec<Complex64> = l_kernel.iter()
207        .map(|&x| Complex64::new(x, 0.0))
208        .collect();
209    fft_ws.fft3d(&mut l_complex);
210    let ee2: Vec<f64> = l_complex.iter().map(|c| c.re).collect();
211    (fft_ws, k_kernel, ee2)
212}
qsm_core/inversion/admm.rs

qsm_core/inversion/
admm.rs
