Merge pull request #18293 from QuLogic/scatter3d-color

dopplershift · web-flow · commit 52f7f4252e4a · 2020-08-29T15:50:05.000-06:00
Fix scatter3d color/linewidth re-projection
diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
@@ -488,6 +488,7 @@ def set_3d_properties(self, zs, zdir):
         self._facecolor3d = self.get_facecolor()
         self._edgecolor3d = self.get_edgecolor()
         self._sizes3d = self.get_sizes()
+        self._linewidth3d = self.get_linewidth()
         self.stale = True
 
     def do_3d_projection(self, renderer):
@@ -496,13 +497,10 @@ def do_3d_projection(self, renderer):
 
         fcs = (_zalpha(self._facecolor3d, vzs) if self._depthshade else
                self._facecolor3d)
-        fcs = mcolors.to_rgba_array(fcs, self._alpha)
-        self.set_facecolors(fcs)
-
         ecs = (_zalpha(self._edgecolor3d, vzs) if self._depthshade else
                self._edgecolor3d)
-
         sizes = self._sizes3d
+        lws = self._linewidth3d
 
         # Sort the points based on z coordinates
         # Performance optimization: Create a sorted index array and reorder
@@ -513,11 +511,14 @@ def do_3d_projection(self, renderer):
         vzs = vzs[z_markers_idx]
         vxs = vxs[z_markers_idx]
         vys = vys[z_markers_idx]
-        if self._depthshade:
+        if len(fcs) > 1:
             fcs = fcs[z_markers_idx]
+        if len(ecs) > 1:
             ecs = ecs[z_markers_idx]
         if len(sizes) > 1:
             sizes = sizes[z_markers_idx]
+        if len(lws) > 1:
+            lws = lws[z_markers_idx]
         vps = np.column_stack((vxs, vys))
 
         fcs = mcolors.to_rgba_array(fcs, self._alpha)
@@ -526,6 +527,7 @@ def do_3d_projection(self, renderer):
         self.set_edgecolors(ecs)
         self.set_facecolors(fcs)
         self.set_sizes(sizes)
+        self.set_linewidth(lws)
 
         PathCollection.set_offsets(self, vps)
 
diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py
@@ -1,4 +1,5 @@
 import functools
+import itertools
 
 import pytest
 
@@ -242,39 +243,60 @@ def test_scatter3d_color():
                color='b', marker='s')
 
 
-def test_scatter3d_depthshade_false():
-    """
-    Test that 3d scatter plot doesn't throw
-    IndexError with depthshade=False (issue #18037)
-    """
-    x = y = z = np.arange(16)
-    fig_test = plt.figure()
-    ax_test = fig_test.add_subplot(projection='3d')
-    ax_test.scatter(x, y, z, depthshade=False)
+@pytest.mark.parametrize('depthshade', [True, False])
+@check_figures_equal(extensions=['png'])
+def test_scatter3d_sorting(fig_ref, fig_test, depthshade):
+    """Test that marker properties are correctly sorted."""
 
+    y, x = np.mgrid[:10, :10]
+    z = np.arange(x.size).reshape(x.shape)
 
-@check_figures_equal(extensions=['png'])
-def test_scatter3d_size(fig_ref, fig_test):
-    """Test that large markers in correct position (issue #18135)"""
-    x = np.arange(10)
-    x, y = np.meshgrid(x, x)
-    z = np.arange(100).reshape(10, 10)
-
-    s = np.full(z.shape, 5)
-    s[0, 0] = 100
-    s[-1, 0] = 100
-    s[0, -1] = 100
-    s[-1, -1] = 100
+    sizes = np.full(z.shape, 25)
+    sizes[0::2, 0::2] = 100
+    sizes[1::2, 1::2] = 100
 
-    ax_ref = fig_ref.gca(projection='3d')
-    ax_test = fig_test.gca(projection='3d')
+    facecolors = np.full(z.shape, 'C0')
+    facecolors[:5, :5] = 'C1'
+    facecolors[6:, :4] = 'C2'
+    facecolors[6:, 6:] = 'C3'
 
-    small = np.ma.masked_array(z, s == 100, dtype=float)
-    large = np.ma.masked_array(z, s != 100, dtype=float)
+    edgecolors = np.full(z.shape, 'C4')
+    edgecolors[1:5, 1:5] = 'C5'
+    edgecolors[5:9, 1:5] = 'C6'
+    edgecolors[5:9, 5:9] = 'C7'
 
-    ax_ref.scatter(x, y, large, s=100, c="C0", alpha=1)
-    ax_ref.scatter(x, y, small, s=5, c="C0", alpha=1)
-    ax_test.scatter(x, y, z, s=s, c="C0", alpha=1)
+    linewidths = np.full(z.shape, 2)
+    linewidths[0::2, 0::2] = 5
+    linewidths[1::2, 1::2] = 5
+
+    x, y, z, sizes, facecolors, edgecolors, linewidths = [
+        a.flatten()
+        for a in [x, y, z, sizes, facecolors, edgecolors, linewidths]
+    ]
+
+    ax_ref = fig_ref.gca(projection='3d')
+    sets = (np.unique(a) for a in [sizes, facecolors, edgecolors, linewidths])
+    for s, fc, ec, lw in itertools.product(*sets):
+        subset = (
+            (sizes != s) |
+            (facecolors != fc) |
+            (edgecolors != ec) |
+            (linewidths != lw)
+        )
+        subset = np.ma.masked_array(z, subset, dtype=float)
+
+        # When depth shading is disabled, the colors are passed through as
+        # single-item lists; this triggers single path optimization. The
+        # following reshaping is a hack to disable that, since the optimization
+        # would not occur for the full scatter which has multiple colors.
+        fc = np.repeat(fc, sum(~subset.mask))
+
+        ax_ref.scatter(x, y, subset, s=s, fc=fc, ec=ec, lw=lw, alpha=1,
+                       depthshade=depthshade)
+
+    ax_test = fig_test.gca(projection='3d')
+    ax_test.scatter(x, y, z, s=sizes, fc=facecolors, ec=edgecolors,
+                    lw=linewidths, alpha=1, depthshade=depthshade)
 
 
 @pytest.mark.parametrize('azim', [-50, 130])  # yellow first, blue first
