Merge pull request #182 from christianbrodbeck/data-layers

FIX: showing multiple data layers

Author: larsoner

surfer/tests/test_viz.py

@@ -219,10 +219,27 @@ def time_label(t):
     assert_equal(brain.data_dict['lh']['time_idx'], 2)
     # viewer = TimeViewer(brain)
 
+    # multiple data layers
+    assert_raises(ValueError, brain.add_data, data, vertices=vertices,
+                  time=time[:-1])
     brain.add_data(data, colormap=colormap, vertices=vertices,
                    smoothing_steps=10, time=time, time_label=time_label,
-                   initial_time=.09, remove_existing=True)
+                   initial_time=.09)
     assert_equal(brain.data_dict['lh']['time_idx'], 1)
+    data_dicts = brain._data_dicts['lh']
+    assert_equal(len(data_dicts), 2)
+    assert_equal(data_dicts[0]['time_idx'], 1)
+    assert_equal(data_dicts[1]['time_idx'], 1)
+
+    # shift time in both layers
+    brain.set_data_time_index(0)
+    assert_equal(data_dicts[0]['time_idx'], 0)
+    assert_equal(data_dicts[1]['time_idx'], 0)
+
+    # remove all layers
+    brain.remove_data()
+    assert_equal(brain._data_dicts['lh'], [])
+
     brain.close()
 
 

surfer/viz.py

@@ -500,11 +500,20 @@ def __init__(self, subject_id, hemi, surf, title=None,
         self.contour_list = []
         self.morphometry_list = []
         self.annot_list = []
-        self.data_dict = dict(lh=None, rh=None)
+        self._data_dicts = dict(lh=[], rh=[])
         # note that texts gets treated differently
         self.texts_dict = dict()
+        self._times = None
         self.n_times = None
 
+    @property
+    def data_dict(self):
+        """For backwards compatibility"""
+        lh_list = self._data_dicts['lh']
+        rh_list = self._data_dicts['rh']
+        return dict(lh=lh_list[-1] if lh_list else None,
+                    rh=rh_list[-1] if rh_list else None)
+
     ###########################################################################
     # HELPERS
     def _toggle_render(self, state, views=None):
@@ -1052,20 +1061,38 @@ def add_data(self, array, min=None, max=None, thresh=None,
                     transparent=False, time=0, time_idx=0,
                     vertices=vertices, smooth_mat=smooth_mat)
 
+        # clean up existing data
+        if remove_existing:
+            self.remove_data(hemi)
+
         # Create time array and add label if 2D
         if array.ndim == 2:
+            # check time array
             if time is None:
                 time = np.arange(array.shape[1])
-            self._times = time
-            self.n_times = array.shape[1]
-            if not self.n_times == len(time):
-                raise ValueError('time is not the same length as '
-                                 'array.shape[1]')
+            else:
+                time = np.asarray(time)
+                if time.shape != (array.shape[1],):
+                    raise ValueError('time has shape %s, but need shape %s '
+                                     '(array.shape[1])' %
+                                     (time.shape, (array.shape[1],)))
+
+            if self.n_times is None:
+                self.n_times = len(time)
+                self._times = time
+            elif len(time) != self.n_times:
+                raise ValueError("New n_times is different from previous "
+                                 "n_times")
+            elif not np.array_equal(time, self._times):
+                raise ValueError("Not all time values are consistent with "
+                                 "previously set times.")
+
             # initial time
             if initial_time is None:
                 initial_time_index = None
             else:
                 initial_time_index = self.index_for_time(initial_time)
+
             # time label
             if isinstance(time_label, string_types):
                 time_label_fmt = time_label
@@ -1077,8 +1104,6 @@ def time_label(x):
             data["time_idx"] = 0
             y_txt = 0.05 + 0.05 * bool(colorbar)
         else:
-            self._times = None
-            self.n_times = None
             initial_time_index = None
 
         surfs = []
@@ -1103,10 +1128,7 @@ def time_label(x):
         data['colorbars'] = bars
         data['orig_ctable'] = ct
 
-        if remove_existing and self.data_dict[hemi] is not None:
-            for surf in self.data_dict[hemi]['surfaces']:
-                surf.parent.parent.remove()
-        self.data_dict[hemi] = data
+        self._data_dicts[hemi].append(data)
 
         if initial_time_index is not None:
             self.set_data_time_index(initial_time_index)
@@ -1340,6 +1362,26 @@ def _to_borders(self, label, hemi, borders, restrict_idx=None):
             show[keep_idx] = 1
             label *= show
 
+    def remove_data(self, hemi=None):
+        """Remove data shown with ``Brain.add_data()``.
+
+        Parameters
+        ----------
+        hemi : str | None
+            Hemisphere from which to remove data (default is all shown
+            hemispheres).
+        """
+        hemis = self._check_hemis(hemi)
+        for hemi in hemis:
+            for data in self._data_dicts[hemi]:
+                for surf in data['surfaces']:
+                    surf.parent.parent.remove()
+            self._data_dicts[hemi] = []
+
+        # if no data is left, reset time properties
+        if not self._data_dicts['lh'] and not self._data_dicts['rh']:
+            self.n_times = self._times = None
+
     def remove_labels(self, labels=None, hemi=None):
         """Remove one or more previously added labels from the image.
 
@@ -1797,8 +1839,7 @@ def set_data_time_index(self, time_idx, interpolation='quadratic'):
 
         views = self._toggle_render(False)
         for hemi in ['lh', 'rh']:
-            data = self.data_dict[hemi]
-            if data is not None:
+            for data in self._data_dicts[hemi]:
                 # interpolation
                 if isinstance(time_idx, float):
                     times = np.arange(self.n_times)