Merge pull request #495 from MarcCote/enh_tractogram_operators

MRG: Support + and += operators for Tractogram

This PR adds the functionality of concatenating two Tractogram objects using either `tractogram += other_tractogram` or `tractogram = tractogram1 + tractogram2`.

Author: matthew-brett

nibabel/streamlines/tests/test_tractogram.py

@@ -1,7 +1,10 @@
 import sys
+import copy
 import unittest
 import numpy as np
 import warnings
+import operator
+from collections import defaultdict
 
 from nibabel.testing import assert_arrays_equal
 from nibabel.testing import clear_and_catch_warnings
@@ -16,37 +19,99 @@
 DATA = {}
 
 
+def make_fake_streamline(nb_points, data_per_point_shapes={},
+                         data_for_streamline_shapes={}, rng=None):
+    """ Make a single streamline according to provided requirements. """
+    if rng is None:
+        rng = np.random.RandomState()
+
+    streamline = rng.randn(nb_points, 3).astype("f4")
+
+    data_per_point = {}
+    for k, shape in data_per_point_shapes.items():
+        data_per_point[k] = rng.randn(*((nb_points,) + shape)).astype("f4")
+
+    data_for_streamline = {}
+    for k, shape in data_for_streamline.items():
+        data_for_streamline[k] = rng.randn(*shape).astype("f4")
+
+    return streamline, data_per_point, data_for_streamline
+
+
+def make_fake_tractogram(list_nb_points, data_per_point_shapes={},
+                         data_for_streamline_shapes={}, rng=None):
+    """ Make multiple streamlines according to provided requirements. """
+    all_streamlines = []
+    all_data_per_point = defaultdict(lambda: [])
+    all_data_per_streamline = defaultdict(lambda: [])
+    for nb_points in list_nb_points:
+        data = make_fake_streamline(nb_points, data_per_point_shapes,
+                                    data_for_streamline_shapes, rng)
+        streamline, data_per_point, data_for_streamline = data
+
+        all_streamlines.append(streamline)
+        for k, v in data_per_point.items():
+            all_data_per_point[k].append(v)
+
+        for k, v in data_for_streamline.items():
+            all_data_per_streamline[k].append(v)
+
+    return all_streamlines, all_data_per_point, all_data_per_streamline
+
+
+def make_dummy_streamline(nb_points):
+    """ Make the streamlines that have been used to create test data files."""
+    if nb_points == 1:
+        streamline = np.arange(1*3, dtype="f4").reshape((1, 3))
+        data_per_point = {"fa": np.array([[0.2]], dtype="f4"),
+                          "colors": np.array([(1, 0, 0)]*1, dtype="f4")}
+        data_for_streamline = {"mean_curvature": np.array([1.11], dtype="f4"),
+                               "mean_torsion": np.array([1.22], dtype="f4"),
+                               "mean_colors": np.array([1, 0, 0], dtype="f4")}
+
+    elif nb_points == 2:
+        streamline = np.arange(2*3, dtype="f4").reshape((2, 3))
+        data_per_point = {"fa": np.array([[0.3],
+                                          [0.4]], dtype="f4"),
+                          "colors": np.array([(0, 1, 0)]*2, dtype="f4")}
+        data_for_streamline = {"mean_curvature": np.array([2.11], dtype="f4"),
+                               "mean_torsion": np.array([2.22], dtype="f4"),
+                               "mean_colors": np.array([0, 1, 0], dtype="f4")}
+
+    elif nb_points == 5:
+        streamline = np.arange(5*3, dtype="f4").reshape((5, 3))
+        data_per_point = {"fa": np.array([[0.5],
+                                          [0.6],
+                                          [0.6],
+                                          [0.7],
+                                          [0.8]], dtype="f4"),
+                          "colors": np.array([(0, 0, 1)]*5, dtype="f4")}
+        data_for_streamline = {"mean_curvature": np.array([3.11], dtype="f4"),
+                               "mean_torsion": np.array([3.22], dtype="f4"),
+                               "mean_colors": np.array([0, 0, 1], dtype="f4")}
+
+    return streamline, data_per_point, data_for_streamline
+
+
 def setup():
     global DATA
     DATA['rng'] = np.random.RandomState(1234)
-    DATA['streamlines'] = [np.arange(1*3, dtype="f4").reshape((1, 3)),
-                           np.arange(2*3, dtype="f4").reshape((2, 3)),
-                           np.arange(5*3, dtype="f4").reshape((5, 3))]
-
-    DATA['fa'] = [np.array([[0.2]], dtype="f4"),
-                  np.array([[0.3],
-                            [0.4]], dtype="f4"),
-                  np.array([[0.5],
-                            [0.6],
-                            [0.6],
-                            [0.7],
-                            [0.8]], dtype="f4")]
-
-    DATA['colors'] = [np.array([(1, 0, 0)]*1, dtype="f4"),
-                      np.array([(0, 1, 0)]*2, dtype="f4"),
-                      np.array([(0, 0, 1)]*5, dtype="f4")]
-
-    DATA['mean_curvature'] = [np.array([1.11], dtype="f4"),
-                              np.array([2.11], dtype="f4"),
-                              np.array([3.11], dtype="f4")]
-
-    DATA['mean_torsion'] = [np.array([1.22], dtype="f4"),
-                            np.array([2.22], dtype="f4"),
-                            np.array([3.22], dtype="f4")]
-
-    DATA['mean_colors'] = [np.array([1, 0, 0], dtype="f4"),
-                           np.array([0, 1, 0], dtype="f4"),
-                           np.array([0, 0, 1], dtype="f4")]
+
+    DATA['streamlines'] = []
+    DATA['fa'] = []
+    DATA['colors'] = []
+    DATA['mean_curvature'] = []
+    DATA['mean_torsion'] = []
+    DATA['mean_colors'] = []
+    for nb_points in [1, 2, 5]:
+        data = make_dummy_streamline(nb_points)
+        streamline, data_per_point, data_for_streamline = data
+        DATA['streamlines'].append(streamline)
+        DATA['fa'].append(data_per_point['fa'])
+        DATA['colors'].append(data_per_point['colors'])
+        DATA['mean_curvature'].append(data_for_streamline['mean_curvature'])
+        DATA['mean_torsion'].append(data_for_streamline['mean_torsion'])
+        DATA['mean_colors'].append(data_for_streamline['mean_colors'])
 
     DATA['data_per_point'] = {'colors': DATA['colors'],
                               'fa': DATA['fa']}
@@ -63,17 +128,13 @@ def setup():
                                     affine_to_rasmm=np.eye(4))
 
     DATA['streamlines_func'] = lambda: (e for e in DATA['streamlines'])
-    fa_func = lambda: (e for e in DATA['fa'])
-    colors_func = lambda: (e for e in DATA['colors'])
-    mean_curvature_func = lambda: (e for e in DATA['mean_curvature'])
-    mean_torsion_func = lambda: (e for e in DATA['mean_torsion'])
-    mean_colors_func = lambda: (e for e in DATA['mean_colors'])
-
-    DATA['data_per_point_func'] = {'colors': colors_func,
-                                   'fa': fa_func}
-    DATA['data_per_streamline_func'] = {'mean_curvature': mean_curvature_func,
-                                        'mean_torsion': mean_torsion_func,
-                                        'mean_colors': mean_colors_func}
+    DATA['data_per_point_func'] = {
+        'colors': lambda: (e for e in DATA['colors']),
+        'fa': lambda: (e for e in DATA['fa'])}
+    DATA['data_per_streamline_func'] = {
+        'mean_curvature': lambda: (e for e in DATA['mean_curvature']),
+        'mean_torsion': lambda: (e for e in DATA['mean_torsion']),
+        'mean_colors': lambda: (e for e in DATA['mean_colors'])}
 
     DATA['lazy_tractogram'] = LazyTractogram(DATA['streamlines_func'],
                                              DATA['data_per_streamline_func'],
@@ -130,6 +191,11 @@ def assert_tractogram_equal(t1, t2):
                      t2.data_per_streamline, t2.data_per_point)
 
 
+def extender(a, b):
+    a.extend(b)
+    return a
+
+
 class TestPerArrayDict(unittest.TestCase):
 
     def test_per_array_dict_creation(self):
@@ -181,6 +247,53 @@ def test_getitem(self):
             assert_arrays_equal(sdict[-1][k], v[-1])
             assert_arrays_equal(sdict[[0, -1]][k], v[[0, -1]])
 
+    def test_extend(self):
+        sdict = PerArrayDict(len(DATA['tractogram']),
+                             DATA['data_per_streamline'])
+
+        new_data = {'mean_curvature': 2 * np.array(DATA['mean_curvature']),
+                    'mean_torsion': 3 * np.array(DATA['mean_torsion']),
+                    'mean_colors': 4 * np.array(DATA['mean_colors'])}
+        sdict2 = PerArrayDict(len(DATA['tractogram']),
+                              new_data)
+
+        sdict.extend(sdict2)
+        assert_equal(len(sdict), len(sdict2))
+        for k in DATA['tractogram'].data_per_streamline:
+            assert_arrays_equal(sdict[k][:len(DATA['tractogram'])],
+                                DATA['tractogram'].data_per_streamline[k])
+            assert_arrays_equal(sdict[k][len(DATA['tractogram']):],
+                                new_data[k])
+
+        # Extending with an empty PerArrayDict should change nothing.
+        sdict_orig = copy.deepcopy(sdict)
+        sdict.extend(PerArrayDict())
+        for k in sdict_orig.keys():
+            assert_arrays_equal(sdict[k], sdict_orig[k])
+
+        # Test incompatible PerArrayDicts.
+        # Other dict has more entries.
+        new_data = {'mean_curvature': 2 * np.array(DATA['mean_curvature']),
+                    'mean_torsion': 3 * np.array(DATA['mean_torsion']),
+                    'mean_colors': 4 * np.array(DATA['mean_colors']),
+                    'other': 5 * np.array(DATA['mean_colors'])}
+        sdict2 = PerArrayDict(len(DATA['tractogram']), new_data)
+        assert_raises(ValueError, sdict.extend, sdict2)
+
+        # Other dict has not the same entries (key mistmached).
+        new_data = {'mean_curvature': 2 * np.array(DATA['mean_curvature']),
+                    'mean_torsion': 3 * np.array(DATA['mean_torsion']),
+                    'other': 4 * np.array(DATA['mean_colors'])}
+        sdict2 = PerArrayDict(len(DATA['tractogram']), new_data)
+        assert_raises(ValueError, sdict.extend, sdict2)
+
+        # Other dict has the right number of entries but wrong shape.
+        new_data = {'mean_curvature': 2 * np.array(DATA['mean_curvature']),
+                    'mean_torsion': 3 * np.array(DATA['mean_torsion']),
+                    'mean_colors': 4 * np.array(DATA['mean_torsion'])}
+        sdict2 = PerArrayDict(len(DATA['tractogram']), new_data)
+        assert_raises(ValueError, sdict.extend, sdict2)
+
 
 class TestPerArraySequenceDict(unittest.TestCase):
 
@@ -233,6 +346,62 @@ def test_getitem(self):
             assert_arrays_equal(sdict[-1][k], v[-1])
             assert_arrays_equal(sdict[[0, -1]][k], v[[0, -1]])
 
+    def test_extend(self):
+        total_nb_rows = DATA['tractogram'].streamlines.total_nb_rows
+        sdict = PerArraySequenceDict(total_nb_rows, DATA['data_per_point'])
+
+        # Test compatible PerArraySequenceDicts.
+        list_nb_points = [2, 7, 4]
+        data_per_point_shapes = {"colors": DATA['colors'][0].shape[1:],
+                                 "fa": DATA['fa'][0].shape[1:]}
+        _, new_data, _ = make_fake_tractogram(list_nb_points,
+                                              data_per_point_shapes,
+                                              rng=DATA['rng'])
+        sdict2 = PerArraySequenceDict(np.sum(list_nb_points), new_data)
+
+        sdict.extend(sdict2)
+        assert_equal(len(sdict), len(sdict2))
+        for k in DATA['tractogram'].data_per_point:
+            assert_arrays_equal(sdict[k][:len(DATA['tractogram'])],
+                                DATA['tractogram'].data_per_point[k])
+            assert_arrays_equal(sdict[k][len(DATA['tractogram']):],
+                                new_data[k])
+
+        # Extending with an empty PerArraySequenceDicts should change nothing.
+        sdict_orig = copy.deepcopy(sdict)
+        sdict.extend(PerArraySequenceDict())
+        for k in sdict_orig.keys():
+            assert_arrays_equal(sdict[k], sdict_orig[k])
+
+        # Test incompatible PerArraySequenceDicts.
+        # Other dict has more entries.
+        data_per_point_shapes = {"colors": DATA['colors'][0].shape[1:],
+                                 "fa": DATA['fa'][0].shape[1:],
+                                 "other": (7,)}
+        _, new_data, _ = make_fake_tractogram(list_nb_points,
+                                              data_per_point_shapes,
+                                              rng=DATA['rng'])
+        sdict2 = PerArraySequenceDict(np.sum(list_nb_points), new_data)
+        assert_raises(ValueError, sdict.extend, sdict2)
+
+        # Other dict has not the same entries (key mistmached).
+        data_per_point_shapes = {"colors": DATA['colors'][0].shape[1:],
+                                 "other": DATA['fa'][0].shape[1:]}
+        _, new_data, _ = make_fake_tractogram(list_nb_points,
+                                              data_per_point_shapes,
+                                              rng=DATA['rng'])
+        sdict2 = PerArraySequenceDict(np.sum(list_nb_points), new_data)
+        assert_raises(ValueError, sdict.extend, sdict2)
+
+        # Other dict has the right number of entries but wrong shape.
+        data_per_point_shapes = {"colors": DATA['colors'][0].shape[1:],
+                                 "fa": DATA['fa'][0].shape[1:] + (3,)}
+        _, new_data, _ = make_fake_tractogram(list_nb_points,
+                                              data_per_point_shapes,
+                                              rng=DATA['rng'])
+        sdict2 = PerArraySequenceDict(np.sum(list_nb_points), new_data)
+        assert_raises(ValueError, sdict.extend, sdict2)
+
 
 class TestLazyDict(unittest.TestCase):
 
@@ -570,6 +739,28 @@ def test_tractogram_to_world(self):
         tractogram.affine_to_rasmm = None
         assert_raises(ValueError, tractogram.to_world)
 
+    def test_tractogram_extend(self):
+        # Load tractogram that contains some metadata.
+        t = DATA['tractogram'].copy()
+
+        for op, in_place in ((operator.add, False), (operator.iadd, True),
+                             (extender, True)):
+            first_arg = t.copy()
+            new_t = op(first_arg, t)
+            assert_equal(new_t is first_arg, in_place)
+            assert_tractogram_equal(new_t[:len(t)], DATA['tractogram'])
+            assert_tractogram_equal(new_t[len(t):], DATA['tractogram'])
+
+        # Test extending an empty Tractogram.
+        t = Tractogram()
+        t += DATA['tractogram']
+        assert_tractogram_equal(t, DATA['tractogram'])
+
+        # and the other way around.
+        t = DATA['tractogram'].copy()
+        t += Tractogram()
+        assert_tractogram_equal(t, DATA['tractogram'])
+
 
 class TestLazyTractogram(unittest.TestCase):
 
@@ -580,11 +771,12 @@ def test_lazy_tractogram_creation(self):
         # Streamlines and other data as generators
         streamlines = (x for x in DATA['streamlines'])
         data_per_point = {"colors": (x for x in DATA['colors'])}
-        data_per_streamline = {'mean_torsion': (x for x in DATA['mean_torsion']),
-                               'mean_colors': (x for x in DATA['mean_colors'])}
+        data_per_streamline = {'torsion': (x for x in DATA['mean_torsion']),
+                               'colors': (x for x in DATA['mean_colors'])}
 
         # Creating LazyTractogram with generators is not allowed as
-        # generators get exhausted and are not reusable unlike generator function.
+        # generators get exhausted and are not reusable unlike generator
+        # function.
         assert_raises(TypeError, LazyTractogram, streamlines)
         assert_raises(TypeError, LazyTractogram,
                       data_per_streamline=data_per_streamline)
@@ -610,12 +802,11 @@ def test_lazy_tractogram_creation(self):
 
     def test_lazy_tractogram_from_data_func(self):
         # Create an empty `LazyTractogram` yielding nothing.
-        _empty_data_gen = lambda: iter([])
-
-        tractogram = LazyTractogram.from_data_func(_empty_data_gen)
+        tractogram = LazyTractogram.from_data_func(lambda: iter([]))
         check_tractogram(tractogram)
 
-        # Create `LazyTractogram` from a generator function yielding TractogramItem.
+        # Create `LazyTractogram` from a generator function yielding
+        # TractogramItem.
         data = [DATA['streamlines'], DATA['fa'], DATA['colors'],
                 DATA['mean_curvature'], DATA['mean_torsion'],
                 DATA['mean_colors']]
@@ -641,6 +832,13 @@ def test_lazy_tractogram_getitem(self):
         assert_raises(NotImplementedError,
                       DATA['lazy_tractogram'].__getitem__, 0)
 
+    def test_lazy_tractogram_extend(self):
+        t = DATA['lazy_tractogram'].copy()
+        new_t = DATA['lazy_tractogram'].copy()
+
+        for op in (operator.add, operator.iadd, extender):
+            assert_raises(NotImplementedError, op, new_t, t)
+
     def test_lazy_tractogram_len(self):
         modules = [module_tractogram]  # Modules for which to catch warnings.
         with clear_and_catch_warnings(record=True, modules=modules) as w:
@@ -746,8 +944,8 @@ def test_lazy_tractogram_copy(self):
         # Check we copied the data and not simply created new references.
         assert_true(tractogram is not DATA['lazy_tractogram'])
 
-        # When copying LazyTractogram, the generator function yielding streamlines
-        # should stay the same.
+        # When copying LazyTractogram, the generator function yielding
+        # streamlines should stay the same.
         assert_true(tractogram._streamlines
                     is DATA['lazy_tractogram']._streamlines)
 
@@ -759,12 +957,14 @@ def test_lazy_tractogram_copy(self):
                     is not DATA['lazy_tractogram']._data_per_point)
 
         for key in tractogram.data_per_streamline:
-            assert_true(tractogram.data_per_streamline.store[key]
-                        is DATA['lazy_tractogram'].data_per_streamline.store[key])
+            data = tractogram.data_per_streamline.store[key]
+            expected = DATA['lazy_tractogram'].data_per_streamline.store[key]
+            assert_true(data is expected)
 
         for key in tractogram.data_per_point:
-            assert_true(tractogram.data_per_point.store[key]
-                        is DATA['lazy_tractogram'].data_per_point.store[key])
+            data = tractogram.data_per_point.store[key]
+            expected = DATA['lazy_tractogram'].data_per_point.store[key]
+            assert_true(data is expected)
 
         # The affine should be a copy.
         assert_true(tractogram._affine_to_apply