From 7f70c15d800124b125c5ca5c172d0a60ce8de5c1 Mon Sep 17 00:00:00 2001
From: Ian Hunt-Isaak <ianhuntisaak@gmail.com>
Date: Wed, 24 Feb 2021 16:13:46 -0500
Subject: [PATCH] create set_offsets3d for PathCollection3d

---
 lib/mpl_toolkits/mplot3d/art3d.py      | 15 +++++++++++++++
 lib/mpl_toolkits/tests/test_mplot3d.py | 20 ++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
index 822ff79fcc24..ec68ae180c07 100644
--- 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
index cfcc383db1a4..7a395dbaf374 100644
--- 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.