From 09d83adef8aac5845ac1ba4f6bfdbbc43e6dd657 Mon Sep 17 00:00:00 2001 From: Andrew Van Date: Sat, 25 Apr 2026 01:01:57 -0500 Subject: [PATCH 1/8] :sparkles: Add wk-* CLI suite for warp pipeline operations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Expand warpkit's CLI surface from `medic` + `extract_field_from_maps` to six purpose-built tools, all `wk-` prefixed to avoid collisions with namesakes from FSL/ANTs/AFNI/etc.: * `wk-medic` — full pipeline (renamed from `medic`). * `wk-unwrap-phase` + `wk-compute-fieldmap` — the same pipeline split into reusable stages, addresses #12 (unwrapped phase as a first-class output usable for phase regression, T2* estimation, etc.). * `wk-apply-warp` — resample any 3D/4D image through a displacement map/field, single or per-frame; auto-detects map vs field from intent + shape, with an override flag. * `wk-convert-warp` — interconverts maps <-> fields, ITK/FSL/ANTs/AFNI format conversion, single-frame extraction, and (with `--invert`) inverts the warp along the way (1D inverter for series, full 3D inverter for single frames). Replaces `extract_field_from_maps`. * `wk-compute-jacobian` — per-voxel Jacobian determinant for any warp. Behind the scenes: * Refactored `unwrap_and_compute_field_maps` into reusable `unwrap_phases` + `compute_field_maps` stages; the original wrapper is kept and `medic()` still calls it, so the Python API is unchanged. * Replaced the `Romeo` C++ class facade with free functions in `warpkit_cpp` (`romeo_unwrap3d`, `romeo_unwrap4d`, `romeo_voxelquality`, `calculate_weights`); pyi regenerated via `scripts/regen-stub.sh`. * Fixed a pre-existing bug in `invert_displacement_field` that returned a 5-channel array (`np.pad` was padding the channel axis); padding is now spatial-only and the output is the expected (X, Y, Z, 3). * Added `displacement_field_to_map`, the inverse of `displacement_map_to_field`. * Extracted shared CLI IO helpers (frame splitting, classification, bundle/per-frame write) into a private `warpkit/scripts/_warp_io.py`. CI-relevant: * Removed the script auto-discovery glob in `setup.py`; entry points are now explicit in `[project.scripts]` in `pyproject.toml`. No matrix change. * Dockerfile entrypoint switched from `medic` -> `wk-medic`. * Test count 80 -> 111; coverage 74% -> 90%. Co-Authored-By: Claude Opus 4.7 (1M context) --- CLAUDE.md | 29 +- Dockerfile | 6 +- README.md | 58 +- include/romeo/romeo.h | 119 +-- pyproject.toml | 10 +- setup.py | 8 - src/warpkit.cpp | 44 +- tests/conftest.py | 18 +- tests/test_romeo.py | 42 +- tests/test_scripts.py | 1099 +++++++++++++++++++- tests/test_unwrap.py | 29 +- tests/test_utilities.py | 23 +- warpkit/scripts/_warp_io.py | 140 +++ warpkit/scripts/apply_warp.py | 263 +++++ warpkit/scripts/compute_fieldmap.py | 197 ++++ warpkit/scripts/compute_jacobian.py | 126 +++ warpkit/scripts/convert_warp.py | 267 +++++ warpkit/scripts/extract_field_from_maps.py | 63 -- warpkit/scripts/medic.py | 2 +- warpkit/scripts/unwrap_phase.py | 125 +++ warpkit/unwrap.py | 219 ++-- warpkit/utilities.py | 53 +- warpkit/warpkit_cpp.pyi | 104 +- 23 files changed, 2669 insertions(+), 375 deletions(-) create mode 100644 warpkit/scripts/_warp_io.py create mode 100644 warpkit/scripts/apply_warp.py create mode 100644 warpkit/scripts/compute_fieldmap.py create mode 100644 warpkit/scripts/compute_jacobian.py create mode 100644 warpkit/scripts/convert_warp.py delete mode 100644 warpkit/scripts/extract_field_from_maps.py create mode 100644 warpkit/scripts/unwrap_phase.py diff --git a/CLAUDE.md b/CLAUDE.md index becee05..334ff2c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -15,8 +15,15 @@ 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-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 `wk-convert-warp` and `wk-compute-jacobian` 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 +47,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 +116,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..3bf6b21 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,27 @@ 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-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 +131,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 +142,40 @@ 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 +``` ## 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..71da173 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,14 @@ 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-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..e86335a 100644 --- a/tests/test_scripts.py +++ b/tests/test_scripts.py @@ -1,19 +1,35 @@ """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_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 +40,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 +67,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 +83,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 +105,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 +129,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 +167,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 +189,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 +212,1061 @@ def test_medic_phase_encoding_direction_choices_enforced(argv, capsys): # --------------------------------------------------------------------------- -# extract_field_from_maps --help / choices +# wk-convert-warp --help / argument validation +# --------------------------------------------------------------------------- + + +def test_convert_warp_help(argv, capsys): + argv(["wk-convert-warp", "--help"]) + with pytest.raises(SystemExit) as exc: + convert_warp_main() + assert exc.value.code == 0 + out = capsys.readouterr().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_convert_warp_rejects_bad_format(argv, capsys, tmp_path): + field = _write_nifti(tmp_path / "field.nii", (4, 4, 4, 3)) + argv( + [ + "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: + 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"), + "--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"), + "--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"), + "--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 + ] + ) + 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"), + "--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"), + "--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 + + +def test_convert_warp_rejects_mixed_input_types(argv, capsys, tmp_path): + """All inputs must classify to the same map/field type.""" + maps = _write_nifti(tmp_path / "maps.nii", (4, 4, 4, 5)) # auto -> map + field = _write_nifti(tmp_path / "field.nii", (4, 4, 4, 3)) # auto -> field + argv( + [ + "wk-convert-warp", + "--input", + maps, + field, + "--output", + str(tmp_path / "out.nii"), + ] + ) + with pytest.raises(SystemExit) as exc: + convert_warp_main() + assert exc.value.code == 2 + err = capsys.readouterr().err + assert "mixed map/field" 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 + + +# --------------------------------------------------------------------------- +# 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_extract_field_from_maps_help(argv, capsys): - argv(["extract_field_from_maps", "--help"]) +def test_apply_warp_help(argv, capsys): + argv(["wk-apply-warp", "--help"]) with pytest.raises(SystemExit) as exc: - extract_main() + apply_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 "Resample" in out + assert "--input" in out + assert "--transform" in out assert "--phase-encoding-axis" in out -def test_extract_field_from_maps_rejects_bad_axis(argv, capsys): +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 (last dim != 3 ensures auto-classifies as 'map') + tx = _write_nifti(tmp_path / "tx.nii", (4, 4, 4, 5)) argv( [ - "extract_field_from_maps", - "maps.nii.gz", - "out.nii.gz", + "wk-apply-warp", + "--input", + inp, + "--transform", + tx, + "--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: - extract_main() + 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, + "--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, + "--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, + "--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 (will auto-classify as '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), + "--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 with the "vector" intent set so the + # auto-classifier uses the intent path (not the shape fallback). + 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), + "--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, + "--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), + "--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), + "--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), + "--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), + "--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"), + ] + ) + 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-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"), + "--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), + ] + ) + 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), + "--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. Uses 4 frames so the + last-dim-3 heuristic doesn't auto-classify the input as a single field.""" + 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, + "--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) diff --git a/tests/test_unwrap.py b/tests/test_unwrap.py index a52e961..6cee3a2 100644 --- a/tests/test_unwrap.py +++ b/tests/test_unwrap.py @@ -68,31 +68,32 @@ 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 + import warpkit.warpkit_cpp as cpp - 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") + # 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 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..3fb887d 100644 --- a/tests/test_utilities.py +++ b/tests/test_utilities.py @@ -9,6 +9,7 @@ WARP_ITK_FLIPS, build_low_pass_filter, compute_hausdorff_distance, + compute_jacobian_determinant, convert_warp, corr2_coeff, create_brain_mask, @@ -234,16 +235,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 +251,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 # --------------------------------------------------------------------------- diff --git a/warpkit/scripts/_warp_io.py b/warpkit/scripts/_warp_io.py new file mode 100644 index 0000000..f556d50 --- /dev/null +++ b/warpkit/scripts/_warp_io.py @@ -0,0 +1,140 @@ +"""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 classification, 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 classify(img: nib.Nifti1Image, override: str) -> str: + """Classify a NIfTI as a 1-channel map or a 3-channel field. + + ``override`` may be ``"auto"``, ``"map"``, or ``"field"``. The auto path + uses the NIfTI ``intent_code`` (``"vector"`` → field) and shape (5D or + 4D-with-last==3 → field, else map). + """ + if override != "auto": + return override + intent = img.header.get_intent() if img.header is not None else (None,) + if intent and intent[0] == "vector": + return "field" + if img.ndim == 5: + return "field" + if img.ndim == 4 and img.shape[-1] == 3: + return "field" + return "map" + + +def read_input_frames( + input_paths: list[str], + from_override: str, + parser: argparse.ArgumentParser, +) -> tuple[list[nib.Nifti1Image], str]: + """Load input file(s) and split into a flat list of single-frame images. + + 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. All inputs must classify to the + same type. + """ + type_choices: list[str] = [] + frames: list[nib.Nifti1Image] = [] + for p in input_paths: + img = cast(nib.Nifti1Image, nib.load(p)) + ftype = classify(img, from_override) + type_choices.append(ftype) + if ftype == "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}" + ) + if len(set(type_choices)) > 1: + parser.error( + f"mixed map/field inputs: {type_choices}; all input files must " + "classify as the same type (use --from to override)" + ) + return frames, type_choices[0] + + +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. + """ + data = np.stack([f.get_fdata() for f in frames], axis=-1).astype(np.float32) + return nib.Nifti1Image(data, frames[0].affine, frames[0].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): + 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..723c2ce --- /dev/null +++ b/warpkit/scripts/apply_warp.py @@ -0,0 +1,263 @@ +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 _classify_single_transform(img: nib.Nifti1Image, override: str) -> str: + """Classify a single-file transform as 'map' (1-channel) or 'field' (3-channel).""" + if override != "auto": + return override + intent = img.header.get_intent() if img.header is not None else (None,) + if intent and intent[0] == "vector": + return "field" + if img.ndim == 5: + return "field" + if img.ndim == 4 and img.shape[-1] == 3: + return "field" + return "map" + + +def _build_transform_getter( + transforms: list[nib.Nifti1Image], + transform_type_override: 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, classified_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: + 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] + transform_type = _classify_single_transform(t, transform_type_override) + + 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 is auto-" + "classified as displacement maps (1-channel) or a displacement " + "field (3-channel); 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=("auto", "map", "field"), + default="auto", + help=( + "Override the single-file classifier. 'map' = 1-channel " + "displacement magnitudes along --phase-encoding-axis. 'field' = " + "3-channel displacement vectors. Ignored when --transform has " + "more than one file (always treated as a field series)." + ), + ) + 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_override=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..57b3bed --- /dev/null +++ b/warpkit/scripts/compute_jacobian.py @@ -0,0 +1,126 @@ +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=("auto", "map", "field"), + default="auto", + help="Input type. Default 'auto' uses the NIfTI intent code and shape.", + ) + 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, in_type = read_input_frames(args.input, args.from_type, 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]] + + 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_warp.py b/warpkit/scripts/convert_warp.py new file mode 100644 index 0000000..4942679 --- /dev/null +++ b/warpkit/scripts/convert_warp.py @@ -0,0 +1,267 @@ +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=("auto", "map", "field"), + default="auto", + help="Input type. Default 'auto' uses the NIfTI intent code and shape.", + ) + parser.add_argument( + "--to", + dest="to_type", + choices=("map", "field"), + help="Output type. Defaults to whatever the input classifies as.", + ) + 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. Maps are " + "inverted with the 1D map inverter along --axis; fields are " + "inverted with the full 3D field inverter. Inversion of maps " + "requires --axis (the map's own axis); inversion of fields does " + "not, 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, in_type = read_input_frames(args.input, args.from_type, parser) + 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..becd5cd --- /dev/null +++ b/warpkit/scripts/unwrap_phase.py @@ -0,0 +1,125 @@ +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.") + + 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: + 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..e39ff37 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,77 @@ 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] + + # 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 +956,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..ee7b88c 100644 --- a/warpkit/utilities.py +++ b/warpkit/utilities.py @@ -390,6 +390,50 @@ 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] + return cast( + nib.Nifti1Image, + nib.Nifti1Image(map_data, field_itk.affine, field_itk.header), + ) + + def get_x_orient_transform( img: nib.Nifti1Image, x: str ) -> tuple[Sequence[Sequence[int]], Sequence[Sequence[int]]]: @@ -535,13 +579,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( 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 + """ From af8549dfaed28b74aac5cad9235f2f651e646fd7 Mon Sep 17 00:00:00 2001 From: Andrew Van Date: Sat, 25 Apr 2026 01:13:20 -0500 Subject: [PATCH 2/8] :sparkles: Add wk-convert-fieldmap (mm <-> Hz unit conversion) Adds a seventh wk-* CLI: convert between mm displacement maps/fields and Hz B0 field maps in either direction, using the same per-frame input model as `wk-convert-warp`. * Auto-classifies 1-channel inputs by inference from `--to`: when going to a fieldmap the input is mm displacement; when going to a map/field the input is a Hz fieldmap. * On the mm side, accepts either 1-channel maps or 3-channel fields; field inputs have the PE-axis channel extracted before unit conversion, and field outputs are promoted from the mm map back to 3 channels. * `--total-readout-time` and `--phase-encoding-direction` are required (the conversion is `displacement = fieldmap * trt * voxel_size_along_PE`); `--flip-sign` mirrors the flip-sign branch in `warpkit.distortion.medic`. * Refuses mm <-> mm conversions and same-`--from`-as-`--to` calls, pointing users at `wk-convert-warp` for those cases. Also: * Reuses `_warp_io.read_input_frames` and `write_output` for input splitting and bundle/split output (no new IO code). * 9 new tests cover help, all the validation paths, the mm-Hz-mm round-trip (preserved within 1e-5), and the field-input axis extraction. Total tests 111 -> 120; coverage holds at 90%. Co-Authored-By: Claude Opus 4.7 (1M context) --- CLAUDE.md | 9 +- README.md | 13 ++ pyproject.toml | 1 + tests/test_scripts.py | 260 ++++++++++++++++++++++++++++ warpkit/scripts/convert_fieldmap.py | 227 ++++++++++++++++++++++++ 5 files changed, 506 insertions(+), 4 deletions(-) create mode 100644 warpkit/scripts/convert_fieldmap.py diff --git a/CLAUDE.md b/CLAUDE.md index 334ff2c..e34e47b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -19,10 +19,11 @@ Pre-print: . 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-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 `wk-convert-warp` and `wk-compute-jacobian` live in + `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]`. diff --git a/README.md b/README.md index 3bf6b21..f5601d9 100644 --- a/README.md +++ b/README.md @@ -109,6 +109,7 @@ with same-named tools from FSL/ANTs/AFNI/etc.: | `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` @@ -177,6 +178,18 @@ wk-compute-jacobian \ --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 Vahdeta Suljic <suljic@wustl.edu>, Andrew Van <vanandrew77@gmail.com> diff --git a/pyproject.toml b/pyproject.toml index 71da173..f677180 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ dependencies = [ "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] diff --git a/tests/test_scripts.py b/tests/test_scripts.py index e86335a..deb96c3 100644 --- a/tests/test_scripts.py +++ b/tests/test_scripts.py @@ -21,6 +21,7 @@ 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 @@ -1270,3 +1271,262 @@ def test_compute_jacobian_per_frame_outputs(argv, tmp_path): 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"), + ] + ) + 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"), + "--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"), + "--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). Also asserts that the + 'auto' classifier picks 'map' for 1-channel input + --to=fieldmap.""" + 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), + "--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), + "--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), + "--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 (auto-classify 1-channel as + fieldmap because --to is on the mm side).""" + 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), + "--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) diff --git a/warpkit/scripts/convert_fieldmap.py b/warpkit/scripts/convert_fieldmap.py new file mode 100644 index 0000000..9ee764d --- /dev/null +++ b/warpkit/scripts/convert_fieldmap.py @@ -0,0 +1,227 @@ +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 _resolve_in_type( + from_arg: str, + shape_type: str, + to_type: str, + parser: argparse.ArgumentParser, +) -> str: + """Disambiguate the input type given ``--from``, the shape-based + classification (``"map"`` for 1-channel, ``"field"`` for 3-channel), + and ``--to``.""" + if from_arg != "auto": + if from_arg == "field" and shape_type != "field": + parser.error( + "--from=field expects a 3-channel input but got 1-channel " + "data; pass --from map or --from fieldmap instead." + ) + if from_arg in ("map", "fieldmap") and shape_type != "map": + parser.error( + f"--from={from_arg} expects a 1-channel input but got " + "3-channel field data; pass --from field instead." + ) + return from_arg + if shape_type == "field": + return "field" + # 1-channel input: 'map' (mm) and 'fieldmap' (Hz) look the same on + # disk, so disambiguate by which side --to is on. + return "map" if to_type == "fieldmap" else "fieldmap" + + +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=("auto", "map", "field", "fieldmap"), + default="auto", + help=( + "Input type. 'auto' uses the shape (3-channel = field) plus " + "--to to disambiguate map (mm) from fieldmap (Hz) for " + "1-channel inputs." + ), + ) + 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() + + # The shape-based classifier in _warp_io knows about map/field only; + # remap "fieldmap" to "map" for the splitter's purposes. + from_arg_for_io = "map" if args.from_type == "fieldmap" else args.from_type + frames, shape_type = 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 = _resolve_in_type(args.from_type, shape_type, args.to_type, parser) + + 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.") From 4b8da4893177677be0358e485b9d8511c52d1654 Mon Sep 17 00:00:00 2001 From: Andrew Van Date: Sat, 25 Apr 2026 01:44:27 -0500 Subject: [PATCH 3/8] :bug: Fix convert_warp shape check rejecting valid 5D ANTs/AFNI warps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The 5D branch had an inverted condition: it raised when the input had the canonical ANTs/AFNI single-warp layout (X, Y, Z, 1, 3) — exactly the shape the error message described as required. Rewrite the validation to accept 4D (X, Y, Z, 3) and 5D (X, Y, Z, 1, 3) and reject everything else. The CLIs masked this because read_input_frames splits 5D files into 4D frames upstream, so it only fired on direct Python-API use — including the documented displacement_map_to_field -> displacement_field_to_map roundtrip for the ants/afni formats, which produce a 5D intermediate. Add regression tests covering 5D-accepting, 5D-rejecting, and the full (axis, format) roundtrip matrix for the map<->field pair. Co-Authored-By: Claude Opus 4.7 (1M context) --- tests/test_utilities.py | 39 +++++++++++++++++++++++++++++++++++++++ warpkit/utilities.py | 21 ++++++++++----------- 2 files changed, 49 insertions(+), 11 deletions(-) diff --git a/tests/test_utilities.py b/tests/test_utilities.py index 3fb887d..c1adaeb 100644 --- a/tests/test_utilities.py +++ b/tests/test_utilities.py @@ -13,6 +13,7 @@ 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, @@ -337,3 +338,41 @@ 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) diff --git a/warpkit/utilities.py b/warpkit/utilities.py index ee7b88c..c235a8e 100644 --- a/warpkit/utilities.py +++ b/warpkit/utilities.py @@ -692,19 +692,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) From b039c78c019c5ece11568b7a8828b9eae7a84f1b Mon Sep 17 00:00:00 2001 From: Andrew Van Date: Sat, 25 Apr 2026 01:49:35 -0500 Subject: [PATCH 4/8] :bug: Address PR review: clear stale vector intent + tighter unwrap-phase validation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three intent-leak fixes — a 1-channel map output that inherits a vector intent from an upstream field operation gets re-classified as a field by the auto-classifier, leading to surprising behavior: * displacement_field_to_map copies the intermediate field header verbatim. Copy and reset to "none" so the scalar map is classified as a map. * bundle_frames_to_3d_series stacks scalar frames but reuses the first frame's header. Same fix. * write_output's per-frame branch wrote each frame untouched. For out_type=="map", clear the intent before writing. wk-unwrap-phase now checks --magnitude and --metadata counts match --phase up front, so users get a clean argparse error instead of a downstream zip(strict=True) ValueError. The count check was moved before the JSON sidecar load so a count mismatch isn't masked by a FileNotFoundError. Co-Authored-By: Claude Opus 4.7 (1M context) --- tests/test_scripts.py | 86 +++++++++++++++++++++++++++++++++ tests/test_utilities.py | 13 +++++ warpkit/scripts/_warp_io.py | 14 +++++- warpkit/scripts/unwrap_phase.py | 13 +++++ warpkit/utilities.py | 6 ++- 5 files changed, 130 insertions(+), 2 deletions(-) diff --git a/tests/test_scripts.py b/tests/test_scripts.py index deb96c3..4ed323a 100644 --- a/tests/test_scripts.py +++ b/tests/test_scripts.py @@ -507,6 +507,54 @@ def test_unwrap_phase_te_count_must_match_phase_count(argv, capsys, tmp_path): 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 + + # --------------------------------------------------------------------------- # compute_fieldmap --help / argument validation # --------------------------------------------------------------------------- @@ -1530,3 +1578,41 @@ def test_convert_fieldmap_fieldmap_to_field(argv, tmp_path): 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" diff --git a/tests/test_utilities.py b/tests/test_utilities.py index c1adaeb..4f191a1 100644 --- a/tests/test_utilities.py +++ b/tests/test_utilities.py @@ -376,3 +376,16 @@ def test_displacement_map_field_roundtrip_all_formats(axis, fmt): 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 auto-classifiers will 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" diff --git a/warpkit/scripts/_warp_io.py b/warpkit/scripts/_warp_io.py index f556d50..d1e77b9 100644 --- a/warpkit/scripts/_warp_io.py +++ b/warpkit/scripts/_warp_io.py @@ -88,9 +88,14 @@ 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 the bundled scalar series isn't later + misclassified as a field. """ data = np.stack([f.get_fdata() for f in frames], axis=-1).astype(np.float32) - return nib.Nifti1Image(data, frames[0].affine, frames[0].header) + 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: @@ -132,6 +137,13 @@ def write_output( 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 the file isn't later + # auto-classified 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( diff --git a/warpkit/scripts/unwrap_phase.py b/warpkit/scripts/unwrap_phase.py index becd5cd..2578a56 100644 --- a/warpkit/scripts/unwrap_phase.py +++ b/warpkit/scripts/unwrap_phase.py @@ -76,6 +76,19 @@ def main(): 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 = [] diff --git a/warpkit/utilities.py b/warpkit/utilities.py index c235a8e..41f1c5b 100644 --- a/warpkit/utilities.py +++ b/warpkit/utilities.py @@ -428,9 +428,13 @@ def displacement_field_to_map( f"expected a 4D 3-channel field after itk conversion; got shape {data.shape}" ) map_data = data[..., axis_code] + # Drop the vector intent so the 1-channel result isn't later auto-classified + # as a field by intent-based heuristics. + 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, field_itk.header), + nib.Nifti1Image(map_data, field_itk.affine, map_header), ) From 9164bcb00a0e8afb952f123f1e67bfe24ccfc362 Mon Sep 17 00:00:00 2001 From: Andrew Van Date: Sat, 25 Apr 2026 02:02:03 -0500 Subject: [PATCH 5/8] :fire: Drop --from=auto and --transform-type=auto across CLIs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Auto-detection of map vs. field from NIfTI shape + intent was a hidden source of surprise — a stale vector intent code (or a 4D shape with last==3) would silently flip a map to a field. The user is now required to declare the input type up front: * wk-convert-warp / wk-convert-fieldmap / wk-compute-jacobian: --from is required, choices are {map, field} (plus 'fieldmap' for convert-fieldmap). * wk-apply-warp: --transform-type is required, {map, field}. Multi-file --transform now requires --transform-type=field (was implicit). * _warp_io.classify is gone (only ever served the auto path); the shape-mismatch check inside read_input_frames remains as the user's declared type can still be inconsistent with the actual file. * convert_fieldmap._resolve_in_type collapses to a no-op now that --from is always explicit; the consistency check moved into the shape-split branch of read_input_frames. Test churn is mechanical (every CLI invocation gains an explicit --from or --transform-type). The mixed-input-type test for convert_warp is removed since one --from value now applies to all inputs by construction. Co-Authored-By: Claude Opus 4.7 (1M context) --- tests/test_scripts.py | 97 ++++++++++++++++++++--------- tests/test_utilities.py | 2 +- warpkit/scripts/_warp_io.py | 56 +++++------------ warpkit/scripts/apply_warp.py | 47 ++++++-------- warpkit/scripts/compute_jacobian.py | 9 +-- warpkit/scripts/convert_fieldmap.py | 45 +++---------- warpkit/scripts/convert_warp.py | 9 +-- warpkit/utilities.py | 4 +- 8 files changed, 121 insertions(+), 148 deletions(-) diff --git a/tests/test_scripts.py b/tests/test_scripts.py index 4ed323a..6749e43 100644 --- a/tests/test_scripts.py +++ b/tests/test_scripts.py @@ -285,6 +285,8 @@ def test_convert_warp_requires_axis_for_map_to_field(argv, capsys, tmp_path): maps, "--output", str(tmp_path / "out.nii"), + "--from", + "map", "--to", "field", ] @@ -306,6 +308,8 @@ def test_convert_warp_requires_axis_for_field_to_map(argv, capsys, tmp_path): field, "--output", str(tmp_path / "out.nii"), + "--from", + "field", "--to", "map", ] @@ -326,6 +330,8 @@ def test_convert_warp_frame_out_of_range(argv, capsys, tmp_path): maps, "--output", str(tmp_path / "out.nii"), + "--from", + "map", "--frame", "10", ] @@ -348,6 +354,8 @@ def test_convert_warp_output_count_mismatch(argv, capsys, tmp_path): "--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: @@ -368,6 +376,8 @@ def test_convert_warp_invert_single_map_requires_axis(argv, capsys, tmp_path): single_map, "--output", str(tmp_path / "out.nii"), + "--from", + "map", "--invert", # no --axis, no --to=field ] @@ -391,6 +401,8 @@ def test_convert_warp_invert_multi_frame_field_requires_axis(argv, capsys, tmp_p *fields, "--output", str(tmp_path / "out.nii"), + "--from", + "field", "--invert", # no --axis ] @@ -403,27 +415,6 @@ def test_convert_warp_invert_multi_frame_field_requires_axis(argv, capsys, tmp_p assert "multi-frame" in err -def test_convert_warp_rejects_mixed_input_types(argv, capsys, tmp_path): - """All inputs must classify to the same map/field type.""" - maps = _write_nifti(tmp_path / "maps.nii", (4, 4, 4, 5)) # auto -> map - field = _write_nifti(tmp_path / "field.nii", (4, 4, 4, 3)) # auto -> field - argv( - [ - "wk-convert-warp", - "--input", - maps, - field, - "--output", - str(tmp_path / "out.nii"), - ] - ) - with pytest.raises(SystemExit) as exc: - convert_warp_main() - assert exc.value.code == 2 - err = capsys.readouterr().err - assert "mixed map/field" in err - - # --------------------------------------------------------------------------- # unwrap_phase --help / argument validation # --------------------------------------------------------------------------- @@ -705,7 +696,7 @@ def test_apply_warp_help(argv, capsys): 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 (last dim != 3 ensures auto-classifies as 'map') + # 4D 1-channel map series tx = _write_nifti(tmp_path / "tx.nii", (4, 4, 4, 5)) argv( [ @@ -714,6 +705,8 @@ def test_apply_warp_requires_phase_encoding_axis_for_map(argv, capsys, tmp_path) inp, "--transform", tx, + "--transform-type", + "map", "--output", str(tmp_path / "out.nii"), ] @@ -782,6 +775,8 @@ def test_apply_warp_3d_input_with_series_transform_errors(argv, capsys, tmp_path inp, "--transform", tx, + "--transform-type", + "map", "--phase-encoding-axis", "j", "--output", @@ -806,6 +801,8 @@ def test_apply_warp_frame_count_mismatch_errors(argv, capsys, tmp_path): inp, "--transform", tx, + "--transform-type", + "map", "--phase-encoding-axis", "j", "--output", @@ -833,6 +830,8 @@ def test_apply_warp_series_files_must_be_3channel(argv, capsys, tmp_path): "--transform", bad1, bad2, + "--transform-type", + "field", "--output", str(tmp_path / "out.nii"), ] @@ -958,7 +957,7 @@ def test_apply_warp_happy_path_zero_displacement(argv, tmp_path): in_path = tmp_path / "img.nii" nib.Nifti1Image(img_data, affine).to_filename(str(in_path)) - # a 4D 1-channel zero displacement map (will auto-classify as 'map') + # 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)) @@ -971,6 +970,8 @@ def test_apply_warp_happy_path_zero_displacement(argv, tmp_path): str(in_path), "--transform", str(map_path), + "--transform-type", + "map", "--phase-encoding-axis", "j", "--output", @@ -997,8 +998,7 @@ def test_apply_warp_happy_path_4d_image_with_zero_field(argv, tmp_path): in_path = tmp_path / "img.nii" nib.Nifti1Image(img_data, affine).to_filename(str(in_path)) - # single 4D 3-channel zero field with the "vector" intent set so the - # auto-classifier uses the intent path (not the shape fallback). + # 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", (), "") @@ -1013,6 +1013,8 @@ def test_apply_warp_happy_path_4d_image_with_zero_field(argv, tmp_path): str(in_path), "--transform", str(field_path), + "--transform-type", + "field", "--output", str(out_path), ] @@ -1048,6 +1050,8 @@ def test_apply_warp_happy_path_field_series(argv, tmp_path): str(in_path), "--transform", *field_paths, + "--transform-type", + "field", "--output", str(out_path), ] @@ -1076,6 +1080,8 @@ def test_convert_warp_happy_path_map_field_roundtrip(argv, tmp_path): str(map_path), "--output", str(field_path), + "--from", + "map", "--to", "field", "--axis", @@ -1096,6 +1102,8 @@ def test_convert_warp_happy_path_map_field_roundtrip(argv, tmp_path): str(field_path), "--output", str(back_path), + "--from", + "field", "--to", "map", "--axis", @@ -1125,6 +1133,8 @@ def test_convert_warp_happy_path_format_conversion(argv, tmp_path): str(in_path), "--output", str(out_path), + "--from", + "field", "--to-format", "ants", ] @@ -1156,6 +1166,8 @@ def test_convert_warp_happy_path_invert_zero_map(argv, tmp_path): str(in_path), "--output", str(out_path), + "--from", + "map", "--invert", "--axis", "j", @@ -1194,6 +1206,8 @@ def test_compute_jacobian_requires_axis_for_map(argv, capsys, tmp_path): maps, "--output", str(tmp_path / "out.nii"), + "--from", + "map", ] ) with pytest.raises(SystemExit) as exc: @@ -1212,6 +1226,8 @@ def test_compute_jacobian_rejects_bad_format(argv, capsys, tmp_path): field, "--output", str(tmp_path / "out.nii"), + "--from", + "field", "--from-format", "matlab", # not a valid choice ] @@ -1232,6 +1248,8 @@ def test_compute_jacobian_frame_out_of_range(argv, capsys, tmp_path): maps, "--output", str(tmp_path / "out.nii"), + "--from", + "map", "--axis", "j", "--frame", @@ -1260,6 +1278,8 @@ def test_compute_jacobian_single_field_zero(argv, tmp_path): str(field_path), "--output", str(out_path), + "--from", + "field", ] ) compute_jacobian_main() @@ -1284,6 +1304,8 @@ def test_compute_jacobian_map_series_zero(argv, tmp_path): str(map_path), "--output", str(out_path), + "--from", + "map", "--axis", "j", ] @@ -1295,8 +1317,7 @@ def test_compute_jacobian_map_series_zero(argv, tmp_path): def test_compute_jacobian_per_frame_outputs(argv, tmp_path): - """N output paths → one Jacobian volume per frame. Uses 4 frames so the - last-dim-3 heuristic doesn't auto-classify the input as a single field.""" + """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" @@ -1310,6 +1331,8 @@ def test_compute_jacobian_per_frame_outputs(argv, tmp_path): str(map_path), "--output", *out_paths, + "--from", + "map", "--axis", "j", ] @@ -1349,6 +1372,8 @@ def test_convert_fieldmap_requires_to(argv, capsys, tmp_path): fmap, "--output", str(tmp_path / "out.nii"), + "--from", + "fieldmap", ] ) with pytest.raises(SystemExit) as exc: @@ -1367,6 +1392,8 @@ def test_convert_fieldmap_rejects_bad_pe_direction(argv, capsys, tmp_path): fmap, "--output", str(tmp_path / "out.nii"), + "--from", + "fieldmap", "--to", "map", "--total-readout-time", @@ -1391,6 +1418,8 @@ def test_convert_fieldmap_requires_trt_and_pe(argv, capsys, tmp_path): fmap, "--output", str(tmp_path / "out.nii"), + "--from", + "fieldmap", "--to", "map", ] @@ -1459,8 +1488,7 @@ def test_convert_fieldmap_rejects_same_from_to(argv, capsys, tmp_path): 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). Also asserts that the - 'auto' classifier picks 'map' for 1-channel input + --to=fieldmap.""" + 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 @@ -1475,6 +1503,8 @@ def test_convert_fieldmap_map_to_fieldmap_roundtrip(argv, tmp_path): str(map_path), "--output", str(fmap_path), + "--from", + "map", "--to", "fieldmap", "--total-readout-time", @@ -1496,6 +1526,8 @@ def test_convert_fieldmap_map_to_fieldmap_roundtrip(argv, tmp_path): str(fmap_path), "--output", str(back_path), + "--from", + "fieldmap", "--to", "map", "--total-readout-time", @@ -1528,6 +1560,8 @@ def test_convert_fieldmap_field_to_fieldmap(argv, tmp_path): str(field_path), "--output", str(out_path), + "--from", + "field", "--to", "fieldmap", "--total-readout-time", @@ -1546,8 +1580,7 @@ def test_convert_fieldmap_field_to_fieldmap(argv, tmp_path): def test_convert_fieldmap_fieldmap_to_field(argv, tmp_path): - """Hz fieldmap -> mm 3-channel field (auto-classify 1-channel as - fieldmap because --to is on the mm side).""" + """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 @@ -1562,6 +1595,8 @@ def test_convert_fieldmap_fieldmap_to_field(argv, tmp_path): str(fmap_path), "--output", str(field_path), + "--from", + "fieldmap", "--to", "field", "--total-readout-time", diff --git a/tests/test_utilities.py b/tests/test_utilities.py index 4f191a1..425d410 100644 --- a/tests/test_utilities.py +++ b/tests/test_utilities.py @@ -380,7 +380,7 @@ def test_displacement_map_field_roundtrip_all_formats(axis, fmt): def test_displacement_field_to_map_clears_vector_intent(): """The 1-channel result must not carry the input field's vector intent, - or downstream auto-classifiers will mistake it for a field.""" + 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) diff --git a/warpkit/scripts/_warp_io.py b/warpkit/scripts/_warp_io.py index d1e77b9..17ccca0 100644 --- a/warpkit/scripts/_warp_io.py +++ b/warpkit/scripts/_warp_io.py @@ -2,8 +2,8 @@ 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 classification, frame -splitting, and bundling helpers that the scripts share. +per-frame files" output model. This module hosts the frame splitting and +bundling helpers that the scripts share. """ from __future__ import annotations @@ -15,44 +15,23 @@ import numpy as np -def classify(img: nib.Nifti1Image, override: str) -> str: - """Classify a NIfTI as a 1-channel map or a 3-channel field. - - ``override`` may be ``"auto"``, ``"map"``, or ``"field"``. The auto path - uses the NIfTI ``intent_code`` (``"vector"`` → field) and shape (5D or - 4D-with-last==3 → field, else map). - """ - if override != "auto": - return override - intent = img.header.get_intent() if img.header is not None else (None,) - if intent and intent[0] == "vector": - return "field" - if img.ndim == 5: - return "field" - if img.ndim == 4 and img.shape[-1] == 3: - return "field" - return "map" - - def read_input_frames( input_paths: list[str], - from_override: str, + from_type: str, parser: argparse.ArgumentParser, -) -> tuple[list[nib.Nifti1Image], str]: +) -> list[nib.Nifti1Image]: """Load input file(s) and split into a flat list of single-frame images. - 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. All inputs must classify to the - same type. + ``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. """ - type_choices: list[str] = [] frames: list[nib.Nifti1Image] = [] for p in input_paths: img = cast(nib.Nifti1Image, nib.load(p)) - ftype = classify(img, from_override) - type_choices.append(ftype) - if ftype == "map": + if from_type == "map": if img.ndim == 3: frames.append(img) elif img.ndim == 4: @@ -76,12 +55,7 @@ def read_input_frames( "field input must be 4D (X,Y,Z,3) or 5D (X,Y,Z,T,3); " f"got shape {img.shape} for {p}" ) - if len(set(type_choices)) > 1: - parser.error( - f"mixed map/field inputs: {type_choices}; all input files must " - "classify as the same type (use --from to override)" - ) - return frames, type_choices[0] + return frames def bundle_frames_to_3d_series(frames: list[nib.Nifti1Image]) -> nib.Nifti1Image: @@ -89,8 +63,8 @@ def bundle_frames_to_3d_series(frames: list[nib.Nifti1Image]) -> nib.Nifti1Image 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 the bundled scalar series isn't later - misclassified as a field. + 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()) @@ -138,8 +112,8 @@ def write_output( 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 the file isn't later - # auto-classified as a field. + # 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", (), "") diff --git a/warpkit/scripts/apply_warp.py b/warpkit/scripts/apply_warp.py index 723c2ce..bc3501b 100644 --- a/warpkit/scripts/apply_warp.py +++ b/warpkit/scripts/apply_warp.py @@ -18,35 +18,27 @@ from . import epilog -def _classify_single_transform(img: nib.Nifti1Image, override: str) -> str: - """Classify a single-file transform as 'map' (1-channel) or 'field' (3-channel).""" - if override != "auto": - return override - intent = img.header.get_intent() if img.header is not None else (None,) - if intent and intent[0] == "vector": - return "field" - if img.ndim == 5: - return "field" - if img.ndim == 4 and img.shape[-1] == 3: - return "field" - return "map" - - def _build_transform_getter( transforms: list[nib.Nifti1Image], - transform_type_override: str, + 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, classified_type, getter). The getter is a callable + 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( @@ -65,7 +57,6 @@ def get_series(i: int) -> nib.Nifti1Image: return len(transforms), "field", get_series t = transforms[0] - transform_type = _classify_single_transform(t, transform_type_override) if transform_type == "map": if not phase_encoding_axis: @@ -146,10 +137,11 @@ def main(): nargs="+", required=True, help=( - "One or more displacement transforms. A single file is auto-" - "classified as displacement maps (1-channel) or a displacement " - "field (3-channel); pass multiple 4D fields (X,Y,Z,3) to apply " - "a per-frame field series to a 4D input." + "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( @@ -159,13 +151,12 @@ def main(): ) parser.add_argument( "--transform-type", - choices=("auto", "map", "field"), - default="auto", + choices=("map", "field"), + required=True, help=( - "Override the single-file classifier. 'map' = 1-channel " - "displacement magnitudes along --phase-encoding-axis. 'field' = " - "3-channel displacement vectors. Ignored when --transform has " - "more than one file (always treated as a field series)." + "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( @@ -209,7 +200,7 @@ def main(): # build a per-frame transform getter n_transform, transform_type, get_transform = _build_transform_getter( transforms, - transform_type_override=args.transform_type, + transform_type=args.transform_type, phase_encoding_axis=args.phase_encoding_axis, in_format=args.format, parser=parser, diff --git a/warpkit/scripts/compute_jacobian.py b/warpkit/scripts/compute_jacobian.py index 57b3bed..77736d2 100644 --- a/warpkit/scripts/compute_jacobian.py +++ b/warpkit/scripts/compute_jacobian.py @@ -68,9 +68,9 @@ def main(): parser.add_argument( "--from", dest="from_type", - choices=("auto", "map", "field"), - default="auto", - help="Input type. Default 'auto' uses the NIfTI intent code and shape.", + choices=("map", "field"), + required=True, + help="Input type: 'map' (1-channel along --axis) or 'field' (3-channel).", ) parser.add_argument( "--from-format", @@ -96,7 +96,8 @@ def main(): args = parser.parse_args() setup_logging() - frames, in_type = read_input_frames(args.input, args.from_type, parser) + 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): diff --git a/warpkit/scripts/convert_fieldmap.py b/warpkit/scripts/convert_fieldmap.py index 9ee764d..a0cf441 100644 --- a/warpkit/scripts/convert_fieldmap.py +++ b/warpkit/scripts/convert_fieldmap.py @@ -19,34 +19,6 @@ PE_DIRECTIONS = tuple(AXIS_MAP) -def _resolve_in_type( - from_arg: str, - shape_type: str, - to_type: str, - parser: argparse.ArgumentParser, -) -> str: - """Disambiguate the input type given ``--from``, the shape-based - classification (``"map"`` for 1-channel, ``"field"`` for 3-channel), - and ``--to``.""" - if from_arg != "auto": - if from_arg == "field" and shape_type != "field": - parser.error( - "--from=field expects a 3-channel input but got 1-channel " - "data; pass --from map or --from fieldmap instead." - ) - if from_arg in ("map", "fieldmap") and shape_type != "map": - parser.error( - f"--from={from_arg} expects a 1-channel input but got " - "3-channel field data; pass --from field instead." - ) - return from_arg - if shape_type == "field": - return "field" - # 1-channel input: 'map' (mm) and 'fieldmap' (Hz) look the same on - # disk, so disambiguate by which side --to is on. - return "map" if to_type == "fieldmap" else "fieldmap" - - def main(): parser = argparse.ArgumentParser( description=( @@ -82,12 +54,11 @@ def main(): parser.add_argument( "--from", dest="from_type", - choices=("auto", "map", "field", "fieldmap"), - default="auto", + choices=("map", "field", "fieldmap"), + required=True, help=( - "Input type. 'auto' uses the shape (3-channel = field) plus " - "--to to disambiguate map (mm) from fieldmap (Hz) for " - "1-channel inputs." + "Input type: 'map' (1-channel mm), 'field' (3-channel mm), or " + "'fieldmap' (1-channel Hz)." ), ) parser.add_argument( @@ -141,10 +112,10 @@ def main(): args = parser.parse_args() setup_logging() - # The shape-based classifier in _warp_io knows about map/field only; - # remap "fieldmap" to "map" for the splitter's purposes. + # _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, shape_type = read_input_frames(args.input, from_arg_for_io, parser) + 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): @@ -154,7 +125,7 @@ def main(): ) frames = [frames[args.frame]] - in_type = _resolve_in_type(args.from_type, shape_type, args.to_type, parser) + in_type = args.from_type if in_type == args.to_type: parser.error( diff --git a/warpkit/scripts/convert_warp.py b/warpkit/scripts/convert_warp.py index 4942679..2a457fc 100644 --- a/warpkit/scripts/convert_warp.py +++ b/warpkit/scripts/convert_warp.py @@ -148,9 +148,9 @@ def main(): parser.add_argument( "--from", dest="from_type", - choices=("auto", "map", "field"), - default="auto", - help="Input type. Default 'auto' uses the NIfTI intent code and shape.", + choices=("map", "field"), + required=True, + help="Input type: 'map' (1-channel along --axis) or 'field' (3-channel).", ) parser.add_argument( "--to", @@ -204,7 +204,8 @@ def main(): args = parser.parse_args() setup_logging() - frames, in_type = read_input_frames(args.input, args.from_type, parser) + 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: diff --git a/warpkit/utilities.py b/warpkit/utilities.py index 41f1c5b..8b4c4ab 100644 --- a/warpkit/utilities.py +++ b/warpkit/utilities.py @@ -428,8 +428,8 @@ def displacement_field_to_map( f"expected a 4D 3-channel field after itk conversion; got shape {data.shape}" ) map_data = data[..., axis_code] - # Drop the vector intent so the 1-channel result isn't later auto-classified - # as a field by intent-based heuristics. + # 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( From a3a8f4bd93f12ec8a26e6c4fdeeeeb88bfa15978 Mon Sep 17 00:00:00 2001 From: Andrew Van Date: Sat, 25 Apr 2026 02:08:19 -0500 Subject: [PATCH 6/8] :white_check_mark: Add comprehensive conversion tests Closes coverage gaps left by the auto-removal cleanup. The mm<->Hz path had only roundtrip and flip-sign coverage; format / orientation / inversion paths only had zero-input or roundtrip coverage. New tests target absolute values, sign conventions, and CLI-level behavior. In tests/test_utilities.py: * Hand-computed fmap -> dmap value, parametrized over all 6 PE codes (locks the LPS-x/y vs z sign convention). * Anisotropic-voxel scaling (verifies the right axis's voxel is used). * PE-sign symmetry: '' and '-' produce exact negations. * displacement_map_to_field frame=N selection uses the right frame. * displacement_field_to_map drops off-axis channels (asserted directly). * convert_warp itk->itk roundtrip on an LPS affine (covers the as_reoriented round trip). * compute_jacobian_determinant of a constant translation = 1 (the non-zero counterpart to the existing zero-field test). In tests/test_scripts.py: * wk-convert-warp --frame extracts the right frame; per-frame outputs preserve frame ordering. * wk-convert-fieldmap --frame extracts the right frame, with a hand-computed expected Hz->mm value as a sanity check. * wk-convert-fieldmap with vs without --flip-sign produces exact negations. * 5D ANTs single-warp file roundtrip through wk-convert-warp via files on disk (regression for the 5D shape-validation fix). * CLI Hz->mm->invert->Hz(--flip-sign) chain reproduces the medic non-native fieldmap output bit-for-bit. * invert(invert(map)) approximately recovers the original on a small smooth displacement (loose tolerance, interior only). 173 tests pass total (33 added). Co-Authored-By: Claude Opus 4.7 (1M context) --- tests/test_scripts.py | 358 ++++++++++++++++++++++++++++++++++++++++ tests/test_utilities.py | 136 +++++++++++++++ 2 files changed, 494 insertions(+) diff --git a/tests/test_scripts.py b/tests/test_scripts.py index 6749e43..b457156 100644 --- a/tests/test_scripts.py +++ b/tests/test_scripts.py @@ -1651,3 +1651,361 @@ def test_write_output_per_frame_map_clears_vector_intent(tmp_path): 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) diff --git a/tests/test_utilities.py b/tests/test_utilities.py index 425d410..7c1cee8 100644 --- a/tests/test_utilities.py +++ b/tests/test_utilities.py @@ -389,3 +389,139 @@ def test_displacement_field_to_map_clears_vector_intent(): 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_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) From d65e76714eeea31f477e8c8e005d248f9d07f695 Mon Sep 17 00:00:00 2001 From: Andrew Van Date: Sat, 25 Apr 2026 02:17:38 -0500 Subject: [PATCH 7/8] :white_check_mark: Push CLI + utilities coverage to 99% Adds tests for the previously-uncovered branches across all six CLIs plus a few low-level utility paths. CLI branches covered: * wk-apply-warp: 3D map transform, 5D ANTs single-warp, 5D field series, --reference grid, --transform-type=map with multi-file --transform rejection, --transform-type=field with non-itk --format, invalid field/input/map shape errors. * wk-compute-jacobian: ANTs-format input, --frame selection, bundled vs. per-frame outputs are bit-identical. * wk-convert-warp: --invert + format conversion combo, multi-frame field --invert routing (to default field output AND --to=map output). * wk-convert-fieldmap: multi-file fieldmap input, 5D field input, --frame out-of-range error. * wk-unwrap-phase: --noiseframes strips trailing volumes (asserts shape reduction on the bundled fixture). Utility branches covered: * convert_warp out_type=bogus rejected (separate from in_type branch). * resample_image rejects wrong-channel transform; squeezes 5D transform. * create_brain_mask: positive extra_dilation grows mask, negative erodes. * bundle_frames_to_field_series squeezes 5D singleton inputs. Coverage on warpkit/scripts + warpkit/utilities went from 94% to 99% (201 tests, was 173). Remaining 8 lines are essentially-unreachable defensive raises and one debug branch. Co-Authored-By: Claude Opus 4.7 (1M context) --- tests/test_scripts.py | 748 ++++++++++++++++++++++++++++++++++++++++ tests/test_utilities.py | 69 ++++ 2 files changed, 817 insertions(+) diff --git a/tests/test_scripts.py b/tests/test_scripts.py index b457156..fd03385 100644 --- a/tests/test_scripts.py +++ b/tests/test_scripts.py @@ -2009,3 +2009,751 @@ def test_convert_warp_double_invert_recovers_input(argv, tmp_path): # 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_utilities.py b/tests/test_utilities.py index 7c1cee8..fc375b7 100644 --- a/tests/test_utilities.py +++ b/tests/test_utilities.py @@ -511,6 +511,75 @@ def test_convert_warp_lps_affine_itk_roundtrip(): # --------------------------------------------------------------------------- +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 From a14a6c4859455d8abdb0059e2dfb0e275fb55aa5 Mon Sep 17 00:00:00 2001 From: Andrew Van Date: Sat, 25 Apr 2026 02:30:21 -0500 Subject: [PATCH 8/8] :bug: Address remaining PR review: mask validation, noiseframes guard, help text - compute_field_maps now validates masks shape against unwrapped data (raises ValueError instead of misbehaving deeper in SVD). - wk-unwrap-phase rejects negative --noiseframes and any value that would leave 0 frames per input, with a clean parser.error. - wk-convert-warp --invert help now describes the per-frame routing and the off-axis-component drop for multi-frame field input. - wk-convert-warp --to help no longer claims auto-classification (we dropped --from=auto in 9164bcb); now says it defaults to --from. - New regression tests for each. Co-Authored-By: Claude Opus 4.7 (1M context) --- tests/test_scripts.py | 54 +++++++++++++++++++++++++++++++++ tests/test_unwrap.py | 54 ++++++++++++++++++++++++++++++++- warpkit/scripts/convert_warp.py | 18 +++++++---- warpkit/scripts/unwrap_phase.py | 14 +++++++++ warpkit/unwrap.py | 13 ++++++++ 5 files changed, 146 insertions(+), 7 deletions(-) diff --git a/tests/test_scripts.py b/tests/test_scripts.py index fd03385..0160f1d 100644 --- a/tests/test_scripts.py +++ b/tests/test_scripts.py @@ -546,6 +546,60 @@ def test_unwrap_phase_metadata_count_must_match_phase_count(argv, capsys, tmp_pa 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 # --------------------------------------------------------------------------- diff --git a/tests/test_unwrap.py b/tests/test_unwrap.py index 6cee3a2..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 @@ -88,6 +89,57 @@ def test_romeo_lowercase_bindings_exist(): 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_unwrap3d diff --git a/warpkit/scripts/convert_warp.py b/warpkit/scripts/convert_warp.py index 2a457fc..58a9dc8 100644 --- a/warpkit/scripts/convert_warp.py +++ b/warpkit/scripts/convert_warp.py @@ -156,7 +156,7 @@ def main(): "--to", dest="to_type", choices=("map", "field"), - help="Output type. Defaults to whatever the input classifies as.", + help="Output type. Defaults to --from (no type conversion).", ) parser.add_argument( "--from-format", @@ -188,11 +188,17 @@ def main(): "--invert", action="store_true", help=( - "Invert each frame before any type/format conversion. Maps are " - "inverted with the 1D map inverter along --axis; fields are " - "inverted with the full 3D field inverter. Inversion of maps " - "requires --axis (the map's own axis); inversion of fields does " - "not, but downstream map output still does." + "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( diff --git a/warpkit/scripts/unwrap_phase.py b/warpkit/scripts/unwrap_phase.py index 2578a56..5246eaa 100644 --- a/warpkit/scripts/unwrap_phase.py +++ b/warpkit/scripts/unwrap_phase.py @@ -111,7 +111,21 @@ def main(): 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) diff --git a/warpkit/unwrap.py b/warpkit/unwrap.py index e39ff37..dea909e 100644 --- a/warpkit/unwrap.py +++ b/warpkit/unwrap.py @@ -939,6 +939,19 @@ def compute_field_maps( 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)