qsm_core/pipeline/motion.rs
1//! Motion-correction stage: co-register a multi-echo series onto one echo before field mapping.
2//!
3//! Runs before [`super::run_field_mapping`] and hands it the same shape of data it already takes,
4//! so a pipeline that wants motion correction inserts this and changes nothing else. The
5//! published procedure is exactly this: register every echo to the first TE as a fixed reference
6//! with a rigid algorithm, then carry on.
7//!
8//! The library half lives in [`crate::motion`]; this is the stage wrapper that speaks
9//! [`ScanMetadata`] and [`PipelineError`].
10//!
11//! # EPI distortion correction, if any, goes before this
12//!
13//! A pipeline that unwarps susceptibility-induced EPI distortion must do it **first**, on each
14//! volume in its own acquired geometry, and hand the unwarped series here.
15//!
16//! The reason is that the two corrections do not commute. The phase-encode axis an EPI
17//! displacement runs along is fixed to the gradients, not to the head, so once the head has
18//! rotated that axis points somewhere else *in head coordinates* — a few degrees of rotation
19//! tilts a ten-voxel frontal distortion by a few degrees, which is a fraction of a voxel of
20//! difference that is not rigid. Run this stage on still-distorted volumes and it is fitting a
21//! six-parameter rigid model to volumes that differ non-rigidly; it will converge, and the answer
22//! will be wrong in a way nothing downstream flags. Unwarping first leaves a residual difference
23//! that genuinely is rigid, which is the assumption the NCC optimiser is entitled to make.
24//!
25//! The honest caveat on the other side, which this stage cannot fix and an unwarping stage should
26//! record: a single distortion estimate — one field map, one reversed-PE pair — is measured at
27//! one head position, so applying it to a volume acquired after the head moved applies a stale
28//! estimate. That error is small while the motion is small, and
29//! [`MotionCorrectionResult::max_displacement_mm`] and
30//! [`MotionCorrectionResult::rotation_spread_deg`] are the numbers that say whether it still is.
31//! Resolving it properly means iterating the two stages, which nothing here does.
32//!
33//! Masks compose without either stage knowing about the other: this one intersects whatever mask
34//! it is handed with its own coverage, so an upstream stage that has already shrunk the mask to
35//! where it still has data simply passes the shrunk one in, and the result is the intersection of
36//! all of them.
37//!
38//! # What it can and cannot fix
39//!
40//! It puts every echo's anatomy back in one place, which is what lets a cross-echo fit mean
41//! anything at all. What it cannot undo is that the echoes were *acquired* with B0 in different
42//! places relative to the head, so their dipole kernels genuinely differ:
43//! [`MotionCorrectionResult::bdirs`] reports one direction per echo and
44//! [`MotionCorrectionResult::rotation_spread_deg`] the angle between the extremes. A field map
45//! fitted across them is a weighted average over that spread — the "residual boundary effect"
46//! the UTE work reports after correcting substantial motion, and the reason this stage reports
47//! the spread rather than quietly collapsing it. Below a degree or so it is negligible; a series
48//! with a large spread is really a multi-orientation acquisition, and
49//! [`crate::inversion::cosmos()`] is the honest way to use it.
50
51use super::config::{PipelineError, ScanMetadata};
52use crate::motion::{estimate_motion, uniform_series, MotionParams, MotionSeries};
53
54/// Configuration for the motion-correction stage.
55#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
56#[derive(Clone, Debug, Default)]
57pub struct MotionCorrectionConfig {
58 /// Run the stage at all. `false` passes the echoes through untouched, with identity motion
59 /// reported, so a pipeline can keep one code path whether or not correction is wanted.
60 pub enabled: bool,
61 /// Everything about *how* to register. [`MotionParams::declared_b0`] is overridden by
62 /// [`ScanMetadata::b0_direction`], which is where a pipeline's B0 direction actually comes
63 /// from.
64 pub motion: MotionParams,
65}
66
67/// A motion-corrected echo series, ready for [`super::run_field_mapping`].
68pub struct MotionCorrectionResult {
69 /// Per-echo wrapped phase on the reference echo's grid, moved through the complex domain so
70 /// the wraps survived.
71 pub phases: Vec<Vec<f64>>,
72 /// Per-echo magnitude on the reference echo's grid.
73 pub magnitudes: Vec<Vec<f64>>,
74 /// The input mask restricted to voxels every echo still reaches after being moved.
75 ///
76 /// A rotated echo cannot fill the whole reference grid, and a resampled volume is zero where
77 /// it does not reach. Those zeros are absence of data, not a zero measurement, and feeding
78 /// them to an unwrapper as if they were brain is how a rim of nonsense gets into a field
79 /// map. Pass **this** mask downstream, not the one handed in.
80 pub mask: Vec<u8>,
81 /// Each echo's B0 direction in the reference echo's voxel frame. See the module docs on why
82 /// these differ and what that costs.
83 pub bdirs: Vec<(f64, f64, f64)>,
84 /// The motion that was estimated, for reporting and for the `suspect` flags.
85 pub motion: MotionSeries,
86}
87
88impl MotionCorrectionResult {
89 /// Angle between the extreme B0 directions in the series, in degrees — how far from a single
90 /// dipole-kernel orientation this series actually is.
91 pub fn rotation_spread_deg(&self) -> f64 {
92 self.motion.rotation_spread_deg()
93 }
94
95 /// Largest distance any brain tissue moved during the series, in mm. Compare it against the
96 /// voxel size: below one voxel the correction had nothing to do.
97 pub fn max_displacement_mm(&self) -> f64 {
98 self.motion.max_displacement_mm()
99 }
100}
101
102/// Co-register a multi-echo series onto one of its echoes.
103///
104/// `phases` and `magnitudes` are per-echo and must be the same length; magnitude is required
105/// here rather than optional, because registration is estimated on magnitude and there is
106/// nothing to estimate it from without it. `mask` is the brain mask on the series' grid.
107///
108/// # Errors
109///
110/// [`PipelineError::InvalidInput`] for an empty series, a phase/magnitude count mismatch, or a
111/// reference index that does not name an echo; [`PipelineError::DimensionMismatch`] if any volume
112/// disagrees with `metadata.dims`; [`PipelineError::AlgorithmError`] if a registration could not
113/// be evaluated at all (an empty mask, or volumes with no overlap).
114pub fn run_motion_correction(
115 phases: &[&[f64]],
116 magnitudes: &[&[f64]],
117 mask: &[u8],
118 metadata: &ScanMetadata,
119 config: &MotionCorrectionConfig,
120 progress: &mut dyn FnMut(usize, usize),
121) -> Result<MotionCorrectionResult, PipelineError> {
122 let (nx, ny, nz) = metadata.dims;
123 let n_voxels = nx * ny * nz;
124 let n_echoes = phases.len();
125
126 if n_echoes == 0 {
127 return Err(PipelineError::InvalidInput("no phase echoes provided".into()));
128 }
129 if magnitudes.len() != n_echoes {
130 return Err(PipelineError::InvalidInput(format!(
131 "{n_echoes} phase echoes but {} magnitude echoes",
132 magnitudes.len()
133 )));
134 }
135 for v in phases.iter().chain(magnitudes.iter()) {
136 if v.len() != n_voxels {
137 return Err(PipelineError::DimensionMismatch {
138 expected: n_voxels,
139 got: v.len(),
140 });
141 }
142 }
143 if mask.len() != n_voxels {
144 return Err(PipelineError::DimensionMismatch {
145 expected: n_voxels,
146 got: mask.len(),
147 });
148 }
149 let reference = config
150 .motion
151 .reference
152 .resolve(n_echoes)
153 .ok_or_else(|| {
154 PipelineError::InvalidInput(format!(
155 "reference {:?} does not name one of {n_echoes} echoes",
156 config.motion.reference
157 ))
158 })?;
159
160 // Every echo of one acquisition shares a grid, so one voxel->world affine serves the series.
161 // A diagonal affine built from the voxel sizes is the right one: `ScanMetadata` carries the
162 // B0 direction in *voxel* coordinates already, and passing it as the declared direction is
163 // what makes the recovered per-echo directions come back in the same voxel frame the rest of
164 // the pipeline's dipole kernels are built in. Deriving them from a fabricated affine instead
165 // would silently assume an axial acquisition.
166 let (vsx, vsy, vsz) = metadata.voxel_size;
167 let affine = [
168 vsx, 0.0, 0.0, 0.0, //
169 0.0, vsy, 0.0, 0.0, //
170 0.0, 0.0, vsz, 0.0, //
171 0.0, 0.0, 0.0, 1.0,
172 ];
173 let params = MotionParams {
174 declared_b0: Some(metadata.b0_direction),
175 ..config.motion.clone()
176 };
177
178 if !config.enabled {
179 progress(1, 1);
180 let series = uniform_series(magnitudes, metadata.dims, &affine);
181 let motion = estimate_motion(
182 &series[reference..reference + 1],
183 Some(mask),
184 &MotionParams {
185 reference: crate::motion::MotionReference::First,
186 ..params.clone()
187 },
188 )
189 .ok_or_else(|| PipelineError::AlgorithmError("motion estimation failed".into()))?;
190 // One identity entry, repeated: no registration ran, so every echo is where it was and
191 // the reference echo's own direction is the only one there is.
192 let bdir = motion.volumes[0].bdir;
193 return Ok(MotionCorrectionResult {
194 phases: phases.iter().map(|p| p.to_vec()).collect(),
195 magnitudes: magnitudes.iter().map(|m| m.to_vec()).collect(),
196 mask: mask.to_vec(),
197 bdirs: vec![bdir; n_echoes],
198 motion: MotionSeries {
199 reference: 0,
200 volumes: (0..n_echoes)
201 .map(|index| {
202 let mut v = motion.volumes[0].clone();
203 v.index = index;
204 v
205 })
206 .collect(),
207 },
208 });
209 }
210
211 progress(0, n_echoes + 1);
212 let series = uniform_series(magnitudes, metadata.dims, &affine);
213 let motion = estimate_motion(&series, Some(mask), ¶ms)
214 .ok_or_else(|| PipelineError::AlgorithmError("rigid registration failed".into()))?;
215 progress(1, n_echoes + 1);
216
217 let mut out_phases = Vec::with_capacity(n_echoes);
218 let mut out_mags = Vec::with_capacity(n_echoes);
219 let mut coverage = vec![1u8; n_voxels];
220 for e in 0..n_echoes {
221 let (m, p) = motion
222 .apply_complex(e, magnitudes[e], phases[e])
223 .ok_or_else(|| PipelineError::AlgorithmError(format!("resampling echo {e} failed")))?;
224 let c = motion
225 .coverage(e)
226 .ok_or_else(|| PipelineError::AlgorithmError(format!("coverage for echo {e} failed")))?;
227 for (a, &b) in coverage.iter_mut().zip(c.iter()) {
228 *a &= b;
229 }
230 out_mags.push(m);
231 out_phases.push(p);
232 progress(e + 2, n_echoes + 1);
233 }
234
235 let out_mask: Vec<u8> = mask
236 .iter()
237 .zip(coverage.iter())
238 .map(|(&m, &c)| m & c)
239 .collect();
240
241 Ok(MotionCorrectionResult {
242 phases: out_phases,
243 magnitudes: out_mags,
244 bdirs: motion.bdirs(),
245 mask: out_mask,
246 motion,
247 })
248}
249
250#[cfg(test)]
251mod tests {
252 use super::*;
253 use crate::motion::MotionReference;
254
255 fn metadata(dims: (usize, usize, usize)) -> ScanMetadata {
256 ScanMetadata {
257 dims,
258 voxel_size: (1.0, 1.0, 1.0),
259 echo_times: vec![0.004, 0.012, 0.020, 0.028],
260 field_strength: 3.0,
261 b0_direction: (0.0, 0.0, 1.0),
262 slice_geometry: None,
263 affine: None,
264 }
265 }
266
267 /// Two offset blobs, enough structure for a registration and cheap enough for a unit test.
268 fn blobs(dims: (usize, usize, usize)) -> Vec<f64> {
269 let (nx, ny, nz) = dims;
270 let mut v = vec![0.0f64; nx * ny * nz];
271 for k in 0..nz {
272 for j in 0..ny {
273 for i in 0..nx {
274 let (x, y, z) = (
275 i as f64 / nx as f64,
276 j as f64 / ny as f64,
277 k as f64 / nz as f64,
278 );
279 let mut acc = 0.0;
280 for &(cx, cy, cz, r, val) in &[
281 (0.50, 0.50, 0.50, 0.30, 1.0),
282 (0.38, 0.44, 0.54, 0.10, 2.2),
283 (0.62, 0.58, 0.44, 0.09, 1.6),
284 ] {
285 let d = ((x - cx) / r).powi(2) + ((y - cy) / r).powi(2) + ((z - cz) / r).powi(2);
286 if d <= 1.0 {
287 acc += val * (1.0 - d).sqrt();
288 }
289 }
290 v[i + j * nx + k * nx * ny] = acc;
291 }
292 }
293 }
294 v
295 }
296
297 #[test]
298 fn disabled_passes_the_series_through_untouched() {
299 let dims = (20, 20, 20);
300 let meta = metadata(dims);
301 let mag = blobs(dims);
302 let pha: Vec<f64> = (0..mag.len()).map(|i| (i as f64 * 0.37).sin() * 3.0).collect();
303 let mask = vec![1u8; mag.len()];
304 let mags: Vec<&[f64]> = vec![&mag; 4];
305 let phas: Vec<&[f64]> = vec![&pha; 4];
306
307 let out = run_motion_correction(
308 &phas,
309 &mags,
310 &mask,
311 &meta,
312 &MotionCorrectionConfig::default(),
313 &mut |_, _| {},
314 )
315 .unwrap();
316 assert_eq!(out.phases.len(), 4);
317 assert_eq!(out.magnitudes.len(), 4);
318 assert_eq!(out.bdirs, vec![(0.0, 0.0, 1.0); 4]);
319 assert_eq!(out.mask, mask, "a disabled stage must not shrink the mask");
320 for e in 0..4 {
321 assert_eq!(out.phases[e], pha, "echo {e} phase must be untouched");
322 assert_eq!(out.magnitudes[e], mag, "echo {e} magnitude must be untouched");
323 assert_eq!(out.motion.volumes[e].index, e);
324 assert_eq!(out.motion.volumes[e].rotation_deg, 0.0);
325 }
326 assert_eq!(out.rotation_spread_deg(), 0.0);
327 assert_eq!(out.max_displacement_mm(), 0.0);
328 }
329
330 #[test]
331 fn bad_inputs_are_rejected() {
332 let dims = (16, 16, 16);
333 let meta = metadata(dims);
334 let mag = blobs(dims);
335 let mask = vec![1u8; mag.len()];
336 let cfg = MotionCorrectionConfig {
337 enabled: true,
338 ..Default::default()
339 };
340
341 assert!(matches!(
342 run_motion_correction(&[], &[], &mask, &meta, &cfg, &mut |_, _| {}),
343 Err(PipelineError::InvalidInput(_))
344 ));
345 assert!(
346 matches!(
347 run_motion_correction(
348 &[mag.as_slice()], &[mag.as_slice(), mag.as_slice()],
349 &mask, &meta, &cfg, &mut |_, _| {},
350 ),
351 Err(PipelineError::InvalidInput(_))
352 ),
353 "phase and magnitude counts must agree"
354 );
355 assert!(
356 matches!(
357 run_motion_correction(
358 &[&mag[..mag.len() - 1]], &[&mag[..mag.len() - 1]],
359 &mask, &meta, &cfg, &mut |_, _| {},
360 ),
361 Err(PipelineError::DimensionMismatch { .. })
362 ),
363 "a volume that disagrees with metadata.dims"
364 );
365 assert!(
366 matches!(
367 run_motion_correction(
368 &[mag.as_slice()], &[mag.as_slice()],
369 &mask[..10], &meta, &cfg, &mut |_, _| {},
370 ),
371 Err(PipelineError::DimensionMismatch { .. })
372 ),
373 "a mask that disagrees with metadata.dims"
374 );
375 let out_of_range = MotionCorrectionConfig {
376 enabled: true,
377 motion: MotionParams {
378 reference: MotionReference::Index(2),
379 ..Default::default()
380 },
381 };
382 assert!(
383 matches!(
384 run_motion_correction(
385 &[mag.as_slice()], &[mag.as_slice()],
386 &mask, &meta, &out_of_range, &mut |_, _| {},
387 ),
388 Err(PipelineError::InvalidInput(_))
389 ),
390 "a reference index past the end of the series"
391 );
392 }
393
394 /// An enabled stage moves the echoes onto the reference's grid, shrinks the mask to what
395 /// every echo still reaches, and reports one B0 direction per echo.
396 #[test]
397 fn an_enabled_stage_registers_and_shrinks_the_mask() {
398 use crate::geometry::resample_complex_onto;
399
400 let dims = (32, 32, 32);
401 let meta = metadata(dims);
402 let affine = [
403 1.0, 0.0, 0.0, 0.0, //
404 0.0, 1.0, 0.0, 0.0, //
405 0.0, 0.0, 1.0, 0.0, //
406 0.0, 0.0, 0.0, 1.0,
407 ];
408 let mag = blobs(dims);
409 let pha: Vec<f64> = (0..mag.len()).map(|i| (i as f64 * 0.37).sin() * 3.0).collect();
410 let mask = vec![1u8; mag.len()];
411
412 // Echo 1 is the reference; echoes 2-4 are shifted progressively along x, which takes a
413 // slab off the edge of the grid and so must shrink the mask.
414 let mut mags = vec![mag.clone()];
415 let mut phas = vec![pha.clone()];
416 for e in 1..4 {
417 let shift = e as f64;
418 let w = [
419 1.0, 0.0, 0.0, shift, //
420 0.0, 1.0, 0.0, 0.0, //
421 0.0, 0.0, 1.0, 0.0, //
422 0.0, 0.0, 0.0, 1.0,
423 ];
424 let (m, p) =
425 resample_complex_onto(&mag, &pha, dims, &w, dims, &affine).unwrap();
426 mags.push(m);
427 phas.push(p);
428 }
429 let mag_refs: Vec<&[f64]> = mags.iter().map(|m| m.as_slice()).collect();
430 let pha_refs: Vec<&[f64]> = phas.iter().map(|p| p.as_slice()).collect();
431
432 let mut steps = Vec::new();
433 let out = run_motion_correction(
434 &pha_refs,
435 &mag_refs,
436 &mask,
437 &meta,
438 &MotionCorrectionConfig {
439 enabled: true,
440 ..Default::default()
441 },
442 &mut |a, b| steps.push((a, b)),
443 )
444 .unwrap();
445
446 assert_eq!(out.bdirs.len(), 4);
447 assert_eq!(steps.last(), Some(&(5, 5)), "progress should reach its total");
448 let kept = out.mask.iter().filter(|&&m| m != 0).count();
449 println!("mask kept {kept} of {}", mask.len());
450 assert!(kept < mask.len(), "shifted echoes cannot cover the whole grid");
451 assert!(kept > mask.len() / 2);
452
453 // Pure translations, so every echo's B0 direction is still the declared one.
454 for (e, b) in out.bdirs.iter().enumerate() {
455 let err = (1.0f64 - b.2.abs()).abs();
456 assert!(err < 1e-3, "echo {e}: bdir {b:?} should still be +z after a translation");
457 }
458 assert!(out.rotation_spread_deg() < 0.5, "spread {}", out.rotation_spread_deg());
459 assert!(
460 out.max_displacement_mm() > 2.0,
461 "the series moved 3 mm, got {}",
462 out.max_displacement_mm()
463 );
464
465 // And the corrected echoes agree with the reference, which is the point.
466 for e in 1..4 {
467 let worst = out.magnitudes[e]
468 .iter()
469 .zip(mag.iter())
470 .zip(out.mask.iter())
471 .filter(|(_, &m)| m != 0)
472 .fold(0.0f64, |w, ((&got, &want), _)| w.max((got - want).abs()));
473 println!("echo {e}: worst magnitude difference from the reference {worst:.4}");
474 assert!(worst < 0.1, "echo {e} did not come back into register: {worst}");
475 }
476 }
477}