Backport PR #18136: Sort 3d sizes along with other properties

jkseppan · meeseeksmachine · commit 0628e6fc30b5 · 2020-08-02T08:25:17.000Z
diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
@@ -481,6 +481,7 @@ def set_3d_properties(self, zs, zdir):
         self._offsets3d = juggle_axes(xs, ys, np.atleast_1d(zs), zdir)
         self._facecolor3d = self.get_facecolor()
         self._edgecolor3d = self.get_edgecolor()
+        self._sizes3d = self.get_sizes()
         self.stale = True
 
     def do_3d_projection(self, renderer):
@@ -495,6 +496,8 @@ def do_3d_projection(self, renderer):
         ecs = (_zalpha(self._edgecolor3d, vzs) if self._depthshade else
                self._edgecolor3d)
 
+        sizes = self._sizes3d
+
         # Sort the points based on z coordinates
         # Performance optimization: Create a sorted index array and reorder
         # points and point properties according to the index array
@@ -506,13 +509,16 @@ def do_3d_projection(self, renderer):
         vys = vys[z_markers_idx]
         fcs = fcs[z_markers_idx]
         ecs = ecs[z_markers_idx]
+        if len(sizes) > 1:
+            sizes = sizes[z_markers_idx]
         vps = np.column_stack((vxs, vys))
 
         fcs = mcolors.to_rgba_array(fcs, self._alpha)
         ecs = mcolors.to_rgba_array(ecs, self._alpha)
 
         self.set_edgecolors(ecs)
         self.set_facecolors(fcs)
+        self.set_sizes(sizes)
 
         PathCollection.set_offsets(self, vps)
 
diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py
@@ -242,6 +242,30 @@ def test_scatter3d_color():
                color='b', marker='s')
 
 
+@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
+
+    ax_ref = fig_ref.gca(projection='3d')
+    ax_test = fig_test.gca(projection='3d')
+
+    small = np.ma.masked_array(z, s == 100, dtype=float)
+    large = np.ma.masked_array(z, s != 100, dtype=float)
+
+    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)
+
+
 @pytest.mark.parametrize('azim', [-50, 130])  # yellow first, blue first
 @check_figures_equal(extensions=['png'])
 def test_marker_draw_order_data_reversed(fig_test, fig_ref, azim):
