ENH: Flesh out NdGrid implementation

Author: effigies

nibabel/pointset.py

@@ -1,3 +1,11 @@
+import operator as op
+from functools import reduce
+
+import numpy as np
+
+from nibabel.affines import apply_affine
+
+
 class Pointset:
     @property
     def n_coords(self):
@@ -51,6 +59,32 @@ class NdGrid(Pointset):
         number of coordinates in each dimension of grid
     """
 
+    def __init__(self, shape, affines):
+        self.shape = tuple(shape)
+        try:
+            self._affines = dict(affines)
+        except (TypeError, ValueError):
+            self._affines = {'world': np.array(affines)}
+        if 'voxels' not in self._affines:
+            self._affines['voxels'] = np.eye(4, dtype=np.uint8)
+
     def get_affine(self, name=None):
         """4x4 array"""
-        raise NotImplementedError
+        if name is None:
+            name = next(iter(self._affines))
+        return self._affines[name]
+
+    def get_coords(self, name=None):
+        if name is None:
+            name = next(iter(self._affines))
+        aff = self.get_affine(name)
+        dt = np.result_type(*(np.min_scalar_type(dim) for dim in self.shape))
+        # This is pretty wasteful; we almost certainly want instead an
+        # object that will retrieve a coordinate when indexed, but where
+        # np.array(obj) returns this
+        ijk_coords = np.array(list(np.ndindex(self.shape)), dtype=dt)
+        return apply_affine(aff, ijk_coords)
+
+    @property
+    def n_coords(self):
+        return reduce(op.mul, self.shape)