ENH: Add a shell data property to DWI data class

src/nifreeze/data/dmri.py

@@ -242,6 +242,42 @@ def to_nifti(self, filename: Path | str, insert_b0: bool = False) -> None:
         np.savetxt(bvecs_file, self.gradients[:3, ...].T, fmt="%.6f")
         np.savetxt(bvals_file, self.gradients[:3, ...], fmt="%.6f")
 
+    def shells(
+        self,
+        num_bins: int = DEFAULT_NUM_BINS,
+        multishell_nonempty_bin_count_thr: int = DEFAULT_MULTISHELL_BIN_COUNT_THR,
+        bval_cap: int = DEFAULT_HIGHB_THRESHOLD,
+    ) -> list:
+        """Get the shell data according to the b-value groups.
+
+        Bin the shell data according to the b-value groups found by `~find_shelling_scheme`.
+
+        Parameters
+        ----------
+        num_bins : :obj:`int`, optional
+            Number of bins.
+        multishell_nonempty_bin_count_thr : :obj:`int`, optional
+            Bin count to consider a multi-shell scheme.
+        bval_cap : :obj:`int`, optional
+            Maximum b-value to be considered in a multi-shell scheme.
+
+        Returns
+        -------
+        :obj:`list`
+            Tuples of binned b-values and corresponding shell data.
+        """
+
+        _, bval_groups, bval_estimated = find_shelling_scheme(
+            self.gradients[-1, ...],
+            num_bins=num_bins,
+            multishell_nonempty_bin_count_thr=multishell_nonempty_bin_count_thr,
+            bval_cap=bval_cap,
+        )
+        indices = [
+            np.hstack(np.where(np.isin(self.gradients[-1, ...], bvals))) for bvals in bval_groups
+        ]
+        return [(bval_estimated[idx], *self[indices]) for idx, indices in enumerate(indices)]
+
 
 def load(
     filename: Path | str,

test/test_data_dmri.py

@@ -182,6 +182,46 @@ def test_equality_operator(tmp_path):
     assert round_trip_dwi_obj == dwi_obj
 
 
+def test_shells(datadir):
+    dwi_h5 = load(datadir / "dwi.h5")
+    num_bins = 3
+
+    _, expected_bval_groups, expected_bval_est = find_shelling_scheme(
+        dwi_h5.gradients[-1, ...], num_bins=num_bins
+    )
+
+    indices = [
+        np.hstack(np.where(np.isin(dwi_h5.gradients[-1, ...], bvals)))
+        for bvals in expected_bval_groups
+    ]
+    expected_dwi_data = [dwi_h5.dataobj[..., idx] for idx in indices]
+    expected_motion_affines = [
+        dwi_h5.motion_affines[idx] if dwi_h5.motion_affines else None for idx in indices
+    ]
+    expected_gradients = [dwi_h5.gradients[..., idx] for idx in indices]
+
+    shell_data = dwi_h5.shells(num_bins=num_bins)
+    obtained_bval_est, obtained_dwi_data, obtained_motion_affines, obtained_gradients = zip(
+        *shell_data, strict=True
+    )
+
+    assert len(shell_data) == num_bins
+    assert list(obtained_bval_est) == expected_bval_est
+    assert all(
+        np.allclose(arr1, arr2)
+        for arr1, arr2 in zip(list(obtained_dwi_data), expected_dwi_data, strict=True)
+    )
+    assert all(
+        (arr1 is None and arr2 is None)
+        or (arr1 is not None and arr2 is not None and np.allclose(arr1, arr2))
+        for arr1, arr2 in zip(list(obtained_motion_affines), expected_motion_affines, strict=True)
+    )
+    assert all(
+        np.allclose(arr1, arr2)
+        for arr1, arr2 in zip(list(obtained_gradients), expected_gradients, strict=True)
+    )
+
+
 @pytest.mark.parametrize(
     ("bvals", "exp_scheme", "exp_bval_groups", "exp_bval_estimated"),
     [