create set_offsets3d for PathCollection3d

ianhi · ianhi · commit 7f70c15d8001 · 2021-02-24T16:20:19.000-05:00
diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
@@ -533,6 +533,21 @@ def set_sort_zpos(self, val):
         self._sort_zpos = val
         self.stale = True
 
+    def get_offsets3d(self):
+        """Return the 3d offsets for the collection as numpy array."""
+        return np.array(self._offsets3d).T
+
+    def set_offsets3d(self, offsets):
+        """
+        Set the offsets for the collection.
+
+        Parameters
+        ----------
+        offsets : (N, 3) array-like
+        """
+        self._offsets3d = np.asanyarray(offsets).T
+        self.stale = True
+
     def set_3d_properties(self, zs, zdir):
         # Force the collection to initialize the face and edgecolors
         # just in case it is a scalarmappable with a colormap.
diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py
@@ -1116,6 +1116,26 @@ def test_line3d_set_get_data_3d():
     np.testing.assert_array_equal((x2, y2, z2), line.get_data_3d())
 
 
+def test_PathCollection3d_set_get_offsets_3d():
+    orig = np.array([
+            [0, 1, 1],
+            [2, 3, 2]
+    ])
+
+    new = np.array([
+            [6, 7, 4],
+            [8, 9, 5],
+            [10, 11, 6]
+    ])
+
+    fig = plt.figure()
+    ax = fig.add_subplot(projection='3d')
+    scat = ax.scatter(*orig.T)
+    np.testing.assert_array_equal(orig, scat.get_offsets3d())
+    scat.set_offsets3d(new)
+    np.testing.assert_array_equal(new, scat.get_offsets3d())
+
+
 @check_figures_equal(extensions=["png"])
 def test_inverted(fig_test, fig_ref):
     # Plot then invert.
