Add Tractogram concatenation

Author: MarcCote

nibabel/streamlines/tests/test_tractogram.py

@@ -570,6 +570,17 @@ 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()
+        new_t = DATA['tractogram'].copy()
+
+        # Double the tractogram.
+        new_t += t
+        assert_equal(len(new_t), 2*len(t))
+        assert_tractogram_equal(new_t[:len(t)], DATA['tractogram'])
+        assert_tractogram_equal(new_t[len(t):], DATA['tractogram'])
+
 
 class TestLazyTractogram(unittest.TestCase):
 
@@ -641,6 +652,11 @@ 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()
+        assert_raises(NotImplementedError, new_t.__iadd__, 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:

nibabel/streamlines/tractogram.py

@@ -112,6 +112,28 @@ def __setitem__(self, key, value):
 
         self.store[key] = value
 
+    def extend(self, other):
+        if len(self) != len(other):
+            msg = ("Size mismatched between the two PerArrayDict objects."
+                   " This PerArrayDict has {0} elements whereas the other "
+                   " has {1} elements.").format(len(self), len(other))
+            raise ValueError(msg)
+
+        if sorted(self.keys()) != sorted(other.keys()):
+            msg = ("Key mismatched between the two PerArrayDict objects."
+                   " This PerArrayDict contains '{0}' whereas the other "
+                   " contains '{1}'.").format(sorted(self.keys()),
+                                              sorted(other.keys()))
+            raise ValueError(msg)
+
+        self.n_rows += other.n_rows
+        for key in self.keys():
+            self[key] = np.concatenate([self[key], other[key]])
+
+    def __iadd__(self, other):
+        self.extend(other)
+        return self
+
 
 class PerArraySequenceDict(PerArrayDict):
     """ Dictionary for which key access can do slicing on the values.
@@ -136,6 +158,28 @@ def __setitem__(self, key, value):
 
         self.store[key] = value
 
+    def extend(self, other):
+        if len(self) != len(other):
+            msg = ("Size mismatched between the two PerArrayDict objects."
+                   " This PerArrayDict has {0} elements whereas the other "
+                   " has {1} elements.").format(len(self), len(other))
+            raise ValueError(msg)
+
+        if sorted(self.keys()) != sorted(other.keys()):
+            msg = ("Key mismatched between the two PerArrayDict objects."
+                   " This PerArrayDict contains '{0}' whereas the other "
+                   " contains '{1}'.").format(sorted(self.keys()),
+                                              sorted(other.keys()))
+            raise ValueError(msg)
+
+        self.n_rows += other.n_rows
+        for key in self.keys():
+            self[key].extend(other[key])
+
+    def __iadd__(self, other):
+        self.extend(other)
+        return self
+
 
 class LazyDict(collections.MutableMapping):
     """ Dictionary of generator functions.
@@ -418,6 +462,16 @@ def to_world(self, lazy=False):
 
         return self.apply_affine(self.affine_to_rasmm, lazy=lazy)
 
+    def extend(self, other):
+        # TODO: Make sure the other tractogram is compatible.
+        self.streamlines.extend(other.streamlines)
+        self.data_per_streamline += other.data_per_streamline
+        self.data_per_point += other.data_per_point
+
+    def __iadd__(self, other):
+        self.extend(other)
+        return self
+
 
 class LazyTractogram(Tractogram):
     """ Lazy container for streamlines and their data information.
@@ -653,6 +707,9 @@ def _gen_data():
     def __getitem__(self, idx):
         raise NotImplementedError('LazyTractogram does not support indexing.')
 
+    def extend(self, other):
+        raise NotImplementedError('LazyTractogram does not support concatenation.')
+
     def __iter__(self):
         count = 0
         for tractogram_item in self.data: