ENH: Add CoordinateImage slicing by parcel name

Author: effigies

nibabel/coordimage.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 import nibabel as nib
 from nibabel.fileslice import fill_slicer
 import nibabel.pointset as ps
@@ -16,6 +18,26 @@ def __init__(self, data, coordaxis, header=None):
         self.coordaxis = coordaxis
         self.header = header
 
+    @property
+    def shape(self):
+        return self.data.shape
+
+    def __getitem__(self, slicer):
+        if isinstance(slicer, str):
+            slicer = self.coordaxis.get_indices(slicer)
+        elif isinstance(slicer, list):
+            slicer = np.hstack([self.coordaxis.get_indices(sub) for sub in slicer])
+
+        if isinstance(slicer, range):
+            slicer = slice(slicer.start, slicer.stop, slicer.step)
+
+        data = self.data
+        if not isinstance(slicer, slice):
+            data = np.asanyarray(data)
+        return self.__class__(data[slicer],
+                              self.coordaxis[slicer],
+                              header=self.header.copy())
+
     @classmethod
     def from_image(klass, img):
         coordaxis = CoordinateAxis.from_header(img.header)
@@ -53,7 +75,7 @@ def __getitem__(self, slicer):
         Return a sub-sampled CoordinateAxis containing structures
         matching the indices provided.
         """
-        if slicer is Ellipsis or slicer == slice(None):
+        if slicer is Ellipsis or isinstance(slicer, slice) and slicer == slice(None):
             return self
         elif isinstance(slicer, slice):
             slicer = fill_slicer(slicer, len(self))

nibabel/tests/test_coordimage.py

@@ -68,4 +68,11 @@ def test_Cifti2Image_as_CoordImage():
     assert len(subaxis) == len(caxis) - 200
     assert len(subaxis.parcels) == len(caxis.parcels)
 
-    caxis.get_indices("CIFTI_STRUCTURE_CORTEX_LEFT")
+    lh_img = cimg["CIFTI_STRUCTURE_CORTEX_LEFT"]
+    assert len(lh_img.coordaxis.parcels) == 1
+    assert lh_img.shape == (29696, 1)
+
+    # # Not working yet.
+    # cortex_img = cimg[["CIFTI_STRUCTURE_CORTEX_LEFT", "CIFTI_STRUCTURE_CORTEX_RIGHT"]]
+    # assert len(cortex_img.coordaxis.parcels) == 2
+    # assert cortex_img.shape == (59412, 1)