Supports extending empty tractograms

Author: MarcCote

nibabel/streamlines/tests/test_tractogram.py

@@ -1,4 +1,5 @@
 import sys
+import copy
 import unittest
 import numpy as np
 import warnings
@@ -264,10 +265,13 @@ def test_extend(self):
             assert_arrays_equal(sdict[k][len(DATA['tractogram']):],
                                 new_data[k])
 
-        # Test incompatible PerArrayDicts.
-        # Other dict is missing entries.
-        assert_raises(ValueError, sdict.extend, PerArrayDict())
+        # Extending with an empty PerArrayDicts 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']),
@@ -356,10 +360,13 @@ def test_extend(self):
             assert_arrays_equal(sdict[k][len(DATA['tractogram']):],
                                 new_data[k])
 
-        # Test incompatible PerArrayDicts.
-        # Other dict is missing entries.
-        assert_raises(ValueError, sdict.extend, PerArraySequenceDict())
+        # Extending with an empty PerArrayDicts 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 PerArrayDicts.
         # Other dict has more entries.
         data_per_point_shapes = {"colors": DATA['colors'][0].shape[1:],
                                  "fa": DATA['fa'][0].shape[1:],
@@ -728,6 +735,16 @@ def test_tractogram_extend(self):
             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):
 

nibabel/streamlines/tractogram.py

@@ -59,6 +59,9 @@ def __getitem__(self, key):
         # Key was not a valid index/slice after all.
         return self.store[key]  # Will raise the proper error.
 
+    def __contains__(self, key):
+        return key in self.store
+
     def __delitem__(self, key):
         del self.store[key]
 
@@ -90,7 +93,7 @@ class PerArrayDict(SliceableDataDict):
         Positional and keyword arguments, passed straight through the ``dict``
         constructor.
     """
-    def __init__(self, n_rows=None, *args, **kwargs):
+    def __init__(self, n_rows=0, *args, **kwargs):
         self.n_rows = n_rows
         super(PerArrayDict, self).__init__(*args, **kwargs)
 
@@ -105,13 +108,17 @@ def __setitem__(self, key, value):
             raise ValueError("data_per_streamline must be a 2D array.")
 
         # We make sure there is the right amount of values
-        if self.n_rows is not None and len(value) != self.n_rows:
+        if self.n_rows > 0 and len(value) != self.n_rows:
             msg = ("The number of values ({0}) should match n_elements "
                    "({1}).").format(len(value), self.n_rows)
             raise ValueError(msg)
 
         self.store[key] = value
 
+    def _extend_entry(self, key, value):
+        """ Appends the `value` to the entry specified by `key`. """
+        self[key] = np.concatenate([self[key], value])
+
     def extend(self, other):
         """ Appends the elements of another :class:`PerArrayDict`.
 
@@ -131,16 +138,20 @@ def extend(self, other):
         -----
         The keys in both dictionaries must be the same.
         """
-        if sorted(self.keys()) != sorted(other.keys()):
+        if (len(self) > 0 and len(other) > 0
+            and sorted(self.keys()) != sorted(other.keys())):
             msg = ("Entry 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]])
+        for key in other.keys():
+            if key not in self:
+                self[key] = other[key]
+            else:
+                self._extend_entry(key, other[key])
 
 
 class PerArraySequenceDict(PerArrayDict):
@@ -158,43 +169,16 @@ def __setitem__(self, key, value):
         value = ArraySequence(value)
 
         # We make sure there is the right amount of data.
-        if (self.n_rows is not None and
-                value.total_nb_rows != self.n_rows):
+        if self.n_rows > 0 and value.total_nb_rows != self.n_rows:
             msg = ("The number of values ({0}) should match "
                    "({1}).").format(value.total_nb_rows, self.n_rows)
             raise ValueError(msg)
 
         self.store[key] = value
 
-    def extend(self, other):
-        """ Appends the elements of another :class:`PerArraySequenceDict`.
-
-        That is, for each entry in this dictionary, we append the elements
-        coming from the other dictionary at the corresponding entry.
-
-        Parameters
-        ----------
-        other : :class:`PerArraySequenceDict` object
-            Its data will be appended to the data of this dictionary.
-
-        Returns
-        -------
-        None
-
-        Notes
-        -----
-        The keys in both dictionaries must be the same.
-        """
-        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 _extend_entry(self, key, value):
+        """ Appends the `value` to the entry specified by `key`. """
+        self[key].extend(value)
 
 
 class LazyDict(collections.MutableMapping):