Update lib/mpl_toolkits/mplot3d/art3d.py

apaszke · eric-wieser · web-flow · commit 41f36ac99249 · 2020-03-22T16:29:57.000+01:00
Co-Authored-By: Eric Wieser &lt;wieser.eric@gmail.com&gt;
diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
@@ -615,11 +615,16 @@ def get_vector(self, segments3d):
             self._segments = segments3d
         else:
             num_faces = len(segments3d)
-            max_verts = max((len(face) for face in segments3d), default=0)
-            padded = np.full((num_faces, max_verts, 3), np.nan)
+            num_verts = np.array([len(face) for face in segments3d], dtype=np.intp)
+            max_verts = num_verts.max(initial=0)
+            padded = np.empty((num_faces, max_verts, 3))
             for i, face in enumerate(segments3d):
                 padded[i, :len(face)] = face
-            self._segments = np.ma.masked_invalid(padded)
+            mask = np.arange(max_verts) >= num_verts[:,None]
+            mask = mask[..., None]  # add a component axis
+            # ma.array does not broadcast the mask for us
+            mask = np.broadcast_to(mask, padded.shape)
+            self._segments = np.ma.array(padded, mask=mask)
 
     def set_verts(self, verts, closed=True):
         """Set 3D vertices."""
