diff --git a/CLAUDE.md b/CLAUDE.md index becee05..e34e47b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -15,8 +15,16 @@ Pre-print: . ## Layout - `warpkit/` — Python package. Public entry points: `warpkit.distortion.medic`, - `warpkit.utilities.*`. CLI scripts in `warpkit/scripts/` are auto-discovered - by `setup.py` and exposed as `medic` and `extract_field_from_maps`. + `warpkit.utilities.*`. CLI scripts live in `warpkit/scripts/` and are + registered explicitly in `[project.scripts]` in `pyproject.toml`. All CLIs + ship with a `wk-` prefix to avoid colliding with same-named tools from + FSL/ANTs/AFNI/etc.: `wk-medic`, `wk-unwrap-phase`, `wk-compute-fieldmap`, + `wk-apply-warp`, `wk-convert-warp`, `wk-convert-fieldmap`, + `wk-compute-jacobian`. Adding a new CLI means adding a new file under + `warpkit/scripts/` *and* a new line to `[project.scripts]` — there is + no longer any auto-discovery. Shared IO helpers used by the + `wk-convert-*` and `wk-compute-jacobian` scripts live in + `warpkit/scripts/_warp_io.py` (private; not a CLI). - `warpkit/warpkit_cpp.pyi` + `warpkit/py.typed` — type info for the compiled extension, shipped via `MANIFEST.in` and `[tool.setuptools.package-data]`. Regenerate the stub after pybind11 binding changes via the wrapper script — @@ -40,6 +48,9 @@ uv sync --group dev --config-setting editable_mode=strict # tests uv run pytest -q +# coverage (matches the CI `coverage` job) +uv run coverage run && uv run coverage report -m + # lint + types (matches pre-commit; never bypass with --no-verify) uvx ruff check uvx ruff format @@ -106,9 +117,12 @@ and call out anything CI-relevant (wheel matrix, pybind11 ABI, ITK). ## CI specifics -GitHub Actions builds wheels for Python 3.11–3.14 on `ubuntu-latest` and -`macos-latest` via cibuildwheel. `pyproject.toml`'s `[tool.cibuildwheel]` -skips `*musllinux*` and free-threaded builds (`cp313t-*`, `cp314t-*`); -re-enabling free-threaded support requires auditing the pybind11 + ITK code -paths for the no-GIL ABI. PyPI publish and the GHCR Docker image only run on -a published GitHub release. +GitHub Actions builds wheels for Python 3.11–3.14 on `ubuntu-latest`, +`ubuntu-24.04-arm`, and `macos-latest` via cibuildwheel. +`pyproject.toml`'s `[tool.cibuildwheel]` skips `*musllinux*` and the +`cp314t-*` free-threaded build; re-enabling free-threaded support requires +auditing the pybind11 + ITK code paths for the no-GIL ABI. The sdist job +also runs `uv run coverage run` then `coverage report -m` — keep coverage +healthy when adding code (the `[tool.coverage.report]` config in +`pyproject.toml` omits the test files). PyPI publish and the GHCR Docker +image only run on a published GitHub release. diff --git a/Dockerfile b/Dockerfile index 3c6f356..95c5184 100644 --- a/Dockerfile +++ b/Dockerfile @@ -21,11 +21,11 @@ ENV UV_PROJECT_ENVIRONMENT=/opt/warpkit/.venv ENV UV_PYTHON_PREFERENCE=only-managed RUN cd /opt/warpkit && uv sync --group dev --config-setting editable_mode=strict -v -# put the project venv on PATH so `medic` and friends resolve +# put the project venv on PATH so `wk-medic` and friends resolve ENV PATH=/opt/warpkit/.venv/bin:${PATH} # test warpkit RUN cd /opt/warpkit && uv run pytest -s -v -# set medic script as entrypoint -ENTRYPOINT ["medic"] +# set wk-medic script as entrypoint +ENTRYPOINT ["wk-medic"] diff --git a/README.md b/README.md index fd59cfd..f5601d9 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ Pre-built wheels are published for Linux (x86_64) and macOS (universal2). If `pi docker run -it --rm ghcr.io/vanandrew/warpkit:latest --help ``` -The image's entrypoint is the `medic` CLI. +The image's entrypoint is the `wk-medic` CLI. ### From source @@ -99,12 +99,28 @@ displacement_field = displacement_map_to_field( ### CLI -A `medic` script is installed on `PATH`. Acquisition parameters can come either from BIDS sidecars or from the command line — pick one. +All warpkit CLIs are installed on `PATH` with a `wk-` prefix to avoid colliding +with same-named tools from FSL/ANTs/AFNI/etc.: + +| Command | Purpose | +| ---------------------- | ------------------------------------------------------------------------------------- | +| `wk-medic` | End-to-end MEDIC pipeline: phase + magnitude → field maps + displacement maps. | +| `wk-unwrap-phase` | Stage 1: ROMEO multi-echo phase unwrapping → unwrapped phase + per-frame masks. | +| `wk-compute-fieldmap` | Stage 2: take stage-1 outputs → native + displacement + undistorted-space field maps. | +| `wk-apply-warp` | Resample an image through a displacement map / field (single or per-frame series). | +| `wk-convert-warp` | Convert between maps ↔ fields and between ITK / FSL / ANTs / AFNI; `--invert` warps. | +| `wk-convert-fieldmap` | Convert between mm displacement maps/fields and Hz B0 field maps. | +| `wk-compute-jacobian` | Per-voxel Jacobian determinant (1 = no change, <1 = compression, >1 = expansion). | + +`wk-medic` runs the full pipeline; `wk-unwrap-phase` + `wk-compute-fieldmap` +run the same thing in two stages so you can inspect/reuse the unwrapped +phase. Acquisition parameters can come either from BIDS sidecars or from the +command line — pick one. From BIDS sidecars: ```bash -medic \ +wk-medic \ --magnitude mag_e1.nii.gz mag_e2.nii.gz mag_e3.nii.gz \ --phase phase_e1.nii.gz phase_e2.nii.gz phase_e3.nii.gz \ --metadata mag_e1.json mag_e2.json mag_e3.json \ @@ -116,7 +132,7 @@ medic \ Or by passing acquisition parameters directly: ```bash -medic \ +wk-medic \ --magnitude mag_e1.nii.gz mag_e2.nii.gz mag_e3.nii.gz \ --phase phase_e1.nii.gz phase_e2.nii.gz phase_e3.nii.gz \ --TEs 14.2 38.93 63.66 \ @@ -127,7 +143,52 @@ medic \ `--TEs` is in **milliseconds**, `--total-readout-time` in **seconds**, and `--phase-encoding-direction` is one of `i, j, k, i-, j-, k-, x, y, z, x-, y-, z-`. -Run `medic --help` for the full option list (noise-frame trimming, CPU count, debug mode, etc.). +Run any `wk-*` CLI with `--help` for the full option list. + +#### Common follow-on workflows + +Apply a MEDIC displacement-map series to the BOLD it was derived from +(per-frame distortion correction): + +```bash +wk-apply-warp \ + --input bold.nii.gz \ + --transform sub-01_run-01_displacementmaps.nii \ + --phase-encoding-axis j \ + --output bold_corrected.nii.gz +``` + +Convert MEDIC's per-frame displacement maps into per-frame ANTs-format +displacement fields: + +```bash +wk-convert-warp \ + --input sub-01_run-01_displacementmaps.nii \ + --to field --axis j --to-format ants \ + --output field_{0..14}.nii.gz +``` + +Compute the per-frame Jacobian determinant (volume-change map) of those +displacement maps: + +```bash +wk-compute-jacobian \ + --input sub-01_run-01_displacementmaps.nii \ + --axis j \ + --output sub-01_run-01_jacobian.nii +``` + +Convert MEDIC's mm displacement maps to a Hz B0 field map (or back — +this CLI handles either direction, with maps or fields on the mm side): + +```bash +wk-convert-fieldmap \ + --input sub-01_run-01_displacementmaps.nii \ + --to fieldmap \ + --total-readout-time 0.0501 \ + --phase-encoding-direction j- \ + --output sub-01_run-01_fieldmap.nii +``` ## Authors diff --git a/include/romeo/romeo.h b/include/romeo/romeo.h index 21a443f..5f72a4a 100644 --- a/include/romeo/romeo.h +++ b/include/romeo/romeo.h @@ -20,67 +20,28 @@ namespace py = pybind11; namespace romeo { -// Python-facing facade over the pure-C++ ROMEO implementation. +// Python-facing entry points over the pure-C++ ROMEO implementation. // -// Method names retain the `romeo_*` prefix from the original Julia-backed -// pybind class so call sites in warpkit/unwrap.py read naturally. -template -class Romeo { - public: - Romeo() = default; - ~Romeo() = default; - - // Port of ROMEO.jl `calculateweights(phase; weights=:romeo, ...)` — the - // only weight preset we support. Exposed primarily so Python tests can - // validate the internal machinery against literal goldens from - // ROMEO.jl test/specialcases.jl. Not used by warpkit itself. - // - // `phase` is a column-major (nx, ny, nz) array. `mag`, `phase2`, `mask` may - // be 0-sized / empty arrays to indicate "not provided"; `TEs` is required - // only when `phase2` is provided (length 2: [te_phase, te_phase2]). - // Returns a (3, nx, ny, nz) uint8 array. - py::array_t calculate_weights(py::array_t phase, - py::array_t mag, - py::array_t phase2, - py::array_t TEs, - py::array_t mask); - - py::array_t romeo_voxelquality(py::array_t phase, - py::array_t TEs, - py::array_t mag); - - py::array_t romeo_unwrap3D(py::array_t phase, - std::string weights, - py::array_t mag, - py::array_t mask, - bool correctglobal = false, - int maxseeds = 1, - bool merge_regions = false, - bool correct_regions = false); - - py::array_t romeo_unwrap4D(py::array_t phase, - py::array_t TEs, - std::string weights, - py::array_t mag, - py::array_t mask, - bool correctglobal = false, - int maxseeds = 1, - bool merge_regions = false, - bool correct_regions = false); -}; - -// ---------------------------------------------------------------------------- -// Method implementations (Phase 1: all throw; wired to the public entry points -// under include/romeo/*.h so Phase 2+ can fill them in without touching this -// file's signatures). -// ---------------------------------------------------------------------------- - +// Function names retain the `romeo_*` prefix from the original Julia-backed +// API so call sites in warpkit/unwrap.py read naturally. They are stateless +// free functions; pybind11 binds them as module-level functions in +// `warpkit_cpp` (no class wrapper). + +// Port of ROMEO.jl `calculateweights(phase; weights=:romeo, ...)` — the +// only weight preset we support. Exposed primarily so Python tests can +// validate the internal machinery against literal goldens from +// ROMEO.jl test/specialcases.jl. Not used by warpkit itself. +// +// `phase` is a column-major (nx, ny, nz) array. `mag`, `phase2`, `mask` may +// be 0-sized / empty arrays to indicate "not provided"; `TEs` is required +// only when `phase2` is provided (length 2: [te_phase, te_phase2]). +// Returns a (3, nx, ny, nz) uint8 array. template -py::array_t Romeo::calculate_weights(py::array_t phase, - py::array_t mag, - py::array_t phase2, - py::array_t TEs, - py::array_t mask) { +py::array_t calculate_weights(py::array_t phase, + py::array_t mag, + py::array_t phase2, + py::array_t TEs, + py::array_t mask) { if (phase.ndim() != 3) throw std::invalid_argument("calculate_weights: phase must be 3D"); const auto nx = static_cast(phase.shape(0)); const auto ny = static_cast(phase.shape(1)); @@ -118,9 +79,9 @@ py::array_t Romeo::calculate_weights(py::ar } template -py::array_t Romeo::romeo_voxelquality(py::array_t phase, - py::array_t TEs, - py::array_t mag) { +py::array_t romeo_voxelquality(py::array_t phase, + py::array_t TEs, + py::array_t mag) { if (phase.ndim() != 4) throw std::invalid_argument("romeo_voxelquality: phase must be 4D"); const auto nx = static_cast(phase.shape(0)); const auto ny = static_cast(phase.shape(1)); @@ -150,14 +111,14 @@ py::array_t Romeo::romeo_voxelquality(py::array_t -py::array_t Romeo::romeo_unwrap3D(py::array_t phase, - std::string weights, - py::array_t mag, - py::array_t mask, - bool correctglobal, - int maxseeds, - bool merge_regions, - bool correct_regions) { +py::array_t romeo_unwrap3D(py::array_t phase, + std::string weights, + py::array_t mag, + py::array_t mask, + bool correctglobal = false, + int maxseeds = 1, + bool merge_regions = false, + bool correct_regions = false) { if (weights != "romeo") throw std::invalid_argument("romeo_unwrap3D: only the \"romeo\" weight preset is supported."); if (merge_regions || correct_regions) @@ -198,15 +159,15 @@ py::array_t Romeo::romeo_unwrap3D(py::array_t -py::array_t Romeo::romeo_unwrap4D(py::array_t phase, - py::array_t TEs, - std::string weights, - py::array_t mag, - py::array_t mask, - bool correctglobal, - int maxseeds, - bool merge_regions, - bool correct_regions) { +py::array_t romeo_unwrap4D(py::array_t phase, + py::array_t TEs, + std::string weights, + py::array_t mag, + py::array_t mask, + bool correctglobal = false, + int maxseeds = 1, + bool merge_regions = false, + bool correct_regions = false) { if (weights != "romeo") throw std::invalid_argument("romeo_unwrap4D: only the \"romeo\" weight preset is supported."); if (merge_regions || correct_regions) diff --git a/pyproject.toml b/pyproject.toml index e787137..f677180 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ authors = [{ name = "Andrew Van", email = "vanandrew77@gmail.com" }] keywords = ["neuroimaging"] classifiers = ["Programming Language :: Python :: 3"] urls = { github = "https://github.com/vanandrew/warpkit" } -dynamic = ["version", "scripts"] +dynamic = ["version"] dependencies = [ "nibabel >= 4.0.2", "numpy >= 1.23.3", @@ -17,6 +17,15 @@ dependencies = [ "indexed-gzip >= 1.7.0", ] +[project.scripts] +"wk-medic" = "warpkit.scripts.medic:main" +"wk-unwrap-phase" = "warpkit.scripts.unwrap_phase:main" +"wk-compute-fieldmap" = "warpkit.scripts.compute_fieldmap:main" +"wk-apply-warp" = "warpkit.scripts.apply_warp:main" +"wk-convert-warp" = "warpkit.scripts.convert_warp:main" +"wk-convert-fieldmap" = "warpkit.scripts.convert_fieldmap:main" +"wk-compute-jacobian" = "warpkit.scripts.compute_jacobian:main" + [dependency-groups] dev = [ "coverage[toml] >= 5.5", diff --git a/setup.py b/setup.py index 3eb13d6..43f76d2 100644 --- a/setup.py +++ b/setup.py @@ -8,7 +8,6 @@ THISDIR = Path(__file__).parent cmake_dir = (Path(THISDIR)).absolute().as_posix() -scripts_path = THISDIR / "warpkit" / "scripts" IS_CIBUILDWHEEL = os.environ.get("CIBUILDWHEEL", "0") == "1" IS_MACOS = os.environ.get("RUNNER_OS", "0") == "macOS" @@ -32,11 +31,4 @@ ) ], cmdclass={"build_ext": BuildExtension}, - entry_points={ - "console_scripts": [ - f"{f.stem}=warpkit.scripts.{f.stem}:main" - for f in scripts_path.glob("*.py") - if f.name not in "__init__.py" - ] - }, ) diff --git a/src/warpkit.cpp b/src/warpkit.cpp index 6a8e900..fc74d85 100644 --- a/src/warpkit.cpp +++ b/src/warpkit.cpp @@ -6,28 +6,28 @@ namespace py = pybind11; PYBIND11_MODULE(warpkit_cpp, m) { - using Romeo32 = romeo::Romeo; - - py::class_(m, "Romeo") - .def(py::init<>()) - .def("calculate_weights", &Romeo32::calculate_weights, - "ROMEO edge-weight map (3, nx, ny, nz) uint8. Exposed for port validation; not used by warpkit.", - py::arg("phase"), - py::arg("mag") = py::array_t(), - py::arg("phase2") = py::array_t(), - py::arg("tes") = py::array_t(), - py::arg("mask") = py::array_t(), - py::return_value_policy::move) - .def("romeo_voxelquality", &Romeo32::romeo_voxelquality, - "Compute a per-voxel quality map from multi-echo phase/magnitude", py::arg("phase"), py::arg("tes"), - py::arg("mag"), py::return_value_policy::move) - .def("romeo_unwrap3d", &Romeo32::romeo_unwrap3D, "3D ROMEO phase unwrap", py::arg("phase"), py::arg("weights"), - py::arg("mag"), py::arg("mask"), py::arg("correct_global") = true, py::arg("maxseeds") = 1, - py::arg("merge_regions") = false, py::arg("correct_regions") = false, py::return_value_policy::move) - .def("romeo_unwrap4d", &Romeo32::romeo_unwrap4D, "4D (multi-echo) ROMEO phase unwrap", py::arg("phase"), - py::arg("tes"), py::arg("weights"), py::arg("mag"), py::arg("mask"), py::arg("correct_global") = true, - py::arg("maxseeds") = 1, py::arg("merge_regions") = false, py::arg("correct_regions") = false, - py::return_value_policy::move); + m.def("calculate_weights", &romeo::calculate_weights, + "ROMEO edge-weight map (3, nx, ny, nz) uint8. Exposed for port validation; not used by warpkit.", + py::arg("phase"), + py::arg("mag") = py::array_t(), + py::arg("phase2") = py::array_t(), + py::arg("tes") = py::array_t(), + py::arg("mask") = py::array_t(), + py::return_value_policy::move); + + m.def("romeo_voxelquality", &romeo::romeo_voxelquality, + "Compute a per-voxel quality map from multi-echo phase/magnitude", py::arg("phase"), py::arg("tes"), + py::arg("mag"), py::return_value_policy::move); + + m.def("romeo_unwrap3d", &romeo::romeo_unwrap3D, "3D ROMEO phase unwrap", py::arg("phase"), + py::arg("weights"), py::arg("mag"), py::arg("mask"), py::arg("correct_global") = true, + py::arg("maxseeds") = 1, py::arg("merge_regions") = false, py::arg("correct_regions") = false, + py::return_value_policy::move); + + m.def("romeo_unwrap4d", &romeo::romeo_unwrap4D, "4D (multi-echo) ROMEO phase unwrap", py::arg("phase"), + py::arg("tes"), py::arg("weights"), py::arg("mag"), py::arg("mask"), py::arg("correct_global") = true, + py::arg("maxseeds") = 1, py::arg("merge_regions") = false, py::arg("correct_regions") = false, + py::return_value_policy::move); m.def("invert_displacement_map", &invert_displacement_map, "Invert a displacement map", py::arg("displacement_map"), py::arg("origin"), py::arg("direction"), py::arg("spacing"), py::arg("axis") = 1, diff --git a/tests/conftest.py b/tests/conftest.py index 41e1186..b3bf045 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,15 +8,16 @@ # get this directory THISDIR = Path(__file__).parent +TEST_DATA_DIR = THISDIR / "data" / "test_data" # fixture for test data @fixture(scope="session") def test_data(): # get mag and phase data - mag = sorted(Path(THISDIR, "data", "test_data").glob("*mag*.nii.gz")) - phase = sorted(Path(THISDIR, "data", "test_data").glob("*phase*.nii.gz")) - sidecar = sorted(Path(THISDIR, "data", "test_data").glob("*mag*.json")) + mag = sorted(TEST_DATA_DIR.glob("*mag*.nii.gz")) + phase = sorted(TEST_DATA_DIR.glob("*phase*.nii.gz")) + sidecar = sorted(TEST_DATA_DIR.glob("*mag*.json")) metadata = [] for s in sidecar: with s.open() as f: @@ -28,3 +29,14 @@ def test_data(): "total_readout_time": metadata[0]["TotalReadoutTime"], "phase_encoding_direction": metadata[0]["PhaseEncodingDirection"], } + + +@fixture(scope="session") +def test_data_paths(): + """File paths for the bundled BIDS-style MEDIC test data, suitable for + passing directly to a CLI.""" + return { + "mag": [str(p) for p in sorted(TEST_DATA_DIR.glob("*mag*.nii.gz"))], + "phase": [str(p) for p in sorted(TEST_DATA_DIR.glob("*phase*.nii.gz"))], + "metadata": [str(p) for p in sorted(TEST_DATA_DIR.glob("*mag*.json"))], + } diff --git a/tests/test_romeo.py b/tests/test_romeo.py index aaab499..f4939ad 100644 --- a/tests/test_romeo.py +++ b/tests/test_romeo.py @@ -49,14 +49,6 @@ def mag4d() -> np.ndarray: return _load_nii(ROMEO_TEST_DATA / "Mag.nii") -@pytest.fixture(scope="module") -def romeo(): - """Handle to the ROMEO C++ context.""" - from warpkit.warpkit_cpp import Romeo - - return Romeo() - - # --------------------------------------------------------------------------- # Literal goldens — ported from ROMEO.jl test/dsp_tests.jl # --------------------------------------------------------------------------- @@ -79,14 +71,16 @@ def romeo(): "wrap-up-idx0", ], ) -def test_unwrap_1d_literals(romeo, wrapped): +def test_unwrap_1d_literals(wrapped): """ROMEO.jl's unwrap!() reshapes <=2D input to 3D before running. Reproduce by feeding a (N, 1, 1) volume through the 3D entry point.""" + from warpkit.warpkit_cpp import romeo_unwrap3d + expected = np.array([0.1, 0.2, 0.3, 0.4], dtype=np.float32) phase = np.asarray(wrapped, dtype=np.float32).reshape(-1, 1, 1) mag = np.ones_like(phase) mask = np.ones(phase.shape, dtype=bool) - unwrapped = romeo.romeo_unwrap3d(phase, "romeo", mag, mask).reshape(-1) + unwrapped = romeo_unwrap3d(phase, "romeo", mag, mask).reshape(-1) np.testing.assert_allclose(unwrapped, expected, atol=1e-5) @@ -115,11 +109,13 @@ def test_unwrap_1d_literals(romeo, wrapped): ], ids=["linearity-border", "nan-neighbor"], ) -def test_weight_calc_literals(romeo, phase, expected): +def test_weight_calc_literals(phase, expected): """Port of the weight_test() assertions from specialcases.jl.""" + from warpkit.warpkit_cpp import calculate_weights + phase_arr = np.asarray(phase, dtype=np.float32).reshape(-1, 1, 1) # (3, nx, ny, nz) uint8 — we check dim-0 edges along the length-4 x-axis. - weights = romeo.calculate_weights(phase_arr) + weights = calculate_weights(phase_arr) np.testing.assert_array_equal( weights[0, :, 0, 0], np.asarray(expected, dtype=np.uint8) ) @@ -135,14 +131,16 @@ def _rem2pi_nearest(x: np.ndarray) -> np.ndarray: return x - 2 * np.pi * np.round(x / (2 * np.pi)) -def test_unwrap3d_property(romeo, phase4d, mag4d): +def test_unwrap3d_property(phase4d, mag4d): """ROMEO's 3D unwrap must differ from wrapped input only by multiples of 2π.""" + from warpkit.warpkit_cpp import romeo_unwrap3d + echo = 2 # index for the 3rd echo, matching Julia's `echo = 3` (1-based) wrapped = np.ascontiguousarray(phase4d[..., echo]) mag = np.ascontiguousarray(mag4d[..., echo]) mask = np.ones(wrapped.shape, dtype=bool) - unwrapped = romeo.romeo_unwrap3d(wrapped, "romeo", mag, mask) + unwrapped = romeo_unwrap3d(wrapped, "romeo", mag, mask) assert unwrapped.shape == wrapped.shape assert not np.array_equal(unwrapped, wrapped), "unwrap returned the input unchanged" @@ -151,12 +149,14 @@ def test_unwrap3d_property(romeo, phase4d, mag4d): np.testing.assert_allclose(residual, 0.0, atol=1e-5) -def test_unwrap4d_property(romeo, phase4d, mag4d): +def test_unwrap4d_property(phase4d, mag4d): """Same 2π-modulo invariant across every echo of the 4D multi-echo unwrap.""" + from warpkit.warpkit_cpp import romeo_unwrap4d + tes = np.array([4.0, 8.0, 12.0], dtype=np.float32) # matches voxelquality.jl mask = np.ones(phase4d.shape[:3], dtype=bool) - unwrapped = romeo.romeo_unwrap4d(phase4d, tes, "romeo", mag4d, mask) + unwrapped = romeo_unwrap4d(phase4d, tes, "romeo", mag4d, mask) assert unwrapped.shape == phase4d.shape assert np.isfinite(unwrapped).all() @@ -172,17 +172,19 @@ def test_unwrap4d_property(romeo, phase4d, mag4d): # --------------------------------------------------------------------------- -def test_voxelquality_behavior(romeo, phase4d, mag4d): +def test_voxelquality_behavior(phase4d, mag4d): """ Mirror the three-variant comparison from voxelquality.jl. ROMEO's voxelquality entry point requires tes and mag, so we reuse the 4D volume and vary the echo time ordering to produce distinct qmaps. """ + from warpkit.warpkit_cpp import romeo_voxelquality + tes = np.array([4.0, 8.0, 12.0], dtype=np.float32) - qm_uniform_mag = romeo.romeo_voxelquality(phase4d, tes, np.ones_like(mag4d)) - qm_real_mag = romeo.romeo_voxelquality(phase4d, tes, mag4d) - qm_reordered = romeo.romeo_voxelquality(phase4d, tes[::-1].copy(), mag4d) + qm_uniform_mag = romeo_voxelquality(phase4d, tes, np.ones_like(mag4d)) + qm_real_mag = romeo_voxelquality(phase4d, tes, mag4d) + qm_reordered = romeo_voxelquality(phase4d, tes[::-1].copy(), mag4d) for qmap, label in [ (qm_uniform_mag, "uniform-mag"), diff --git a/tests/test_scripts.py b/tests/test_scripts.py index 14def7a..0160f1d 100644 --- a/tests/test_scripts.py +++ b/tests/test_scripts.py @@ -1,19 +1,36 @@ """CLI argument-validation tests. These exercise argparse + the post-parse mutex/either-or logic in -``medic.main`` and the help text of ``extract_field_from_maps.main``. They -do not touch the heavy unwrap/inversion pipeline — every test here triggers -``parser.error`` (which raises ``SystemExit(2)``) before the data load. +``medic.main``, ``unwrap_phase.main``, ``compute_fieldmap.main``, +``apply_warp.main`` and ``convert_warp.main``. They do not touch the heavy +unwrap/inversion/resample pipeline — every test here triggers +``parser.error`` (which raises ``SystemExit(2)``) before the data load (or +before the resample, in apply_warp's case where reaching parser.error +requires loading the inputs first). """ from __future__ import annotations import json import sys +from typing import cast +import nibabel as nib +import numpy as np import pytest -from warpkit.scripts.extract_field_from_maps import main as extract_main +from warpkit.scripts.apply_warp import main as apply_warp_main +from warpkit.scripts.compute_fieldmap import main as compute_fieldmap_main +from warpkit.scripts.compute_jacobian import main as compute_jacobian_main +from warpkit.scripts.convert_fieldmap import main as convert_fieldmap_main +from warpkit.scripts.convert_warp import main as convert_warp_main from warpkit.scripts.medic import main as medic_main +from warpkit.scripts.unwrap_phase import main as unwrap_phase_main + + +def _load(path) -> nib.Nifti1Image: + """Type-narrowed nib.load for tests; the bundled data and our synthetic + NIfTIs are concretely Nifti1Image.""" + return cast(nib.Nifti1Image, nib.load(str(path))) @pytest.fixture @@ -24,13 +41,21 @@ def _set(args): return _set +def _write_nifti(path, shape, affine=None, dtype=np.float32): + """Write a zero-filled NIfTI of the requested shape; return the path.""" + if affine is None: + affine = np.eye(4) + nib.Nifti1Image(np.zeros(shape, dtype=dtype), affine).to_filename(str(path)) + return str(path) + + # --------------------------------------------------------------------------- # medic --help / --version # --------------------------------------------------------------------------- def test_medic_help(argv, capsys): - argv(["medic", "--help"]) + argv(["wk-medic", "--help"]) with pytest.raises(SystemExit) as exc: medic_main() assert exc.value.code == 0 @@ -43,12 +68,12 @@ def test_medic_help(argv, capsys): def test_medic_version(argv, capsys): - argv(["medic", "--version"]) + argv(["wk-medic", "--version"]) with pytest.raises(SystemExit) as exc: medic_main() assert exc.value.code == 0 out = capsys.readouterr().out - assert "medic" in out + assert "wk-medic" in out # --------------------------------------------------------------------------- @@ -59,7 +84,7 @@ def test_medic_version(argv, capsys): def test_medic_requires_acquisition_args(argv, capsys): argv( [ - "medic", + "wk-medic", "--magnitude", "m.nii", "--phase", @@ -81,7 +106,7 @@ def test_medic_requires_acquisition_args(argv, capsys): def test_medic_metadata_and_direct_args_conflict(argv, capsys): argv( [ - "medic", + "wk-medic", "--magnitude", "m.nii", "--phase", @@ -105,7 +130,7 @@ def test_medic_metadata_and_direct_args_conflict(argv, capsys): def test_medic_te_count_must_match_phase_count(argv, capsys, tmp_path): argv( [ - "medic", + "wk-medic", "--magnitude", "m1.nii", "m2.nii", @@ -143,7 +168,7 @@ def test_medic_metadata_te_count_must_match_phase_count(argv, capsys, tmp_path): ) argv( [ - "medic", + "wk-medic", "--magnitude", "m.nii", "--phase", @@ -165,7 +190,7 @@ def test_medic_metadata_te_count_must_match_phase_count(argv, capsys, tmp_path): def test_medic_phase_encoding_direction_choices_enforced(argv, capsys): argv( [ - "medic", + "wk-medic", "--magnitude", "m.nii", "--phase", @@ -188,34 +213,2601 @@ def test_medic_phase_encoding_direction_choices_enforced(argv, capsys): # --------------------------------------------------------------------------- -# extract_field_from_maps --help / choices +# wk-convert-warp --help / argument validation # --------------------------------------------------------------------------- -def test_extract_field_from_maps_help(argv, capsys): - argv(["extract_field_from_maps", "--help"]) +def test_convert_warp_help(argv, capsys): + argv(["wk-convert-warp", "--help"]) with pytest.raises(SystemExit) as exc: - extract_main() + convert_warp_main() assert exc.value.code == 0 out = capsys.readouterr().out - assert "extracts a displacement field" in out - # dash-form flags - assert "--frame-number" in out - assert "--phase-encoding-axis" in out + assert "Interconvert" in out + assert "--from" in out + assert "--to" in out + assert "--from-format" in out + assert "--to-format" in out + assert "--axis" in out + assert "--frame" in out + assert "--invert" in out -def test_extract_field_from_maps_rejects_bad_axis(argv, capsys): +def test_convert_warp_rejects_bad_format(argv, capsys, tmp_path): + field = _write_nifti(tmp_path / "field.nii", (4, 4, 4, 3)) argv( [ - "extract_field_from_maps", - "maps.nii.gz", - "out.nii.gz", - "--phase-encoding-axis", + "wk-convert-warp", + "--input", + field, + "--output", + str(tmp_path / "out.nii"), + "--to-format", + "matlab", # not a valid choice + ] + ) + with pytest.raises(SystemExit) as exc: + convert_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "invalid choice" in err + + +def test_convert_warp_rejects_bad_axis(argv, capsys, tmp_path): + maps = _write_nifti(tmp_path / "maps.nii", (4, 4, 4, 5)) + argv( + [ + "wk-convert-warp", + "--input", + maps, + "--output", + str(tmp_path / "out.nii"), + "--to", + "field", + "--axis", "diagonal", ] ) with pytest.raises(SystemExit) as exc: - extract_main() + convert_warp_main() assert exc.value.code == 2 err = capsys.readouterr().err assert "invalid choice" in err + + +def test_convert_warp_requires_axis_for_map_to_field(argv, capsys, tmp_path): + """Map-to-field conversion requires --axis.""" + maps = _write_nifti(tmp_path / "maps.nii", (4, 4, 4, 5)) + argv( + [ + "wk-convert-warp", + "--input", + maps, + "--output", + str(tmp_path / "out.nii"), + "--from", + "map", + "--to", + "field", + ] + ) + with pytest.raises(SystemExit) as exc: + convert_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "--axis is required" in err + + +def test_convert_warp_requires_axis_for_field_to_map(argv, capsys, tmp_path): + """Field-to-map conversion requires --axis.""" + field = _write_nifti(tmp_path / "field.nii", (4, 4, 4, 3)) + argv( + [ + "wk-convert-warp", + "--input", + field, + "--output", + str(tmp_path / "out.nii"), + "--from", + "field", + "--to", + "map", + ] + ) + with pytest.raises(SystemExit) as exc: + convert_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "--axis is required" in err + + +def test_convert_warp_frame_out_of_range(argv, capsys, tmp_path): + maps = _write_nifti(tmp_path / "maps.nii", (4, 4, 4, 5)) + argv( + [ + "wk-convert-warp", + "--input", + maps, + "--output", + str(tmp_path / "out.nii"), + "--from", + "map", + "--frame", + "10", + ] + ) + with pytest.raises(SystemExit) as exc: + convert_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "out of range" in err + + +def test_convert_warp_output_count_mismatch(argv, capsys, tmp_path): + """Output paths must equal frame count or be a single bundle path.""" + maps = _write_nifti(tmp_path / "maps.nii", (4, 4, 4, 5)) # 5 frames + argv( + [ + "wk-convert-warp", + "--input", + maps, + "--output", + str(tmp_path / "o1.nii"), + str(tmp_path / "o2.nii"), # only 2 outputs for 5 frames + "--from", + "map", + ] + ) + with pytest.raises(SystemExit) as exc: + convert_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "must be 1" in err and "one per frame" in err + + +def test_convert_warp_invert_single_map_requires_axis(argv, capsys, tmp_path): + """Inverting a single-frame map always requires --axis (the map's own axis), + even when the output is also a map (no map<->field conversion).""" + single_map = _write_nifti(tmp_path / "map.nii", (4, 4, 4)) # 3D = single map + argv( + [ + "wk-convert-warp", + "--input", + single_map, + "--output", + str(tmp_path / "out.nii"), + "--from", + "map", + "--invert", + # no --axis, no --to=field + ] + ) + with pytest.raises(SystemExit) as exc: + convert_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "--axis is required" in err + + +def test_convert_warp_invert_multi_frame_field_requires_axis(argv, capsys, tmp_path): + """Multi-frame inversion routes through the 1D map inverter, which needs + --axis even for a field input (to know which channel to invert along).""" + # 4 separate 4D fields = a 4-frame field series + fields = [_write_nifti(tmp_path / f"f{i}.nii", (4, 4, 4, 3)) for i in range(4)] + argv( + [ + "wk-convert-warp", + "--input", + *fields, + "--output", + str(tmp_path / "out.nii"), + "--from", + "field", + "--invert", + # no --axis + ] + ) + with pytest.raises(SystemExit) as exc: + convert_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "--axis is required" in err + assert "multi-frame" in err + + +# --------------------------------------------------------------------------- +# unwrap_phase --help / argument validation +# --------------------------------------------------------------------------- + + +def test_unwrap_phase_help(argv, capsys): + argv(["wk-unwrap-phase", "--help"]) + with pytest.raises(SystemExit) as exc: + unwrap_phase_main() + assert exc.value.code == 0 + out = capsys.readouterr().out + assert "phase unwrapping" in out + assert "--out-prefix" in out + # readout/PE flags are MEDIC-only; unwrap_phase should not advertise them + assert "--total-readout-time" not in out + assert "--phase-encoding-direction" not in out + + +def test_unwrap_phase_requires_tes(argv, capsys): + argv( + [ + "wk-unwrap-phase", + "--magnitude", + "m.nii", + "--phase", + "p.nii", + "--out-prefix", + "out", + ] + ) + with pytest.raises(SystemExit) as exc: + unwrap_phase_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "either --metadata or --TEs" in err + + +def test_unwrap_phase_metadata_and_tes_conflict(argv, capsys): + argv( + [ + "wk-unwrap-phase", + "--magnitude", + "m.nii", + "--phase", + "p.nii", + "--out-prefix", + "out", + "--metadata", + "m.json", + "--TEs", + "14.0", + ] + ) + with pytest.raises(SystemExit) as exc: + unwrap_phase_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "mutually exclusive" in err + + +def test_unwrap_phase_te_count_must_match_phase_count(argv, capsys, tmp_path): + argv( + [ + "wk-unwrap-phase", + "--magnitude", + "m1.nii", + "m2.nii", + "--phase", + "p1.nii", + "p2.nii", + "--TEs", + "14.0", # only 1 TE for 2 phase files + "--out-prefix", + str(tmp_path / "out"), + ] + ) + with pytest.raises(SystemExit) as exc: + unwrap_phase_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "must match" in err + + +def test_unwrap_phase_magnitude_count_must_match_phase_count(argv, capsys, tmp_path): + """Mismatched --magnitude / --phase counts should raise a clean parser + error, not a downstream zip(strict=True) ValueError.""" + argv( + [ + "wk-unwrap-phase", + "--magnitude", + "m1.nii", # only 1 mag for 2 phase files + "--phase", + "p1.nii", + "p2.nii", + "--TEs", + "14.0", + "28.0", + "--out-prefix", + str(tmp_path / "out"), + ] + ) + with pytest.raises(SystemExit) as exc: + unwrap_phase_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "--magnitude" in err and "--phase" in err + + +def test_unwrap_phase_metadata_count_must_match_phase_count(argv, capsys, tmp_path): + argv( + [ + "wk-unwrap-phase", + "--magnitude", + "m1.nii", + "m2.nii", + "--phase", + "p1.nii", + "p2.nii", + "--metadata", + "e1.json", # only 1 sidecar for 2 echoes + "--out-prefix", + str(tmp_path / "out"), + ] + ) + with pytest.raises(SystemExit) as exc: + unwrap_phase_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "--metadata" in err and "--phase" in err + + +def test_unwrap_phase_noiseframes_negative(argv, capsys, tmp_path): + """A negative --noiseframes must fail with a clean parser error.""" + mag = _write_nifti(tmp_path / "m.nii", (4, 4, 4, 5)) + phase = _write_nifti(tmp_path / "p.nii", (4, 4, 4, 5)) + argv( + [ + "wk-unwrap-phase", + "--magnitude", + mag, + "--phase", + phase, + "--TEs", + "14.0", + "--out-prefix", + str(tmp_path / "out"), + "-f", + "-1", + ] + ) + with pytest.raises(SystemExit) as exc: + unwrap_phase_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "non-negative" in err + + +def test_unwrap_phase_noiseframes_consumes_all_frames(argv, capsys, tmp_path): + """--noiseframes >= n_frames would leave 0 volumes; reject with a clean + parser.error instead of producing an empty 4D series that crashes + unwrap_phases at frames[0].""" + mag = _write_nifti(tmp_path / "m.nii", (4, 4, 4, 5)) + phase = _write_nifti(tmp_path / "p.nii", (4, 4, 4, 5)) + argv( + [ + "wk-unwrap-phase", + "--magnitude", + mag, + "--phase", + phase, + "--TEs", + "14.0", + "--out-prefix", + str(tmp_path / "out"), + "-f", + "5", + ] + ) + with pytest.raises(SystemExit) as exc: + unwrap_phase_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "0 frames" in err + + +# --------------------------------------------------------------------------- +# compute_fieldmap --help / argument validation +# --------------------------------------------------------------------------- + + +def test_compute_fieldmap_help(argv, capsys): + argv(["wk-compute-fieldmap", "--help"]) + with pytest.raises(SystemExit) as exc: + compute_fieldmap_main() + assert exc.value.code == 0 + out = capsys.readouterr().out + assert "field map" in out + assert "--unwrapped" in out + assert "--masks" in out + # post-unwrap distortion correction needs the readout/PE info + assert "--total-readout-time" in out + assert "--phase-encoding-direction" in out + + +def test_compute_fieldmap_requires_acquisition_args(argv, capsys): + argv( + [ + "wk-compute-fieldmap", + "--magnitude", + "m.nii", + "--unwrapped", + "u.nii", + "--masks", + "masks.nii", + "--out-prefix", + "out", + ] + ) + with pytest.raises(SystemExit) as exc: + compute_fieldmap_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "either --metadata or all of" in err + assert "--TEs" in err + assert "--total-readout-time" in err + assert "--phase-encoding-direction" in err + + +def test_compute_fieldmap_metadata_and_direct_args_conflict(argv, capsys): + argv( + [ + "wk-compute-fieldmap", + "--magnitude", + "m.nii", + "--unwrapped", + "u.nii", + "--masks", + "masks.nii", + "--out-prefix", + "out", + "--metadata", + "m.json", + "--TEs", + "14.0", + ] + ) + with pytest.raises(SystemExit) as exc: + compute_fieldmap_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "mutually exclusive" in err + assert "--TEs" in err + + +def test_compute_fieldmap_input_count_must_match(argv, capsys, tmp_path): + """Mag/unwrapped/TE counts must all line up echo-for-echo.""" + argv( + [ + "wk-compute-fieldmap", + "--magnitude", + "m1.nii", + "m2.nii", + "--unwrapped", + "u1.nii", # only 1 unwrapped for 2 mags + "--masks", + "masks.nii", + "--TEs", + "14.0", + "38.0", + "--total-readout-time", + "0.05", + "--phase-encoding-direction", + "j", + "--out-prefix", + str(tmp_path / "out"), + ] + ) + with pytest.raises(SystemExit) as exc: + compute_fieldmap_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "must match" in err + + +def test_compute_fieldmap_phase_encoding_direction_choices_enforced(argv, capsys): + argv( + [ + "wk-compute-fieldmap", + "--magnitude", + "m.nii", + "--unwrapped", + "u.nii", + "--masks", + "masks.nii", + "--TEs", + "14.0", + "--total-readout-time", + "0.05", + "--phase-encoding-direction", + "diagonal", # not a valid choice + "--out-prefix", + "out", + ] + ) + with pytest.raises(SystemExit) as exc: + compute_fieldmap_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "invalid choice" in err + + +# --------------------------------------------------------------------------- +# apply_warp --help / argument validation +# +# The validation tests that need to reach `parser.error` after the data load +# (e.g. transform/input frame-count mismatch) write tiny synthetic NIfTIs to +# tmp_path so we don't depend on the heavy MEDIC fixtures. +# --------------------------------------------------------------------------- + + +def test_apply_warp_help(argv, capsys): + argv(["wk-apply-warp", "--help"]) + with pytest.raises(SystemExit) as exc: + apply_warp_main() + assert exc.value.code == 0 + out = capsys.readouterr().out + assert "Resample" in out + assert "--input" in out + assert "--transform" in out + assert "--phase-encoding-axis" in out + + +def test_apply_warp_requires_phase_encoding_axis_for_map(argv, capsys, tmp_path): + """Map-type transforms require --phase-encoding-axis.""" + inp = _write_nifti(tmp_path / "in.nii", (4, 4, 4)) + # 4D 1-channel map series + tx = _write_nifti(tmp_path / "tx.nii", (4, 4, 4, 5)) + argv( + [ + "wk-apply-warp", + "--input", + inp, + "--transform", + tx, + "--transform-type", + "map", + "--output", + str(tmp_path / "out.nii"), + ] + ) + with pytest.raises(SystemExit) as exc: + apply_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "--phase-encoding-axis" in err + + +def test_apply_warp_rejects_bad_axis(argv, capsys, tmp_path): + inp = _write_nifti(tmp_path / "in.nii", (4, 4, 4)) + tx = _write_nifti(tmp_path / "tx.nii", (4, 4, 4, 5)) + argv( + [ + "wk-apply-warp", + "--input", + inp, + "--transform", + tx, + "--phase-encoding-axis", + "diagonal", + "--output", + str(tmp_path / "out.nii"), + ] + ) + with pytest.raises(SystemExit) as exc: + apply_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "invalid choice" in err + + +def test_apply_warp_rejects_bad_format(argv, capsys, tmp_path): + inp = _write_nifti(tmp_path / "in.nii", (4, 4, 4)) + tx = _write_nifti(tmp_path / "tx.nii", (4, 4, 4, 3)) # 4D field + argv( + [ + "wk-apply-warp", + "--input", + inp, + "--transform", + tx, + "--format", + "matlab", # not a valid choice + "--output", + str(tmp_path / "out.nii"), + ] + ) + with pytest.raises(SystemExit) as exc: + apply_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "invalid choice" in err + + +def test_apply_warp_3d_input_with_series_transform_errors(argv, capsys, tmp_path): + """3D input + N-frame map transform is ambiguous — should be rejected.""" + inp = _write_nifti(tmp_path / "in.nii", (4, 4, 4)) # 3D input + tx = _write_nifti(tmp_path / "tx.nii", (4, 4, 4, 5)) # 5-frame map series + argv( + [ + "wk-apply-warp", + "--input", + inp, + "--transform", + tx, + "--transform-type", + "map", + "--phase-encoding-axis", + "j", + "--output", + str(tmp_path / "out.nii"), + ] + ) + with pytest.raises(SystemExit) as exc: + apply_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "input is 3D" in err + + +def test_apply_warp_frame_count_mismatch_errors(argv, capsys, tmp_path): + """4D input frames must match transform frames when transform is a series.""" + inp = _write_nifti(tmp_path / "in.nii", (4, 4, 4, 7)) # 7 input frames + tx = _write_nifti(tmp_path / "tx.nii", (4, 4, 4, 5)) # 5 transform frames + argv( + [ + "wk-apply-warp", + "--input", + inp, + "--transform", + tx, + "--transform-type", + "map", + "--phase-encoding-axis", + "j", + "--output", + str(tmp_path / "out.nii"), + ] + ) + with pytest.raises(SystemExit) as exc: + apply_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "they must match" in err + + +def test_apply_warp_series_files_must_be_3channel(argv, capsys, tmp_path): + """Multiple --transform files are treated as a field series; each must be 4D last==3.""" + inp = _write_nifti(tmp_path / "in.nii", (4, 4, 4, 2)) + # 4D last dim != 3 → not a valid field + bad1 = _write_nifti(tmp_path / "bad1.nii", (4, 4, 4, 5)) + bad2 = _write_nifti(tmp_path / "bad2.nii", (4, 4, 4, 5)) + argv( + [ + "wk-apply-warp", + "--input", + inp, + "--transform", + bad1, + bad2, + "--transform-type", + "field", + "--output", + str(tmp_path / "out.nii"), + ] + ) + with pytest.raises(SystemExit) as exc: + apply_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "4D" in err and "3-channel" in err + + +# --------------------------------------------------------------------------- +# CLI happy-path tests +# +# Unlike the validation tests above, these run each CLI to completion and +# assert on the output files. They use the bundled MEDIC fixture for the +# unwrap/medic/fieldmap pipeline (test_data_paths from conftest.py) and tiny +# synthetic NIfTIs for apply-warp / convert-warp (no need for real data +# there). These are the regression backstop for the script bodies that +# argparse-only tests can't reach. +# --------------------------------------------------------------------------- + + +def test_medic_happy_path(argv, test_data_paths, tmp_path): + """Run wk-medic end-to-end on the bundled fixture; assert all three + output NIfTIs exist with the expected shape and finite values.""" + out_prefix = tmp_path / "run" + argv( + [ + "wk-medic", + "--magnitude", + *test_data_paths["mag"], + "--phase", + *test_data_paths["phase"], + "--metadata", + *test_data_paths["metadata"], + "--out-prefix", + str(out_prefix), + "-n", + "1", + ] + ) + medic_main() + for suffix in ( + "_fieldmaps_native.nii", + "_displacementmaps.nii", + "_fieldmaps.nii", + ): + out = _load(f"{out_prefix}{suffix}") + data = out.get_fdata() + assert out.ndim == 4 + assert out.shape[:3] == (64, 64, 40) + assert out.shape[3] == 15 + assert np.isfinite(data).all() + + +def test_unwrap_phase_then_compute_fieldmap_happy_path(argv, test_data_paths, tmp_path): + """Run wk-unwrap-phase, then wk-compute-fieldmap on its outputs. Together + these cover the post-medic-split CLI flow.""" + unwrap_prefix = tmp_path / "unwrap" + argv( + [ + "wk-unwrap-phase", + "--magnitude", + *test_data_paths["mag"], + "--phase", + *test_data_paths["phase"], + "--metadata", + *test_data_paths["metadata"], + "--out-prefix", + str(unwrap_prefix), + "-n", + "1", + ] + ) + unwrap_phase_main() + + # one unwrapped phase per echo + one masks file + unwrapped = sorted(tmp_path.glob("unwrap_unwrapped_echo-*.nii")) + assert len(unwrapped) == 3 + masks_path = tmp_path / "unwrap_masks.nii" + assert masks_path.exists() + for u in unwrapped: + img = _load(str(u)) + assert img.shape == (64, 64, 40, 15) + assert np.isfinite(img.get_fdata()).all() + + fmap_prefix = tmp_path / "fmap" + argv( + [ + "wk-compute-fieldmap", + "--magnitude", + *test_data_paths["mag"], + "--unwrapped", + *[str(u) for u in unwrapped], + "--masks", + str(masks_path), + "--metadata", + *test_data_paths["metadata"], + "--out-prefix", + str(fmap_prefix), + "-n", + "1", + ] + ) + compute_fieldmap_main() + for suffix in ( + "_fieldmaps_native.nii", + "_displacementmaps.nii", + "_fieldmaps.nii", + ): + out = _load(f"{fmap_prefix}{suffix}") + assert out.shape == (64, 64, 40, 15) + assert np.isfinite(out.get_fdata()).all() + + +def test_apply_warp_happy_path_zero_displacement(argv, tmp_path): + """Identity resample: applying a zero displacement map to a non-trivial + image gives back the same image (within ITK resampling tolerance).""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + rng = np.random.default_rng(0) + img_data = rng.random((8, 8, 8), dtype=np.float32) + in_path = tmp_path / "img.nii" + nib.Nifti1Image(img_data, affine).to_filename(str(in_path)) + + # a 4D 1-channel zero displacement map + zero_map = np.zeros((8, 8, 8, 1), dtype=np.float32) + map_path = tmp_path / "map.nii" + nib.Nifti1Image(zero_map, affine).to_filename(str(map_path)) + + out_path = tmp_path / "out.nii" + argv( + [ + "wk-apply-warp", + "--input", + str(in_path), + "--transform", + str(map_path), + "--transform-type", + "map", + "--phase-encoding-axis", + "j", + "--output", + str(out_path), + ] + ) + apply_warp_main() + out = _load(str(out_path)) + out_data = out.get_fdata() + # Output is broadcast over the single transform frame, so it carries a + # length-1 time axis; squeeze for the value comparison. + assert out_data.squeeze().shape == img_data.shape + np.testing.assert_allclose(out_data.squeeze(), img_data, atol=1e-3) + + +def test_apply_warp_happy_path_4d_image_with_zero_field(argv, tmp_path): + """4D BOLD-like input + single 3-channel zero field broadcasts across + timepoints; identity resample should preserve every frame. Exercises the + single-field branch of the transform getter and the per-frame resample + loop.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + rng = np.random.default_rng(2) + img_data = rng.random((8, 8, 8, 4), dtype=np.float32) + in_path = tmp_path / "img.nii" + nib.Nifti1Image(img_data, affine).to_filename(str(in_path)) + + # single 4D 3-channel zero field + zero_field = np.zeros((8, 8, 8, 3), dtype=np.float32) + field_img = nib.Nifti1Image(zero_field, affine) + field_img.header.set_intent("vector", (), "") + field_path = tmp_path / "field.nii" + field_img.to_filename(str(field_path)) + + out_path = tmp_path / "out.nii" + argv( + [ + "wk-apply-warp", + "--input", + str(in_path), + "--transform", + str(field_path), + "--transform-type", + "field", + "--output", + str(out_path), + ] + ) + apply_warp_main() + out = _load(str(out_path)) + assert out.shape == img_data.shape + np.testing.assert_allclose(out.get_fdata(), img_data, atol=1e-3) + + +def test_apply_warp_happy_path_field_series(argv, tmp_path): + """Multi-file --transform = a per-frame field series. Exercises the + series branch of the transform getter.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + rng = np.random.default_rng(3) + img_data = rng.random((8, 8, 8, 3), dtype=np.float32) + in_path = tmp_path / "img.nii" + nib.Nifti1Image(img_data, affine).to_filename(str(in_path)) + + field_paths = [] + for i in range(3): + path = tmp_path / f"field_{i}.nii" + nib.Nifti1Image(np.zeros((8, 8, 8, 3), dtype=np.float32), affine).to_filename( + str(path) + ) + field_paths.append(str(path)) + + out_path = tmp_path / "out.nii" + argv( + [ + "wk-apply-warp", + "--input", + str(in_path), + "--transform", + *field_paths, + "--transform-type", + "field", + "--output", + str(out_path), + ] + ) + apply_warp_main() + out = _load(str(out_path)) + assert out.shape == img_data.shape + np.testing.assert_allclose(out.get_fdata(), img_data, atol=1e-3) + + +def test_convert_warp_happy_path_map_field_roundtrip(argv, tmp_path): + """Map -> field -> map round-trip preserves values within the + convert_warp orient-roundtrip jitter.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + rng = np.random.default_rng(0) + map_data = rng.random((6, 6, 6, 4), dtype=np.float32) + map_path = tmp_path / "maps.nii" + nib.Nifti1Image(map_data, affine).to_filename(str(map_path)) + + # maps -> bundled 5D field series (single output path) + field_path = tmp_path / "fields.nii" + argv( + [ + "wk-convert-warp", + "--input", + str(map_path), + "--output", + str(field_path), + "--from", + "map", + "--to", + "field", + "--axis", + "j", + ] + ) + convert_warp_main() + fields = _load(str(field_path)) + assert fields.shape == (6, 6, 6, 4, 3) + assert fields.header.get_intent()[0] == "vector" + + # fields -> maps (single bundled output path) + back_path = tmp_path / "maps_back.nii" + argv( + [ + "wk-convert-warp", + "--input", + str(field_path), + "--output", + str(back_path), + "--from", + "field", + "--to", + "map", + "--axis", + "j", + ] + ) + convert_warp_main() + back = _load(str(back_path)) + assert back.shape == map_data.shape + np.testing.assert_allclose(back.get_fdata(), map_data, atol=1e-3) + + +def test_convert_warp_happy_path_format_conversion(argv, tmp_path): + """itk -> ants format conversion: ants ch0/ch1 are sign-flipped vs itk, + output is 5D with a singleton 4th axis.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + rng = np.random.default_rng(1) + field_data = rng.random((6, 6, 6, 3), dtype=np.float32) + in_path = tmp_path / "field_itk.nii" + nib.Nifti1Image(field_data, affine).to_filename(str(in_path)) + + out_path = tmp_path / "field_ants.nii" + argv( + [ + "wk-convert-warp", + "--input", + str(in_path), + "--output", + str(out_path), + "--from", + "field", + "--to-format", + "ants", + ] + ) + convert_warp_main() + ants = _load(str(out_path)) + # ants single-field convention is 5D with shape (X, Y, Z, 1, 3) + assert ants.shape == (6, 6, 6, 1, 3) + ants_squeezed = ants.get_fdata().squeeze(axis=3) + # ants flip array is (-1, -1, 1) — channels 0 and 1 negated, 2 unchanged + np.testing.assert_allclose(ants_squeezed[..., 0], -field_data[..., 0], atol=1e-5) + np.testing.assert_allclose(ants_squeezed[..., 1], -field_data[..., 1], atol=1e-5) + np.testing.assert_allclose(ants_squeezed[..., 2], field_data[..., 2], atol=1e-5) + + +def test_convert_warp_happy_path_invert_zero_map(argv, tmp_path): + """Inverting a zero displacement map gives back zero (the identity warp + is its own inverse). Exercises the multi-frame map invert route.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + map_data = np.zeros((6, 6, 6, 3), dtype=np.float32) + in_path = tmp_path / "maps.nii" + nib.Nifti1Image(map_data, affine).to_filename(str(in_path)) + + out_path = tmp_path / "maps_inv.nii" + argv( + [ + "wk-convert-warp", + "--input", + str(in_path), + "--output", + str(out_path), + "--from", + "map", + "--invert", + "--axis", + "j", + ] + ) + convert_warp_main() + inv = _load(str(out_path)) + assert inv.shape == map_data.shape + np.testing.assert_allclose(inv.get_fdata(), 0.0, atol=1e-5) + + +# --------------------------------------------------------------------------- +# wk-compute-jacobian — argument validation + happy paths +# --------------------------------------------------------------------------- + + +def test_compute_jacobian_help(argv, capsys): + argv(["wk-compute-jacobian", "--help"]) + with pytest.raises(SystemExit) as exc: + compute_jacobian_main() + assert exc.value.code == 0 + out = capsys.readouterr().out + assert "Jacobian" in out + assert "--input" in out + assert "--axis" in out + assert "--frame" in out + + +def test_compute_jacobian_requires_axis_for_map(argv, capsys, tmp_path): + """Map inputs need --axis to promote to a 3-channel field.""" + maps = _write_nifti(tmp_path / "maps.nii", (4, 4, 4, 5)) + argv( + [ + "wk-compute-jacobian", + "--input", + maps, + "--output", + str(tmp_path / "out.nii"), + "--from", + "map", + ] + ) + with pytest.raises(SystemExit) as exc: + compute_jacobian_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "--axis is required" in err + + +def test_compute_jacobian_rejects_bad_format(argv, capsys, tmp_path): + field = _write_nifti(tmp_path / "field.nii", (4, 4, 4, 3)) + argv( + [ + "wk-compute-jacobian", + "--input", + field, + "--output", + str(tmp_path / "out.nii"), + "--from", + "field", + "--from-format", + "matlab", # not a valid choice + ] + ) + with pytest.raises(SystemExit) as exc: + compute_jacobian_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "invalid choice" in err + + +def test_compute_jacobian_frame_out_of_range(argv, capsys, tmp_path): + maps = _write_nifti(tmp_path / "maps.nii", (4, 4, 4, 5)) + argv( + [ + "wk-compute-jacobian", + "--input", + maps, + "--output", + str(tmp_path / "out.nii"), + "--from", + "map", + "--axis", + "j", + "--frame", + "10", + ] + ) + with pytest.raises(SystemExit) as exc: + compute_jacobian_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "out of range" in err + + +def test_compute_jacobian_single_field_zero(argv, tmp_path): + """Jacobian of a zero displacement field is identically 1.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + field_data = np.zeros((6, 6, 6, 3), dtype=np.float32) + field_path = tmp_path / "field.nii" + nib.Nifti1Image(field_data, affine).to_filename(str(field_path)) + + out_path = tmp_path / "jdet.nii" + argv( + [ + "wk-compute-jacobian", + "--input", + str(field_path), + "--output", + str(out_path), + "--from", + "field", + ] + ) + compute_jacobian_main() + j = _load(str(out_path)) + assert j.shape == (6, 6, 6) + np.testing.assert_allclose(j.get_fdata(), 1.0, atol=1e-5) + + +def test_compute_jacobian_map_series_zero(argv, tmp_path): + """Multi-frame zero map series → Jacobian of 1 per frame, bundled into a + 4D output. Exercises the map-promotion path and bundle write.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + map_data = np.zeros((6, 6, 6, 4), dtype=np.float32) + map_path = tmp_path / "maps.nii" + nib.Nifti1Image(map_data, affine).to_filename(str(map_path)) + + out_path = tmp_path / "jdet.nii" + argv( + [ + "wk-compute-jacobian", + "--input", + str(map_path), + "--output", + str(out_path), + "--from", + "map", + "--axis", + "j", + ] + ) + compute_jacobian_main() + j = _load(str(out_path)) + assert j.shape == (6, 6, 6, 4) + np.testing.assert_allclose(j.get_fdata(), 1.0, atol=1e-5) + + +def test_compute_jacobian_per_frame_outputs(argv, tmp_path): + """N output paths → one Jacobian volume per frame.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + map_data = np.zeros((6, 6, 6, 4), dtype=np.float32) + map_path = tmp_path / "maps.nii" + nib.Nifti1Image(map_data, affine).to_filename(str(map_path)) + + out_paths = [str(tmp_path / f"j{i}.nii") for i in range(4)] + argv( + [ + "wk-compute-jacobian", + "--input", + str(map_path), + "--output", + *out_paths, + "--from", + "map", + "--axis", + "j", + ] + ) + compute_jacobian_main() + for p in out_paths: + j = _load(p) + assert j.shape == (6, 6, 6) + np.testing.assert_allclose(j.get_fdata(), 1.0, atol=1e-5) + + +# --------------------------------------------------------------------------- +# wk-convert-fieldmap — argument validation + happy paths +# --------------------------------------------------------------------------- + + +def test_convert_fieldmap_help(argv, capsys): + argv(["wk-convert-fieldmap", "--help"]) + with pytest.raises(SystemExit) as exc: + convert_fieldmap_main() + assert exc.value.code == 0 + out = capsys.readouterr().out + assert "mm displacement" in out and "Hz" in out + assert "--total-readout-time" in out + assert "--phase-encoding-direction" in out + assert "--to" in out + assert "--flip-sign" in out + + +def test_convert_fieldmap_requires_to(argv, capsys, tmp_path): + """--to is required (no sensible default).""" + fmap = _write_nifti(tmp_path / "fmap.nii", (4, 4, 4)) + argv( + [ + "wk-convert-fieldmap", + "--input", + fmap, + "--output", + str(tmp_path / "out.nii"), + "--from", + "fieldmap", + ] + ) + with pytest.raises(SystemExit) as exc: + convert_fieldmap_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "--to" in err + + +def test_convert_fieldmap_rejects_bad_pe_direction(argv, capsys, tmp_path): + fmap = _write_nifti(tmp_path / "fmap.nii", (4, 4, 4)) + argv( + [ + "wk-convert-fieldmap", + "--input", + fmap, + "--output", + str(tmp_path / "out.nii"), + "--from", + "fieldmap", + "--to", + "map", + "--total-readout-time", + "0.05", + "--phase-encoding-direction", + "diagonal", + ] + ) + with pytest.raises(SystemExit) as exc: + convert_fieldmap_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "invalid choice" in err + + +def test_convert_fieldmap_requires_trt_and_pe(argv, capsys, tmp_path): + fmap = _write_nifti(tmp_path / "fmap.nii", (4, 4, 4)) + argv( + [ + "wk-convert-fieldmap", + "--input", + fmap, + "--output", + str(tmp_path / "out.nii"), + "--from", + "fieldmap", + "--to", + "map", + ] + ) + with pytest.raises(SystemExit) as exc: + convert_fieldmap_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "--total-readout-time" in err + assert "--phase-encoding-direction" in err + + +def test_convert_fieldmap_rejects_mm_to_mm(argv, capsys, tmp_path): + """--from=map --to=field is a representation conversion (no unit + crossing). The script should redirect users to wk-convert-warp.""" + maps = _write_nifti(tmp_path / "maps.nii", (4, 4, 4, 5)) + argv( + [ + "wk-convert-fieldmap", + "--input", + maps, + "--output", + str(tmp_path / "out.nii"), + "--from", + "map", + "--to", + "field", + "--total-readout-time", + "0.05", + "--phase-encoding-direction", + "j", + ] + ) + with pytest.raises(SystemExit) as exc: + convert_fieldmap_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "wk-convert-warp" in err + + +def test_convert_fieldmap_rejects_same_from_to(argv, capsys, tmp_path): + fmap = _write_nifti(tmp_path / "fmap.nii", (4, 4, 4)) + argv( + [ + "wk-convert-fieldmap", + "--input", + fmap, + "--output", + str(tmp_path / "out.nii"), + "--from", + "fieldmap", + "--to", + "fieldmap", + "--total-readout-time", + "0.05", + "--phase-encoding-direction", + "j", + ] + ) + with pytest.raises(SystemExit) as exc: + convert_fieldmap_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "are the same" in err + + +def test_convert_fieldmap_map_to_fieldmap_roundtrip(argv, tmp_path): + """mm displacement map -> Hz fieldmap -> mm displacement map preserves + the original (within float roundtrip jitter).""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + rng = np.random.default_rng(0) + map_data = rng.random((6, 6, 6, 4), dtype=np.float32) - 0.5 + map_path = tmp_path / "maps.nii" + nib.Nifti1Image(map_data, affine).to_filename(str(map_path)) + + fmap_path = tmp_path / "fmap.nii" + argv( + [ + "wk-convert-fieldmap", + "--input", + str(map_path), + "--output", + str(fmap_path), + "--from", + "map", + "--to", + "fieldmap", + "--total-readout-time", + "0.05", + "--phase-encoding-direction", + "j", + ] + ) + convert_fieldmap_main() + fmap = _load(str(fmap_path)) + assert fmap.shape == map_data.shape + assert np.isfinite(fmap.get_fdata()).all() + + back_path = tmp_path / "maps_back.nii" + argv( + [ + "wk-convert-fieldmap", + "--input", + str(fmap_path), + "--output", + str(back_path), + "--from", + "fieldmap", + "--to", + "map", + "--total-readout-time", + "0.05", + "--phase-encoding-direction", + "j", + ] + ) + convert_fieldmap_main() + back = _load(str(back_path)) + assert back.shape == map_data.shape + np.testing.assert_allclose(back.get_fdata(), map_data, atol=1e-5) + + +def test_convert_fieldmap_field_to_fieldmap(argv, tmp_path): + """3-channel mm displacement field -> Hz fieldmap. Exercises the + field->map axis-extraction branch upstream of the unit conversion.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + rng = np.random.default_rng(1) + field_data = np.zeros((6, 6, 6, 3), dtype=np.float32) + field_data[..., 1] = rng.random((6, 6, 6), dtype=np.float32) - 0.5 + field_path = tmp_path / "field.nii" + nib.Nifti1Image(field_data, affine).to_filename(str(field_path)) + + out_path = tmp_path / "fmap.nii" + argv( + [ + "wk-convert-fieldmap", + "--input", + str(field_path), + "--output", + str(out_path), + "--from", + "field", + "--to", + "fieldmap", + "--total-readout-time", + "0.05", + "--phase-encoding-direction", + "j", + ] + ) + convert_fieldmap_main() + fmap = _load(str(out_path)) + assert fmap.shape == (6, 6, 6) + assert np.isfinite(fmap.get_fdata()).all() + # zero off-axis channels in the field -> the fieldmap should be a + # nontrivial scalar volume (not all zeros). + assert float(np.abs(fmap.get_fdata()).max()) > 0.0 + + +def test_convert_fieldmap_fieldmap_to_field(argv, tmp_path): + """Hz fieldmap -> mm 3-channel field.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + rng = np.random.default_rng(2) + fmap_data = rng.random((6, 6, 6), dtype=np.float32) - 0.5 + fmap_path = tmp_path / "fmap.nii" + nib.Nifti1Image(fmap_data, affine).to_filename(str(fmap_path)) + + field_path = tmp_path / "field.nii" + argv( + [ + "wk-convert-fieldmap", + "--input", + str(fmap_path), + "--output", + str(field_path), + "--from", + "fieldmap", + "--to", + "field", + "--total-readout-time", + "0.05", + "--phase-encoding-direction", + "j", + ] + ) + convert_fieldmap_main() + field = _load(str(field_path)) + assert field.shape == (6, 6, 6, 3) + # j-axis has the displacement; off-axis channels are zero. + fdata = field.get_fdata() + assert float(np.abs(fdata[..., 1]).max()) > 0.0 + np.testing.assert_allclose(fdata[..., 0], 0.0, atol=1e-5) + np.testing.assert_allclose(fdata[..., 2], 0.0, atol=1e-5) + + +# --------------------------------------------------------------------------- +# _warp_io intent-leak regression tests +# --------------------------------------------------------------------------- + + +def _vector_intent_frame(shape=(4, 4, 4)) -> nib.Nifti1Image: + """Construct a 3D scalar frame whose header carries a stale vector intent + (the typical leak from an upstream field operation).""" + img = nib.Nifti1Image(np.zeros(shape, dtype=np.float32), np.eye(4)) + cast(nib.Nifti1Header, img.header).set_intent("vector", (), "") + return img + + +def test_bundle_frames_to_3d_series_clears_vector_intent(): + from warpkit.scripts._warp_io import bundle_frames_to_3d_series + + frames = [_vector_intent_frame() for _ in range(3)] + bundled = bundle_frames_to_3d_series(frames) + assert bundled.shape == (4, 4, 4, 3) + assert bundled.header.get_intent()[0] != "vector" + + +def test_write_output_per_frame_map_clears_vector_intent(tmp_path): + """Per-frame map outputs must round-trip without a stale vector intent.""" + import argparse + + from warpkit.scripts._warp_io import write_output + + frames = [_vector_intent_frame() for _ in range(2)] + out_paths = [str(tmp_path / "f1.nii"), str(tmp_path / "f2.nii")] + + write_output(frames, out_paths, "map", argparse.ArgumentParser()) + + for p in out_paths: + loaded = _load(p) + assert loaded.header.get_intent()[0] != "vector" + + +# --------------------------------------------------------------------------- +# CLI frame selection: --frame extracts the right frame, per-frame outputs +# preserve frame ordering. +# --------------------------------------------------------------------------- + + +def test_convert_warp_frame_extracts_right_frame(argv, tmp_path): + """--frame N picks the Nth frame of a 4D map series (not just any frame).""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + map_data = np.stack( + [np.full((4, 4, 4), float(i), dtype=np.float32) for i in range(5)], + axis=-1, + ) + map_path = tmp_path / "maps.nii" + nib.Nifti1Image(map_data, affine).to_filename(str(map_path)) + + out_path = tmp_path / "frame.nii" + argv( + [ + "wk-convert-warp", + "--input", + str(map_path), + "--output", + str(out_path), + "--from", + "map", + "--frame", + "3", + ] + ) + convert_warp_main() + out = _load(str(out_path)) + np.testing.assert_allclose(out.get_fdata(), 3.0, atol=1e-6) + + +def test_convert_warp_per_frame_outputs_preserve_ordering(argv, tmp_path): + """N --output paths for a 4D map series produce one file per frame with + frame i landing in the i-th output path (not shuffled).""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + map_data = np.stack( + [np.full((4, 4, 4), float(i), dtype=np.float32) for i in range(4)], + axis=-1, + ) + map_path = tmp_path / "maps.nii" + nib.Nifti1Image(map_data, affine).to_filename(str(map_path)) + + out_paths = [tmp_path / f"f{i}.nii" for i in range(4)] + argv( + [ + "wk-convert-warp", + "--input", + str(map_path), + "--output", + *[str(p) for p in out_paths], + "--from", + "map", + ] + ) + convert_warp_main() + for i, p in enumerate(out_paths): + np.testing.assert_allclose(_load(str(p)).get_fdata(), float(i), atol=1e-6) + + +def test_convert_fieldmap_frame_extracts_right_frame(argv, tmp_path): + """wk-convert-fieldmap --frame N selects the Nth frame before the Hz->mm + conversion.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + fmap_data = np.stack( + [np.full((4, 4, 4), 10.0 * (i + 1), dtype=np.float32) for i in range(3)], + axis=-1, + ) + fmap_path = tmp_path / "fmap.nii" + nib.Nifti1Image(fmap_data, affine).to_filename(str(fmap_path)) + + out_path = tmp_path / "dmap.nii" + argv( + [ + "wk-convert-fieldmap", + "--input", + str(fmap_path), + "--output", + str(out_path), + "--from", + "fieldmap", + "--to", + "map", + "--total-readout-time", + "0.05", + "--phase-encoding-direction", + "k", # axis 2 → no LPS-x/y flip; positive sign for "k" + "--frame", + "1", + ] + ) + convert_fieldmap_main() + # Frame 1 has fmap=20 Hz; voxel=2mm, trt=0.05s → expected disp = 20*0.05*2 = +2.0 mm + np.testing.assert_allclose(_load(str(out_path)).get_fdata(), 2.0, atol=1e-5) + + +# --------------------------------------------------------------------------- +# CLI flip-sign behavior: --flip-sign exactly negates the no-flip output +# --------------------------------------------------------------------------- + + +def test_convert_fieldmap_flip_sign_exact_negation(argv, tmp_path): + """Running map -> fieldmap twice (with and without --flip-sign) on the + same input must produce results that are exact negations.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + rng = np.random.default_rng(0) + map_data = (rng.standard_normal((6, 6, 6)) * 1.5).astype(np.float32) + map_path = tmp_path / "map.nii" + nib.Nifti1Image(map_data, affine).to_filename(str(map_path)) + + common = [ + "--input", + str(map_path), + "--from", + "map", + "--to", + "fieldmap", + "--total-readout-time", + "0.05", + "--phase-encoding-direction", + "j", + ] + + out_plain = tmp_path / "fmap_plain.nii" + argv(["wk-convert-fieldmap", *common, "--output", str(out_plain)]) + convert_fieldmap_main() + out_flip = tmp_path / "fmap_flip.nii" + argv(["wk-convert-fieldmap", *common, "--output", str(out_flip), "--flip-sign"]) + convert_fieldmap_main() + + np.testing.assert_allclose( + _load(str(out_flip)).get_fdata(), + -_load(str(out_plain)).get_fdata(), + atol=1e-6, + ) + + +# --------------------------------------------------------------------------- +# CLI 5D ANTs/AFNI single-warp file roundtrip +# --------------------------------------------------------------------------- + + +def test_convert_warp_5d_ants_file_roundtrip(argv, tmp_path): + """A 5D ANTs single-warp file (X,Y,Z,1,3) must load through wk-convert-warp, + convert to itk, and convert back to ants without losing data.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + rng = np.random.default_rng(0) + ants_5d = rng.standard_normal((5, 5, 5, 1, 3)).astype(np.float32) + in_path = tmp_path / "ants.nii" + img = nib.Nifti1Image(ants_5d, affine) + cast(nib.Nifti1Header, img.header).set_intent("vector", (), "") + img.to_filename(str(in_path)) + + itk_path = tmp_path / "itk.nii" + argv( + [ + "wk-convert-warp", + "--input", + str(in_path), + "--output", + str(itk_path), + "--from", + "field", + "--from-format", + "ants", + "--to-format", + "itk", + ] + ) + convert_warp_main() + + back_path = tmp_path / "ants_back.nii" + argv( + [ + "wk-convert-warp", + "--input", + str(itk_path), + "--output", + str(back_path), + "--from", + "field", + "--from-format", + "itk", + "--to-format", + "ants", + ] + ) + convert_warp_main() + back = _load(str(back_path)).get_fdata() + # Result should match the original 5D ANTs file (within orient-roundtrip noise). + assert back.shape == ants_5d.shape + np.testing.assert_allclose(back, ants_5d, atol=1e-5) + + +# --------------------------------------------------------------------------- +# CLI medic-chain reconstruction: +# fieldmap_native (Hz, distorted) -> mm -> invert -> Hz (with --flip-sign) +# must reproduce the medic non-native fieldmap output bit-for-bit. +# --------------------------------------------------------------------------- + + +def test_cli_chain_reproduces_medic_non_native_fieldmap(argv, tmp_path): + """Hz->mm->invert->Hz(--flip-sign) chain via the CLIs matches what + warpkit.distortion.medic does internally (distortion.py:122-149). + Skips the correlation-based safety negation, which doesn't fire on + smooth synthetic input.""" + from scipy.ndimage import gaussian_filter + from warpkit.utilities import ( + displacement_maps_to_field_maps, + field_maps_to_displacement_maps, + invert_displacement_maps, + ) + + pe = "j" + trt = 0.05 + affine = np.diag([2.5, 2.5, 2.5, 1.0]) + rng = np.random.default_rng(0) + base = gaussian_filter(rng.standard_normal((10, 10, 10)) * 20.0, sigma=2.0) + fmap_data = np.stack([base, 0.95 * base], axis=-1).astype(np.float32) + fmap_path = tmp_path / "fmap_native.nii" + nib.Nifti1Image(fmap_data, affine).to_filename(str(fmap_path)) + + # 1) Reference: medic's internal chain, frame-by-frame. + medic_frames = [] + for i in range(fmap_data.shape[-1]): + frame_4d = nib.Nifti1Image(fmap_data[..., i : i + 1], affine) + inv_disp = field_maps_to_displacement_maps(frame_4d, trt, pe) + disp = invert_displacement_maps(inv_disp, pe) + recon = displacement_maps_to_field_maps(disp, trt, pe, flip_sign=True) + medic_frames.append(recon.get_fdata()) + medic_out = np.stack(medic_frames, axis=-1).squeeze() + + # 2) CLI chain: Hz -> mm + dmap_distorted = tmp_path / "dmap_distorted.nii" + argv( + [ + "wk-convert-fieldmap", + "--input", + str(fmap_path), + "--output", + str(dmap_distorted), + "--from", + "fieldmap", + "--to", + "map", + "--total-readout-time", + str(trt), + "--phase-encoding-direction", + pe, + ] + ) + convert_fieldmap_main() + + # 3) invert + dmap_undistorted = tmp_path / "dmap_undistorted.nii" + argv( + [ + "wk-convert-warp", + "--input", + str(dmap_distorted), + "--output", + str(dmap_undistorted), + "--from", + "map", + "--invert", + "--axis", + pe, + ] + ) + convert_warp_main() + + # 4) mm -> Hz with --flip-sign + fmap_recon = tmp_path / "fmap_recon.nii" + argv( + [ + "wk-convert-fieldmap", + "--input", + str(dmap_undistorted), + "--output", + str(fmap_recon), + "--from", + "map", + "--to", + "fieldmap", + "--flip-sign", + "--total-readout-time", + str(trt), + "--phase-encoding-direction", + pe, + ] + ) + convert_fieldmap_main() + + cli_out = _load(str(fmap_recon)).get_fdata() + assert cli_out.shape == medic_out.shape + np.testing.assert_allclose(cli_out, medic_out, atol=1e-5) + + +# --------------------------------------------------------------------------- +# CLI inverse self-consistency: invert(invert(disp)) ≈ disp on small smooth disp. +# --------------------------------------------------------------------------- + + +def test_convert_warp_double_invert_recovers_input(argv, tmp_path): + """Double inversion of a small smooth displacement map approximately + recovers the original (an invertibility / discretization sanity check).""" + from scipy.ndimage import gaussian_filter + + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + rng = np.random.default_rng(0) + base = gaussian_filter(rng.standard_normal((12, 12, 12)) * 0.3, sigma=2.5) + map_data = base.astype(np.float32) + in_path = tmp_path / "map.nii" + nib.Nifti1Image(map_data, affine).to_filename(str(in_path)) + + inv_path = tmp_path / "map_inv.nii" + argv( + [ + "wk-convert-warp", + "--input", + str(in_path), + "--output", + str(inv_path), + "--from", + "map", + "--invert", + "--axis", + "j", + ] + ) + convert_warp_main() + + inv2_path = tmp_path / "map_inv2.nii" + argv( + [ + "wk-convert-warp", + "--input", + str(inv_path), + "--output", + str(inv2_path), + "--from", + "map", + "--invert", + "--axis", + "j", + ] + ) + convert_warp_main() + recovered = _load(str(inv2_path)).get_fdata() + # Discretization error grows with displacement magnitude; trim a 2-voxel + # border and use a loose tolerance. + interior = (slice(2, -2),) * 3 + np.testing.assert_allclose(recovered[interior], map_data[interior], atol=5e-2) + + +# --------------------------------------------------------------------------- +# wk-apply-warp coverage gaps: 3D map, 5D ANTs single field, 5D field series, +# --reference, multi-file + --transform-type=map error, non-itk --format. +# --------------------------------------------------------------------------- + + +def test_apply_warp_multi_file_with_transform_type_map_errors(argv, capsys, tmp_path): + """A multi-file --transform with --transform-type=map is incompatible + (only fields can be passed as a multi-file series).""" + inp = _write_nifti(tmp_path / "in.nii", (4, 4, 4, 2)) + f1 = _write_nifti(tmp_path / "f1.nii", (4, 4, 4, 3)) + f2 = _write_nifti(tmp_path / "f2.nii", (4, 4, 4, 3)) + argv( + [ + "wk-apply-warp", + "--input", + inp, + "--transform", + f1, + f2, + "--transform-type", + "map", + "--phase-encoding-axis", + "j", + "--output", + str(tmp_path / "out.nii"), + ] + ) + with pytest.raises(SystemExit) as exc: + apply_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "incompatible" in err and "field" in err + + +def test_apply_warp_invalid_map_shape_errors(argv, capsys, tmp_path): + """A 5D --transform with --transform-type=map is rejected (maps are 3D + or 4D scalar).""" + inp = _write_nifti(tmp_path / "in.nii", (4, 4, 4)) + bad_map = _write_nifti(tmp_path / "bad.nii", (4, 4, 4, 1, 3)) + argv( + [ + "wk-apply-warp", + "--input", + inp, + "--transform", + bad_map, + "--transform-type", + "map", + "--phase-encoding-axis", + "j", + "--output", + str(tmp_path / "out.nii"), + ] + ) + with pytest.raises(SystemExit) as exc: + apply_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "displacement map must be 3D or 4D" in err + + +def test_apply_warp_invalid_input_shape_errors(argv, capsys, tmp_path): + """A 5D --input is rejected (input must be 3D or 4D).""" + inp = _write_nifti(tmp_path / "in.nii", (4, 4, 4, 2, 2)) # 5D + tx = _write_nifti(tmp_path / "tx.nii", (4, 4, 4, 3)) + argv( + [ + "wk-apply-warp", + "--input", + inp, + "--transform", + tx, + "--transform-type", + "field", + "--output", + str(tmp_path / "out.nii"), + ] + ) + with pytest.raises(SystemExit) as exc: + apply_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "input image must be 3D or 4D" in err + + +def test_apply_warp_3d_map_transform(argv, tmp_path): + """A 3D single-frame displacement map (not 4D) routes through the cached + single-frame branch.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + rng = np.random.default_rng(0) + img_data = rng.random((6, 6, 6), dtype=np.float32) + in_path = tmp_path / "img.nii" + nib.Nifti1Image(img_data, affine).to_filename(str(in_path)) + + map_path = tmp_path / "map.nii" + nib.Nifti1Image(np.zeros((6, 6, 6), dtype=np.float32), affine).to_filename( + str(map_path) + ) + + out_path = tmp_path / "out.nii" + argv( + [ + "wk-apply-warp", + "--input", + str(in_path), + "--transform", + str(map_path), + "--transform-type", + "map", + "--phase-encoding-axis", + "j", + "--output", + str(out_path), + ] + ) + apply_warp_main() + out = _load(str(out_path)) + np.testing.assert_allclose(out.get_fdata(), img_data, atol=1e-3) + + +def test_apply_warp_5d_ants_single_field(argv, tmp_path): + """5D (X,Y,Z,1,3) zero ANTs single-warp file. Identity resample preserves + the input.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + rng = np.random.default_rng(0) + img_data = rng.random((6, 6, 6), dtype=np.float32) + in_path = tmp_path / "img.nii" + nib.Nifti1Image(img_data, affine).to_filename(str(in_path)) + + field_5d = np.zeros((6, 6, 6, 1, 3), dtype=np.float32) + field_path = tmp_path / "ants.nii" + field_img = nib.Nifti1Image(field_5d, affine) + cast(nib.Nifti1Header, field_img.header).set_intent("vector", (), "") + field_img.to_filename(str(field_path)) + + out_path = tmp_path / "out.nii" + argv( + [ + "wk-apply-warp", + "--input", + str(in_path), + "--transform", + str(field_path), + "--transform-type", + "field", + "--format", + "ants", + "--output", + str(out_path), + ] + ) + apply_warp_main() + out = _load(str(out_path)) + np.testing.assert_allclose(out.get_fdata().squeeze(), img_data, atol=1e-3) + + +def test_apply_warp_5d_field_series_per_frame(argv, tmp_path): + """A 5D (X,Y,Z,T,3) field series in a single file applied to a 4D input; + each frame uses its own field. Zero fields → identity per frame.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + rng = np.random.default_rng(0) + img_data = rng.random((5, 5, 5, 3), dtype=np.float32) + in_path = tmp_path / "img.nii" + nib.Nifti1Image(img_data, affine).to_filename(str(in_path)) + + field_5d = np.zeros((5, 5, 5, 3, 3), dtype=np.float32) + field_img = nib.Nifti1Image(field_5d, affine) + cast(nib.Nifti1Header, field_img.header).set_intent("vector", (), "") + field_path = tmp_path / "field5d.nii" + field_img.to_filename(str(field_path)) + + out_path = tmp_path / "out.nii" + argv( + [ + "wk-apply-warp", + "--input", + str(in_path), + "--transform", + str(field_path), + "--transform-type", + "field", + "--output", + str(out_path), + ] + ) + apply_warp_main() + out = _load(str(out_path)) + assert out.shape == img_data.shape + np.testing.assert_allclose(out.get_fdata(), img_data, atol=1e-3) + + +def test_apply_warp_with_explicit_reference(argv, tmp_path): + """--reference uses an explicit grid (not the input). The output adopts + the reference's shape and affine.""" + in_affine = np.diag([2.0, 2.0, 2.0, 1.0]) + ref_affine = np.diag([2.0, 2.0, 2.0, 1.0]) + ref_affine[:3, 3] = [1.0, 2.0, 3.0] # different origin + + rng = np.random.default_rng(0) + img = rng.random((6, 6, 6), dtype=np.float32) + in_path = tmp_path / "img.nii" + nib.Nifti1Image(img, in_affine).to_filename(str(in_path)) + + ref = np.zeros((4, 4, 4), dtype=np.float32) + ref_path = tmp_path / "ref.nii" + nib.Nifti1Image(ref, ref_affine).to_filename(str(ref_path)) + + field = np.zeros((6, 6, 6, 3), dtype=np.float32) + field_img = nib.Nifti1Image(field, in_affine) + cast(nib.Nifti1Header, field_img.header).set_intent("vector", (), "") + field_path = tmp_path / "field.nii" + field_img.to_filename(str(field_path)) + + out_path = tmp_path / "out.nii" + argv( + [ + "wk-apply-warp", + "--input", + str(in_path), + "--reference", + str(ref_path), + "--transform", + str(field_path), + "--transform-type", + "field", + "--output", + str(out_path), + ] + ) + apply_warp_main() + out = _load(str(out_path)) + # Output adopts the reference grid shape (4,4,4), not the input (6,6,6). + assert out.shape == (4, 4, 4) + + +def test_apply_warp_field_with_non_itk_format(argv, tmp_path): + """A zero field in fsl/ants/afni format should still resample to identity + after the format conversion in convert_warp.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + rng = np.random.default_rng(0) + img_data = rng.random((5, 5, 5), dtype=np.float32) + in_path = tmp_path / "img.nii" + nib.Nifti1Image(img_data, affine).to_filename(str(in_path)) + + field_path = tmp_path / "field_fsl.nii" + nib.Nifti1Image(np.zeros((5, 5, 5, 3), dtype=np.float32), affine).to_filename( + str(field_path) + ) + + out_path = tmp_path / "out.nii" + argv( + [ + "wk-apply-warp", + "--input", + str(in_path), + "--transform", + str(field_path), + "--transform-type", + "field", + "--format", + "fsl", + "--output", + str(out_path), + ] + ) + apply_warp_main() + out = _load(str(out_path)) + np.testing.assert_allclose(out.get_fdata().squeeze(), img_data, atol=1e-3) + + +# --------------------------------------------------------------------------- +# wk-compute-jacobian coverage gaps: non-itk format, --frame, 5D series. +# --------------------------------------------------------------------------- + + +def test_compute_jacobian_field_with_ants_format(argv, tmp_path): + """ANTs-format zero field still routes through convert_warp before the + Jacobian computation; result is identically 1.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + field_5d = np.zeros((6, 6, 6, 1, 3), dtype=np.float32) + field_img = nib.Nifti1Image(field_5d, affine) + cast(nib.Nifti1Header, field_img.header).set_intent("vector", (), "") + field_path = tmp_path / "ants.nii" + field_img.to_filename(str(field_path)) + + out_path = tmp_path / "jdet.nii" + argv( + [ + "wk-compute-jacobian", + "--input", + str(field_path), + "--output", + str(out_path), + "--from", + "field", + "--from-format", + "ants", + ] + ) + compute_jacobian_main() + np.testing.assert_allclose(_load(str(out_path)).get_fdata(), 1.0, atol=1e-5) + + +def test_compute_jacobian_with_frame(argv, tmp_path): + """--frame N picks the Nth field of a 5D field series.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + field_5d = np.zeros((5, 5, 5, 3, 3), dtype=np.float32) + field_img = nib.Nifti1Image(field_5d, affine) + cast(nib.Nifti1Header, field_img.header).set_intent("vector", (), "") + field_path = tmp_path / "field5d.nii" + field_img.to_filename(str(field_path)) + + out_path = tmp_path / "jdet.nii" + argv( + [ + "wk-compute-jacobian", + "--input", + str(field_path), + "--output", + str(out_path), + "--from", + "field", + "--frame", + "1", + ] + ) + compute_jacobian_main() + j = _load(str(out_path)) + assert j.shape == (5, 5, 5) + np.testing.assert_allclose(j.get_fdata(), 1.0, atol=1e-5) + + +def test_compute_jacobian_bundled_vs_per_frame_outputs_match(argv, tmp_path): + """Bundled (single output) and per-frame (N outputs) writes of the same + input produce identical data.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + rng = np.random.default_rng(0) + map_data = (rng.standard_normal((5, 5, 5, 3)) * 0.1).astype(np.float32) + map_path = tmp_path / "maps.nii" + nib.Nifti1Image(map_data, affine).to_filename(str(map_path)) + + bundled_path = tmp_path / "jdet_bundled.nii" + argv( + [ + "wk-compute-jacobian", + "--input", + str(map_path), + "--output", + str(bundled_path), + "--from", + "map", + "--axis", + "j", + ] + ) + compute_jacobian_main() + bundled = _load(str(bundled_path)).get_fdata() + assert bundled.shape == (5, 5, 5, 3) + + per_frame_paths = [tmp_path / f"jdet_{i}.nii" for i in range(3)] + argv( + [ + "wk-compute-jacobian", + "--input", + str(map_path), + "--output", + *[str(p) for p in per_frame_paths], + "--from", + "map", + "--axis", + "j", + ] + ) + compute_jacobian_main() + for i, p in enumerate(per_frame_paths): + np.testing.assert_allclose(_load(str(p)).get_fdata(), bundled[..., i], atol=0) + + +# --------------------------------------------------------------------------- +# wk-convert-warp coverage gaps: invert + format conversion + axis combo, +# multi-frame field series invert. +# --------------------------------------------------------------------------- + + +def test_convert_warp_invert_field_with_format_conversion(argv, tmp_path): + """Single-frame field input in fsl format + --invert. Routes through + convert_warp(fsl→itk) before inversion. Zero field inverts to zero.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + field = np.zeros((5, 5, 5, 3), dtype=np.float32) + in_path = tmp_path / "field_fsl.nii" + nib.Nifti1Image(field, affine).to_filename(str(in_path)) + + out_path = tmp_path / "field_inv.nii" + argv( + [ + "wk-convert-warp", + "--input", + str(in_path), + "--output", + str(out_path), + "--from", + "field", + "--from-format", + "fsl", + "--to-format", + "itk", + "--invert", + ] + ) + convert_warp_main() + out = _load(str(out_path)) + np.testing.assert_allclose(out.get_fdata(), 0.0, atol=1e-5) + + +def test_convert_warp_multi_frame_field_invert_routes_through_map_inverter( + argv, tmp_path +): + """Multi-frame field input + --invert: route extracts per-frame channel + along --axis, runs the 1D map inverter, then promotes maps back to fields + for the field output (default --to matches --from). Zero in, zero out.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + fields = [] + for i in range(3): + path = tmp_path / f"f{i}.nii" + nib.Nifti1Image(np.zeros((5, 5, 5, 3), dtype=np.float32), affine).to_filename( + str(path) + ) + fields.append(str(path)) + + out_path = tmp_path / "inv.nii" + argv( + [ + "wk-convert-warp", + "--input", + *fields, + "--output", + str(out_path), + "--from", + "field", + "--invert", + "--axis", + "j", + ] + ) + convert_warp_main() + loaded = _load(str(out_path)) + # Output is a 5D field series (3 frames, 3 channels). + assert loaded.shape == (5, 5, 5, 3, 3) + assert loaded.header.get_intent()[0] == "vector" + np.testing.assert_allclose(loaded.get_fdata(), 0.0, atol=1e-5) + + +def test_convert_warp_multi_frame_field_invert_to_map_output(argv, tmp_path): + """Multi-frame field input + --invert + --to=map: maps come out of the + inverter (1-channel along --axis) and stay as a map series. Zero in, zero + out, and the map series header has no vector intent.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + fields = [] + for i in range(3): + path = tmp_path / f"f{i}.nii" + nib.Nifti1Image(np.zeros((5, 5, 5, 3), dtype=np.float32), affine).to_filename( + str(path) + ) + fields.append(str(path)) + + out_path = tmp_path / "inv_maps.nii" + argv( + [ + "wk-convert-warp", + "--input", + *fields, + "--output", + str(out_path), + "--from", + "field", + "--to", + "map", + "--invert", + "--axis", + "j", + ] + ) + convert_warp_main() + loaded = _load(str(out_path)) + # Map series: (X, Y, Z, T) without channel dim; intent is not 'vector'. + assert loaded.shape == (5, 5, 5, 3) + assert loaded.header.get_intent()[0] != "vector" + np.testing.assert_allclose(loaded.get_fdata(), 0.0, atol=1e-5) + + +# --------------------------------------------------------------------------- +# wk-convert-fieldmap coverage gaps: multi-file fieldmap input, 5D field input. +# --------------------------------------------------------------------------- + + +def test_convert_fieldmap_multi_file_input(argv, tmp_path): + """Multiple --input fieldmap files are flattened into a frame series.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + fmaps = [] + for i in range(3): + path = tmp_path / f"fmap_{i}.nii" + nib.Nifti1Image( + np.full((4, 4, 4), 10.0 * (i + 1), dtype=np.float32), affine + ).to_filename(str(path)) + fmaps.append(str(path)) + + out_path = tmp_path / "dmap.nii" + argv( + [ + "wk-convert-fieldmap", + "--input", + *fmaps, + "--output", + str(out_path), + "--from", + "fieldmap", + "--to", + "map", + "--total-readout-time", + "0.05", + "--phase-encoding-direction", + "k", # voxel z=2, no LPS flip → disp = +fmap*0.05*2 + ] + ) + convert_fieldmap_main() + bundled = _load(str(out_path)).get_fdata() + # Three 1-channel maps stacked into a 4D series; frame i has fmap=10*(i+1). + assert bundled.shape == (4, 4, 4, 3) + for i in range(3): + np.testing.assert_allclose( + bundled[..., i], 10.0 * (i + 1) * 0.05 * 2.0, atol=1e-5 + ) + + +def test_convert_fieldmap_5d_field_input(argv, tmp_path): + """5D (X,Y,Z,T,3) field series input gets split into per-frame fields + before unit conversion.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + field_5d = np.zeros((4, 4, 4, 2, 3), dtype=np.float32) + field_5d[..., 0, 1] = 1.0 # frame 0: 1mm displacement on j channel + field_5d[..., 1, 1] = 2.0 # frame 1: 2mm displacement on j channel + field_path = tmp_path / "field5d.nii" + img = nib.Nifti1Image(field_5d, affine) + cast(nib.Nifti1Header, img.header).set_intent("vector", (), "") + img.to_filename(str(field_path)) + + out_path = tmp_path / "fmap.nii" + argv( + [ + "wk-convert-fieldmap", + "--input", + str(field_path), + "--output", + str(out_path), + "--from", + "field", + "--to", + "fieldmap", + "--total-readout-time", + "0.05", + "--phase-encoding-direction", + "j", + ] + ) + convert_fieldmap_main() + fmap = _load(str(out_path)).get_fdata() + # Two frames bundled into 4D scalar series. + assert fmap.shape == (4, 4, 4, 2) + # |fmap[i]| / |dmap[i]| = 1 / (trt * voxel) = 1 / 0.1 = 10 Hz/mm + # frame 0: 1mm → 10 Hz; frame 1: 2mm → 20 Hz (sign depends on PE convention) + np.testing.assert_allclose(np.abs(fmap[..., 0]), 10.0, atol=1e-4) + np.testing.assert_allclose(np.abs(fmap[..., 1]), 20.0, atol=1e-4) + + +# --------------------------------------------------------------------------- +# input-shape error messages from _warp_io +# --------------------------------------------------------------------------- + + +def test_warp_io_map_input_5d_errors(argv, capsys, tmp_path): + """--from=map with a 5D file errors out (maps are 3D or 4D).""" + bad = _write_nifti(tmp_path / "bad.nii", (4, 4, 4, 1, 3)) + argv( + [ + "wk-convert-warp", + "--input", + bad, + "--output", + str(tmp_path / "out.nii"), + "--from", + "map", + ] + ) + with pytest.raises(SystemExit) as exc: + convert_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "map input must be 3D or 4D" in err + + +def test_warp_io_field_input_4d_wrong_channel_errors(argv, capsys, tmp_path): + """--from=field with a 4D file whose last dim isn't 3 errors out.""" + bad = _write_nifti(tmp_path / "bad.nii", (4, 4, 4, 7)) + argv( + [ + "wk-convert-warp", + "--input", + bad, + "--output", + str(tmp_path / "out.nii"), + "--from", + "field", + ] + ) + with pytest.raises(SystemExit) as exc: + convert_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "field input must be 4D" in err + + +def test_apply_warp_invalid_field_shape_errors(argv, capsys, tmp_path): + """A 4D --transform whose last dim != 3 with --transform-type=field is + rejected (fields must be 4D X,Y,Z,3 or 5D X,Y,Z,T,3).""" + inp = _write_nifti(tmp_path / "in.nii", (4, 4, 4)) + bad_field = _write_nifti(tmp_path / "bad.nii", (4, 4, 4, 7)) + argv( + [ + "wk-apply-warp", + "--input", + inp, + "--transform", + bad_field, + "--transform-type", + "field", + "--output", + str(tmp_path / "out.nii"), + ] + ) + with pytest.raises(SystemExit) as exc: + apply_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "displacement field must be 4D" in err + + +def test_convert_fieldmap_frame_out_of_range_errors(argv, capsys, tmp_path): + """wk-convert-fieldmap --frame N where N >= number of frames is rejected.""" + fmap = _write_nifti(tmp_path / "fmap.nii", (4, 4, 4, 3)) + argv( + [ + "wk-convert-fieldmap", + "--input", + fmap, + "--output", + str(tmp_path / "out.nii"), + "--from", + "fieldmap", + "--to", + "map", + "--total-readout-time", + "0.05", + "--phase-encoding-direction", + "j", + "--frame", + "10", + ] + ) + with pytest.raises(SystemExit) as exc: + convert_fieldmap_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "out of range" in err + + +def test_unwrap_phase_noiseframes_strips_trailing_volumes( + argv, test_data_paths, tmp_path +): + """--noiseframes N drops the last N volumes from each input. Run on the + bundled fixture (15 volumes) with -f 1 and assert the unwrapped output + has 14 volumes.""" + out_prefix = tmp_path / "noise" + argv( + [ + "wk-unwrap-phase", + "--magnitude", + *test_data_paths["mag"], + "--phase", + *test_data_paths["phase"], + "--metadata", + *test_data_paths["metadata"], + "--out-prefix", + str(out_prefix), + "-n", + "1", + "-f", + "1", # drop last frame + ] + ) + unwrap_phase_main() + unwrapped = sorted(tmp_path.glob("noise_unwrapped_echo-*.nii")) + assert len(unwrapped) == 3 + for u in unwrapped: + assert _load(str(u)).shape == (64, 64, 40, 14) # was 15, minus 1 + + +def test_bundle_frames_to_field_series_squeezes_5d_singleton_input(tmp_path): + """When per-frame fields are themselves 5D (X,Y,Z,1,3), the bundler must + squeeze the singleton 4th axis before stacking — otherwise the result + would be 6D.""" + from warpkit.scripts._warp_io import bundle_frames_to_field_series + + frames = [] + for _ in range(2): + data = np.zeros((4, 4, 4, 1, 3), dtype=np.float32) + frames.append(nib.Nifti1Image(data, np.eye(4))) + bundled = bundle_frames_to_field_series(frames) + assert bundled.shape == (4, 4, 4, 2, 3) + + +def test_apply_warp_with_5d_field_with_unsupported_last_axis_errors( + argv, capsys, tmp_path +): + """A 5D --transform whose last dim != 3 with --transform-type=field is + rejected. Covers the 5D-but-not-3-channel branch in the field validator.""" + inp = _write_nifti(tmp_path / "in.nii", (4, 4, 4)) + bad = _write_nifti(tmp_path / "bad5d.nii", (4, 4, 4, 1, 7)) + argv( + [ + "wk-apply-warp", + "--input", + inp, + "--transform", + bad, + "--transform-type", + "field", + "--output", + str(tmp_path / "out.nii"), + ] + ) + with pytest.raises(SystemExit) as exc: + apply_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "displacement field must be 4D" in err diff --git a/tests/test_unwrap.py b/tests/test_unwrap.py index a52e961..94deaf7 100644 --- a/tests/test_unwrap.py +++ b/tests/test_unwrap.py @@ -1,7 +1,8 @@ +import nibabel as nib import numpy as np import pytest from numpy.testing import assert_allclose -from warpkit.unwrap import compute_offset, reject_outliers +from warpkit.unwrap import compute_field_maps, compute_offset, reject_outliers # --------------------------------------------------------------------------- # reject_outliers: median + MAD, threshold m=2.0 @@ -68,31 +69,83 @@ def test_compute_offset_recovers_known_2pi_shift(): # --------------------------------------------------------------------------- -# Romeo bindings — light smoke test that the new lowercase names exist with -# the expected signature shape (the heavy property tests live in test_romeo). +# Romeo bindings — light smoke test that the lowercase free-function names +# exist at module level (the heavy property tests live in test_romeo). # --------------------------------------------------------------------------- def test_romeo_lowercase_bindings_exist(): - from warpkit.warpkit_cpp import Romeo - - romeo = Romeo() - # the Python-visible names are lowercase post-rename; the original - # uppercase names should not exist as attributes. - assert hasattr(romeo, "romeo_unwrap3d") - assert hasattr(romeo, "romeo_unwrap4d") - assert hasattr(romeo, "romeo_voxelquality") - assert not hasattr(romeo, "romeo_unwrap3D") - assert not hasattr(romeo, "romeo_unwrap4D") + import warpkit.warpkit_cpp as cpp + + # Python-visible names are lowercase free functions on the module; + # there is no Romeo class wrapper, and the original uppercase names + # should not appear at module level either. + assert hasattr(cpp, "romeo_unwrap3d") + assert hasattr(cpp, "romeo_unwrap4d") + assert hasattr(cpp, "romeo_voxelquality") + assert hasattr(cpp, "calculate_weights") + assert not hasattr(cpp, "Romeo") + assert not hasattr(cpp, "romeo_unwrap3D") + assert not hasattr(cpp, "romeo_unwrap4D") + + +def _make_field_inputs(spatial=(4, 4, 4), n_frames=2, n_echoes=3): + """Build minimal valid ``compute_field_maps`` inputs (unwrapped per-echo + 4D images + per-frame masks). Caller can swap the masks to break the + expected shape.""" + affine = np.eye(4) + unwrapped = [ + nib.Nifti1Image( + np.zeros((*spatial, n_frames), dtype=np.float32), + affine, + ) + for _ in range(n_echoes) + ] + mag = [ + nib.Nifti1Image( + np.ones((*spatial, n_frames), dtype=np.float32), + affine, + ) + for _ in range(n_echoes) + ] + masks = nib.Nifti1Image( + np.ones((*spatial, n_frames), dtype=np.int8), + affine, + ) + tes = [10.0, 20.0, 30.0] + return unwrapped, mag, masks, tes + + +def test_compute_field_maps_rejects_3d_masks(): + """A masks input that's missing the time axis must fail loudly, not + silently broadcast or crash deeper inside the SVD pass.""" + unwrapped, mag, _masks, tes = _make_field_inputs() + bad_masks = nib.Nifti1Image(np.ones((4, 4, 4), dtype=np.int8), np.eye(4)) + with pytest.raises(ValueError, match="masks must have shape"): + compute_field_maps(unwrapped, bad_masks, mag, tes) + + +def test_compute_field_maps_rejects_mismatched_frame_count(): + """masks frame count must match the unwrapped time dimension.""" + unwrapped, mag, _masks, tes = _make_field_inputs(n_frames=2) + bad_masks = nib.Nifti1Image(np.ones((4, 4, 4, 5), dtype=np.int8), np.eye(4)) + with pytest.raises(ValueError, match="masks must have shape"): + compute_field_maps(unwrapped, bad_masks, mag, tes) + + +def test_compute_field_maps_rejects_mismatched_spatial_shape(): + unwrapped, mag, _masks, tes = _make_field_inputs(spatial=(4, 4, 4)) + bad_masks = nib.Nifti1Image(np.ones((5, 4, 4, 2), dtype=np.int8), np.eye(4)) + with pytest.raises(ValueError, match="masks must have shape"): + compute_field_maps(unwrapped, bad_masks, mag, tes) def test_romeo_unwrap3d_rejects_unknown_weight_preset(): """`weights` is a preset name string; only "romeo" is supported.""" - from warpkit.warpkit_cpp import Romeo + from warpkit.warpkit_cpp import romeo_unwrap3d - romeo = Romeo() phase = np.zeros((3, 3, 3), dtype=np.float32) mag = np.ones_like(phase) mask = np.ones(phase.shape, dtype=bool) with pytest.raises(Exception, match='only the "romeo" weight preset'): - romeo.romeo_unwrap3d(phase, "ramen", mag, mask) + romeo_unwrap3d(phase, "ramen", mag, mask) diff --git a/tests/test_utilities.py b/tests/test_utilities.py index a071eb2..fc375b7 100644 --- a/tests/test_utilities.py +++ b/tests/test_utilities.py @@ -9,9 +9,11 @@ WARP_ITK_FLIPS, build_low_pass_filter, compute_hausdorff_distance, + compute_jacobian_determinant, convert_warp, corr2_coeff, create_brain_mask, + displacement_field_to_map, displacement_map_to_field, displacement_maps_to_field_maps, field_maps_to_displacement_maps, @@ -234,16 +236,11 @@ def _zero_3vec_field(shape=(8, 8, 8)) -> nib.Nifti1Image: def test_invert_displacement_field_zero(): - """Inverting a zero field gives back zero (within numerical noise). - - Note: the implementation pads every axis (incl. the 3-channel axis) and - only un-pads the spatial dims, so the output channel axis ends up size 5 - rather than 3. We assert spatial shape match and zero values; fixing the - output-shape quirk is out of scope here. - """ + """Inverting a zero field gives back zero (within numerical noise) and + preserves the 3-channel last axis.""" field = _zero_3vec_field() inverted = invert_displacement_field(field) - assert inverted.shape[:3] == field.shape[:3] + assert inverted.shape == field.shape assert_allclose(inverted.get_fdata(), 0.0, atol=1e-5) @@ -255,6 +252,17 @@ def test_invert_displacement_maps_zero(): assert_allclose(inverted.get_fdata(), 0.0, atol=1e-5) +def test_compute_jacobian_determinant_zero_field(): + """The Jacobian determinant of a zero displacement field is identically 1 + (the identity transform doesn't change volume). Exercises the + compute_jacobian_determinant_cpp binding.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + field = nib.Nifti1Image(np.zeros((6, 6, 6, 3), dtype=np.float32), affine) + jdet = compute_jacobian_determinant(field) + assert jdet.shape == (6, 6, 6) + assert_allclose(jdet.get_fdata(), 1.0, atol=1e-5) + + # --------------------------------------------------------------------------- # Hausdorff distance # --------------------------------------------------------------------------- @@ -330,3 +338,259 @@ def test_convert_warp_rejects_unknown_type(): warp = nib.Nifti1Image(np.zeros((4, 4, 4, 3), dtype=np.float32), affine) with pytest.raises(ValueError, match="not recognized"): convert_warp(warp, in_type="bogus", out_type="itk") + + +def test_convert_warp_accepts_5d_ants_layout(): + """ANTs/AFNI single warps are 5D with shape (X, Y, Z, 1, 3); convert_warp + must accept that layout.""" + affine = np.eye(4) + rng = np.random.default_rng(0) + data_5d = rng.standard_normal((4, 4, 4, 1, 3)).astype(np.float32) + warp = nib.Nifti1Image(data_5d, affine) + out = convert_warp(warp, in_type="ants", out_type="itk") + # ants -> itk flips x and y (WARP_ITK_FLIPS["ants"] = [-1, -1, 1]) + expected = data_5d[..., 0, :].copy() + expected[..., 0] *= -1 + expected[..., 1] *= -1 + assert_allclose(out.get_fdata().squeeze(), expected, atol=1e-5) + + +def test_convert_warp_rejects_bad_5d_shape(): + """A 5D warp without the (X,Y,Z,1,3) layout (e.g. multi-frame 5D series) + must be rejected — the per-frame split should have happened upstream.""" + affine = np.eye(4) + bad = nib.Nifti1Image(np.zeros((4, 4, 4, 2, 3), dtype=np.float32), affine) + with pytest.raises(ValueError, match="singleton 4th axis"): + convert_warp(bad, "itk", "itk") + + +@pytest.mark.parametrize("fmt", ["itk", "fsl", "ants", "afni"]) +@pytest.mark.parametrize("axis", ["x", "y", "z", "x-", "y-", "z-"]) +def test_displacement_map_field_roundtrip_all_formats(axis, fmt): + """displacement_map_to_field followed by displacement_field_to_map must + recover the original 1-channel map, for every (axis, format) pair — + including the ants/afni formats whose intermediate field is 5D.""" + rng = np.random.default_rng(7) + data = rng.standard_normal((4, 4, 4)).astype(np.float32) + dmap = nib.Nifti1Image(data, np.eye(4)) + field = displacement_map_to_field(dmap, axis=axis, format=fmt, frame=0) + back = displacement_field_to_map(field, axis=axis, format=fmt) + assert_allclose(back.get_fdata(), data, atol=1e-5) + + +def test_displacement_field_to_map_clears_vector_intent(): + """The 1-channel result must not carry the input field's vector intent, + or downstream type checks may mistake it for a field.""" + affine = np.eye(4) + data = np.zeros((4, 4, 4), dtype=np.float32) + dmap = nib.Nifti1Image(data, affine) + field = displacement_map_to_field(dmap, axis="y", format="itk", frame=0) + # The intermediate field carries the vector intent (set by convert_warp). + assert field.header.get_intent()[0] == "vector" + back = displacement_field_to_map(field, axis="y", format="itk") + assert back.header.get_intent()[0] != "vector" + + +# --------------------------------------------------------------------------- +# Quantitative Hz <-> mm tests with hand-computed expected values. +# Formula: displacement = fieldmap * total_readout_time * voxel_size, where +# voxel_size carries the LPS-x/y sign flip (axes 0/1) and the negative-PE flip. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "pe_dir, expected_sign", + [ + ("i", -1.0), # axis 0: itk LPS x flip + ("i-", +1.0), # axis 0: "-" cancels itk flip + ("j", -1.0), # axis 1: itk LPS y flip + ("j-", +1.0), # axis 1: "-" cancels itk flip + ("k", +1.0), # axis 2: no LPS flip + ("k-", -1.0), # axis 2: just "-" + ], +) +def test_field_maps_to_displacement_maps_known_value(pe_dir, expected_sign): + """For a constant fmap with isotropic voxels, the displacement equals + fmap * trt * voxel * (sign per PE convention).""" + voxel = 2.0 + affine = np.diag([voxel, voxel, voxel, 1.0]) + fmap_value = 10.0 # Hz + trt = 0.05 # s + fmap_data = np.full((4, 4, 4), fmap_value, dtype=np.float32) + fmap = nib.Nifti1Image(fmap_data, affine) + dmap = field_maps_to_displacement_maps(fmap, trt, pe_dir) + expected = fmap_value * trt * voxel * expected_sign + assert_allclose(dmap.get_fdata(), expected, atol=1e-5) + + +@pytest.mark.parametrize("axis", ["i", "j", "k"]) +def test_field_maps_to_displacement_maps_anisotropic_voxels(axis): + """The voxel size used for scaling must come from the PE-axis dimension, + not a different axis.""" + voxels = (1.0, 2.0, 3.0) + affine = np.diag([*voxels, 1.0]) + fmap = nib.Nifti1Image(np.full((4, 4, 4), 10.0, dtype=np.float32), affine) + dmap = field_maps_to_displacement_maps(fmap, 0.05, axis) + axis_idx = AXIS_MAP[axis] + expected_mag = 10.0 * 0.05 * voxels[axis_idx] + assert_allclose(np.abs(dmap.get_fdata()), expected_mag, atol=1e-5) + + +@pytest.mark.parametrize("axis", ["i", "j", "k"]) +def test_field_maps_to_displacement_maps_pe_sign_negation(axis): + """For the same fmap and voxels, '' and '-' produce results + that are exact negations of each other.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + rng = np.random.default_rng(0) + fmap = nib.Nifti1Image(rng.standard_normal((4, 4, 4)).astype(np.float32), affine) + pos = field_maps_to_displacement_maps(fmap, 0.05, axis) + neg = field_maps_to_displacement_maps(fmap, 0.05, f"{axis}-") + assert_allclose(pos.get_fdata(), -neg.get_fdata(), atol=1e-6) + + +# --------------------------------------------------------------------------- +# displacement_map_to_field frame selection +# --------------------------------------------------------------------------- + + +def test_displacement_map_to_field_frame_selects_correct_frame(): + """A 4D displacement map series + frame=N yields a 4D field whose + PE-axis channel matches frame N of the source.""" + affine = np.eye(4) + rng = np.random.default_rng(0) + map_data = rng.standard_normal((3, 3, 3, 4)).astype(np.float32) + dmap = nib.Nifti1Image(map_data, affine) + for frame in range(map_data.shape[-1]): + field = displacement_map_to_field(dmap, axis="z", format="itk", frame=frame) + # itk format with axis=z (no sign flips) → channel 2 == source frame + assert_allclose(field.get_fdata()[..., 2], map_data[..., frame], atol=1e-6) + + +# --------------------------------------------------------------------------- +# displacement_field_to_map drops off-axis channels (documented behavior) +# --------------------------------------------------------------------------- + + +def test_displacement_field_to_map_drops_off_axis_channels(): + """For an EPI distortion warp all displacement is along the PE axis, so + extracting one channel is exact. Confirm that a field with non-zero values + on every channel has the off-axis values dropped.""" + affine = np.eye(4) + rng = np.random.default_rng(0) + field_data = rng.standard_normal((4, 4, 4, 3)).astype(np.float32) + field = nib.Nifti1Image(field_data, affine) + # itk format = identity → extracted map equals the chosen channel verbatim. + extracted_y = displacement_field_to_map(field, axis="y", format="itk").get_fdata() + assert_allclose(extracted_y, field_data[..., 1], atol=1e-6) + + +# --------------------------------------------------------------------------- +# convert_warp roundtrip across non-canonical orientations +# --------------------------------------------------------------------------- + + +def test_convert_warp_lps_affine_itk_roundtrip(): + """itk -> itk on an image with an LPS affine must preserve data exactly + (the as_reoriented round trip should land back on the original grid).""" + # LPS affine: x and y axes flipped vs RAS + affine = np.diag([-2.0, -2.0, 2.0, 1.0]) + affine[:3, 3] = [10.0, 20.0, -5.0] + rng = np.random.default_rng(0) + data = rng.standard_normal((4, 4, 4, 3)).astype(np.float32) + warp = nib.Nifti1Image(data, affine) + out = convert_warp(warp, in_type="itk", out_type="itk") + assert_allclose(out.get_fdata(), data, atol=1e-5) + # Affine and shape must be preserved on the round trip. + assert out.affine is not None and warp.affine is not None + assert_allclose(out.affine, warp.affine) + assert out.shape == warp.shape + + +# --------------------------------------------------------------------------- +# compute_jacobian_determinant on a non-zero translation field +# --------------------------------------------------------------------------- + + +def test_convert_warp_rejects_unknown_output_type(): + """A bogus out_type triggers the post-flip ValueError (separate from the + in_type unknown branch).""" + affine = np.eye(4) + warp = nib.Nifti1Image(np.zeros((4, 4, 4, 3), dtype=np.float32), affine) + with pytest.raises(ValueError, match="not recognized"): + convert_warp(warp, in_type="itk", out_type="bogus") + + +def test_resample_image_rejects_transform_with_wrong_last_axis(): + """resample_image checks the transform data has size 3 in the last axis + after the optional 5D squeeze.""" + from warpkit.utilities import resample_image + + affine = np.eye(4) + ref = nib.Nifti1Image(np.zeros((4, 4, 4), dtype=np.float32), affine) + inp = nib.Nifti1Image(np.zeros((4, 4, 4), dtype=np.float32), affine) + bad_transform = nib.Nifti1Image(np.zeros((4, 4, 4, 7), dtype=np.float32), affine) + with pytest.raises(ValueError, match="size 3 in last axis"): + resample_image(ref, inp, bad_transform) + + +def test_resample_image_squeezes_5d_transform(): + """A 5D ANTs/AFNI transform (X,Y,Z,1,3) is squeezed to 4D before resample. + Identity transform → output equals input.""" + from warpkit.utilities import resample_image + + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + rng = np.random.default_rng(0) + img_data = rng.random((6, 6, 6), dtype=np.float32) + inp = nib.Nifti1Image(img_data, affine) + ref = nib.Nifti1Image(np.zeros((6, 6, 6), dtype=np.float32), affine) + transform = nib.Nifti1Image(np.zeros((6, 6, 6, 1, 3), dtype=np.float32), affine) + out = resample_image(ref, inp, transform) + assert_allclose(out.get_fdata().squeeze(), img_data, atol=1e-3) + + +def test_create_brain_mask_negative_extra_dilation_erodes(): + """A negative extra_dilation triggers the erosion branch in + create_brain_mask.""" + rng = np.random.default_rng(0) + shape = (16, 16, 16) + coords = np.indices(shape).astype(np.float32) + center = np.array([7.5, 7.5, 7.5]) + r2 = sum((coords[i] - center[i]) ** 2 for i in range(3)) + mag = np.exp(-r2 / 30.0).astype(np.float32) + rng.normal(0, 0.01, shape).astype( + np.float32 + ) + mask_default = create_brain_mask(mag) + mask_eroded = create_brain_mask(mag, extra_dilation=-2) + # Erosion strictly shrinks the mask. + assert mask_eroded.sum() < mask_default.sum() + + +def test_create_brain_mask_extra_dilation_grows_mask(): + """A positive extra_dilation strictly grows the mask.""" + rng = np.random.default_rng(0) + shape = (16, 16, 16) + coords = np.indices(shape).astype(np.float32) + center = np.array([7.5, 7.5, 7.5]) + r2 = sum((coords[i] - center[i]) ** 2 for i in range(3)) + mag = np.exp(-r2 / 30.0).astype(np.float32) + rng.normal(0, 0.01, shape).astype( + np.float32 + ) + mask_default = create_brain_mask(mag) + mask_dilated = create_brain_mask(mag, extra_dilation=3) + assert mask_dilated.sum() > mask_default.sum() + + +def test_compute_jacobian_determinant_constant_translation_is_one(): + """A spatially constant displacement is a pure translation and has unit + Jacobian determinant. Complements the all-zero test by exercising the + non-zero code path.""" + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + field = np.zeros((8, 8, 8, 3), dtype=np.float32) + field[..., 0] = 3.5 # constant x-shift + field[..., 1] = -1.0 # constant y-shift + field[..., 2] = 0.5 # constant z-shift + img = nib.Nifti1Image(field, affine) + jdet = compute_jacobian_determinant(img) + # Trim a 1-voxel border to skip any boundary differencing artifacts. + interior = jdet.get_fdata()[1:-1, 1:-1, 1:-1] + assert_allclose(interior, 1.0, atol=1e-4) diff --git a/warpkit/scripts/_warp_io.py b/warpkit/scripts/_warp_io.py new file mode 100644 index 0000000..17ccca0 --- /dev/null +++ b/warpkit/scripts/_warp_io.py @@ -0,0 +1,126 @@ +"""Shared IO helpers for the warpkit warp-CLI scripts. + +Both ``wk-convert-warp`` and ``wk-compute-jacobian`` accept the same +"1+ files of maps or fields" input model and the same "1 bundled file or N +per-frame files" output model. This module hosts the frame splitting and +bundling helpers that the scripts share. +""" + +from __future__ import annotations + +import argparse +from typing import cast + +import nibabel as nib +import numpy as np + + +def read_input_frames( + input_paths: list[str], + from_type: str, + parser: argparse.ArgumentParser, +) -> list[nib.Nifti1Image]: + """Load input file(s) and split into a flat list of single-frame images. + + ``from_type`` is ``"map"`` (1-channel) or ``"field"`` (3-channel) — the + user-supplied input type from ``--from``. Each input may be a single 3D + map, a 4D map series, a 4D field, or a 5D field (singleton or + multi-frame). The returned frames are 3D for maps and 4D + ``(X, Y, Z, 3)`` for fields. + """ + frames: list[nib.Nifti1Image] = [] + for p in input_paths: + img = cast(nib.Nifti1Image, nib.load(p)) + if from_type == "map": + if img.ndim == 3: + frames.append(img) + elif img.ndim == 4: + for i in range(img.shape[-1]): + fd = np.asarray(img.dataobj[..., i]) + frames.append(nib.Nifti1Image(fd, img.affine, img.header)) + else: + parser.error( + f"map input must be 3D or 4D; got shape {img.shape} for {p}" + ) + else: # field + if img.ndim == 4 and img.shape[-1] == 3: + frames.append(img) + elif img.ndim == 5 and img.shape[-1] == 3: + # ANTs/AFNI single = (X, Y, Z, 1, 3); series = (X, Y, Z, T, 3) + for i in range(img.shape[3]): + fd = np.asarray(img.dataobj[..., i, :]) + frames.append(nib.Nifti1Image(fd, img.affine, img.header)) + else: + parser.error( + "field input must be 4D (X,Y,Z,3) or 5D (X,Y,Z,T,3); " + f"got shape {img.shape} for {p}" + ) + return frames + + +def bundle_frames_to_3d_series(frames: list[nib.Nifti1Image]) -> nib.Nifti1Image: + """Stack 3D scalar frames into a 4D ``(X, Y, Z, T)`` series. + + Used for both 1-channel displacement maps and scalar Jacobian fields. + The frame headers may have inherited a vector intent code from an upstream + field operation; clear it so downstream tools don't misread the bundled + scalar series as a field. + """ + data = np.stack([f.get_fdata() for f in frames], axis=-1).astype(np.float32) + header = cast(nib.Nifti1Header, frames[0].header.copy()) + header.set_intent("none", (), "") + return nib.Nifti1Image(data, frames[0].affine, header) + + +def bundle_frames_to_field_series(frames: list[nib.Nifti1Image]) -> nib.Nifti1Image: + """Stack 4D ``(X, Y, Z, 3)`` field frames into a 5D ``(X, Y, Z, T, 3)`` + series.""" + arrs: list[np.ndarray] = [] + for f in frames: + a = f.get_fdata() + if a.ndim == 5 and a.shape[3] == 1: + a = a[..., 0, :] + arrs.append(a) + data = np.stack(arrs, axis=-2).astype(np.float32) + out = nib.Nifti1Image(data, frames[0].affine, frames[0].header) + cast(nib.Nifti1Header, out.header).set_intent("vector", (), "") + return out + + +def write_output( + frames: list[nib.Nifti1Image], + out_paths: list[str], + out_type: str, + parser: argparse.ArgumentParser, +) -> None: + """Write per-frame images either bundled into one file (when exactly one + output path is given for >1 frames) or one file per frame. + + ``out_type`` is ``"map"`` or ``"field"`` for warp-style outputs; the + Jacobian script passes ``"map"`` since a scalar Jacobian is shaped like + a 1-channel map. + """ + n = len(frames) + n_out = len(out_paths) + if n_out == 1 and n > 1: + bundled = ( + bundle_frames_to_3d_series(frames) + if out_type == "map" + else bundle_frames_to_field_series(frames) + ) + bundled.to_filename(out_paths[0]) + elif n_out == n: + for path, img in zip(out_paths, frames, strict=True): + # Per-frame map outputs may carry a vector intent inherited from an + # upstream field operation — drop it so downstream tools don't + # misread the file as a field. + if out_type == "map": + header = cast(nib.Nifti1Header, img.header.copy()) + header.set_intent("none", (), "") + img = nib.Nifti1Image(np.asarray(img.dataobj), img.affine, header) + img.to_filename(path) + else: + parser.error( + f"got {n_out} --output path(s) for {n} frame(s); must be 1 " + f"(bundle into a single file) or {n} (one per frame)" + ) diff --git a/warpkit/scripts/apply_warp.py b/warpkit/scripts/apply_warp.py new file mode 100644 index 0000000..bc3501b --- /dev/null +++ b/warpkit/scripts/apply_warp.py @@ -0,0 +1,254 @@ +import argparse +from collections.abc import Callable +from typing import cast + +import nibabel as nib +import numpy as np + +from warpkit import __version__ +from warpkit.utilities import ( + AXIS_MAP, + WARP_ITK_FLIPS, + convert_warp, + displacement_map_to_field, + resample_image, + setup_logging, +) + +from . import epilog + + +def _build_transform_getter( + transforms: list[nib.Nifti1Image], + transform_type: str, + phase_encoding_axis: str | None, + in_format: str, + parser: argparse.ArgumentParser, +) -> tuple[int, str, Callable[[int], nib.Nifti1Image]]: + """Validate and wrap the user-supplied transform inputs. + + Returns (frame_count, transform_type, getter). The getter is a callable + that takes a 0-indexed frame number and returns an itk-format + ``Nifti1Image`` ready for ``resample_image``. Single-frame transforms are + cached on first access. + """ + if len(transforms) > 1: + if transform_type != "field": + parser.error( + "--transform-type=map is incompatible with a multi-file " + "--transform series (each file must be a 3-channel field). " + "Pass a single 4D map series or use --transform-type field." + ) + for t in transforms: + if t.ndim != 4 or t.shape[-1] != 3: + parser.error( + "when --transform is a series, each file must be a 4D " + f"3-channel field (X,Y,Z,3); got shape {t.shape}" + ) + cache: list[nib.Nifti1Image | None] = [None] * len(transforms) + + def get_series(i: int) -> nib.Nifti1Image: + entry = cache[i] + if entry is None: + entry = convert_warp(transforms[i], in_type=in_format, out_type="itk") + cache[i] = entry + return entry + + return len(transforms), "field", get_series + + t = transforms[0] + + if transform_type == "map": + if not phase_encoding_axis: + parser.error( + "--phase-encoding-axis is required when the transform is a " + "1-channel displacement map." + ) + axis = cast(str, phase_encoding_axis) + if t.ndim == 3: + cached = displacement_map_to_field(t, axis=axis, format="itk", frame=0) + return 1, "map", lambda _i: cached + if t.ndim == 4: + n = t.shape[-1] + if n == 1: + cached = displacement_map_to_field(t, axis=axis, format="itk", frame=0) + return 1, "map", lambda _i: cached + return ( + n, + "map", + lambda i: displacement_map_to_field( + t, axis=axis, format="itk", frame=i + ), + ) + parser.error( + f"displacement map must be 3D or 4D; got {t.ndim}D shape {t.shape}" + ) + + # field + if t.ndim == 4 and t.shape[-1] == 3: + cached = convert_warp(t, in_type=in_format, out_type="itk") + return 1, "field", lambda _i: cached + if t.ndim == 5 and t.shape[-1] == 3: + if t.shape[3] == 1: + # ANTs / AFNI single-field convention (X,Y,Z,1,3) + cached = convert_warp(t, in_type=in_format, out_type="itk") + return 1, "field", lambda _i: cached + n = t.shape[3] + + def get_5d_frame(i: int) -> nib.Nifti1Image: + frame_data = np.asarray(t.dataobj[..., i, :]) + frame_img = nib.Nifti1Image(frame_data, t.affine, t.header) + return convert_warp(frame_img, in_type=in_format, out_type="itk") + + return n, "field", get_5d_frame + parser.error( + "displacement field must be 4D (X,Y,Z,3) or 5D (X,Y,Z,T,3); " + f"got shape {t.shape}" + ) + + +def main(): + parser = argparse.ArgumentParser( + description=( + "Resample an image through a displacement transform. Supports " + "single-frame and time-series resampling with either 1-channel " + "displacement maps (warpkit/medic output, along a single phase-" + "encoding axis) or 3-channel displacement fields " + "(ITK/FSL/ANTs/AFNI). When the transform has N frames, frame i " + "of the input is resampled with frame i of the transform; a " + "single-frame transform broadcasts across all input frames." + ), + epilog=f"{epilog} 04/24/2026", + ) + parser.add_argument( + "--version", action="version", version=f"%(prog)s {__version__}" + ) + parser.add_argument( + "--input", + required=True, + help="Image to resample (3D or 4D).", + ) + parser.add_argument( + "--reference", + help="Reference 3D grid for the output. Defaults to the input (frame 0 if 4D).", + ) + parser.add_argument( + "--transform", + nargs="+", + required=True, + help=( + "One or more displacement transforms. A single file may be a " + "displacement map (1-channel along --phase-encoding-axis) or a " + "displacement field (3-channel); the type must be declared via " + "--transform-type. Pass multiple 4D fields (X,Y,Z,3) to apply a " + "per-frame field series to a 4D input." + ), + ) + parser.add_argument( + "--output", + required=True, + help="Output NIfTI path.", + ) + parser.add_argument( + "--transform-type", + choices=("map", "field"), + required=True, + help=( + "Transform type. 'map' = 1-channel displacement magnitudes along " + "--phase-encoding-axis. 'field' = 3-channel displacement vectors. " + "Multi-file --transform must be 'field'." + ), + ) + parser.add_argument( + "--phase-encoding-axis", + choices=list(AXIS_MAP), + metavar="AXIS", + help=( + "Axis the 1-channel displacement maps are along (one of: " + f"{', '.join(AXIS_MAP)}). Required when the transform is a " + "displacement map." + ), + ) + parser.add_argument( + "--format", + choices=list(WARP_ITK_FLIPS), + default="itk", + help="Format of the input transform when it is a 3-channel field (default: itk).", + ) + + args = parser.parse_args() + + setup_logging() + print(f"wk-apply-warp: {args}") + + input_img = cast(nib.Nifti1Image, nib.load(args.input)) + transforms = [cast(nib.Nifti1Image, nib.load(p)) for p in args.transform] + + # determine the reference grid + reference_img: nib.Nifti1Image + if args.reference: + reference_img = cast(nib.Nifti1Image, nib.load(args.reference)) + elif input_img.ndim == 3: + reference_img = input_img + else: + reference_img = nib.Nifti1Image( + np.asarray(input_img.dataobj[..., 0]), + input_img.affine, + input_img.header, + ) + + # build a per-frame transform getter + n_transform, transform_type, get_transform = _build_transform_getter( + transforms, + transform_type=args.transform_type, + phase_encoding_axis=args.phase_encoding_axis, + in_format=args.format, + parser=parser, + ) + + # input frame count + if input_img.ndim == 3: + n_input = 1 + elif input_img.ndim == 4: + n_input = input_img.shape[-1] + else: + parser.error( + f"input image must be 3D or 4D; got {input_img.ndim}D shape {input_img.shape}" + ) + + # compatibility checks + if n_transform > 1 and n_input == 1: + parser.error( + f"got a {n_transform}-frame transform but input is 3D; input " + "must be 4D when applying a series of transforms." + ) + if n_transform > 1 and n_transform != n_input: + parser.error( + f"transform has {n_transform} frame(s) but input has {n_input}; " + "they must match (or pass a single-frame transform to broadcast)." + ) + + print( + f" input frames: {n_input}; transform frames: {n_transform} " + f"(type={transform_type})" + ) + + # resample + if n_input == 1: + out_img = resample_image(reference_img, input_img, get_transform(0)) + else: + out_frames = [] + for i in range(n_input): + frame_data = np.asarray(input_img.dataobj[..., i]) + frame_img = nib.Nifti1Image(frame_data, input_img.affine, input_img.header) + t_idx = i if n_transform > 1 else 0 + resampled = resample_image(reference_img, frame_img, get_transform(t_idx)) + out_frames.append(resampled.get_fdata()) + if (i + 1) % 10 == 0 or (i + 1) == n_input: + print(f" resampled frame {i + 1}/{n_input}") + out_data = np.stack(out_frames, axis=-1).astype(np.float32) + out_img = nib.Nifti1Image(out_data, reference_img.affine, reference_img.header) + + print(f"Saving resampled image to {args.output}...") + out_img.to_filename(args.output) + print("Done.") diff --git a/warpkit/scripts/compute_fieldmap.py b/warpkit/scripts/compute_fieldmap.py new file mode 100644 index 0000000..cf0cb19 --- /dev/null +++ b/warpkit/scripts/compute_fieldmap.py @@ -0,0 +1,197 @@ +import argparse +import json +from typing import cast + +import nibabel as nib +import numpy as np + +from warpkit import __version__ +from warpkit.unwrap import compute_field_maps +from warpkit.utilities import ( + displacement_maps_to_field_maps, + field_maps_to_displacement_maps, + invert_displacement_maps, + setup_logging, +) + +from . import epilog + +PE_DIRECTIONS = ("i", "j", "k", "i-", "j-", "k-", "x", "y", "z", "x-", "y-", "z-") + + +def main(): + parser = argparse.ArgumentParser( + description=( + "Compute B0 field maps and EPI distortion-correction displacement " + "maps from previously unwrapped multi-echo phase. The unwrapped " + "phase and masks are the outputs of `wk-unwrap-phase`. This is " + "the post-unwrap half of `wk-medic` and writes the same three " + "NIfTIs: native-space field map (Hz), displacement maps (mm) and " + "undistorted-space field map (Hz)." + ), + epilog=f"{epilog} 04/24/2026", + ) + parser.add_argument( + "--version", action="version", version=f"%(prog)s {__version__}" + ) + parser.add_argument( + "--magnitude", + nargs="+", + required=True, + help="Magnitude data, one per echo (used as regression weights).", + ) + parser.add_argument( + "--unwrapped", + nargs="+", + required=True, + help="Unwrapped phase data, one per echo (output of `wk-unwrap-phase`).", + ) + parser.add_argument( + "--masks", + required=True, + help="Per-frame masks NIfTI (output of `wk-unwrap-phase`).", + ) + parser.add_argument( + "--metadata", + nargs="+", + help=( + "BIDS-style JSON sidecar for each echo. EchoTime (s) is read per " + "file; TotalReadoutTime and PhaseEncodingDirection are taken from " + "the first. Mutually exclusive with --TEs / " + "--total-readout-time / --phase-encoding-direction." + ), + ) + parser.add_argument( + "--TEs", + dest="tes", + nargs="+", + type=float, + help=( + "Echo times in milliseconds, one per echo (must match " + "--unwrapped order). Required unless --metadata is given." + ), + ) + parser.add_argument( + "--total-readout-time", + type=float, + help="Total readout time in seconds. Required unless --metadata is given.", + ) + parser.add_argument( + "--phase-encoding-direction", + choices=PE_DIRECTIONS, + metavar="DIR", + help=( + f"Phase encoding direction (one of: {', '.join(PE_DIRECTIONS)}). " + "Required unless --metadata is given." + ), + ) + parser.add_argument( + "--out-prefix", + required=True, + help="Prefix for output field maps and displacement maps.", + ) + parser.add_argument( + "--border-filt", + nargs=2, + type=int, + default=(1, 5), + metavar=("PASS1", "PASS2"), + help="SVD components for the two-pass border filter (default: 1 5).", + ) + parser.add_argument( + "--svd-filt", + type=int, + default=10, + help="SVD components for global denoising (default: 10).", + ) + parser.add_argument( + "-n", "--n-cpus", type=int, default=4, help="Number of CPUs to use." + ) + + args = parser.parse_args() + + direct_args = { + "--TEs": args.tes, + "--total-readout-time": args.total_readout_time, + "--phase-encoding-direction": args.phase_encoding_direction, + } + direct_supplied = [name for name, val in direct_args.items() if val is not None] + + if args.metadata and direct_supplied: + parser.error( + "--metadata is mutually exclusive with " + f"{', '.join(direct_supplied)}; pass one or the other, not both." + ) + if not args.metadata and len(direct_supplied) != len(direct_args): + missing = [name for name in direct_args if name not in direct_supplied] + parser.error( + "either --metadata or all of --TEs, --total-readout-time, and " + f"--phase-encoding-direction must be provided (missing: {', '.join(missing)})." + ) + + echo_times: list[float] + total_readout_time: float + phase_encoding_direction: str + if args.metadata: + metadatas = [] + for j in args.metadata: + with open(j) as f: + metadatas.append(json.load(f)) + echo_times = [float(m["EchoTime"]) * 1000 for m in metadatas] + total_readout_time = float(metadatas[0]["TotalReadoutTime"]) + phase_encoding_direction = str(metadatas[0]["PhaseEncodingDirection"]) + else: + echo_times = cast(list[float], args.tes) + total_readout_time = cast(float, args.total_readout_time) + phase_encoding_direction = cast(str, args.phase_encoding_direction) + + if len(echo_times) != len(args.unwrapped) or len(echo_times) != len(args.magnitude): + parser.error( + f"got {len(echo_times)} echo time(s), {len(args.unwrapped)} " + f"--unwrapped file(s), and {len(args.magnitude)} --magnitude " + "file(s); all three must match." + ) + + setup_logging() + print(f"wk-compute-fieldmap: {args}") + + mag_imgs = [cast(nib.Nifti1Image, nib.load(m)) for m in args.magnitude] + unwrapped_imgs = [cast(nib.Nifti1Image, nib.load(u)) for u in args.unwrapped] + masks_img = cast(nib.Nifti1Image, nib.load(args.masks)) + + fmaps_native = compute_field_maps( + unwrapped_imgs, + masks_img, + mag_imgs, + echo_times, + border_filt=tuple(args.border_filt), + svd_filt=args.svd_filt, + n_cpus=args.n_cpus, + ) + + # convert native-space field maps to displacement maps in distorted space, + # invert to get distorted -> undistorted, then re-derive an undistorted- + # space field map. Mirrors warpkit.distortion.medic. + inv_displacement_maps = field_maps_to_displacement_maps( + fmaps_native, total_readout_time, phase_encoding_direction + ) + dmaps = invert_displacement_maps(inv_displacement_maps, phase_encoding_direction) + fmaps = displacement_maps_to_field_maps( + dmaps, total_readout_time, phase_encoding_direction, flip_sign=True + ) + + # sign flip if undistorted-space fmap correlates negatively with native + if ( + np.corrcoef( + fmaps.dataobj[..., 0].ravel(), + fmaps_native.dataobj[..., 0].ravel(), + )[0, 1] + < 0 + ): + fmaps = nib.Nifti1Image(fmaps.get_fdata() * -1, fmaps.affine, fmaps.header) + + print("Saving field maps and displacement maps to file...") + fmaps_native.to_filename(f"{args.out_prefix}_fieldmaps_native.nii") + dmaps.to_filename(f"{args.out_prefix}_displacementmaps.nii") + fmaps.to_filename(f"{args.out_prefix}_fieldmaps.nii") + print("Done.") diff --git a/warpkit/scripts/compute_jacobian.py b/warpkit/scripts/compute_jacobian.py new file mode 100644 index 0000000..77736d2 --- /dev/null +++ b/warpkit/scripts/compute_jacobian.py @@ -0,0 +1,127 @@ +import argparse +from typing import cast + +import nibabel as nib + +from warpkit import __version__ +from warpkit.utilities import ( + AXIS_MAP, + WARP_ITK_FLIPS, + compute_jacobian_determinant, + convert_warp, + displacement_map_to_field, + setup_logging, +) + +from . import epilog +from ._warp_io import read_input_frames, write_output + + +def _frame_to_itk_field( + img: nib.Nifti1Image, + in_type: str, + in_format: str, + axis: str | None, +) -> nib.Nifti1Image: + """Coerce a single map or field frame into an itk-format 3-channel field.""" + if in_type == "map": + assert axis is not None + return displacement_map_to_field(img, axis=axis, format="itk", frame=0) + if in_format == "itk": + return img + return cast(nib.Nifti1Image, convert_warp(img, in_type=in_format, out_type="itk")) + + +def main(): + parser = argparse.ArgumentParser( + description=( + "Compute the Jacobian determinant of a displacement warp. The " + "input may be 1-channel displacement maps along --axis, or " + "3-channel displacement fields in any of the ITK / FSL / ANTs / " + "AFNI conventions; single-frame and multi-frame series are both " + "accepted with the same input model as wk-convert-warp. The " + "output is one scalar (Jacobian) volume per input frame: pass " + "one --output path to bundle into a 4D series, or pass N output " + "paths for one file per frame. A Jacobian of 1 means no local " + "volume change; <1 = compression, >1 = expansion." + ), + epilog=f"{epilog} 04/25/2026", + ) + parser.add_argument( + "--version", action="version", version=f"%(prog)s {__version__}" + ) + parser.add_argument( + "--input", + nargs="+", + required=True, + help="Input map(s) or field(s). One or more files; any 4D/5D series is split into frames.", + ) + parser.add_argument( + "--output", + nargs="+", + required=True, + help=( + "Output path(s). Pass a single path to bundle Jacobian volumes " + "into a 4D series, or pass one path per frame." + ), + ) + parser.add_argument( + "--from", + dest="from_type", + choices=("map", "field"), + required=True, + help="Input type: 'map' (1-channel along --axis) or 'field' (3-channel).", + ) + parser.add_argument( + "--from-format", + choices=tuple(WARP_ITK_FLIPS), + default="itk", + help="Input field format (itk/fsl/ants/afni). Used only when --from=field.", + ) + parser.add_argument( + "--axis", + choices=tuple(AXIS_MAP), + metavar="AXIS", + help=( + f"Axis the 1-channel maps are along (one of: {', '.join(AXIS_MAP)}). " + "Required when the input is a displacement map." + ), + ) + parser.add_argument( + "--frame", + type=int, + help="Optional: compute the Jacobian of a single 0-indexed frame from the input.", + ) + + args = parser.parse_args() + setup_logging() + + frames = read_input_frames(args.input, args.from_type, parser) + in_type = args.from_type + + if args.frame is not None: + if args.frame < 0 or args.frame >= len(frames): + parser.error( + f"--frame {args.frame} is out of range; input has " + f"{len(frames)} frame(s)" + ) + frames = [frames[args.frame]] + + if in_type == "map" and not args.axis: + parser.error("--axis is required when the input is a displacement map.") + + in_fmt_label = args.from_format if in_type == "field" else "n/a" + print( + f"wk-compute-jacobian: {len(frames)} frame(s); {in_type}({in_fmt_label}) " + "-> jacobian" + ) + + jacobians: list[nib.Nifti1Image] = [] + for img in frames: + field_itk = _frame_to_itk_field(img, in_type, args.from_format, args.axis) + jacobians.append(compute_jacobian_determinant(field_itk)) + + # Output is a 3D scalar per frame, so use the same packing as 1-channel + # maps (4D series when bundled into a single file). + write_output(jacobians, args.output, "map", parser) + print("Done.") diff --git a/warpkit/scripts/convert_fieldmap.py b/warpkit/scripts/convert_fieldmap.py new file mode 100644 index 0000000..a0cf441 --- /dev/null +++ b/warpkit/scripts/convert_fieldmap.py @@ -0,0 +1,198 @@ +import argparse + +import nibabel as nib + +from warpkit import __version__ +from warpkit.utilities import ( + AXIS_MAP, + WARP_ITK_FLIPS, + displacement_field_to_map, + displacement_map_to_field, + displacement_maps_to_field_maps, + field_maps_to_displacement_maps, + setup_logging, +) + +from . import epilog +from ._warp_io import read_input_frames, write_output + +PE_DIRECTIONS = tuple(AXIS_MAP) + + +def main(): + parser = argparse.ArgumentParser( + description=( + "Convert between mm displacement maps/fields and Hz B0 field " + "maps. The conversion uses the EPI total readout time and the " + "phase-encoding direction (with sign): " + "displacement = fieldmap * total_readout_time * voxel_size. " + "Accepts the same per-frame input model as wk-convert-warp " + "(1+ files, 3D/4D/5D series) and writes one output volume per " + "input frame (bundled into a single file if given one --output, " + "or one per frame if given N)." + ), + epilog=f"{epilog} 04/25/2026", + ) + parser.add_argument( + "--version", action="version", version=f"%(prog)s {__version__}" + ) + parser.add_argument( + "--input", + nargs="+", + required=True, + help=( + "Input file(s). 1-channel (3D/4D) for displacement maps (mm) or " + "B0 field maps (Hz); 4D/5D 3-channel for displacement fields (mm)." + ), + ) + parser.add_argument( + "--output", + nargs="+", + required=True, + help="Output path(s). Pass one for a bundled series, or N for one per frame.", + ) + parser.add_argument( + "--from", + dest="from_type", + choices=("map", "field", "fieldmap"), + required=True, + help=( + "Input type: 'map' (1-channel mm), 'field' (3-channel mm), or " + "'fieldmap' (1-channel Hz)." + ), + ) + parser.add_argument( + "--to", + dest="to_type", + choices=("map", "field", "fieldmap"), + required=True, + help="Output type.", + ) + parser.add_argument( + "--total-readout-time", + type=float, + help="Total readout time in seconds. Required.", + ) + parser.add_argument( + "--phase-encoding-direction", + choices=PE_DIRECTIONS, + metavar="DIR", + help=( + f"Phase encoding direction with sign (one of: " + f"{', '.join(PE_DIRECTIONS)}). Required." + ), + ) + parser.add_argument( + "--from-format", + choices=tuple(WARP_ITK_FLIPS), + default="itk", + help="Input field format. Used only when --from=field.", + ) + parser.add_argument( + "--to-format", + choices=tuple(WARP_ITK_FLIPS), + default="itk", + help="Output field format. Used only when --to=field.", + ) + parser.add_argument( + "--flip-sign", + action="store_true", + help=( + "Flip the sign of the resulting fieldmap (only used when " + "--to=fieldmap). Mirrors the flip-sign branch in " + "warpkit.distortion.medic." + ), + ) + parser.add_argument( + "--frame", + type=int, + help="Optional: convert a single 0-indexed frame from the input.", + ) + + args = parser.parse_args() + setup_logging() + + # _warp_io.read_input_frames knows about map/field only; remap "fieldmap" + # to "map" since on disk a Hz field map is shaped like a 1-channel map. + from_arg_for_io = "map" if args.from_type == "fieldmap" else args.from_type + frames = read_input_frames(args.input, from_arg_for_io, parser) + + if args.frame is not None: + if args.frame < 0 or args.frame >= len(frames): + parser.error( + f"--frame {args.frame} is out of range; input has " + f"{len(frames)} frame(s)" + ) + frames = [frames[args.frame]] + + in_type = args.from_type + + if in_type == args.to_type: + parser.error( + f"--from={in_type} and --to={args.to_type} are the same; use " + "wk-convert-warp for representation/format conversions on the " + "mm side." + ) + + crosses_units = (in_type == "fieldmap") != (args.to_type == "fieldmap") + if not crosses_units: + parser.error( + "wk-convert-fieldmap converts between mm (map/field) and Hz " + "(fieldmap); both --from and --to are on the mm side. Use " + "wk-convert-warp instead." + ) + + if args.total_readout_time is None or not args.phase_encoding_direction: + parser.error( + "--total-readout-time and --phase-encoding-direction are " + "required for mm <-> Hz conversion." + ) + + print( + f"wk-convert-fieldmap: {len(frames)} frame(s); " + f"{in_type} -> {args.to_type} " + f"(trt={args.total_readout_time}s, pe={args.phase_encoding_direction})" + ) + + converted: list[nib.Nifti1Image] = [] + for img in frames: + if args.to_type == "fieldmap": + # mm side -> Hz fieldmap + if in_type == "field": + map_img = displacement_field_to_map( + img, + axis=args.phase_encoding_direction, + format=args.from_format, + ) + else: + map_img = img + converted.append( + displacement_maps_to_field_maps( + map_img, + args.total_readout_time, + args.phase_encoding_direction, + flip_sign=args.flip_sign, + ) + ) + else: + # Hz fieldmap -> mm side + map_img = field_maps_to_displacement_maps( + img, + args.total_readout_time, + args.phase_encoding_direction, + ) + if args.to_type == "field": + converted.append( + displacement_map_to_field( + map_img, + axis=args.phase_encoding_direction, + format=args.to_format, + frame=0, + ) + ) + else: + converted.append(map_img) + + out_writer_type = "field" if args.to_type == "field" else "map" + write_output(converted, args.output, out_writer_type, parser) + print("Done.") diff --git a/warpkit/scripts/convert_warp.py b/warpkit/scripts/convert_warp.py new file mode 100644 index 0000000..58a9dc8 --- /dev/null +++ b/warpkit/scripts/convert_warp.py @@ -0,0 +1,274 @@ +import argparse + +import nibabel as nib +import numpy as np + +from warpkit import __version__ +from warpkit.utilities import ( + AXIS_MAP, + WARP_ITK_FLIPS, + convert_warp, + displacement_field_to_map, + displacement_map_to_field, + invert_displacement_field, + invert_displacement_maps, + setup_logging, +) + +from . import epilog +from ._warp_io import read_input_frames, write_output + + +def _invert_frames( + frames: list[nib.Nifti1Image], + in_type: str, + in_format: str, + axis: str | None, + verbose: bool, +) -> tuple[list[nib.Nifti1Image], str, str]: + """Invert each frame and return ``(frames, post_type, post_format)``. + + Routing is by frame count, not by input type, because the 1D map inverter + is markedly faster per frame than the full 3D field inverter: + + * **Single frame** uses :func:`invert_displacement_field`. A map input is + first promoted to a 3-channel itk field via + :func:`displacement_map_to_field`. The result is always in itk field + form. + * **Multi-frame** stacks all frames into a single 4D ``(X, Y, Z, T)`` map + and runs :func:`invert_displacement_maps` once. A field input first has + its ``axis`` channel extracted (off-axis channels are dropped — fine + for the EPI-distortion case, where displacement is along the + phase-encoding axis). The result is always in 1-channel map form. + + The returned ``post_type`` / ``post_format`` are the actual representation + of the inverted frames (which may differ from the input's), so the + downstream conversion stage can route correctly. + """ + n = len(frames) + if n == 1: + img = frames[0] + if in_type == "map": + assert axis is not None + field_itk = displacement_map_to_field(img, axis=axis, format="itk", frame=0) + else: + field_itk = ( + img + if in_format == "itk" + else convert_warp(img, in_type=in_format, out_type="itk") + ) + return [invert_displacement_field(field_itk, verbose=verbose)], "field", "itk" + + # N > 1: per-frame -> 1-channel maps -> single batched inversion. + assert axis is not None + if in_type == "field": + map_frames = [ + displacement_field_to_map(f, axis=axis, format=in_format) for f in frames + ] + else: + map_frames = frames + template = map_frames[0] + stacked = np.stack([f.get_fdata() for f in map_frames], axis=-1).astype(np.float32) + stacked_img = nib.Nifti1Image(stacked, template.affine, template.header) + inverted_4d = invert_displacement_maps(stacked_img, axis=axis, verbose=verbose) + inv_data = np.asarray(inverted_4d.dataobj) + inverted = [ + nib.Nifti1Image(inv_data[..., i], inverted_4d.affine, inverted_4d.header) + for i in range(n) + ] + return inverted, "map", in_format + + +def _convert_frames( + frames: list[nib.Nifti1Image], + in_type: str, + out_type: str, + in_format: str, + out_format: str, + axis: str | None, +) -> list[nib.Nifti1Image]: + converted: list[nib.Nifti1Image] = [] + for img in frames: + if in_type == "map" and out_type == "map": + converted.append(img) + elif in_type == "field" and out_type == "field": + if in_format == out_format: + converted.append(img) + else: + converted.append( + convert_warp(img, in_type=in_format, out_type=out_format) + ) + elif in_type == "map" and out_type == "field": + assert axis is not None + converted.append( + displacement_map_to_field(img, axis=axis, format=out_format, frame=0) + ) + else: # field -> map + assert axis is not None + converted.append( + displacement_field_to_map(img, axis=axis, format=in_format) + ) + return converted + + +def main(): + parser = argparse.ArgumentParser( + description=( + "Interconvert displacement maps and displacement fields, convert " + "displacement fields between ITK / FSL / ANTs / AFNI format " + "conventions, and (with --invert) invert the warp along the way. " + "A single input file may be a 3D map, a 4D map series, a 4D " + "field, or a 5D field (single or series); multiple input files " + "are read as a flat series. Outputs are either bundled into one " + "file (pass one --output path) or written one-per-frame (pass N " + "output paths). Replaces the older extract_field_from_maps " + "script." + ), + epilog=f"{epilog} 04/25/2026", + ) + parser.add_argument( + "--version", action="version", version=f"%(prog)s {__version__}" + ) + parser.add_argument( + "--input", + nargs="+", + required=True, + help="Input map(s) or field(s). One or more files; any 4D/5D series is split into frames.", + ) + parser.add_argument( + "--output", + nargs="+", + required=True, + help=( + "Output path(s). Pass a single path to bundle all frames into " + "one file (4D for maps, 5D for fields), or pass one path per " + "frame for split output." + ), + ) + parser.add_argument( + "--from", + dest="from_type", + choices=("map", "field"), + required=True, + help="Input type: 'map' (1-channel along --axis) or 'field' (3-channel).", + ) + parser.add_argument( + "--to", + dest="to_type", + choices=("map", "field"), + help="Output type. Defaults to --from (no type conversion).", + ) + parser.add_argument( + "--from-format", + choices=tuple(WARP_ITK_FLIPS), + default="itk", + help="Input field format (itk/fsl/ants/afni). Used only when --from=field.", + ) + parser.add_argument( + "--to-format", + choices=tuple(WARP_ITK_FLIPS), + default="itk", + help="Output field format (itk/fsl/ants/afni). Used only when --to=field.", + ) + parser.add_argument( + "--axis", + choices=tuple(AXIS_MAP), + metavar="AXIS", + help=( + f"Axis the 1-channel maps are along (one of: {', '.join(AXIS_MAP)}). " + "Required when converting between map and field." + ), + ) + parser.add_argument( + "--frame", + type=int, + help="Optional: extract a single 0-indexed frame from the input series.", + ) + parser.add_argument( + "--invert", + action="store_true", + help=( + "Invert each frame before any type/format conversion. Routing is " + "by frame count, not by --from: a single-frame input is inverted " + "with the full 3D field inverter (a map is first promoted to a " + "field via --axis), and a multi-frame input is routed through " + "the per-frame 1D map inverter. For a multi-frame field input, " + "only the --axis component is inverted and off-axis components " + "are dropped (fine for EPI distortion correction, where " + "displacement is along the phase-encoding axis). --axis is " + "required for any map input and for multi-frame inversion; " + "single-frame field inversion does not require --axis, but " + "downstream map output still does." + ), + ) + parser.add_argument( + "--verbose", + action="store_true", + help="Pass-through to the underlying inverter (no effect without --invert).", + ) + + args = parser.parse_args() + setup_logging() + + frames = read_input_frames(args.input, args.from_type, parser) + in_type = args.from_type + out_type = args.to_type or in_type + + if args.frame is not None: + if args.frame < 0 or args.frame >= len(frames): + parser.error( + f"--frame {args.frame} is out of range; input has " + f"{len(frames)} frame(s)" + ) + frames = [frames[args.frame]] + + # --axis is required for map<->field conversion AND for inversion when + # the chosen inversion routing needs it: the single-frame route promotes + # a map to a field (needs axis), and the multi-frame route always runs + # the 1D map inverter (needs axis for map<->axis-channel). + multi_frame = len(frames) > 1 + needs_axis = ( + (in_type == "map" and out_type == "field") + or (in_type == "field" and out_type == "map") + or (args.invert and (in_type == "map" or multi_frame)) + ) + if needs_axis and not args.axis: + parser.error( + "--axis is required when converting between maps and fields, " + "when inverting a single-frame map, or when inverting a " + "multi-frame series (the multi-frame inverter operates along a " + "single axis)." + ) + + in_fmt_label = args.from_format if in_type == "field" else "n/a" + out_fmt_label = args.to_format if out_type == "field" else "n/a" + invert_label = " (inverted)" if args.invert else "" + print( + f"wk-convert-warp: {len(frames)} frame(s); " + f"{in_type}({in_fmt_label}) -> {out_type}({out_fmt_label}){invert_label}" + ) + + # post_type / post_format track the actual representation of the frames + # after the inversion stage, which may differ from the user-declared + # input (e.g. multi-frame field input emerges as 1-channel maps). + post_type, post_format = in_type, args.from_format + if args.invert: + frames, post_type, post_format = _invert_frames( + frames, + in_type=in_type, + in_format=args.from_format, + axis=args.axis, + verbose=args.verbose, + ) + + converted = _convert_frames( + frames, + in_type=post_type, + out_type=out_type, + in_format=post_format, + out_format=args.to_format, + axis=args.axis, + ) + + write_output(converted, args.output, out_type, parser) + print("Done.") diff --git a/warpkit/scripts/extract_field_from_maps.py b/warpkit/scripts/extract_field_from_maps.py deleted file mode 100644 index 2a52b12..0000000 --- a/warpkit/scripts/extract_field_from_maps.py +++ /dev/null @@ -1,63 +0,0 @@ -import argparse -from typing import cast - -import nibabel as nib - -from warpkit import __version__ -from warpkit.utilities import AXIS_MAP, WARP_ITK_FLIPS, displacement_map_to_field - -from . import epilog - - -def main(): - parser = argparse.ArgumentParser( - description="This program extracts a displacement field from a series of displacement maps.", - epilog=f"{epilog} 12/14/2022", - ) - parser.add_argument( - "--version", action="version", version=f"%(prog)s {__version__}" - ) - parser.add_argument("maps", help="Displacement maps to extract field from.") - parser.add_argument("field", help="Displacement field to write out.") - parser.add_argument( - "-n", - "--frame-number", - type=int, - default=0, - help="Frame number to extract field from. 0-indexed. By default 0th frame.", - ) - parser.add_argument( - "-p", - "--phase-encoding-axis", - default="j", - choices=list(AXIS_MAP), - help="The phase encoding axis of the data. Default is j.", - ) - parser.add_argument( - "-f", - "--format", - default="itk", - choices=list(WARP_ITK_FLIPS), - ) - - # parse arguments - args = parser.parse_args() - - # load the displacement maps - maps_img = cast(nib.Nifti1Image, nib.load(args.maps)) - - # grab the map specified by frame_number - selected_map_data = maps_img.dataobj[:, :, :, args.frame_number] - - # make a new nifti image with the selected map - selected_map_img = nib.Nifti1Image( - selected_map_data, maps_img.affine, maps_img.header - ) - - # transform map to field (specified by phase_encoding_axis and file format) - field_img = displacement_map_to_field( - selected_map_img, args.phase_encoding_axis, args.format - ) - - # save the field - field_img.to_filename(args.field) diff --git a/warpkit/scripts/medic.py b/warpkit/scripts/medic.py index b3465f4..9433726 100644 --- a/warpkit/scripts/medic.py +++ b/warpkit/scripts/medic.py @@ -15,7 +15,7 @@ def main(): parser = argparse.ArgumentParser( - description="Multi-Echo DIstortion Correction", epilog=f"{epilog} 12/09/2022" + description="Multi-Echo DIstortion Correction", epilog=f"{epilog}" ) parser.add_argument( "--version", action="version", version=f"%(prog)s {__version__}" diff --git a/warpkit/scripts/unwrap_phase.py b/warpkit/scripts/unwrap_phase.py new file mode 100644 index 0000000..5246eaa --- /dev/null +++ b/warpkit/scripts/unwrap_phase.py @@ -0,0 +1,152 @@ +import argparse +import json +from typing import cast + +import nibabel as nib + +from warpkit import __version__ +from warpkit.unwrap import unwrap_phases +from warpkit.utilities import setup_logging + +from . import epilog + + +def main(): + parser = argparse.ArgumentParser( + description=( + "ROMEO multi-echo phase unwrapping (the unwrap stage of MEDIC). " + "Outputs one unwrapped phase NIfTI per echo plus the per-frame " + "automask. Pair with `wk-compute-fieldmap` to obtain a " + "native-space B0 field map." + ), + epilog=f"{epilog} 04/24/2026", + ) + parser.add_argument( + "--version", action="version", version=f"%(prog)s {__version__}" + ) + parser.add_argument("--magnitude", nargs="+", required=True, help="Magnitude data") + parser.add_argument("--phase", nargs="+", required=True, help="Phase data") + parser.add_argument( + "--metadata", + nargs="+", + help=( + "BIDS-style JSON sidecar for each echo. EchoTime (s) is read per " + "file. Mutually exclusive with --TEs." + ), + ) + parser.add_argument( + "--TEs", + dest="tes", + nargs="+", + type=float, + help=( + "Echo times in milliseconds, one per echo (must match --phase " + "order). Required unless --metadata is given." + ), + ) + parser.add_argument( + "--out-prefix", + required=True, + help="Prefix for output unwrapped phase and mask files.", + ) + parser.add_argument( + "-f", "--noiseframes", type=int, default=0, help="Number of noise frames" + ) + parser.add_argument( + "-n", "--n-cpus", type=int, default=4, help="Number of CPUs to use." + ) + parser.add_argument( + "--debug", + 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() + + if args.metadata and args.tes is not None: + parser.error( + "--metadata is mutually exclusive with --TEs; pass one or the " + "other, not both." + ) + if not args.metadata and args.tes is None: + parser.error("either --metadata or --TEs must be provided.") + + if len(args.magnitude) != len(args.phase): + parser.error( + f"got {len(args.magnitude)} --magnitude file(s) but " + f"{len(args.phase)} --phase file(s); they must match (one " + "mag/phase pair per echo)." + ) + if args.metadata is not None and len(args.metadata) != len(args.phase): + parser.error( + f"got {len(args.metadata)} --metadata file(s) but " + f"{len(args.phase)} --phase file(s); they must match (one " + "sidecar per echo)." + ) + + echo_times: list[float] + if args.metadata: + metadatas = [] + for j in args.metadata: + with open(j) as f: + metadatas.append(json.load(f)) + echo_times = [float(m["EchoTime"]) * 1000 for m in metadatas] + else: + echo_times = cast(list[float], args.tes) + + if len(echo_times) != len(args.phase): + parser.error( + f"got {len(echo_times)} echo time(s) but --phase has " + f"{len(args.phase)} file(s); they must match." + ) + + setup_logging() + print(f"wk-unwrap-phase: {args}") + + mag_data = [cast(nib.Nifti1Image, nib.load(m)) for m in args.magnitude] + phase_data = [cast(nib.Nifti1Image, nib.load(p)) for p in args.phase] + + if args.noiseframes < 0: + parser.error(f"--noiseframes must be non-negative; got {args.noiseframes}.") + if args.noiseframes > 0: + for label, imgs, paths in ( + ("phase", phase_data, args.phase), + ("magnitude", mag_data, args.magnitude), + ): + for img, path in zip(imgs, paths, strict=True): + n_frames = img.shape[-1] if img.ndim == 4 else 1 + if args.noiseframes >= n_frames: + parser.error( + f"--noiseframes={args.noiseframes} would leave 0 " + f"frames in {label} file '{path}' (has {n_frames} " + "frame(s))." + ) + print(f"Removing {args.noiseframes} noise frames from the end of each file...") + mag_data = [ + nib.Nifti1Image(m.dataobj[..., : -args.noiseframes], m.affine, m.header) + for m in mag_data + ] + phase_data = [ + nib.Nifti1Image(p.dataobj[..., : -args.noiseframes], p.affine, p.header) + for p in phase_data + ] + + unwrapped_imgs, masks_img = unwrap_phases( + phase_data, + mag_data, + echo_times, + n_cpus=args.n_cpus, + debug=args.debug, + wrap_limit=args.wrap_limit, + ) + + print("Saving unwrapped phase images and masks to file...") + for i_echo, img in enumerate(unwrapped_imgs, start=1): + img.to_filename(f"{args.out_prefix}_unwrapped_echo-{i_echo:02d}.nii") + masks_img.to_filename(f"{args.out_prefix}_masks.nii") + print("Done.") diff --git a/warpkit/unwrap.py b/warpkit/unwrap.py index 1b90a5f..dea909e 100644 --- a/warpkit/unwrap.py +++ b/warpkit/unwrap.py @@ -23,7 +23,7 @@ get_largest_connected_component, rescale_phase, ) -from .warpkit_cpp import Romeo +from .warpkit_cpp import romeo_unwrap3d, romeo_unwrap4d, romeo_voxelquality FMAP_PROPORTION_HEURISTIC = 0.25 FMAP_AMBIGUIOUS_HEURISTIC = 0.5 @@ -58,9 +58,8 @@ def get_dual_echo_fieldmap(phases, tes, mags, mask): np.ndarray of shape (x, y, z, echo) Unwrapped phases """ - romeo = Romeo() # unwrap the phases - unwrapped_phases = romeo.romeo_unwrap4d( # type: ignore + unwrapped_phases = romeo_unwrap4d( phase=phases, tes=tes, weights="romeo", @@ -115,11 +114,10 @@ def mcpc_3d_s( npt.NDArray[np.float32] Unwrapped difference in phase """ - romeo = Romeo() signal_diff = mag0 * mag1 * np.exp(1j * (phase1 - phase0)) mag_diff = np.abs(signal_diff) phase_diff = np.angle(signal_diff) - unwrapped_diff = romeo.romeo_unwrap3d( # type: ignore + unwrapped_diff = romeo_unwrap3d( phase=phase_diff, weights="romeo", mag=mag_diff, @@ -253,8 +251,6 @@ def unwrap_phase( npt.NDArray[np.int8] mask """ - romeo = Romeo() - if idx is not None: logging.info(f"Processing frame: {idx}") @@ -263,9 +259,9 @@ def unwrap_phase( # the theory goes like this, the magnitude/otsu base mask can be too aggressive occasionally # and the voxel quality mask can get extra voxels that are not brain, but is noisy # so we combine the two masks to get a better mask - vq = romeo.romeo_voxelquality( + vq = romeo_voxelquality( phase_data, tes, np.ones(shape=mag_data.shape, dtype=np.float32) - ) # type: ignore + ) vq_mask = vq > threshold_otsu(vq) strel = generate_binary_structure(3, 2) @@ -326,7 +322,7 @@ def unwrap_phase( phase_data -= phase_offset[..., np.newaxis] # unwrap the phase data - unwrapped = romeo.romeo_unwrap4d( # type: ignore + unwrapped = romeo_unwrap4d( phase=phase_data, tes=tes, weights="romeo", @@ -638,59 +634,58 @@ def svd_filtering( ] -def unwrap_and_compute_field_maps( +def unwrap_phases( phase: list[nib.Nifti1Image], mag: list[nib.Nifti1Image], tes: list[float] | tuple[float] | npt.NDArray[np.float32], mask: nib.Nifti1Image | SimpleNamespace | None = None, automask: bool = True, automask_dilation: int = 3, - border_size: int = 5, - border_filt: tuple[int, int] = (1, 5), - svd_filt: int = 10, frames: list[int] | None = None, n_cpus: int = 4, debug: bool = False, wrap_limit: bool = False, -) -> nib.Nifti1Image: - """Unwrap phase of data weighted by magnitude data and compute field maps. This makes a call - to the ROMEO phase unwrapping algorithm for each frame. To learn more about ROMEO, see this paper: +) -> tuple[list[nib.Nifti1Image], nib.Nifti1Image]: + """Unwrap multi-echo phase per frame and enforce temporal consistency. - Dymerska, B., Eckstein, K., Bachrata, B., Siow, B., Trattnig, S., Shmueli, K., Robinson, S.D., 2020. - Phase Unwrapping with a Rapid Opensource Minimum Spanning TreE AlgOrithm (ROMEO). - Magnetic Resonance in Medicine. https://doi.org/10.1002/mrm.28563 + Calls ROMEO via the warpkit C++ bindings. The returned unwrapped phases are + useful outside of distortion correction (e.g. phase regression, T2* + estimation). Pair with :func:`compute_field_maps` to reconstruct the + native-space B0 field map. Parameters ---------- - phase : List[nib.Nifti1Image] - Phases to unwrap - mag : List[nib.Nifti1Image] - Magnitudes associated with each phase - tes : Union[List[float], Tuple[float], npt.NDArray[np.float32]] - Echo times associated with each phase (in ms) + phase : list[nib.Nifti1Image] + Phase images, one per echo. Each may be 3D or 4D. + mag : list[nib.Nifti1Image] + Magnitude images matched to ``phase``. + tes : list[float] | tuple[float] | npt.NDArray + Echo times in milliseconds, one per echo. mask : nib.Nifti1Image, optional - Boolean mask, by default None + Boolean mask. Ignored when ``automask`` is True. automask : bool, optional - Automatically generate a mask (ignore mask option), by default True + Auto-generate the mask, by default True. automask_dilation : int, optional - Number of dilation iterations applied to the automask, by default 3 - border_size : int, optional - Size of border in automask, by default 5 - border_filt : Tuple[int, int], optional - Number of SVD components for each step of border filtering, by default (1, 5) - svd_filt : int, optional - Number of SVD components to use for filtering of field maps, by default 30 - frames : List[int], optional - Only process these frame indices, by default None (which means all frames) + Dilation iterations for the auto mask, by default 3. + frames : list[int], optional + Subset of frame indices to process, by default None (all frames). n_cpus : int, optional - Number of CPUs to use, by default 4 + CPU parallelism for the unwrap loop, by default 4. debug : bool, optional - Debug mode, by default False + 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 ------- + list[nib.Nifti1Image] + Per-echo unwrapped phase images (radians), each of shape + ``(x, y, z, n_frames)``. nib.Nifti1Image - Field maps in Hz + Per-frame masks, shape ``(x, y, z, n_frames)``, dtype int8. 0 = outside, + 1 = brain core, 2 = dilated border (consumed by SVD filtering in + :func:`compute_field_maps`). """ # check tes if < 0.1, tell user they probably need to convert to ms if np.min(tes) < 0.1: @@ -725,8 +720,6 @@ def unwrap_and_compute_field_maps( # check if data is 4D or 3D if len(phase[0].shape) == 3: - # set total number of frames to 1 - n_frames = 1 # convert data to 4D phase = [ nib.Nifti1Image(p.get_fdata()[..., np.newaxis], p.affine, p.header) @@ -736,12 +729,12 @@ def unwrap_and_compute_field_maps( nib.Nifti1Image(m.get_fdata()[..., np.newaxis], m.affine, m.header) for m in mag ] + if frames is None: + frames = [0] elif len(phase[0].shape) == 4: # if frames is None, set it to all frames if frames is None: frames = list(range(phase[0].shape[-1])) - # get the total number of frames - n_frames = len(frames) else: raise ValueError("Data must be 3D or 4D.") # frames should be a list at this point @@ -753,10 +746,8 @@ def unwrap_and_compute_field_maps( "Number of echo times must equal number of mag and phase images." ) - # allocate space for field maps and unwrapped - field_maps = np.zeros((*phase[0].shape[:3], n_frames), dtype=np.float32) - unwrapped = np.zeros((*phase[0].shape[:3], len(tes), n_frames), dtype=np.float32) - # array for storing auto-generated masks + # allocate space for unwrapped phases and per-frame masks + unwrapped = np.zeros((*phase[0].shape[:3], len(tes), len(frames)), dtype=np.float32) new_masks = np.zeros((*mag[0].shape[:3], len(frames)), dtype=np.int8) # FOR DEBUGGING @@ -885,13 +876,90 @@ def post_temporal_consistency_check(idx, result): "masks.nii" ) - # compute field maps on temporally consistent unwrapped phase - def field_map_iterator(field_maps, unwrapped, mag, tes): + # split the (x, y, z, echo, frame) array into a list of 4D Niftis (one per echo) + unwrapped_imgs = [ + nib.Nifti1Image(unwrapped[:, :, :, i_echo, :], phase[0].affine, phase[0].header) + for i_echo in range(unwrapped.shape[-2]) + ] + masks_img = nib.Nifti1Image(new_masks, mag[0].affine, mag[0].header) + return unwrapped_imgs, masks_img + + +def compute_field_maps( + unwrapped: list[nib.Nifti1Image], + masks: nib.Nifti1Image, + mag: list[nib.Nifti1Image], + tes: list[float] | tuple[float] | npt.NDArray[np.float32], + border_filt: tuple[int, int] = (1, 5), + svd_filt: int = 10, + n_cpus: int = 4, +) -> nib.Nifti1Image: + """Compute native-space B0 field maps from unwrapped phases. + + Performs the second half of MEDIC: weighted echo regression per frame to + yield a field map in Hz, followed by border-aware + global SVD filtering + using the per-frame masks from :func:`unwrap_phases`. + + Parameters + ---------- + unwrapped : list[nib.Nifti1Image] + Per-echo unwrapped phase images (radians), each shape + ``(x, y, z, n_frames)``. Order must match ``mag`` and ``tes``. + masks : nib.Nifti1Image + Per-frame masks, shape ``(x, y, z, n_frames)``, with values 0/1/2 as + produced by :func:`unwrap_phases`. + mag : list[nib.Nifti1Image] + Magnitude images, used as regression weights. + tes : list[float] | tuple[float] | npt.NDArray + Echo times in milliseconds. + border_filt : tuple[int, int], optional + SVD components for the two-pass border filter, by default ``(1, 5)``. + svd_filt : int, optional + SVD components for global denoising, by default 10. + n_cpus : int, optional + CPU parallelism for the field map loop, by default 4. + + Returns + ------- + nib.Nifti1Image + Native-space field map in Hz, shape ``(x, y, z, n_frames)``. + """ + if len(unwrapped) != len(mag) or len(unwrapped) != len(tes): + raise ValueError( + "Number of unwrapped phase images, magnitude images, and echo " + "times must all match." + ) + + tes = cast(npt.NDArray[np.float32], np.array(tes, dtype=np.float32)) + + # stack per-echo 4D unwrapped images into (x, y, z, n_echoes, n_frames) + unwrapped_arr = np.stack( + [np.asarray(u.dataobj, dtype=np.float32) for u in unwrapped], axis=-2 + ) + new_masks = np.asarray(masks.dataobj, dtype=np.int8) + n_frames = unwrapped_arr.shape[-1] + + if ( + new_masks.ndim != 4 + or new_masks.shape[:3] != unwrapped_arr.shape[:3] + or new_masks.shape[-1] != n_frames + ): + raise ValueError( + "masks must have shape (x, y, z, n_frames) matching the spatial " + "dimensions and frame count of the unwrapped data; got masks shape " + f"{new_masks.shape} and expected " + f"({unwrapped_arr.shape[0]}, {unwrapped_arr.shape[1]}, " + f"{unwrapped_arr.shape[2]}, {n_frames})." + ) + + # allocate output + field_maps = np.zeros((*unwrapped_arr.shape[:3], n_frames), dtype=np.float32) + + def field_map_iterator(unwrapped_arr, mag, tes): logging.info("Running field map computation...") - # convert tes to a matrix tes_mat = tes[:, np.newaxis] - for frame_num in range(unwrapped.shape[-1]): - yield (unwrapped[..., frame_num], mag, tes.shape[0], tes_mat, frame_num) + for frame_num in range(unwrapped_arr.shape[-1]): + yield (unwrapped_arr[..., frame_num], mag, tes.shape[0], tes_mat, frame_num) def post_field_map(idx, result): logging.info(f"Field map computation for frame {idx} complete.") @@ -901,22 +969,64 @@ def post_field_map(idx, result): ncpus=n_cpus, type="thread", fn=compute_field_map, - iterator=field_map_iterator(field_maps, unwrapped, mag, tes), + iterator=field_map_iterator(unwrapped_arr, mag, tes), post_fn=post_field_map, ) - # if border voxels defined, use that information to stabilize border regions using SVD filtering - # this will probably kill any respiration signals in these voxels but improve the - # temporal stability of the field maps in these regions (and could we really resolve - # respiration in those voxels any way? probably not...) + # border-aware + global SVD filtering of the field maps svd_filtering( field_maps, new_masks, - phase[0].header.get_zooms()[0], # type: ignore + unwrapped[0].header.get_zooms()[0], # type: ignore n_frames, border_filt, svd_filt, ) - # return the field map as a nifti image - return nib.Nifti1Image(field_maps[..., frames], phase[0].affine, phase[0].header) + return nib.Nifti1Image(field_maps, unwrapped[0].affine, unwrapped[0].header) + + +def unwrap_and_compute_field_maps( + phase: list[nib.Nifti1Image], + mag: list[nib.Nifti1Image], + tes: list[float] | tuple[float] | npt.NDArray[np.float32], + mask: nib.Nifti1Image | SimpleNamespace | None = None, + automask: bool = True, + automask_dilation: int = 3, + border_size: int = 5, + border_filt: tuple[int, int] = (1, 5), + svd_filt: int = 10, + 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. + + Thin wrapper around :func:`unwrap_phases` followed by + :func:`compute_field_maps`. See those functions for parameter and return + semantics. ``border_size`` is accepted for backwards compatibility but is + currently unused. Returns the native-space field map in Hz. + """ + del border_size # accepted for back-compat; not consumed by either stage + unwrapped_imgs, masks_img = unwrap_phases( + phase, + mag, + tes, + mask=mask, + automask=automask, + automask_dilation=automask_dilation, + frames=frames, + n_cpus=n_cpus, + debug=debug, + wrap_limit=wrap_limit, + ) + return compute_field_maps( + unwrapped_imgs, + masks_img, + mag, + tes, + border_filt=border_filt, + svd_filt=svd_filt, + n_cpus=n_cpus, + ) diff --git a/warpkit/utilities.py b/warpkit/utilities.py index dd78702..8b4c4ab 100644 --- a/warpkit/utilities.py +++ b/warpkit/utilities.py @@ -390,6 +390,54 @@ def displacement_map_to_field( return convert_warp(warp, in_type="itk", out_type=format) +def displacement_field_to_map( + displacement_field: nib.Nifti1Image, + axis: str = "y", + format: str = "itk", +) -> nib.Nifti1Image: + """Extract a 1-channel displacement map (along ``axis``) from a 3-channel + displacement field. Inverse of :func:`displacement_map_to_field`. + + The off-axis channels are dropped — for an EPI distortion-correction warp + this is exact because all displacement is along the phase-encoding axis. + + Parameters + ---------- + displacement_field : nib.Nifti1Image + Displacement field, in ``format`` convention. 4D (X, Y, Z, 3) or 5D + (X, Y, Z, 1, 3) with a singleton 4th axis (ANTs/AFNI single-field). + axis : str, optional + Axis to extract along, by default "y". + format : str, optional + Format of the input field (one of ``itk``, ``fsl``, ``ants``, + ``afni``), by default ``itk``. + + Returns + ------- + nib.Nifti1Image + 1-channel displacement map, shape ``(X, Y, Z)``. + """ + axis_code = AXIS_MAP[axis] + # Round-trip through itk so we can index the channel axis directly. + field_itk = convert_warp(displacement_field, in_type=format, out_type="itk") + data = field_itk.get_fdata() + if data.ndim == 5 and data.shape[3] == 1: + data = data[..., 0, :] + if data.ndim != 4 or data.shape[-1] != 3: + raise ValueError( + f"expected a 4D 3-channel field after itk conversion; got shape {data.shape}" + ) + map_data = data[..., axis_code] + # Drop the vector intent so downstream tools don't misread the 1-channel + # result as a field. + map_header = cast(nib.Nifti1Header, field_itk.header.copy()) + map_header.set_intent("none", (), "") + return cast( + nib.Nifti1Image, + nib.Nifti1Image(map_data, field_itk.affine, map_header), + ) + + def get_x_orient_transform( img: nib.Nifti1Image, x: str ) -> tuple[Sequence[Sequence[int]], Sequence[Sequence[int]]]: @@ -535,13 +583,16 @@ def invert_displacement_field( # split affine into components translations, rotations, zooms, _ = decompose44(displacement_field_ras.affine) - # pad array with edge values so edge effects of inverse are avoided - mod_data = np.pad(data, pad_width=1) + # Pad spatial dims only — leaving the 3-channel axis at size 3 — so we + # can avoid edge effects of the inverse without inflating the channel + # count. (Earlier versions padded all 4 axes which produced a 5-channel + # output.) + mod_data = np.pad(data, ((1, 1), (1, 1), (1, 1), (0, 0))) # invert displacement field new_data = invert_displacement_field_cpp( mod_data, translations, rotations, zooms, verbose=verbose - )[1 : data.shape[0] + 1, 1 : data.shape[1] + 1, 1 : data.shape[2] + 1] + )[1 : data.shape[0] + 1, 1 : data.shape[1] + 1, 1 : data.shape[2] + 1, :] # make new image inv_displacement_field = nib.Nifti1Image( @@ -645,19 +696,18 @@ def convert_warp( nib.Nifti1Image Output warp in desired format. """ - # check data shape - if len(in_warp.shape) != 4: - if len(in_warp.shape) != 5: - raise ValueError("Input warp must be 4D or 5D.") - else: - if in_warp.shape[3] == 1 and in_warp.shape[-1]: - raise ValueError( - "Input warp must have singleton dimension in 4th axis and size in last axis." - ) - else: - # check last axis size + # check data shape: 4D (X,Y,Z,3) or 5D (X,Y,Z,1,3) (ANTs/AFNI single warp). + if in_warp.ndim == 4: if in_warp.shape[-1] != 3: raise ValueError("Warp must have size 3 in last axis.") + elif in_warp.ndim == 5: + if in_warp.shape[3] != 1 or in_warp.shape[-1] != 3: + raise ValueError( + "5D warp must have singleton 4th axis and size 3 in last axis " + f"(got shape {tuple(in_warp.shape)})." + ) + else: + raise ValueError("Input warp must be 4D or 5D.") # get input in RAS orientation to_canonical, from_canonical = get_ras_orient_transform(in_warp) diff --git a/warpkit/warpkit_cpp.pyi b/warpkit/warpkit_cpp.pyi index d76a9b1..b1ee43a 100644 --- a/warpkit/warpkit_cpp.pyi +++ b/warpkit/warpkit_cpp.pyi @@ -6,65 +6,27 @@ import numpy import numpy.typing __all__: list[str] = [ - "Romeo", + "calculate_weights", "compute_hausdorff_distance", "compute_jacobian_determinant", "invert_displacement_field", "invert_displacement_map", "resample", + "romeo_unwrap3d", + "romeo_unwrap4d", + "romeo_voxelquality", ] -class Romeo: - def __init__(self) -> None: ... - def calculate_weights( - self, - phase: typing.Annotated[numpy.typing.ArrayLike, numpy.float32], - mag: typing.Annotated[numpy.typing.ArrayLike, numpy.float32] = ..., - phase2: typing.Annotated[numpy.typing.ArrayLike, numpy.float32] = ..., - tes: typing.Annotated[numpy.typing.ArrayLike, numpy.float32] = ..., - mask: typing.Annotated[numpy.typing.ArrayLike, numpy.bool_] = ..., - ) -> numpy.typing.NDArray[numpy.uint8]: - """ - ROMEO edge-weight map (3, nx, ny, nz) uint8. Exposed for port validation; not used by warpkit. - """ - def romeo_unwrap3d( - self, - phase: typing.Annotated[numpy.typing.ArrayLike, numpy.float32], - weights: str, - mag: typing.Annotated[numpy.typing.ArrayLike, numpy.float32], - mask: typing.Annotated[numpy.typing.ArrayLike, numpy.bool_], - correct_global: bool = True, - maxseeds: typing.SupportsInt | typing.SupportsIndex = 1, - merge_regions: bool = False, - correct_regions: bool = False, - ) -> numpy.typing.NDArray[numpy.float32]: - """ - 3D ROMEO phase unwrap - """ - def romeo_unwrap4d( - self, - phase: typing.Annotated[numpy.typing.ArrayLike, numpy.float32], - tes: typing.Annotated[numpy.typing.ArrayLike, numpy.float32], - weights: str, - mag: typing.Annotated[numpy.typing.ArrayLike, numpy.float32], - mask: typing.Annotated[numpy.typing.ArrayLike, numpy.bool_], - correct_global: bool = True, - maxseeds: typing.SupportsInt | typing.SupportsIndex = 1, - merge_regions: bool = False, - correct_regions: bool = False, - ) -> numpy.typing.NDArray[numpy.float32]: - """ - 4D (multi-echo) ROMEO phase unwrap - """ - def romeo_voxelquality( - self, - phase: typing.Annotated[numpy.typing.ArrayLike, numpy.float32], - tes: typing.Annotated[numpy.typing.ArrayLike, numpy.float32], - mag: typing.Annotated[numpy.typing.ArrayLike, numpy.float32], - ) -> numpy.typing.NDArray[numpy.float32]: - """ - Compute a per-voxel quality map from multi-echo phase/magnitude - """ +def calculate_weights( + phase: typing.Annotated[numpy.typing.ArrayLike, numpy.float32], + mag: typing.Annotated[numpy.typing.ArrayLike, numpy.float32] = ..., + phase2: typing.Annotated[numpy.typing.ArrayLike, numpy.float32] = ..., + tes: typing.Annotated[numpy.typing.ArrayLike, numpy.float32] = ..., + mask: typing.Annotated[numpy.typing.ArrayLike, numpy.bool_] = ..., +) -> numpy.typing.NDArray[numpy.uint8]: + """ + ROMEO edge-weight map (3, nx, ny, nz) uint8. Exposed for port validation; not used by warpkit. + """ def compute_hausdorff_distance( image1: typing.Annotated[numpy.typing.ArrayLike, numpy.float64], @@ -132,3 +94,41 @@ def resample( """ Resample an image with transform """ + +def romeo_unwrap3d( + phase: typing.Annotated[numpy.typing.ArrayLike, numpy.float32], + weights: str, + mag: typing.Annotated[numpy.typing.ArrayLike, numpy.float32], + mask: typing.Annotated[numpy.typing.ArrayLike, numpy.bool_], + correct_global: bool = True, + maxseeds: typing.SupportsInt | typing.SupportsIndex = 1, + merge_regions: bool = False, + correct_regions: bool = False, +) -> numpy.typing.NDArray[numpy.float32]: + """ + 3D ROMEO phase unwrap + """ + +def romeo_unwrap4d( + phase: typing.Annotated[numpy.typing.ArrayLike, numpy.float32], + tes: typing.Annotated[numpy.typing.ArrayLike, numpy.float32], + weights: str, + mag: typing.Annotated[numpy.typing.ArrayLike, numpy.float32], + mask: typing.Annotated[numpy.typing.ArrayLike, numpy.bool_], + correct_global: bool = True, + maxseeds: typing.SupportsInt | typing.SupportsIndex = 1, + merge_regions: bool = False, + correct_regions: bool = False, +) -> numpy.typing.NDArray[numpy.float32]: + """ + 4D (multi-echo) ROMEO phase unwrap + """ + +def romeo_voxelquality( + phase: typing.Annotated[numpy.typing.ArrayLike, numpy.float32], + tes: typing.Annotated[numpy.typing.ArrayLike, numpy.float32], + mag: typing.Annotated[numpy.typing.ArrayLike, numpy.float32], +) -> numpy.typing.NDArray[numpy.float32]: + """ + Compute a per-voxel quality map from multi-echo phase/magnitude + """