From 9fbf5de8ff3d0b950330e70e392e6159c81e490a Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Fri, 30 Aug 2024 17:46:43 +0200 Subject: [PATCH 01/12] enh: prepare a utility that generates ANTs' command lines --- src/eddymotion/conftest.py | 54 ++++++ src/eddymotion/data/fileA.nii.gz | 0 src/eddymotion/data/fileB.nii.gz | 0 src/eddymotion/data/maskA.nii.gz | 0 src/eddymotion/registration/ants.py | 169 ++++++++++++++++++ .../registration/config/b0-to-b0_level0.json | 25 +++ 6 files changed, 248 insertions(+) create mode 100644 src/eddymotion/conftest.py create mode 100644 src/eddymotion/data/fileA.nii.gz create mode 100644 src/eddymotion/data/fileB.nii.gz create mode 100644 src/eddymotion/data/maskA.nii.gz create mode 100644 src/eddymotion/registration/config/b0-to-b0_level0.json diff --git a/src/eddymotion/conftest.py b/src/eddymotion/conftest.py new file mode 100644 index 00000000..5d47bfc3 --- /dev/null +++ b/src/eddymotion/conftest.py @@ -0,0 +1,54 @@ +# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*- +# vi: set ft=python sts=4 ts=4 sw=4 et: +# +# Copyright 2021 The NiPreps Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# We support and encourage derived works from this project, please read +# about our expectations at +# +# https://www.nipreps.org/community/licensing/ +# +"""py.test configuration.""" + +import os +from pathlib import Path + +import nibabel +import numpy +import pytest + +test_data_env = os.getenv("TEST_DATA_HOME", str(Path.home() / "eddymotion-tests")) +test_output_dir = os.getenv("TEST_OUTPUT_DIR") +test_workdir = os.getenv("TEST_WORK_DIR") + +_datadir = (Path(__file__).parent.parent / "eddymotion" / "data").absolute() + + +def pytest_report_header(config): + return f"""\ +TEST_DATA_HOME={test_data_env}. +TEST_OUTPUT_DIR={test_output_dir or ' (output files will be discarded)'}. +TEST_WORK_DIR={test_workdir or ' (intermediate files will be discarded)'}. +""" + + +@pytest.fixture(autouse=True) +def doctest_imports(doctest_namespace): + """Populates doctests with some conveniency imports.""" + doctest_namespace["np"] = numpy + doctest_namespace["nb"] = nibabel + doctest_namespace["os"] = os + doctest_namespace["Path"] = Path + doctest_namespace["repodata"] = _datadir diff --git a/src/eddymotion/data/fileA.nii.gz b/src/eddymotion/data/fileA.nii.gz new file mode 100644 index 00000000..e69de29b diff --git a/src/eddymotion/data/fileB.nii.gz b/src/eddymotion/data/fileB.nii.gz new file mode 100644 index 00000000..e69de29b diff --git a/src/eddymotion/data/maskA.nii.gz b/src/eddymotion/data/maskA.nii.gz new file mode 100644 index 00000000..e69de29b diff --git a/src/eddymotion/registration/ants.py b/src/eddymotion/registration/ants.py index f238174d..5bc6a3fc 100644 --- a/src/eddymotion/registration/ants.py +++ b/src/eddymotion/registration/ants.py @@ -23,6 +23,9 @@ """Using ANTs for image registration.""" from collections import namedtuple +from json import loads +from pathlib import Path +from warnings import warn import nibabel as nb import nitransforms as nt @@ -86,6 +89,172 @@ def _prepare_registration_data(dwframe, predicted, affine, vol_idx, dirname, reg return fixed, moving +def _get_ants_settings(settings="b0-to-b0_level0"): + return Path( + pkg_fn( + "eddymotion.registration", + f"config/{settings}.json", + ) + ) + + +def generate_command( + fixed_path, + moving_path, + fixedmask_path=None, + movingmask_path=None, + init_affine=None, + default="b0-to-b0_level0", + **kwargs, +): + """ + Generate an ANTs' command line. + + Examples + -------- + >>> generate_command( + ... fixed_path=repodata / 'fileA.nii.gz', + ... moving_path=repodata / 'fileB.nii.gz', + ... ) # doctest: +NORMALIZE_WHITESPACE + 'antsRegistration --collapse-output-transforms 1 --dimensionality 3 \ + --initialize-transforms-per-stage 0 --interpolation Linear --output transform \ + --transform Rigid[ 20.0 ] \ + --metric GC[ /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + 1, 5, Random, 0.1 ] \ + --convergence [ 20, 1e-05, 2 ] --smoothing-sigmas 2.0vox --shrink-factors 2 \ + --use-histogram-matching 1 -v --winsorize-image-intensities [ 0.01, 0.998 ] \ + --write-composite-transform 0' + + >>> generate_command( + ... fixed_path=repodata / 'fileA.nii.gz', + ... moving_path=repodata / 'fileB.nii.gz', + ... default="dwi-to-b0_level0", + ... ) # doctest: +NORMALIZE_WHITESPACE + 'antsRegistration --collapse-output-transforms 1 --dimensionality 3 \ + --initialize-transforms-per-stage 0 --interpolation Linear --output transform \ + --transform Rigid[ 0.01 ] --metric Mattes[ \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + 1, 32, Regular, 0.2 \ + ] --convergence [ 100x50, 1e-05, 10 ] --smoothing-sigmas 2.0x0.0vox \ + --shrink-factors 2x1 --use-histogram-matching 1 --transform Rigid[ 0.001 ] \ + --metric Mattes[ \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + 1, 32, Random, 0.1 \ + ] --convergence [ 25, 1e-06, 2 ] --smoothing-sigmas 0.0vox --shrink-factors 1 \ + --use-histogram-matching 1 -v --winsorize-image-intensities [ 0.0001, 0.9998 ] \ + --write-composite-transform 0' + + >>> generate_command( + ... fixed_path=repodata / 'fileA.nii.gz', + ... moving_path=repodata / 'fileB.nii.gz', + ... fixedmask_path=repodata / 'maskA.nii.gz', + ... default="dwi-to-b0_level0", + ... ) # doctest: +NORMALIZE_WHITESPACE + 'antsRegistration --collapse-output-transforms 1 --dimensionality 3 \ + --initialize-transforms-per-stage 0 --interpolation Linear --output transform \ + --transform Rigid[ 0.01 ] --metric Mattes[ \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + 1, 32, Regular, 0.2 ] \ + --convergence [ 100x50, 1e-05, 10 ] --smoothing-sigmas 2.0x0.0vox --shrink-factors 2x1 \ + --use-histogram-matching 1 --masks [ \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/maskA.nii.gz, NULL ] \ + --transform Rigid[ 0.001 ] --metric Mattes[ \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + 1, 32, Random, 0.1 ] \ + --convergence [ 25, 1e-06, 2 ] --smoothing-sigmas 0.0vox --shrink-factors 1 \ + --use-histogram-matching 1 --masks [ \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/maskA.nii.gz, NULL ] \ + -v --winsorize-image-intensities [ 0.0001, 0.9998 ] --write-composite-transform 0' + + >>> generate_command( + ... fixed_path=repodata / 'fileA.nii.gz', + ... moving_path=repodata / 'fileB.nii.gz', + ... default="dwi-to-b0_level0", + ... ) # doctest: +NORMALIZE_WHITESPACE + 'antsRegistration --collapse-output-transforms 1 --dimensionality 3 \ + --initialize-transforms-per-stage 0 --interpolation Linear --output transform \ + --transform Rigid[ 0.01 ] --metric Mattes[ \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + 1, 32, Regular, 0.2 \ + ] --convergence [ 100x50, 1e-05, 10 ] --smoothing-sigmas 2.0x0.0vox \ + --shrink-factors 2x1 --use-histogram-matching 1 --transform Rigid[ 0.001 ] \ + --metric Mattes[ \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + 1, 32, Random, 0.1 \ + ] --convergence [ 25, 1e-06, 2 ] --smoothing-sigmas 0.0vox --shrink-factors 1 \ + --use-histogram-matching 1 -v --winsorize-image-intensities [ 0.0001, 0.9998 ] \ + --write-composite-transform 0' + + >>> generate_command( + ... fixed_path=repodata / 'fileA.nii.gz', + ... moving_path=repodata / 'fileB.nii.gz', + ... fixedmask_path=[repodata / 'maskA.nii.gz'], + ... default="dwi-to-b0_level0", + ... ) # doctest: +NORMALIZE_WHITESPACE + 'antsRegistration --collapse-output-transforms 1 --dimensionality 3 \ + --initialize-transforms-per-stage 0 --interpolation Linear --output transform \ + --transform Rigid[ 0.01 ] --metric Mattes[ \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + 1, 32, Regular, 0.2 ] \ + --convergence [ 100x50, 1e-05, 10 ] --smoothing-sigmas 2.0x0.0vox --shrink-factors 2x1 \ + --use-histogram-matching 1 --masks [ NULL, NULL ] \ + --transform Rigid[ 0.001 ] --metric Mattes[ \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + 1, 32, Random, 0.1 ] \ + --convergence [ 25, 1e-06, 2 ] --smoothing-sigmas 0.0vox --shrink-factors 1 \ + --use-histogram-matching 1 --masks [ \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/maskA.nii.gz, NULL ] \ + -v --winsorize-image-intensities [ 0.0001, 0.9998 ] --write-composite-transform 0' + + """ + + settings = loads(_get_ants_settings(default).read_text()) | kwargs + + nlevels = len(settings["metric"]) + fixed_path = Path(fixed_path).absolute() + moving_path = Path(moving_path).absolute() + + if fixedmask_path is not None: + if isinstance(fixedmask_path, (str, Path)): + fixedmask_path = [str(fixedmask_path)] * nlevels + elif len(fixedmask_path) < nlevels: + fixedmask_path = ["NULL"] * (nlevels - len(fixedmask_path)) + fixedmask_path + elif len(fixedmask_path) > nlevels: + warn("More fixed mask paths than levels", stacklevel=1) + fixedmask_path = fixedmask_path[:nlevels] + + settings["fixed_image_masks"] = [str(f) for f in fixedmask_path] + + if movingmask_path is not None: + if isinstance(movingmask_path, (str, Path)): + movingmask_path = [movingmask_path] * nlevels + elif len(movingmask_path) < nlevels: + movingmask_path = ["NULL"] * (nlevels - len(movingmask_path)) + movingmask_path + elif len(movingmask_path) > nlevels: + warn("More moving mask paths than levels", stacklevel=1) + movingmask_path = movingmask_path[:nlevels] + + settings["moving_image_masks"] = [str(m) for m in movingmask_path] + + if init_affine is not None: + settings["initial_moving_transform"] = str(init_affine) + + return Registration( + fixed_image=str(fixed_path), + moving_image=str(moving_path), + **settings, + ).cmdline + + def _run_registration( fixed, moving, diff --git a/src/eddymotion/registration/config/b0-to-b0_level0.json b/src/eddymotion/registration/config/b0-to-b0_level0.json new file mode 100644 index 00000000..067bb827 --- /dev/null +++ b/src/eddymotion/registration/config/b0-to-b0_level0.json @@ -0,0 +1,25 @@ +{ + "collapse_output_transforms": true, + "convergence_threshold": [ 1E-5 ], + "convergence_window_size": [ 2 ], + "dimension": 3, + "initialize_transforms_per_stage": false, + "interpolation": "Linear", + "metric": [ "GC" ], + "metric_weight": [ 1.0 ], + "number_of_iterations": [[ 20 ]], + "radius_or_number_of_bins": [ 5 ], + "sampling_percentage": [ 0.1 ], + "sampling_strategy": [ "Random" ], + "shrink_factors": [[ 2 ]], + "sigma_units": [ "vox" ], + "smoothing_sigmas": [[ 2.0 ]], + "transform_parameters": [[ 20.0 ]], + "transforms": [ "Rigid" ], + "use_histogram_matching": [ true ], + "verbose": true, + "winsorize_lower_quantile": 0.01, + "winsorize_upper_quantile": 0.998, + "write_composite_transform": false + } + \ No newline at end of file From 8645f4c087e890aa74010712035a7f8906b47912 Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Sun, 1 Sep 2024 12:15:39 +0200 Subject: [PATCH 02/12] enh: add utilities for registration --- src/eddymotion/registration/utils.py | 59 ++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) create mode 100644 src/eddymotion/registration/utils.py diff --git a/src/eddymotion/registration/utils.py b/src/eddymotion/registration/utils.py new file mode 100644 index 00000000..23d0cfc9 --- /dev/null +++ b/src/eddymotion/registration/utils.py @@ -0,0 +1,59 @@ +# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*- +# vi: set ft=python sts=4 ts=4 sw=4 et: +# +# Copyright 2024 The NiPreps Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# We support and encourage derived works from this project, please read +# about our expectations at +# +# https://www.nipreps.org/community/licensing/ +# +"""Utilities to aid in performing and evaluating image registration.""" + +from __future__ import annotations + +from itertools import product + +import nibabel as nb +import nitransforms as nt +import numpy as np + + +def displacements_within_mask( + maskimg: nb.spatialimages.SpatialImage, + test_xfm: nt.base.BaseTransform, + reference_xfm: nt.base.BaseTransform | None = None, +) -> np.ndarray: + maskdata = np.asanyarray(maskimg.dataobj) > 0 + xyz = nb.affines.apply_affine( + maskimg.affine, + np.argwhere(maskdata), + ) + targets = test_xfm.map(xyz) + + diffs = targets - xyz if reference_xfm is None else targets - reference_xfm.map(xyz) + return np.linalg.norm(diffs, axis=-1) + + +def displacement_framewise( + img: nb.spatialimages.SpatialImage, + test_xfm: nt.base.BaseTransform, + radius: float = 50.0, +): + affine = img.affine + center_ijk = 0.5 * (np.array(img.shape[:3]) - 1) + center_xyz = nb.affines.apply_affine(affine, center_ijk) + fd_coords = np.array(list(product(*((radius, -radius),) * 3))) + center_xyz + return np.mean(np.linalg.norm(test_xfm.map(fd_coords) - fd_coords, axis=-1)) From 8b2df679baff99652be8cc44dcd893ba627320de Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Sun, 1 Sep 2024 12:16:35 +0200 Subject: [PATCH 03/12] enh: script to select ANTs' hyperparameters with SMAC --- scripts/optimize_registration.py | 201 +++++++++++++++++++++++++++++++ 1 file changed, 201 insertions(+) create mode 100644 scripts/optimize_registration.py diff --git a/scripts/optimize_registration.py b/scripts/optimize_registration.py new file mode 100644 index 00000000..c67cadf6 --- /dev/null +++ b/scripts/optimize_registration.py @@ -0,0 +1,201 @@ +# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*- +# vi: set ft=python sts=4 ts=4 sw=4 et: +# +# Copyright 2024 The NiPreps Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# We support and encourage derived works from this project, please read +# about our expectations at +# +# https://www.nipreps.org/community/licensing/ +# +"""Optimize ANTs' configurations.""" + +import asyncio +import random +from pathlib import Path +from shutil import rmtree +from tempfile import mkdtemp + +import nibabel as nb +import nitransforms as nt +import numpy as np +from ConfigSpace import Configuration, ConfigurationSpace +from smac import HyperparameterOptimizationFacade, Scenario + +from eddymotion.registration import ants as erants +from eddymotion.registration import utils + +# When inside ipython / jupyter +# import nest_asyncio +# nest_asyncio.apply() + +## Generate config dictionary +configdict = { + # "convergence_threshold": (1e-5, 1e-6), + "winsorize_lower_quantile": (0.001, 0.2), + "winsorize_upper_quantile": (0.98, 0.999), + "transform_parameters": (0.1, 50.0), + "smoothing_sigmas": (0.0, 8.0), + "shrink_factors": (1, 4), + "radius_or_number_of_bins": (3, 7), + "sampling_percentage": (0.1, 0.8), + "metric": ["GC", "CC"], + "sampling_strategy": ["Random", "Regular"], +} +paramspace = ConfigurationSpace(configdict) + + +async def ants(cmd, cwd, stdout, stderr, semaphore): + async with semaphore: + proc = await asyncio.create_subprocess_shell( + cmd, + cwd=cwd, + stdout=stdout, + stderr=stderr, + ) + returncode = await proc.wait() + return returncode + + +DATASET_PATH = Path("/data/datasets/hcph") +DWI_RUN = ( + DATASET_PATH / "sub-001" / "ses-038" / "dwi" / "sub-001_ses-038_acq-highres_dir-PA_dwi.nii.gz" +) + +WORKDIR = Path.home() / "tmp" / "eddymotiondev" +WORKDIR.mkdir(parents=True, exist_ok=True) + +nii = nb.load(DWI_RUN) + +bvecs = np.loadtxt(DWI_RUN.parent / DWI_RUN.name.replace(".nii.gz", ".bvec")).T +bvals = np.loadtxt(DWI_RUN.parent / DWI_RUN.name.replace(".nii.gz", ".bval")) +b0mask = bvals < 50 +data_b0s = nii.get_fdata()[..., b0mask] +brainmask_path = WORKDIR / "b0_desc-brain_mask.nii.gz" +brainmask_nii = nb.load(brainmask_path) +brainmask = np.asanyarray(brainmask_nii.dataobj) > 1e-5 + +ribbon_nii = nb.load(WORKDIR / "ribbon.nii.gz") + +fixed_path = WORKDIR / "first-b0.nii.gz" +refnii = nb.four_to_three(nii)[0] +refnii.to_filename(fixed_path) + +EXPERIMENTDIR = WORKDIR / "smac" +if EXPERIMENTDIR.exists(): + rmtree(EXPERIMENTDIR, ignore_errors=True) + +EXPERIMENTDIR.mkdir(parents=True, exist_ok=True) + + +async def train_coro(config: Configuration, conversions: int = 50, seed: int = 0) -> float: + random.seed(seed) + + tmp_folder = Path(mkdtemp(prefix="i-", dir=EXPERIMENTDIR)) + + ref_xfms = [] + tasks = [] + semaphore = asyncio.Semaphore(12) + for i in range(conversions): + r_x, r_y, r_z = ( + random.uniform(-0.1, 0.1), + random.uniform(-0.1, 0.1), + random.uniform(-0.1, 0.1), + ) + t_x, t_y, t_z = ( + random.uniform(-1.0, 1.0), + random.uniform(-1.0, 1.0), + random.uniform(-1.0, 1.0), + ) + T = nb.affines.from_matvec( + nb.eulerangles.euler2mat(x=r_x, y=r_y, z=r_z), + (t_x, t_y, t_z), + ) + xfm = nt.linear.Affine(T, reference=refnii) + ref_xfms.append(xfm) + + moving_path = tmp_folder / f"test-{i:02d}.nii.gz" + (~xfm).apply(refnii, reference=refnii).to_filename(moving_path) + + align_kwargs = {k: config[k] for k in configdict.keys()} + + cmdline = erants.generate_command( + fixed_path, + moving_path, + fixedmask_path=brainmask_path, + output_transform_prefix=f"conversion-{i:02d}", + **align_kwargs, + ) + + tasks.append( + ants( + cmdline, + cwd=str(tmp_folder), + stdout=Path(tmp_folder / f"ants-{i:04d}.out").open("w+"), + stderr=Path(tmp_folder / f"ants-{i:04d}.err").open("w+"), + semaphore=semaphore, + ) + ) + + results = await asyncio.gather(*tasks, return_exceptions=True) + + diff = [] + for i, r in enumerate(results): + if r: + return 1e6 + + fixednii = nb.load(fixed_path) + movingnii = nb.load(tmp_folder / f"test-{i:02d}.nii.gz") + xform = nt.linear.Affine( + nt.io.itk.ITKLinearTransform.from_filename( + tmp_folder / f"conversion-{i:02d}0GenericAffine.mat" + ).to_ras( + reference=fixednii, + moving=movingnii, + ), + ) + disps = utils.displacements_within_mask( + ribbon_nii, + xform, + ref_xfms[i], + ) + diff.append( + 0.20 * np.percentile(disps, 95) + + 0.8 * np.median(disps) + - 0.2 * np.percentile(disps, 5), + ) + + error = np.mean(diff) + + with Path(tmp_folder / f"ants-{i:04d}.out").open("a") as f: + f.write(f"Iteration error={error}\n\n") + f.write(f"Tested parameters:\n{config}") + + return error + + +def train(config: Configuration, conversions: int = 50, seed: int = 0) -> float: + loop = asyncio.get_event_loop() + return loop.run_until_complete(train_coro(config, conversions, seed)) + + +# Scenario object specifying the optimization environment +scenario = Scenario(paramspace, n_trials=200) + +# Use SMAC to find the best configuration/hyperparameters +smac = HyperparameterOptimizationFacade(scenario, train) +incumbent = smac.optimize() + +print(incumbent) From 178d24fddd510c3c14e9da2dfc1d8a54ea594457 Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Sun, 1 Sep 2024 12:18:01 +0200 Subject: [PATCH 04/12] enh: update ANTs' support --- src/eddymotion/registration/ants.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/src/eddymotion/registration/ants.py b/src/eddymotion/registration/ants.py index 5bc6a3fc..42615eb7 100644 --- a/src/eddymotion/registration/ants.py +++ b/src/eddymotion/registration/ants.py @@ -34,6 +34,14 @@ from nitransforms.linear import Affine from pkg_resources import resource_filename as pkg_fn +PARAMETERS_SINGLE_LIST = { + "radius_or_number_of_bins", + "sampling_percentage", + "metric", + "sampling_strategy", +} +PARAMETERS_DOUBLE_LIST = {"shrink_factors", "smoothing_sigmas", "transform_parameters"} + def _to_nifti(data, affine, filename, clip=True): data = np.squeeze(data) @@ -217,7 +225,15 @@ def generate_command( """ - settings = loads(_get_ants_settings(default).read_text()) | kwargs + settings = loads(_get_ants_settings(default).read_text()) + + for key, value in kwargs.items(): + if key in PARAMETERS_SINGLE_LIST: + value = [value] + elif key in PARAMETERS_DOUBLE_LIST: + value = [[value]] + + settings[key] = value nlevels = len(settings["metric"]) fixed_path = Path(fixed_path).absolute() From 514f3ac3a528fbac2e4ccedb9d22f680f24b26e5 Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Tue, 3 Sep 2024 10:52:16 +0200 Subject: [PATCH 05/12] enh: improve test dataset --- test/test_registration.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/test/test_registration.py b/test/test_registration.py index a7415859..b8a1ca64 100644 --- a/test/test_registration.py +++ b/test/test_registration.py @@ -33,7 +33,7 @@ from nipype.interfaces.ants.registration import Registration from pkg_resources import resource_filename as pkg_fn -from eddymotion.data.dmri import DWI +from eddymotion.registration.utils import displacements_within_mask @pytest.mark.parametrize("r_x", [0.0, 0.1, 0.3]) @@ -42,19 +42,17 @@ @pytest.mark.parametrize("t_x", [0.0, 1.0]) @pytest.mark.parametrize("t_y", [0.0, 1.0]) @pytest.mark.parametrize("t_z", [0.0, 1.0]) -def test_ANTs_config_b0(datadir, tmp_path, r_x, r_y, r_z, t_x, t_y, t_z): +# @pytest.mark.parametrize("dataset", ["hcph", "dwi"]) +@pytest.mark.parametrize("dataset", ["dwi"]) # only dwi for now +def test_ANTs_config_b0(datadir, tmp_path, dataset, r_x, r_y, r_z, t_x, t_y, t_z): """Check that the registration parameters for b=0 gives a good estimate of known affine""" - fixed = tmp_path / "b0.nii.gz" + fixed = datadir / f"{dataset}-b0_desc-avg.nii.gz" + fixed_mask = datadir / f"{dataset}-b0_desc-brain.nii.gz" moving = tmp_path / "moving.nii.gz" - dwdata = DWI.from_filename(datadir / "dwi.h5") - b0nii = nb.Nifti1Image(dwdata.bzero, dwdata.affine, None) - b0nii.header.set_qform(dwdata.affine, code=1) - b0nii.header.set_sform(dwdata.affine, code=1) - b0nii.to_filename(fixed) - + b0nii = nb.load(fixed) T = from_matvec(euler2mat(x=r_x, y=r_y, z=r_z), (t_x, t_y, t_z)) xfm = nt.linear.Affine(T, reference=b0nii) @@ -68,6 +66,7 @@ def test_ANTs_config_b0(datadir, tmp_path, r_x, r_y, r_z, t_x, t_y, t_z): ), fixed_image=str(fixed.absolute()), moving_image=str(moving.absolute()), + fixed_image_masks=[str(fixed_mask)], random_seed=1234, num_threads=cpu_count(), ) @@ -78,6 +77,7 @@ def test_ANTs_config_b0(datadir, tmp_path, r_x, r_y, r_z, t_x, t_y, t_z): reference=b0nii, ) - coords = xfm.reference.ndcoords.T - rms = np.sqrt(((xfm.map(coords) - xform.map(coords)) ** 2).sum(1)).mean() - assert rms < 0.8 + masknii = nb.load(fixed_mask) + assert displacements_within_mask(masknii, xform, xfm).mean() < 0.5 * np.mean( + b0nii.header.get_zooms()[:3] + ) From b6aa756f083a45584e2f2df56b27206f001f8e22 Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Tue, 3 Sep 2024 10:53:10 +0200 Subject: [PATCH 06/12] enh: improve registration hyperparameters tunning script --- scripts/optimize_registration.py | 92 +++++++++++-------- .../registration/config/b0-to-b0_level0.json | 16 ++-- 2 files changed, 60 insertions(+), 48 deletions(-) diff --git a/scripts/optimize_registration.py b/scripts/optimize_registration.py index c67cadf6..1a9afbe5 100644 --- a/scripts/optimize_registration.py +++ b/scripts/optimize_registration.py @@ -23,7 +23,9 @@ """Optimize ANTs' configurations.""" import asyncio +import logging import random +from os import getenv from pathlib import Path from shutil import rmtree from tempfile import mkdtemp @@ -37,10 +39,14 @@ from eddymotion.registration import ants as erants from eddymotion.registration import utils +logger = logging.getLogger("ants-optimization") + # When inside ipython / jupyter # import nest_asyncio # nest_asyncio.apply() +TIME_PENALTY_WEIGHT = 0.1 + ## Generate config dictionary configdict = { # "convergence_threshold": (1e-5, 1e-6), @@ -49,9 +55,9 @@ "transform_parameters": (0.1, 50.0), "smoothing_sigmas": (0.0, 8.0), "shrink_factors": (1, 4), - "radius_or_number_of_bins": (3, 7), + "radius_or_number_of_bins": (3, 12), "sampling_percentage": (0.1, 0.8), - "metric": ["GC", "CC"], + "metric": ["GC"], "sampling_strategy": ["Random", "Regular"], } paramspace = ConfigurationSpace(configdict) @@ -69,30 +75,11 @@ async def ants(cmd, cwd, stdout, stderr, semaphore): return returncode -DATASET_PATH = Path("/data/datasets/hcph") -DWI_RUN = ( - DATASET_PATH / "sub-001" / "ses-038" / "dwi" / "sub-001_ses-038_acq-highres_dir-PA_dwi.nii.gz" -) +DATASET_PATH = Path(getenv("TEST_DATA_HOME", f"{getenv('HOME')}/.cache/eddymotion-tests")) WORKDIR = Path.home() / "tmp" / "eddymotiondev" WORKDIR.mkdir(parents=True, exist_ok=True) -nii = nb.load(DWI_RUN) - -bvecs = np.loadtxt(DWI_RUN.parent / DWI_RUN.name.replace(".nii.gz", ".bvec")).T -bvals = np.loadtxt(DWI_RUN.parent / DWI_RUN.name.replace(".nii.gz", ".bval")) -b0mask = bvals < 50 -data_b0s = nii.get_fdata()[..., b0mask] -brainmask_path = WORKDIR / "b0_desc-brain_mask.nii.gz" -brainmask_nii = nb.load(brainmask_path) -brainmask = np.asanyarray(brainmask_nii.dataobj) > 1e-5 - -ribbon_nii = nb.load(WORKDIR / "ribbon.nii.gz") - -fixed_path = WORKDIR / "first-b0.nii.gz" -refnii = nb.four_to_three(nii)[0] -refnii.to_filename(fixed_path) - EXPERIMENTDIR = WORKDIR / "smac" if EXPERIMENTDIR.exists(): rmtree(EXPERIMENTDIR, ignore_errors=True) @@ -100,10 +87,16 @@ async def ants(cmd, cwd, stdout, stderr, semaphore): EXPERIMENTDIR.mkdir(parents=True, exist_ok=True) -async def train_coro(config: Configuration, conversions: int = 50, seed: int = 0) -> float: +async def train_coro( + config: Configuration, + conversions: int = 100, + seed: int = 0, + verbose: bool = False, +) -> float: random.seed(seed) tmp_folder = Path(mkdtemp(prefix="i-", dir=EXPERIMENTDIR)) + align_kwargs = {k: config[k] for k in configdict.keys()} ref_xfms = [] tasks = [] @@ -123,19 +116,23 @@ async def train_coro(config: Configuration, conversions: int = 50, seed: int = 0 nb.eulerangles.euler2mat(x=r_x, y=r_y, z=r_z), (t_x, t_y, t_z), ) + + fixed_path = ( + DATASET_PATH / f"{'dwi' if i < (conversions // 2) else 'hcph'}-b0_desc-avg.nii.gz" + ) + brainmask_path = DATASET_PATH / fixed_path.name.replace("desc-avg", "desc-brain_mask") + refnii = nb.load(fixed_path) xfm = nt.linear.Affine(T, reference=refnii) ref_xfms.append(xfm) - moving_path = tmp_folder / f"test-{i:02d}.nii.gz" + moving_path = tmp_folder / f"test-{i:04d}.nii.gz" (~xfm).apply(refnii, reference=refnii).to_filename(moving_path) - align_kwargs = {k: config[k] for k in configdict.keys()} - cmdline = erants.generate_command( fixed_path, moving_path, fixedmask_path=brainmask_path, - output_transform_prefix=f"conversion-{i:02d}", + output_transform_prefix=f"conversion-{i:04d}", **align_kwargs, ) @@ -152,41 +149,53 @@ async def train_coro(config: Configuration, conversions: int = 50, seed: int = 0 results = await asyncio.gather(*tasks, return_exceptions=True) diff = [] + times = [] for i, r in enumerate(results): if r: return 1e6 + fixed_path = ( + DATASET_PATH / f"{'dwi' if i < (conversions // 2) else 'hcph'}-b0_desc-avg.nii.gz" + ) + brainmask_path = DATASET_PATH / fixed_path.name.replace("desc-avg", "desc-brain_mask") + fixednii = nb.load(fixed_path) - movingnii = nb.load(tmp_folder / f"test-{i:02d}.nii.gz") + movingnii = nb.load(tmp_folder / f"test-{i:04d}.nii.gz") xform = nt.linear.Affine( nt.io.itk.ITKLinearTransform.from_filename( - tmp_folder / f"conversion-{i:02d}0GenericAffine.mat" + tmp_folder / f"conversion-{i:04d}0GenericAffine.mat" ).to_ras( reference=fixednii, moving=movingnii, ), ) disps = utils.displacements_within_mask( - ribbon_nii, + nb.load(brainmask_path), xform, ref_xfms[i], ) - diff.append( - 0.20 * np.percentile(disps, 95) - + 0.8 * np.median(disps) - - 0.2 * np.percentile(disps, 5), - ) + diff.append(np.mean(disps)) + + # Parse log -- Total elapsed time: 1.0047e+00 + for line in reversed(Path(tmp_folder / f"ants-{i:04d}.out").read_text().splitlines()): + if line.strip().startswith("Total elapsed time:"): + times.append(float(line.strip().split(" ")[-1])) - error = np.mean(diff) + maxdiff = max(diff) + meantime = np.mean(times) + error = (1.0 - TIME_PENALTY_WEIGHT) * maxdiff + TIME_PENALTY_WEIGHT * meantime - with Path(tmp_folder / f"ants-{i:04d}.out").open("a") as f: - f.write(f"Iteration error={error}\n\n") - f.write(f"Tested parameters:\n{config}") + logger.warning( + f"Max. error ({conversions} it.): {error:0.3f} " + f"({maxdiff:0.2f} mm | {meantime:0.2f} s)." + ) + if verbose: + logger.warning(f"\n\nParameters:\n{align_kwargs}" f"\n\nConversions folder: {tmp_folder}.") return error -def train(config: Configuration, conversions: int = 50, seed: int = 0) -> float: +def train(config: Configuration, conversions: int = 100, seed: int = 0) -> float: loop = asyncio.get_event_loop() return loop.run_until_complete(train_coro(config, conversions, seed)) @@ -199,3 +208,6 @@ def train(config: Configuration, conversions: int = 50, seed: int = 0) -> float: incumbent = smac.optimize() print(incumbent) + +loop = asyncio.get_event_loop() +loop.run_until_complete(train_coro(incumbent, conversions=200, verbose=True)) diff --git a/src/eddymotion/registration/config/b0-to-b0_level0.json b/src/eddymotion/registration/config/b0-to-b0_level0.json index 067bb827..c22cb6c2 100644 --- a/src/eddymotion/registration/config/b0-to-b0_level0.json +++ b/src/eddymotion/registration/config/b0-to-b0_level0.json @@ -1,25 +1,25 @@ { "collapse_output_transforms": true, - "convergence_threshold": [ 1E-5 ], - "convergence_window_size": [ 2 ], + "convergence_threshold": [ 1E-6 ], + "convergence_window_size": [ 4 ], "dimension": 3, "initialize_transforms_per_stage": false, "interpolation": "Linear", "metric": [ "GC" ], "metric_weight": [ 1.0 ], "number_of_iterations": [[ 20 ]], - "radius_or_number_of_bins": [ 5 ], + "radius_or_number_of_bins": [ 7 ], "sampling_percentage": [ 0.1 ], - "sampling_strategy": [ "Random" ], - "shrink_factors": [[ 2 ]], + "sampling_strategy": [ "Regular" ], + "shrink_factors": [[ 1 ]], "sigma_units": [ "vox" ], - "smoothing_sigmas": [[ 2.0 ]], + "smoothing_sigmas": [[ 0.5 ]], "transform_parameters": [[ 20.0 ]], "transforms": [ "Rigid" ], "use_histogram_matching": [ true ], "verbose": true, - "winsorize_lower_quantile": 0.01, - "winsorize_upper_quantile": 0.998, + "winsorize_lower_quantile": 0.075, + "winsorize_upper_quantile": 0.982, "write_composite_transform": false } \ No newline at end of file From b8a2c840878ba9c21c4ff8e3f3b059d9e1de9eb0 Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Thu, 19 Sep 2024 08:24:11 +0200 Subject: [PATCH 07/12] enh: new b0-to-b0 registration generalizing well --- scripts/optimize_registration.py | 145 ++++++++++-------- .../registration/config/b0-to-b0_level0.json | 44 ++++-- test/test_registration.py | 10 +- 3 files changed, 110 insertions(+), 89 deletions(-) diff --git a/scripts/optimize_registration.py b/scripts/optimize_registration.py index 1a9afbe5..1890dcfd 100644 --- a/scripts/optimize_registration.py +++ b/scripts/optimize_registration.py @@ -24,17 +24,16 @@ import asyncio import logging -import random from os import getenv from pathlib import Path from shutil import rmtree -from tempfile import mkdtemp import nibabel as nb import nitransforms as nt import numpy as np from ConfigSpace import Configuration, ConfigurationSpace from smac import HyperparameterOptimizationFacade, Scenario +from smac.utils.configspace import get_config_hash from eddymotion.registration import ants as erants from eddymotion.registration import utils @@ -46,17 +45,18 @@ # nest_asyncio.apply() TIME_PENALTY_WEIGHT = 0.1 +SEED = 2390232 ## Generate config dictionary configdict = { # "convergence_threshold": (1e-5, 1e-6), "winsorize_lower_quantile": (0.001, 0.2), - "winsorize_upper_quantile": (0.98, 0.999), - "transform_parameters": (0.1, 50.0), - "smoothing_sigmas": (0.0, 8.0), + "winsorize_upper_quantile": (0.9, 0.999), + "transform_parameters": (0.01, 20.0), + "smoothing_sigmas": (0.0, 4.0), "shrink_factors": (1, 4), - "radius_or_number_of_bins": (3, 12), - "sampling_percentage": (0.1, 0.8), + "radius_or_number_of_bins": (3, 6), + "sampling_percentage": (0.1, 0.6), "metric": ["GC"], "sampling_strategy": ["Random", "Regular"], } @@ -86,77 +86,75 @@ async def ants(cmd, cwd, stdout, stderr, semaphore): EXPERIMENTDIR.mkdir(parents=True, exist_ok=True) +rng = np.random.default_rng(SEED) +MOTION_PARAMETERS = np.hstack( + (rng.uniform(-0.4, 0.4, size=(60, 3)), rng.uniform(-2.0, 2.0, size=(60, 3))) +) +CONVERSIONS = [ + nb.affines.from_matvec(nb.eulerangles.euler2mat(*parameters[:3]), parameters[3:]) + for parameters in MOTION_PARAMETERS +] + +REFERENCES = ( + DATASET_PATH / "dwi-b0_desc-avg.nii.gz", + DATASET_PATH / "hcph-b0_desc-avg.nii.gz", +) + async def train_coro( config: Configuration, - conversions: int = 100, seed: int = 0, verbose: bool = False, ) -> float: - random.seed(seed) - - tmp_folder = Path(mkdtemp(prefix="i-", dir=EXPERIMENTDIR)) + tmp_folder = EXPERIMENTDIR / get_config_hash(config) + tmp_folder.mkdir(parents=True, exist_ok=True) align_kwargs = {k: config[k] for k in configdict.keys()} ref_xfms = [] tasks = [] - semaphore = asyncio.Semaphore(12) - for i in range(conversions): - r_x, r_y, r_z = ( - random.uniform(-0.1, 0.1), - random.uniform(-0.1, 0.1), - random.uniform(-0.1, 0.1), - ) - t_x, t_y, t_z = ( - random.uniform(-1.0, 1.0), - random.uniform(-1.0, 1.0), - random.uniform(-1.0, 1.0), - ) - T = nb.affines.from_matvec( - nb.eulerangles.euler2mat(x=r_x, y=r_y, z=r_z), - (t_x, t_y, t_z), - ) - - fixed_path = ( - DATASET_PATH / f"{'dwi' if i < (conversions // 2) else 'hcph'}-b0_desc-avg.nii.gz" - ) - brainmask_path = DATASET_PATH / fixed_path.name.replace("desc-avg", "desc-brain_mask") - refnii = nb.load(fixed_path) - xfm = nt.linear.Affine(T, reference=refnii) - ref_xfms.append(xfm) - - moving_path = tmp_folder / f"test-{i:04d}.nii.gz" - (~xfm).apply(refnii, reference=refnii).to_filename(moving_path) - - cmdline = erants.generate_command( - fixed_path, - moving_path, - fixedmask_path=brainmask_path, - output_transform_prefix=f"conversion-{i:04d}", - **align_kwargs, - ) + semaphore = asyncio.Semaphore(18) + nconv = len(CONVERSIONS) + for i, T in enumerate(CONVERSIONS): + for j in (0, 1): + fixed_path = REFERENCES[j] + brainmask_path = DATASET_PATH / fixed_path.name.replace("desc-avg", "desc-brain_mask") + refnii = nb.load(fixed_path) + xfm = nt.linear.Affine(T, reference=refnii) + ref_xfms.append(xfm) + + index = i * len(REFERENCES) + j + moving_path = tmp_folder / f"test-{index:04d}.nii.gz" + (~xfm).apply(refnii, reference=refnii).to_filename(moving_path) + + cmdline = erants.generate_command( + fixed_path, + moving_path, + fixedmask_path=brainmask_path, + output_transform_prefix=f"conversion-{index:04d}", + **align_kwargs, + ) - tasks.append( - ants( - cmdline, - cwd=str(tmp_folder), - stdout=Path(tmp_folder / f"ants-{i:04d}.out").open("w+"), - stderr=Path(tmp_folder / f"ants-{i:04d}.err").open("w+"), - semaphore=semaphore, + tasks.append( + ants( + cmdline, + cwd=str(tmp_folder), + stdout=Path(tmp_folder / f"ants-{index:04d}.out").open("w+"), + stderr=Path(tmp_folder / f"ants-{index:04d}.err").open("w+"), + semaphore=semaphore, + ) ) - ) results = await asyncio.gather(*tasks, return_exceptions=True) diff = [] times = [] + start = [] for i, r in enumerate(results): if r: return 1e6 - fixed_path = ( - DATASET_PATH / f"{'dwi' if i < (conversions // 2) else 'hcph'}-b0_desc-avg.nii.gz" - ) + j = i % 2 + fixed_path = REFERENCES[j] brainmask_path = DATASET_PATH / fixed_path.name.replace("desc-avg", "desc-brain_mask") fixednii = nb.load(fixed_path) @@ -169,35 +167,46 @@ async def train_coro( moving=movingnii, ), ) + + masknii = nb.load(brainmask_path) + initial_error = utils.displacements_within_mask( + masknii, + ref_xfms[i], + ) + disps = utils.displacements_within_mask( - nb.load(brainmask_path), + masknii, xform, ref_xfms[i], ) - diff.append(np.mean(disps)) + diff.append(np.percentile(disps, 95)) + start.append(np.percentile(initial_error, 95)) # Parse log -- Total elapsed time: 1.0047e+00 for line in reversed(Path(tmp_folder / f"ants-{i:04d}.out").read_text().splitlines()): if line.strip().startswith("Total elapsed time:"): times.append(float(line.strip().split(" ")[-1])) - maxdiff = max(diff) + meandiff = np.mean(diff) meantime = np.mean(times) - error = (1.0 - TIME_PENALTY_WEIGHT) * maxdiff + TIME_PENALTY_WEIGHT * meantime + error = ((1.0 - TIME_PENALTY_WEIGHT) * meandiff + TIME_PENALTY_WEIGHT * meantime) / np.mean( + start + ) - logger.warning( - f"Max. error ({conversions} it.): {error:0.3f} " - f"({maxdiff:0.2f} mm | {meantime:0.2f} s)." + logger.info( + f"Normalized objective ({nconv} it.): {error:0.3f} " + f"({meandiff:0.2f} mm | {meantime:0.2f} s). " + f"Avg. p95 initial error: {np.mean(start):0.2f} mm." ) if verbose: - logger.warning(f"\n\nParameters:\n{align_kwargs}" f"\n\nConversions folder: {tmp_folder}.") + logger.info(f"\n\nParameters:\n{align_kwargs}" f"\n\nConversions folder: {tmp_folder}.") return error -def train(config: Configuration, conversions: int = 100, seed: int = 0) -> float: +def train(config: Configuration, seed: int = 0) -> float: loop = asyncio.get_event_loop() - return loop.run_until_complete(train_coro(config, conversions, seed)) + return loop.run_until_complete(train_coro(config, seed)) # Scenario object specifying the optimization environment @@ -210,4 +219,4 @@ def train(config: Configuration, conversions: int = 100, seed: int = 0) -> float print(incumbent) loop = asyncio.get_event_loop() -loop.run_until_complete(train_coro(incumbent, conversions=200, verbose=True)) +loop.run_until_complete(train_coro(incumbent, verbose=True)) diff --git a/src/eddymotion/registration/config/b0-to-b0_level0.json b/src/eddymotion/registration/config/b0-to-b0_level0.json index c22cb6c2..bab78e1f 100644 --- a/src/eddymotion/registration/config/b0-to-b0_level0.json +++ b/src/eddymotion/registration/config/b0-to-b0_level0.json @@ -1,25 +1,37 @@ { "collapse_output_transforms": true, - "convergence_threshold": [ 1E-6 ], - "convergence_window_size": [ 4 ], + "convergence_threshold": [ 1E-6, 1E-7 ], + "convergence_window_size": [ 4, 2 ], "dimension": 3, "initialize_transforms_per_stage": false, "interpolation": "Linear", - "metric": [ "GC" ], - "metric_weight": [ 1.0 ], - "number_of_iterations": [[ 20 ]], - "radius_or_number_of_bins": [ 7 ], - "sampling_percentage": [ 0.1 ], - "sampling_strategy": [ "Regular" ], - "shrink_factors": [[ 1 ]], - "sigma_units": [ "vox" ], - "smoothing_sigmas": [[ 0.5 ]], - "transform_parameters": [[ 20.0 ]], - "transforms": [ "Rigid" ], - "use_histogram_matching": [ true ], + "metric": [ "GC", "GC" ], + "metric_weight": [ 1.0, 1.0 ], + "number_of_iterations": [ + [ 20 ], + [ 10 ] + ], + "radius_or_number_of_bins": [ 3, 5 ], + "sampling_percentage": [ 0.4, 0.2 ], + "sampling_strategy": [ "Random", "Random" ], + "shrink_factors": [ + [ 3 ], + [ 1 ] + ], + "sigma_units": [ "vox", "vox" ], + "smoothing_sigmas": [ + [ 2.71 ], + [ 0.0 ] + ], + "transform_parameters": [ + [ 12.0 ], + [ 1.0 ] + ], + "transforms": [ "Rigid", "Rigid" ], + "use_histogram_matching": [ true, true ], "verbose": true, - "winsorize_lower_quantile": 0.075, - "winsorize_upper_quantile": 0.982, + "winsorize_lower_quantile": 0.063, + "winsorize_upper_quantile": 0.991, "write_composite_transform": false } \ No newline at end of file diff --git a/test/test_registration.py b/test/test_registration.py index b8a1ca64..ae45cb8a 100644 --- a/test/test_registration.py +++ b/test/test_registration.py @@ -42,8 +42,8 @@ @pytest.mark.parametrize("t_x", [0.0, 1.0]) @pytest.mark.parametrize("t_y", [0.0, 1.0]) @pytest.mark.parametrize("t_z", [0.0, 1.0]) -# @pytest.mark.parametrize("dataset", ["hcph", "dwi"]) -@pytest.mark.parametrize("dataset", ["dwi"]) # only dwi for now +@pytest.mark.parametrize("dataset", ["hcph", "dwi"]) +# @pytest.mark.parametrize("dataset", ["dwi"]) def test_ANTs_config_b0(datadir, tmp_path, dataset, r_x, r_y, r_z, t_x, t_y, t_z): """Check that the registration parameters for b=0 gives a good estimate of known affine""" @@ -62,7 +62,7 @@ def test_ANTs_config_b0(datadir, tmp_path, dataset, r_x, r_y, r_z, t_x, t_y, t_z terminal_output="file", from_file=pkg_fn( "eddymotion.registration", - "config/dwi-to-b0_level0.json", + "config/b0-to-b0_level0.json", ), fixed_image=str(fixed.absolute()), moving_image=str(moving.absolute()), @@ -78,6 +78,6 @@ def test_ANTs_config_b0(datadir, tmp_path, dataset, r_x, r_y, r_z, t_x, t_y, t_z ) masknii = nb.load(fixed_mask) - assert displacements_within_mask(masknii, xform, xfm).mean() < 0.5 * np.mean( - b0nii.header.get_zooms()[:3] + assert displacements_within_mask(masknii, xform, xfm).mean() < ( + 0.6 * np.mean(b0nii.header.get_zooms()[:3]) ) From 946af08a6e3dd99ce0e753a4b42112db799f45af Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Fri, 20 Sep 2024 09:24:42 +0200 Subject: [PATCH 08/12] enh: optimize multi-level registration schemes and update b0-to-b0 --- scripts/optimize_registration.py | 16 +-- src/eddymotion/registration/ants.py | 116 +++++++++++++----- .../registration/config/b0-to-b0_level0.json | 8 +- 3 files changed, 95 insertions(+), 45 deletions(-) diff --git a/scripts/optimize_registration.py b/scripts/optimize_registration.py index 1890dcfd..25b6e2a5 100644 --- a/scripts/optimize_registration.py +++ b/scripts/optimize_registration.py @@ -50,14 +50,14 @@ ## Generate config dictionary configdict = { # "convergence_threshold": (1e-5, 1e-6), - "winsorize_lower_quantile": (0.001, 0.2), - "winsorize_upper_quantile": (0.9, 0.999), - "transform_parameters": (0.01, 20.0), - "smoothing_sigmas": (0.0, 4.0), - "shrink_factors": (1, 4), - "radius_or_number_of_bins": (3, 6), - "sampling_percentage": (0.1, 0.6), - "metric": ["GC"], + # "winsorize_lower_quantile": (0.001, 0.2), + # "winsorize_upper_quantile": (0.9, 0.999), + "transform_parameters": (0.01, 2.0), + "smoothing_sigmas": (0.0, 1.0), + "shrink_factors": (1, 2), + "radius_or_number_of_bins": (3, 5), + "sampling_percentage": (0.1, 0.4), + # "metric": ["GC"], "sampling_strategy": ["Random", "Regular"], } paramspace = ConfigurationSpace(configdict) diff --git a/src/eddymotion/registration/ants.py b/src/eddymotion/registration/ants.py index 42615eb7..b390adfa 100644 --- a/src/eddymotion/registration/ants.py +++ b/src/eddymotion/registration/ants.py @@ -34,6 +34,19 @@ from nitransforms.linear import Affine from pkg_resources import resource_filename as pkg_fn +PARAMETERS_SINGLE_VALUE = { + "collapse_output_transforms", + "dimension", + "initial_moving_transform", + "initialize_transforms_per_stage", + "interpolation", + "output_transform_prefix", + "verbose", + "winsorize_lower_quantile", + "winsorize_upper_quantile", + "write_composite_transform", +} + PARAMETERS_SINGLE_LIST = { "radius_or_number_of_bins", "sampling_percentage", @@ -106,6 +119,35 @@ def _get_ants_settings(settings="b0-to-b0_level0"): ) +def _massage_mask_path(mask_path, nlevels): + """ + Generate nipype-compatible masks paths. + + Examples + -------- + >>> _massage_mask_path("/some/path", 2) + ['/some/path', '/some/path'] + + >>> _massage_mask_path(["/some/path"] * 2, 2) + ['/some/path', '/some/path'] + + >>> _massage_mask_path(["/some/path"] * 2, 4) + ['NULL', 'NULL', '/some/path', '/some/path'] + + >>> _massage_mask_path(["/some/path"] * 2, 1) + ['/some/path'] + + """ + if isinstance(mask_path, (str, Path)): + return [str(mask_path)] * nlevels + if len(mask_path) < nlevels: + return ["NULL"] * (nlevels - len(mask_path)) + mask_path + if len(mask_path) > nlevels: + warn("More mask paths than levels", stacklevel=1) + return mask_path[:nlevels] + return mask_path + + def generate_command( fixed_path, moving_path, @@ -126,12 +168,19 @@ def generate_command( ... ) # doctest: +NORMALIZE_WHITESPACE 'antsRegistration --collapse-output-transforms 1 --dimensionality 3 \ --initialize-transforms-per-stage 0 --interpolation Linear --output transform \ - --transform Rigid[ 20.0 ] \ + --transform Rigid[ 12.0 ] \ + --metric GC[ /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ + /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + 1, 3, Random, 0.4 ] \ + --convergence [ 20, 1e-06, 4 ] --smoothing-sigmas 2.71vox --shrink-factors 3 \ + --use-histogram-matching 1 \ + --transform Rigid[ 1.96 ] \ --metric GC[ /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ - 1, 5, Random, 0.1 ] \ - --convergence [ 20, 1e-05, 2 ] --smoothing-sigmas 2.0vox --shrink-factors 2 \ - --use-histogram-matching 1 -v --winsorize-image-intensities [ 0.01, 0.998 ] \ + 1, 4, Random, 0.18 ] \ + --convergence [ 10, 1e-07, 2 ] --smoothing-sigmas 0.0vox --shrink-factors 2 \ + --use-histogram-matching 1 \ + -v --winsorize-image-intensities [ 0.063, 0.991 ] \ --write-composite-transform 0' >>> generate_command( @@ -225,48 +274,49 @@ def generate_command( """ - settings = loads(_get_ants_settings(default).read_text()) + # Bootstrap settings from defaults file and override with single-valued parameters in args + settings = loads(_get_ants_settings(default).read_text()) | { + k: kwargs.pop(k) for k in PARAMETERS_SINGLE_VALUE if k in kwargs + } + # Determine number of levels and assert consistency of levels + levels = {len(settings[p]) for p in PARAMETERS_SINGLE_LIST if p in settings} + nlevels = levels.pop() + if levels: + raise RuntimeError(f"Malformed settings file (levels: {levels})") + + # Override list (and nested-list) parameters for key, value in kwargs.items(): - if key in PARAMETERS_SINGLE_LIST: + if key in PARAMETERS_DOUBLE_LIST: value = [value] - elif key in PARAMETERS_DOUBLE_LIST: - value = [[value]] - - settings[key] = value + elif key not in PARAMETERS_SINGLE_LIST: + continue - nlevels = len(settings["metric"]) - fixed_path = Path(fixed_path).absolute() - moving_path = Path(moving_path).absolute() + if levels == 1: + settings[key] = [value] + else: + settings[key][-1] = value + # Set fixed masks if provided if fixedmask_path is not None: - if isinstance(fixedmask_path, (str, Path)): - fixedmask_path = [str(fixedmask_path)] * nlevels - elif len(fixedmask_path) < nlevels: - fixedmask_path = ["NULL"] * (nlevels - len(fixedmask_path)) + fixedmask_path - elif len(fixedmask_path) > nlevels: - warn("More fixed mask paths than levels", stacklevel=1) - fixedmask_path = fixedmask_path[:nlevels] - - settings["fixed_image_masks"] = [str(f) for f in fixedmask_path] + settings["fixed_image_masks"] = [ + str(p) for p in _massage_mask_path(fixedmask_path, nlevels) + ] + # Set moving masks if provided if movingmask_path is not None: - if isinstance(movingmask_path, (str, Path)): - movingmask_path = [movingmask_path] * nlevels - elif len(movingmask_path) < nlevels: - movingmask_path = ["NULL"] * (nlevels - len(movingmask_path)) + movingmask_path - elif len(movingmask_path) > nlevels: - warn("More moving mask paths than levels", stacklevel=1) - movingmask_path = movingmask_path[:nlevels] - - settings["moving_image_masks"] = [str(m) for m in movingmask_path] + settings["moving_image_masks"] = [ + str(p) for p in _massage_mask_path(movingmask_path, nlevels) + ] + # Set initalizing affine if provided if init_affine is not None: settings["initial_moving_transform"] = str(init_affine) + # Generate command line with nipype and return return Registration( - fixed_image=str(fixed_path), - moving_image=str(moving_path), + fixed_image=str(Path(fixed_path).absolute()), + moving_image=str(Path(moving_path).absolute()), **settings, ).cmdline diff --git a/src/eddymotion/registration/config/b0-to-b0_level0.json b/src/eddymotion/registration/config/b0-to-b0_level0.json index bab78e1f..5035c4c2 100644 --- a/src/eddymotion/registration/config/b0-to-b0_level0.json +++ b/src/eddymotion/registration/config/b0-to-b0_level0.json @@ -11,12 +11,12 @@ [ 20 ], [ 10 ] ], - "radius_or_number_of_bins": [ 3, 5 ], - "sampling_percentage": [ 0.4, 0.2 ], + "radius_or_number_of_bins": [ 3, 4 ], + "sampling_percentage": [ 0.4, 0.18 ], "sampling_strategy": [ "Random", "Random" ], "shrink_factors": [ [ 3 ], - [ 1 ] + [ 2 ] ], "sigma_units": [ "vox", "vox" ], "smoothing_sigmas": [ @@ -25,7 +25,7 @@ ], "transform_parameters": [ [ 12.0 ], - [ 1.0 ] + [ 1.96 ] ], "transforms": [ "Rigid", "Rigid" ], "use_histogram_matching": [ true, true ], From 50074d3d84fe5454b2cfd92d0cd6e8c3495e86ec Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Mon, 30 Sep 2024 23:19:19 +0200 Subject: [PATCH 09/12] fix: doctest in ants module --- pyproject.toml | 2 +- src/eddymotion/registration/ants.py | 46 ++++++++++++++--------------- 2 files changed, 24 insertions(+), 24 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b8c95e5b..b4f68910 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -164,7 +164,7 @@ known-first-party=["eddymotion"] pythonpath = "src/ test/" norecursedirs = [".*", "_*"] addopts = "-v --doctest-modules" -doctest_optionflags = "ALLOW_UNICODE NORMALIZE_WHITESPACE" +doctest_optionflags = "ALLOW_UNICODE NORMALIZE_WHITESPACE ELLIPSIS" env = "PYTHONHASHSEED=0" filterwarnings = ["ignore::DeprecationWarning"] diff --git a/src/eddymotion/registration/ants.py b/src/eddymotion/registration/ants.py index b390adfa..103ce112 100644 --- a/src/eddymotion/registration/ants.py +++ b/src/eddymotion/registration/ants.py @@ -169,14 +169,14 @@ def generate_command( 'antsRegistration --collapse-output-transforms 1 --dimensionality 3 \ --initialize-transforms-per-stage 0 --interpolation Linear --output transform \ --transform Rigid[ 12.0 ] \ - --metric GC[ /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + --metric GC[ .../fileA.nii.gz, \ + .../fileB.nii.gz, \ 1, 3, Random, 0.4 ] \ --convergence [ 20, 1e-06, 4 ] --smoothing-sigmas 2.71vox --shrink-factors 3 \ --use-histogram-matching 1 \ --transform Rigid[ 1.96 ] \ - --metric GC[ /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + --metric GC[ .../fileA.nii.gz, \ + .../fileB.nii.gz, \ 1, 4, Random, 0.18 ] \ --convergence [ 10, 1e-07, 2 ] --smoothing-sigmas 0.0vox --shrink-factors 2 \ --use-histogram-matching 1 \ @@ -191,14 +191,14 @@ def generate_command( 'antsRegistration --collapse-output-transforms 1 --dimensionality 3 \ --initialize-transforms-per-stage 0 --interpolation Linear --output transform \ --transform Rigid[ 0.01 ] --metric Mattes[ \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + .../fileA.nii.gz, \ + .../fileB.nii.gz, \ 1, 32, Regular, 0.2 \ ] --convergence [ 100x50, 1e-05, 10 ] --smoothing-sigmas 2.0x0.0vox \ --shrink-factors 2x1 --use-histogram-matching 1 --transform Rigid[ 0.001 ] \ --metric Mattes[ \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + .../fileA.nii.gz, \ + .../fileB.nii.gz, \ 1, 32, Random, 0.1 \ ] --convergence [ 25, 1e-06, 2 ] --smoothing-sigmas 0.0vox --shrink-factors 1 \ --use-histogram-matching 1 -v --winsorize-image-intensities [ 0.0001, 0.9998 ] \ @@ -213,19 +213,19 @@ def generate_command( 'antsRegistration --collapse-output-transforms 1 --dimensionality 3 \ --initialize-transforms-per-stage 0 --interpolation Linear --output transform \ --transform Rigid[ 0.01 ] --metric Mattes[ \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + .../fileA.nii.gz, \ + .../fileB.nii.gz, \ 1, 32, Regular, 0.2 ] \ --convergence [ 100x50, 1e-05, 10 ] --smoothing-sigmas 2.0x0.0vox --shrink-factors 2x1 \ --use-histogram-matching 1 --masks [ \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/maskA.nii.gz, NULL ] \ + .../maskA.nii.gz, NULL ] \ --transform Rigid[ 0.001 ] --metric Mattes[ \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + .../fileA.nii.gz, \ + .../fileB.nii.gz, \ 1, 32, Random, 0.1 ] \ --convergence [ 25, 1e-06, 2 ] --smoothing-sigmas 0.0vox --shrink-factors 1 \ --use-histogram-matching 1 --masks [ \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/maskA.nii.gz, NULL ] \ + .../maskA.nii.gz, NULL ] \ -v --winsorize-image-intensities [ 0.0001, 0.9998 ] --write-composite-transform 0' >>> generate_command( @@ -236,14 +236,14 @@ def generate_command( 'antsRegistration --collapse-output-transforms 1 --dimensionality 3 \ --initialize-transforms-per-stage 0 --interpolation Linear --output transform \ --transform Rigid[ 0.01 ] --metric Mattes[ \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + .../fileA.nii.gz, \ + .../fileB.nii.gz, \ 1, 32, Regular, 0.2 \ ] --convergence [ 100x50, 1e-05, 10 ] --smoothing-sigmas 2.0x0.0vox \ --shrink-factors 2x1 --use-histogram-matching 1 --transform Rigid[ 0.001 ] \ --metric Mattes[ \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + .../fileA.nii.gz, \ + .../fileB.nii.gz, \ 1, 32, Random, 0.1 \ ] --convergence [ 25, 1e-06, 2 ] --smoothing-sigmas 0.0vox --shrink-factors 1 \ --use-histogram-matching 1 -v --winsorize-image-intensities [ 0.0001, 0.9998 ] \ @@ -258,18 +258,18 @@ def generate_command( 'antsRegistration --collapse-output-transforms 1 --dimensionality 3 \ --initialize-transforms-per-stage 0 --interpolation Linear --output transform \ --transform Rigid[ 0.01 ] --metric Mattes[ \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + .../fileA.nii.gz, \ + .../fileB.nii.gz, \ 1, 32, Regular, 0.2 ] \ --convergence [ 100x50, 1e-05, 10 ] --smoothing-sigmas 2.0x0.0vox --shrink-factors 2x1 \ --use-histogram-matching 1 --masks [ NULL, NULL ] \ --transform Rigid[ 0.001 ] --metric Mattes[ \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileA.nii.gz, \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/fileB.nii.gz, \ + .../fileA.nii.gz, \ + .../fileB.nii.gz, \ 1, 32, Random, 0.1 ] \ --convergence [ 25, 1e-06, 2 ] --smoothing-sigmas 0.0vox --shrink-factors 1 \ --use-histogram-matching 1 --masks [ \ - /data/home/oesteban/workspace/eddymotion/src/eddymotion/data/maskA.nii.gz, NULL ] \ + .../maskA.nii.gz, NULL ] \ -v --winsorize-image-intensities [ 0.0001, 0.9998 ] --write-composite-transform 0' """ From c2994a75abccd045f2c0859d4451290a259538f7 Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Tue, 1 Oct 2024 11:39:36 +0200 Subject: [PATCH 10/12] fix: address review comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Jon Haitz Legarreta Gorroño --- pyproject.toml | 9 +- src/eddymotion/registration/ants.py | 166 ++++++++++++++++++++------- src/eddymotion/registration/utils.py | 57 ++++++++- test/test_registration.py | 9 ++ 4 files changed, 195 insertions(+), 46 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b4f68910..4a096105 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,8 +66,6 @@ plotting = ["nilearn"] resmon = ["psutil >=5.4"] -popylar = ["popylar >= 0.2"] - test = [ "coverage", "pytest >= 4.4", @@ -76,10 +74,15 @@ test = [ "pytest-xdist >= 1.28" ] +antsopt = [ + "ConfigSpace", + "smac", +] + # Aliases docs = ["eddymotion[doc]"] tests = ["eddymotion[test]"] -all = ["eddymotion[doc,test,dev,plotting,resmon,popylar]"] +all = ["eddymotion[doc,test,dev,plotting,resmon,antsopt]"] [project.scripts] eddymotion = "eddymotion.cli.run:main" diff --git a/src/eddymotion/registration/ants.py b/src/eddymotion/registration/ants.py index 103ce112..925ed1a3 100644 --- a/src/eddymotion/registration/ants.py +++ b/src/eddymotion/registration/ants.py @@ -22,6 +22,8 @@ # """Using ANTs for image registration.""" +from __future__ import annotations + from collections import namedtuple from json import loads from pathlib import Path @@ -56,7 +58,24 @@ PARAMETERS_DOUBLE_LIST = {"shrink_factors", "smoothing_sigmas", "transform_parameters"} -def _to_nifti(data, affine, filename, clip=True): +def _to_nifti( + data: np.ndarray, affine: np.ndarray, filename: str | Path, clip: bool = True +) -> None: + """ + Save data as a NIfTI file, optionally applying clipping. + + Parameters + ---------- + data : :obj:`~numpy.ndarray` + The image data to be saved. + affine : :obj:`~numpy.ndarray` + The affine transformation matrix. + filename : :obj:`os.pathlike` + The file path where the NIfTI file will be saved. + clip : :obj:`bool`, optional + Whether to apply clipping to the data before saving, by default True. + + """ data = np.squeeze(data) if clip: from eddymotion.data.filtering import advanced_clip @@ -72,34 +91,42 @@ def _to_nifti(data, affine, filename, clip=True): nii.to_filename(filename) -def _prepare_registration_data(dwframe, predicted, affine, vol_idx, dirname, reg_target_type): - """Prepare the registration data: save the fixed and moving images to disk. +def _prepare_registration_data( + dwframe: np.ndarray, + predicted: np.ndarray, + affine: np.ndarray, + vol_idx: int, + dirname: Path | str, + reg_target_type: str, +) -> tuple[Path, Path]: + """ + Prepare the registration data: save the fixed and moving images to disk. Parameters ---------- - dwframe : :obj:`numpy.ndarray` + dwframe : :obj:`~numpy.ndarray` DWI data object. - predicted : :obj:`numpy.ndarray` + predicted : :obj:`~numpy.ndarray` Predicted data. affine : :obj:`numpy.ndarray` Affine transformation matrix. - vol_idx : :obj:`int + vol_idx : :obj:`int` DWI volume index. - dirname : :obj:`Path` + dirname : :obj:`os.pathlike` Directory name where the data is saved. reg_target_type : :obj:`str` Target registration type. Returns ------- - fixed : :obj:`Path` + fixed : :obj:`~pathlib.Path` Fixed image filename. - moving : :obj:`Path` + moving : :obj:`~pathlib.Path` Moving image filename. """ - moving = dirname / f"moving{vol_idx:05d}.nii.gz" - fixed = dirname / f"fixed{vol_idx:05d}.nii.gz" + moving = Path(dirname) / f"moving{vol_idx:05d}.nii.gz" + fixed = Path(dirname) / f"fixed{vol_idx:05d}.nii.gz" _to_nifti(dwframe, affine, moving) _to_nifti( predicted, @@ -110,7 +137,29 @@ def _prepare_registration_data(dwframe, predicted, affine, vol_idx, dirname, reg return fixed, moving -def _get_ants_settings(settings="b0-to-b0_level0"): +def _get_ants_settings(settings: str = "b0-to-b0_level0") -> Path: + """ + Retrieve the path to ANTs settings configuration file. + + Parameters + ---------- + settings : :obj:`str`, optional + Name of the settings configuration, by default ``"b0-to-b0_level0"``. + + Returns + ------- + :obj:`~pathlib.Path` + The path to the configuration file. + + Examples + -------- + >>> _get_ants_settings() + PosixPath('.../config/b0-to-b0_level0.json') + + >>> _get_ants_settings("b0-to-b0_level1") + PosixPath('.../config/b0-to-b0_level1.json') + + """ return Path( pkg_fn( "eddymotion.registration", @@ -119,10 +168,22 @@ def _get_ants_settings(settings="b0-to-b0_level0"): ) -def _massage_mask_path(mask_path, nlevels): +def _massage_mask_path(mask_path: str | Path | list[str], nlevels: int) -> list[str]: """ Generate nipype-compatible masks paths. + Parameters + ---------- + mask_path : :obj:`os.pathlike` or :obj:`list` + Path(s) to the mask file(s). + nlevels : :obj:`int` + Number of registration levels. + + Returns + ------- + :obj:`list` + A list of mask paths formatted for *Nipype*. + Examples -------- >>> _massage_mask_path("/some/path", 2) @@ -134,9 +195,6 @@ def _massage_mask_path(mask_path, nlevels): >>> _massage_mask_path(["/some/path"] * 2, 4) ['NULL', 'NULL', '/some/path', '/some/path'] - >>> _massage_mask_path(["/some/path"] * 2, 1) - ['/some/path'] - """ if isinstance(mask_path, (str, Path)): return [str(mask_path)] * nlevels @@ -149,17 +207,39 @@ def _massage_mask_path(mask_path, nlevels): def generate_command( - fixed_path, - moving_path, - fixedmask_path=None, - movingmask_path=None, - init_affine=None, - default="b0-to-b0_level0", - **kwargs, -): + fixed_path: str | Path, + moving_path: str | Path, + fixedmask_path: str | Path | list[str] | None = None, + movingmask_path: str | Path | list[str] | None = None, + init_affine: str | Path | None = None, + default: str = "b0-to-b0_level0", + **kwargs: dict, +) -> str: """ Generate an ANTs' command line. + Parameters + ---------- + fixed_path : :obj:`os.pathlike` + Path to the fixed image. + moving_path : :obj:`os.pathlike` + Path to the moving image. + fixedmask_path : :obj:`os.pathlike` or :obj:`list`, optional + Path to the fixed image mask, by default None. + movingmask_path : :obj:`os.pathlike` or :obj:`list`, optional + Path to the moving image mask, by default None. + init_affine : :obj:`os.pathlike`, optional + Initial affine transformation, by default None. + default : :obj:`str`, optional + Default settings configuration, by default "b0-to-b0_level0". + **kwargs : :obj:`dict` + Additional parameters for ANTs registration. + + Returns + ------- + :obj:`str` + The ANTs registration command line string. + Examples -------- >>> generate_command( @@ -322,21 +402,22 @@ def generate_command( def _run_registration( - fixed, - moving, - bmask_img, - em_affines, - affine, - shape, - bval, - fieldmap, - i_iter, - vol_idx, - dirname, - reg_target_type, - align_kwargs, -): - """Register the moving image to the fixed image. + fixed: Path, + moving: Path, + bmask_img: nb.spatialimages.SpatialImage, + em_affines: np.ndarray, + affine: np.ndarray, + shape: tuple[int, int, int], + bval: int, + fieldmap: nb.spatialimages.SpatialImage, + i_iter: int, + vol_idx: int, + dirname: Path, + reg_target_type: str, + align_kwargs: dict, +) -> nt.base.BaseTransform: + """ + Register the moving image to the fixed image. Parameters ---------- @@ -344,7 +425,7 @@ def _run_registration( Fixed image filename. moving : :obj:`Path` Moving image filename. - bmask_img : :class:`~nibabel.nifti1.Nifti1Image` + bmask_img : :class:`~nibabel.spatialimages.SpatialImage` Brainmask image. em_affines : :obj:`numpy.ndarray` Estimated eddy motion affine transformation matrices. @@ -354,7 +435,7 @@ def _run_registration( Shape of the DWI frame. bval : :obj:`int` b-value of the corresponding DWI volume. - fieldmap : :class:`~nibabel.nifti1.Nifti1Image` + fieldmap : :class:`~nibabel.spatialimages.SpatialImage` Fieldmap. i_iter : :obj:`int` Iteration number. @@ -369,8 +450,9 @@ def _run_registration( Returns ------- - xform : :class:`~nitransforms.linear.Affine` + xform : :obj:`~nitransforms.base.BaseTransform` Registration transformation. + """ if isinstance(reg_target_type, str): diff --git a/src/eddymotion/registration/utils.py b/src/eddymotion/registration/utils.py index 23d0cfc9..59380319 100644 --- a/src/eddymotion/registration/utils.py +++ b/src/eddymotion/registration/utils.py @@ -20,7 +20,14 @@ # # https://www.nipreps.org/community/licensing/ # -"""Utilities to aid in performing and evaluating image registration.""" +""" +Utilities to aid in performing and evaluating image registration. + +This module provides functions to compute displacements of image coordinates +under a transformation, useful for assessing the accuracy of image registration +processes. + +""" from __future__ import annotations @@ -36,13 +43,38 @@ def displacements_within_mask( test_xfm: nt.base.BaseTransform, reference_xfm: nt.base.BaseTransform | None = None, ) -> np.ndarray: + """ + Compute the distance between voxel coordinates mapped through two transforms. + + Parameters + ---------- + maskimg : :obj:`~nibabel.spatialimages.SpatialImage` + A mask image that defines the region of interest. Voxel coordinates + within the mask are transformed. + test_xfm : :obj:`~nitransforms.base.BaseTransform` + The transformation to test. This transformation is applied to the + voxel coordinates. + reference_xfm : :obj:`~nitransforms.base.BaseTransform`, optional + A reference transformation to compare with. If ``None``, the identity + transformation is assumed (no transformation). + + Returns + ------- + :obj:`~numpy.ndarray` + An array of displacements (in mm) for each voxel within the mask. + + """ + # Mask data as boolean (True for voxels inside the mask) maskdata = np.asanyarray(maskimg.dataobj) > 0 + # Convert voxel coordinates to world coordinates using affine transform xyz = nb.affines.apply_affine( maskimg.affine, np.argwhere(maskdata), ) + # Apply the test transformation targets = test_xfm.map(xyz) + # Compute the difference (displacement) between the test and reference transformations diffs = targets - xyz if reference_xfm is None else targets - reference_xfm.map(xyz) return np.linalg.norm(diffs, axis=-1) @@ -52,8 +84,31 @@ def displacement_framewise( test_xfm: nt.base.BaseTransform, radius: float = 50.0, ): + """ + Compute the framewise displacement (FD) for a given transformation. + + Parameters + ---------- + img : :obj:`~nibabel.spatialimages.SpatialImage` + The reference image. Used to extract the center coordinates. + test_xfm : :obj:`~nitransforms.base.BaseTransform` + The transformation to test. Applied to coordinates around the image center. + radius : :obj:`float`, optional + The radius (in mm) of the spherical neighborhood around the center of the image. + Default is 50.0 mm. + + Returns + ------- + :obj:`float` + The average framewise displacement (FD) for the test transformation. + + """ affine = img.affine + # Compute the center of the image in voxel space center_ijk = 0.5 * (np.array(img.shape[:3]) - 1) + # Convert to world coordinates center_xyz = nb.affines.apply_affine(affine, center_ijk) + # Generate coordinates of points at radius distance from center fd_coords = np.array(list(product(*((radius, -radius),) * 3))) + center_xyz + # Compute the average displacement from the test transformation return np.mean(np.linalg.norm(test_xfm.map(fd_coords) - fd_coords, axis=-1)) diff --git a/test/test_registration.py b/test/test_registration.py index ae45cb8a..4f193911 100644 --- a/test/test_registration.py +++ b/test/test_registration.py @@ -33,6 +33,7 @@ from nipype.interfaces.ants.registration import Registration from pkg_resources import resource_filename as pkg_fn +from eddymotion.registration.ants import _massage_mask_path from eddymotion.registration.utils import displacements_within_mask @@ -81,3 +82,11 @@ def test_ANTs_config_b0(datadir, tmp_path, dataset, r_x, r_y, r_z, t_x, t_y, t_z assert displacements_within_mask(masknii, xform, xfm).mean() < ( 0.6 * np.mean(b0nii.header.get_zooms()[:3]) ) + + +def test_massage_mask_path(): + """Test the case where a warning must be issued.""" + with pytest.warns(UserWarning, match="More mask paths than levels"): + maskpath = _massage_mask_path(["/some/path"] * 2, 1) + + assert maskpath == ["/some/path"] From 55a0cf2e4d7e5189435fb5a1f68d174d9c4fe319 Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Tue, 1 Oct 2024 15:16:57 +0200 Subject: [PATCH 11/12] enh: use fast config in integration test --- test/test_integration.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_integration.py b/test/test_integration.py index 485a99cf..1e5f1a4f 100644 --- a/test/test_integration.py +++ b/test/test_integration.py @@ -74,7 +74,7 @@ def test_proximity_estimator_trivial_model(datadir): models=("b0",), seed=None, align_kwargs={ - "fixed_modality": "dwi", + "fixed_modality": "b0", "moving_modality": "b0", "num_threads": min(cpu_count(), 8), }, From 1e00683072afa56a8a6e1cbda93cd336bc721746 Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Tue, 1 Oct 2024 15:33:49 +0200 Subject: [PATCH 12/12] sty: revise argument name MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Jon Haitz Legarreta Gorroño --- src/eddymotion/registration/utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/eddymotion/registration/utils.py b/src/eddymotion/registration/utils.py index 59380319..0ac54240 100644 --- a/src/eddymotion/registration/utils.py +++ b/src/eddymotion/registration/utils.py @@ -39,7 +39,7 @@ def displacements_within_mask( - maskimg: nb.spatialimages.SpatialImage, + mask_img: nb.spatialimages.SpatialImage, test_xfm: nt.base.BaseTransform, reference_xfm: nt.base.BaseTransform | None = None, ) -> np.ndarray: @@ -48,7 +48,7 @@ def displacements_within_mask( Parameters ---------- - maskimg : :obj:`~nibabel.spatialimages.SpatialImage` + mask_img : :obj:`~nibabel.spatialimages.SpatialImage` A mask image that defines the region of interest. Voxel coordinates within the mask are transformed. test_xfm : :obj:`~nitransforms.base.BaseTransform` @@ -65,10 +65,10 @@ def displacements_within_mask( """ # Mask data as boolean (True for voxels inside the mask) - maskdata = np.asanyarray(maskimg.dataobj) > 0 + maskdata = np.asanyarray(mask_img.dataobj) > 0 # Convert voxel coordinates to world coordinates using affine transform xyz = nb.affines.apply_affine( - maskimg.affine, + mask_img.affine, np.argwhere(maskdata), ) # Apply the test transformation