From 4fc3073f63644fb6e11b96d79d6cf7ef2eab0631 Mon Sep 17 00:00:00 2001 From: Scott Shambaugh Date: Thu, 21 Dec 2023 14:56:29 -0700 Subject: [PATCH 1/9] get_offsets3d and set_offsets3d --- .../3d_collections_offset3d.rst | 8 ++ lib/mpl_toolkits/mplot3d/art3d.py | 128 ++++++++++++++++-- 2 files changed, 122 insertions(+), 14 deletions(-) create mode 100644 doc/users/next_whats_new/3d_collections_offset3d.rst diff --git a/doc/users/next_whats_new/3d_collections_offset3d.rst b/doc/users/next_whats_new/3d_collections_offset3d.rst new file mode 100644 index 000000000000..af63d0bdaf55 --- /dev/null +++ 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 index d06d157db4ce..a821f43ed9cd 100644 --- 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() + 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._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,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 _get_data_scale(X, Y, Z): """ @@ -940,14 +964,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 +984,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 +1032,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 +1120,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, @@ -1471,6 +1519,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', axlim_clip=False): """ @@ -1526,6 +1598,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, From 9abbdb25776d47ee2113cc632b82a30f8dbf00d3 Mon Sep 17 00:00:00 2001 From: Scott Shambaugh Date: Thu, 21 Dec 2023 18:17:25 -0700 Subject: [PATCH 2/9] Whats new --- doc/users/next_whats_new/3d_collections_offset3d.rst | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/doc/users/next_whats_new/3d_collections_offset3d.rst b/doc/users/next_whats_new/3d_collections_offset3d.rst index af63d0bdaf55..cbbf36b24f26 100644 --- a/doc/users/next_whats_new/3d_collections_offset3d.rst +++ b/doc/users/next_whats_new/3d_collections_offset3d.rst @@ -1,8 +1,8 @@ -3D Collections have `set_offset3d` and `get_offset3d` methods -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +3D Collections have ``set_offset3d`` and ``get_offset3d`` methods +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -All 3D Collections (`Patch3DCollection`, `Path3DCollection`, -`Poly3DCollection`) now 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. From fdd99e81bc36cb2c55514855c19a3d0c00c7ad06 Mon Sep 17 00:00:00 2001 From: Tim Hoffmann <2836374+timhoffm@users.noreply.github.com> Date: Sun, 10 Mar 2024 21:33:02 +0100 Subject: [PATCH 3/9] Style fixes / improvements --- lib/mpl_toolkits/mplot3d/art3d.py | 27 +++++++++++++++------------ 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py index a821f43ed9cd..d80271b7332c 100644 --- a/lib/mpl_toolkits/mplot3d/art3d.py +++ b/lib/mpl_toolkits/mplot3d/art3d.py @@ -784,9 +784,9 @@ 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 = super().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 = [] @@ -856,7 +856,8 @@ def set_offsets3d(self, offsets, zdir='z'): return _set_offsets3d(self, offsets, zdir) def get_offsets3d(self): - """Return the 3d offsets for the collection. + """ + Return the 3d offsets for the collection. Returns ------- @@ -1128,14 +1129,15 @@ def set_offsets3d(self, offsets, zdir='z'): ---------- 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'. + 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. + """ + Return the 3d offsets for the collection. Returns ------- @@ -1527,14 +1529,15 @@ def set_offsets3d(self, offsets, zdir='z'): ---------- 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'. + 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. + """ + Return the 3d offsets for the collection. Returns ------- @@ -1606,8 +1609,8 @@ def _set_offsets3d(col_3d, offsets, zdir='z'): ---------- 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'. + 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) From cbdcebbece23f305e8571126b203d183d28ecb76 Mon Sep 17 00:00:00 2001 From: Scott Shambaugh <14363975+scottshambaugh@users.noreply.github.com> Date: Mon, 11 Mar 2024 11:21:23 -0600 Subject: [PATCH 4/9] Apply suggestions from code review Co-authored-by: Tim Hoffmann <2836374+timhoffm@users.noreply.github.com> --- lib/mpl_toolkits/mplot3d/art3d.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py index d80271b7332c..176583dc46ec 100644 --- a/lib/mpl_toolkits/mplot3d/art3d.py +++ b/lib/mpl_toolkits/mplot3d/art3d.py @@ -1605,6 +1605,11 @@ 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_offests3d()`` methods. + Parameters ---------- offsets : (N, 3) or (3,) array-like @@ -1624,7 +1629,20 @@ def _set_offsets3d(col_3d, offsets, zdir='z'): def _get_offsets3d(col3d): - """Return the offsets for the collection.""" + """ + 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_offests3d()`` methods. + + Usage pattern:: + + def get_offsets3d(self): + return _get_offsets3d(self) + + """ # Default to zeros in the no-offset (None) case return np.zeros((1, 3)) if col3d._offsets3d is None else col3d._offsets3d From 58327131dcc002ddb4ea8eeff35c076e94e837b1 Mon Sep 17 00:00:00 2001 From: Scott Shambaugh <14363975+scottshambaugh@users.noreply.github.com> Date: Mon, 11 Mar 2024 11:27:38 -0600 Subject: [PATCH 5/9] Code review modifications --- lib/mpl_toolkits/mplot3d/art3d.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py index 176583dc46ec..6eb0765f770b 100644 --- a/lib/mpl_toolkits/mplot3d/art3d.py +++ b/lib/mpl_toolkits/mplot3d/art3d.py @@ -1606,7 +1606,7 @@ 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_offests3d()`` methods. @@ -1631,20 +1631,19 @@ def _set_offsets3d(col_3d, offsets, zdir='z'): 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_offests3d()`` methods. - + Usage pattern:: - + def get_offsets3d(self): return _get_offsets3d(self) """ - # Default to zeros in the no-offset (None) case - return np.zeros((1, 3)) if col3d._offsets3d is None else col3d._offsets3d + return col3d._offsets3d def _zalpha( From 71397f8a0d73d1b2d4fbef0b2697814b4c68f153 Mon Sep 17 00:00:00 2001 From: Scott Shambaugh <14363975+scottshambaugh@users.noreply.github.com> Date: Mon, 11 Mar 2024 14:51:24 -0600 Subject: [PATCH 6/9] Fix docs typos --- lib/mpl_toolkits/mplot3d/art3d.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py index 6eb0765f770b..07a01a27849a 100644 --- a/lib/mpl_toolkits/mplot3d/art3d.py +++ b/lib/mpl_toolkits/mplot3d/art3d.py @@ -1608,7 +1608,7 @@ def _set_offsets3d(col_3d, offsets, zdir='z'): .. note:: Since 3D collections have no common 3D base class, this function - factors out the common code for ``set_offests3d()`` methods. + factors out the common code for ``set_offsets3d()`` methods. Parameters ---------- @@ -1635,7 +1635,7 @@ def _get_offsets3d(col3d): .. note:: Since 3D collections have no common 3D base class, this function - factors out the common code for ``get_offests3d()`` methods. + factors out the common code for ``get_offsets3d()`` methods. Usage pattern:: From a2bec12ea7ccb960b33bd1d6a7ef2618eff10c59 Mon Sep 17 00:00:00 2001 From: Scott Shambaugh <14363975+scottshambaugh@users.noreply.github.com> Date: Mon, 11 Mar 2024 15:22:26 -0600 Subject: [PATCH 7/9] Fix docs build? --- lib/mpl_toolkits/mplot3d/art3d.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py index 07a01a27849a..eecee5b24e7a 100644 --- a/lib/mpl_toolkits/mplot3d/art3d.py +++ b/lib/mpl_toolkits/mplot3d/art3d.py @@ -1605,10 +1605,8 @@ 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. + Note: Since 3D collections have no common 3D base class, this function + factors out the common code for ``set_offsets3d()`` methods. Parameters ---------- @@ -1632,15 +1630,13 @@ 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. + Note: Since 3D collections have no common 3D base class, this function + factors out the common code for ``get_offsets3d()`` methods. - Usage pattern:: + Usage pattern:: - def get_offsets3d(self): - return _get_offsets3d(self) + def get_offsets3d(self): + return _get_offsets3d(self) """ return col3d._offsets3d From fc034a70b470d15974678bdb8584de2f7dadb2f6 Mon Sep 17 00:00:00 2001 From: Scott Shambaugh <14363975+scottshambaugh@users.noreply.github.com> Date: Mon, 11 Mar 2024 16:32:03 -0600 Subject: [PATCH 8/9] Fix docs build? --- lib/mpl_toolkits/mplot3d/art3d.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py index eecee5b24e7a..d007ad10a98a 100644 --- a/lib/mpl_toolkits/mplot3d/art3d.py +++ b/lib/mpl_toolkits/mplot3d/art3d.py @@ -1606,7 +1606,8 @@ 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. + factors out the common code for `set_offsets3d` methods of different 3D + collections. Parameters ---------- @@ -1631,7 +1632,8 @@ 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. + factors out the common code for `get_offsets3d` methods of different 3D + collections. Usage pattern:: From dcdcb8266b2aaf26b33af84ade30f02c8b701495 Mon Sep 17 00:00:00 2001 From: Scott Shambaugh Date: Fri, 6 Mar 2026 16:01:05 -0700 Subject: [PATCH 9/9] Poly3DCollection handling --- .../3d_collections_offset3d.rst | 15 +- lib/mpl_toolkits/mplot3d/art3d.py | 14 +- lib/mpl_toolkits/mplot3d/tests/test_art3d.py | 153 ++++++++++++++++++ 3 files changed, 172 insertions(+), 10 deletions(-) diff --git a/doc/users/next_whats_new/3d_collections_offset3d.rst b/doc/users/next_whats_new/3d_collections_offset3d.rst index cbbf36b24f26..a9e732c12bc3 100644 --- a/doc/users/next_whats_new/3d_collections_offset3d.rst +++ b/doc/users/next_whats_new/3d_collections_offset3d.rst @@ -1,8 +1,11 @@ -3D Collections have ``set_offset3d`` and ``get_offset3d`` methods -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +3D Collections have ``set_offsets3d`` and ``get_offsets3d`` 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. +``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 index d007ad10a98a..462c2d3913b4 100644 --- a/lib/mpl_toolkits/mplot3d/art3d.py +++ b/lib/mpl_toolkits/mplot3d/art3d.py @@ -1374,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): @@ -1399,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 diff --git a/lib/mpl_toolkits/mplot3d/tests/test_art3d.py b/lib/mpl_toolkits/mplot3d/tests/test_art3d.py index aca943f9e0c0..f6083a01f78b 100644 --- a/lib/mpl_toolkits/mplot3d/tests/test_art3d.py +++ b/lib/mpl_toolkits/mplot3d/tests/test_art3d.py @@ -3,8 +3,10 @@ import pytest import matplotlib.pyplot as plt +from matplotlib.testing.decorators import check_figures_equal from matplotlib.backend_bases import MouseEvent +from mpl_toolkits.mplot3d import art3d from mpl_toolkits.mplot3d.art3d import ( get_dir_vector, Line3DCollection, @@ -117,3 +119,154 @@ def test_generate_normals(): ax = fig.add_subplot(projection='3d') ax.add_collection3d(shape) plt.draw() + + +# --- set_offsets3d / get_offsets3d tests --- + +def test_offsets3d_patch3d_roundtrip(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + col = ax.scatter([0, 1], [0, 1], [0, 1]) + offsets = np.array([[1, 2, 3], [4, 5, 6]], dtype=float) + col.set_offsets3d(offsets) + xs, ys, zs = col.get_offsets3d() + nptest.assert_array_almost_equal(xs, offsets[:, 0]) + nptest.assert_array_almost_equal(ys, offsets[:, 1]) + nptest.assert_array_almost_equal(zs, offsets[:, 2]) + + +def test_offsets3d_poly3d_roundtrip(): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + verts = [[[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0]]] + poly = Poly3DCollection(verts) + ax.add_collection3d(poly) + offsets = np.array([[1, 2, 3]], dtype=float) + poly.set_offsets3d(offsets) + xs, ys, zs = poly.get_offsets3d() + nptest.assert_array_almost_equal(xs, offsets[:, 0]) + nptest.assert_array_almost_equal(ys, offsets[:, 1]) + nptest.assert_array_almost_equal(zs, offsets[:, 2]) + + +def test_offsets3d_single_offset_broadcast(): + """A (3,) offset should be broadcast to (1, 3).""" + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + col = ax.scatter([0], [0], [0]) + col.set_offsets3d([1, 2, 3]) + xs, ys, zs = col.get_offsets3d() + nptest.assert_array_almost_equal(xs, [1]) + nptest.assert_array_almost_equal(ys, [2]) + nptest.assert_array_almost_equal(zs, [3]) + + +@pytest.mark.parametrize("zdir", ['x', 'y', 'z']) +def test_offsets3d_zdir_roundtrip(zdir): + fig = plt.figure() + ax = fig.add_subplot(projection='3d') + col = ax.scatter([0], [0], [0]) + col.set_offsets3d([[1, 2, 3]], zdir=zdir) + xs, ys, zs = col.get_offsets3d() + # The offsets are juggled according to zdir, so just verify + # the values are preserved (in some axis order) + result = sorted([xs[0], ys[0], zs[0]]) + assert result == [1.0, 2.0, 3.0] + + +@check_figures_equal() +def test_offsets3d_scatter_set(fig_test, fig_ref): + """set_offsets3d moves scatter points to new positions.""" + # Reference: scatter at target positions directly + ax_ref = fig_ref.add_subplot(projection='3d') + ax_ref.scatter([1, 2], [3, 4], [5, 6]) + ax_ref.set_xlim(0, 3) + ax_ref.set_ylim(2, 5) + ax_ref.set_zlim(4, 7) + + # Test: scatter at origin, then move with set_offsets3d + ax_test = fig_test.add_subplot(projection='3d') + col = ax_test.scatter([0, 0], [0, 0], [0, 0]) + col.set_offsets3d([[1, 3, 5], [2, 4, 6]]) + ax_test.set_xlim(0, 3) + ax_test.set_ylim(2, 5) + ax_test.set_zlim(4, 7) + + +@check_figures_equal() +def test_offsets3d_poly3d_translate(fig_test, fig_ref): + """Poly3DCollection: set_offsets3d translates polygon faces.""" + square = np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0]], + dtype=float) + offset = np.array([2, 3, 4], dtype=float) + + # Reference: polygon placed at target position directly + ax_ref = fig_ref.add_subplot(projection='3d') + poly_ref = Poly3DCollection([square + offset], alpha=0.5) + ax_ref.add_collection3d(poly_ref) + ax_ref.set_xlim(0, 5) + ax_ref.set_ylim(0, 5) + ax_ref.set_zlim(0, 5) + + # Test: polygon at origin, moved with set_offsets3d + ax_test = fig_test.add_subplot(projection='3d') + poly_test = Poly3DCollection([square], alpha=0.5) + ax_test.add_collection3d(poly_test) + poly_test.set_offsets3d([offset]) + ax_test.set_xlim(0, 5) + ax_test.set_ylim(0, 5) + ax_test.set_zlim(0, 5) + + +@check_figures_equal() +def test_offsets3d_poly3d_multiple(fig_test, fig_ref): + """Poly3DCollection: N offsets translate N faces independently.""" + square = np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0]], + dtype=float) + offsets = np.array([[0, 0, 0], [3, 3, 3]], dtype=float) + + # Reference: two polygons placed at target positions directly + ax_ref = fig_ref.add_subplot(projection='3d') + poly_ref = Poly3DCollection([square + offsets[0], square + offsets[1]], + alpha=0.5) + ax_ref.add_collection3d(poly_ref) + ax_ref.set_xlim(-1, 5) + ax_ref.set_ylim(-1, 5) + ax_ref.set_zlim(-1, 5) + + # Test: two polygons at origin, moved with set_offsets3d + ax_test = fig_test.add_subplot(projection='3d') + poly_test = Poly3DCollection([square, square], alpha=0.5) + ax_test.add_collection3d(poly_test) + poly_test.set_offsets3d(offsets) + ax_test.set_xlim(-1, 5) + ax_test.set_ylim(-1, 5) + ax_test.set_zlim(-1, 5) + + +@check_figures_equal() +def test_offsets3d_poly3d_broadcast_single(fig_test, fig_ref): + """Poly3DCollection: 1 offset broadcasts to all N faces.""" + square = np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0]], + dtype=float) + offset = np.array([2, 2, 2], dtype=float) + + # Reference: two polygons shifted by same offset directly + ax_ref = fig_ref.add_subplot(projection='3d') + poly_ref = Poly3DCollection( + [square + offset, square + np.array([0, 3, 0]) + offset], + alpha=0.5) + ax_ref.add_collection3d(poly_ref) + ax_ref.set_xlim(0, 5) + ax_ref.set_ylim(0, 6) + ax_ref.set_zlim(0, 5) + + # Test: two polygons, single offset broadcast to both + ax_test = fig_test.add_subplot(projection='3d') + poly_test = Poly3DCollection( + [square, square + np.array([0, 3, 0])], alpha=0.5) + ax_test.add_collection3d(poly_test) + poly_test.set_offsets3d([offset]) + ax_test.set_xlim(0, 5) + ax_test.set_ylim(0, 6) + ax_test.set_zlim(0, 5)