From 2002a911868171ba2d7db8e587faf92db4414461 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?= Date: Thu, 11 Apr 2024 20:47:08 -0400 Subject: [PATCH] ENH: Implement DWI equality explicitly Implement DWI equality explicitly: account for the cases where the properties can be `None` by creating a helper method and exclude the `_filepath` property as it is created using a randomly generated dirname. --- src/eddymotion/data/dmri.py | 22 +++++-- test/test_dmri.py | 128 ++++++++++++++++++++++++++++++++++++ 2 files changed, 143 insertions(+), 7 deletions(-) diff --git a/src/eddymotion/data/dmri.py b/src/eddymotion/data/dmri.py index d3f28269..6e46f0b1 100644 --- a/src/eddymotion/data/dmri.py +++ b/src/eddymotion/data/dmri.py @@ -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.""" diff --git a/test/test_dmri.py b/test/test_dmri.py index 004a4e13..dde0a44d 100644 --- a/test/test_dmri.py +++ b/test/test_dmri.py @@ -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""" @@ -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