Merge pull request #328 from jhlegarreta/sty/gradients-misc-changes

STY: Miscellaneous changes to gradients indexing and variable naming

Author: oesteban

docs/notebooks/data_structures.ipynb

@@ -219,7 +219,7 @@
     "data, _, grad = dwi[b2000_gradientmask]\n",
     "\n",
     "# Check correct filtering\n",
-    "data.shape[-1] == b2000_gradientmask.sum(), set(grad[-1, ...]) == {2000.0}"
+    "data.shape[-1] == b2000_gradientmask.sum(), set(grad[-1, :]) == {2000.0}"
    ]
   },
   {
@@ -240,7 +240,7 @@
    ],
    "source": [
     "# Plot volume 10 of b=2000 shell\n",
-    "plot_dwi(data[..., 10], dwi.affine, gradient=grad[..., 10]);"
+    "plot_dwi(data[..., 10], dwi.affine, gradient=grad[:, 10]);"
    ]
   }
  ],

src/nifreeze/data/dmri.py

@@ -157,11 +157,11 @@ def from_filename(cls, filename: Path | str) -> Self:
 
     @property
     def bvals(self):
-        return self.gradients[..., -1]
+        return self.gradients[:, -1]
 
     @property
     def bvecs(self):
-        return self.gradients[..., :-1]
+        return self.gradients[:, :-1]
 
     def get_shells(
         self,

src/nifreeze/data/filtering.py

@@ -254,21 +254,21 @@ def dwi_select_shells(
         Shell mask.
     """
 
-    bvalues = gradients[:, -1]
-    bcenter = bvalues[index]
+    bvals = gradients[:, -1]
+    bcenter = bvals[index]
 
-    shellmask = np.ones(len(bvalues), dtype=bool)
+    shellmask = np.ones(len(bvals), dtype=bool)
     shellmask[index] = False  # Drop the held-out index
 
     if atol_low is None and atol_high is None:
         return shellmask
 
     atol_low = 0 if atol_low is None else atol_low
-    atol_high = gradients[:, -1].max() if atol_high is None else atol_high
+    atol_high = bvals.max() if atol_high is None else atol_high
 
     # Keep only b-values within the range defined by atol_high and atol_low
-    shellmask[bvalues > (bcenter + atol_high)] = False
-    shellmask[bvalues < (bcenter - atol_low)] = False
+    shellmask[bvals > (bcenter + atol_high)] = False
+    shellmask[bvals < (bcenter - atol_low)] = False
 
     if not shellmask.sum():
         raise RuntimeError(f"Shell corresponding to index {index} (b={bcenter}) is empty.")