ENH: Get the nd_sort method mostly working w/ basic tests

Author: moloney

nibabel/metasum.py

@@ -100,6 +100,15 @@ def __getitem__(self, value):
         ba = self._val_bitarrs[value]
         return list(self._extract_indices(ba))
 
+    def first(self, value):
+        '''Return the first index where this value appears'''
+        if self._const_val == value:
+            return 0
+        idx = self._unique_vals.get(value)
+        if idx is not None:
+            return idx
+        return self._val_bitarrs[value].index(True)
+
     def values(self):
         '''Generate each unique value that has been seen'''
         if self._const_val is not _NoValue:
@@ -339,8 +348,15 @@ def get_block_size(self):
                 return None
         return block_size
 
-    def is_subpartition(self, other):
-        ''''''
+    def is_orthogonal(self, other, size=1):
+        '''Check our value's indices overlaps each from `other` exactly `size` times
+        '''
+        other_bas = {v: other.get_mask(v) for v in other.values()}
+        for val in self.values():
+            for other_val, other_ba in other_bas.items():
+                if self.count(val, mask=other_ba) != size:
+                    return False
+        return True
 
     def _extract_indices(self, ba):
         '''Generate integer indices from bitarray representation'''
@@ -416,7 +432,7 @@ def _extend_const(self, other):
 
 
 class DimTypes(IntEnum):
-    '''Enmerate the three types of nD dimensions'''
+    '''Enumerate the three types of nD dimensions'''
     SLICE = 1
     TIME = 2
     PARAM = 3
@@ -556,8 +572,9 @@ def nd_sort(self, dims):
             last_dim = dim
 
         # Pull out info about different types of dims
-        n_slices = None
-        n_vol = None
+        n_input = self._n_input
+        total_vol = None
+        slice_dim = None
         time_dim = None
         param_dims = []
         n_params = []
@@ -568,9 +585,10 @@ def nd_sort(self, dims):
             dim_vidx = self._v_idxs[dim.key]
             dim_type = dim.dim_type
             if dim_type is DimTypes.SLICE:
+                slice_dim = dim
                 n_slices = len(dim_vidx)
-                n_vol = dim_vidx.get_block_size()
-                if n_vol is None:
+                total_vol = dim_vidx.get_block_size()
+                if total_vol is None:
                     raise NdSortError("There are missing or extra slices")
                 shape.append(n_slices)
                 curr_size *= n_slices
@@ -583,29 +601,39 @@ def nd_sort(self, dims):
                 n_param = len(dim_vidx)
                 n_params.append(n_param)
                 total_params *= n_param
-        if n_vol is None:
-            n_vol = self._n_input
+        if total_vol is None:
+            total_vol = n_input
 
-        # Size of the time dimension must be infered from the size of the other dims
+        # Size of the time dimension must be inferred from the size of the other dims
         n_time = 1
+        prev_dim = slice_dim
         if time_dim is not None:
-            n_time, rem = divmod(n_vol, total_params)
+            n_time, rem = divmod(total_vol, total_params)
             if rem != 0:
-                raise NdSortError(f"The combined parameters don't evenly divide inputs")
+                raise NdSortError("The combined parameters don't evenly divide inputs")
             shape.append(n_time)
             curr_size *= n_time
+            prev_dim = time_dim
 
-        # Complete the "shape", and do a more detailed check that our param dims make sense
+        # Complete the "shape", and do a more detailed check that our dims make sense
         for dim, n_param in zip(param_dims, n_params):
             dim_vidx = self._v_idxs[dim.key]
-            if dim_vidx.get_block_size() != curr_size:
+            if dim_vidx.get_block_size() != n_input // n_param:
                 raise NdSortError(f"The parameter {dim.key} doesn't evenly divide inputs")
+            if prev_dim is not None and prev_dim.dim_type != DimTypes.TIME:
+                count_per = (curr_size // shape[-1]) * (n_input // (curr_size * n_param))
+                if not self._v_idxs[prev_dim.key].is_orthogonal(dim_vidx,  count_per):
+                    raise NdSortError("The dimensions are not orthogonal")
             shape.append(n_param)
             curr_size *= n_param
+            prev_dim = dim
 
         # Extract dim keys for each input and do the actual sort
         sort_keys = [(idx, tuple(self.get_val(idx, dim.key) for dim in reversed(dims)))
-                     for idx in range(self._n_input)]
+                     for idx in range(n_input)]
         sort_keys.sort(key=lambda x: x[1])
 
-        # TODO: Finish this
+        # TODO: If we have non-singular time dimension we need to do some additional
+        #       validation checks here after sorting.
+
+        return tuple(shape), [x[0] for x in sort_keys]

nibabel/tests/test_metasum.py

@@ -1,6 +1,9 @@
-from ..metasum import MetaSummary, ValueIndices
+import random
 
 import pytest
+import numpy as np
+
+from ..metasum import DimIndex, DimTypes, MetaSummary, ValueIndices
 
 
 vidx_test_patterns = ([0] * 8,
@@ -14,15 +17,15 @@
 
 @pytest.mark.parametrize("in_list", vidx_test_patterns)
 def test_value_indices_basics(in_list):
-    '''Test we can roundtrip list -> ValueIndices -> list'''
+    '''Test basic ValueIndices behavior'''
     vidx = ValueIndices(in_list)
     assert vidx.n_input == len(in_list)
     assert len(vidx) == len(set(in_list))
     assert sorted(vidx.values()) == sorted(list(set(in_list)))
     for val in vidx.values():
         assert vidx.count(val) == in_list.count(val)
         for in_idx in vidx[val]:
-            assert in_list[in_idx] == val
+            assert in_list[in_idx] == val == vidx.get_value(in_idx)
     out_list = vidx.to_list()
     assert in_list == out_list
 
@@ -78,3 +81,73 @@ def test_meta_summary_basics(in_dicts):
         assert out_dict == in_dicts[in_idx]
         for key, in_val in in_dicts[in_idx].items():
             assert in_val == msum.get_val(in_idx, key)
+
+
+def _make_nd_meta(shape, dim_info, const_meta=None):
+    if const_meta is None:
+        const_meta = {'series_number': '5'}
+    meta_seq = []
+    for nd_idx in np.ndindex(*shape):
+        curr_meta = {}
+        curr_meta.update(const_meta)
+        for dim, dim_idx in zip(dim_info, nd_idx):
+            curr_meta[dim.key] = dim_idx
+        meta_seq.append(curr_meta)
+    return meta_seq
+
+
+ndsort_test_args = (((3,),
+                     (DimIndex(DimTypes.SLICE, 'slice_location'),),
+                     None),
+                    ((3, 5),
+                     (DimIndex(DimTypes.SLICE, 'slice_location'),
+                      DimIndex(DimTypes.TIME, 'acq_time')),
+                     None),
+                    ((3, 5),
+                     (DimIndex(DimTypes.SLICE, 'slice_location'),
+                      DimIndex(DimTypes.PARAM, 'inversion_time')),
+                     None),
+                    ((3, 5, 7),
+                     (DimIndex(DimTypes.SLICE, 'slice_location'),
+                      DimIndex(DimTypes.TIME, 'acq_time'),
+                      DimIndex(DimTypes.PARAM, 'echo_time')),
+                     None),
+                    ((3, 5, 7),
+                     (DimIndex(DimTypes.SLICE, 'slice_location'),
+                      DimIndex(DimTypes.PARAM, 'inversion_time'),
+                      DimIndex(DimTypes.PARAM, 'echo_time')),
+                     None),
+                    ((5, 3),
+                     (DimIndex(DimTypes.TIME, 'acq_time'),
+                      DimIndex(DimTypes.PARAM, 'echo_time')),
+                     None),
+                    ((3, 5, 7),
+                     (DimIndex(DimTypes.TIME, 'acq_time'),
+                      DimIndex(DimTypes.PARAM, 'inversion_time'),
+                      DimIndex(DimTypes.PARAM, 'echo_time')),
+                     None),
+                    ((5, 7),
+                     (DimIndex(DimTypes.PARAM, 'inversion_time'),
+                      DimIndex(DimTypes.PARAM, 'echo_time')),
+                     None),
+                    ((5, 7, 3),
+                     (DimIndex(DimTypes.PARAM, 'inversion_time'),
+                      DimIndex(DimTypes.PARAM, 'echo_time'),
+                      DimIndex(DimTypes.PARAM, 'repetition_time')),
+                     None),
+                    )
+
+
+@pytest.mark.parametrize("shape,dim_info,const_meta", ndsort_test_args)
+def test_ndsort(shape, dim_info, const_meta):
+    meta_seq = _make_nd_meta(shape, dim_info, const_meta)
+    rand_idx_seq = [(i, m) for i, m in enumerate(meta_seq)]
+    # TODO: Use some pytest plugin to manage randomness?  Just use fixed seed?
+    random.shuffle(rand_idx_seq)
+    rand_idx = [x[0] for x in rand_idx_seq]
+    rand_seq = [x[1] for x in rand_idx_seq]
+    msum = MetaSummary()
+    for meta in rand_seq:
+        msum.append(meta)
+    out_shape, out_idxs = msum.nd_sort(dim_info)
+    assert shape == out_shape