get_offsets3d and set_offsets3d

scottshambaugh · scottshambaugh · commit bbd244bd0949 · 2023-12-21T14:56:29.000-07:00
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,8 @@
+3D Collections have `set_offset3d` and `get_offset3d` methods
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+All 3D Collections (`Patch3DCollection`, `Path3DCollection`,
+`Poly3DCollection`) now have `set_offset3d` and `get_offset3d` methods
+which allow you to set and get the offset of the collection in data
+coordinates. In other words, this allows you to set and get the position of the
+of the collection points.
diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
@@ -598,19 +598,20 @@ 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.
         self.update_scalarmappable()
-        offsets = self.get_offsets()
+        offsets = super().get_offsets()
         if len(offsets) > 0:
             xs, ys = offsets.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.stale = True
 
     def do_3d_projection(self):
-        xs, ys, zs = self._offsets3d
+        xs, ys, zs = self.get_offsets3d()
         vxs, vys, vzs, vis = proj3d.proj_transform_clip(xs, ys, zs,
                                                         self.axes.M)
         self._vzs = vzs
@@ -642,6 +643,29 @@ 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)
 
 class Path3DCollection(PathCollection):
     """
@@ -696,13 +720,13 @@ 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.
         self.update_scalarmappable()
-        offsets = self.get_offsets()
+        offsets = super().get_offsets()
         if len(offsets) > 0:
             xs, ys = offsets.T
         else:
             xs = []
             ys = []
-        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.
@@ -715,7 +739,6 @@ def set_3d_properties(self, zs, zdir):
         # 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
@@ -751,7 +774,7 @@ def set_depthshade(self, depthshade):
         self.stale = True
 
     def do_3d_projection(self):
-        xs, ys, zs = self._offsets3d
+        xs, ys, zs = self.get_offsets3d()
         vxs, vys, vzs, vis = proj3d.proj_transform_clip(xs, ys, zs,
                                                         self.axes.M)
         # Sort the points based on z coordinates
@@ -818,6 +841,30 @@ 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 patch_collection_2d_to_3d(col, zs=0, zdir='z', depthshade=True):
     """
@@ -1113,6 +1160,30 @@ 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'}
+            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 poly_collection_2d_to_3d(col, zs=0, zdir='z'):
     """
@@ -1167,6 +1238,34 @@ 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.
+
+    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.
+    """
+    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."""
+    # Default to zeros in the no-offset (None) case
+    return np.zeros((1, 3)) if col3d._offsets3d is None else col3d._offsets3d
+
+
 def _zalpha(colors, zs):
     """Modify the alphas of the color list according to depth."""
     # FIXME: This only works well if the points for *zs* are well-spaced
