Autoscale Line3DCollection and Poly3DCollection

scottshambaugh · scottshambaugh · commit 195a208dc1e1 · 2024-06-15T21:45:52.000-06:00
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -2614,7 +2614,7 @@ def tricontourf(self, *args, zdir='z', offset=None, **kwargs):
         self._auto_scale_contourf(X, Y, Z, zdir, levels, had_data)
         return cset
 
-    def add_collection3d(self, col, zs=0, zdir='z'):
+    def add_collection3d(self, col, zs=0, zdir='z', autolim=True):
         """
         Add a 3D collection object to the plot.
 
@@ -2627,7 +2627,20 @@ def add_collection3d(self, col, zs=0, zdir='z'):
         - `.PolyCollection`
         - `.LineCollection`
         - `.PatchCollection`
+
+        Parameters
+        ----------
+        col : `.Collection`
+            A 2D collection object.
+        zs : float or array-like, default: 0
+            The z-positions to be used for the 2D objects.
+        zdir : {'x', 'y', 'z'}, default: 'z'
+            The direction to use for the z-positions.
+        autolim : bool, default: True
+            Whether to update the data limits.
         """
+        had_data = self.has_data()
+
         zvals = np.atleast_1d(zs)
         zsortval = (np.min(zvals) if zvals.size
                     else 0)  # FIXME: arbitrary default
@@ -2645,6 +2658,17 @@ def add_collection3d(self, col, zs=0, zdir='z'):
             art3d.patch_collection_2d_to_3d(col, zs=zs, zdir=zdir)
             col.set_sort_zpos(zsortval)
 
+        # FIXME: Implement auto-scaling function for Patch3DCollection
+        # Currently unable to do so due to issues with Patch3DCollection
+        # See https://github.com/matplotlib/matplotlib/issues/14298 for details
+        if autolim:
+            if isinstance(col, art3d.Line3DCollection):
+                self.auto_scale_xyz(*np.array(col._segments3d).transpose(),
+                                    had_data=had_data)
+
+            if isinstance(col, art3d.Poly3DCollection):
+                self.auto_scale_xyz(*col._vec[:-1], had_data=had_data)
+
         collection = super().add_collection(col)
         return collection
 
diff --git a/lib/mpl_toolkits/mplot3d/tests/baseline_images/test_axes3d/voxels-named-colors.png b/lib/mpl_toolkits/mplot3d/tests/baseline_images/test_axes3d/voxels-named-colors.png
diff --git a/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py b/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py
@@ -962,8 +962,8 @@ def test_poly3dcollection_closed():
                                 facecolor=(0.5, 0.5, 1, 0.5), closed=True)
     c2 = art3d.Poly3DCollection([poly2], linewidths=3, edgecolor='k',
                                 facecolor=(1, 0.5, 0.5, 0.5), closed=False)
-    ax.add_collection3d(c1)
-    ax.add_collection3d(c2)
+    ax.add_collection3d(c1, autolim=False)
+    ax.add_collection3d(c2, autolim=False)
 
 
 def test_poly_collection_2d_to_3d_empty():
@@ -996,8 +996,8 @@ def test_poly3dcollection_alpha():
     c2.set_facecolor((1, 0.5, 0.5))
     c2.set_edgecolor('k')
     c2.set_alpha(0.5)
-    ax.add_collection3d(c1)
-    ax.add_collection3d(c2)
+    ax.add_collection3d(c1, autolim=False)
+    ax.add_collection3d(c2, autolim=False)
 
 
 @mpl3d_image_comparison(['add_collection3d_zs_array.png'], style='mpl20')
@@ -1056,6 +1056,32 @@ def test_add_collection3d_zs_scalar():
     ax.set_zlim(0, 2)
 
 
+def test_line3dCollection_autoscaling():
+    fig = plt.figure()
+    ax = fig.add_subplot(projection='3d')
+
+    lines = [[(0, 0, 0), (1, 4, 2)],
+             [(1, 1, 3), (2, 0, 2)],
+             [(1, 0, 4), (1, 4, 5)]]
+
+    lc = art3d.Line3DCollection(lines)
+    ax.add_collection3d(lc)
+    assert np.allclose(ax.get_xlim3d(), (-0.041666666666666664, 2.0416666666666665))
+    assert np.allclose(ax.get_ylim3d(), (-0.08333333333333333, 4.083333333333333))
+    assert np.allclose(ax.get_zlim3d(), (-0.10416666666666666, 5.104166666666667))
+
+
+def test_poly3dCollection_autoscaling():
+    fig = plt.figure()
+    ax = fig.add_subplot(projection='3d')
+    poly = np.array([[0, 0, 0], [1, 1, 3], [1, 0, 4]])
+    col = art3d.Poly3DCollection([poly])
+    ax.add_collection3d(col)
+    assert np.allclose(ax.get_xlim3d(), (-0.020833333333333332, 1.0208333333333333))
+    assert np.allclose(ax.get_ylim3d(), (-0.020833333333333332, 1.0208333333333333))
+    assert np.allclose(ax.get_zlim3d(), (-0.0833333333333333, 4.083333333333333))
+
+
 @mpl3d_image_comparison(['axes3d_labelpad.png'],
                         remove_text=False, style='mpl20')
 def test_axes3d_labelpad():
