fill_between extended to 3D

scottshambaugh · scottshambaugh · commit 6a130374c4d8 · 2024-05-14T10:42:50.000-06:00
diff --git a/doc/users/next_whats_new/fill_between_3d.rst b/doc/users/next_whats_new/fill_between_3d.rst
@@ -0,0 +1,26 @@
+Fill between 3D lines
+---------------------
+
+The `.Axes.fill_between` method has been extended to 3D, allowing to fill the
+surface between two 3D lines with polygons through the `.Axes3D.fill_between`
+function.
+
+.. plot::
+    :include-source:
+    :alt: Example of fill_between
+
+    N = 50
+    theta = np.linspace(0, 2*np.pi, N)
+
+    xs1 = np.cos(theta)
+    ys1 = np.sin(theta)
+    zs1 = 0.1 * np.sin(6 * theta)
+
+    xs2 = 0.6 * np.cos(theta)
+    ys2 = 0.6 * np.sin(theta)
+    zs2 = 2  # Note that scalar values work in addition to length N arrays
+
+    fig = plt.figure()
+    ax = fig.add_subplot(111, projection='3d')
+    ax.fill_between(xs1, ys1, zs1, xs2, ys2, zs2,
+                    alpha=0.5, edgecolor='k')
diff --git a/galleries/examples/mplot3d/fillbetween3d.py b/galleries/examples/mplot3d/fillbetween3d.py
@@ -0,0 +1,34 @@
+"""
+========================
+Fill under 3D line graph
+========================
+
+Demonstrate how to create polygons which fill the space under a line
+graph. In this example polygons are semi-transparent, creating a sort
+of 'jagged stained glass' effect.
+"""
+
+import math
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+gamma = np.vectorize(math.gamma)
+N = 31
+x = np.linspace(0., 10., N)
+lambdas = range(1, 9)
+
+ax = plt.figure().add_subplot(projection='3d')
+
+facecolors = plt.colormaps['viridis_r'](np.linspace(0, 1, len(lambdas)))
+
+for i, l in enumerate(lambdas):
+    # Note fill_between can take coordinates as length N vectors, or scalars
+    ax.fill_between(x, l, l**x * np.exp(-l) / gamma(x + 1),
+                    x, l, 0,
+                    facecolors=facecolors[i], alpha=.7)
+
+ax.set(xlim=(0, 10), ylim=(1, 9), zlim=(0, 0.35),
+       xlabel='x', ylabel=r'$\lambda$', zlabel='probability')
+
+plt.show()
diff --git a/galleries/examples/mplot3d/polys3d.py b/galleries/examples/mplot3d/polys3d.py
@@ -1,47 +1,36 @@
 """
-=============================================
-Generate polygons to fill under 3D line graph
-=============================================
+====================
+Generate 3D polygons
+====================
 
-Demonstrate how to create polygons which fill the space under a line
-graph. In this example polygons are semi-transparent, creating a sort
-of 'jagged stained glass' effect.
+Demonstrate how to create polygons in 3D. Here we stack 3 hexagons, somewhat
+like a traffic light.
 """
 
-import math
-
 import matplotlib.pyplot as plt
 import numpy as np
 
-from matplotlib.collections import PolyCollection
-
-# Fixing random state for reproducibility
-np.random.seed(19680801)
+from mpl_toolkits.mplot3d.art3d import Poly3DCollection
 
+# Coordinates of a hexagon
+angles = np.linspace(0, 2 * np.pi, 6, endpoint=False)
+x = np.cos(angles)
+y = np.sin(angles)
+zs = [0, 1, 2]
 
-def polygon_under_graph(x, y):
-    """
-    Construct the vertex list which defines the polygon filling the space under
-    the (x, y) line graph. This assumes x is in ascending order.
-    """
-    return [(x[0], 0.), *zip(x, y), (x[-1], 0.)]
+# Close the hexagon by repeating the first vertex
+x = np.append(x, x[0])
+y = np.append(y, y[0])
 
+verts = []
+for z in zs:
+    verts.append(list(zip(x, y, np.full_like(x, z))))
 
 ax = plt.figure().add_subplot(projection='3d')
 
-x = np.linspace(0., 10., 31)
-lambdas = range(1, 9)
-
-# verts[i] is a list of (x, y) pairs defining polygon i.
-gamma = np.vectorize(math.gamma)
-verts = [polygon_under_graph(x, l**x * np.exp(-l) / gamma(x + 1))
-         for l in lambdas]
-facecolors = plt.colormaps['viridis_r'](np.linspace(0, 1, len(verts)))
-
-poly = PolyCollection(verts, facecolors=facecolors, alpha=.7)
-ax.add_collection3d(poly, zs=lambdas, zdir='y')
-
-ax.set(xlim=(0, 10), ylim=(1, 9), zlim=(0, 0.35),
-       xlabel='x', ylabel=r'$\lambda$', zlabel='probability')
+poly = Poly3DCollection(verts, facecolors=['g', 'y', 'r'], alpha=.7)
+ax.add_collection3d(poly)
+ax.auto_scale_xyz(x, y, zs)
+ax.set_aspect('equal')
 
 plt.show()
diff --git a/galleries/users_explain/toolkits/mplot3d.rst b/galleries/users_explain/toolkits/mplot3d.rst
@@ -111,6 +111,16 @@ See `.Axes3D.contourf` for API documentation.
    The feature demoed in the second contourf3d example was enabled as a
    result of a bugfix for version 1.1.0.
 
+.. _fillbetween3d:
+
+Fill between plots
+==================
+See `.Axes3D.fill_between` for API documentation.
+
+.. figure:: /gallery/mplot3d/images/sphx_glr_fillbetween3d_001.png
+   :target: /gallery/mplot3d/fillbetween3d.html
+   :align: center
+
 .. _polygon3d:
 
 Polygon plots
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -1904,6 +1904,80 @@ def plot(self, xs, ys, *args, zdir='z', **kwargs):
 
     plot3D = plot
 
+    def fill_between(self, xs1, ys1, zs1, xs2, ys2, zs2, *, where=None, **kwargs):
+        """
+        Fill the area between two 2D or 3D curves.
+
+        The curves are defined by the points (*xs1*, *ys1*, *zs1*) and
+        (*xs2*, *ys2*, *zs2*). This creates one or multiple quadrangle
+        polygons that are filled. All points must be the same length N, or a
+        single value to be used for all points.
+
+        Parameters
+        ----------
+        xs1, ys1, zs1 : float or 1D array-like
+            x, y, and z  coordinates of vertices for 1st line.
+
+        xs2, ys2, zs2 : float or 1D array-like
+            x, y, and z coordinates of vertices for 2nd line.
+
+        where : array of bool (length N), optional
+            Define *where* to exclude some regions from being filled. The
+            filled regions are defined by the coordinates ``pts[where]``,
+            for all x, y, and z pts. More precisely, fill between ``pts[i]``
+            and ``pts[i+1]`` if ``where[i] and where[i+1]``. Note that this
+            definition implies that an isolated *True* value between two
+            *False* values in *where* will not result in filling. Both sides of
+            the *True* position remain unfilled due to the adjacent *False*
+            values.
+
+        **kwargs
+            All other keyword arguments are passed on to `.Poly3DCollection`.
+
+        Returns
+        -------
+        `.Poly3DCollection`
+            A `.Poly3DCollection` containing the plotted polygons.
+
+        """
+        had_data = self.has_data()
+        xs1, ys1, zs1, xs2, ys2, zs2 = cbook._broadcast_with_masks(xs1, ys1, zs1,
+                                                                   xs2, ys2, zs2)
+
+        if where is None:
+            where = True
+        else:
+            where = np.asarray(where, dtype=bool)
+            if where.size != xs1.size:
+                raise ValueError(f"where size ({where.size}) does not match "
+                                 f"size ({xs1.size})")
+        where = where & ~np.isnan(xs1)  # NaNs were broadcast in _broadcast_with_masks
+
+        polys = []
+        for idx0, idx1 in cbook.contiguous_regions(where):
+            xs1slice = xs1[idx0:idx1]
+            ys1slice = ys1[idx0:idx1]
+            zs1slice = zs1[idx0:idx1]
+            xs2slice = xs2[idx0:idx1]
+            ys2slice = ys2[idx0:idx1]
+            zs2slice = zs2[idx0:idx1]
+
+            if not len(xs1slice):
+                continue
+
+            for i in range(len(xs1slice) - 1):
+                poly = [(xs1slice[i], ys1slice[i], zs1slice[i]),
+                        (xs1slice[i+1], ys1slice[i+1], zs1slice[i+1]),
+                        (xs2slice[i+1], ys2slice[i+1], zs2slice[i+1]),
+                        (xs2slice[i], ys2slice[i], zs2slice[i])]
+                polys.append(poly)
+
+        polyc = art3d.Poly3DCollection(polys, **kwargs)
+        self.add_collection(polyc)
+
+        self.auto_scale_xyz([xs1, xs2], [ys1, ys2], [zs1, zs2], had_data)
+        return polyc
+
     def plot_surface(self, X, Y, Z, *, norm=None, vmin=None,
                      vmax=None, lightsource=None, **kwargs):
         """
diff --git a/lib/mpl_toolkits/mplot3d/tests/baseline_images/test_axes3d/fill_between.png b/lib/mpl_toolkits/mplot3d/tests/baseline_images/test_axes3d/fill_between.png
diff --git a/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py b/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py
@@ -592,6 +592,27 @@ def test_plot_3d_from_2d():
     ax.plot(xs, ys, zs=0, zdir='y')
 
 
+@mpl3d_image_comparison(['fill_between.png'], style='mpl20')
+def test_fill_between():
+    fig = plt.figure()
+    ax = fig.add_subplot(projection='3d')
+
+    theta = np.linspace(0, 2*np.pi, 50)
+
+    xs1 = np.cos(theta)
+    ys1 = np.sin(theta)
+    zs1 = 0.1 * np.sin(6 * theta)
+
+    xs2 = 0.6 * np.cos(theta)
+    ys2 = 0.6 * np.sin(theta)
+    zs2 = 2
+
+    where = (theta < np.pi/2) | (theta > 3*np.pi/2)
+
+    ax.fill_between(xs1, ys1, zs1, xs2, ys2, zs2,
+                    where=where, alpha=0.5, edgecolor='k')
+
+
 @mpl3d_image_comparison(['surface3d.png'], style='mpl20')
 def test_surface3d():
     # Remove this line when this test image is regenerated.
