✨ Detect and correct global 2π branch errors in MCPC-3D-S

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,179 @@ 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):
+    """Stub _evaluate_branch so only the decision rule is under test.
+
+    ``scores`` are intercepts in radians. The selector consults nothing else --
+    no field prior, no tiebreak -- so the remaining tuple slots are unused.
+    """
+
+    def fake(n_wraps, *args, **kwargs):
+        return scores[n_wraps], None, None, 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_defers_when_two_branches_fit(monkeypatch):
+    """The ds007637 case: N=0 and N=-1 are both through-origin. That is a
+    genuine alias -- both explain the phase exactly -- so the selector must not
+    choose between them. Deferring leaves correct_global's answer in place.
+
+    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: 1.2e-03, 0: 1.9e-03, 1: 0.9324})
+    best, _ = _sel()
+    assert best == 0
+
+
+def test_branch_selector_defers_even_when_zero_looks_wrong(monkeypatch):
+    """N=0 carries an intercept but *two* alternatives fit. Neither is
+    preferable on the evidence, so change nothing rather than pick one.
+
+    This is the case that used to consult the field prior. Measurement showed
+    the prior was right about as often as it was wrong, so it was removed.
+    """
+    _patch_scores(monkeypatch, {-1: 1.0e-7, 0: 1.8402, 1: 2.0e-7})
+    best, _ = _sel()
+    assert best == 0
+
+
+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})
+    best, _ = _sel()
+    assert best == 0
+
+
+def test_branch_selector_noop_when_nothing_fits(monkeypatch):
+    """Degenerate/failed fit: every candidate is bad, so change nothing."""
+    _patch_scores(monkeypatch, {-1: 0.0, 0: 0.0, 1: 0.0})
+    best, _ = _sel()
+    assert best == 0
+
+    _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 test_romeo_unwrap3d_rejects_unknown_weight_preset():
     """`weights` is a preset name string; only "romeo" is supported."""
     from warpkit.warpkit_cpp import romeo_unwrap3d

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(

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:

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: