Backport PR #18929: FIX: make sure scalarmappable updates are handled correctly in 3D

QuLogic · QuLogic · commit f679026f3979 · 2020-11-11T15:47:58.000-05:00
diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
@@ -264,6 +264,8 @@ def do_3d_projection(self, renderer):
         """
         Project the points according to renderer matrix.
         """
+        # see _update_scalarmappable docstring for why this must be here
+        _update_scalarmappable(self)
         xyslist = [
             proj3d.proj_trans_points(points, renderer.M) for points in
             self._segments3d]
@@ -418,6 +420,8 @@ def set_3d_properties(self, zs, zdir):
         self.stale = True
 
     def do_3d_projection(self, renderer):
+        # see _update_scalarmappable docstring for why this must be here
+        _update_scalarmappable(self)
         xs, ys, zs = self._offsets3d
         vxs, vys, vzs, vis = proj3d.proj_transform_clip(xs, ys, zs, renderer.M)
 
@@ -486,6 +490,8 @@ def set_3d_properties(self, zs, zdir):
         self.stale = True
 
     def do_3d_projection(self, renderer):
+        # see _update_scalarmappable docstring for why this must be here
+        _update_scalarmappable(self)
         xs, ys, zs = self._offsets3d
         vxs, vys, vzs, vis = proj3d.proj_transform_clip(xs, ys, zs, renderer.M)
 
@@ -528,6 +534,77 @@ def do_3d_projection(self, renderer):
         return np.min(vzs) if vzs.size else np.nan
 
 
+def _update_scalarmappable(sm):
+    """
+    Update a 3D ScalarMappable.
+
+    With ScalarMappable objects if the data, colormap, or norm are
+    changed, we need to update the computed colors.  This is handled
+    by the base class method update_scalarmappable.  This method works
+    by, detecting if work needs to be done, and if so stashing it on
+    the ``self._facecolors`` attribute.
+
+    With 3D collections we internally sort the components so that
+    things that should be "in front" are rendered later to simulate
+    having a z-buffer (in addition to doing the projections).  This is
+    handled in the ``do_3d_projection`` methods which are called from the
+    draw method of the 3D Axes.  These methods:
+
+    - do the projection from 3D -> 2D
+    - internally sort based on depth
+    - stash the results of the above in the 2D analogs of state
+    - return the z-depth of the whole artist
+
+    the last step is so that we can, at the Axes level, sort the children by
+    depth.
+
+    The base `draw` method of the 2D artists unconditionally calls
+    update_scalarmappable and rely on the method's internal caching logic to
+    lazily evaluate.
+
+    These things together mean you can have the sequence of events:
+
+    - we create the artist, do the color mapping and stash the results
+      in a 3D specific state.
+    - change something about the ScalarMappable that marks it as in
+      need of an update (`ScalarMappable.changed` and friends).
+    - We call do_3d_projection and shuffle the stashed colors into the
+      2D version of face colors
+    - the draw method calls the update_scalarmappable method which
+      overwrites our shuffled colors
+    - we get a render that is wrong
+    - if we re-render (either with a second save or implicitly via
+      tight_layout / constrained_layout / bbox_inches='tight' (ex via
+      inline's defaults)) we again shuffle the 3D colors
+    - because the CM is not marked as changed update_scalarmappable is
+      a no-op and we get a correct looking render.
+
+    This function is an internal helper to:
+
+    - sort out if we need to do the color mapping at all (has data!)
+    - sort out if update_scalarmappable is going to be a no-op
+    - copy the data over from the 2D -> 3D version
+
+    This must be called first thing in do_3d_projection to make sure that
+    the correct colors get shuffled.
+
+    Parameters
+    ----------
+    sm : ScalarMappable
+        The ScalarMappable to update and stash the 3D data from
+
+    """
+    if sm._A is None:
+        return
+    copy_state = sm._update_dict['array']
+    ret = sm.update_scalarmappable()
+    if copy_state:
+        if sm._is_filled:
+            sm._facecolor3d = sm._facecolors
+        elif sm._is_stroked:
+            sm._edgecolor3d = sm._edgecolors
+
+
 def patch_collection_2d_to_3d(col, zs=0, zdir='z', depthshade=True):
     """
     Convert a :class:`~matplotlib.collections.PatchCollection` into a
@@ -650,8 +727,8 @@ def set_3d_properties(self):
         self.update_scalarmappable()
         self._sort_zpos = None
         self.set_zsort('average')
-        self._facecolors3d = PolyCollection.get_facecolor(self)
-        self._edgecolors3d = PolyCollection.get_edgecolor(self)
+        self._facecolor3d = PolyCollection.get_facecolor(self)
+        self._edgecolor3d = PolyCollection.get_edgecolor(self)
         self._alpha3d = PolyCollection.get_alpha(self)
         self.stale = True
 
@@ -664,17 +741,15 @@ def do_3d_projection(self, renderer):
         """
         Perform the 3D projection for this object.
         """
-        # FIXME: This may no longer be needed?
-        if self._A is not None:
-            self.update_scalarmappable()
-            self._facecolors3d = self._facecolors
+        # see _update_scalarmappable docstring for why this must be here
+        _update_scalarmappable(self)
 
         txs, tys, tzs = proj3d._proj_transform_vec(self._vec, renderer.M)
         xyzlist = [(txs[sl], tys[sl], tzs[sl]) for sl in self._segslices]
 
         # This extra fuss is to re-order face / edge colors
-        cface = self._facecolors3d
-        cedge = self._edgecolors3d
+        cface = self._facecolor3d
+        cedge = self._edgecolor3d
         if len(cface) != len(xyzlist):
             cface = cface.repeat(len(xyzlist), axis=0)
         if len(cedge) != len(xyzlist):
@@ -699,8 +774,8 @@ def do_3d_projection(self, renderer):
         else:
             PolyCollection.set_verts(self, segments_2d, self._closed)
 
-        if len(self._edgecolors3d) != len(cface):
-            self._edgecolors2d = self._edgecolors3d
+        if len(self._edgecolor3d) != len(cface):
+            self._edgecolors2d = self._edgecolor3d
 
         # Return zorder value
         if self._sort_zpos is not None:
@@ -717,23 +792,23 @@ def do_3d_projection(self, renderer):
 
     def set_facecolor(self, colors):
         PolyCollection.set_facecolor(self, colors)
-        self._facecolors3d = PolyCollection.get_facecolor(self)
+        self._facecolor3d = PolyCollection.get_facecolor(self)
 
     def set_edgecolor(self, colors):
         PolyCollection.set_edgecolor(self, colors)
-        self._edgecolors3d = PolyCollection.get_edgecolor(self)
+        self._edgecolor3d = PolyCollection.get_edgecolor(self)
 
     def set_alpha(self, alpha):
         # docstring inherited
         artist.Artist.set_alpha(self, alpha)
         try:
-            self._facecolors3d = mcolors.to_rgba_array(
-                self._facecolors3d, self._alpha)
+            self._facecolor3d = mcolors.to_rgba_array(
+                self._facecolor3d, self._alpha)
         except (AttributeError, TypeError, IndexError):
             pass
         try:
             self._edgecolors = mcolors.to_rgba_array(
-                    self._edgecolors3d, self._alpha)
+                    self._edgecolor3d, self._alpha)
         except (AttributeError, TypeError, IndexError):
             pass
         self.stale = True
diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py
@@ -106,7 +106,7 @@ def test_bar3d_lightsource():
     # the top facecolors compared to the default, and that those colors are
     # precisely the colors from the colormap, due to the illumination parallel
     # to the z-axis.
-    np.testing.assert_array_equal(color, collection._facecolors3d[1::6])
+    np.testing.assert_array_equal(color, collection._facecolor3d[1::6])
 
 
 @mpl3d_image_comparison(['contour3d.png'])
@@ -1087,3 +1087,25 @@ def test_colorbar_pos():
     fig.canvas.draw()
     # check that actually on the bottom
     assert cbar.ax.get_position().extents[1] < 0.2
+
+
+@pytest.mark.style('default')
+@check_figures_equal(extensions=["png"])
+def test_scalarmap_update(fig_test, fig_ref):
+
+    x, y, z = np.array((list(itertools.product(*[np.arange(0, 5, 1),
+                                                 np.arange(0, 5, 1),
+                                                 np.arange(0, 5, 1)])))).T
+    c = x + y
+
+    # test
+    ax_test = fig_test.add_subplot(111, projection='3d')
+    sc_test = ax_test.scatter(x, y, z, c=c, s=40, cmap='viridis')
+    # force a draw
+    fig_test.canvas.draw()
+    # mark it as "stale"
+    sc_test.changed()
+
+    # ref
+    ax_ref = fig_ref.add_subplot(111, projection='3d')
+    sc_ref = ax_ref.scatter(x, y, z, c=c, s=40, cmap='viridis')
