Merge 188732e into f94bea2

surfer/tests/test_viz.py

@@ -120,6 +120,7 @@ def test_annot():
     brain = Brain(*std_args)
     for a, b, p in zip(annots, borders, alphas):
         brain.add_annotation(a, b, p)
+    brain.set_surface('white')
     assert_raises(ValueError, brain.add_annotation, 'aparc', borders=-1)
     brain.close()
 
@@ -148,6 +149,7 @@ def test_data():
     reg_file = pjoin(data_dir, 'register.dat')
     surf_data = io.project_volume_data(mri_file, "lh", reg_file)
     brain.add_data(surf_data, -.7, .7, colormap="jet", alpha=.7)
+    brain.set_surface('white')
     brain.add_data([], vertices=np.array([], int))
     brain.close()
 
@@ -191,6 +193,7 @@ def test_label():
     brain.add_label("V1", color="steelblue", alpha=.6)
     brain.add_label("V2", color="#FF6347", alpha=.6)
     brain.add_label("entorhinal", color=(.2, 1, .5), alpha=.6)
+    brain.set_surface('white')
 
     # remove labels
     brain.remove_labels('V1')
@@ -244,6 +247,9 @@ def time_label(t):
     assert_equal(data_dicts[0]['time_idx'], 0)
     assert_equal(data_dicts[1]['time_idx'], 0)
 
+    # change surface
+    brain.set_surface('white')
+
     # remove all layers
     brain.remove_data()
     assert_equal(brain._data_dicts['lh'], [])

surfer/viz.py

@@ -493,9 +493,10 @@ def __init__(self, subject_id, hemi, surf, title=None,
         self.brains = [b['brain'] for b in brains]
         self.brain_matrix = np.array(brain_matrix)
         self.subjects_dir = subjects_dir
+        self.surf = surf
         # Initialize the overlay and label dictionaries
         self.foci_dict = dict()
-        self.labels_dict = dict()
+        self._label_dicts = dict()
         self.overlays_dict = dict()
         self.contour_list = []
         self.morphometry_list = []
@@ -514,6 +515,12 @@ def data_dict(self):
         return dict(lh=lh_list[-1] if lh_list else None,
                     rh=rh_list[-1] if rh_list else None)
 
+    @property
+    def labels_dict(self):
+        """For backwards compatibility"""
+        return {key: data['surfaces'] for key, data in
+                self._label_dicts.items()}
+
     ###########################################################################
     # HELPERS
     def _toggle_render(self, state, views=None):
@@ -1109,6 +1116,7 @@ def time_label(x):
         else:
             initial_time_index = None
 
+        meshes = []
         surfs = []
         bars = []
         views = self._toggle_render(False)
@@ -1118,7 +1126,8 @@ def time_label(x):
                                               smooth_mat, min, max, thresh,
                                               lut, colormap, alpha, time,
                                               time_label, colorbar)
-                s, ct, bar = out
+                mesh, s, ct, bar = out
+                meshes.append(mesh)
                 surfs.append(s)
                 bars.append(bar)
                 row, col = np.unravel_index(bi, self.brain_matrix.shape)
@@ -1127,6 +1136,7 @@ def time_label(x):
                                   name="time_label", row=row, col=col,
                                   font_size=time_label_size,
                                   justification='right')
+        data['meshes'] = meshes
         data['surfaces'] = surfs
         data['colorbars'] = bars
         data['orig_ctable'] = ct
@@ -1195,7 +1205,6 @@ def add_annotation(self, annot, borders=True, alpha=1, hemi=None,
                 a['surface'].remove()
             self.annot_list = []
 
-        al = self.annot_list
         for hemi, filepath in zip(hemis, filepaths):
             # Read in the data
             labels, cmap, _ = nib.freesurfer.read_annot(filepath,
@@ -1223,8 +1232,8 @@ def add_annotation(self, annot, borders=True, alpha=1, hemi=None,
 
             for brain in self._brain_list:
                 if brain['hemi'] == hemi:
-                    al.append(brain['brain'].add_annotation(annot, ids, cmap))
-        self.annot_list = al
+                    self.annot_list.append(
+                        brain['brain'].add_annotation(annot, ids, cmap))
         self._toggle_render(True, views)
 
     def add_label(self, label, color=None, alpha=1, scalar_thresh=None,
@@ -1325,23 +1334,27 @@ def add_label(self, label, color=None, alpha=1, scalar_thresh=None,
         label[ids] = 1
 
         # make sure we have a unique name
-        if label_name in self.labels_dict:
+        if label_name in self._label_dicts:
             i = 2
             name = label_name + '_%i'
-            while name % i in self.labels_dict:
+            while name % i in self._label_dicts:
                 i += 1
             label_name = name % i
 
         self._to_borders(label, hemi, borders, restrict_idx=ids)
 
         # make a list of all the plotted labels
-        ll = []
+        surfaces = []
+        meshes = []
         views = self._toggle_render(False)
         for brain in self._brain_list:
             if brain['hemi'] == hemi:
-                ll.append(brain['brain'].add_label(label, label_name,
-                          color, alpha))
-        self.labels_dict[label_name] = ll
+                mesh, surf = brain['brain'].add_label(label, label_name, color,
+                                                      alpha)
+                surfaces.append(surf)
+                meshes.append(mesh)
+        self._label_dicts[label_name] = {'hemi': hemi, 'meshes': meshes,
+                                         'surfaces': surfaces}
         self._toggle_render(True, views)
 
     def _to_borders(self, label, hemi, borders, restrict_idx=None):
@@ -1402,14 +1415,18 @@ def remove_labels(self, labels=None, hemi=None):
                  "and will be removed in PySurfer 0.9", DeprecationWarning)
 
         if labels is None:
-            labels = self.labels_dict.keys()
-        elif isinstance(labels, str):
-            labels = [labels]
+            labels_ = self._label_dicts.keys()
+        else:
+            labels_ = [labels] if isinstance(labels, str) else labels
+            missing = [key for key in labels_ if key not in self._label_dicts]
+            if missing:
+                raise ValueError("labels=%r contains unknown labels: %s" %
+                                 (labels, ', '.join(map(repr, missing))))
 
-        for key in labels:
-            label = self.labels_dict.pop(key)
-            for ll in label:
-                ll.remove()
+        for key in labels_:
+            data = self._label_dicts.pop(key)
+            for surf in data['surfaces']:
+                surf.remove()
 
     def add_morphometry(self, measure, grayscale=False, hemi=None,
                         remove_existing=True, colormap=None,
@@ -1745,6 +1762,42 @@ def set_distance(self, distance=None):
                 mlab.view(distance=distance, figure=f)
         return distance
 
+    def set_surface(self, surface):
+        if self.surf == surface:
+            return
+
+        views = self._toggle_render(False)
+        for h in self.brains:
+            try:
+                h._geo.surf = surface
+                h._geo.load_geometry()
+            except IOError:
+                h._geo.surf = h.surf
+                self._toggle_render(True)
+                raise
+            # collect meshes: main surface mesh
+            meshes = [h._geo_mesh]
+            # data overlays
+            for data in self._data_dicts[h.hemi]:
+                meshes.extend(data['meshes'])
+            # labels
+            for data in self._label_dicts.values():
+                if data['hemi'] == h.hemi:
+                    meshes.extend(data['meshes'])
+            # annotations
+            for data in self.annot_list:
+                if data['hemi'] == h.hemi:
+                    meshes.append(data['mesh'])
+
+            for mesh in meshes:
+                mesh.data.points = h._geo.coords
+                mesh.data.point_data.normals = h._geo.nn
+            for mesh in meshes:
+                mesh.update()
+            h.surf = surface
+        self.surf = surface
+        self._toggle_render(True, views)
+
     @verbose
     def scale_data_colormap(self, fmin, fmid, fmax, transparent, verbose=None):
         """Scale the data colormap.
@@ -2740,7 +2793,7 @@ def add_data(self, array, mlab_plot, vertices, smooth_mat, min, max,
         else:
             bar = None
 
-        return surf, orig_ctable, bar
+        return mesh, surf, orig_ctable, bar
 
     def add_annotation(self, annot, ids, cmap):
         """Add an annotation file"""
@@ -2759,8 +2812,8 @@ def add_annotation(self, annot, ids, cmap):
         lut_manager.load_lut_from_list(cmap / 255.)
 
         # Set the brain attributes
-        annot = dict(surface=surf, name=annot, colormap=cmap)
-        return annot
+        return dict(surface=surf, name=annot, colormap=cmap, mesh=mesh,
+                    hemi=self.hemi)
 
     def add_label(self, label, label_name, color, alpha):
         """Add an ROI label to the image"""
@@ -2776,7 +2829,7 @@ def add_label(self, label, label_name, color, alpha):
         cmap = np.array([(0, 0, 0, 0,), color])
         lut_manager = surf.module_manager.scalar_lut_manager
         lut_manager.load_lut_from_list(cmap)
-        return surf
+        return mesh, surf
 
     def add_morphometry(self, morph_data, colormap, measure,
                         min, max, colorbar):