diff --git a/.gitignore b/.gitignore index eaf10a2..bfeff3c 100644 --- a/.gitignore +++ b/.gitignore @@ -151,3 +151,6 @@ warpkit/build-tmp*/ warpkit/share/ # ignore .secrets + +# macOS +.DS_Store diff --git a/README.md b/README.md index 72a9787..733568f 100644 --- a/README.md +++ b/README.md @@ -66,9 +66,9 @@ from warpkit.utilities import displacement_map_to_field # each list entry is a different echo phases = [nib.load(p) for p in phases_paths] magnitudes = [nib.load(p) for p in magnitude_paths] -TEs = [TE1, TE2, ...] # milliseconds -total_readout_time = ... # seconds -phase_encoding_direction = ... # one of i, j, k, i-, j-, k-, x, y, z, x-, y-, z- +TEs = [TE1, TE2, ...] # milliseconds +total_readout_time = ... # seconds +phase_encoding_direction = ... # one of i, j, k, i-, j-, k-, x, y, z, x-, y-, z- field_maps_native, displacement_maps, field_maps = medic( phases, magnitudes, TEs, total_readout_time, phase_encoding_direction diff --git a/tests/conftest.py b/tests/conftest.py index b3bf045..7779913 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -9,6 +9,7 @@ # get this directory THISDIR = Path(__file__).parent TEST_DATA_DIR = THISDIR / "data" / "test_data" +BRANCH_FLIP_DIR = THISDIR / "data" / "branch_flip" # fixture for test data @@ -31,6 +32,37 @@ def test_data(): } +@fixture(scope="session") +def branch_flip_data(): + """Two frames from `ds006131` sub-20828 run-01 that expose a real + ``correct_global`` branch flip. + + That subject's global field sits at 17.7 Hz against a 20.22 Hz half-wrap + (TEs 14.2/38.93/63.66 ms), so ROMEO's median of rounded wrap counts is on a + knife edge and tips over on individual frames. Frame 0 here is the original + frame 43 and is healthy; frame 1 is the original frame 44, where the + dual-echo field flips a full 40.44 Hz wrap. + + Cropped to the brain bounding box, which preserves the behaviour exactly. + It is *not* decimated: resampling changes which voxels vote in the ballot + and tips it the other way, which would destroy the very thing under test. + """ + mag = sorted(BRANCH_FLIP_DIR.glob("*part-mag*.nii.gz")) + phase = sorted(BRANCH_FLIP_DIR.glob("*part-phase*.nii.gz")) + sidecar = sorted(BRANCH_FLIP_DIR.glob("*part-mag*.json")) + metadata = [] + for s in sidecar: + with s.open() as f: + metadata.append(load(f)) + return { + "phase": [cast(Nifti1Image, nib.load(str(p))) for p in phase], + "mag": [cast(Nifti1Image, nib.load(str(m))) for m in mag], + "tes": [m["EchoTime"] * 1000 for m in metadata], + # frame index -> branch the selector must choose + "expected_branch": {0: 0, 1: 1}, + } + + @fixture(scope="session") def test_data_paths(): """File paths for the bundled BIDS-style MEDIC test data, suitable for diff --git a/tests/data/branch_flip/sub-20828_echo-1_part-mag_bold.json b/tests/data/branch_flip/sub-20828_echo-1_part-mag_bold.json new file mode 100644 index 0000000..f133cc6 --- /dev/null +++ b/tests/data/branch_flip/sub-20828_echo-1_part-mag_bold.json @@ -0,0 +1,5 @@ +{ + "EchoTime": 0.0142, + "RepetitionTime": 1.761, + "_comment": "ds006131 sub-20828 run-01, frames 43 (healthy) and 44 (correct_global flip)" +} \ No newline at end of file diff --git a/tests/data/branch_flip/sub-20828_echo-1_part-mag_bold.nii.gz b/tests/data/branch_flip/sub-20828_echo-1_part-mag_bold.nii.gz new file mode 100644 index 0000000..4987a17 Binary files /dev/null and b/tests/data/branch_flip/sub-20828_echo-1_part-mag_bold.nii.gz differ diff --git a/tests/data/branch_flip/sub-20828_echo-1_part-phase_bold.nii.gz b/tests/data/branch_flip/sub-20828_echo-1_part-phase_bold.nii.gz new file mode 100644 index 0000000..0816784 Binary files /dev/null and b/tests/data/branch_flip/sub-20828_echo-1_part-phase_bold.nii.gz differ diff --git a/tests/data/branch_flip/sub-20828_echo-2_part-mag_bold.json b/tests/data/branch_flip/sub-20828_echo-2_part-mag_bold.json new file mode 100644 index 0000000..14a9926 --- /dev/null +++ b/tests/data/branch_flip/sub-20828_echo-2_part-mag_bold.json @@ -0,0 +1,5 @@ +{ + "EchoTime": 0.03893, + "RepetitionTime": 1.761, + "_comment": "ds006131 sub-20828 run-01, frames 43 (healthy) and 44 (correct_global flip)" +} \ No newline at end of file diff --git a/tests/data/branch_flip/sub-20828_echo-2_part-mag_bold.nii.gz b/tests/data/branch_flip/sub-20828_echo-2_part-mag_bold.nii.gz new file mode 100644 index 0000000..1a19a82 Binary files /dev/null and b/tests/data/branch_flip/sub-20828_echo-2_part-mag_bold.nii.gz differ diff --git a/tests/data/branch_flip/sub-20828_echo-2_part-phase_bold.nii.gz b/tests/data/branch_flip/sub-20828_echo-2_part-phase_bold.nii.gz new file mode 100644 index 0000000..a7e0fa0 Binary files /dev/null and b/tests/data/branch_flip/sub-20828_echo-2_part-phase_bold.nii.gz differ diff --git a/tests/data/branch_flip/sub-20828_echo-3_part-mag_bold.json b/tests/data/branch_flip/sub-20828_echo-3_part-mag_bold.json new file mode 100644 index 0000000..285b9df --- /dev/null +++ b/tests/data/branch_flip/sub-20828_echo-3_part-mag_bold.json @@ -0,0 +1,5 @@ +{ + "EchoTime": 0.06366, + "RepetitionTime": 1.761, + "_comment": "ds006131 sub-20828 run-01, frames 43 (healthy) and 44 (correct_global flip)" +} \ No newline at end of file diff --git a/tests/data/branch_flip/sub-20828_echo-3_part-mag_bold.nii.gz b/tests/data/branch_flip/sub-20828_echo-3_part-mag_bold.nii.gz new file mode 100644 index 0000000..f386004 Binary files /dev/null and b/tests/data/branch_flip/sub-20828_echo-3_part-mag_bold.nii.gz differ diff --git a/tests/data/branch_flip/sub-20828_echo-3_part-phase_bold.nii.gz b/tests/data/branch_flip/sub-20828_echo-3_part-phase_bold.nii.gz new file mode 100644 index 0000000..95f4072 Binary files /dev/null and b/tests/data/branch_flip/sub-20828_echo-3_part-phase_bold.nii.gz differ diff --git a/tests/test_unwrap.py b/tests/test_unwrap.py index 94deaf7..0c7e6ca 100644 --- a/tests/test_unwrap.py +++ b/tests/test_unwrap.py @@ -2,7 +2,14 @@ import numpy as np import pytest from numpy.testing import assert_allclose -from warpkit.unwrap import compute_field_maps, compute_offset, reject_outliers +from warpkit import unwrap as unwrap_mod +from warpkit.unwrap import ( + _branch_intercept_step, + _select_branch, + compute_field_maps, + compute_offset, + reject_outliers, +) # --------------------------------------------------------------------------- # reject_outliers: median + MAD, threshold m=2.0 @@ -140,6 +147,278 @@ def test_compute_field_maps_rejects_mismatched_spatial_shape(): compute_field_maps(unwrapped, bad_masks, mag, tes) +# --------------------------------------------------------------------------- +# Global 2*pi branch selection (replaces the field-magnitude heuristic cascade) +# --------------------------------------------------------------------------- + + +def _patch_scores(monkeypatch, scores, fields=None): + """Stub _evaluate_branch so only the decision rule is under test. + + ``scores`` are intercepts in radians; ``fields`` are per-branch field maps + (constant-valued), consulted only when the consistency test ties. + """ + fields = fields or dict.fromkeys(scores, 0.0) + + def fake(n_wraps, *args, **kwargs): + fmap = np.full((2, 2, 2), fields[n_wraps], dtype=np.float32) + return scores[n_wraps], None, fmap, None + + monkeypatch.setattr(unwrap_mod, "_evaluate_branch", fake) + + +def _sel(te0: float = 12.0, te1: float = 28.97): + """Call the selector with the bundled protocol's TEs by default. + + ``_evaluate_branch`` is stubbed out by ``_patch_scores`` in every test that + uses this, so the array arguments are never read -- they just need to be the + right type. Defaults give the bundled protocol, k = 0.7071. + """ + vol = np.zeros((2, 2, 2), dtype=np.float32) + flags = np.ones((2, 2, 2), dtype=bool) + return _select_branch( + vol, # unwrapped_diff + vol, # phase0 + vol, # phase1 + np.ones((2, 2, 2), dtype=np.float32), # mag0 + vol, # mag1 + np.float32(te0), + np.float32(te1), + flags, # mask + flags, # score_mask + ) + + +# --------------------------------------------------------------------------- +# _branch_intercept_step: how far apart candidate branches sit, from TEs alone +# --------------------------------------------------------------------------- + + +def test_branch_intercept_step_matches_measured_offset(): + """The bundled protocol's wrong branches measure an intercept of 1.8402 rad. + The analytic step must reproduce that without touching any data.""" + assert _branch_intercept_step(np.float32(12.0), np.float32(28.97)) == pytest.approx( + 1.8402, rel=1e-4 + ) + + +def test_branch_intercept_step_matches_ds007637(): + assert _branch_intercept_step( + np.float32(14.20), np.float32(38.93) + ) == pytest.approx(2.6754, rel=1e-4) + + +@pytest.mark.parametrize(("te0", "te1"), [(10.0, 20.0), (15.0, 20.0), (20.0, 30.0)]) +def test_branch_intercept_step_zero_for_integer_ratio(te0, te1): + """te0/dTE integer -> the branch cannot move the offset at all.""" + assert _branch_intercept_step(np.float32(te0), np.float32(te1)) == pytest.approx( + 0.0 + ) + + +def test_branch_intercept_step_rejects_nonincreasing_tes(): + assert _branch_intercept_step(np.float32(20.0), np.float32(20.0)) == 0.0 + + +def test_branch_selector_noop_for_integer_te_ratio(monkeypatch): + """With an integer te0/dTE the selector must not act, even on a score + spread that would otherwise look decisive -- the branch is a no-op there, + so any apparent difference is numerical noise.""" + _patch_scores(monkeypatch, {-1: 1.84, 0: 1.84, 1: 1.0e-7}) + best, _ = _sel(te0=10.0, te1=20.0) + assert best == 0 + + +def test_branch_selector_corrects_inconsistent_zero(monkeypatch): + """N=0 carries an intercept and exactly one alternative does not: move.""" + _patch_scores(monkeypatch, {-1: 1.8402, 0: 1.8402, 1: 1.0e-7}) + best, _ = _sel() + assert best == 1 + + +def test_branch_selector_tie_keeps_zero_when_its_field_is_smaller(monkeypatch): + """A healthy ds006131 frame: N=0 and N=-1 are both through-origin, so the + field prior decides. N=0 has the smaller |field|, so nothing moves. + + A bare argmin on the intercepts would flip between them on numerical noise + and inject a full wrap of field into the time series. + """ + _patch_scores( + monkeypatch, + {-1: 9.12e-08, 0: 3.86e-08, 1: 0.9324}, + fields={-1: -24.29, 0: 16.15, 1: 16.15}, + ) + best, _ = _sel(te0=14.2, te1=38.93) + assert best == 0 + + +def test_branch_selector_tie_recovers_correct_global_flip(monkeypatch): + """The ds006131 sub-20828 failure, frame 44. + + correct_global's ballot tipped over the half-wrap boundary and the + dual-echo field flipped a full wrap, from +16.2 Hz to -22.8 Hz. Branch 0 + and branch +1 are both through-origin so the intercept cannot rank them, + but +1 restores +17.7 Hz against 0's -22.8 Hz and the prior picks it. + + This is the case that motivated keeping the tiebreaker: 19 of 243 frames in + that run, each a single-frame ~40 Hz excursion in the field time series. + """ + _patch_scores( + monkeypatch, + {-1: 0.9324, 0: 8.64e-08, 1: 6.43e-08}, + fields={-1: -22.78, 0: -22.78, 1: 17.66}, + ) + best, _ = _sel(te0=14.2, te1=38.93) + assert best == 1 + + +def test_branch_selector_noop_when_all_candidates_tie(monkeypatch): + """Integer TE0/dTE makes the branch a no-op; scores are all ~equal.""" + _patch_scores( + monkeypatch, + {-1: 1.1e-8, 0: 2.0e-8, 1: 1.9e-8}, + fields={-1: -58.8, 0: 0.1, 1: 58.9}, + ) + best, _ = _sel() + assert best == 0 + + +def test_branch_selector_noop_when_every_candidate_fits_perfectly(monkeypatch): + """All intercepts are exactly 0, so every candidate is through-origin. The + observed scale collapses to 0, leaving no wrong-branch magnitude to calibrate + a cutoff against, so the selector bails before the consistency test runs at + all -- the field prior never gets to break this tie.""" + _patch_scores(monkeypatch, {-1: 0.0, 0: 0.0, 1: 0.0}) + best, _ = _sel() + assert best == 0 + + +def test_branch_selector_noop_when_nothing_fits(monkeypatch): + """Every candidate carries roughly a full step of intercept, so none of them + explains the phase. A failed fit is not evidence for any branch, so change + nothing rather than take the least-bad one.""" + _patch_scores(monkeypatch, {-1: 1.71, 0: 1.80, 1: 1.74}) + best, _ = _sel() + assert best == 0 + + +def test_branch_selector_recovers_injected_wrap(test_data): + """End-to-end: shift the unwrapped difference by a known number of wraps and + confirm the selector undoes it, landing on the same phase offset.""" + from warpkit.utilities import create_brain_mask, rescale_phase + from warpkit.warpkit_cpp import romeo_unwrap3d + + phase, mag, tes = test_data["phase"], test_data["mag"], test_data["tes"] + tes = np.asarray(tes, dtype=np.float32) + raw = np.stack([p.dataobj[..., 0] for p in phase], axis=-1) + mn = min(float(np.asarray(p.dataobj[..., 0]).min()) for p in phase) + mx = max(float(np.asarray(p.dataobj[..., 0]).max()) for p in phase) + ph = rescale_phase(raw, min=mn, max=mx).astype(np.float32) + mg = np.stack([m.dataobj[..., 0] for m in mag], axis=-1).astype(np.float32) + + mag0, mag1 = mg[..., 0], mg[..., 1] + phase0, phase1 = ph[..., 0], ph[..., 1] + mask = create_brain_mask(mag0, 3) + score_mask = create_brain_mask(mag0, -2) + + signal_diff = mag0 * mag1 * np.exp(1j * (phase1 - phase0)) + unwrapped_diff = romeo_unwrap3d( + phase=np.angle(signal_diff).astype(np.float32), + weights="romeo", + mag=np.abs(signal_diff).astype(np.float32), + mask=mask, + correct_global=True, + ) + + for injected in (-1, 0, 1): + best, scores = _select_branch( + unwrapped_diff + 2 * np.pi * injected, + phase0, + phase1, + mag0, + mag1, + tes[0], + tes[1], + mask, + score_mask, + ) + assert best == -injected, f"injected {injected}, got {best} (scores={scores})" + + +def _branch_flip_frame(branch_flip_data, index): + """Reconstruct one fixture frame's selector inputs, as mcpc_3d_s would.""" + from warpkit.utilities import create_brain_mask, rescale_phase + from warpkit.warpkit_cpp import romeo_unwrap3d + + phase, mag = branch_flip_data["phase"], branch_flip_data["mag"] + tes = np.asarray(branch_flip_data["tes"], dtype=np.float32) + mn = min(float(np.asarray(p.dataobj[..., 0]).min()) for p in phase) + mx = max(float(np.asarray(p.dataobj[..., 0]).max()) for p in phase) + ph = rescale_phase( + np.stack([p.dataobj[..., index] for p in phase], axis=-1), min=mn, max=mx + ).astype(np.float32) + mg = np.stack([m.dataobj[..., index] for m in mag], axis=-1).astype(np.float32) + mag0, mag1 = mg[..., 0], mg[..., 1] + phase0, phase1 = ph[..., 0], ph[..., 1] + mask = create_brain_mask(mag0, 3) + score_mask = create_brain_mask(mag0, -2) + signal_diff = mag0 * mag1 * np.exp(1j * (phase1 - phase0)) + unwrapped_diff = romeo_unwrap3d( + phase=np.angle(signal_diff).astype(np.float32), + weights="romeo", + mag=np.abs(signal_diff).astype(np.float32), + mask=mask, + correct_global=True, + ) + # ordered to splat straight into _select_branch / _evaluate_branch + return ( + unwrapped_diff, + phase0, + phase1, + mag0, + mag1, + tes[0], + tes[1], + mask, + score_mask, + ) + + +@pytest.mark.parametrize("index", [0, 1]) +def test_branch_selector_fixes_correct_global_flip(branch_flip_data, index): + """Regression: `ds006131` sub-20828, a real ``correct_global`` branch flip. + + Frame 0 is healthy and must be left alone; frame 1 is one of the 19 frames + (of 243) where ROMEO's ballot tipped and the dual-echo field jumped a full + 40.44 Hz wrap. The selector must return +1 there to undo it. + + Before the field prior was in place this run produced 34 wrap-sized steps + in the field time series with a 35.9 Hz maximum; afterwards, zero, with a + 1.5 Hz maximum. + """ + args = _branch_flip_frame(branch_flip_data, index) + best, scores = _select_branch(*args) + expected = branch_flip_data["expected_branch"][index] + assert best == expected, f"frame {index}: got {best}, want {expected} ({scores})" + + +def test_branch_flip_frame_is_a_genuine_tie(branch_flip_data): + """The intercept test alone cannot fix the flip -- it needs the field prior. + + On the broken frame, branch 0 and branch +1 are *both* through-origin at + ~1e-7 while branch -1 carries most of a step. Any rule that only classifies + by intercept has to defer here, which would leave the 40 Hz error in place. + This is the evidence for keeping the tiebreaker, so pin it. + """ + args = _branch_flip_frame(branch_flip_data, 1) + te0, te1 = args[5], args[6] + step = _branch_intercept_step(te0, te1) + scores = {n: unwrap_mod._evaluate_branch(n, *args)[0] for n in (-1, 0, 1)} + cutoff = min(max(scores.values()), step) / 2 + fits = sorted(n for n, s in scores.items() if s < cutoff) + assert fits == [0, 1], f"expected a 0/+1 tie, got {fits} from {scores}" + + def test_romeo_unwrap3d_rejects_unknown_weight_preset(): """`weights` is a preset name string; only "romeo" is supported.""" from warpkit.warpkit_cpp import romeo_unwrap3d diff --git a/warpkit/distortion.py b/warpkit/distortion.py index 7296f22..a0f6797 100644 --- a/warpkit/distortion.py +++ b/warpkit/distortion.py @@ -19,7 +19,6 @@ def medic( svd_filt: int = 10, n_cpus: int = 4, debug: bool = False, - wrap_limit: bool = False, ) -> tuple[nib.Nifti1Image, nib.Nifti1Image, nib.Nifti1Image]: """This runs Multi-Echo DIstortion Correction (MEDIC) on a set of phase and magnitude images. @@ -101,7 +100,6 @@ def medic( frames=frames, n_cpus=n_cpus, debug=debug, - wrap_limit=wrap_limit, ) except IndexError as e: raise IndexError( diff --git a/warpkit/scripts/medic.py b/warpkit/scripts/medic.py index 84ab22d..6080d11 100644 --- a/warpkit/scripts/medic.py +++ b/warpkit/scripts/medic.py @@ -42,7 +42,6 @@ def medic( metadata: Sequence[PathLike[str] | str] | None = None, noise_frames: int = 0, n_cpus: int = 4, - wrap_limit: bool = False, debug: bool = False, ) -> MedicResult: """Run the full MEDIC pipeline and write the three output NIfTIs. @@ -99,7 +98,6 @@ def medic( border_filt=(1000, 1000), svd_filt=1000, debug=True, - wrap_limit=wrap_limit, ) else: fmaps_native, dmaps, fmaps = _medic_distortion( @@ -111,7 +109,6 @@ def medic( n_cpus=n_cpus, svd_filt=10, border_size=5, - wrap_limit=wrap_limit, ) return write_medic_outputs(out_prefix, fmaps_native, dmaps, fmaps) @@ -155,11 +152,6 @@ def main(): ) add_n_cpus_arg(parser) parser.add_argument("--debug", action="store_true", help="Debug mode") - parser.add_argument( - "--wrap-limit", - action="store_true", - help="Turns off some heuristics for phase unwrapping", - ) args = parser.parse_args() setup_logging() @@ -176,7 +168,6 @@ def main(): metadata=args.metadata, noise_frames=args.noiseframes, n_cpus=args.n_cpus, - wrap_limit=args.wrap_limit, debug=args.debug, ) except ValueError as e: diff --git a/warpkit/scripts/unwrap_phase.py b/warpkit/scripts/unwrap_phase.py index 2664a63..673a892 100644 --- a/warpkit/scripts/unwrap_phase.py +++ b/warpkit/scripts/unwrap_phase.py @@ -42,7 +42,6 @@ def unwrap_phase( metadata: Sequence[PathLike[str] | str] | None = None, noise_frames: int = 0, n_cpus: int = 4, - wrap_limit: bool = False, debug: bool = False, ) -> UnwrapPhaseResult: """Run ROMEO multi-echo phase unwrapping. @@ -87,7 +86,6 @@ def unwrap_phase( list(tes_ms), n_cpus=n_cpus, debug=debug, - wrap_limit=wrap_limit, ) out_prefix_str = str(out_prefix) @@ -148,11 +146,6 @@ def main(): action="store_true", help="Skip the temporal consistency pass and dump intermediate files.", ) - parser.add_argument( - "--wrap-limit", - action="store_true", - help="Turn off some heuristics for phase unwrapping.", - ) args = parser.parse_args() setup_logging() @@ -167,7 +160,6 @@ def main(): metadata=args.metadata, noise_frames=args.noiseframes, n_cpus=args.n_cpus, - wrap_limit=args.wrap_limit, debug=args.debug, ) except ValueError as e: diff --git a/warpkit/unwrap.py b/warpkit/unwrap.py index c49f12d..10aebf2 100644 --- a/warpkit/unwrap.py +++ b/warpkit/unwrap.py @@ -25,8 +25,11 @@ ) from .warpkit_cpp import romeo_unwrap3d, romeo_unwrap4d, romeo_voxelquality -FMAP_PROPORTION_HEURISTIC = 0.25 -FMAP_AMBIGUIOUS_HEURISTIC = 0.5 +# Candidate global 2*pi branches scanned by the intercept selector. One wrap is +# 1/dTE Hz of global field (~59 Hz for a 17 ms echo spacing), and ROMEO's +# ``correct_global`` already pins |median field| below 1/(2*dTE), so +/-1 covers +# everything reachable in practice. +BRANCH_CANDIDATES = (-1, 0, 1) def reject_outliers(data, m=2.0): @@ -76,15 +79,209 @@ def get_dual_echo_fieldmap(phases, tes, mags, mask): return fieldmap, unwrapped_phases +def _evaluate_branch( + n_wraps: int, + unwrapped_diff: npt.NDArray[np.float32], + phase0: npt.NDArray[np.float32], + phase1: npt.NDArray[np.float32], + mag0: npt.NDArray[np.float32], + mag1: npt.NDArray[np.float32], + te0: np.float32 | float, + te1: np.float32 | float, + mask: npt.NDArray[np.bool_], + score_mask: npt.NDArray[np.bool_], +): + """Reconstruct one candidate global 2*pi branch and score how well it fits. + + Shifting the unwrapped phase difference by ``2*pi*n_wraps`` moves the + MCPC-3D-S phase offset by ``wrap(2*pi*n_wraps*te0/dTE)``. That lands as the + same additive constant on every echo, so the echoes no longer extrapolate + back through zero at TE=0. Fitting a line through the two echoes and + reading off its intercept measures that constant directly: the right branch + gives ~0, a wrong branch gives roughly one ``_branch_intercept_step``. + + The intercept is taken as a median over ``score_mask`` because the error is + a single global constant, not a per-voxel effect. + + Returns + ------- + intercept : float + |median intercept| in radians over ``score_mask``. Zero means the + echoes are proportional to TE, which is what a correct offset gives. + offset : npt.NDArray[np.float32] + Phase offset implied by this branch. + fieldmap : npt.NDArray[np.float32] + Dual-echo field map in Hz implied by this branch. + unwrapped_phases : npt.NDArray[np.float32] + Dual-echo unwrapped phases implied by this branch. + """ + shifted = unwrapped_diff + 2 * np.pi * n_wraps + offset = np.angle(np.exp(1j * (phase0 - ((te0 * shifted) / (te1 - te0))))) + proposed_phases = ( + np.stack([phase0, phase1], axis=-1) - offset[..., np.newaxis] + ).astype(np.float32) + fieldmap, unwrapped_phases = get_dual_echo_fieldmap( + proposed_phases, + np.array([te0, te1], dtype=np.float32), + np.stack([mag0, mag1], axis=-1).astype(np.float32), + mask, + ) + y0 = unwrapped_phases[score_mask, 0].astype(np.float64) + y1 = unwrapped_phases[score_mask, 1].astype(np.float64) + if y0.size == 0: + return float("inf"), offset, fieldmap, unwrapped_phases + slope = (y1 - y0) / (float(te1) - float(te0)) + intercept = float(abs(np.median(y0 - slope * float(te0)))) + return intercept, offset, fieldmap, unwrapped_phases + + +def _weighted_median(values: npt.NDArray, weights: npt.NDArray) -> float: + """Weighted median of ``values``. Used with magnitude-squared weights so + low-SNR voxels do not sway the global field estimate. + + A median, not a mean or an M-estimator. The field distribution over the + brain is wide (~2 wraps) and right-skewed, and the 1/(2*dTE) threshold this + feeds sits *inside* its bulk, between the 50th and 75th percentiles. So the + question the threshold asks -- is most of the brain within half a wrap of + zero -- is a question about the 50% point, and the median is the statistic + that answers it. Anything pulled toward the mean reads ~2 Hz higher and + moves the estimate toward the boundary for no gain in accuracy: a Huber + M-estimator tried here cut the worst-case margin from 1.5 Hz to 0.05 Hz. + + Nor an argmax. The threshold comparison has to give the same answer on + consecutive frames, and a median moves smoothly with the data; field + distributions are often bimodal and a near-tie between peaks makes an + argmax flip frame to frame. A half-sample mode was tried and reverted for + exactly that reason. + """ + order = np.argsort(values) + values, weights = values[order], weights[order] + cumulative = np.cumsum(weights) + if cumulative[-1] <= 0: + return float(np.median(values)) + return float(values[np.searchsorted(cumulative, 0.5 * cumulative[-1])]) + + +def _branch_intercept_step(te0: np.float32 | float, te1: np.float32 | float) -> float: + """Intercept, in radians, that one wrap of branch error introduces. + + Candidate branches are spaced exactly this far apart in intercept, so it is + the natural scale for "this branch does not fit". It depends only on the + echo times -- no data required. + + Returns 0.0 when te0/dTE is an integer: there the branch does not move the + phase offset at all, so the candidates are indistinguishable (and the choice + cannot affect the output either). + """ + dte = float(te1) - float(te0) + if dte <= 0: + return 0.0 + return float(abs(np.angle(np.exp(1j * 2 * np.pi * float(te0) / dte)))) + + +def _select_branch( + unwrapped_diff: npt.NDArray[np.float32], + phase0: npt.NDArray[np.float32], + phase1: npt.NDArray[np.float32], + mag0: npt.NDArray[np.float32], + mag1: npt.NDArray[np.float32], + te0: np.float32 | float, + te1: np.float32 | float, + mask: npt.NDArray[np.bool_], + score_mask: npt.NDArray[np.bool_], +) -> tuple[int, dict[int, float]]: + """Pick the global 2*pi branch, in two stages. + + 1. **Consistency.** A branch carrying a leftover intercept is not a valid + explanation of the data at all, so discard those. If exactly one + candidate survives it is the answer, whatever its field looks like. If + none survives, return 0 -- a failed fit is not evidence for any branch. + 2. **Prior, only to break a tie.** Depending on TE0/dTE several branches can + be exactly through-origin; the alias is genuinely reachable and no + statistic computed from the phase can separate them. Among the survivors, + take the smallest weighted-median field. + + Stage 2 is the same prior ``correct_global`` already applies, but on a much + better conditioned statistic: weighted by ``mag1**2``, over an eroded brain + mask, on the field itself rather than an unweighted median of rounded wrap + counts over a dilated mask. That difference is the whole point. + ``correct_global``'s ballot becomes unstable when the global field sits near + 1/(2*dTE), and then it flips frame to frame. + + The weight is the *second* echo's magnitude because the field comes from + ``phase1 - phase0``, whose noise is dominated by the weaker echo. Weighting + on ``mag0`` instead lets voxels with fast T2* decay -- a healthy echo 0 and a + collapsed echo 1, i.e. exactly the air-tissue interfaces -- carry full + weight. Measured over 29 runs, this does not move the decision boundary + (that is fixed at 1/(2*dTE)) and barely moves the margin; what it buys is + stability. The estimator's offset varies by 0.43 Hz across subjects under + ``mag1**2`` against 0.99 Hz under ``mag0**2``, and that between-subject + spread, not the per-subject margin, is what predicts a wrap flip. + + Measured on `ds006131` sub-20828 (k = 0.5742, wrap 40.44 Hz, half-wrap + 20.22 Hz), whose field sits at 17.7 Hz -- 2.5 Hz under the boundary. On 19 + of 243 frames ``correct_global``'s median jumps to +/-1 and the dual-echo + field flips from +16.2 Hz to -22.8 Hz, a full wrap, for a single frame at a + time. The intercept test detects that something moved (the fitting set goes + from {-1, 0} to {0, +1}) but cannot rank the two survivors. The prior can: + branch +1 restores +17.7 Hz against branch 0's -22.8 Hz, so it fixes all 19 + and leaves the other 224 untouched. + + The prior is consulted *only* between branches that already fit, so it can + never select something the data contradicts. + """ + evaluated = { + n: _evaluate_branch( + n, unwrapped_diff, phase0, phase1, mag0, mag1, te0, te1, mask, score_mask + ) + for n in BRANCH_CANDIDATES + } + scores = {n: e[0] for n, e in evaluated.items()} + # Calibrate "how big is a wrong branch" two ways and take the stricter. + # The analytic scale depends only on the TEs, so it still calibrates when + # every candidate happens to fit; the observed scale stays honest when + # ROMEO absorbs part of the intercept into its own 2*pi steps and the real + # penalty comes out smaller than theory predicts. Taking the minimum makes + # the consistency test harder to pass, which biases toward doing nothing. + step = _branch_intercept_step(te0, te1) + if step <= 0: + # te0/dTE is an integer: the branch does not move the phase offset, so + # every candidate returns the same result and there is nothing to pick. + return 0, scores + scale = min(max(scores.values()), step) + if not np.isfinite(scale) or scale <= 0: + return 0, scores + # Nearest-neighbour on the lattice: candidate intercepts sit at 0 or at one + # ``scale``, so the boundary is the midpoint. Not a tuned threshold -- the + # two clusters are separated by 6-9 orders of magnitude on measured data, so + # any boundary strictly inside (0, scale) gives the same answer. The + # midpoint is simply the one that needs no justifying, and it keeps a branch + # that is most of a wrap out from ever counting as a fit. + cutoff = scale / 2 + consistent = [n for n, s in scores.items() if s < cutoff] + if not consistent: + # nothing fits; a failed fit is not evidence for any branch + return 0, scores + if len(consistent) == 1: + return consistent[0], scores + # Tie: every survivor explains the phase equally well, so fall back to the + # prior and take the one closest to zero global field. + weights = np.square(mag1[score_mask].astype(np.float64)) + fields = { + n: _weighted_median(evaluated[n][2][score_mask].astype(np.float64), weights) + for n in consistent + } + return min(consistent, key=lambda n: abs(fields[n])), scores + + def mcpc_3d_s( mag0: npt.NDArray[np.float32], mag1: npt.NDArray[np.float32], phase0: npt.NDArray[np.float32], phase1: npt.NDArray[np.float32], - te0: npt.NDArray[np.float32], - te1: npt.NDArray[np.float32], + te0: np.float32 | float, + te1: np.float32 | float, mask: npt.NDArray[np.bool_], - wrap_limit: bool = False, ): """Apply the MCPC-3D-S algorithm to compute the phase offset. @@ -98,14 +295,12 @@ def mcpc_3d_s( Phase image for the first echo phase1 : npt.NDArray[np.float32] Phase image for the second echo - te0 : npt.NDArray[np.float32] + te0 : np.float32 | float Echo time for the first echo - te1 : npt.NDArray[np.float32] + te1 : np.float32 | float Echo time for the second echo mask : npt.NDArray[np.bool_] Mask of voxels to use for unwrapping - wrap_limit : bool, optional - Limit the phase wrapping, by default False Returns ------- @@ -125,88 +320,17 @@ def mcpc_3d_s( correct_global=True, ) voxel_mask = create_brain_mask(mag0, -2) - phases = np.stack([phase0, phase1], axis=-1) - mags = np.stack([mag0, mag1], axis=-1) - tes = np.array([te0, te1]) - all_tes = np.array([0.0, te0, te1]) - proposed_offset = np.angle( - np.exp(1j * (phase0 - ((te0 * unwrapped_diff) / (te1 - te0)))) - ) - # get the new phases - proposed_phases = phases - proposed_offset[..., np.newaxis] - - # compute the fieldmap - proposed_fieldmap, proposed_unwrapped_phases = get_dual_echo_fieldmap( - proposed_phases, tes, mags, mask + n_wraps, scores = _select_branch( + unwrapped_diff, phase0, phase1, mag0, mag1, te0, te1, mask, voxel_mask + ) + logging.info( + "branch selection: n=%+d (intercepts=%s)", + n_wraps, + {n: f"{s_:.4e}" for n, s_ in scores.items()}, ) - # check if the proposed fieldmap is below 10 - logging.info(f"proposed_fieldmap: {proposed_fieldmap[voxel_mask].mean()}") - if proposed_fieldmap[voxel_mask].mean() < -10: - unwrapped_diff += 2 * np.pi - # check if the propossed fieldmap is between -10 and 0 - elif proposed_fieldmap[voxel_mask].mean() < 0 and not wrap_limit: - # look at proportion of voxels that are positive - voxel_prop = ( - np.count_nonzero(proposed_fieldmap[voxel_mask] > 0) - / proposed_fieldmap[voxel_mask].shape[0] - ) - - # if the proportion of positive voxels is less than 0.25, then add 2pi - if voxel_prop < FMAP_PROPORTION_HEURISTIC: - unwrapped_diff += 2 * np.pi - elif voxel_prop < FMAP_AMBIGUIOUS_HEURISTIC: - # compute mean of phase offset - mean_phase_offset = proposed_offset[voxel_mask].mean() - # print(f"mean_phase_offset: {mean_phase_offset}") - # if less than -1 then - if mean_phase_offset < -1: - phase_fits = np.concatenate( - ( - np.zeros((*proposed_unwrapped_phases.shape[:-1], 1)), - proposed_unwrapped_phases, - ), - axis=-1, - ) - _, residuals_1, _, _, _ = np.polyfit( - all_tes, phase_fits[voxel_mask, :].T, 1, full=True - ) - - # check if adding 2pi makes it better - new_proposed_offset = np.angle( - np.exp( - 1j - * ( - phase0 - - ((te0 * (unwrapped_diff + 2 * np.pi)) / (te1 - te0)) - ) - ) - ) - new_proposed_phases = phases - new_proposed_offset[..., np.newaxis] - new_proposed_fieldmap, new_proposed_unwrapped_phases = ( - get_dual_echo_fieldmap(new_proposed_phases, tes, mags, mask) - ) - # fit linear model to the proposed phases - new_phase_fits = np.concatenate( - ( - np.zeros((*new_proposed_unwrapped_phases.shape[:-1], 1)), - new_proposed_unwrapped_phases, - ), - axis=-1, - ) - _, residuals_2, _, _, _ = np.polyfit( - all_tes, new_phase_fits[voxel_mask, :].T, 1, full=True - ) - if ( - np.isclose( - residuals_1.mean(), residuals_2.mean(), atol=1e-3, rtol=1e-3 - ) - and new_proposed_fieldmap[voxel_mask].mean() > 0 - ): - unwrapped_diff += 2 * np.pi - else: - unwrapped_diff -= 2 * np.pi + unwrapped_diff = unwrapped_diff + 2 * np.pi * n_wraps # compute the phase offset return np.angle( @@ -222,7 +346,6 @@ def unwrap_phase( automask: bool = True, automask_dilation: int = 3, idx: int | None = None, - wrap_limit: bool = False, debug: bool = False, ) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.int8]]: """Unwraps the phase for a single frame of ME-EPI data. @@ -308,7 +431,6 @@ def unwrap_phase( tes[0], tes[1], mask_data, - wrap_limit=wrap_limit, ) if debug: global affine @@ -644,7 +766,6 @@ def unwrap_phases( frames: list[int] | None = None, n_cpus: int = 4, debug: bool = False, - wrap_limit: bool = False, ) -> tuple[list[nib.Nifti1Image], nib.Nifti1Image]: """Unwrap multi-echo phase per frame and enforce temporal consistency. @@ -674,8 +795,6 @@ def unwrap_phases( debug : bool, optional Skip the temporal consistency pass and dump intermediate files, by default False. - wrap_limit : bool, optional - Disable some MCPC-3D-S heuristics, by default False. Returns ------- @@ -810,7 +929,6 @@ def phase_iterator(phase, mag, tes, mask, frames, automask, automask_dilation): automask, automask_dilation, idx, - wrap_limit, debug, ) @@ -1001,7 +1119,6 @@ def unwrap_and_compute_field_maps( frames: list[int] | None = None, n_cpus: int = 4, debug: bool = False, - wrap_limit: bool = False, ) -> nib.Nifti1Image: """Unwrap phase and compute native-space field maps in a single call. @@ -1021,7 +1138,6 @@ def unwrap_and_compute_field_maps( frames=frames, n_cpus=n_cpus, debug=debug, - wrap_limit=wrap_limit, ) return compute_field_maps( unwrapped_imgs, diff --git a/warpkit/utilities.py b/warpkit/utilities.py index 6925002..3f97ff5 100644 --- a/warpkit/utilities.py +++ b/warpkit/utilities.py @@ -141,7 +141,7 @@ def corr2_coeff(a: npt.NDArray, b: npt.NDArray) -> npt.NDArray: def rescale_phase( - data: npt.NDArray[Any], min: int = -4096, max: int = 4096 + data: npt.NDArray[Any], min: float = -4096, max: float = 4096 ) -> npt.NDArray[Any]: """Rescale phase data to [-pi, pi] @@ -151,9 +151,9 @@ def rescale_phase( ---------- data : npt.NDArray[Any] phase data to be rescaled - min : int, optional + min : float, optional min value that should be mapped to -pi, by default -4096 - max : int, optional + max : float, optional max value that should be mapped to pi, by default 4096 Returns