Repository: vispy/vispy

Issue:
SurfacePlot color change slower with set_data versus mesh_data.set_vertex_colors
Sample code below: a SurfacePlot whose color is updated when "O" is pressed using `mesh_data.set_vertex_colors`.

Problem: the `mesh_data.set_vertex_colors` works if called before adding the surface to the view, but not in later calls.

Temporary fix (thanks @djhoese ): add `surf.mesh_data_changed()` after changing the color.

PS: using `surf.set_data(colors=c)` instead of `surf.mesh_data.set_vertex_colors(c); surf.mesh_data_changed()` also works, but for some reason is much slower

```
import vispy, vispy.scene
import numpy as np

shape =  (200, 200)
Z = np.random.random(shape)*5
B = np.random.random(shape)

##


canvas = vispy.scene.SceneCanvas(keys='interactive', show=True)
view = canvas.central_widget.add_view()

surf = vispy.scene.visuals.SurfacePlot(z=Z, color=(1, 1, 1, 1))

# Add color
def setColor(C):
	C = np.array(C)
	c = vispy.color.get_colormap("hsl").map(C/np.max(C)).reshape(C.shape + (-1,))
	c = c.flatten().tolist()
	c=list(map(lambda x,y,z,w:(x,y,z,w), c[0::4],c[1::4],c[2::4],c[3::4]))
	surf.mesh_data.set_vertex_colors(c)
setColor(B)

# cycle colors on key press
col_cycle_i = 0
def cycle_colors(event):
	key = event.key.name
	if key != "O": return
	global col_cycle_i
	col_cycle_i = (col_cycle_i+1)%2
	print('Setting new coli', col_cycle_i)
	if col_cycle_i == 0:
		setColor(B)
	elif col_cycle_i == 1:
		setColor(Z)
		
canvas.events.key_press.connect(cycle_colors)
view.add(surf)
axis = vispy.scene.visuals.XYZAxis(parent=view.scene)
view.camera = 'fly'
```


Discussion:
Thanks for this. As mentioned on gitter, the `set_data(colors=c)` works but is much slower than touching the mesh_data directly. This is probably because the surface plot is recreating/re-setting the data for the mesh every time it is run. We should be able to optimize that away if we detect that only a color change is needed.

Resolving patch:
diff --git a/vispy/geometry/meshdata.py b/vispy/geometry/meshdata.py
--- a/vispy/geometry/meshdata.py
+++ b/vispy/geometry/meshdata.py
@@ -411,7 +411,7 @@ def set_vertex_colors(self, colors, indexed=None):
         indexed : str | None
             Should be 'faces' if colors are indexed by faces.
         """
-        colors = _fix_colors(np.asarray(colors))
+        colors = _fix_colors(colors)
         if indexed is None:
             if colors.ndim != 2:
                 raise ValueError('colors must be 2D if indexed is None')
diff --git a/vispy/visuals/surface_plot.py b/vispy/visuals/surface_plot.py
--- a/vispy/visuals/surface_plot.py
+++ b/vispy/visuals/surface_plot.py
@@ -54,36 +54,19 @@ def __init__(self, x=None, y=None, z=None, colors=None, **kwargs):
         MeshVisual.__init__(self, **kwargs)
         self.set_data(x, y, z, colors)
 
-    def set_data(self, x=None, y=None, z=None, colors=None):
-        """Update the data in this surface plot.
-
-        Parameters
-        ----------
-        x : ndarray | None
-            1D/2D array of values specifying the x positions of vertices in
-            the grid. In case 1D array given as input, the values will be
-            replicated to fill the 2D array of size(z). If None, values will be
-            assumed to be integers.
-        y : ndarray | None
-            1D/2D array of values specifying the x positions of vertices in
-            the grid. In case 1D array given as input, the values will be
-            replicated to fill the 2D array of size(z). If None, values will be
-            assumed to be integers.
-        z : ndarray
-            2D array of height values for each grid vertex.
-        colors : ndarray
-            (width, height, 4) array of vertex colors.
-        """
+    def _update_x_data(self, x):
         if x is not None:
             if self._x is None or len(x) != len(self._x):
                 self.__vertices = None
             self._x = x
 
+    def _update_y_data(self, y):
         if y is not None:
             if self._y is None or len(y) != len(self._y):
                 self.__vertices = None
             self._y = y
 
+    def _update_z_data(self, z):
         if z is not None:
             if self._x is not None and z.shape[0] != len(self._x):
                 raise TypeError('Z values must have shape (len(x), len(y))')
@@ -94,64 +77,106 @@ def set_data(self, x=None, y=None, z=None, colors=None):
                     self._z.shape != self.__vertices.shape[:2]):
                 self.__vertices = None
 
-        if self._z is None:
-            return
-
-        update_mesh = False
+    def _update_mesh_vertices(self, x_is_new, y_is_new, z_is_new):
         new_vertices = False
+        update_vertices = False
+        update_faces = False
 
         # Generate vertex and face array
         if self.__vertices is None:
-            new_vertices = True
             self.__vertices = np.empty((self._z.shape[0], self._z.shape[1], 3),
                                        dtype=np.float32)
-            self.generate_faces()
-            self.__meshdata.set_faces(self.__faces)
-            update_mesh = True
+            self.__faces = self._generate_faces()
+            new_vertices = True
+            update_faces = True
 
         # Copy x, y, z data into vertex array
-        if new_vertices or x is not None:
-            if x is None:
-                if self._x is None:
-                    x = np.arange(self._z.shape[0])
-                else:
-                    x = self._x
+        if new_vertices or x_is_new:
+            if not x_is_new and self._x is None:
+                x = np.arange(self._z.shape[0])
+            else:
+                x = self._x
             if x.ndim == 1:
                 x = x.reshape(len(x), 1)
             # Copy the 2D data into the appropriate slice
             self.__vertices[:, :, 0] = x
-            update_mesh = True
-
-        if new_vertices or y is not None:
-            if y is None:
-                if self._y is None:
-                    y = np.arange(self._z.shape[1])
-                else:
-                    y = self._y
+            update_vertices = True
 
+        if new_vertices or y_is_new:
+            if not y_is_new and self._y is None:
+                y = np.arange(self._z.shape[1])
+            else:
+                y = self._y
             if y.ndim == 1:
                 y = y.reshape(1, len(y))
-
             # Copy the 2D data into the appropriate slice
             self.__vertices[:, :, 1] = y
-            update_mesh = True
+            update_vertices = True
 
-        if new_vertices or z is not None:
+        if new_vertices or z_is_new:
             self.__vertices[..., 2] = self._z
-            update_mesh = True
+            update_vertices = True
+        return update_faces, update_vertices
 
-        if colors is not None:
-            self.__meshdata.set_vertex_colors(colors)
-            update_mesh = True
+    def _prepare_mesh_colors(self, colors):
+        if colors is None:
+            return
+        colors = np.asarray(colors)
+        if colors.ndim == 3:
+            # convert (width, height, 4) to (num_verts, 4)
+            vert_shape = self.__vertices.shape
+            num_vertices = vert_shape[0] * vert_shape[1]
+            colors = colors.reshape(num_vertices, 3)
+        return colors
+
+    def set_data(self, x=None, y=None, z=None, colors=None):
+        """Update the data in this surface plot.
+
+        Parameters
+        ----------
+        x : ndarray | None
+            1D/2D array of values specifying the x positions of vertices in
+            the grid. In case 1D array given as input, the values will be
+            replicated to fill the 2D array of size(z). If None, values will be
+            assumed to be integers.
+        y : ndarray | None
+            1D/2D array of values specifying the x positions of vertices in
+            the grid. In case 1D array given as input, the values will be
+            replicated to fill the 2D array of size(z). If None, values will be
+            assumed to be integers.
+        z : ndarray
+            2D array of height values for each grid vertex.
+        colors : ndarray
+            (width, height, 4) array of vertex colors.
+        """
+        self._update_x_data(x)
+        self._update_y_data(y)
+        self._update_z_data(z)
+
+        if self._z is None:
+            # no mesh data to plot so no need to update
+            return
 
-        # Update MeshData
-        if update_mesh:
+        update_faces, update_vertices = self._update_mesh_vertices(
+            x is not None,
+            y is not None,
+            z is not None
+        )
+
+        colors = self._prepare_mesh_colors(colors)
+        update_colors = colors is not None
+        if update_colors:
+            self.__meshdata.set_vertex_colors(colors)
+        if update_faces:
+            self.__meshdata.set_faces(self.__faces)
+        if update_vertices:
             self.__meshdata.set_vertices(
                 self.__vertices.reshape(self.__vertices.shape[0] *
                                         self.__vertices.shape[1], 3))
+        if update_faces or update_vertices or update_colors:
             MeshVisual.set_data(self, meshdata=self.__meshdata)
 
-    def generate_faces(self):
+    def _generate_faces(self):
         cols = self._z.shape[1] - 1
         rows = self._z.shape[0] - 1
         faces = np.empty((cols * rows * 2, 3), dtype=np.uint)
@@ -164,4 +189,4 @@ def generate_faces(self):
             faces[start:start + cols] = rowtemplate1 + row * (cols + 1)
             faces[start + cols:start + (cols * 2)] =\
                 rowtemplate2 + row * (cols + 1)
-        self.__faces = faces
+        return faces


Test patch: