Skip to content
This repository has been archived by the owner on Dec 20, 2024. It is now read-only.

Commit

Permalink
Merge pull request #161 from jhlegarreta/OverloadDWIEquality
Browse files Browse the repository at this point in the history
ENH: Implement DWI equality explicitly
  • Loading branch information
effigies authored May 20, 2024
2 parents de1b535 + 2002a91 commit 356431a
Show file tree
Hide file tree
Showing 2 changed files with 144 additions and 8 deletions.
24 changes: 16 additions & 8 deletions src/eddymotion/data/dmri.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,34 +40,42 @@ def _data_repr(value):
return f"<{'x'.join(str(v) for v in value.shape)} ({value.dtype})>"


def _cmp(lh, rh):
if isinstance(lh, np.ndarray) and isinstance(rh, np.ndarray):
return np.allclose(lh, rh)

return lh == rh


@attr.s(slots=True)
class DWI:
"""Data representation structure for dMRI data."""

dataobj = attr.ib(default=None, repr=_data_repr)
dataobj = attr.ib(default=None, repr=_data_repr, eq=attr.cmp_using(eq=_cmp))
"""A numpy ndarray object for the data array, without *b=0* volumes."""
affine = attr.ib(default=None, repr=_data_repr)
affine = attr.ib(default=None, repr=_data_repr, eq=attr.cmp_using(eq=_cmp))
"""Best affine for RAS-to-voxel conversion of coordinates (NIfTI header)."""
brainmask = attr.ib(default=None, repr=_data_repr)
brainmask = attr.ib(default=None, repr=_data_repr, eq=attr.cmp_using(eq=_cmp))
"""A boolean ndarray object containing a corresponding brainmask."""
bzero = attr.ib(default=None, repr=_data_repr)
bzero = attr.ib(default=None, repr=_data_repr, eq=attr.cmp_using(eq=_cmp))
"""
A *b=0* reference map, preferably obtained by some smart averaging.
If the :math:`B_0` fieldmap is set, this *b=0* reference map should also
be unwarped.
"""
gradients = attr.ib(default=None, repr=_data_repr)
gradients = attr.ib(default=None, repr=_data_repr, eq=attr.cmp_using(eq=_cmp))
"""A 2D numpy array of the gradient table in RAS+B format."""
em_affines = attr.ib(default=None)
em_affines = attr.ib(default=None, eq=attr.cmp_using(eq=_cmp))
"""
List of :obj:`nitransforms.linear.Affine` objects that bring
DWIs (i.e., no b=0) into alignment.
"""
fieldmap = attr.ib(default=None, repr=_data_repr)
fieldmap = attr.ib(default=None, repr=_data_repr, eq=attr.cmp_using(eq=_cmp))
"""A 3D displacements field to unwarp susceptibility distortions."""
_filepath = attr.ib(
factory=lambda: Path(mkdtemp()) / "em_cache.h5",
repr=False,
eq=False,
)
"""A path to an HDF5 file to store the whole dataset."""

Expand Down Expand Up @@ -239,6 +247,6 @@ def load(

if fmap_file:
fmapimg = nb.load(fmap_file)
retval.fieldmap = fmapimg.get_fdata(fmapimg, dtype="float32")
retval.fieldmap = fmapimg.get_fdata(dtype="float32")

return retval
128 changes: 128 additions & 0 deletions test/test_dmri.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,12 +22,87 @@
#
"""Unit tests exercising the dMRI data structure."""

import nibabel as nb
import numpy as np
import pytest

from eddymotion.data.dmri import load


def _create_dwi_random_dataobj():
rng = np.random.default_rng(1234)

n_gradients = 10
b0s = 1
volumes = n_gradients + b0s
b0_thres = 50
bvals = np.hstack([b0s * [0], n_gradients * [1000]])
bvecs = np.hstack([np.zeros((3, b0s)), rng.random((3, n_gradients))])

vol_size = (34, 36, 24)

dwi_dataobj = rng.random((*vol_size, volumes), dtype="float32")
affine = np.eye(4, dtype="float32")
brainmask_dataobj = rng.random(vol_size, dtype="float32")
b0_dataobj = rng.random(vol_size, dtype="float32")
gradients = np.vstack([bvecs, bvals[np.newaxis, :]], dtype="float32")
fieldmap_dataobj = rng.random(vol_size, dtype="float32")

return (
dwi_dataobj,
affine,
brainmask_dataobj,
b0_dataobj,
gradients,
fieldmap_dataobj,
b0_thres,
)


def _create_dwi_random_data(
dwi_dataobj,
affine,
brainmask_dataobj,
b0_dataobj,
fieldmap_dataobj,
):
dwi = nb.Nifti1Image(dwi_dataobj, affine)
brainmask = nb.Nifti1Image(brainmask_dataobj, affine)
b0 = nb.Nifti1Image(b0_dataobj, affine)
fieldmap = nb.Nifti1Image(fieldmap_dataobj, affine)

return dwi, brainmask, b0, fieldmap


def _serialize_dwi_data(
dwi,
brainmask,
b0,
gradients,
fieldmap,
_tmp_path,
):
dwi_fname = _tmp_path / "dwi.nii.gz"
brainmask_fname = _tmp_path / "brainmask.nii.gz"
b0_fname = _tmp_path / "b0.nii.gz"
gradients_fname = _tmp_path / "gradients.txt"
fieldmap_fname = _tmp_path / "fieldmap.nii.gz"

nb.save(dwi, dwi_fname)
nb.save(brainmask, brainmask_fname)
nb.save(b0, b0_fname)
np.savetxt(gradients_fname, gradients.T)
nb.save(fieldmap, fieldmap_fname)

return (
dwi_fname,
brainmask_fname,
b0_fname,
gradients_fname,
fieldmap_fname,
)


def test_load(datadir, tmp_path):
"""Check that the registration parameters for b=0
gives a good estimate of known affine"""
Expand Down Expand Up @@ -65,3 +140,56 @@ def test_load(datadir, tmp_path):
assert np.allclose(dwi_h5.dataobj, dwi_from_nifti2.dataobj)
assert np.allclose(dwi_h5.bzero, dwi_from_nifti2.bzero)
assert np.allclose(dwi_h5.gradients, dwi_from_nifti2.gradients)


def test_equality_operator(tmp_path):
# Create some random data
(
dwi_dataobj,
affine,
brainmask_dataobj,
b0_dataobj,
gradients,
fieldmap_dataobj,
b0_thres,
) = _create_dwi_random_dataobj()

dwi, brainmask, b0, fieldmap = _create_dwi_random_data(
dwi_dataobj,
affine,
brainmask_dataobj,
b0_dataobj,
fieldmap_dataobj,
)

(
dwi_fname,
brainmask_fname,
b0_fname,
gradients_fname,
fieldmap_fname,
) = _serialize_dwi_data(
dwi,
brainmask,
b0,
gradients,
fieldmap,
tmp_path,
)

dwi_obj = load(
dwi_fname,
gradients_file=gradients_fname,
b0_file=b0_fname,
brainmask_file=brainmask_fname,
fmap_file=fieldmap_fname,
b0_thres=b0_thres,
)
hdf5_filename = tmp_path / "test_dwi.h5"
dwi_obj.to_filename(hdf5_filename)

round_trip_dwi_obj = load(hdf5_filename)

# Symmetric equality
assert dwi_obj == round_trip_dwi_obj
assert round_trip_dwi_obj == dwi_obj

0 comments on commit 356431a

Please sign in to comment.