matplotlib · scottshambaugh · Dec 21, 2023 · Dec 22, 2023 · Mar 10, 2024 · Mar 11, 2024
diff --git a/doc/users/next_whats_new/3d_collections_offset3d.rst b/doc/users/next_whats_new/3d_collections_offset3d.rst
@@ -0,0 +1,11 @@
+3D Collections have ``set_offsets3d`` and ``get_offsets3d`` methods
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+All 3D Collections (``Patch3DCollection``, ``Path3DCollection``,
+``Poly3DCollection``) now have ``set_offsets3d`` and ``get_offsets3d`` methods
+which allow you to set and get the offsets of the collection in data
+coordinates.
+
+For ``Patch3DCollection`` and ``Path3DCollection`` (e.g. from `~.Axes3D.scatter`),
+this sets the position of each element. For ``Poly3DCollection``, the offsets
+translate the polygon faces in 3D space.
diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
@@ -784,24 +784,24 @@ def set_3d_properties(self, zs, zdir, axlim_clip=False):
         # Force the collection to initialize the face and edgecolors
         # just in case it is a scalarmappable with a colormap.
         self.update_scalarmappable()
-        offsets = self.get_offsets()
-        if len(offsets) > 0:
-            xs, ys = offsets.T
+        offsets2d = super().get_offsets()
+        if len(offsets2d) > 0:
+            xs, ys = offsets2d.T
         else:
             xs = []
             ys = []
-        self._offsets3d = juggle_axes(xs, ys, np.atleast_1d(zs), zdir)
+        self._zdir = zdir
+        self.set_offsets3d(np.ma.column_stack((xs, ys, np.atleast_1d(zs))), zdir)
         self._z_markers_idx = slice(-1)
         self._vzs = None
         self._axlim_clip = axlim_clip
         self.stale = True
 
     def do_3d_projection(self):
+        xs, ys, zs = self.get_offsets3d()
         if self._axlim_clip:
-            mask = _viewlim_mask(*self._offsets3d, self.axes)
-            xs, ys, zs = np.ma.array(self._offsets3d, mask=mask)
-        else:
-            xs, ys, zs = self._offsets3d
+            mask = _viewlim_mask(xs, ys, zs, self.axes)
+            xs, ys, zs = np.ma.array((xs, ys, zs), mask=mask)
         vxs, vys, vzs, vis = proj3d._proj_transform_clip(xs, ys, zs,
                                                          self.axes.M,
                                                          self.axes._focal_length)
@@ -841,6 +841,31 @@ def get_edgecolor(self):
             return self.get_facecolor()
         return self._maybe_depth_shade_and_sort_colors(super().get_edgecolor())
 
+    def set_offsets3d(self, offsets, zdir='z'):
+        """
+        Set the 3d offsets for the collection.
+
+        Parameters
+        ----------
+        offsets : (N, 3) or (3,) array-like
+            The offsets to be set.
+        zdir : {'x', 'y', 'z'}
+            The axis in which to place the offsets. Default: 'z'.
+            See `.get_dir_vector` for a description of the values.
+        """
+        return _set_offsets3d(self, offsets, zdir)
+
+    def get_offsets3d(self):
+        """
+        Return the 3d offsets for the collection.
+
+        Returns
+        -------
+        offsets : (N, 3) array
+            The offsets for the collection.
+        """
+        return _get_offsets3d(self)
+
 
 def _get_data_scale(X, Y, Z):
     """
@@ -940,14 +965,14 @@ def set_3d_properties(self, zs, zdir, axlim_clip=False):
         # Force the collection to initialize the face and edgecolors
         # just in case it is a scalarmappable with a colormap.
         self.update_scalarmappable()
-        offsets = self.get_offsets()
+        offsets = super().get_offsets()
         if len(offsets) > 0:
             xs, ys = offsets.T
         else:
             xs = []
             ys = []
         self._zdir = zdir
-        self._offsets3d = juggle_axes(xs, ys, np.atleast_1d(zs), zdir)
+        self.set_offsets3d(np.ma.column_stack((xs, ys, np.atleast_1d(zs))), zdir)
         # In the base draw methods we access the attributes directly which
         # means we cannot resolve the shuffling in the getter methods like
         # we do for the edge and face colors.
@@ -960,7 +985,6 @@ def set_3d_properties(self, zs, zdir, axlim_clip=False):
         # Grab the current sizes and linewidths to preserve them.
         self._sizes3d = self._sizes
         self._linewidths3d = np.array(self._linewidths)
-        xs, ys, zs = self._offsets3d
 
         # Sort the points based on z coordinates
         # Performance optimization: Create a sorted index array and reorder
@@ -1009,17 +1033,18 @@ def set_depthshade(
         self.stale = True
 
     def do_3d_projection(self):
+        offsets3d = self.get_offsets3d()
         mask = False
-        for xyz in self._offsets3d:
+        for xyz in offsets3d:
             if np.ma.isMA(xyz):
                 mask = mask | xyz.mask
         if self._axlim_clip:
-            mask = mask | _viewlim_mask(*self._offsets3d, self.axes)
+            mask = mask | _viewlim_mask(*offsets3d, self.axes)
             mask = np.broadcast_to(mask,
-                                   (len(self._offsets3d), *self._offsets3d[0].shape))
-            xyzs = np.ma.array(self._offsets3d, mask=mask)
+                                   (len(offsets3d), *offsets3d[0].shape))
+            xyzs = np.ma.array(offsets3d, mask=mask)
         else:
-            xyzs = self._offsets3d
+            xyzs = offsets3d
         vxs, vys, vzs, vis = proj3d._proj_transform_clip(*xyzs,
                                                          self.axes.M,
                                                          self.axes._focal_length)
@@ -1096,6 +1121,31 @@ def get_edgecolor(self):
             return self.get_facecolor()
         return self._maybe_depth_shade_and_sort_colors(super().get_edgecolor())
 
+    def set_offsets3d(self, offsets, zdir='z'):
+        """
+        Set the 3d offsets for the collection.
+
+        Parameters
+        ----------
+        offsets : (N, 3) or (3,) array-like
+            The offsets to be set.
+        zdir : {'x', 'y', 'z'}, default: 'z'
+            The axis in which to place the offsets.
+            See `.get_dir_vector` for a description of the values.
+        """
+        return _set_offsets3d(self, offsets, zdir)
+
+    def get_offsets3d(self):
+        """
+        Return the 3d offsets for the collection.
+
+        Returns
+        -------
+        offsets : (N, 3) array
+            The offsets for the collection.
+        """
+        return _get_offsets3d(self)
+
 
 def patch_collection_2d_to_3d(
     col,
@@ -1324,6 +1374,7 @@ def set_3d_properties(self, axlim_clip=False):
         self._facecolor3d = PolyCollection.get_facecolor(self)
         self._edgecolor3d = PolyCollection.get_edgecolor(self)
         self._alpha3d = PolyCollection.get_alpha(self)
+        self._offsets3d = None
         self.stale = True
 
     def set_sort_zpos(self, val):
@@ -1349,18 +1400,23 @@ def do_3d_projection(self):
             if self._edge_is_mapped:
                 self._edgecolor3d = self._edgecolors
 
+        faces = self._faces
+        if getattr(self, '_offsets3d', None) is not None:
+            offsets_arr = np.column_stack(self._offsets3d)
+            faces = faces + offsets_arr[:, np.newaxis, :]  # broadcast per-face
+
         needs_masking = np.any(self._invalid_vertices)
-        num_faces = len(self._faces)
+        num_faces = len(faces)
         mask = self._invalid_vertices
 
         # Some faces might contain masked vertices, so we want to ignore any
         # errors that those might cause
         with np.errstate(invalid='ignore', divide='ignore'):
-            pfaces = proj3d._proj_transform_vectors(self._faces, self.axes.M)
+            pfaces = proj3d._proj_transform_vectors(faces, self.axes.M)
 
         if self._axlim_clip:
-            viewlim_mask = _viewlim_mask(self._faces[..., 0], self._faces[..., 1],
-                                         self._faces[..., 2], self.axes)
+            viewlim_mask = _viewlim_mask(faces[..., 0], faces[..., 1],
+                                         faces[..., 2], self.axes)
             if np.any(viewlim_mask):
                 needs_masking = True
                 mask = mask | viewlim_mask
@@ -1471,6 +1527,31 @@ def get_edgecolor(self):
             self.do_3d_projection()
         return np.asarray(self._edgecolors2d)
 
+    def set_offsets3d(self, offsets, zdir='z'):
+        """
+        Set the 3d offsets for the collection.
+
+        Parameters
+        ----------
+        offsets : (N, 3) or (3,) array-like
+            The offsets to be set.
+        zdir : {'x', 'y', 'z'}, default: 'z'
+            The axis in which to place the offsets.
+            See `.get_dir_vector` for a description of the values.
+        """
+        return _set_offsets3d(self, offsets, zdir)
+
+    def get_offsets3d(self):
+        """
+        Return the 3d offsets for the collection.
+
+        Returns
+        -------
+        offsets : (N, 3) array
+            The offsets for the collection.
+        """
+        return _get_offsets3d(self)
+
 
 def poly_collection_2d_to_3d(col, zs=0, zdir='z', axlim_clip=False):
     """
@@ -1526,6 +1607,49 @@ def rotate_axes(xs, ys, zs, zdir):
         return xs, ys, zs
 
 
+def _set_offsets3d(col_3d, offsets, zdir='z'):
+    """
+    Set the 3d offsets for the collection.
+
+    Note: Since 3D collections have no common 3D base class, this function
+    factors out the common code for `set_offsets3d` methods of different 3D
+    collections.
+
+    Parameters
+    ----------
+    offsets : (N, 3) or (3,) array-like
+        The offsets to be set.
+    zdir : {'x', 'y', 'z'}, default: 'z'
+        The axis in which to place the offsets.
+        See `.get_dir_vector` for a description of the values.
+    """
+    offsets = np.asanyarray(offsets)
+    if offsets.shape == (3,):  # Broadcast (3,) -> (1, 3) but nothing else.
+        offsets = offsets[None, :]
+    xs = np.asanyarray(col_3d.convert_xunits(offsets[:, 0]), float)
+    ys = np.asanyarray(col_3d.convert_yunits(offsets[:, 1]), float)
+    zs = np.asanyarray(col_3d.convert_yunits(offsets[:, 2]), float)
+    col_3d._offsets3d = juggle_axes(xs, ys, np.atleast_1d(zs), zdir)
+    col_3d.stale = True
+
+
+def _get_offsets3d(col3d):
+    """
+    Return the offsets for the collection.
+
+    Note: Since 3D collections have no common 3D base class, this function
+    factors out the common code for `get_offsets3d` methods of different 3D
+    collections.
+
+    Usage pattern::
+
+        def get_offsets3d(self):
+            return _get_offsets3d(self)
+
+    """
+    return col3d._offsets3d
+
+
 def _zalpha(
     colors,
     zs,