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..a9e732c12bc3 --- /dev/null +++ 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 index d06d157db4ce..462c2d3913b4 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() - 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, 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)